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
`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_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.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_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__)