[systalyze-ai]

Description

Adds mesh_ring_axis to arrange an even-sized logical mesh axis into physical rings built from k-by-2 torus blocks. This provides a physical closing link when the axis is mapped to a slice dimension without wraparound.

The change affects device placement only. AOT HLO was byte-identical by md5 with the option disabled and enabled.

Performance

• On v6e-128 with DeepSeek-V4-Flash 284B LoRA at seq 16384, enabling the option changed step time from 4.632 to 4.138 s/step (1.12x).

Setup: ici_expert_parallelism=16, ici_tensor_parallelism=8, ici_fsdp_parallelism=1, max_target_length=16384, per_device_batch_size=0.125 (global batch 16), LIBTPU_INIT_ARGS=--xla_tpu_scoped_vmem_limit_kib=98304, dynamic-splash attention path enabled. mesh_ring_axis=tensor places the 8-way tensor axis on physical rings.

Tests

• The command python3 -m pytest tests/unit/max_utils_test.py -k RingAxisDeviceMeshTest passed 3 tests.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces the mesh_ring_axis configuration to lay out a logical mesh axis as a physical TPU ring for single-slice setups, implementing the create_ring_axis_device_mesh utility and corresponding unit tests. The review feedback highlights a potential issue in create_ring_axis_device_mesh where the hardcoded transposition check can cause a ValueError on certain valid TPU topologies (such as a 2x16 grid with ring=8), and suggests a more robust, dynamic tiling check.

[gemini-code-assist]

The current implementation of create_ring_axis_device_mesh only transposes the physical grid if physical.shape[0] > physical.shape[1]. However, this hardcoded check is too restrictive and can cause a ValueError on valid TPU topologies (such as a 2x16 grid with ring=8, where ring_height=4 and ring_width=2). In this case, the height of the block (4) is greater than the height of the grid (2), so tiling fails in the default orientation, but would succeed if the grid were transposed to 16x2.

To make the mesh creation more robust, we should dynamically check if transposing the physical grid is required to make it divisible by ring_height and ring_width.

  ring_height, ring_width = ring // 2, 2
  # Determine if transposing is needed/preferred to allow valid tiling.
  can_tile_normal = (physical.shape[0] % ring_height == 0) and (physical.shape[1] % ring_width == 0)
  can_tile_transposed = (physical.shape[1] % ring_height == 0) and (physical.shape[0] % ring_width == 0)
  if not can_tile_normal and can_tile_transposed:
    physical = physical.T
  elif can_tile_normal and can_tile_transposed and physical.shape[0] > physical.shape[1]:
    physical = physical.T

[arminghobadi]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces the mesh_ring_axis configuration option, allowing a logical mesh axis to be laid out as a physical TPU ring for single-slice setups. It implements the device mesh construction logic in max_utils.py and integrates it into maxtext_utils.py, accompanied by comprehensive unit tests. The review feedback highlights two key improvements: first, replacing the serpentine traversal of blocks with a simple nested loop to preserve the physical orthogonality and locality of other logical axes (such as FSDP); second, initializing the physical device array with None using np.full instead of np.empty to avoid potential undefined behavior from uninitialized object pointers.

[gemini-code-assist]

The serpentine traversal of blocks (for j in columns if i % 2 == 0 else reversed(columns)) reverses the column order of rings on odd rows. When the 2D grid of rings is reshaped to map to other logical axes (such as fsdp or expert), this reversal breaks the physical orthogonality and locality of those axes.\n\nSpecifically, a logical axis like fsdp will end up grouping a ring at column k on an even row with a ring at column width - 1 - k on an odd row, causing FSDP groups to span across the entire physical width of the torus. This significantly increases communication latency and reduces bandwidth for FSDP collectives.\n\nUsing a simple, non-reversing nested loop ensures that the other logical axes remain orthogonal and physically local (e.g., FSDP groups will consist of vertically adjacent blocks).

  for i in range(physical.shape[0] // ring_height):\n    for j in range(physical.shape[1] // ring_width):\n      left = list(physical[i * ring_height : (i + 1) * ring_height, ring_width * j])\n      right = list(physical[i * ring_height : (i + 1) * ring_height, ring_width * j + 1])\n      rings.append(left + right[::-1])

[gemini-code-assist]

Using np.empty with dtype=object creates an array with uninitialized object pointers (which can contain arbitrary memory garbage or NULL pointers). If the physical coordinates of the devices do not perfectly and densely cover the bounding box extent (e.g., due to a faulty or sliced node layout), accessing the unpopulated elements can lead to undefined behavior or crashes.\n\nIt is much safer and more explicit to initialize the array with None using np.full.

Suggested change

	physical = np.empty(tuple(extent), dtype=object)
	physical = np.full(tuple(extent), None, dtype=object)

[RSudarsanan]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces the mesh_ring_axis configuration option to allow laying out a logical mesh axis as a physical TPU ring. It implements create_ring_axis_device_mesh to construct the device mesh and integrates it into the device mesh creation process. Additionally, unit tests are added to verify the functionality. A review comment points out a bug in the newly added test_transposes_grid_when_needed_for_tiling test, where a mismatch between the number of devices (32) and the specified logical mesh size (128) will cause a ValueError during execution. A fix is suggested to adjust the logical mesh size to [1, 1, 8, 4].

[gemini-code-assist]

The test test_transposes_grid_when_needed_for_tiling is currently broken and will raise a ValueError during execution.

Cause of the Issue

The test initializes devices with 32 elements (a 2x16 grid), but passes [1, 1, 8, 16] as the ici_parallelism argument to create_ring_axis_device_mesh. The product of [1, 1, 8, 16] is 128.
Inside create_ring_axis_device_mesh, the code attempts to reshape the constructed mesh of size 32 into [*other_axes, ring], which evaluates to [1, 1, 16, 8] (size 128). This mismatch causes NumPy to raise a ValueError: cannot reshape array of size 32 into shape (1,1,16,8).

Solution

To align the logical mesh size with the actual number of devices (32), update the ici_parallelism argument to [1, 1, 8, 4] and assert that the resulting mesh shape is (1, 1, 8, 4).

Suggested change

	def test_transposes_grid_when_needed_for_tiling(self):
	"""Test that a 2x16 grid is transposed to 16x2 when ring=8."""
	devices = [self.FakeDevice(x * 16 + y, (x, y, 0)) for x in range(2) for y in range(16)]
	mesh = max_utils.create_ring_axis_device_mesh([1, 1, 8, 16], self.AXES, devices, "tensor")
	self.assertEqual(mesh.shape, (1, 1, 8, 16))
	self.assertEqual(sorted(d.id for d in mesh.flatten()), sorted(d.id for d in devices))
	def test_transposes_grid_when_needed_for_tiling(self):
	"""Test that a 2x16 grid is transposed to 16x2 when ring=8."""
	devices = [self.FakeDevice(x * 16 + y, (x, y, 0)) for x in range(2) for y in range(16)]
	mesh = max_utils.create_ring_axis_device_mesh([1, 1, 8, 4], self.AXES, devices, "tensor")
	self.assertEqual(mesh.shape, (1, 1, 8, 4))
	self.assertEqual(sorted(d.id for d in mesh.flatten()), sorted(d.id for d in devices))

[RSudarsanan]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces the mesh_ring_axis configuration option, allowing a logical mesh axis to be laid out as a physical TPU ring for single-slice configurations. It implements the mesh creation logic in max_utils.py and integrates it into maxtext_utils.py, accompanied by unit tests. The review feedback suggests simplifying the NumPy array initialization in create_ring_axis_device_mesh and correcting the parameters in the transposition unit test to ensure it properly triggers and validates the transposition logic.

[gemini-code-assist]

This loop to populate the mesh array can be simplified by directly converting the rings list of lists into a NumPy array. This is more idiomatic and concise.

Suggested change

	mesh = np.empty((len(rings), ring), dtype=object)
	for i, devices_in_ring in enumerate(rings):
	mesh[i] = devices_in_ring
	mesh = np.array(rings, dtype=object)

[gemini-code-assist]

This test is named test_transposes_grid_when_needed_for_tiling and its docstring mentions testing transposition for ring=8, but the current implementation uses ring=2 and doesn't actually trigger the transposition logic.

To ensure the transposition logic is covered by tests, I suggest updating the test to use parameters that force a transpose, as described in the docstring.

Suggested change

	def test_transposes_grid_when_needed_for_tiling(self):
	"""Test that a 2x16 grid is transposed to 16x2 when ring=8."""
	devices = [self.FakeDevice(x * 16 + y, (x, y, 0)) for x in range(2) for y in range(16)]
	mesh = max_utils.create_ring_axis_device_mesh([1, 1, 2, 16], self.AXES, devices, "tensor")
	self.assertEqual(mesh.shape, (1, 1, 2, 16))
	self.assertEqual(sorted(d.id for d in mesh.flatten()), sorted(d.id for d in devices))
	def test_transposes_grid_when_needed_for_tiling(self):
	"""Test that a 2x16 grid is transposed to 16x2 when ring=8."""
	devices = [self.FakeDevice(x * 16 + y, (x, y, 0)) for x in range(2) for y in range(16)]
	# A 2x16 physical grid can only be tiled by 4x2 blocks (for ring=8) if transposed.
	mesh = max_utils.create_ring_axis_device_mesh([1, 1, 8, 4], self.AXES, devices, "tensor")
	self.assertEqual(mesh.shape, (1, 1, 8, 4))
	self.assertEqual(sorted(d.id for d in mesh.flatten()), sorted(d.id for d in devices))

[RSudarsanan]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces the mesh_ring_axis configuration option to lay out a logical mesh axis as a physical TPU ring, implementing the create_ring_axis_device_mesh utility along with corresponding unit tests. The review feedback suggests several key improvements: enhancing the transpose logic in create_ring_axis_device_mesh to align the physical grid with non-trivial logical axes to prevent scrambling and reduce communication latency, adding a defensive check to ensure the physical device grid is fully populated, and adding a unit test to verify the alignment-aware transposing behavior.

[gemini-code-assist]

The current transpose logic only considers whether the physical grid can be tiled by ring_height x ring_width blocks. However, it does not consider the sizes of the other non-trivial logical axes in other_axes.

If there are exactly two non-trivial other axes, we should prefer transposing the physical grid if it aligns the resulting ring grid shape (R_H, R_W) with the target logical axes shape. Otherwise, the logical axes will be scrambled across physical dimensions, violating physical contiguity and leading to suboptimal communication latency.

  # Determine if transposing is needed/preferred to allow valid tiling and logical axis alignment.
  can_tile_normal = (physical.shape[0] % ring_height == 0) and (physical.shape[1] % ring_width == 0)
  can_tile_transposed = (physical.shape[1] % ring_height == 0) and (physical.shape[0] % ring_width == 0)
  
  other_axes = [size for i, size in enumerate(ici_parallelism) if i != axis]
  non_trivial_other_axes = [size for size in other_axes if size > 1]
  if len(non_trivial_other_axes) == 2:
    target_shape = tuple(non_trivial_other_axes)
    normal_ring_shape = (physical.shape[0] // ring_height, physical.shape[1] // ring_width)
    transposed_ring_shape = (physical.shape[1] // ring_height, physical.shape[0] // ring_width)
    if can_tile_transposed and transposed_ring_shape == target_shape:
      physical = physical.T
    elif can_tile_normal and normal_ring_shape == target_shape:
      pass
    elif not can_tile_normal and can_tile_transposed:
      physical = physical.T
    elif can_tile_normal and can_tile_transposed and physical.shape[0] > physical.shape[1]:
      physical = physical.T
  else:
    if not can_tile_normal and can_tile_transposed:
      physical = physical.T
    elif can_tile_normal and can_tile_transposed and physical.shape[0] > physical.shape[1]:
      physical = physical.T

[gemini-code-assist]

If the physical grid of devices is not fully populated (e.g., due to missing coordinates or an incomplete grid), some elements in the physical array will remain None. This will silently propagate into the final device mesh, leading to cryptic errors later when JAX attempts to use the mesh.

Adding a defensive check to ensure all coordinates are populated with valid devices will prevent this.

Suggested change

	physical = np.full(tuple(extent), None, dtype=object)
	for device, coord in zip(devices, coords):
	physical[tuple(coord - origin)] = device
	physical = physical.reshape([dim for dim in extent if dim > 1])
	physical = np.full(tuple(extent), None, dtype=object)
	for device, coord in zip(devices, coords):
	physical[tuple(coord - origin)] = device
	physical = physical.reshape([dim for dim in extent if dim > 1])
	if None in physical:
	raise ValueError("The physical device grid is not fully populated (contains missing devices/coordinates).")

[RSudarsanan]

@gemini-code-assist can you check now. It should be fixed.

[gemini-code-assist]

Thank you for the update. I have reviewed the changes. The defensive check correctly ensures that the physical device grid is fully populated before proceeding, which prevents silent propagation of None values and potential runtime errors. The implementation looks solid.

[gemini-code-assist]

Adding a test case to verify that the physical grid is transposed to align with the non-trivial logical axes (e.g., transposing a 4x32 grid to 32x4 to match [8, 2] logical axes) ensures that the alignment-aware transposing logic works correctly and prevents future regressions.

  def test_transposes_grid_when_needed_for_tiling(self):
    """Test that a 2x16 grid is transposed to 16x2 when ring=8."""
    devices = [self.FakeDevice(x * 16 + y, (x, y, 0)) for x in range(2) for y in range(16)]
    # A 2x16 physical grid can only be tiled by 4x2 blocks (for ring=8) if transposed.
    mesh = max_utils.create_ring_axis_device_mesh([1, 1, 8, 4], self.AXES, devices, "tensor")
    self.assertEqual(mesh.shape, (1, 1, 8, 4))
    self.assertEqual(sorted(d.id for d in mesh.flatten()), sorted(d.id for d in devices))

  def test_transposes_grid_to_align_with_logical_axes(self):
    """Test that a 4x32 grid is transposed to 32x4 to align with other logical axes [8, 2]."""
    devices = [self.FakeDevice(x * 32 + y, (x, y, 0)) for x in range(4) for y in range(32)]
    # other_axes (excluding tensor) is [8, 1, 2], non-trivial: [8, 2]
    # Without alignment-aware transposing, the 4x32 grid is tiled as (1, 16) rings,
    # which scrambles the [8, 2] logical axes.
    # With alignment-aware transposing, it is transposed to 32x4, tiled as (8, 2) rings,
    # matching the logical axes perfectly.
    mesh = max_utils.create_ring_axis_device_mesh([8, 1, 8, 2], self.AXES, devices, "tensor")
    self.assertEqual(mesh.shape, (8, 1, 8, 2))
    self.assertEqual(sorted(d.id for d in mesh.flatten()), sorted(d.id for d in devices))

[RSudarsanan]

@gemini-code-assist We fixed it, can you check again. Thanks!

[gemini-code-assist]

Thank you for the update. I have reviewed the changes, and the additional test case correctly validates the alignment-aware transposing logic for the 4x32 physical grid. The implementation looks solid and addresses the requirements for physical ring arrangement on the torus. No further issues identified.

[richjames0]

I think there's an issue here in test_transposes_grid_when_needed_for_tiling which will fail due to a mismatch in the number of devices.

It appears that an earlier commit in this PR noticed that and fixed it but then it's subsequently been reverted and now will fail again

[codecov]

Codecov Report

❌ Patch coverage is 61.01695% with 23 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/max_utils.py	66.00%	11 Missing and 6 partials ⚠️
src/maxtext/utils/maxtext_utils.py	33.33%	5 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[systalyze-ai]

I think there's an issue here in test_transposes_grid_when_needed_for_tiling which will fail due to a mismatch in the number of devices.

It appears that an earlier commit in this PR noticed that and fixed it but then it's subsequently been reverted and now will fail again

Thanks @richjames0 for catching that, its been fixed. The test now matches its device count.