diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index a4cf6877e3..67f1f02f88 100644 --- a/src/maxtext/configs/base.yml +++ b/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) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 1f8659f7e0..f08a42585a 100644 --- a/src/maxtext/configs/types.py +++ b/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, diff --git a/src/maxtext/utils/max_utils.py b/src/maxtext/utils/max_utils.py index 57cc8d49c4..efaef34fa0 100644 --- a/src/maxtext/utils/max_utils.py +++ b/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 = [] diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index f7c15e7d90..8af8978ed5 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ 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 = { diff --git a/tests/unit/max_utils_test.py b/tests/unit/max_utils_test.py index 3842bed1d4..f84719a9f2 100644 --- a/tests/unit/max_utils_test.py +++ b/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."""