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

## Module Contents

### Functions

| Name                                                                                                                     | Description                                                           |
| ------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------- |
| [`_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

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

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

### API

```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._TE_PATCHES_APPLIED = False
```

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