> This page is for version 0.4.0 · 26.04.
> 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.models.step5.parallelization

FSDP sidecar that keeps the Step 5 final norm in an FP32 unit.

## Module Contents

### Classes

| Name                                                                                                       | Description                                            |
| ---------------------------------------------------------------------------------------------------------- | ------------------------------------------------------ |
| [`Step5ModelParallelizer`](#nemo_automodel-components-models-step5-parallelization-Step5ModelParallelizer) | Keep the final decoder norm in its own FP32 FSDP unit. |

### Data

[`PARALLELIZER`](#nemo_automodel-components-models-step5-parallelization-PARALLELIZER)

[`_FP32_NORM_POLICY`](#nemo_automodel-components-models-step5-parallelization-_FP32_NORM_POLICY)

### API

```python
class nemo_automodel.components.models.step5.parallelization.Step5ModelParallelizer()
```

**Bases:** `ModelParallelizer`

Keep the final decoder norm in its own FP32 FSDP unit.

```python
nemo_automodel.components.models.step5.parallelization.Step5ModelParallelizer._fully_shard_module(
    module: torch.nn.Module,
    kwargs = {}
) -> torch.nn.Module
```

Shard a module, first giving any contained final norm an FP32 unit.

Decoder blocks already group their strict-FP32 norms by dtype. The
final norm is owned by the root unit, which uses the BF16 parameter
policy, so it is wrapped separately before that root.

**Parameters:**

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

Module to shard in place.

---

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

`fully_shard` keyword arguments (mesh, mp_policy,
offload_policy, reshard_after_forward, ignored_params).

---

**Returns:** `nn.Module`

The input module with FSDP applied.

```python
nemo_automodel.components.models.step5.parallelization.PARALLELIZER = Step5ModelParallelizer()
```

```python
nemo_automodel.components.models.step5.parallelization._FP32_NORM_POLICY = MixedPrecisionPolicy(param_dtype=(torch.float32), reduce_dtype=(torch.float32), ...
```