> 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                                                                          |
| ---------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------ |
| [`DeepseekV4ParallelizationStrategy`](#nemo_automodel-components-distributed-parallelizer-DeepseekV4ParallelizationStrategy)       | DeepSeek-V4 keeps a small set of reference-sensitive parameters in fp32.             |
| [`DefaultParallelizationStrategy`](#nemo_automodel-components-distributed-parallelizer-DefaultParallelizationStrategy)             | Default parallelization strategy used by most models.                                |
| [`Gemma4ForConditionalGeneration`](#nemo_automodel-components-distributed-parallelizer-Gemma4ForConditionalGeneration)             | Placeholder when the installed transformers build has no Gemma4.                     |
| [`HunyuanParallelizationStrategy`](#nemo_automodel-components-distributed-parallelizer-HunyuanParallelizationStrategy)             | Parallelization strategy for Hunyuan-style transformer modules used in HunyuanVideo. |
| [`LTX2ParallelizationStrategy`](#nemo_automodel-components-distributed-parallelizer-LTX2ParallelizationStrategy)                   | Parallelization strategy for the LTX-2 video+audio transformer.                      |
| [`NemotronHParallelizationStrategy`](#nemo_automodel-components-distributed-parallelizer-NemotronHParallelizationStrategy)         | Specialized parallelization strategy for NemotronH models.                           |
| [`ParallelizationStrategy`](#nemo_automodel-components-distributed-parallelizer-ParallelizationStrategy)                           | Abstract base class for model parallelization strategies.                            |
| [`Qwen3_5ParallelizationStrategy`](#nemo_automodel-components-distributed-parallelizer-Qwen3_5ParallelizationStrategy)             | Parallelization strategy for Qwen3.5 dense models with mixed-dtype GatedDeltaNet.    |
| [`QwenImageEditParallelizationStrategy`](#nemo_automodel-components-distributed-parallelizer-QwenImageEditParallelizationStrategy) | Shard upstream Qwen image transformer blocks as complete FSDP2 units.                |
| [`WanParallelizationStrategy`](#nemo_automodel-components-distributed-parallelizer-WanParallelizationStrategy)                     | Parallelization strategy for Wan-style transformer modules used in Diffusers.        |

### Functions

| Name                                                                                                                                                         | Description                                                                                      |
| ------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------ |
| [`_apply_bagel_full_layer_activation_checkpointing`](#nemo_automodel-components-distributed-parallelizer-_apply_bagel_full_layer_activation_checkpointing)   | Apply native BAGEL-style activation checkpointing to whole logical layers.                       |
| [`_apply_per_layer_compile`](#nemo_automodel-components-distributed-parallelizer-_apply_per_layer_compile)                                                   | Compile each decoder layer in-place after FSDP2 sharding.                                        |
| [`_apply_qwen_block_activation_checkpointing`](#nemo_automodel-components-distributed-parallelizer-_apply_qwen_block_activation_checkpointing)               | Checkpoint complete Qwen blocks, including both image and text MLPs.                             |
| [`_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`.                         |
| [`_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.                       |
| [`_find_largest_module_list`](#nemo_automodel-components-distributed-parallelizer-_find_largest_module_list)                                                 | Heuristic function to find the largest layer container in a model.                               |
| [`_get_model_layer_group_specs`](#nemo_automodel-components-distributed-parallelizer-_get_model_layer_group_specs)                                           | -                                                                                                |
| [`_get_module_by_fqn`](#nemo_automodel-components-distributed-parallelizer-_get_module_by_fqn)                                                               | -                                                                                                |
| [`_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)                                                 | -                                                                                                |
| [`_is_checkpoint_wrapped`](#nemo_automodel-components-distributed-parallelizer-_is_checkpoint_wrapped)                                                       | -                                                                                                |
| [`_is_transformers_v5_or_higher`](#nemo_automodel-components-distributed-parallelizer-_is_transformers_v5_or_higher)                                         | Check if transformers version is 5.x or higher.                                                  |
| [`_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.                          |
| [`_nemotronh_decoder_blocks`](#nemo_automodel-components-distributed-parallelizer-_nemotronh_decoder_blocks)                                                 | Return `(container, blocks)` for a NemotronH model's decoder blocks.                             |
| [`_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.               |
| [`_unwrap_qwen_checkpointed_block`](#nemo_automodel-components-distributed-parallelizer-_unwrap_qwen_checkpointed_block)                                     | Return the Qwen transformer block held by a checkpoint wrapper.                                  |
| [`_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.                            |
| [`_validate_qwen_transformer_blocks`](#nemo_automodel-components-distributed-parallelizer-_validate_qwen_transformer_blocks)                                 | Validate the upstream Qwen dual-stream block structure.                                          |
| [`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)                                                           | Get the Hugging Face tensor parallel plan from the model.                                        |
| [`get_model_layer_groups`](#nemo_automodel-components-distributed-parallelizer-get_model_layer_groups)                                                       | Return transformer layers grouped by model role (e.g. `language`, `vision`).                     |
| [`get_parallelization_strategy`](#nemo_automodel-components-distributed-parallelizer-get_parallelization_strategy)                                           | Get the appropriate parallelization strategy for the given model.                                |
| [`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. |
| [`register_parallel_strategy`](#nemo_automodel-components-distributed-parallelizer-register_parallel_strategy)                                               | Decorator to register out-of-tree parallelism strategies.                                        |
| [`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 attention heads and key value heads are divisible by TP size                       |
| [`validate_tp_mesh_for_nemotron_nas`](#nemo_automodel-components-distributed-parallelizer-validate_tp_mesh_for_nemotron_nas)                                 | Validate that a Nemotron-NAS model can be tensor-parallel sharded.                               |

### Data

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

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

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

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

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

[`_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.DeepseekV4ParallelizationStrategy()
```

**Bases:** [DefaultParallelizationStrategy](#nemo_automodel-components-distributed-parallelizer-DefaultParallelizationStrategy)

DeepSeek-V4 keeps a small set of reference-sensitive parameters in fp32.

```python
nemo_automodel.components.distributed.parallelizer.DeepseekV4ParallelizationStrategy.parallelize(
    model,
    device_mesh,
    dp_shard_cp_mesh_name = 'dp_shard_cp',
    kwargs = {}
)
```

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

**Bases:** [ParallelizationStrategy](#nemo_automodel-components-distributed-parallelizer-ParallelizationStrategy)

Default parallelization strategy used by most models.

```python
nemo_automodel.components.distributed.parallelizer.DefaultParallelizationStrategy.parallelize(
    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.shared.multimodal_fsdp.FrozenMultimodalSharding = 'root',
    fully_shard_fn = None
) -> torch.nn.Module
```

Apply the default parallelization flow.

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

Placeholder when the installed transformers build has no Gemma4.

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

**Bases:** [ParallelizationStrategy](#nemo_automodel-components-distributed-parallelizer-ParallelizationStrategy)

Parallelization strategy for Hunyuan-style transformer modules used in HunyuanVideo.

```python
nemo_automodel.components.distributed.parallelizer.HunyuanParallelizationStrategy.parallelize(
    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 = True,
    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',
    kwargs = {}
) -> torch.nn.Module
```

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

**Bases:** [HunyuanParallelizationStrategy](#nemo_automodel-components-distributed-parallelizer-HunyuanParallelizationStrategy)

Parallelization strategy for the LTX-2 video+audio transformer.

`LTX2VideoTransformer3DModel` exposes its layers as `transformer_blocks`
but names the attention/FFN submodules `attn1`/`attn2`/`ff`, which the
Default strategy's submodule-level activation checkpointing does not
recognize — leaving attention and MLP activations un-checkpointed and OOM-ing
on the long combined video+audio token sequence. Wrapping each whole block
(as the HunyuanVideo strategy does) restores the expected memory profile.

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

**Bases:** [ParallelizationStrategy](#nemo_automodel-components-distributed-parallelizer-ParallelizationStrategy)

Specialized parallelization strategy for NemotronH models.

```python
nemo_automodel.components.distributed.parallelizer.NemotronHParallelizationStrategy.parallelize(
    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',
    reshard_after_forward: bool | None = None,
    kwargs = {}
) -> torch.nn.Module
```

Apply NemotronH-specific parallelization.

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

Abstract

Abstract base class for model parallelization strategies.

```python
nemo_automodel.components.distributed.parallelizer.ParallelizationStrategy.parallelize(
    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',
    reshard_after_forward: bool | None = None,
    activation_checkpointing_scope: nemo_automodel.components.distributed.config.ActivationCheckpointingScope | None = 'all',
    frozen_multimodal_sharding: nemo_automodel.shared.multimodal_fsdp.FrozenMultimodalSharding = 'root',
    kwargs = {}
) -> torch.nn.Module
```

abstract

Apply parallelization strategy to the model.

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

**Bases:** [DefaultParallelizationStrategy](#nemo_automodel-components-distributed-parallelizer-DefaultParallelizationStrategy)

Parallelization strategy for Qwen3.5 dense models with mixed-dtype GatedDeltaNet.

Qwen3.5 has linear\_attn layers with float32 params (A\_log, norm) alongside
bfloat16 params. Overrides the FSDP sharding step to use fully\_shard\_by\_dtype
per layer, and sets the CP mesh on CPAwareGatedDeltaNet modules.

```python
nemo_automodel.components.distributed.parallelizer.Qwen3_5ParallelizationStrategy.parallelize(
    model,
    device_mesh,
    dp_shard_cp_mesh_name = 'dp_shard_cp',
    kwargs = {}
)
```

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

**Bases:** [DefaultParallelizationStrategy](#nemo_automodel-components-distributed-parallelizer-DefaultParallelizationStrategy)

Shard upstream Qwen image transformer blocks as complete FSDP2 units.

Applies whole-block activation checkpointing that covers the attention,
image-MLP, and text-MLP branches of each dual-stream block, then delegates
TP planning and FSDP2 sharding to the default strategy.

```python
nemo_automodel.components.distributed.parallelizer.QwenImageEditParallelizationStrategy.parallelize(
    model: torch.nn.Module,
    args = (),
    kwargs = {}
) -> torch.nn.Module
```

Apply Qwen block checkpointing followed by the standard FSDP2 flow.

**Parameters:**

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

Upstream Qwen image transformer. Tensor layouts are unchanged
by sharding; each dual-stream block consumes image tensors of
shape \[batch, image\_tokens, hidden] and text tensors of shape
\[batch, text\_tokens, hidden].

---

**`*args`** — default: ()

Positional arguments forwarded to the default strategy.

---

**`**kwargs`** — default: \{}

Keyword arguments accepted by
:meth:`DefaultParallelizationStrategy.parallelize`.

---

**Returns:** `nn.Module`

The same upstream model with its parameters represented by FSDP2

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

**Bases:** [ParallelizationStrategy](#nemo_automodel-components-distributed-parallelizer-ParallelizationStrategy)

Parallelization strategy for Wan-style transformer modules used in Diffusers.

Applies TP to condition embedders, FFN projections in each block, and final projection,
then applies FSDP sharding similarly to other strategies.

```python
nemo_automodel.components.distributed.parallelizer.WanParallelizationStrategy.parallelize(
    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',
    kwargs = {}
) -> torch.nn.Module
```

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

Apply native BAGEL-style activation checkpointing to whole logical layers.

```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 in DefaultParallelizationStrategy).  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._apply_qwen_block_activation_checkpointing(
    model: torch.nn.Module
) -> None
```

Checkpoint complete Qwen blocks, including both image and text MLPs.

**Parameters:**

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

Upstream Qwen image transformer. Each block consumes image
hidden states with shape \[batch, image\_tokens, hidden], text hidden
states with shape \[batch, text\_tokens, hidden], a text mask with
shape \[batch, text\_tokens], timestep embeddings with shape \[batch,
hidden], and rotary tensors whose leading layout is owned by
Diffusers. Its block outputs preserve the image/text layouts.

---

```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._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._find_largest_module_list(
    model: torch.nn.Module
) -> typing.Union[torch.nn.ModuleList, torch.nn.ModuleDict] | None
```

Heuristic function to find the largest layer container in a model.

This function recursively traverses the model to find all nn.ModuleList and
pipeline-split nn.ModuleDict instances and returns the one with the most
modules. This is useful as a fallback when the model architecture is unknown,
since transformer layers are typically organized in ModuleLists. Pipeline
splitting converts ModuleLists to ModuleDicts keyed by original layer index.

**Parameters:**

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

The model to search through.

---

**Returns:** `Union[nn.ModuleList, nn.ModuleDict] | None`

Optional\[Union\[nn.ModuleList, nn.ModuleDict]]: The largest layer container found, or None.

```python
nemo_automodel.components.distributed.parallelizer._get_model_layer_group_specs() -> typing.Dict[typing.Any, typing.Dict[str, typing.List[str]]]
```

```python
nemo_automodel.components.distributed.parallelizer._get_module_by_fqn(
    module: torch.nn.Module,
    fqn: str
) -> torch.nn.Module | None
```

```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._is_checkpoint_wrapped(
    module: torch.nn.Module
) -> bool
```

```python
nemo_automodel.components.distributed.parallelizer._is_transformers_v5_or_higher() -> bool
```

Check if transformers version is 5.x or higher.

```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._nemotronh_decoder_blocks(
    model: torch.nn.Module
) -> tuple[torch.nn.Module, list[torch.nn.Module]]
```

Return `(container, blocks)` for a NemotronH model's decoder blocks.

Two distinct classes share the name `NemotronHForCausalLM`:

* the HF model keeps its blocks in `model.backbone.layers` (an `nn.ModuleList`), while
* the native Nemotron-V3 model (`NemotronV3Model`) keeps them in `model.model.layers`
  (an `nn.ModuleDict` keyed `"0".."N-1"`).

`container` is the underlying `ModuleList`/`ModuleDict` (so callers can write rewrapped
blocks back into the model), and `blocks` is the ordered list of block modules.

```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._unwrap_qwen_checkpointed_block(
    block: torch.nn.Module
) -> torch.nn.Module
```

Return the Qwen transformer block held by a checkpoint wrapper.

```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._validate_qwen_transformer_blocks(
    model: torch.nn.Module
) -> torch.nn.ModuleList
```

Validate the upstream Qwen dual-stream block structure.

**Parameters:**

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

Upstream Qwen image transformer whose `transformer_blocks`
container must hold dual-stream image/text blocks.

---

**Returns:** `nn.ModuleList`

The model's `transformer_blocks` ModuleList. Returned modules alias

```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,
    fully_shard_fn = None,
    frozen_multimodal_sharding: nemo_automodel.shared.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.

---

**`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.shared.multimodal_fsdp.FrozenMultimodalSharding = 'root'
)
```

Apply parallelisms and activation checkpointing to the model.

Enhanced version that uses a strategy pattern for different model parallelization approaches:

* Automatic strategy selection based on model type
* Polymorphic parallelization strategies for different model families
* 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"`).

---

**Returns:**

The parallelized model.

```python
nemo_automodel.components.distributed.parallelizer.get_hf_tp_shard_plan(
    model
)
```

Get the Hugging Face tensor parallel plan from the model.

This function:

* Retrieves TP strategies from model class, instance, and inner model levels.
* Handles special cases for `embed_tokens` and `lm_head` for speed up.
* Converts string-based parallel styles to DTensor parallelization strategies.

Taken and modified from: [https://github.com/NVIDIA/NeMo/blob/6c6169db01bcca73ae8ad3ac35242fadbb9a78ba/nemo/lightning/pytorch/strategies/utils.py#L532](https://github.com/NVIDIA/NeMo/blob/6c6169db01bcca73ae8ad3ac35242fadbb9a78ba/nemo/lightning/pytorch/strategies/utils.py#L532)

**Parameters:**

**`model`**

A Hugging Face model instance

---

**Returns:**

A dictionary mapping model component paths to their parallelization strategies

**Raises:**

* `AssertionError`: If no TP plan is found

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

Get the appropriate parallelization strategy for the given model.

```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'
)
```

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
:func:`_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".

---

```python
nemo_automodel.components.distributed.parallelizer.register_parallel_strategy(
    arg = None,
    name: str | None = None
)
```

Decorator to register out-of-tree parallelism strategies.

Supports:

* @register\_parallel\_strategy(name="CustomModelName")

```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,
    tp_mesh
)
```

Validate that attention heads and key value heads are divisible by TP size

```python
nemo_automodel.components.distributed.parallelizer.validate_tp_mesh_for_nemotron_nas(
    model,
    tp_size
)
```

Validate that a Nemotron-NAS model can be tensor-parallel sharded.

```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.PARALLELIZATION_STRATEGIES: Dict[str, ParallelizationStrategy] = {'NemotronHForCausalLM': NemotronHParallelizationStrategy(), 'DeepseekV4ForCausa...
```

```python
nemo_automodel.components.distributed.parallelizer._BAGEL_FULL_LAYER_CHECKPOINT_MODULE_LISTS = ('model.language_model.model.layers', 'model.vit_model.vision_model.encoder.laye...
```

```python
nemo_automodel.components.distributed.parallelizer._DEFAULT_STRATEGY = DefaultParallelizationStrategy()
```

```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__)
```