nemo_automodel.components.models.common

Attention backend (“torch”, “te”, “sdpa”, “flex”, “eager”, “tilelang”, “cudnn”, “fa4”, or “magi”). “fa4” selects native FlashAttention-4 (CuTe) and requires an INSTALL_FA4=true image. HF-dispatched models use “flash_attention_4”. A rank-2 binary padding mask is converted to varlen metadata, while dense/block masks and indexed document masks without explicit packed metadata are rejected. Inputs must use BSHD layout; pre-packed THD is unsupported. Packed Qwen3Next and DeepSeek V3.2 sparse attention do not support FA4. For DeepSeek V4, “tilelang” enables the TileLang sparse attention, indexer, and Sinkhorn kernels together; “cudnn” keeps the TileLang indexer/Sinkhorn and selects FlashMLA forward with cuDNN backward. For GLM DSA, “tilelang” and “cudnn” select their respective packed sparse-attention kernels. For Qwen3.8-Flash-Next, “flex” selects FlexAttention sparse GQA on CUDA BF16; CPU execution retains the PyTorch numerical oracle.

Benchmark-only. Requires fake_balanced_gate=True with fake_gate_noise=0.0, where routing metadata (tokens per expert, permuted token counts) is identical for every microbatch. Skips the per-microbatch device-to-host reads of that metadata (.tolist() / count_nonzero / dispatcher size checks) by caching the first microbatch’s values, removing recurring host-sync stalls on the hot path. Never enable with a learned gate: cached metadata would go stale and silently corrupt expert dispatch.

torch.compile(fullgraph) the attention module’s forward — both the DeepSeek-V3 MLA and standard GQA attention (e.g. Qwen3-MoE) honor it. Requires attn=“sdpa”, linear=“torch”, rms_norm=“torch”, rope_fusion=False.

torch.compile the mHC (hyper-connection) cores of models that opt in (currently DeepSeek V4 / V4.1): coefficient projection, stream collapse and stream expand, one fused kernel each, static shapes (a new sequence length compiles them again; past dynamo’s recompile limit a core runs eager). Same lazy once-per-process pattern as compile_situ; numerics are allclose to eager but not bitwise-identical.

torch.compile the fp32 RMSNorm chain of models that opt in (currently Kimi K3 and DeepSeek V4.1), fusing cast/pow/mean/rsqrt/mul into one kernel. Same lazy once-per-process pattern as compile_situ; numerics are allclose to eager but not bitwise-identical.

torch.compile the fp32 router-weight multiply applied to expert outputs (GroupedExperts / GroupedExpertsDeepEP / GroupedExpertsTE, all MoE models), so each pass is one fused kernel instead of the eager per-chunk cast/multiply/ cast/slice-assign loop. Same lazy once-per-process pattern as compile_situ; numerics are allclose to eager but not bitwise-identical.

torch.compile the SiTU expert activation of models that opt in (currently Kimi K3): the weighted-SiTU forward and backward run as one fused whole-tensor kernel each (low-precision in/out, fp32 math inside, no chunk loop), the dense / shared-expert SituAndMul uses a compiled core, and the attn-res mixing chain is compiled as well. Compiled numerics are allclose to eager but not bitwise-identical.

Scoped partial CUDA-graph configuration.

MoE token dispatcher. “torch” uses DTensor all-gather/reduce-scatter, “deepep” uses DeepEP, “hybridep” uses HybridEP, “uccl_ep” uses UCCL-EP, and “mok” uses Mixture-of-Kittens’ fused communication and SwiGLU megakernel.

Whether DeepEP/UCCL-EP dispatch and combine should return asynchronously and allocate their outputs on the communication stream.

HybridEP only, dynamic routing. Run HybridEP dispatch in its non-blocking mode with output buffers sized to the first microbatch’s permuted row count times this factor (EP-group max, 4-token aligned) instead of letting every dispatch drain the compute stream to read the exact count. Overflow trips a device-side assert. None (default) keeps the blocking reference path. The capacity is calibrated once per communication manager from its first real batch and kept for the run; with dispatcher_share_token_dispatcher=True (default) that is one calibration, taken on the first local MoE layer’s first microbatch, for every layer, so choose a factor that covers the largest routing skew expected across layers and batches. Not supported for experts with expert_bias (GroupedExpertsDeepEP rejects the combination).

HybridEP only. Declare that every EP rank dispatches the same row count, so the per-dispatch EP-group max all-reduce (and its host sync) that derives the pad size is skipped and the count is only aligned. Default False; keep it False for variable-length or in-batch-packed inputs.

Send compact top-k expert indices to HybridEP instead of the dense Boolean routing map. This can reduce metadata traffic, but remains opt-in because its latency depends on topology and model shape.

Optional number of dispatch permutation blocks. None preserves HybridEP’s default.

Optional number of combine unpermutation blocks. None preserves HybridEP’s default.

Optional number of SMs used by HybridEP routing-metadata preprocessing.

Fuse HybridEP token permutation into dispatch and unpermutation into combine. Has no effect for other dispatcher backends.

HybridEP only; effective under full-block activation checkpointing (not installed for checkpoint_moe_only regions, which end at the combine the recompute never re-runs). "none" (default) recomputes the whole MoE layer in the backward. "combine" keeps a clone of every MoE layer’s combine output ([tokens, hidden] in the activation dtype) from the checkpoint forward and returns it in the recompute instead of re-running the unpermute, the EP-group barrier and the combine all-to-all. "moe_tail" additionally builds the down projection’s autograd node in the recompute without running its GEMM (its output only fed the skipped combine; torch_mm experts without expert bias and with the router weight applied inside the activation) (a post-expert callback that reads the output keeps the GEMM). Costs one [tokens, hidden] activation per MoE layer per in-flight microbatch; gradients are unchanged.

Whether flex token dispatchers share a communication manager instance across MoE layers.

Removed and ignored. Logs a warning if set; configure “dispatcher” and “experts” explicitly instead.

Whether to enable FSDP2 optimizations.

Whether to enable HuggingFace state dict adapter.

MoE expert GEMM backend. “torch” uses per-expert loop, “te” uses TE GroupedLinear, “torch_mm” uses torch._grouped_mm, “gmm” is a deprecated compatibility alias for “torch_mm”, and “torch_mm_mxfp8” uses torch._grouped_mm dispatch but routes the expert grouped GEMMs through torchao’s MXFP8 scaled grouped GEMM (training-only; GB200/sm_100+ with torchao installed, else falls back to torch._grouped_mm at runtime).

If True, replace the learned Gate with FakeBalancedGate that assigns tokens to experts without learned routing weights.

Noise level [0, 1] for FakeBalancedGate. When > 0, uses biased topk selection seeded from the input content so routing varies dynamically across training steps (like real Gate) while remaining deterministic for activation checkpointing recompute (same input = same routing). Only used when fake_balanced_gate=True.

Optional dtype override for the gate computation. Accepts torch.dtype or string (e.g., “torch.float32”, “float32”).

Linear layer backend (“torch”, “te”, or “quack”).

Mixture-of-Kittens tuning settings used when dispatcher="mok".

RMSNorm backend (“torch”, “torch_fp32”, “te”, or “quack”).

Rotary embedding backend (“torch”, “quack” or “triton”). QuACK is currently integrated for Llama-family rotary embeddings; “triton” selects the fused Triton rotation of DeepSeek V4 / V4.1 (other models keep the torch path).

Whether to use fused RoPE (requires TE).

run the shared experts of MoE models (the shared MoE class, e.g. DeepSeek V4 / V4.1, and Kimi K3) on a side CUDA stream so their GEMMs overlap the expert-parallel dispatch / combine communication of the routed path; numerics unchanged. Default False.

Sparse-attention backend. “generic” preserves each model’s existing sparse mask plus attn path; “msa” selects the optional SM100 MSA kernels a model provides for its sparse layers only.

Prefer TE 2.19 native grouped-tensor GEMMs for experts="te". Requires dispatcher="deepep", "hybridep", or "uccl_ep". Device and precision support follows TE’s GroupedLinear contract. Parameters and checkpoints retain their existing layouts.

Per-layer modules to capture with Transformer Engine, following Megatron Core’s module names. AutoModel supports whole attn, moe_router, and moe_preprocess scopes, plus the AutoModel-specific narrow te_dpa scope. An empty list disables CUDA graphs.

Routing all-gather chunk size in bytes.

Communication SMs reserved during the backward megakernel.

Communication SMs reserved during the forward megakernel.

Routed-token ring-buffer capacity and a multiple of minibatch_size.

Routed-token overlap granularity, divisible by 256.

Routed-expert compute precision. "bf16" passes the materialized BF16 weights directly to MoK; "mxfp8" quantizes them with MoK’s E4M3 data and E8M0 block-scale format.

Worst-case fraction of routes sent to one EP rank.

The language-model head module, or None on a pipeline stage that does not own it.

Final hidden states, shaped [T, H] or [B, S, H].

0 to project every position; a positive int to keep the last N positions; or a tensor of position indices.

Whether the inputs were THD/packed; if so, a 2D logits (and hidden-states) result is unsqueezed back to a leading batch dim of 1.

Project in fp32 and cast the result back to the input dtype. Ignored when lm_head is None.

Attach the final hidden states to the output.

A HuggingFace model config object.

Backend implementation (“te”, “torch”, or “quack”)

Input features

Output features

Whether to use bias

Device to create module on (None uses PyTorch default, typically CPU)

Parameter dtype

Backend implementation (“te”, “torch”, “torch_fp32”, or “quack”)

• “te”: Transformer Engine fused RMSNorm kernel
• “torch”: PyTorch native nn.RMSNorm (computes in input dtype)
• “torch_fp32”: torch.compiled fp32 RMSNorm for training stability
• “quack”: QuACK CuTe DSL RMSNorm kernel

Normalized dimension

Epsilon for numerical stability

Device to create module on (None uses PyTorch default, typically CPU)

Parameter dtype