Function create_device_mesh

axlearn/common/utils.py:1767–1888

Constructs a device mesh. If `mesh_shape` is specified as a `HybridMeshShape`, we use the `ici_mesh_shape` and `dcn_mesh_shape` directly to construct the mesh. If `mesh_shape` is specified as a `MeshShape`, we first determine whether we are running in a TPU or GPU environment.

(
    mesh_shape: Union[MeshShape, HybridMeshShape],
    *,
    devices: Optional[Sequence[Any]] = None,
)

                        a[a_i + 1, a_j + 1],
                        a[a_i + 1, a_j],
                    ]
                )
    else:
        raise ValueError(f"The target mesh shape {shape} is not implemented.")
    return np.reshape(np.array(b), shape)


def _maybe_get_special_mesh(
    mesh_shape: MeshShape, *, devices: np.ndarray
) -> Optional[tuple[int, int]]:
    """Checks if any of the special mesh shapes are applicable."""
    if int(np.prod(mesh_shape)) != 256:
        return None
    if getattr(devices[0], "device_kind", None) not in [
        "TPU v5e",
        "TPU v6e",
        "TPU v6 lite",
        "TPU v5 lite",
    ]:
        return None

    filtered_mesh = tuple(filter(lambda x: x != 1, mesh_shape))
    target_shapes = [(64, 4), (32, 8)]
    return None if filtered_mesh not in target_shapes else filtered_mesh


def build_standard_mesh(mesh_shape: MeshShape, *, devices: np.ndarray) -> np.ndarray:
    logging.info("Building device mesh.")
    mesh_shape = infer_mesh_shape(mesh_shape, num_devices=devices.size)
    try:
        if (shape := _maybe_get_special_mesh(mesh_shape, devices=devices)) is not None:
            # If any of the special mesh shapes is applicable, use them.
            mesh = mesh_utils.create_device_mesh([16, 16], devices=devices)
            mesh = _reshape_mesh_to_rings(mesh, shape=shape)
            mesh = mesh.reshape(mesh_shape)
            logging.log_first_n(logging.INFO, "Using custom mesh: %s", 1, str(mesh))
            return mesh
        return mesh_utils.create_device_mesh(
            mesh_shape,
            devices=devices,
            # Set allow_split_physical_axes to True to split physical axes.
            # Reference: jax.experimental.mesh_utils.create_device_mesh docs.
            allow_split_physical_axes=True,
        )
    except NotImplementedError as e:
        logging.warning(
            "mesh_utils.create_device_mesh cannot handle shape %s: %s. "
            "Falling back to the naive mesh. Performance may be reduced.",
            mesh_shape,
            e,
        )
        return devices.reshape(mesh_shape)


def create_hybrid_device_mesh(
    mesh_shape: HybridMeshShape,
    *,
    devices: Sequence[Any],