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

## Module Contents

### Classes

| Name                                                                                         | Description                                          |
| -------------------------------------------------------------------------------------------- | ---------------------------------------------------- |
| [`ModelParallelizer`](#nemo_automodel-components-distributed-parallelizer-ModelParallelizer) | Single model-owned parallelization sidecar contract. |

### Functions

| Name                                                                                                                                                         | Description                                                                                      |
| ------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------ |
| [`_apply_per_layer_compile`](#nemo_automodel-components-distributed-parallelizer-_apply_per_layer_compile)                                                   | Compile each decoder layer in-place after FSDP2 sharding.                                        |
| [`_attention_config`](#nemo_automodel-components-distributed-parallelizer-_attention_config)                                                                 | Resolve the language-model config structurally for TP validation.                                |
| [`_attention_is_head_sharded`](#nemo_automodel-components-distributed-parallelizer-_attention_is_head_sharded)                                               | Return True when the TP plan column-wise shards any QKV attention projection.                    |
| [`_dedupe_layers`](#nemo_automodel-components-distributed-parallelizer-_dedupe_layers)                                                                       | -                                                                                                |
| [`_derive_megatron_fsdp_unit_modules`](#nemo_automodel-components-distributed-parallelizer-_derive_megatron_fsdp_unit_modules)                               | Derive the MegatronFSDP wrap classes from a model's `_no_split_modules`.                         |
| [`_discover_layer_groups`](#nemo_automodel-components-distributed-parallelizer-_discover_layer_groups)                                                       | Discover language, vision, and audio layer containers structurally.                              |
| [`_extend_layers`](#nemo_automodel-components-distributed-parallelizer-_extend_layers)                                                                       | -                                                                                                |
| [`_extract_model_layer_groups`](#nemo_automodel-components-distributed-parallelizer-_extract_model_layer_groups)                                             | Extract transformer layers grouped by model role.                                                |
| [`_extract_model_layers`](#nemo_automodel-components-distributed-parallelizer-_extract_model_layers)                                                         | Extract layers from different model architectures for parallelization.                           |
| [`_filter_layer_groups_for_activation_checkpointing`](#nemo_automodel-components-distributed-parallelizer-_filter_layer_groups_for_activation_checkpointing) | Select trainable activation-checkpointed layers from grouped model layers.                       |
| [`_get_input_output_embeddings`](#nemo_automodel-components-distributed-parallelizer-_get_input_output_embeddings)                                           | Resolve the optional Hugging Face embedding getters.                                             |
| [`_get_parallel_plan`](#nemo_automodel-components-distributed-parallelizer-_get_parallel_plan)                                                               | Select the tensor-parallel plan for the given model.                                             |
| [`_has_trainable_parameters`](#nemo_automodel-components-distributed-parallelizer-_has_trainable_parameters)                                                 | -                                                                                                |
| [`_kv_sharing_survives_checkpoint_replay`](#nemo_automodel-components-distributed-parallelizer-_kv_sharing_survives_checkpoint_replay)                       | Return whether whole-block activation checkpointing is safe for a KV-shared model.               |
| [`_megatron_fsdp_compat_kwargs`](#nemo_automodel-components-distributed-parallelizer-_megatron_fsdp_compat_kwargs)                                           | Translate the config precision controls to the Megatron-FSDP 0.5.0 API.                          |
| [`_patch_dtensor_spec_hash_for_symint`](#nemo_automodel-components-distributed-parallelizer-_patch_dtensor_spec_hash_for_symint)                             | Fix a crash when torch.compile + DTensor are used together.                                      |
| [`_reduce_attrs`](#nemo_automodel-components-distributed-parallelizer-_reduce_attrs)                                                                         | -                                                                                                |
| [`_should_use_hf_native_gradient_checkpointing`](#nemo_automodel-components-distributed-parallelizer-_should_use_hf_native_gradient_checkpointing)           | Return whether HF-native gradient checkpointing can preserve AutoModel's AC scope.               |
| [`_tp_plan_owner`](#nemo_automodel-components-distributed-parallelizer-_tp_plan_owner)                                                                       | Return the most specific nested Hugging Face TP-plan provider.                                   |
| [`_update_attention_head_counts_for_tp`](#nemo_automodel-components-distributed-parallelizer-_update_attention_head_counts_for_tp)                           | After TP sharding, the Q/K/V outputs are split across ranks (each rank has                       |
| [`_uses_custom_moe_modules`](#nemo_automodel-components-distributed-parallelizer-_uses_custom_moe_modules)                                                   | Return whether `model` contains Automodel's grouped custom-MoE layer.                            |
| [`apply_fsdp2_sharding_recursively`](#nemo_automodel-components-distributed-parallelizer-apply_fsdp2_sharding_recursively)                                   | Recursively apply FSDP2 sharding to modules, with optimizations for ModuleList.                  |
| [`apply_selective_activation_checkpointing`](#nemo_automodel-components-distributed-parallelizer-apply_selective_activation_checkpointing)                   | Apply selective activation checkpointing to `model` end to end.                                  |
| [`fsdp2_strategy_parallelize`](#nemo_automodel-components-distributed-parallelizer-fsdp2_strategy_parallelize)                                               | Apply parallelisms and activation checkpointing to the model.                                    |
| [`get_hf_tp_shard_plan`](#nemo_automodel-components-distributed-parallelizer-get_hf_tp_shard_plan)                                                           | Translate a model-owned Hugging Face TP plan without importing model classes.                    |
| [`get_model_layer_groups`](#nemo_automodel-components-distributed-parallelizer-get_model_layer_groups)                                                       | Return transformer layers grouped by model role (e.g. `language`, `vision`).                     |
| [`import_class_from_path`](#nemo_automodel-components-distributed-parallelizer-import_class_from_path)                                                       | Import a class from a string path (e.g. 'torch.optim.AdamW').                                    |
| [`import_classes_from_paths`](#nemo_automodel-components-distributed-parallelizer-import_classes_from_paths)                                                 | Helper function to import classes from string paths.                                             |
| [`megatron_fsdp_strategy_parallelize`](#nemo_automodel-components-distributed-parallelizer-megatron_fsdp_strategy_parallelize)                               | Apply tensor/data parallelism (MegatronFSDP) and optional activation-checkpointing to the model. |
| [`translate_to_torch_parallel_style`](#nemo_automodel-components-distributed-parallelizer-translate_to_torch_parallel_style)                                 | Translates string descriptions to parallelism plans.                                             |
| [`unshard_fsdp2_model`](#nemo_automodel-components-distributed-parallelizer-unshard_fsdp2_model)                                                             | Explicitly unshard and then reshard the FSDP2 modules. Useful for logprob inference.             |
| [`validate_tp_mesh`](#nemo_automodel-components-distributed-parallelizer-validate_tp_mesh)                                                                   | Validate that every sharded language attention shape divides the TP size.                        |

### Data

[`HAVE_MEGATRON_FSDP`](#nemo_automodel-components-distributed-parallelizer-HAVE_MEGATRON_FSDP)

[`MegatronFSDPMixedPrecisionPolicy`](#nemo_automodel-components-distributed-parallelizer-MegatronFSDPMixedPrecisionPolicy)

[`_MEGATRON_FSDP_050_REQUIRED_MSG`](#nemo_automodel-components-distributed-parallelizer-_MEGATRON_FSDP_050_REQUIRED_MSG)

[`_megatron_fsdp_nan_check_noop_warned`](#nemo_automodel-components-distributed-parallelizer-_megatron_fsdp_nan_check_noop_warned)

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

### API

```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
nemo_automodel.components.distributed.parallelizer._apply_per_layer_compile(
    model: torch.nn.Module
) -> None
```

Compile each decoder layer in-place after FSDP2 sharding.

Compiles at decoder-layer granularity (not sub-module) so that AOT autograd traces
the joint fwd+bwd graph under the training loop's enable\_grad context.  Sub-module
compile (e.g. on mlp alone) would be traced during activation checkpointing's first
forward pass which runs under no\_grad, producing a forward-only graph that drops
LoRA and other trainable-parameter gradients.

Prerequisite: NO\_REENTRANT checkpoint\_wrapper must already be applied to self\_attn
and mlp before FSDP2 sharding (done by `ModelParallelizer`).  This
function only handles the compile step.

Whole-block selective-AC wrappers (tagged with `SELECTIVE_AC_WRAPPER_FLAG`)
are compiled OUTER -- the wrapper itself is compiled so the selective policy
is traced and the partitioner honors its recompute tags. Other layer-level
CheckpointWrappers (e.g. the PP path) are unwrapped and the decoder layer is
compiled directly.

nn.Module.compile() is used instead of torch.compile() to compile in-place without
introducing an \_orig\_mod wrapper, which would add a key prefix and break checkpoint
loading.

\_patch\_dtensor\_spec\_hash\_for\_symint() is called to allow torch.compile with dynamic
shapes to coexist with DTensor's lru\_cache-based sharding propagation.

```python
nemo_automodel.components.distributed.parallelizer._attention_config(
    model: torch.nn.Module
)
```

Resolve the language-model config structurally for TP validation.

```python
nemo_automodel.components.distributed.parallelizer._attention_is_head_sharded(
    model_parallel_plan: dict
) -> bool
```

Return True when the TP plan column-wise shards any QKV attention projection.

When Q/K/V projections use `ColwiseParallel` with sharded output (the
default), each TP rank holds `num_heads / tp_size` heads and the model
config / layer attributes must be updated accordingly.

Plans that keep attention replicated (e.g. Phi-3 with `RowwiseParallel`
on fused QKV and `Replicate` output) should *not* trigger a head-count
update.

```python
nemo_automodel.components.distributed.parallelizer._dedupe_layers(
    layers: typing.Sequence[torch.nn.Module]
) -> typing.List[torch.nn.Module]
```

```python
nemo_automodel.components.distributed.parallelizer._derive_megatron_fsdp_unit_modules(
    model: torch.nn.Module
) -> list[type[torch.nn.Module]]
```

Derive the MegatronFSDP wrap classes from a model's `_no_split_modules`.

Used when a config does not specify `megatron_fsdp_unit_modules`. HF
`PreTrainedModel` and the NeMo custom models both define `_no_split_modules`
as a list of block class *names* (for example `["LlamaDecoderLayer"]`).
Walking `model.modules()` and matching `type(module).__name__` against those
names resolves the actual instantiated classes, so the result is correct for
both the HF backend and the NeMo-custom backend (which use distinct classes that
share the same name). For VLM/MoE models whose top-level `_no_split_modules`
lists several block classes (for example vision and language towers), every
matching class found anywhere in the module tree is collected.

**Parameters:**

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

The (already TP-parallelized) model to be wrapped by MegatronFSDP.

---

**Returns:** `list[type[nn.Module]]`

The de-duplicated list of submodule classes to wrap as MegatronFSDP units,

**Raises:**

* `ValueError`: If the model does not expose a non-empty `_no_split_modules`,
  or if none of those names match an instantiated submodule. Raised with
  an actionable message instead of letting MegatronFSDP later fail with
  `ZeroDivisionError` (`total_fsdp_module=0`) when zero modules are wrapped.

```python
nemo_automodel.components.distributed.parallelizer._discover_layer_groups(
    model: torch.nn.Module
) -> typing.Dict[str, typing.List[torch.nn.Module]]
```

Discover language, vision, and audio layer containers structurally.

```python
nemo_automodel.components.distributed.parallelizer._extend_layers(
    layers: typing.List[torch.nn.Module],
    modules: typing.Sequence[torch.nn.Module]
) -> None
```

```python
nemo_automodel.components.distributed.parallelizer._extract_model_layer_groups(
    model: torch.nn.Module
) -> typing.Dict[str, typing.List[torch.nn.Module]]
```

Extract transformer layers grouped by model role.

```python
nemo_automodel.components.distributed.parallelizer._extract_model_layers(
    model: torch.nn.Module
) -> typing.List[torch.nn.Module]
```

Extract layers from different model architectures for parallelization.

This compatibility wrapper flattens grouped language/vision/audio layers.
New activation-checkpointing code should use `_extract_model_layer_groups`
so scope decisions can be explicit.

```python
nemo_automodel.components.distributed.parallelizer._filter_layer_groups_for_activation_checkpointing(
    layer_groups: typing.Dict[str, typing.List[torch.nn.Module]],
    activation_checkpointing_scope: nemo_automodel.components.distributed.config.ActivationCheckpointingScope | None = 'all'
) -> typing.Tuple[typing.List[torch.nn.Module], typing.Tuple[str, ...]]
```

Select trainable activation-checkpointed layers from grouped model layers.

```python
nemo_automodel.components.distributed.parallelizer._get_input_output_embeddings(
    model: torch.nn.Module
) -> tuple[torch.nn.Module | None, torch.nn.Module | None]
```

Resolve the optional Hugging Face embedding getters.

```python
nemo_automodel.components.distributed.parallelizer._get_parallel_plan(
    model: torch.nn.Module,
    sequence_parallel: bool = False,
    tp_shard_plan: typing.Union[typing.Dict[str, torch.distributed.tensor.parallel.ParallelStyle], str] | None = None,
    tp_size: int = 1
) -> typing.Dict[str, torch.distributed.tensor.parallel.ParallelStyle]
```

Select the tensor-parallel plan for the given model.

Priority order:

1. If `tp_shard_plan` is provided as a dict or import path, use it.
2. If the model type exists in `PARALLELIZE_FUNCTIONS`, use its optimised plan; on failure, fall back to HF plan.
3. Otherwise, prefer the model's HF-native `_tp_plan` (via `get_hf_tp_shard_plan`).
4. Otherwise, fall back to the default base plan.

When `tp_size > 1` and the model falls through to path 4 *and* the
model class was loaded from a custom-code source (HF's
`trust_remote_code=True` path, where the dynamic class lives under
`transformers_modules.*`), this raises `ValueError` instead of
returning the default base plan. On recent PyTorch the default plan's
placements do not populate `shard_order` and trip an internal assert in
`torch.distributed.tensor._redistribute` on the first weight
redistribute, which surfaces to the user as an opaque PyTorch internal
error. Custom-code architectures are the only known-broken case (see
[https://github.com/NVIDIA-NeMo/Automodel/issues/2243](https://github.com/NVIDIA-NeMo/Automodel/issues/2243)); known HF
architectures that happen to fall through (e.g. Mixtral) are left on the
default plan with a warning, since they have been working in practice.

When the model *did* define a `_tp_plan` but `get_hf_tp_shard_plan`
raised while translating it (e.g. styles nemo does not recognize), the
translator's error message is folded into the `ValueError` as a
diagnostic so the user can tell whether to add a `_tp_plan` from
scratch or fix the styles in the one they already have.

```python
nemo_automodel.components.distributed.parallelizer._has_trainable_parameters(
    module: torch.nn.Module
) -> bool
```

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

Return whether whole-block activation checkpointing is safe for a KV-shared model.

`checkpoint_wrapper` replays a whole decoder block during backward with the
arguments the forward saw. Unlike HF's `GradientCheckpointingLayer.__call__`
it cannot drop `past_key_values` from that replay, so every layer that
writes to a cache writes to it a second time. A KV-shared model then needs
both halves to hold:

* the layers that populate the cache -- the *non*-shared ones, which are what
  call `Cache.update()` -- must not accumulate on the replay, or the
  recomputed K/V stops matching the forward;
* the shared layers must still read the K/V their source layer produced.

Neither holds for a model backed by an accumulating `Cache`: the second
`Cache.update()` grows the entry and backward dies with a
`CheckpointError` about changed tensor metadata (observed on native HF
`Gemma3nForCausalLM` with `use_cache=True`). KV-shared models therefore
stay on `apply_submodule_checkpointing`, which leaves attention unwrapped,
by default.

A model that satisfies both halves opts in by setting the class attribute
`kv_sharing_survives_checkpoint_replay = True`. Gemma4 E2B/E4B qualify: a
pass-through holder stands in for the cache, and the shared layers read a
separate store that the replay does not disturb (see
`gemma4_moe/model.py`).

**Parameters:**

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

The model about to be checkpointed.

---

**Returns:** `bool`

Whether the model declares its KV sharing safe under whole-block replay.

```python
nemo_automodel.components.distributed.parallelizer._megatron_fsdp_compat_kwargs(
    shard_fn,
    grad_reduce_in_fp32: bool,
    preserve_fp32_weights: bool,
    check_for_nan_in_grad: bool,
    report_nan_in_param_grad: bool
) -> typing.Dict[str, typing.Any]
```

Translate the config precision controls to the Megatron-FSDP 0.5.0 API.

megatron-fsdp==0.5.0, the only supported release, expresses precision
through a `MixedPrecisionPolicy` plus a more expensive per-parameter NaN
reporter. The reporter stays a separate opt-in rather than being silently
enabled from the legacy buffer-check setting; because 0.5.0 has no
buffer-level NaN check at all, a truthy `check_for_nan_in_grad` is
dropped with a one-time warning that points at `report_nan_in_param_grad`
as the opt-in replacement. Any other `fully_shard` signature — older or
newer releases alike — fails loudly instead of guessing a translation.

```python
nemo_automodel.components.distributed.parallelizer._patch_dtensor_spec_hash_for_symint() -> None
```

Fix a crash when torch.compile + DTensor are used together.

Problem: torch.compile traces with symbolic shapes (SymInt). DTensorSpec hashes
its shape to cache sharding decisions, but SymInt is not hashable -> crash.

Fix: if hashing the shape fails, fall back to hashing only (mesh, placements).
Cache hits are slightly reduced but correctness is unaffected.

```python
nemo_automodel.components.distributed.parallelizer._reduce_attrs(
    model: torch.nn.Module,
    fqns: typing.Sequence[str]
) -> typing.List[torch.nn.Module]
```

```python
nemo_automodel.components.distributed.parallelizer._should_use_hf_native_gradient_checkpointing(
    model: torch.nn.Module,
    layer_groups: typing.Dict[str, typing.List[torch.nn.Module]],
    scopes: typing.Tuple[str, ...],
    enable_compile: bool = False
) -> bool
```

Return whether HF-native gradient checkpointing can preserve AutoModel's AC scope.

```python
nemo_automodel.components.distributed.parallelizer._tp_plan_owner(
    model: torch.nn.Module
) -> tuple[str, torch.nn.Module] | None
```

Return the most specific nested Hugging Face TP-plan provider.

```python
nemo_automodel.components.distributed.parallelizer._update_attention_head_counts_for_tp(
    model: torch.nn.Module,
    tp_size: int
) -> None
```

After TP sharding, the Q/K/V outputs are split across ranks (each rank has
num\_heads/tp\_size heads). Update the config and each attention layer's
num\_heads / num\_key\_value\_heads so the forward uses the local head count
instead of the global one (avoids shape mismatches in .view()).

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

Return whether `model` contains Automodel's grouped custom-MoE layer.

```python
nemo_automodel.components.distributed.parallelizer.apply_fsdp2_sharding_recursively(
    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,
    model_parallelizer: nemo_automodel.components.distributed.parallelizer.ModelParallelizer | None = None,
    frozen_multimodal_sharding: nemo_automodel.components.distributed.multimodal_fsdp.FrozenMultimodalSharding = 'root',
    ignored_multimodal_params: set[torch.nn.Parameter] | None = None
) -> None
```

Recursively apply FSDP2 sharding to modules, with optimizations for ModuleList.

This utility function traverses a model hierarchy and applies FSDP2 sharding
to each module. For ModuleList instances (commonly used for transformer layers),
it applies an optimization where the last layer doesn't reshard after forward
since FSDP2 will prefetch it immediately.

Handles both single-level and nested ModuleList/ModuleDict structures. If a
ModuleList contains other ModuleLists, it will recurse into them instead of trying
to wrap them (since ModuleList doesn't have a forward method).

Note:
This function modifies the module in-place by replacing modules with their
FSDP2-subclassed versions.

**Parameters:**

**`module`** `nn.Module`

The module to apply FSDP sharding to.

---

**`mesh`** `DeviceMesh`

The device mesh for FSDP sharding.

---

**`mp_policy`** `Optional[MixedPrecisionPolicy]`

Mixed precision policy for FSDP.

---

**`offload_policy`** `Optional[OffloadPolicy]` — default: None

CPU offload policy for FSDP.
Defaults to None.

---

**`enable_fsdp2_prefetch`** `bool` — default: True

Enable explicit forward/backward prefetch chains.

---

**`fsdp2_backward_prefetch_depth`** `int` — default: 2

Backward prefetch depth.

---

**`fsdp2_forward_prefetch_depth`** `int` — default: 1

Forward prefetch depth.

---

**`reshard_after_forward`** `Optional[bool]` — default: None

Optional override for each layer's
`fully_shard` reshard behavior.

---

**`model_parallelizer`** `ModelParallelizer | None` — default: None

Optional model sidecar that owns the
FSDP primitive.

---

**`frozen_multimodal_sharding`** `FrozenMultimodalSharding` — default: 'root'

Whether fully frozen multimodal modules are
owned by the root FSDP unit, sharded per layer, or replicated.

---

**`ignored_multimodal_params`** `set[nn.Parameter] | None` — default: None

Accumulator for replicated frozen multimodal
parameters that must be ignored by ancestor FSDP roots.

---

```python
nemo_automodel.components.distributed.parallelizer.apply_selective_activation_checkpointing(
    model: torch.nn.Module,
    enable_compile: bool = False,
    activation_checkpointing_scope: nemo_automodel.components.distributed.config.ActivationCheckpointingScope | None = 'all'
) -> None
```

Apply selective activation checkpointing to `model` end to end.

Standalone entry point (detects KV-sharing, disables `use_cache`, and
wraps transformer blocks) for paths where the FSDP2 parallelize flow is
skipped -- notably single-GPU training.

**Parameters:**

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

The model to checkpoint.

---

**`enable_compile`** `bool` — default: False

Whether per-layer `torch.compile` will be applied.

---

**`activation_checkpointing_scope`** `ActivationCheckpointingScope | None` — default: 'all'

Which extracted layer groups to wrap.

---

```python
nemo_automodel.components.distributed.parallelizer.fsdp2_strategy_parallelize(
    model,
    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 parallelisms and activation checkpointing to the model.

Compatibility entry point for the model-owned parallelization contract:

* Resolves an optional sidecar from the model class
* Uses the generic infrastructure implementation when no sidecar is supplied
* Custom parallel plan support (dict or string path)
* Sequence parallel support
* Activation checkpointing for linear layers
* Model validation (attention heads divisible by TP size)
* Better fallback logic

NOTE: The passed-in model preferably should be on meta device. Otherwise,
the model must fit on GPU or CPU memory.

**Parameters:**

**`model`**

The model to be parallelized.

---

**`device_mesh`** `DeviceMesh`

The device mesh for distributed training.

---

**`mp_policy`** `Optional[MixedPrecisionPolicy]` — default: None

Mixed precision policy for model parallelism.

---

**`offload_policy`** `Optional[OffloadPolicy]` — default: None

The offload policy for FSDP.

---

**`sequence_parallel`** `bool` — default: False

Whether to use sequence parallelism. Defaults to False.

---

**`activation_checkpointing`** `bool` — default: False

Whether to use activation checkpointing. Defaults to False.

---

**`tp_shard_plan`** `Optional[Union[Dict[str, ParallelStyle], str]]` — default: None

Custom tensor parallel plan for the model. Can be:

* A dictionary mapping module names to parallel styles
* A string path to a dictionary or function that returns a dictionary
  If provided, this takes precedence over automatic plan generation.

---

**`dp_replicate_mesh_name`** `str` — default: 'dp\_replicate'

Key name for the data parallel replicate mesh in device\_mesh.
Used when data parallel replicate is enabled. Defaults to "dp\_replicate".

---

**`dp_shard_cp_mesh_name`** `str` — default: 'dp\_shard\_cp'

Key name for the data parallel shard + context parallel mesh in device\_mesh.
Used when data parallel shard is enabled. Defaults to "dp\_shard\_cp".

---

**`tp_mesh_name`** `str` — default: 'tp'

Key name for the tensor parallel mesh in device\_mesh.
Defaults to "tp".

---

**`frozen_multimodal_sharding`** `FrozenMultimodalSharding` — default: 'root'

Whether fully frozen multimodal modules are
owned by the root FSDP unit (`"root"`), sharded normally
(`"per_layer"`), or excluded from FSDP and copied on every rank
(`"replicate"`).

---

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

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

---

**Returns:** `nn.Module`

The parallelized model.

```python
nemo_automodel.components.distributed.parallelizer.get_hf_tp_shard_plan(
    model: torch.nn.Module
) -> dict[str, torch.distributed.tensor.parallel.ParallelStyle]
```

Translate a model-owned Hugging Face TP plan without importing model classes.

```python
nemo_automodel.components.distributed.parallelizer.get_model_layer_groups(
    model: torch.nn.Module
) -> typing.Dict[str, typing.List[torch.nn.Module]]
```

Return transformer layers grouped by model role (e.g. `language`, `vision`).

Public accessor over the per-model layer-group mapping used for FSDP wrapping
and activation checkpointing, so other components (e.g. the MoE parallelizer)
can consume the grouping without importing private helpers.

**Parameters:**

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

Root model to extract grouped layers from.

---

**Returns:** `Dict[str, List[nn.Module]]`

Mapping from group name to the list of transformer blocks in that group.

```python
nemo_automodel.components.distributed.parallelizer.import_class_from_path(
    name: str
) -> typing.Any
```

Import a class from a string path (e.g. 'torch.optim.AdamW').

**Parameters:**

**`full_path`**

Full path to class including module path and class name

---

**Returns:** `Any`

The imported class object

```python
nemo_automodel.components.distributed.parallelizer.import_classes_from_paths(
    class_paths: typing.List[str]
)
```

Helper function to import classes from string paths.

**Parameters:**

**`class_paths`** `List[str]`

The list of string paths to the classes.

---

**Returns:**

List of imported classes.

```python
nemo_automodel.components.distributed.parallelizer.megatron_fsdp_strategy_parallelize(
    model,
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    optimizer = None,
    megatron_fsdp_unit_modules: typing.List[str] | None = None,
    tp_shard_plan: typing.Dict[str, typing.Union[torch.distributed.tensor.parallel.RowwiseParallel, torch.distributed.tensor.parallel.ColwiseParallel, torch.distributed.tensor.parallel.SequenceParallel]] | 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,
    dp_shard_dim: str = 'dp',
    tp_dim: str = 'tp',
    reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None
)
```

Apply tensor/data parallelism (MegatronFSDP) and optional activation-checkpointing to the model.

NOTE: The passed-in model should preferably reside on the meta device.
Otherwise, ensure the model fits into available GPU or CPU memory.

NOTE: The user must ensure that the provided tp\_shard\_plan is compatible
with the model architecture.

**Parameters:**

**`model`**

The model to be parallelized.

---

**`device_mesh`** `DeviceMesh`

The device mesh describing the physical devices
used for distributed training.

---

**`megatron_fsdp_unit_modules`** `Optional[List[str]]` — default: None

Class paths of the sub-modules that
should become individual MegatronFSDP units. When None or empty, the wrap classes
are auto-derived from the model's `_no_split_modules` (see
`_derive_megatron_fsdp_unit_modules`).

---

**`tp_shard_plan`** `Optional[Dict[str, Union[RowwiseParallel, ColwiseParallel, SequenceParallel]]]` — default: None

A tensor-parallel sharding plan.
Keys are module names; values specify the parallel style to apply
(e.g., RowwiseParallel, ColwiseParallel, SequenceParallel).

---

**`zero_dp_strategy`** `int` — default: 3

The zero-DP strategy to use.

---

**`init_fsdp_with_meta_device`** `bool` — default: False

If True, construct the model on a
meta device first and materialize weights lazily to reduce memory
fragmentation.

---

**`grad_reduce_in_fp32`** `bool` — default: False

Reduce gradients in FP32 irrespective of the
parameter precision to improve numerical stability.

---

**`preserve_fp32_weights`** `bool` — default: False

Keep a master FP32 copy of weights when
training in reduced precision (e.g., FP16/BF16).

---

**`overlap_grad_reduce`** `bool` — default: True

If True, overlap gradient reduction with
backward computation.

---

**`overlap_param_gather`** `bool` — default: True

If True, overlap parameter gathering with
forward computation.

---

**`check_for_nan_in_grad`** `bool` — default: 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 per process. Enable `report_nan_in_param_grad` to restore
gradient NaN checking.

---

**`report_nan_in_param_grad`** `bool` — default: False

Whether Megatron-FSDP should perform
its precise per-parameter gradient NaN check. This is the 0.5.0
replacement for `check_for_nan_in_grad` and is disabled by default
because it can significantly reduce training throughput.

---

**`average_in_collective`** `bool` — default: False

Perform gradient averaging inside the
collective operation instead of dividing afterward.

---

**`disable_bucketing`** `bool` — default: False

Disable gradient bucketing; gradients are
reduced immediately as they are produced.

---

**`calculate_per_token_loss`** `bool` — default: False

Compute loss normalized by the number of
tokens instead of the number of sequences.

---

**`keep_fp8_transpose_cache`** `bool` — default: False

Retain the FP8
transpose cache when using a custom MegatronFSDP wrapper.

---

**`nccl_ub`** `bool` — default: False

Enable NCCL user-buffer API (experimental) for reduced
latency on some networks.

---

**`fsdp_double_buffer`** `bool` — default: False

Enable double buffering of parameters to
overlap communication and computation in MegatronFSDP.

---

**`dp_shard_dim`** `str` — default: 'dp'

Key name for the data parallel mesh in device\_mesh.
Defaults to "dp".

---

**`tp_dim`** `str` — default: 'tp'

Key name for the tensor parallel mesh in device\_mesh.
Defaults to "tp".

---

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

Optional callback that re-resolves parameter
trainability after tensor-parallel surgery and immediately before
Megatron-FSDP construction.

---

```python
nemo_automodel.components.distributed.parallelizer.translate_to_torch_parallel_style(
    style: str
)
```

Translates string descriptions to parallelism plans.

In model configurations, we use a neutral type (string) to specify parallel
styles, here we translate them into torch.distributed tensor-parallel
types.

```python
nemo_automodel.components.distributed.parallelizer.unshard_fsdp2_model(
    model: torch.nn.Module
) -> typing.Generator[None, None, None]
```

Explicitly unshard and then reshard the FSDP2 modules. Useful for logprob inference.

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

Validate that every sharded language attention shape divides the TP size.

```python
nemo_automodel.components.distributed.parallelizer.HAVE_MEGATRON_FSDP = True
```

```python
nemo_automodel.components.distributed.parallelizer.MegatronFSDPMixedPrecisionPolicy = UnavailableMeta('MixedPrecisionPolicy', (), {'_msg': _MEGATRON_FSDP_050_REQUIRED...
```

```python
nemo_automodel.components.distributed.parallelizer._MEGATRON_FSDP_050_REQUIRED_MSG = 'megatron_fsdp.MixedPrecisionPolicy could not be imported: NeMo Automodel requir...
```

```python
nemo_automodel.components.distributed.parallelizer._megatron_fsdp_nan_check_noop_warned = False
```

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