> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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`](#nemo_automodel-shared-torch_patches-_widest_float_dtype)                             | Return the float dtype among `dtypes` that every other one converts into losslessly. |
| [`apply_torch_patches`](#nemo_automodel-shared-torch_patches-apply_torch_patches)                             | Apply small, version/packaging-specific torch monkey patches.                        |
| [`patch_fsdp_accumulated_grad_guard`](#nemo_automodel-shared-torch_patches-patch_fsdp_accumulated_grad_guard) | Guard FSDP2 post-backward against params that were never unsharded.                  |
| [`patch_fsdp_uniform_reduce_dtype`](#nemo_automodel-shared-torch_patches-patch_fsdp_uniform_reduce_dtype)     | Give every FSDP2 reduce-scatter group local gradients of one dtype.                  |
| [`patch_fsdp_unused_param_reduction`](#nemo_automodel-shared-torch_patches-patch_fsdp_unused_param_reduction) | Backport FSDP2 unused-parameter reduction when the public API is absent.             |

### Data

[`_TORCH_PATCHES_APPLIED`](#nemo_automodel-shared-torch_patches-_TORCH_PATCHES_APPLIED)

[`_logger`](#nemo_automodel-shared-torch_patches-_logger)

### API

```python
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

```python
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.

```python
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.

```python
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.

```python
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.

```python
nemo_automodel.shared.torch_patches._TORCH_PATCHES_APPLIED = False
```

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