nemo_automodel.shared.torch_patches

Torch compatibility patches.

These patches are intentionally NOT applied at import nemo_automodel time to keep tokenizer-only imports lightweight. Call apply_torch_patches() from code paths that already depend on torch (training / distributed / dataloading).

Module Contents

Functions

Name	Description
`_widest_float_dtype`	Return the float dtype among `dtypes` that every other one converts into losslessly.
`apply_torch_patches`	Apply small, version/packaging-specific torch monkey patches.
`patch_fsdp_accumulated_grad_guard`	Guard FSDP2 post-backward against params that were never unsharded.
`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

_TORCH_PATCHES_APPLIED

_logger

API

nemo_automodel.shared.torch_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.shared.torch_patches.apply_torch_patches() -> None

Apply small, version/packaging-specific torch monkey patches.

This function is idempotent and safe to call multiple times.

nemo_automodel.shared.torch_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.

nemo_automodel.shared.torch_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 :func: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.shared.torch_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.shared.torch_patches._TORCH_PATCHES_APPLIED = False

nemo_automodel.shared.torch_patches._logger = logging.getLogger(__name__)