nemo_automodel.components.distributed.model_parallelizer

Resolve and execute model-owned parallelization sidecars.

Module Contents

Classes

Name	Description
`ModelParallelizer`	Single model-owned parallelization sidecar contract.

Functions

Name	Description
`_apply_model_parallelizer`	Execute shared strategy dispatch for one model-owned parallelizer.
`_parallelize_ddp`	-
`_parallelize_fsdp2`	-
`_parallelize_megatron_fsdp`	-
`_parallelize_moe`	-
`_parallelize_unsharded_fsdp2`	-
`compile_parallelized_model`	Compile FSDP2 layers after parallelization when requested.
`get_model_parallelizer`	Return the class-owned sidecar, or the shared default implementation.
`parallelize_model`	Apply all requested parallelisms through the model-owned contract.

Data

_DEFAULT_PARALLELIZER

API

class nemo_automodel.components.distributed.parallelizer.ModelParallelizer()

Single model-owned parallelization sidecar contract.

nemo_automodel.components.distributed.parallelizer.ModelParallelizer._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,
    tp_shard_plan: typing.Union[typing.Dict[str, torch.distributed.tensor.parallel.ParallelStyle], str] | None = None,
    dp_replicate_mesh_name: str = 'dp_replicate',
    dp_shard_cp_mesh_name: str = 'dp_shard_cp',
    tp_mesh_name: str = 'tp',
    enable_async_tensor_parallel: bool = False,
    enable_compile: bool = False,
    enable_fsdp2_prefetch: bool = True,
    fsdp2_backward_prefetch_depth: int = 2,
    fsdp2_forward_prefetch_depth: int = 1,
    reshard_after_forward: bool | None = None,
    activation_checkpointing_scope: nemo_automodel.components.distributed.config.ActivationCheckpointingScope | None = 'all',
    frozen_multimodal_sharding: nemo_automodel.components.distributed.multimodal_fsdp.FrozenMultimodalSharding = 'root',
    reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None
) -> torch.nn.Module

Apply the shared dense FSDP2 implementation.

nemo_automodel.components.distributed.parallelizer.ModelParallelizer._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

Wrap the model’s submodules into FSDP2 units.

Model-owned sidecar strategies deriving from this class override this hook to change how parameters are grouped into FSDP units without reimplementing the surrounding TP/AC/mixed-precision flow. Override _fully_shard_module when a model needs a different primitive.

nemo_automodel.components.distributed.parallelizer.ModelParallelizer._apply_moe(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

Apply the expert-parallel flow selected by mesh_context.ep_size > 1.

The shared flow walks one decoder layers stack. A model whose MoE blocks live in several towers overrides this hook and shards them itself.

nemo_automodel.components.distributed.parallelizer.ModelParallelizer._apply_unsharded(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

Prepare model for FSDP2 when the mesh has no sharding dimension (one data-parallel rank).

The shared flow skips fully_shard, installs activation checkpointing on one decoder layers stack and wraps the forward in autocast. A model whose blocks live in several towers overrides this hook.

nemo_automodel.components.distributed.parallelizer.ModelParallelizer._fully_shard_module(
    module: torch.nn.Module,
    kwargs = {}
) -> torch.nn.Module

Apply the FSDP2 primitive used by this model sidecar.

nemo_automodel.components.distributed.parallelizer.ModelParallelizer._fully_shard_untied_input_output_embeddings(
    model: torch.nn.Module,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None,
    input_reshard_after_forward: bool
) -> None

Give large trainable untied embedding tables independent FSDP buffers.

The generic dense path otherwise leaves both tables in the root FSDP unit. With fp32 gradient reduction, that unit allocates one contiguous reduce-scatter input containing both gradients. Keeping the two trainable leaf modules in separate FSDP units bounds that allocation by the larger table instead of their sum. This pass skips tied tables to preserve aliasing, and frozen tables because they have no gradient communication buffer to split. Tables already wrapped as FSDP units retain their existing reshard policy. With an explicit reshard_after_forward=True, a container-hosted head reshards while a newly split top-level head stays gathered. FusedLinearCrossEntropy handles the sharded weight with correct gradients, at the cost of an extra all-gather.

Parameters:

model

nn.Module

Model whose input and output embedding modules may be sharded.

mesh

DeviceMesh

Device mesh that owns the FSDP shards.

mp_policy

MixedPrecisionPolicy | None

Mixed-precision policy used by the surrounding FSDP units.

offload_policy

OffloadPolicy | None

Optional offload policy used by the surrounding FSDP units.

input_reshard_after_forward

bool

Whether the input embedding unit reshards its parameters after forward.

nemo_automodel.components.distributed.parallelizer.ModelParallelizer._prepare_model(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> None

Validate model constraints and install runtime mesh state before strategy dispatch.

A sidecar can reject unsupported topologies or set model-owned context-parallel state here, before any executor changes parameter ownership. The default needs no preparation; all sidecars retain the inherited public entry point.

nemo_automodel.components.distributed.parallelizer.ModelParallelizer._use_full_layer_activation_checkpointing(
    model: torch.nn.Module
) -> bool

Return whether this model safely opts into whole-layer checkpointing.

nemo_automodel.components.distributed.parallelizer.ModelParallelizer._validate_tp_mesh(
    model: torch.nn.Module,
    tp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> None

Validate the model’s attention topology against its TP mesh.

nemo_automodel.components.distributed.parallelizer.ModelParallelizer.parallelize(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

Apply every requested parallelism and return the parallelized model.

nemo_automodel.components.distributed.model_parallelizer._apply_model_parallelizer(
    parallelizer: nemo_automodel.components.distributed.parallelizer.ModelParallelizer,
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

Execute shared strategy dispatch for one model-owned parallelizer.

nemo_automodel.components.distributed.model_parallelizer._parallelize_ddp(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

nemo_automodel.components.distributed.model_parallelizer._parallelize_fsdp2(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext,
    parallelizer: nemo_automodel.components.distributed.parallelizer.ModelParallelizer
) -> torch.nn.Module

nemo_automodel.components.distributed.model_parallelizer._parallelize_megatron_fsdp(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

nemo_automodel.components.distributed.model_parallelizer._parallelize_moe(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext,
    parallelizer: nemo_automodel.components.distributed.parallelizer.ModelParallelizer
) -> torch.nn.Module

nemo_automodel.components.distributed.model_parallelizer._parallelize_unsharded_fsdp2(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext,
    parallelizer: nemo_automodel.components.distributed.parallelizer.ModelParallelizer
) -> torch.nn.Module

nemo_automodel.components.distributed.model_parallelizer.compile_parallelized_model(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> None

Compile FSDP2 layers after parallelization when requested.

nemo_automodel.components.distributed.model_parallelizer.get_model_parallelizer(
    model: torch.nn.Module
) -> nemo_automodel.components.distributed.parallelizer.ModelParallelizer

Return the class-owned sidecar, or the shared default implementation.

nemo_automodel.components.distributed.model_parallelizer.parallelize_model(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

Apply all requested parallelisms through the model-owned contract.

nemo_automodel.components.distributed.model_parallelizer._DEFAULT_PARALLELIZER = ModelParallelizer()