> 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.fsdp2

## Module Contents

### Classes

| Name                                                                        | Description                                                 |
| --------------------------------------------------------------------------- | ----------------------------------------------------------- |
| [`FSDP2Manager`](#nemo_automodel-components-distributed-fsdp2-FSDP2Manager) | Deprecated compatibility wrapper for FSDP2 parallelization. |

### Functions

| Name                                                                                                                            | Description                                                                     |
| ------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`_patch_is_packed_sequence_for_training`](#nemo_automodel-components-distributed-fsdp2-_patch_is_packed_sequence_for_training) | Eliminate CPU-GPU sync from flash attention for standard (non-packed) training. |
| [`fsdp2_sharding_enabled`](#nemo_automodel-components-distributed-fsdp2-fsdp2_sharding_enabled)                                 | Report whether `ModelParallelizer` shards the model for this mesh.              |

### API

```python
class nemo_automodel.components.distributed.fsdp2.FSDP2Manager(
    config: nemo_automodel.components.distributed.config.FSDP2Config,
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    moe_config: nemo_automodel.components.distributed.config.MoEParallelizerConfig | None = None
)
```

Deprecated compatibility wrapper for FSDP2 parallelization.

.. deprecated:: 0.7
Pass `FSDP2Config` through the config-driven infrastructure and
provide model-specific behavior with `ModelParallelizer`. This
compatibility class is scheduled for removal in 0.8.

This manager applies parallelization to the model using a prescribed
TP sharding plan. It supports mixed precision and CPU offloading options.

The device mesh must be created externally and passed in.

**Parameters:**

**`config`** `FSDP2Config`

Configuration for FSDP2 distributed training.

---

**`device_mesh`** `DeviceMesh`

Device mesh for distributed operations.

---

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

Optional device mesh for expert parallelism.

---

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

Optional expert-parallel policy included in the model's
`MeshContext`.

---

```python
nemo_automodel.components.distributed.fsdp2.FSDP2Manager.__getattr__(
    name
)
```

```python
nemo_automodel.components.distributed.fsdp2.FSDP2Manager.maybe_compile(
    model
)
```

Apply per-layer compile after sharding, alongside whole-model compile\_model().

```python
nemo_automodel.components.distributed.fsdp2.FSDP2Manager.parallelize(
    model: torch.nn.Module,
    reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None
) -> torch.nn.Module
```

Parallelizes the given model using FSDP2 and TP sharding strategies.

**Parameters:**

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

The model to be parallelized.

---

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

Optional callback that re-resolves parameter
trainability after model surgery and before FSDP construction.

---

**Returns:** `nn.Module`

The parallelized model.

```python
nemo_automodel.components.distributed.fsdp2._patch_is_packed_sequence_for_training() -> None
```

Eliminate CPU-GPU sync from flash attention for standard (non-packed) training.

transformers.\_is\_packed\_sequence() returns a GPU bool scalar when batch\_size==1,
which causes Python's `if` to call aten::is\_nonzero — a CPU-GPU sync — once per
attention layer per forward pass.  With FSDP+TP+gradient-checkpointing this fires
hundreds of times per iteration.

For standard (non-packed) training sequences are never packed, so returning the
Python False immediately is both correct and avoids the sync.  Do NOT apply this
patch when using packed-sequence training (multiple sequences concatenated into one
tensor with position\_ids that reset to 0 mid-sequence).

```python
nemo_automodel.components.distributed.fsdp2.fsdp2_sharding_enabled(
    device_mesh: torch.distributed.device_mesh.DeviceMesh
) -> bool
```

Report whether `ModelParallelizer` shards the model for this mesh.

Parallelization is skipped on a single-rank world or a single-element mesh, which
also skips every side effect of `fully_shard` — most importantly the
`MixedPrecisionPolicy` cast of parameters to the compute dtype. Callers that
depend on that cast must check this instead of assuming FSDP2 is active.

**Parameters:**

**`device_mesh`** `DeviceMesh`

Device mesh from the model's parallelization context.

---

**Returns:** `bool`

True when `fully_shard` is applied, False when parallelization is skipped.