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`	Model-owned parallelization of `PraxisPolicy`, the `parallelizer` of the model class.

Functions

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

API

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

Bases: 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).

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.

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.

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.

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.

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.

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.

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

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

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.

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 | NoneDefaults to None

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

activation_checkpointing

boolDefaults to True

Recompute the blocks’ activations in the backward pass.

ignore_router_for_ac

boolDefaults to True

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

reshard_after_forward

boolDefaults to True

Free each unit’s gathered parameters after its forward.

offload_policy

OffloadPolicy | NoneDefaults to None

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