Class HybridMeshShape

axlearn/common/utils.py:95–116

A mesh shape for hybrid (i.e., ICI and DCN) parallelism. For example, with mesh axes (data, model): - Pure fsdp on a v4-8: HybridMeshShape(ici_mesh_shape=(1, 4), dcn_mesh_shape=(1, 1)) - Two-way data parallelism over 2 H100 nodes, and fsdp within-node: HybridMeshShape(ic

@dataclasses.dataclass
class HybridMeshShape:
    """A mesh shape for hybrid (i.e., ICI and DCN) parallelism.

    For example, with mesh axes (data, model):
    - Pure fsdp on a v4-8:
        HybridMeshShape(ici_mesh_shape=(1, 4), dcn_mesh_shape=(1, 1))
    - Two-way data parallelism over 2 H100 nodes, and fsdp within-node:
        HybridMeshShape(ici_mesh_shape=(1, 8), dcn_mesh_shape=(2, 1))
    """

    ici_mesh_shape: MeshShape
    dcn_mesh_shape: MeshShape

    def __post_init__(self) -> None:
        if len(self.ici_mesh_shape) != len(self.dcn_mesh_shape):
            raise ValueError(
                f"{self.ici_mesh_shape=} should have the same length as {self.dcn_mesh_shape}."
            )

    def __len__(self):
        assert len(self.ici_mesh_shape) == len(self.dcn_mesh_shape)
        return len(self.ici_mesh_shape)


# "device" = Accelerator memory, e.g. HBM.