nemo_automodel.shared.torch_patches | NVIDIA NeMo AutoModel

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.

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._TORCH_PATCHES_APPLIED = False

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