nemo_automodel.components.distributed.parallelizer_utils

Module Contents

Functions

Name	Description
`_call_fully_shard`	-
`_fully_shard`	-
`_get_module_from_path`	-
`_group_params_by_dtype`	-
`_make_compute_dtype_fn`	Build the per-parameter compute dtype resolver used to group FSDP units.
`_mp_policy_with_param_dtype`	-
`configure_fsdp_unused_param_reduction`	Reduce zero gradients for FSDP parameters unused on a local CP rank.
`fully_shard_across_blocks`	Put every block’s instance of the named fp32 parameter holders into one FSDP unit.
`fully_shard_by_dtype`	Fully shard a module so every parameter computes in its required dtype.
`get_internal_fsdp_mp_policy`	Clone an FSDP policy without imposing an external output dtype.
`iter_maximal_uniform_dtype_subtrees`	Traverse `module` and yield maximal submodules whose entire subtree has a unified dtype.
`reject_unsupported_mtp_cp`	Reject enabled MTP when the model has not declared CP support.
`reject_unsupported_mtp_cp_pp`	Reject unsupported MTP+CP+PP before parallelizing the model.

Data

UniformSubtreeItem

logger

API

nemo_automodel.components.distributed.parallelizer_utils._call_fully_shard(
    module: torch.nn.Module,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None,
    reshard_after_forward: bool | int | None = None,
    ignored_params: set[torch.nn.Parameter] | None = None,
    shard_module: typing.Callable[..., None]
) -> None

nemo_automodel.components.distributed.parallelizer_utils._fully_shard(
    module: torch.nn.Module,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None,
    reshard_after_forward: bool | int | None = None,
    ignored_params: set[torch.nn.Parameter] | None = None,
    shard_module: typing.Callable[..., None]
) -> None

nemo_automodel.components.distributed.parallelizer_utils._get_module_from_path(
    layer: torch.nn.Module,
    path: str
) -> torch.nn.Module

nemo_automodel.components.distributed.parallelizer_utils._group_params_by_dtype(
    layer: torch.nn.Module,
    dtype_of: typing.Callable[[torch.Tensor], torch.dtype] | None = None,
    ignored_params: set[torch.nn.Parameter] | None = None
) -> typing.Dict[torch.dtype, typing.List[torch.nn.Parameter]]

nemo_automodel.components.distributed.parallelizer_utils._make_compute_dtype_fn(
    module: torch.nn.Module,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None,
    fp32_compute_module_names: typing.Tuple[str, ...],
    ignored_params: set[torch.nn.Parameter] | None = None
) -> typing.Callable[[torch.Tensor], torch.dtype]

Build the per-parameter compute dtype resolver used to group FSDP units.

The compute dtype of a floating tensor is resolved by precedence:

Pinned fp32 — the tensor’s name matches fp32_compute_module_names (from the model’s _keep_in_fp32_modules_strict). Authoritative, works even from-scratch / quantized where there is no checkpoint to read.
HF-recorded dtype — tensor._hf_compute_dtype, the checkpoint’s original dtype recorded at load time (see _restore_loaded_model_dtype). This makes any checkpoint-loaded model keep its intrinsically-fp32 params in fp32 compute automatically, even after storage was upcast for fp32 master weights.
Fallback — when the tensor carries no compute hint, an fp32 storage under a lower-precision policy is an fp32 master weight and computes in mp_policy.param_dtype (the requested compute dtype, typically bf16); any other storage keeps its own dtype (and so does the fp32 case when no policy is given). Resolved per-param — a single genuinely lower-precision sibling (e.g. Qwen3.5-MoE’s bf16 shared_expert_gate) no longer forces the layer’s fp32 master weights into fp32 compute. Intrinsic fp32 is already covered by #1/#2; the grouping still keeps each FSDP unit uniform in trainable storage dtype. See NVIDIA-NeMo/Automodel#3327.

Non-floating tensors always keep their storage dtype.

nemo_automodel.components.distributed.parallelizer_utils._mp_policy_with_param_dtype(
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None,
    param_dtype: torch.dtype
) -> torch.distributed.fsdp.MixedPrecisionPolicy | None

nemo_automodel.components.distributed.parallelizer_utils.configure_fsdp_unused_param_reduction(
    module: torch.nn.Module
) -> int

Reduce zero gradients for FSDP parameters unused on a local CP rank.

Packed or modality-dependent context-parallel batches may execute a module on only a subset of ranks. FSDP must still issue the same reduce-scatter sequence everywhere; otherwise a rank with grad is None can omit a collective and discard peer contributions. PyTorch’s public API fills the missing local contribution with zero, analogous to DDP unused-parameter handling. AutoModel keeps a compatibility fallback for supported PyTorch versions that predate that public API.

Parameters:

module

nn.Module

Root module containing the FSDP units to configure.

Returns: int

Number of FSDP units configured.

nemo_automodel.components.distributed.parallelizer_utils.fully_shard_across_blocks(
    blocks: typing.Sequence[torch.nn.Module],
    module_names: typing.Sequence[str],
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None,
    shard_module: typing.Callable[..., None] = fully_shard
) -> set[torch.nn.Parameter]

Put every block’s instance of the named fp32 parameter holders into one FSDP unit.

FSDP2 registers a unit’s post-backward hook on the forward inputs of its modules. A module whose forward ignores its input (DeepSeek-V4’s attention sinks: DeepseekV4FP32Parameter returns its weight and drops the reference tensor) therefore never triggers it, and FSDP2’s root post-backward callback reduce-scatters such units one by one after the whole backward — on the critical path, CPU-bound, with the GPU idle (40 units ≈ 40 serialised collectives per step). One unit for all blocks’ instances turns that into a single reduce-scatter; reshard_after_forward=False keeps the few kilobytes unsharded so the step also pays one all-gather instead of one per block and per activation-checkpoint recompute. That is also why reshard_after_forward is hard-coded here rather than threaded from the caller like every other fully_shard call in the parallelizer: a resharded unit would be re-gathered inside every checkpoint recompute, which is exactly the per-block cost this unit removes, and the memory it keeps resident is a few kilobytes. Numerics are unchanged: the parameters stay FSDP-sharded fp32 DTensors with the same reduce dtype.

Parameters:

blocks

Sequence[nn.Module]

Decoder blocks (checkpoint wrappers included) whose submodules are searched.

module_names

Sequence[str]

Dotted module names relative to a block ("attn.sinks_param"); each name becomes one cross-block unit. Sourced from the model’s _fsdp_shard_across_blocks.

mesh

DeviceMesh

FSDP mesh of the blocks.

mp_policy

MixedPrecisionPolicy | None

The blocks’ mixed-precision policy; the unit computes and reduces in fp32 and returns its natural fp32 outputs (the policy’s output_dtype is cleared, as for every internal unit).

offload_policy

OffloadPolicy | None

Parameter/gradient offload policy.

shard_module

Callable[..., None]Defaults to fully_shard

FSDP primitive (fully_shard or a model parallelizer’s override).

Returns: set[nn.Parameter]

The parameters now owned by the cross-block units. Callers pass them as ignored_params

Raises:

ValueError: when a matched module holds a non-fp32 parameter — the unit’s policy is fp32 only.

nemo_automodel.components.distributed.parallelizer_utils.fully_shard_by_dtype(
    module: torch.nn.Module,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None,
    fp32_compute_module_names: typing.Tuple[str, ...] = (),
    reshard_after_forward: bool | int | None = None,
    ignored_params: set[torch.nn.Parameter] | None = None,
    model_parallelizer: nemo_automodel.components.distributed.parallelizer.ModelParallelizer | None = None
) -> None

Fully shard a module so every parameter computes in its required dtype.

The intent is simple: compute everything in mp_policy.param_dtype (e.g. bf16) except parameters that must stay in fp32 — their FSDP unit gets param_dtype=fp32 while the rest of the module computes in the policy dtype. A parameter “must stay fp32” if it is pinned via fp32_compute_module_names or HF stored it in fp32 (see _make_compute_dtype_fn for the full precedence). This decouples compute dtype from storage dtype, so fp32 master weights (uniform fp32 storage) still compute in bf16 for the bulk.

Implementation: group the module’s parameters by their resolved compute dtype and shard so each FSDP unit is compute-dtype-uniform. The three cases below differ only in sharding granularity:

1 compute dtype -> shard the whole module once.
2 compute dtypes -> shard the minority-dtype subtrees on their own, then shard the parent with the majority dtype (keeps the bulk as one FSDP unit).
3+ compute dtypes -> shard every maximal compute-dtype-uniform subtree on its own.

Dtype-specific child units are internal to the enclosing module’s forward, so they preserve their module’s natural output dtype. Any enclosing FSDP boundary created by this function retains the caller’s output_dtype as its external contract.

Parameters:

fp32_compute_module_names

Tuple[str, ...]Defaults to ()

Parameter/buffer name substrings that must compute in fp32 (e.g. ("_fp32_params",) for Qwen3.5’s GatedDeltaNet fp32 holder). Sourced from the model’s _keep_in_fp32_modules_strict. Matched callable modules must cast their own inputs when required; their nested FP32 FSDP units preserve the parent activation dtype at the module boundary.

reshard_after_forward

bool | int | NoneDefaults to None

Optional FSDP2 reshard override for this module. None leaves the caller’s default FSDP2 behavior unchanged.

ignored_params

set[nn.Parameter] | NoneDefaults to None

Parameters already owned by another FSDP or parallelism unit. They are excluded from dtype grouping and forwarded to the enclosing FSDP unit.

model_parallelizer

ModelParallelizer | NoneDefaults to None

Optional model sidecar that owns the FSDP primitive.

nemo_automodel.components.distributed.parallelizer_utils.get_internal_fsdp_mp_policy(
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None
) -> torch.distributed.fsdp.MixedPrecisionPolicy | None

Clone an FSDP policy without imposing an external output dtype.

Internal FSDP units are implementation details inside a parent module’s forward. Their outputs may feed an unwrapped sibling before another FSDP input cast, so they preserve the wrapped module’s natural output dtype.

Parameters:

mp_policy

MixedPrecisionPolicy | None

Mixed-precision policy inherited from the enclosing FSDP boundary, or None when mixed precision is disabled.

Returns: MixedPrecisionPolicy | None

A cloned policy with output_dtype=None, or None when no policy

nemo_automodel.components.distributed.parallelizer_utils.iter_maximal_uniform_dtype_subtrees(
    module: torch.nn.Module,
    include_buffers: bool = True,
    tensor_pred: typing.Callable[[torch.Tensor], bool] | None = None,
    dtype_of: typing.Callable[[torch.Tensor], torch.dtype] | None = None,
    return_paths: bool = False
) -> typing.Iterator[nemo_automodel.components.distributed.parallelizer_utils.UniformSubtreeItem]

Traverse module and yield maximal submodules whose entire subtree has a unified dtype.

include_buffers: include buffers in dtype unification checks.
tensor_pred: predicate to choose which tensors to consider (default: all). Example: tensor_pred=torch.is_floating_point (to consider only FP tensors)
dtype_of: maps a tensor to the dtype used for unification (default: its storage dtype t.dtype). Pass a custom function to group by compute dtype rather than storage dtype.
return_paths: if True, yields (qualified_name, module, dtype); else (module, dtype).

Notes:

If a module subtree has no tensors passing tensor_pred, it is ignored.
Maximality ensures no yielded module is a strict child of another yielded module.

nemo_automodel.components.distributed.parallelizer_utils.reject_unsupported_mtp_cp(
    model: torch.nn.Module
) -> None

Reject enabled MTP when the model has not declared CP support.

nemo_automodel.components.distributed.parallelizer_utils.reject_unsupported_mtp_cp_pp(
    model: torch.nn.Module,
    pp_size: int | None = None
) -> None

Reject unsupported MTP+CP+PP before parallelizing the model.

Call only when context parallelism is active.

Parameters:

model

nn.Module

Model whose MTP capabilities are checked.

pp_size

int | NoneDefaults to None

Pipeline mesh size, when available. The legacy stage hook also detects trimmed stages when the supplied mesh omits PP.

nemo_automodel.components.distributed.parallelizer_utils.UniformSubtreeItem = Union[Tuple[nn.Module, torch.dtype], Tuple[str, nn.Module, torch.dtype]]

nemo_automodel.components.distributed.parallelizer_utils.logger = logging.getLogger(__name__)