> 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.beam.parallelization

Model-owned FSDP2 parallelization for Beam: isolate the fp32 router and preserve fp32 RoPE inputs.

Every FSDP unit goes through `BeamModelParallelizer._fully_shard_module`: Beam blocks disable FSDP's automatic
forward-input cast and cast hidden states themselves, so the fp32 RoPE tables reach attention unchanged. Decoder
blocks are sharded into dtype-uniform units, so the fp32 router (`_keep_in_fp32_modules_strict`) gets its own
unit: the MoE/EP path does this through the shared `fully_shard_by_dtype` block sharding, the dense FSDP2 path
through `BeamModelParallelizer._apply_fsdp_sharding`.

## Module Contents

### Classes

| Name                                                                                                        | Description                                                                                    |
| ----------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------- |
| [`BeamModelParallelizer`](#nemo_automodel-components-models-beam-parallelization-BeamModelParallelizer)     | Shard Beam with an fp32 router unit and fp32 RoPE tables kept across block boundaries.         |
| [`_DenseBlockParallelizer`](#nemo_automodel-components-models-beam-parallelization-_DenseBlockParallelizer) | Dense FSDP2 walk helper: Beam blocks go through `fully_shard_by_dtype`, other units unchanged. |

### Functions

| Name                                                                                                    | Description                                                                      |
| ------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------- |
| [`beam_fully_shard_unit`](#nemo_automodel-components-models-beam-parallelization-beam_fully_shard_unit) | Shard one FSDP unit without casting a Beam block's FP32 rotary inputs.           |
| [`is_beam_block`](#nemo_automodel-components-models-beam-parallelization-is_beam_block)                 | True for a Beam decoder `Block`, also when wrapped for activation checkpointing. |

### Data

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

### API

```python
class nemo_automodel.components.models.beam.parallelization.BeamModelParallelizer()
```

**Bases:** `ModelParallelizer`

Shard Beam with an fp32 router unit and fp32 RoPE tables kept across block boundaries.

```python
nemo_automodel.components.models.beam.parallelization.BeamModelParallelizer._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
```

Run the shared dense FSDP2 walk with Beam blocks sharded into dtype-uniform units.

```python
nemo_automodel.components.models.beam.parallelization.BeamModelParallelizer._fully_shard_module(
    module: torch.nn.Module,
    kwargs: typing.Any = {}
) -> torch.nn.Module
```

Apply `beam_fully_shard_unit` to every FSDP unit.

```python
class nemo_automodel.components.models.beam.parallelization._DenseBlockParallelizer(
    unit_parallelizer: nemo_automodel.components.models.beam.parallelization.BeamModelParallelizer,
    fp32_compute_module_names: tuple[str, ...]
)
```

**Bases:** `ModelParallelizer`

Dense FSDP2 walk helper: Beam blocks go through `fully_shard_by_dtype`, other units unchanged.

The MoE path already shards each block with `fully_shard_by_dtype`; the dense walk wraps each layer with a
single `_fully_shard_module` call, which this helper turns into the same dtype-aware block sharding.

```python
nemo_automodel.components.models.beam.parallelization._DenseBlockParallelizer._fully_shard_module(
    module: torch.nn.Module,
    kwargs: typing.Any = {}
) -> torch.nn.Module
```

Shard a Beam block into dtype-uniform units; shard any other module as one Beam unit.

```python
nemo_automodel.components.models.beam.parallelization.beam_fully_shard_unit(
    module: torch.nn.Module,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None = None,
    kwargs: typing.Any = {}
) -> torch.nn.Module
```

Shard one FSDP unit without casting a Beam block's FP32 rotary inputs.

For Beam blocks, clone the policy with cast_forward_inputs=False and record param_dtype on the unwrapped
block, whose forward casts hidden states itself. Other modules use the policy unchanged.

**Parameters:**

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

Module to wrap; a Beam block may be wrapped for activation checkpointing.

---

**`mp_policy`** `MixedPrecisionPolicy | None` — default: None

Mixed-precision policy for this unit.

---

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

Remaining `fully_shard` arguments (mesh, reshard_after_forward, offload_policy, ...).

---

**Returns:** `nn.Module`

The sharded module.

```python
nemo_automodel.components.models.beam.parallelization.is_beam_block(
    module: torch.nn.Module
) -> bool
```

True for a Beam decoder `Block`, also when wrapped for activation checkpointing.

```python
nemo_automodel.components.models.beam.parallelization.PARALLELIZER = BeamModelParallelizer()
```