> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.distributed

## Subpackages

* **[`nemo_automodel.components.distributed.blockdiag_cp`](/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp)**
* **[`nemo_automodel.components.distributed.context_parallel`](/nemo-automodel/nemo_automodel/components/distributed/context_parallel)**
* **[`nemo_automodel.components.distributed.pipelining`](/nemo-automodel/nemo_automodel/components/distributed/pipelining)**

## Submodules

* **[`nemo_automodel.components.distributed.activation_checkpointing`](/nemo-automodel/nemo_automodel/components/distributed/activation_checkpointing)**
* **[`nemo_automodel.components.distributed.config`](/nemo-automodel/nemo_automodel/components/distributed/config)**
* **[`nemo_automodel.components.distributed.cp_vision_frame_shard`](/nemo-automodel/nemo_automodel/components/distributed/cp_vision_frame_shard)**
* **[`nemo_automodel.components.distributed.ddp`](/nemo-automodel/nemo_automodel/components/distributed/ddp)**
* **[`nemo_automodel.components.distributed.fsdp2`](/nemo-automodel/nemo_automodel/components/distributed/fsdp2)**
* **[`nemo_automodel.components.distributed.fsdp_patches`](/nemo-automodel/nemo_automodel/components/distributed/fsdp_patches)**
* **[`nemo_automodel.components.distributed.grad_utils`](/nemo-automodel/nemo_automodel/components/distributed/grad_utils)**
* **[`nemo_automodel.components.distributed.init_utils`](/nemo-automodel/nemo_automodel/components/distributed/init_utils)**
* **[`nemo_automodel.components.distributed.megatron_fsdp`](/nemo-automodel/nemo_automodel/components/distributed/megatron_fsdp)**
* **[`nemo_automodel.components.distributed.mesh`](/nemo-automodel/nemo_automodel/components/distributed/mesh)**
* **[`nemo_automodel.components.distributed.mesh_utils`](/nemo-automodel/nemo_automodel/components/distributed/mesh_utils)**
* **[`nemo_automodel.components.distributed.model_parallelizer`](/nemo-automodel/nemo_automodel/components/distributed/model_parallelizer)**
* **[`nemo_automodel.components.distributed.multimodal_fsdp`](/nemo-automodel/nemo_automodel/components/distributed/multimodal_fsdp)**
* **[`nemo_automodel.components.distributed.optimized_tp_plans`](/nemo-automodel/nemo_automodel/components/distributed/optimized_tp_plans)**
* **[`nemo_automodel.components.distributed.parallel_styles`](/nemo-automodel/nemo_automodel/components/distributed/parallel_styles)**
* **[`nemo_automodel.components.distributed.parallelizer`](/nemo-automodel/nemo_automodel/components/distributed/parallelizer)**
* **[`nemo_automodel.components.distributed.parallelizer_utils`](/nemo-automodel/nemo_automodel/components/distributed/parallelizer_utils)**
* **[`nemo_automodel.components.distributed.recompute_replay`](/nemo-automodel/nemo_automodel/components/distributed/recompute_replay)**
* **[`nemo_automodel.components.distributed.tensor_utils`](/nemo-automodel/nemo_automodel/components/distributed/tensor_utils)**
* **[`nemo_automodel.components.distributed.thd_utils`](/nemo-automodel/nemo_automodel/components/distributed/thd_utils)**
* **[`nemo_automodel.components.distributed.tp_replicas`](/nemo-automodel/nemo_automodel/components/distributed/tp_replicas)**
* **[`nemo_automodel.components.distributed.utils`](/nemo-automodel/nemo_automodel/components/distributed/utils)**

## Package Contents

### Classes

| Name                                                                                                       | Description                                                       |
| ---------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------- |
| [`DDPConfig`](#nemo_automodel-components-distributed-config-DDPConfig)                                     | Additional configuration for DDP distributed training.            |
| [`DistInfo`](#nemo_automodel-components-distributed-init_utils-DistInfo)                                   | Holds information about the distributed training environment.     |
| [`DistributedSetup`](#nemo_automodel-components-distributed-config-DistributedSetup)                       | Resolved distributed topology and execution policies.             |
| [`FSDP2Config`](#nemo_automodel-components-distributed-config-FSDP2Config)                                 | Additional configuration for FSDP2 distributed training.          |
| [`MegatronFSDPConfig`](#nemo_automodel-components-distributed-config-MegatronFSDPConfig)                   | Additional configuration for MegatronFSDP distributed training.   |
| [`MeshContext`](#nemo_automodel-components-distributed-mesh-MeshContext)                                   | Resolved runtime context for model parallelization.               |
| [`MoEParallelizerConfig`](#nemo_automodel-components-distributed-config-MoEParallelizerConfig)             | Configuration for MoE model parallelization (EP + FSDP settings). |
| [`ModelParallelizer`](#nemo_automodel-components-distributed-parallelizer-ModelParallelizer)               | Single model-owned parallelization sidecar contract.              |
| [`MultimodalDistributedConfig`](#nemo_automodel-components-distributed-config-MultimodalDistributedConfig) | Distributed policies for resolved multimodal modules.             |
| [`ParallelismSizes`](#nemo_automodel-components-distributed-mesh-ParallelismSizes)                         | Build-time requested parallelism sizes.                           |
| [`PipelineConfig`](#nemo_automodel-components-distributed-pipelining-config-PipelineConfig)                | Configuration for pipeline parallel training.                     |

### Functions

| Name                                                                                                 | Description                                                                          |
| ---------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------ |
| [`initialize_distributed`](#nemo_automodel-components-distributed-init_utils-initialize_distributed) | Initialize the torch.distributed environment and core model parallel infrastructure. |

### API

```python
class nemo_automodel.components.distributed.config.DDPConfig(
    activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False,
    activation_checkpointing_scope: nemo_automodel.components.distributed.config.ActivationCheckpointingScope = 'all',
    broadcast_buffers: bool = False,
    find_unused_parameters: bool = False,
    static_graph: bool = False,
    bucket_cap_mb: float | None = None,
    gradient_as_bucket_view: bool = False,
    autocast_dtype: torch.dtype | None = None
)
```

Dataclass

Additional configuration for DDP distributed training.

Note: DDP does not support tensor parallelism, pipeline parallelism, or expert parallelism.
Only dp\_size is relevant (inferred from world\_size).

**`activation_checkpointing`** `ActivationCheckpointingMode = False`

Enable activation checkpointing. `True` or
`"full"` keeps the existing full activation checkpointing behavior. `"selective"` wraps transformer
blocks with PyTorch selective activation checkpointing.

---

**`activation_checkpointing_scope`** `ActivationCheckpointingScope = 'all'`

Which extracted
layer groups activation checkpointing should wrap. `"all"`
selects every extracted group. Scoped values such as
`"language"`, `"vision"`, and `"multimodal"` are filtered
to trainable layers before generic wrapping.

---

**`autocast_dtype`** `dtype | None = None`

If set, recipes can wrap the forward pass in
`torch.autocast(device_type="cuda", dtype=autocast_dtype)`. Set to `None` to disable.
Can be set from YAML as a string (e.g. `autocast_dtype: bfloat16`).

---

**`broadcast_buffers`** `bool = False`

Synchronize module buffers before each forward.

---

**`bucket_cap_mb`** `float | None = None`

DDP gradient bucket size in MiB. `None` uses PyTorch's default.

---

**`find_unused_parameters`** `bool = False`

Forwarded to PyTorch DDP for models with
conditionally unused trainable parameters.

---

**`gradient_as_bucket_view`** `bool = False`

Make gradients views into DDP buckets after the first iteration.

---

**`static_graph`** `bool = False`

Tell DDP the used/unused parameter set is stable.

---

```python
nemo_automodel.components.distributed.config.DDPConfig.__post_init__()
```

```python
nemo_automodel.components.distributed.config.DDPConfig.to_dict() -> typing.Dict[str, typing.Any]
```

Convert config to dictionary.

```python
class nemo_automodel.components.distributed.init_utils.DistInfo(
    backend: str,
    rank: int,
    world_size: int,
    device: torch.device,
    is_main: bool
)
```

Dataclass

Holds information about the distributed training environment.

**`backend`** `str`

The backend used for torch.distributed (e.g., 'nccl').

---

**`device`** `device`

The device assigned to the current process.

---

**`is_main`** `bool`

True if the process is the main process (rank 0).

---

**`rank`** `int`

The rank of the current process.

---

**`world_size`** `int`

The total number of processes.

---

```python
class nemo_automodel.components.distributed.config.DistributedSetup(
    mesh_context: 'MeshContext',
    strategy_config: nemo_automodel.components.distributed.config.DistributedStrategyConfig | None = None,
    pipeline_config: 'PipelineConfig | None' = None,
    moe_parallel_config: 'MoEParallelizerConfig | None' = None,
    activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False
)
```

Dataclass

Resolved distributed topology and execution policies.

**`activation_checkpointing`** `ActivationCheckpointingMode = False`

---

**`mesh_context`** `'MeshContext'`

---

**`moe_parallel_config`** `'MoEParallelizerConfig | None' = None`

---

**`pipeline_config`** `'PipelineConfig | None' = None`

---

**`strategy_config`** `DistributedStrategyConfig | None = None`

---

```python
nemo_automodel.components.distributed.config.DistributedSetup.__post_init__() -> None
```

Keep the compatibility bundle and its MeshContext policy in sync.

```python
nemo_automodel.components.distributed.config.DistributedSetup.build(
    strategy: str | nemo_automodel.components.distributed.config.DistributedStrategyConfig = 'fsdp2',
    parallelism_sizes: 'ParallelismSizes | None' = None,
    pipeline_config: 'PipelineConfig | dict | None' = None,
    moe_parallel_config: 'MoEParallelizerConfig | dict | None' = None,
    activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False,
    world_size: int | None = None,
    timeout_minutes: int | None = None,
    ranks: list[int] | tuple[int, ...] | None = None
) -> 'DistributedSetup'
```

classmethod

Create a resolved distributed setup from sizes and policy configs.

Intentionally, this function is forgiving wrt the input types, allowing
strings for the strategy and dicts for the pipeline and MoE configs.

```python
class nemo_automodel.components.distributed.config.FSDP2Config(
    sequence_parallel: bool = False,
    tp_plan: dict | str | None = None,
    patch_is_packed_sequence: bool = False,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None = (lambda: MixedPrecisionPoli...,
    offload_policy: torch.distributed.fsdp.CPUOffloadPolicy | None = None,
    autocast_dtype: torch.dtype | None = None,
    activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False,
    activation_checkpointing_scope: nemo_automodel.components.distributed.config.ActivationCheckpointingScope = 'all',
    defer_fsdp_grad_sync: bool = True,
    reshard_after_forward: bool | None = None,
    enable_async_tensor_parallel: bool = False,
    enable_compile: bool = False,
    enable_fsdp2_prefetch: bool = False,
    fsdp2_backward_prefetch_depth: int = 2,
    fsdp2_forward_prefetch_depth: int = 1,
    multimodal: nemo_automodel.components.distributed.config.MultimodalDistributedConfig = MultimodalDistributedConfig()
)
```

Dataclass

Additional configuration for FSDP2 distributed training.

Note: Size parameters (dp\_size, dp\_replicate\_size, tp\_size, pp\_size, cp\_size, ep\_size)
are grouped separately in `ParallelismSizes`.

**`activation_checkpointing`** `ActivationCheckpointingMode = False`

Enable activation checkpointing. `True` or
`"full"` keeps the existing full activation checkpointing behavior. `"selective"` wraps transformer
blocks with PyTorch selective activation checkpointing.

---

**`activation_checkpointing_scope`** `ActivationCheckpointingScope = 'all'`

Which extracted
layer groups activation checkpointing should wrap. `"all"`
selects every extracted group. Scoped values such as
`"language"`, `"vision"`, and `"multimodal"` are filtered
to trainable layers before generic wrapping.

---

**`autocast_dtype`** `dtype | None = None`

If set, wraps the forward pass in
`torch.autocast(device_type="cuda", dtype=autocast_dtype)`.  Use with
`output_dtype=float32` in mp\_policy to keep the residual stream in fp32
while running matmuls in lower precision.  Set to `None` to disable.
Can be set from YAML as a string (e.g. `autocast_dtype: bfloat16`).

---

**`defer_fsdp_grad_sync`** `bool = True`

Defer FSDP gradient sync to final micro-batch.

---

**`enable_async_tensor_parallel`** `bool = False`

Enable async tensor parallelism via
`torch._inductor.config._micro_pipeline_tp`.  Overlaps ReduceScatter with
compute in row-parallel layers.  Requires `sequence_parallel=True` (forced
automatically with a warning if not set).  Also enables symmetric memory for
the TP group.

---

**`enable_compile`** `bool = False`

Apply per-layer `torch.compile` to transformer decoder
layers (with NO\_REENTRANT activation checkpointing inside each compiled layer).
Skips whole-model compile so that checkpoint loading does not produce
`_orig_mod` key-prefix mismatches.

---

**`enable_fsdp2_prefetch`** `bool = False`

Enable explicit forward/backward prefetch chains
between FSDP2 sharded layers.  Default `True`.

---

**`fsdp2_backward_prefetch_depth`** `int = 2`

Number of FSDP units to prefetch during
backward pass.  `2` hides AllGather behind compute; `1` reduces peak
memory at a small throughput cost.  Default `2`.

---

**`fsdp2_forward_prefetch_depth`** `int = 1`

Number of FSDP units to prefetch during
forward pass.  Default `1`.

---

**`mp_policy`** `MixedPrecisionPolicy | None`

MixedPrecisionPolicy for FSDP2.
If `None` (default), uses bf16 forward/backward compute with fp32
gradient reduction, preserving the dtypes of module inputs and outputs.
Pair this with `model.torch_dtype: float32` for
the Megatron-style fp32 master-weights pattern. Override from YAML
using the `_target_` pattern::

mp\_policy:
*target*: torch.distributed.fsdp.MixedPrecisionPolicy
param\_dtype: bfloat16
reduce\_dtype: float32
output\_dtype: null
cast\_forward\_inputs: false

See `docs/guides/mixed-precision-training.md` for the full set of recommended
patterns and the bf16-storage trap.

---

**`multimodal`** `MultimodalDistributedConfig = field(default_factory=MultimodalDistributedConfig)`

Policies for resolved multimodal modules that use the text
model's distributed axes.

---

**`offload_policy`** `CPUOffloadPolicy | None = None`

CPUOffloadPolicy for CPU offloading.

---

**`patch_is_packed_sequence`** `bool = False`

Patch `transformers._is_packed_sequence` to always
return Python `False`. This does two things: (1) removes a CPU-GPU sync per
attention layer (`aten::is_nonzero` triggered by HF when batch\_size==1), and
(2) ensures static attention shapes for `torch.compile`. Safe for standard
(non-packed) training only. Disable if using packed-sequence training
(position\_ids that reset to 0 mid-sequence). Default `False`.

---

**`reshard_after_forward`** `bool | None = None`

Override layer-level FSDP2 resharding.
`None` preserves AutoModel's heuristic: pipeline-parallel layers do
not reshard after forward, while non-pipeline layers reshard all but
the last layer. Set `False` for a ZeRO-2-like benchmark where
gathered parameters stay resident after forward. Set `True` to force
resharding everywhere, including pipeline-parallel layers, which may
reduce throughput by adding per-microbatch all-gathers.

---

**`sequence_parallel`** `bool = False`

Enable sequence parallelism in TP plan.

---

**`tp_plan`** `dict | str | None = None`

Custom TP plan. If None, auto-selected based on model type.

---

```python
nemo_automodel.components.distributed.config.FSDP2Config.__post_init__() -> None
```

```python
nemo_automodel.components.distributed.config.FSDP2Config.to_dict() -> typing.Dict[str, typing.Any]
```

Convert config to dictionary (shallow, preserves policy objects).

```python
class nemo_automodel.components.distributed.config.MegatronFSDPConfig(
    megatron_fsdp_unit_modules: list[str] | None = None,
    zero_dp_strategy: int = 3,
    init_fsdp_with_meta_device: bool = False,
    grad_reduce_in_fp32: bool = False,
    preserve_fp32_weights: bool = False,
    overlap_grad_reduce: bool = True,
    overlap_param_gather: bool = True,
    check_for_nan_in_grad: bool = True,
    report_nan_in_param_grad: bool = False,
    average_in_collective: bool = False,
    disable_bucketing: bool = False,
    calculate_per_token_loss: bool = False,
    keep_fp8_transpose_cache: bool = False,
    nccl_ub: bool = False,
    fsdp_double_buffer: bool = False,
    activation_checkpointing: bool = False
)
```

Dataclass

Additional configuration for MegatronFSDP distributed training.

Note: Size parameters (dp\_size, tp\_size, cp\_size) are grouped separately in
`ParallelismSizes`. MegatronFSDP does not
support pp\_size, dp\_replicate\_size, or ep\_size.

**`activation_checkpointing`** `bool = False`

Enable activation checkpointing for transformer
MLP layers to save memory.

---

**`average_in_collective`** `bool = False`

Average in collective if True.

---

**`calculate_per_token_loss`** `bool = False`

Calculate per token loss if True.

---

**`check_for_nan_in_grad`** `bool = True`

Legacy buffer-level gradient NaN check.
BREAKING CHANGE on megatron-fsdp 0.5.0: this flag is a no-op,
preserved only for config compatibility. 0.5.0 removed the
buffer-level NaN check entirely, so gradient NaN checking is now OFF
regardless of this value; a truthy value is dropped with a one-time
warning. The default is kept True for config compatibility, but has
no effect. To restore gradient NaN checking, enable
report\_nan\_in\_param\_grad instead.

---

**`disable_bucketing`** `bool = False`

Disable bucketing if True.

---

**`fsdp_double_buffer`** `bool = False`

Use double buffer if True.

---

**`grad_reduce_in_fp32`** `bool = False`

Reduce gradients in fp32 if True.

---

**`init_fsdp_with_meta_device`** `bool = False`

Initialize MegatronFSDP with meta device if True.

---

**`keep_fp8_transpose_cache`** `bool = False`

Keep fp8 transpose cache when using custom FSDP if True.

---

**`megatron_fsdp_unit_modules`** `list[str] | None = None`

Class paths of the submodules to wrap
as individual MegatronFSDP units. When `None` (the default), the wrap classes
are auto-derived from the model's `_no_split_modules` so the real instantiated
block classes are used regardless of backend (HF or NeMo-custom).

---

**`nccl_ub`** `bool = False`

Use NCCL UBs if True.

---

**`overlap_grad_reduce`** `bool = True`

Overlap gradient reduction if True.

---

**`overlap_param_gather`** `bool = True`

Overlap parameter gathering if True.

---

**`preserve_fp32_weights`** `bool = False`

Preserve fp32 weights if True.

---

**`report_nan_in_param_grad`** `bool = False`

Enable megatron-fsdp 0.5.0's precise
per-parameter gradient NaN check. This is the replacement for the
removed check\_for\_nan\_in\_grad and is OFF by default; enabling it can
significantly reduce training throughput.

---

**`zero_dp_strategy`** `int = 3`

Data parallel sharding strategy.

---

```python
nemo_automodel.components.distributed.config.MegatronFSDPConfig.to_dict() -> typing.Dict[str, typing.Any]
```

Convert config to dictionary (shallow, preserves objects).

```python
class nemo_automodel.components.distributed.mesh.MeshContext(
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None,
    moe_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None,
    process_group: torch.distributed.ProcessGroup | None = None,
    strategy_config: nemo_automodel.components.distributed.config.DistributedStrategyConfig | None = None,
    moe_parallel_config: nemo_automodel.components.distributed.config.MoEParallelizerConfig | None = None,
    activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False,
    reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None,
    cp_ring_degree: int = 1,
    cp_ulysses_degree: int | None = None
)
```

Dataclass

Resolved runtime context for model parallelization.

Parallelism sizes (`tp_size`, `pp_size`, etc.) are **not** stored as
fields; they are `@property` accessors that read directly from the
attached `DeviceMesh` / `moe_mesh`.  When no mesh is present the
properties return safe defaults (`1` for sizes, `None` for dp / hsdp).

All `DeviceMesh` objects passed in must use axis names from
`MeshAxisName`; a `ValueError` is raised on construction if
any unknown name is encountered.

## Lifecycle

1. Recipes parse YAML to obtain sizes and strategy configs.
2. Sizes are passed to `build` to build `DeviceMesh`
   objects.
3. `MeshContext` is created with those meshes; axis names are
   validated automatically in `__post_init__`.

Alternatively, `from_meshes` constructs an instance directly from
`DeviceMesh` objects (used by `NeMoAutoModel.from_pretrained`).

**`activation_checkpointing`** `ActivationCheckpointingMode = False`

Resolved activation-checkpointing mode.

---

**`cp_ring_degree`** `int = 1`

Ring-attention factor of the context-parallel axis.

---

**`cp_size`** `int`

Context-parallel degree (from `device_mesh`, default `1`).

---

**`cp_ulysses_degree`** `int | None = None`

Ulysses factor of the context-parallel axis. `None`
lets the consumer use the full context-parallel axis.

---

**`device_mesh`** `Optional[DeviceMesh] = field(default=None, repr=False)`

Device mesh for distributed training.

---

**`dp_replicate_size`** `int | None`

HSDP replication degree (from `device_mesh`, default `None`).

---

**`dp_shard_size`** `int`

DP shard degree (from `device_mesh`, default `1`).

---

**`dp_size`** `int | None`

Data-parallel degree (from `device_mesh`, default `None`).

---

**`ep_size`** `int`

Expert-parallel degree (from `moe_mesh`, default `1`).

---

**`moe_mesh`** `Optional[DeviceMesh] = field(default=None, repr=False)`

MoE-specific device mesh.

---

**`moe_parallel_config`** `MoEParallelizerConfig | None = None`

Optional MoE-specific parallelization policy.

---

**`pp_enabled`** `bool`

`True` when `pp_size > 1`.

---

**`pp_size`** `int`

Pipeline-parallel degree (from `device_mesh`, default `1`).

---

**`process_group`** `ProcessGroup | None = field(default=None, repr=False)`

Optional model-local group for recipe-level collectives
that must not involve ranks outside this mesh.

---

**`reapply_trainability`** `Callable[[nn.Module], None] | None = field(default=None, repr=False)`

Optional callback run after model surgery and before
distributed wrapping.

---

**`strategy_config`** `DistributedStrategyConfig | None = None`

Resolved DDP, FSDP2, or Megatron-FSDP policy.

---

**`tp_size`** `int`

Tensor-parallel degree (from `device_mesh`, default `1`).

---

```python
nemo_automodel.components.distributed.mesh.MeshContext.__post_init__() -> None
```

```python
nemo_automodel.components.distributed.mesh.MeshContext._dp_axis_names() -> typing.Tuple[str, ...]
```

DP axis names for FSDP mesh slicing.

```python
nemo_automodel.components.distributed.mesh.MeshContext.build(
    strategy_config: nemo_automodel.components.distributed.config.DistributedStrategyConfig,
    parallelism_sizes: nemo_automodel.components.distributed.mesh.ParallelismSizes | None = None,
    moe_parallel_config: nemo_automodel.components.distributed.config.MoEParallelizerConfig | None = None,
    activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False,
    reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None,
    cp_ring_degree: int = 1,
    cp_ulysses_degree: int | None = None,
    world_size: int | None = None,
    timeout_minutes: int | None = None,
    ranks: typing.Sequence[int] | None = None
) -> nemo_automodel.components.distributed.mesh.MeshContext
```

classmethod

Build a resolved `MeshContext` from sizes and execution policy.

**Parameters:**

**`strategy_config`** `DistributedStrategyConfig`

Already-instantiated distributed strategy config.

---

**`parallelism_sizes`** `ParallelismSizes | None` — default: None

Requested data, tensor, pipeline, context, and expert
parallelism sizes. If `None`, defaults to no parallelism with
DP inferred from `world_size`.

---

**`moe_parallel_config`** `MoEParallelizerConfig | None` — default: None

Optional MoE-specific parallelization policy.

---

**`activation_checkpointing`** `ActivationCheckpointingMode` — default: False

Activation-checkpointing mode applied by
the model sidecar.

---

**`reapply_trainability`** `Callable[[nn.Module], None] | None` — default: None

Optional callback run after model surgery and
before distributed wrapping.

---

**`cp_ring_degree`** `int` — default: 1

Ring-attention factor of the context-parallel axis.

---

**`cp_ulysses_degree`** `int | None` — default: None

Ulysses factor of the context-parallel axis.

---

**`world_size`** `int | None` — default: None

Total process count. If `None`, inferred from the
distributed environment.

---

**`timeout_minutes`** `int | None` — default: None

Optional timeout for process groups created by
`DeviceMesh` axes.

---

**`ranks`** `Sequence[int] | None` — default: None

Optional ordered global ranks to use for the mesh. When
omitted, the mesh uses every rank in the default process group.

---

```python
nemo_automodel.components.distributed.mesh.MeshContext.from_meshes(
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh],
    moe_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None,
    strategy_config: nemo_automodel.components.distributed.config.DistributedStrategyConfig | None = None,
    moe_parallel_config: nemo_automodel.components.distributed.config.MoEParallelizerConfig | None = None,
    activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False,
    reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None,
    cp_ring_degree: int = 1,
    cp_ulysses_degree: int | None = None
) -> nemo_automodel.components.distributed.mesh.MeshContext
```

classmethod

Build a `MeshContext` from `DeviceMesh` objects.

This is the entry-point used by `NeMoAutoModel.from_pretrained` /
`from_config` where the caller has raw meshes rather than a parsed
YAML config.

**Parameters:**

**`device_mesh`** `Optional[DeviceMesh]`

Primary device mesh.

---

**`moe_mesh`** `Optional[DeviceMesh]` — default: None

Optional MoE-specific device mesh.

---

**`strategy_config`** `DistributedStrategyConfig | None` — default: None

Resolved DDP, FSDP2, or Megatron-FSDP policy.

---

**`moe_parallel_config`** `MoEParallelizerConfig | None` — default: None

Optional MoE-specific parallelization policy.

---

**`activation_checkpointing`** `ActivationCheckpointingMode` — default: False

Activation-checkpointing mode applied by
the model sidecar.

---

**`reapply_trainability`** `Callable[[nn.Module], None] | None` — default: None

Optional callback run after model surgery and
before distributed wrapping.

---

**`cp_ring_degree`** `int` — default: 1

Ring-attention factor of the context-parallel axis.

---

**`cp_ulysses_degree`** `int | None` — default: None

Ulysses factor of the context-parallel axis.

---

```python
nemo_automodel.components.distributed.mesh.MeshContext.parallelize_axis_kwargs() -> typing.Dict[str, object]
```

Axis-name kwargs for `parallelize_fn` (EP/FSDP, no `pp_axis_name`).

```python
nemo_automodel.components.distributed.mesh.MeshContext.pipeline_axis_kwargs() -> typing.Dict[str, object]
```

Axis-name kwargs for `AutoPipeline`.

```python
class nemo_automodel.components.distributed.config.MoEParallelizerConfig(
    ignore_router_for_ac: bool = True,
    reshard_after_forward: bool = False,
    lm_head_precision: typing.Union[str, torch.dtype] | None = None,
    wrap_outer_model: bool = True,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None = None
)
```

Dataclass

Configuration for MoE model parallelization (EP + FSDP settings).

**`ignore_router_for_ac`** `bool = True`

---

**`lm_head_precision`** `Union[str, dtype] | None = None`

---

**`mp_policy`** `MixedPrecisionPolicy | None = None`

---

**`reshard_after_forward`** `bool = False`

---

**`wrap_outer_model`** `bool = True`

---

```python
nemo_automodel.components.distributed.config.MoEParallelizerConfig.to_dict() -> typing.Dict[str, typing.Any]
```

```python
class nemo_automodel.components.distributed.parallelizer.ModelParallelizer()
```

Single model-owned parallelization sidecar contract.

```python
nemo_automodel.components.distributed.parallelizer.ModelParallelizer._apply(
    model: torch.nn.Module,
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None = None,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None = None,
    sequence_parallel: bool = False,
    activation_checkpointing: bool = False,
    tp_shard_plan: typing.Union[typing.Dict[str, torch.distributed.tensor.parallel.ParallelStyle], str] | None = None,
    dp_replicate_mesh_name: str = 'dp_replicate',
    dp_shard_cp_mesh_name: str = 'dp_shard_cp',
    tp_mesh_name: str = 'tp',
    enable_async_tensor_parallel: bool = False,
    enable_compile: bool = False,
    enable_fsdp2_prefetch: bool = True,
    fsdp2_backward_prefetch_depth: int = 2,
    fsdp2_forward_prefetch_depth: int = 1,
    reshard_after_forward: bool | None = None,
    activation_checkpointing_scope: nemo_automodel.components.distributed.config.ActivationCheckpointingScope | None = 'all',
    frozen_multimodal_sharding: nemo_automodel.components.distributed.multimodal_fsdp.FrozenMultimodalSharding = 'root',
    reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None
) -> torch.nn.Module
```

Apply the shared dense FSDP2 implementation.

```python
nemo_automodel.components.distributed.parallelizer.ModelParallelizer._apply_fsdp_sharding(
    module: torch.nn.Module,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None = None,
    enable_fsdp2_prefetch: bool = True,
    fsdp2_backward_prefetch_depth: int = 2,
    fsdp2_forward_prefetch_depth: int = 1,
    reshard_after_forward: bool | None = None,
    frozen_multimodal_sharding: nemo_automodel.components.distributed.multimodal_fsdp.FrozenMultimodalSharding = 'root',
    ignored_multimodal_params: set[torch.nn.Parameter] | None = None
) -> None
```

Wrap the model's submodules into FSDP2 units.

Model-owned sidecar strategies deriving from this class override this
hook to change how parameters are grouped into FSDP units without
reimplementing the surrounding TP/AC/mixed-precision flow. Override
`_fully_shard_module` when a model needs a different primitive.

```python
nemo_automodel.components.distributed.parallelizer.ModelParallelizer._fully_shard_module(
    module: torch.nn.Module,
    kwargs = {}
) -> torch.nn.Module
```

Apply the FSDP2 primitive used by this model sidecar.

```python
nemo_automodel.components.distributed.parallelizer.ModelParallelizer._fully_shard_untied_input_output_embeddings(
    model: torch.nn.Module,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None,
    input_reshard_after_forward: bool
) -> None
```

Give large trainable untied embedding tables independent FSDP buffers.

The generic dense path otherwise leaves both tables in the root FSDP unit.
With fp32 gradient reduction, that unit allocates one contiguous
reduce-scatter input containing both gradients. Keeping the two trainable
leaf modules in separate FSDP units bounds that allocation by the larger
table instead of their sum. This pass skips tied tables to preserve aliasing,
and frozen tables because they have no gradient communication buffer
to split. Tables already wrapped as FSDP units retain their existing
reshard policy. With an explicit `reshard_after_forward=True`, a
container-hosted head reshards while a newly split top-level head stays
gathered. FusedLinearCrossEntropy handles the sharded weight with correct
gradients, at the cost of an extra all-gather.

**Parameters:**

**`model`** `nn.Module`

Model whose input and output embedding modules may be sharded.

---

**`mesh`** `DeviceMesh`

Device mesh that owns the FSDP shards.

---

**`mp_policy`** `MixedPrecisionPolicy | None`

Mixed-precision policy used by the surrounding FSDP units.

---

**`offload_policy`** `OffloadPolicy | None`

Optional offload policy used by the surrounding FSDP
units.

---

**`input_reshard_after_forward`** `bool`

Whether the input embedding unit reshards
its parameters after forward.

---

```python
nemo_automodel.components.distributed.parallelizer.ModelParallelizer._use_full_layer_activation_checkpointing(
    model: torch.nn.Module
) -> bool
```

Return whether this model safely opts into whole-layer checkpointing.

```python
nemo_automodel.components.distributed.parallelizer.ModelParallelizer._validate_tp_mesh(
    model: torch.nn.Module,
    tp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> None
```

Validate the model's attention topology against its TP mesh.

```python
nemo_automodel.components.distributed.parallelizer.ModelParallelizer.parallelize(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module
```

Apply every requested parallelism and return the parallelized model.

```python
class nemo_automodel.components.distributed.config.MultimodalDistributedConfig(
    frozen_sharding: nemo_automodel.components.distributed.multimodal_fsdp.FrozenMultimodalSharding = 'root',
    vision: nemo_automodel.components.distributed.config.MultimodalVisionConfig = MultimodalVisionConfig()
)
```

Dataclass

Distributed policies for resolved multimodal modules.

**`frozen_sharding`** `FrozenMultimodalSharding = 'root'`

Controls fully frozen multimodal modules such as
vision/audio towers and projectors. `"root"` (default) keeps
their parameters in an always-run outer FSDP root, which is safe
when modality execution differs across ranks. `"per_layer"`
uses normal layer/container FSDP units and requires every rank in
the FSDP group to execute or skip the module identically on every
microbatch. `"replicate"` excludes the frozen parameters from
FSDP roots so each rank keeps a full copy. Modules with any
trainable parameters use normal layer/container sharding
regardless of this setting.

---

**`vision`** `MultimodalVisionConfig = field(default_factory=MultimodalVisionConfig)`

Policies that apply specifically to resolved vision modules.

---

```python
nemo_automodel.components.distributed.config.MultimodalDistributedConfig.__post_init__() -> None
```

```python
class nemo_automodel.components.distributed.mesh.ParallelismSizes(
    dp_size: int | None = None,
    dp_replicate_size: int | None = None,
    tp_size: int = 1,
    pp_size: int = 1,
    cp_size: int = 1,
    ep_size: int = 1
)
```

Dataclass

Build-time requested parallelism sizes.

This is durable user intent, not runtime topology. `MeshContext` derives
its size properties from live `DeviceMesh` objects after build.

**`cp_size`** `int = 1`

---

**`dp_replicate_size`** `int | None = None`

---

**`dp_size`** `int | None = None`

---

**`ep_size`** `int = 1`

---

**`pp_size`** `int = 1`

---

**`tp_size`** `int = 1`

---

```python
class nemo_automodel.components.distributed.pipelining.config.PipelineConfig(
    pp_schedule: str | None = '1f1b',
    pp_schedule_csv: str | None = None,
    pp_microbatch_size: int = 1,
    pp_batch_size: int = 1,
    layers_per_stage: int | None = None,
    round_virtual_stages_to_pp_multiple: typing.Literal['up', 'down'] | None = None,
    module_fqns_per_model_part: typing.List[typing.List[str]] | None = None,
    patch_inner_model: bool = True,
    patch_causal_lm_model: bool = True,
    patch_stage_backward_maybe_with_nosync: bool = False,
    dtype: torch.dtype | None = None,
    scale_grads_in_schedule: bool = False,
    loss_fn: typing.Callable | None = None,
    pp_seq_len: int | None = None,
    pp_recv_buffer_pool: bool = False,
    pp_recv_buffer_pool_slack: int = 2
)
```

Dataclass

Configuration for pipeline parallel training.

Note: Device mesh (world\_mesh, moe\_mesh) is passed separately on the
from\_pretrained/from\_config method signature. Pipeline parallelism is
enabled when pp\_size > 1. Axis names are inferred automatically from
the device mesh structure.

**`dtype`** `dtype | None = None`

Data type for pipeline computation.
If None, uses the model's default dtype.

---

**`layers_per_stage`** `int | None = None`

Number of transformer layers per
pipeline stage. If None, layers are split evenly across stages.

---

**`loss_fn`** `Callable | None = None`

Loss function used for pipeline training.
Required when pipeline is enabled. The function signature should be
compatible with the model's output format.

---

**`module_fqns_per_model_part`** `List[List[str]] | None = None`

Explicit
specification of which module FQNs belong to each model part/stage.
If provided, overrides automatic layer splitting.

---

**`patch_causal_lm_model`** `bool = True`

Apply pipeline patches to the CausalLM
wrapper model. Defaults to True.

---

**`patch_inner_model`** `bool = True`

Apply pipeline patches to the inner model
(e.g., the base transformer in a CausalLM wrapper). Defaults to True.

---

**`patch_stage_backward_maybe_with_nosync`** `bool = False`

Patch stage backward to
use no\_sync context for gradient accumulation efficiency. Useful
when combining PP with FSDP.

---

**`pp_batch_size`** `int = 1`

Total batch size per pipeline stage. Must be
divisible by pp\_microbatch\_size.

---

**`pp_microbatch_size`** `int = 1`

Size of each microbatch for pipeline execution.
pp\_batch\_size must be divisible by pp\_microbatch\_size.

---

**`pp_recv_buffer_pool`** `bool = False`

Pool the per-microbatch P2P recv buffers into
a ring sized to the 1F1B in-flight depth instead of pre-allocating one
buffer set per microbatch and direction. Cuts pipeline-stage buffer
memory roughly by num\_microbatches / (num\_stages - stage\_index + slack)
on early/middle stages. For native interleaved 1F1B, receives are pooled using the lowered
schedule's actual lifetimes. Inputs are copied into independent
leaves before forward, and gradients are copied before backward
to protect checkpoint recomputation and escaping aliases.
Other schedules are ignored with a warning. Defaults to False.

---

**`pp_recv_buffer_pool_slack`** `int = 2`

Extra buffer sets beyond the 1F1B
in-flight depth when pp\_recv\_buffer\_pool is enabled. Defaults to 2.

---

**`pp_schedule`** `str | None = '1f1b'`

Pipeline schedule type. Supported values:
"1f1b" (one-forward-one-backward), "gpipe", "interleaved\_1f1b",
"looped\_bfs", "dfs", "v\_schedule", "zero\_bubble". Defaults to "1f1b".

---

**`pp_schedule_csv`** `str | None = None`

Path to a CSV file defining a custom
pipeline schedule. If provided, overrides pp\_schedule.

---

**`pp_seq_len`** `int | None = None`

Sequence length hint for pipeline parallel
shape inference. If None, it will be inferred from the dataset config.

---

**`round_virtual_stages_to_pp_multiple`** `Literal['up', 'down'] | None = None`

When using virtual stages (interleaved schedules), round the number
of virtual stages to a multiple of pp\_size. "up" rounds up, "down"
rounds down. If None, no rounding is applied.

---

**`scale_grads_in_schedule`** `bool = False`

Scale gradients within the pipeline
schedule (by 1/n\_microbatches). If False, gradients must be scaled
externally. Defaults to False.

---

```python
nemo_automodel.components.distributed.pipelining.config.PipelineConfig.__post_init__() -> None
```

Validate the recv-buffer-pool knobs.

```python
nemo_automodel.components.distributed.pipelining.config.PipelineConfig.to_dict() -> typing.Dict[str, typing.Any]
```

Convert config to dictionary.

```python
nemo_automodel.components.distributed.init_utils.initialize_distributed(
    backend,
    timeout_minutes = 1
)
```

Initialize the torch.distributed environment and core model parallel infrastructure.

This function sets the device based on the local rank, configures the process group,
and calls torch.distributed.init\_process\_group with the appropriate parameters.
It also registers a cleanup function to properly destroy the process group at exit.

**Parameters:**

**`backend`** `str`

The backend to use for torch.distributed (e.g., 'nccl').

---

**`timeout_minutes`** `int` — default: 1

Timeout (in minutes) for distributed initialization. Defaults to 1.

---

**Returns:**

An instance containing the distributed environment configuration.