nemo_automodel.components.distributed.fsdp_patches

PyTorch FSDP compatibility patches owned by distributed infrastructure.

Module Contents

Functions

Name	Description
`_fsdp_state_param_groups`	Return an FSDP2 state’s parameter groups across PyTorch versions that keep one or several.
`_has_local_unsharded_grad`	Whether this rank’s backward produced a gradient for `fsdp_param` since its last reduction.
`_restore_unused_param_grads`	Reset the gradients that a reduction created for parameters no rank of its domain used.
`_widest_float_dtype`	Return the float dtype among `dtypes` that every other one converts into losslessly.
`patch_fsdp_accumulated_grad_guard`	Guard FSDP2 post-backward against params that were never unsharded.
`patch_fsdp_restore_unused_param_grads`	Keep `grad is None` for parameters that no rank of their reduction used.
`patch_fsdp_uniform_reduce_dtype`	Give every FSDP2 reduce-scatter group local gradients of one dtype.
`patch_fsdp_unused_param_reduction`	Backport FSDP2 unused-parameter reduction when the public API is absent.

Data

_FSDP_ACCUMULATED_GRAD_REMOVED_VERSION

API

nemo_automodel.components.distributed.fsdp_patches._fsdp_state_param_groups(
    state: typing.Any
) -> list[typing.Any]

Return an FSDP2 state’s parameter groups across PyTorch versions that keep one or several.

nemo_automodel.components.distributed.fsdp_patches._has_local_unsharded_grad(
    fsdp_param: typing.Any
) -> bool

Whether this rank’s backward produced a gradient for fsdp_param since its last reduction.

nemo_automodel.components.distributed.fsdp_patches._restore_unused_param_grads(
    param_groups: list[typing.Any]
) -> None

Reset the gradients that a reduction created for parameters no rank of its domain used.

Parameters:

param_groups

list[Any]

FSDP2 parameter groups holding a usage probe from this backward, in the same order on every rank of their reduction domains.

nemo_automodel.components.distributed.fsdp_patches._widest_float_dtype(
    dtypes: collections.abc.Iterable[typing.Any]
) -> typing.Any

Return the float dtype among dtypes that every other one converts into losslessly.

Parameters:

dtypes

Iterable[Any]

Gradient dtypes from a single reduce-scatter group.

Returns: Any

The dtype with the largest element size; ties resolve to float32 over

nemo_automodel.components.distributed.fsdp_patches.patch_fsdp_accumulated_grad_guard() -> None

Guard FSDP2 post-backward against params that were never unsharded.

PyTorch FSDP2 creates _unsharded_param lazily from an FSDP unit’s forward pre-hook. If a separately wrapped unit is skipped by the batch (for example a vision tower on text-only data), deferred post-backward can dereference that missing field. Missing lazy state means there is no unsharded grad to upcast, so the exact missing-field case can return early. Applies only before torch 2.15, which removed the upcast step this guards.

nemo_automodel.components.distributed.fsdp_patches.patch_fsdp_restore_unused_param_grads() -> None

Keep grad is None for parameters that no rank of their reduction used.

The unused-parameter fill (PyTorch’s public API or patch_fsdp_unused_param_reduction) gives a locally unused parameter a zero contribution so every rank issues the same reduce-scatter. When no rank of the parameter’s reduction domain used it, that reduction turns a missing gradient into an exact zero, which an optimizer treats as a real gradient: AdamW advances its step, applies decoupled weight decay and keeps moving the weight with residual momentum. Without the fill the gradient stays None and the optimizer skips the parameter.

Each reducing post-backward records which parameters had a local gradient, and the root’s final backward callback reduces those flags over the reduction domain and resets the sharded gradients this reduction created for parameters no rank used. Usage comes from gradient presence, never from a zero value. Install it after the fill so the probe runs before the zero fill; the patch is process-global and idempotent. Under HSDP with a deferred all-reduce, a reduction that folds into a pending partial reduce is not treated as creating the gradient, so a window in which no rank used the parameter in any micro-batch keeps the reduced zero there.

Raises:

RuntimeError: If the installed PyTorch lacks the FSDP2 internals the patch wraps.

nemo_automodel.components.distributed.fsdp_patches.patch_fsdp_uniform_reduce_dtype() -> None

Give every FSDP2 reduce-scatter group local gradients of one dtype.

Gradient accumulation leaves a group holding reduce_dtype accumulations for the parameters used so far, while any parameter whose gradient joins later — a locally unused parameter zero-filled by PyTorch’s public API or patch_fsdp_unused_param_reduction, or one whose gradient lands after its group’s post-backward already ran — contributes param_dtype. foreach_reduce then aborts with FSDP reduce-scatter expects uniform gradient dtype.

Normalize and widen gradients at the last possible moment, inside foreach_reduce itself. That placement matters:

FSDPParam normally unwraps gradients through _get_grad_inner_tensor. PyTorch versions whose public unused-parameter API appends zeros_like(unsharded_param) directly can still leave a DTensor in this list, so unwrap that residual value before sizing the reduce-scatter buffer;
FSDP2’s own bookkeeping (unsharded_param.grad / unsharded_accumulated_grad) is left exactly as upstream leaves it, so no later reader of that state sees anything unusual;
foreach_reduce immediately copies these gradients into a reduce_dtype buffer anyway, so widening first changes no value.

Uniform groups are passed straight through, so the upstream assertion still fires for genuinely inconsistent gradients such as fp8 weights that fail to produce higher-precision ones. The patch is process-global and idempotent.

nemo_automodel.components.distributed.fsdp_patches.patch_fsdp_unused_param_reduction() -> None

Backport FSDP2 unused-parameter reduction when the public API is absent.

The patch is process-global and idempotent. It only fills a missing local gradient with zeros immediately before FSDP2 post-backward reduction, so ranks that skipped a parameter still participate in the same collective as ranks that used it. Callers must first prefer the public FSDPModule.set_reduce_scatter_unused_params API.

Raises:

RuntimeError: If the installed PyTorch exposes neither the public API nor the compatible private FSDP2 implementation.

nemo_automodel.components.distributed.fsdp_patches._FSDP_ACCUMULATED_GRAD_REMOVED_VERSION = (2, 15)