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

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.

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

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