> This page is for version 0.7.0 · 26.10.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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.te_patches

Transformer Engine compatibility patches.

Runtime monkey-patches applied directly to TE classes in memory so they
take effect immediately in the current process.

Call `apply_te_patches()` early in the process, before TE optimizers are used.

## Module Contents

### Classes

| Name                                                                           | Description                                                                               |
| ------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------------- |
| [`_FusedAdamGroupPlan`](#nemo_automodel-shared-te_patches-_FusedAdamGroupPlan) | The multi-tensor lists of one parameter group, in the order `FusedAdam.step` builds them. |

### Functions

| Name                                                                                                                     | Description                                                                              |
| ------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------- |
| [`_apply_fused_adam_cached_lists_patch`](#nemo_automodel-shared-te_patches-_apply_fused_adam_cached_lists_patch)         | Reuse FusedAdam's multi-tensor lists across steps instead of rebuilding them every step. |
| [`_apply_fused_adam_empty_shard_patch`](#nemo_automodel-shared-te_patches-_apply_fused_adam_empty_shard_patch)           | Keep zero-element FSDP shards out of TE's native multi-tensor kernel.                    |
| [`_apply_fused_adam_in_place_load_patch`](#nemo_automodel-shared-te_patches-_apply_fused_adam_in_place_load_patch)       | Load FusedAdam state into the optimizer's existing state buffers.                        |
| [`_apply_fused_adam_quantized_tensor_patch`](#nemo_automodel-shared-te_patches-_apply_fused_adam_quantized_tensor_patch) | Patch FusedAdam.\_initialize_state to handle QuantizedTensor params.                     |
| [`apply_te_patches`](#nemo_automodel-shared-te_patches-apply_te_patches)                                                 | Apply all Transformer Engine runtime patches.                                            |

### Data

[`_FUSED_ADAM_NO_PLAN`](#nemo_automodel-shared-te_patches-_FUSED_ADAM_NO_PLAN)

[`_TE_PATCHES_APPLIED`](#nemo_automodel-shared-te_patches-_TE_PATCHES_APPLIED)

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

### API

```python
class nemo_automodel.shared.te_patches._FusedAdamGroupPlan(
    params: list
)
```

The multi-tensor lists of one parameter group, in the order `FusedAdam.step` builds them.

**`__slots__`**

---

**`f16_m`** `list = []`

---

**`f16_master`** `list = []`

---

**`f16_params`** `list = []`

---

**`f16_v`** `list = []`

---

**`f32_m`** `list = []`

---

**`f32_params`** `list = []`

---

**`f32_v`** `list = []`

---

**`n_params`** `= len(params)`

---

**`skipped`** `list = []`

---

```python
nemo_automodel.shared.te_patches._FusedAdamGroupPlan.still_valid(
    group: dict
) -> bool
```

True when the group's parameters and gradient presence match the plan.

```python
nemo_automodel.shared.te_patches._apply_fused_adam_cached_lists_patch() -> None
```

Reuse FusedAdam's multi-tensor lists across steps instead of rebuilding them every step.

`FusedAdam.step` walks every parameter on every step — three DTensor `.data` dispatches,
the optimizer-state lookups, the dtype branches and the DTensor-parity asserts — and only
then launches the fused kernels. Under FSDP2 every parameter is a DTensor and a large MoE
model has on the order of a thousand sharded parameter tensors per rank, so this is tens of
milliseconds of Python per step during which the GPU is idle (DeepSeek-V4.1 Flash on 16
GB200 nodes: `optimizer` timer 85 ms, Adam kernels 27 ms). Parameters, optimizer states
and master weights are the same tensor objects on every step; only the gradient objects
change (FSDP2 installs a fresh sharded gradient after each reduce-scatter). The patch builds
the per-group plan once — the same lists in the same order as the original loop — and on
later steps only re-reads `param.grad`, so the kernels receive identical inputs.

Anything the plan does not cover falls back to the original `step` for that call:
capturable mode, decoupled gradients, a grad scaler or closure, FP8 / block-scaled
parameters, optimizer states that need an unscale pass (fp16 or non-fused bf16 moments),
uninitialised states (the first step), or a gradient that appears or disappears. The plan is
dropped on `load_state_dict` and `add_param_group`; a group that can never take the cached
path (FP8 parameters, fp16 moments or master weights, mixed fp16 + bf16) is remembered as such,
so the original `step` runs for it without a per-parameter walk first.

Deliberately *not* checked between steps: parameter dtype and shape, the identity of
`state["exp_avg"]` / `state["exp_avg_sq"]` / `state["master_param"]`, and the optimizer's
`master_weights` / `store_param_remainders` settings. The plan holds the state tensors by
object, exactly as the original loop would have fetched them, so anything that replaces a
state tensor without going through `load_state_dict` or `add_param_group` (a script
assigning `optimizer.state[p]["exp_avg"] = ...`, or a future TE change) would update a stale
tensor. Nothing in the training paths here does that; it is the one invariant the caller owns.

```python
nemo_automodel.shared.te_patches._apply_fused_adam_empty_shard_patch() -> None
```

Keep zero-element FSDP shards out of TE's native multi-tensor kernel.

TE's multi_tensor_apply metadata writer can overrun its final slot when a
zero-element tensor occupies that slot (TE issue #3202).  Filter complete
parameter/state columns at the Python/native boundary so state_dict keys,
parameter groups, and the update for every nonempty shard stay intact.

```python
nemo_automodel.shared.te_patches._apply_fused_adam_in_place_load_patch() -> None
```

Load FusedAdam state into the optimizer's existing state buffers.

TE's `FusedAdam.load_state_dict` first runs `torch.optim.Optimizer.load_state_dict`, which casts a copy of
every state tensor to the parameter dtype, and then resets each parameter's state and allocates new buffers while
the loaded tensors are still referenced. With fp32 master weights (bf16 parameter plus int16 remainder) and fp32
moments, a resume then holds about 22 bytes per parameter instead of the 10 bytes of optimizer state, which does
not fit large sharded models. The patched method restores the parameter-group hyperparameters, including
FusedAdam's per-group `step`, and copies each loaded state into the existing buffer with `set_scaled_state`;
a buffer is allocated only for a state that does not exist yet.

```python
nemo_automodel.shared.te_patches._apply_fused_adam_quantized_tensor_patch() -> None
```

Patch FusedAdam.\_initialize_state to handle QuantizedTensor params.

TE's FusedAdam uses `torch.zeros(param.shape, ...)` /
`torch.empty(param.shape, ...)` in `_initialize_state`, which fails for
QuantizedTensor parameters because their `.shape` does not carry the
correct metadata for allocation.  The fix dequantizes the param first and
uses `torch.zeros_like` / `torch.empty_like` instead.

The fix was merged upstream in TE 2.12 via
[https://github.com/NVIDIA/TransformerEngine/pull/2535](https://github.com/NVIDIA/TransformerEngine/pull/2535).

```python
nemo_automodel.shared.te_patches.apply_te_patches() -> None
```

Apply all Transformer Engine runtime patches.

This function is idempotent and safe to call multiple times.

```python
nemo_automodel.shared.te_patches._FUSED_ADAM_NO_PLAN = object()
```

```python
nemo_automodel.shared.te_patches._TE_PATCHES_APPLIED = False
```

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