nemo_automodel.shared.te_patches

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

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

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.

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.