nemo_automodel.components.distributed.mesh_utils

Device mesh construction and access utilities for distributed training.

Module Contents

Classes

Name	Description
`_MeshSpec`	Named mesh shape plus derived flattened axes.

Functions

Name	Description
`_can_reuse_default_group`	Return whether a derived mesh dimension spans the default process group.
`_create_device_meshes`	Create raw device meshes based on distributed config type.
`_create_fsdp2_device_mesh`	Create the FSDP2 root mesh and optional MoE mesh.
`_create_megatron_fsdp_device_mesh`	Create the Megatron FSDP mesh.
`_create_moe_mesh`	-
`_degree`	-
`_infer_dp_size`	-
`_init_named_mesh`	-
`_mesh_device_type`	-
`_nccl_backend_override`	Create per-axis NCCL options for DeviceMesh subgroups.
`_register_flattened_axes`	-
`_register_flattened_axes_for_rank_subset`	Register flattened axes without slicing a mesh on non-member ranks.
`_require_size_one`	-
`_unflatten_compat`	Unflatten a mesh with its NCCL timeout, including the PyTorch 2.9 fallback.
`_validate_mesh_spec`	-
`create_ring_ulysses_mesh`	Derive a 2D `("ring", "ulysses")` view of the mesh’s `"cp"` axis.
`get_flat_mesh`	Access a 1D submesh by parallelism name (e.g. `"dp"`, `"tp"`, `"dp_cp"`).
`get_fsdp_dp_mesh`	Return the DP mesh for FSDP2 without losing the original root mesh.
`get_submesh`	Access a submesh by parallelism dim names.

Data

__all__

API

class nemo_automodel.components.distributed.mesh_utils._MeshSpec(
    shape: tuple[int, ...],
    axes: tuple[nemo_automodel.components.distributed.mesh.MeshAxisName, ...],
    flattened_axes: dict[nemo_automodel.components.distributed.mesh.MeshAxisName, tuple[nemo_automodel.components.distributed.mesh.MeshAxisName, ...]] = dict()
)

Dataclass

Named mesh shape plus derived flattened axes.

axes

tuple[MeshAxisName, ...]

flattened_axes

dict[MeshAxisName, tuple[MeshAxisName, ...]] = field(default_factory=dict)

shape

tuple[int, ...]

nemo_automodel.components.distributed.mesh_utils._can_reuse_default_group(
    group_size: int
) -> bool

Return whether a derived mesh dimension spans the default process group.

nemo_automodel.components.distributed.mesh_utils._create_device_meshes(
    strategy_config: nemo_automodel.components.distributed.config.DistributedStrategyConfig,
    parallelism: nemo_automodel.components.distributed.mesh.ParallelismSizes,
    world_size: int,
    timeout_minutes: int | None = None,
    ranks: list[int] | tuple[int, ...] | None = None
) -> tuple[torch.distributed.device_mesh.DeviceMesh | None, torch.distributed.device_mesh.DeviceMesh | None]

Create raw device meshes based on distributed config type.

nemo_automodel.components.distributed.mesh_utils._create_fsdp2_device_mesh(
    parallelism: nemo_automodel.components.distributed.mesh.ParallelismSizes,
    world_size: int,
    timeout_minutes: int | None = None,
    ranks: list[int] | tuple[int, ...] | None = None
) -> tuple[torch.distributed.device_mesh.DeviceMesh, torch.distributed.device_mesh.DeviceMesh | None]

Create the FSDP2 root mesh and optional MoE mesh.

nemo_automodel.components.distributed.mesh_utils._create_megatron_fsdp_device_mesh(
    parallelism: nemo_automodel.components.distributed.mesh.ParallelismSizes,
    world_size: int,
    timeout_minutes: int | None = None,
    ranks: list[int] | tuple[int, ...] | None = None
) -> torch.distributed.device_mesh.DeviceMesh

Create the Megatron FSDP mesh.

nemo_automodel.components.distributed.mesh_utils._create_moe_mesh(
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    ep_shard_size: int,
    ep_size: int,
    pp_size: int = 1,
    timeout_minutes: int | None = None,
    ranks: list[int] | tuple[int, ...] | None = None
) -> torch.distributed.device_mesh.DeviceMesh

nemo_automodel.components.distributed.mesh_utils._degree(
    value: int | None
) -> int

nemo_automodel.components.distributed.mesh_utils._infer_dp_size(
    dp_size: int | None,
    world_size: int,
    non_dp_size: int,
    expression: str,
    factors: tuple[int, ...]
) -> int

nemo_automodel.components.distributed.mesh_utils._init_named_mesh(
    spec: nemo_automodel.components.distributed.mesh_utils._MeshSpec,
    timeout_minutes: int | None = None,
    ranks: list[int] | tuple[int, ...] | None = None
) -> torch.distributed.device_mesh.DeviceMesh

nemo_automodel.components.distributed.mesh_utils._mesh_device_type() -> str

nemo_automodel.components.distributed.mesh_utils._nccl_backend_override(
    axes: tuple[str, ...],
    device_type: str,
    timeout_minutes: int | None
)

Create per-axis NCCL options for DeviceMesh subgroups.

init_process_group(timeout=...) configures the default process group, but init_device_mesh creates additional per-axis process groups. Without a backend override those groups keep PyTorch’s default NCCL timeout.

nemo_automodel.components.distributed.mesh_utils._register_flattened_axes(
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    flattened_axes: dict[nemo_automodel.components.distributed.mesh.MeshAxisName, tuple[nemo_automodel.components.distributed.mesh.MeshAxisName, ...]],
    timeout_minutes: int | None = None
) -> None

nemo_automodel.components.distributed.mesh_utils._register_flattened_axes_for_rank_subset(
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    spec: nemo_automodel.components.distributed.mesh_utils._MeshSpec,
    ranks: list[int] | tuple[int, ...],
    timeout_minutes: int | None
) -> None

PyTorch cannot slice an explicit DeviceMesh on ranks outside that mesh. Separate-model jobs still need every global rank to create process groups in the same order, so construct each flattened subgroup explicitly and retain the subgroup local to the current rank.

nemo_automodel.components.distributed.mesh_utils._require_size_one(
    strategy_name: str,
    size: int | None,
    feature_name: str
) -> None

nemo_automodel.components.distributed.mesh_utils._unflatten_compat(
    flat_mesh: torch.distributed.device_mesh.DeviceMesh,
    axis: int,
    sizes: tuple,
    names: tuple,
    timeout_minutes: int | None = None
) -> torch.distributed.device_mesh.DeviceMesh

Unflatten a mesh with its NCCL timeout, including the PyTorch 2.9 fallback.

PyTorch reuses the default process group for a world-sized dimension only when no backend override is supplied. The default group already has the configured distributed timeout, so omitting that redundant override avoids allocating another NCCL communicator.

nemo_automodel.components.distributed.mesh_utils._validate_mesh_spec(
    spec: nemo_automodel.components.distributed.mesh_utils._MeshSpec
) -> None

nemo_automodel.components.distributed.mesh_utils.create_ring_ulysses_mesh(
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    ring_degree: int,
    ulysses_degree: int
) -> torch.distributed.device_mesh.DeviceMesh

Derive a 2D ("ring", "ulysses") view of the mesh’s "cp" axis.

Diffusers’ context-parallel API (ContextParallelConfig.mesh) requires a device mesh with dimensions named "ring" and "ulysses". This reshapes the 1D "cp" axis of an existing root mesh into that layout so context parallelism shares the process groups already created for FSDP2, instead of initializing a second world mesh.

Parameters:

device_mesh

DeviceMesh

Root mesh containing a "cp" dimension (as created by _create_fsdp2_device_mesh).

ring_degree

int

Size of the ring-attention dimension.

ulysses_degree

int

Size of the Ulysses (all-to-all) attention dimension.

Returns: DeviceMesh

A 2D DeviceMesh of shape (ring_degree, ulysses_degree) with dim

Raises:

ValueError: If ring_degree * ulysses_degree does not equal the size of the mesh’s "cp" dimension.

nemo_automodel.components.distributed.mesh_utils.get_flat_mesh(
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    name: str
) -> torch.distributed.device_mesh.DeviceMesh

Access a 1D submesh by parallelism name (e.g. "dp", "tp", "dp_cp").

PyTorch 2.11 deprecates root_mesh["name"] for dimensions created via _flatten(). This reads the _flatten() result directly.

Parameters:

device_mesh

DeviceMesh

Any DeviceMesh (root or submesh).

name

str

Parallelism dimension name.

nemo_automodel.components.distributed.mesh_utils.get_fsdp_dp_mesh(
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    dp_replicate_name: str = MeshAxisName.DP_REPLICATE,
    dp_shard_cp_name: str = MeshAxisName.DP_SHARD_CP
) -> torch.distributed.device_mesh.DeviceMesh

Return the DP mesh for FSDP2 without losing the original root mesh.

get_submesh() may rebuild a fresh DeviceMesh when asked to compose native and flattened dims like ("dp_replicate", "dp_shard_cp"). That is fine for many local operations, but FSDP2 expects its DP mesh to share the same root mesh as TP/EP meshes. On multi-node TP runs this can break group construction in non-obvious ways.

Prefer native dimensions whenever possible:

cp=1, dp_replicate=1 -> device_mesh["dp_shard"]
cp=1, dp_replicate>1 -> device_mesh[("dp_replicate", "dp_shard")]
cp>1, dp_replicate=1 -> device_mesh["dp_shard_cp"]

When both CP and replicated DP are active we fall back to get_submesh() because the composed mesh is genuinely multi-level.

nemo_automodel.components.distributed.mesh_utils.get_submesh(
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    names: tuple
) -> torch.distributed.device_mesh.DeviceMesh

Access a submesh by parallelism dim names.

Handles all cases: single dims, multi-dim slices, and combinations that include _flatten()-created dims (e.g. ("dp_replicate", "dp_shard_cp")). For the latter, finds the parent _flatten() result and calls _unflatten() to decompose it into the requested shape.

Parameters:

device_mesh

DeviceMesh

Any DeviceMesh (root or submesh).

names

tuple

Tuple of dimension names.

nemo_automodel.components.distributed.mesh_utils.__all__ = ['_create_device_meshes', '_create_fsdp2_device_mesh', '_create_megatron_fsdp_de...