nemo_automodel.shared.te_patches
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
Functions
Data
API
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.
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.
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.
Apply all Transformer Engine runtime patches.
This function is idempotent and safe to call multiple times.