> This page is for version 0.7.0 · 26.10.
> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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.

# nemo_automodel.components.models.praxis.parallelize

Parallelization of `PraxisPolicy` for training.

FSDP2 shards every block of the two backbone towers and the action expert over the data-parallel
mesh. With `distributed.ep_size > 1` the MoE feed-forwards are expert parallel through the shared
MoE machinery (expert sharding, the backend's token dispatcher, dispatch replay under activation
checkpointing), as in the language-model MoE path; the walk lives here because the Praxis blocks sit
in two stacks under `backbone` and one under `action_expert` rather than in a single `layers`
list.

## Module Contents

### Classes

| Name                                                                                                      | Description                                                                           |
| --------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- |
| [`PraxisModelParallelizer`](#nemo_automodel-components-models-praxis-parallelize-PraxisModelParallelizer) | Model-owned parallelization of `PraxisPolicy`, the `parallelizer` of the model class. |

### Functions

| Name                                                                                                      | Description                                                                                  |
| --------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------- |
| [`_block_lists`](#nemo_automodel-components-models-praxis-parallelize-_block_lists)                       | The block stacks of the model: encoder tower, decoder tower, action expert.                  |
| [`_expert_meshes`](#nemo_automodel-components-models-praxis-parallelize-_expert_meshes)                   | The expert-parallel mesh and the mesh the experts are FSDP-sharded over; `None` when unused. |
| [`_moe_of`](#nemo_automodel-components-models-praxis-parallelize-_moe_of)                                 | The MoE feed-forward of a backbone block, or `None` for dense and action-expert blocks.      |
| [`_mp_policy`](#nemo_automodel-components-models-praxis-parallelize-_mp_policy)                           | -                                                                                            |
| [`_policy`](#nemo_automodel-components-models-praxis-parallelize-_policy)                                 | -                                                                                            |
| [`replicate_router_biases`](#nemo_automodel-components-models-praxis-parallelize-replicate_router_biases) | Turn every router's selection bias into a DTensor replicated over `mesh`.                    |
| [`shard_policy`](#nemo_automodel-components-models-praxis-parallelize-shard_policy)                       | Shard the two backbone towers and the action expert block by block, then the root.           |

### API

```python
class nemo_automodel.components.models.praxis.parallelize.PraxisModelParallelizer()
```

**Bases:** [ModelParallelizer](/nemo/automodel/v0.7/nemo-automodel/nemo_automodel/components/distributed/parallelizer#nemo_automodel-components-distributed-parallelizer-ModelParallelizer)

Model-owned parallelization of `PraxisPolicy`, the `parallelizer` of the model class.

The shared flows walk one `layers` stack; these hooks walk the Praxis towers with
`shard_policy` instead, with (`_apply_moe`) and without (`_apply`) expert parallelism,
and on one rank (`_apply_unsharded`).

```python
nemo_automodel.components.models.praxis.parallelize.PraxisModelParallelizer._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,
    reshard_after_forward: bool | None = None,
    kwargs = {}
) -> torch.nn.Module
```

Shard `model` for FSDP2 without expert parallelism; the experts stay inside their blocks.

```python
nemo_automodel.components.models.praxis.parallelize.PraxisModelParallelizer._apply_moe(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module
```

Shard `model` for the FSDP2 strategy, and the expert parallelism, of `mesh_context`.

```python
nemo_automodel.components.models.praxis.parallelize.PraxisModelParallelizer._apply_unsharded(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module
```

One data-parallel rank: the per-block FSDP2 wrapping and activation checkpointing over the one-rank mesh.

```python
nemo_automodel.components.models.praxis.parallelize._block_lists(
    policy: nemo_automodel.components.models.praxis.model.PraxisPolicy
) -> tuple[torch.nn.ModuleList, ...]
```

The block stacks of the model: encoder tower, decoder tower, action expert.

```python
nemo_automodel.components.models.praxis.parallelize._expert_meshes(
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> tuple[torch.distributed.device_mesh.DeviceMesh | None, torch.distributed.device_mesh.DeviceMesh | None]
```

The expert-parallel mesh and the mesh the experts are FSDP-sharded over; `None` when unused.

```python
nemo_automodel.components.models.praxis.parallelize._moe_of(
    block: torch.nn.Module
) -> nemo_automodel.components.moe.layers.MoE | None
```

The MoE feed-forward of a backbone block, or `None` for dense and action-expert blocks.

```python
nemo_automodel.components.models.praxis.parallelize._mp_policy(
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None
) -> torch.distributed.fsdp.MixedPrecisionPolicy
```

```python
nemo_automodel.components.models.praxis.parallelize._policy(
    model: torch.nn.Module
) -> nemo_automodel.components.models.praxis.model.PraxisPolicy
```

```python
nemo_automodel.components.models.praxis.parallelize.replicate_router_biases(
    policy: nemo_automodel.components.models.praxis.model.PraxisPolicy,
    mesh: torch.distributed.device_mesh.DeviceMesh
) -> None
```

Turn every router's selection bias into a DTensor replicated over `mesh`.

FSDP2 leaves buffers as plain tensors, and the shared gate steps a plain bias from this rank's
expert load alone. Replicated, its update sums the load over all ranks of `mesh` and applies the
same step on every rank, as the reference trainer does.

**Parameters:**

**`policy`** `PraxisPolicy`

The model; its MoE gates hold `e_score_correction_bias`.

---

**`mesh`** `DeviceMesh`

Data-parallel device mesh.

---

```python
nemo_automodel.components.models.praxis.parallelize.shard_policy(
    policy: nemo_automodel.components.models.praxis.model.PraxisPolicy,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy,
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    activation_checkpointing: bool = True,
    ignore_router_for_ac: bool = True,
    reshard_after_forward: bool = True,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None = None
) -> None
```

Shard the two backbone towers and the action expert block by block, then the root.

Every transformer block is one FSDP2 unit. With `activation_checkpointing` each block recomputes
its activations in the backward pass; for an MoE block the router projection and its top-k
selection are saved instead of recomputed when `ignore_router_for_ac` is set, which keeps the
recompute's routing and expert dispatch identical to the forward's. With a `moe_mesh` whose
`ep` axis has more than one rank the routed experts are expert parallel: every expert lives
whole on one rank of that axis, the backend's token dispatcher exchanges the tokens, and the
experts are FSDP-sharded over the `ep_shard` axis when that axis has more than one rank. The
routers' selection biases become DTensors replicated over `mesh`
(`replicate_router_biases`), so that their load-balancing steps use the expert load of all
ranks.

With LoRA on the action expert the decoder tower never runs, so its units have no unsharded state
while the root accumulates gradients across micro-steps; the FSDP guard patch skips them.

**Parameters:**

**`policy`** `PraxisPolicy`

The model, built on the meta device.

---

**`mesh`** `DeviceMesh`

Data-parallel device mesh of the FSDP shards.

---

**`mp_policy`** `MixedPrecisionPolicy`

Parameter and reduction dtypes of the shards.

---

**`moe_mesh`** `DeviceMesh | None` — default: None

MoE mesh with the `ep` and, optionally, `ep_shard` axes; `None` without expert
parallelism.

---

**`activation_checkpointing`** `bool` — default: True

Recompute the blocks' activations in the backward pass.

---

**`ignore_router_for_ac`** `bool` — default: True

Save the MoE routing across the recompute instead of recomputing it.

---

**`reshard_after_forward`** `bool` — default: True

Free each unit's gathered parameters after its forward.

---

**`offload_policy`** `OffloadPolicy | None` — default: None

FSDP2 offload policy of the units, or `None` for none.

---