nemo_automodel.components.distributed.parallelizer

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

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

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

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

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

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

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

Compile each decoder layer in-place after FSDP2 sharding.

Compiles at decoder-layer granularity (not sub-module) so that AOT autograd traces the joint fwd+bwd graph under the training loop’s enable_grad context. Sub-module compile (e.g. on mlp alone) would be traced during activation checkpointing’s first forward pass which runs under no_grad, producing a forward-only graph that drops LoRA and other trainable-parameter gradients.

Prerequisite: NO_REENTRANT checkpoint_wrapper must already be applied to self_attn and mlp before FSDP2 sharding (done by ModelParallelizer). This function only handles the compile step.

Whole-block selective-AC wrappers (tagged with SELECTIVE_AC_WRAPPER_FLAG) are compiled OUTER — the wrapper itself is compiled so the selective policy is traced and the partitioner honors its recompute tags. Other layer-level CheckpointWrappers (e.g. the PP path) are unwrapped and the decoder layer is compiled directly.

nn.Module.compile() is used instead of torch.compile() to compile in-place without introducing an _orig_mod wrapper, which would add a key prefix and break checkpoint loading.

_patch_dtensor_spec_hash_for_symint() is called to allow torch.compile with dynamic shapes to coexist with DTensor’s lru_cache-based sharding propagation.

Resolve the language-model config structurally for TP validation.

Return True when the TP plan column-wise shards any QKV attention projection.

When Q/K/V projections use ColwiseParallel with sharded output (the default), each TP rank holds num_heads / tp_size heads and the model config / layer attributes must be updated accordingly.

Plans that keep attention replicated (e.g. Phi-3 with RowwiseParallel on fused QKV and Replicate output) should not trigger a head-count update.

Derive the MegatronFSDP wrap classes from a model’s _no_split_modules.

Used when a config does not specify megatron_fsdp_unit_modules. HF PreTrainedModel and the NeMo custom models both define _no_split_modules as a list of block class names (for example ["LlamaDecoderLayer"]). Walking model.modules() and matching type(module).__name__ against those names resolves the actual instantiated classes, so the result is correct for both the HF backend and the NeMo-custom backend (which use distinct classes that share the same name). For VLM/MoE models whose top-level _no_split_modules lists several block classes (for example vision and language towers), every matching class found anywhere in the module tree is collected.

Parameters:

Returns: list[type[nn.Module]]

The de-duplicated list of submodule classes to wrap as MegatronFSDP units,

Raises:

• ValueError: If the model does not expose a non-empty _no_split_modules, or if none of those names match an instantiated submodule. Raised with an actionable message instead of letting MegatronFSDP later fail with ZeroDivisionError (total_fsdp_module=0) when zero modules are wrapped.

Discover language, vision, and audio layer containers structurally.

Extract transformer layers grouped by model role.

Extract layers from different model architectures for parallelization.

This compatibility wrapper flattens grouped language/vision/audio layers. New activation-checkpointing code should use _extract_model_layer_groups so scope decisions can be explicit.

Select trainable activation-checkpointed layers from grouped model layers.

Resolve the optional Hugging Face embedding getters.

Select the tensor-parallel plan for the given model.

Priority order:

1. If tp_shard_plan is provided as a dict or import path, use it.
2. If the model type exists in PARALLELIZE_FUNCTIONS, use its optimised plan; on failure, fall back to HF plan.
3. Otherwise, prefer the model’s HF-native _tp_plan (via get_hf_tp_shard_plan).
4. Otherwise, fall back to the default base plan.

When tp_size > 1 and the model falls through to path 4 and the model class was loaded from a custom-code source (HF’s trust_remote_code=True path, where the dynamic class lives under transformers_modules.*), this raises ValueError instead of returning the default base plan. On recent PyTorch the default plan’s placements do not populate shard_order and trip an internal assert in torch.distributed.tensor._redistribute on the first weight redistribute, which surfaces to the user as an opaque PyTorch internal error. Custom-code architectures are the only known-broken case (see https://github.com/NVIDIA-NeMo/Automodel/issues/2243); known HF architectures that happen to fall through (e.g. Mixtral) are left on the default plan with a warning, since they have been working in practice.

When the model did define a _tp_plan but get_hf_tp_shard_plan raised while translating it (e.g. styles nemo does not recognize), the translator’s error message is folded into the ValueError as a diagnostic so the user can tell whether to add a _tp_plan from scratch or fix the styles in the one they already have.

Return whether whole-block activation checkpointing is safe for a KV-shared model.

checkpoint_wrapper replays a whole decoder block during backward with the arguments the forward saw. Unlike HF’s GradientCheckpointingLayer.__call__ it cannot drop past_key_values from that replay, so every layer that writes to a cache writes to it a second time. A KV-shared model then needs both halves to hold:

• the layers that populate the cache — the non-shared ones, which are what call Cache.update() — must not accumulate on the replay, or the recomputed K/V stops matching the forward;
• the shared layers must still read the K/V their source layer produced.

Neither holds for a model backed by an accumulating Cache: the second Cache.update() grows the entry and backward dies with a CheckpointError about changed tensor metadata (observed on native HF Gemma3nForCausalLM with use_cache=True). KV-shared models therefore stay on apply_submodule_checkpointing, which leaves attention unwrapped, by default.

A model that satisfies both halves opts in by setting the class attribute kv_sharing_survives_checkpoint_replay = True. Gemma4 E2B/E4B qualify: a pass-through holder stands in for the cache, and the shared layers read a separate store that the replay does not disturb (see gemma4_moe/model.py).

Parameters:

Returns: bool

Whether the model declares its KV sharing safe under whole-block replay.

Translate the config precision controls to the Megatron-FSDP 0.5.0 API.

megatron-fsdp==0.5.0, the only supported release, expresses precision through a MixedPrecisionPolicy plus a more expensive per-parameter NaN reporter. The reporter stays a separate opt-in rather than being silently enabled from the legacy buffer-check setting; because 0.5.0 has no buffer-level NaN check at all, a truthy check_for_nan_in_grad is dropped with a one-time warning that points at report_nan_in_param_grad as the opt-in replacement. Any other fully_shard signature — older or newer releases alike — fails loudly instead of guessing a translation.

Fix a crash when torch.compile + DTensor are used together.

Problem: torch.compile traces with symbolic shapes (SymInt). DTensorSpec hashes its shape to cache sharding decisions, but SymInt is not hashable -> crash.

Fix: if hashing the shape fails, fall back to hashing only (mesh, placements). Cache hits are slightly reduced but correctness is unaffected.

Return whether HF-native gradient checkpointing can preserve AutoModel’s AC scope.

Return the most specific nested Hugging Face TP-plan provider.

After TP sharding, the Q/K/V outputs are split across ranks (each rank has num_heads/tp_size heads). Update the config and each attention layer’s num_heads / num_key_value_heads so the forward uses the local head count instead of the global one (avoids shape mismatches in .view()).

Return whether model contains Automodel’s grouped custom-MoE layer.

Recursively apply FSDP2 sharding to modules, with optimizations for ModuleList.

This utility function traverses a model hierarchy and applies FSDP2 sharding to each module. For ModuleList instances (commonly used for transformer layers), it applies an optimization where the last layer doesn’t reshard after forward since FSDP2 will prefetch it immediately.

Handles both single-level and nested ModuleList/ModuleDict structures. If a ModuleList contains other ModuleLists, it will recurse into them instead of trying to wrap them (since ModuleList doesn’t have a forward method).

Note: This function modifies the module in-place by replacing modules with their FSDP2-subclassed versions.

Parameters:

Apply selective activation checkpointing to model end to end.

Standalone entry point (detects KV-sharing, disables use_cache, and wraps transformer blocks) for paths where the FSDP2 parallelize flow is skipped — notably single-GPU training.

Parameters:

Apply parallelisms and activation checkpointing to the model.

Compatibility entry point for the model-owned parallelization contract:

• Resolves an optional sidecar from the model class
• Uses the generic infrastructure implementation when no sidecar is supplied
• Custom parallel plan support (dict or string path)
• Sequence parallel support
• Activation checkpointing for linear layers
• Model validation (attention heads divisible by TP size)
• Better fallback logic

NOTE: The passed-in model preferably should be on meta device. Otherwise, the model must fit on GPU or CPU memory.

Parameters:

Returns: nn.Module

The parallelized model.

Translate a model-owned Hugging Face TP plan without importing model classes.

Return transformer layers grouped by model role (e.g. language, vision).

Public accessor over the per-model layer-group mapping used for FSDP wrapping and activation checkpointing, so other components (e.g. the MoE parallelizer) can consume the grouping without importing private helpers.

Parameters:

Returns: Dict[str, List[nn.Module]]

Mapping from group name to the list of transformer blocks in that group.

Import a class from a string path (e.g. ‘torch.optim.AdamW’).

Parameters:

Returns: Any

The imported class object

Helper function to import classes from string paths.

Parameters:

Returns:

List of imported classes.

Apply tensor/data parallelism (MegatronFSDP) and optional activation-checkpointing to the model.

NOTE: The passed-in model should preferably reside on the meta device. Otherwise, ensure the model fits into available GPU or CPU memory.

NOTE: The user must ensure that the provided tp_shard_plan is compatible with the model architecture.

Parameters:

Translates string descriptions to parallelism plans.

In model configurations, we use a neutral type (string) to specify parallel styles, here we translate them into torch.distributed tensor-parallel types.

Explicitly unshard and then reshard the FSDP2 modules. Useful for logprob inference.

Validate that every sharded language attention shape divides the TP size.