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

## Module Contents

### Classes

| Name                                                                                              | Description                                                                   |
| ------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`MegatronFSDPManager`](#nemo_automodel-components-distributed-megatron_fsdp-MegatronFSDPManager) | Manager for parallelizing models using MegatronFSDP with TP, DP, CP sharding. |

### Functions

| Name                                                                                                                        | Description                                                                       |
| --------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- |
| [`fully_shard_optimizer`](#nemo_automodel-components-distributed-megatron_fsdp-fully_shard_optimizer)                       | Register an already-built optimizer with a MegatronFSDP-wrapped model.            |
| [`maybe_shard_optimizer`](#nemo_automodel-components-distributed-megatron_fsdp-maybe_shard_optimizer)                       | Shard the optimizer with Megatron-FSDP when the strategy requires it.             |
| [`restore_distributed_param_attrs`](#nemo_automodel-components-distributed-megatron_fsdp-restore_distributed_param_attrs)   | Re-apply Megatron-FSDP per-parameter attributes dropped by a post-wrap rebuild.   |
| [`snapshot_distributed_param_attrs`](#nemo_automodel-components-distributed-megatron_fsdp-snapshot_distributed_param_attrs) | Snapshot the per-parameter attributes Megatron-FSDP stamps on distributed params. |

### Data

[`HAS_MEGATRON_FSDP`](#nemo_automodel-components-distributed-megatron_fsdp-HAS_MEGATRON_FSDP)

[`logger`](#nemo_automodel-components-distributed-megatron_fsdp-logger)

### API

```python
class nemo_automodel.components.distributed.megatron_fsdp.MegatronFSDPManager(
    config: nemo_automodel.components.distributed.config.MegatronFSDPConfig,
    device_mesh: torch.distributed.device_mesh.DeviceMesh
)
```

Manager for parallelizing models using MegatronFSDP with TP, DP, CP sharding.

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

The device mesh must be created externally and passed in.

**Parameters:**

**`config`** `MegatronFSDPConfig`

Configuration for MegatronFSDP distributed training.

---

**`device_mesh`** `DeviceMesh`

Device mesh for distributed operations.

---

**`activation_checkpointing`** `= config.activation_checkpointing`

---

**`average_in_collective`** `= config.average_in_collective`

---

**`calculate_per_token_loss`** `= config.calculate_per_token_loss`

---

**`check_for_nan_in_grad`** `= config.check_for_nan_in_grad`

---

**`disable_bucketing`** `= config.disable_bucketing`

---

**`fsdp_double_buffer`** `= config.fsdp_double_buffer`

---

**`grad_reduce_in_fp32`** `= config.grad_reduce_in_fp32`

---

**`init_fsdp_with_meta_device`** `= config.init_fsdp_with_meta_device`

---

**`keep_fp8_transpose_cache`** `= config.keep_fp8_transpose_cache`

---

**`megatron_fsdp_unit_modules`** `= config.megatron_fsdp_unit_modules`

---

**`nccl_ub`** `= config.nccl_ub`

---

**`overlap_grad_reduce`** `= config.overlap_grad_reduce`

---

**`overlap_param_gather`** `= config.overlap_param_gather`

---

**`preserve_fp32_weights`** `= config.preserve_fp32_weights`

---

**`report_nan_in_param_grad`** `= config.report_nan_in_param_grad`

---

**`zero_dp_strategy`** `= config.zero_dp_strategy`

---

```python
nemo_automodel.components.distributed.megatron_fsdp.MegatronFSDPManager.parallelize(
    model,
    optimizer = None
)
```

Parallelizes the given model using MegatronFSDP and TP sharding strategies.

**Parameters:**

**`model`**

The model to be parallelized.

---

**`optimizer`** — default: None

The optimizer for the model. If None, user needs to call
model.finish\_grad\_sync() before optimizer.step(),
model.install\_optimized\_model\_weights() and model.zero\_grad\_buffer()
after optimizer.zero\_grad().

---

**Returns:**

(parallelized\_model, optimizer)

```python
nemo_automodel.components.distributed.megatron_fsdp.fully_shard_optimizer(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    preproc_state_dict_for_dcp_ckpt: bool = True
) -> torch.optim.Optimizer
```

Register an already-built optimizer with a MegatronFSDP-wrapped model.

Megatron-FSDP 0.5.0's `fully_shard_optimizer` recovers the owning
`MegatronFSDP` from a `_megatron_fsdp_model` attribute that
`MegatronFSDP.__init__` stamps onto each distributed `Parameter`. That
attribute is a plain Python attribute and does not survive operations that
rebuild `Parameter` objects (e.g. the dtype/device cast the `from_pretrained`
load path performs after wrapping). The combined `fully_shard(model, optimizer)`
entry point never hits this because it registers the optimizer in the same call,
before any such op runs; the recipe's separate build-model-then-build-optimizer
order does, leaving `fully_shard_optimizer` unable to find the reference and
aborting before the first optimizer step. Re-stamp the reference (mirroring the
wheel's own `__init__` logic) on the current distributed params right before
deferred sharding so the separate sequence matches the combined entry point.

```python
nemo_automodel.components.distributed.megatron_fsdp.maybe_shard_optimizer(
    model_part: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    distributed_config: nemo_automodel.components.distributed.config.DistributedConfig | None,
    allow: bool = True
) -> torch.optim.Optimizer
```

Shard the optimizer with Megatron-FSDP when the strategy requires it.

Returns the optimizer unchanged unless `distributed_config` is a
:class:`MegatronFSDPConfig` running in a distributed (world size > 1) job.

**Parameters:**

**`model_part`** `nn.Module`

The (already sharded) model part the optimizer belongs to.

---

**`optimizer`** `torch.optim.Optimizer`

The optimizer to (optionally) shard.

---

**`distributed_config`** `DistributedConfig | None`

Distributed strategy config; only triggers sharding
when it is a :class:`MegatronFSDPConfig`.

---

**`allow`** `bool` — default: True

Guard for optimizers incompatible with Megatron-FSDP sharding
(e.g. Dion); asserts when sharding would otherwise apply.

---

```python
nemo_automodel.components.distributed.megatron_fsdp.restore_distributed_param_attrs(
    model: torch.nn.Module,
    snapshot: dict[str, dict] | None
) -> None
```

Re-apply Megatron-FSDP per-parameter attributes dropped by a post-wrap rebuild.

Companion to :func:`snapshot_distributed_param_attrs`. For each current parameter
(matched by name, since the rebuild replaced the objects) it restores any snapshot
attribute the rebuilt parameter is missing, following the fix suggested by the
Megatron-FSDP maintainer on NVIDIA/Megatron-LM#5790: only attributes absent on the
new parameter are copied, so genuinely re-derived state is never clobbered.

**Parameters:**

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

The Megatron-FSDP-wrapped model whose parameters were rebuilt.

---

**`snapshot`** `dict[str, dict] | None`

The mapping returned by :func:`snapshot_distributed_param_attrs`, or
`None` (no-op).

---

```python
nemo_automodel.components.distributed.megatron_fsdp.snapshot_distributed_param_attrs(
    model: torch.nn.Module
) -> dict[str, dict] | None
```

Snapshot the per-parameter attributes Megatron-FSDP stamps on distributed params.

`MegatronFSDP.__init__` decorates each distributed `Parameter` with plain
Python attributes that later training steps depend on: `_megatron_fsdp_model`
(owning-model back-ref used by :func:`fully_shard_optimizer`), `_is_shared`
(set on tied parameters so `_grad_acc` routes their gradients through the root
hook instead of double-accumulating), `orig_param`/`megatron_fsdp_dist_index`/
`megatron_fsdp_slice` and the `reset_attribute` closure. These attributes live
in the `Parameter.__dict__` and are silently dropped by any post-wrap operation
that rebuilds `Parameter` objects -- e.g. the `from_pretrained` checkpoint
reload and the `lm_head` re-tie the recipe performs after wrapping. Capture them
here, keyed by parameter name (object identity does not survive the rebuild), so
:func:`restore_distributed_param_attrs` can re-apply them afterwards.

`remove_duplicate=False` is required so tied parameters (e.g. `lm_head.weight`
aliasing `model.embed_tokens.weight`) are captured under every name they appear
under, including the `_is_shared` marker Megatron-FSDP places on the tied alias.

**Parameters:**

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

The (possibly Megatron-FSDP-wrapped) model to snapshot.

---

**Returns:** `dict[str, dict] | None`

A mapping from parameter name to a copy of its `__dict__`, or `None` when

```python
nemo_automodel.components.distributed.megatron_fsdp.HAS_MEGATRON_FSDP = True
```

```python
nemo_automodel.components.distributed.megatron_fsdp.logger = logging.getLogger(__name__)
```