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`	Shard Beam with an fp32 router unit and fp32 RoPE tables kept across block boundaries.
`_DenseBlockParallelizer`	Dense FSDP2 walk helper: Beam blocks go through `fully_shard_by_dtype`, other units unchanged.

Functions

Name	Description
`beam_fully_shard_unit`	Shard one FSDP unit without casting a Beam block’s FP32 rotary inputs.
`is_beam_block`	True for a Beam decoder `Block`, also when wrapped for activation checkpointing.

Data

PARALLELIZER

API

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.

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.

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.

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.

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.

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

Mixed-precision policy for this unit.

**kwargs

AnyDefaults to {}

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

Returns: nn.Module

The sharded module.

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.

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