nemo_automodel.shared.te_patches | NVIDIA NeMo AutoModel

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

Functions

Name	Description
`_apply_fused_adam_empty_shard_patch`	Keep zero-element FSDP shards out of TE’s native multi-tensor kernel.
`_apply_fused_adam_quantized_tensor_patch`	Patch FusedAdam._initialize_state to handle QuantizedTensor params.
`apply_te_patches`	Apply all Transformer Engine runtime patches.

Data

_TE_PATCHES_APPLIED

_logger

API

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.

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.

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.

nemo_automodel.shared.te_patches._TE_PATCHES_APPLIED = False

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