mesh: add mesh_ring_axis to give a mesh axis physical rings on a torus without wraparound

src/maxtext/configs/base.yml

@@ -1128,6 +1128,8 @@ custom_mesh: "" # Available options: ['hybrid_ring_64x4', 'hybrid_ring_32x8']
 allow_split_physical_axes: false
 # Apply transformations to the mesh to optimize for TPU v6e
 optimize_mesh_for_tpu_v6e: false
+# Logical mesh axis to lay out as a physical TPU ring; empty disables.
+mesh_ring_axis: ""
 
 shardy: true # Whether to use shardy XLA backend (default in Jax starting 0.7.0), or GSPMD (to be fully deprecated ~2026)
 

src/maxtext/configs/types.py

@@ -1147,6 +1147,7 @@ class HardwareAndMesh(BaseModel):
       ),
   )
   optimize_mesh_for_tpu_v6e: bool = Field(False, description="Apply transformations to the mesh for TPU v6e.")
+  mesh_ring_axis: str = Field("", description="Logical mesh axis to lay out as a physical TPU ring; empty disables.")
   shardy: bool = Field(True, description="Whether to use shardy XLA backend.")
   pure_nnx_decoder: bool = Field(
       True,

src/maxtext/utils/max_utils.py

@@ -484,6 +484,50 @@ def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_typ
   return parallelism_vals
 
 
+def create_ring_axis_device_mesh(ici_parallelism, mesh_axes, devices, ring_axis):
+  """Creates an ICI device mesh with `ring_axis` groups arranged as physical TPU rings."""
+  if ring_axis not in mesh_axes:
+    raise ValueError(f"mesh_ring_axis={ring_axis} is not one of the mesh axes {mesh_axes}")
+  axis = list(mesh_axes).index(ring_axis)
+  ring = ici_parallelism[axis]
+  if ring % 2:
+    raise ValueError(f"mesh_ring_axis={ring_axis} needs an even parallelism, got {ring}")
+
+  coords = np.array([d.coords for d in devices])
+  origin, extent = coords.min(axis=0), coords.max(axis=0) - coords.min(axis=0) + 1
+  physical = np.empty(tuple(extent), 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 physical.ndim != 2:
+    raise ValueError(f"mesh_ring_axis needs a 2D slice of the torus, got physical shape {tuple(extent)}")
+  # Preserve locality for non-ring axes.
+  if physical.shape[0] > physical.shape[1]:
+    physical = physical.T
+  ring_height, ring_width = ring // 2, 2
+  if physical.shape[0] % ring_height or physical.shape[1] % ring_width:
+    raise ValueError(
+        f"mesh_ring_axis={ring_axis} of size {ring} needs a {ring_height}x{ring_width} block to tile the"
+        f" physical mesh {physical.shape}"
+    )
+
+  rings = []
+  for i in range(physical.shape[0] // ring_height):
+    columns = range(physical.shape[1] // ring_width)
+    for j in columns if i % 2 == 0 else reversed(columns):
+      left = list(physical[i * ring_height : (i + 1) * ring_height, ring_width * j])
+      right = list(physical[i * ring_height : (i + 1) * ring_height, ring_width * j + 1])
+      rings.append(left + right[::-1])
+
+  mesh = np.empty((len(rings), ring), dtype=object)
+  for i, devices_in_ring in enumerate(rings):
+    mesh[i] = devices_in_ring
+  other_axes = [size for i, size in enumerate(ici_parallelism) if i != axis]
+  mesh = np.moveaxis(mesh.reshape([*other_axes, ring]), -1, axis)
+  max_logging.log(f"Laid the '{ring_axis}' mesh axis out as a {ring_height}x{ring_width} physical ring")
+  return mesh
+
+
 def reshape_mesh_to_rings(a, strategy):
   """Reshape device mesh to rings for 64x4 or 32x8 mesh shape"""
   b = []

src/maxtext/utils/maxtext_utils.py

@@ -2193,7 +2193,13 @@ def create_device_mesh(config, devices=None):
 
   allow_split_physical_axes = config.allow_split_physical_axes if config.allow_split_physical_axes else False
 
-  if num_slices > 1:
+  ring_axis = getattr(config, "mesh_ring_axis", "")
+  if ring_axis and num_slices > 1:
+    raise ValueError("mesh_ring_axis lays out the ICI mesh of one slice; it is not implemented for multi-slice.")
+
+  if ring_axis:
+    mesh = max_utils.create_ring_axis_device_mesh(ici_parallelism, config.mesh_axes, devices, ring_axis)
+  elif num_slices > 1:
     dcn_parallelism = getattr(config, "dcn_parallelism", None)
     if dcn_parallelism is None:
       dcn_map = {

tests/unit/max_utils_test.py

@@ -18,6 +18,7 @@
 import time
 import unittest
 from unittest import mock
+import numpy as np
 from flax import linen as nn
 from flax import nnx
 import jax
@@ -135,6 +136,49 @@ def test_invalid_strategy(self):
       max_utils.is_valid_custom_mesh([1, 1, 1, 1, 1, 16, 16, 1], "invalid_strategy")
 
 
+class RingAxisDeviceMeshTest(unittest.TestCase):
+  """Tests ring-axis device mesh creation."""
+
+  class FakeDevice:
+
+    def __init__(self, device_id, coords):
+      self.id = device_id
+      self.coords = coords
+
+  AXES = ["data", "fsdp", "tensor", "expert"]
+  SIZES = [1, 1, 8, 16]
+
+  def _devices(self):
+    return [self.FakeDevice(x * 16 + y, (x, y, 0)) for x in range(8) for y in range(16)]
+
+  def _cyclic_hops(self, group):
+    return [
+        tuple(abs(a - b) for a, b in zip(first.coords, second.coords))
+        for first, second in zip(group, list(group[1:]) + [group[0]])
+    ]
+
+  def test_ring_axis_groups_are_physical_cycles(self):
+    devices = self._devices()
+    mesh = max_utils.create_ring_axis_device_mesh(self.SIZES, self.AXES, devices, "tensor")
+    self.assertEqual(mesh.shape, tuple(self.SIZES))
+    self.assertEqual(sorted(d.id for d in mesh.flatten()), sorted(d.id for d in devices))
+    for tensor_group in mesh.squeeze().T:
+      hops = self._cyclic_hops(tensor_group)
+      self.assertEqual(len(hops), 8)
+      for hop in hops:
+        self.assertEqual(sorted(hop), [0, 0, 1])
+
+  def test_default_mesh_leaves_the_ring_axis_open(self):
+    devices = self._devices()
+    default = np.asarray(devices, dtype=object).reshape(self.SIZES)
+    hops = self._cyclic_hops(default.squeeze()[:, 0])
+    self.assertEqual(sum(sorted(hop) != [0, 0, 1] for hop in hops), 1)
+
+  def test_rejects_odd_ring_axis(self):
+    with self.assertRaises(ValueError):
+      max_utils.create_ring_axis_device_mesh([1, 1, 8, 16], self.AXES, self._devices(), "data")
+
+
 class FillUnspecifiedMeshAxesTest(unittest.TestCase):
   """Tests for fill_unspecified_mesh_axes."""