> This page is for version 0.5.0 · 26.06. > For other versions, use one of these documentation indexes: > - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt > - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt > - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt > - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt > For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # nemo_automodel.components.models.common.utils ## Module Contents ### Classes | Name | Description | | --------------------------------------------------------------------------------------------- | -------------------------------------------------------------- | | [`BackendConfig`](#nemo_automodel-components-models-common-utils-BackendConfig) | Backend configuration for model components. | | [`CudaGraphConfig`](#nemo_automodel-components-models-common-utils-CudaGraphConfig) | Configuration for scoped partial CUDA graphs. | | [`Float32RMSNorm`](#nemo_automodel-components-models-common-utils-Float32RMSNorm) | RMSNorm with explicit fp32 computation for training stability. | | [`MoKBackendConfig`](#nemo_automodel-components-models-common-utils-MoKBackendConfig) | Configuration for the optional Mixture-of-Kittens MoE backend. | | [`TEFp8Config`](#nemo_automodel-components-models-common-utils-TEFp8Config) | Configuration for Transformer Engine FP8 quantization. | | [`_MoKFunctionalConfig`](#nemo_automodel-components-models-common-utils-_MoKFunctionalConfig) | Structural type for the optional `mok.functional.MoKConfig`. | ### Functions | Name | Description | | ----------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------- | | [`_float32_rms_norm_backward`](#nemo_automodel-components-models-common-utils-_float32_rms_norm_backward) | Differentiate the fp32 RMSNorm computation. | | [`_float32_rms_norm_fwd`](#nemo_automodel-components-models-common-utils-_float32_rms_norm_fwd) | Compiled fp32 RMSNorm forward — standalone function to minimize dynamo guards. | | [`_float32_rms_norm_impl`](#nemo_automodel-components-models-common-utils-_float32_rms_norm_impl) | Compute RMSNorm in fp32 with an opaque, device-independent forward. | | [`_float32_rms_norm_meta`](#nemo_automodel-components-models-common-utils-_float32_rms_norm_meta) | Describe the output metadata for inputs x \[..., hidden] and weight \[hidden]. | | [`_float32_rms_norm_setup_context`](#nemo_automodel-components-models-common-utils-_float32_rms_norm_setup_context) | Save the inputs required for the RMSNorm gradient. | | [`_float32_rms_norm_sharding`](#nemo_automodel-components-models-common-utils-_float32_rms_norm_sharding) | Keep the normalized axis complete while allowing leading-axis sharding. | | [`_get_fp32_module_keywords`](#nemo_automodel-components-models-common-utils-_get_fp32_module_keywords) | Collect module name patterns that must remain in fp32. | | [`_get_strict_fp32_module_keywords`](#nemo_automodel-components-models-common-utils-_get_strict_fp32_module_keywords) | - | | [`_has_dtensor_params`](#nemo_automodel-components-models-common-utils-_has_dtensor_params) | Check if any model parameter is a DTensor (FSDP2 sharded). | | [`_make_lazy_te_patcher`](#nemo_automodel-components-models-common-utils-_make_lazy_te_patcher) | Return a callable that patches TE modules exactly once. | | [`_restore_fp32_buffers`](#nemo_automodel-components-models-common-utils-_restore_fp32_buffers) | Cast only matching buffers (not parameters) back to float32. | | [`_restore_fp32_modules`](#nemo_automodel-components-models-common-utils-_restore_fp32_modules) | Cast modules or individual tensors matching *fp32\_keywords* back to float32. | | [`_restore_fp32_tensor_snapshots`](#nemo_automodel-components-models-common-utils-_restore_fp32_tensor_snapshots) | Restore fp32-preserved tensors from pre-cast snapshots. | | [`_snapshot_fp32_tensors`](#nemo_automodel-components-models-common-utils-_snapshot_fp32_tensors) | Clone fp32-preserved tensors before a broad dtype cast. | | [`cast_frozen_modules_to_compute_dtype`](#nemo_automodel-components-models-common-utils-cast_frozen_modules_to_compute_dtype) | Cast frozen floating-point tensors to `compute_dtype`. | | [`cast_model_to_dtype`](#nemo_automodel-components-models-common-utils-cast_model_to_dtype) | Cast model parameters to the target dtype, keeping fp32 modules in full precision. | | [`compute_lm_head_logits`](#nemo_automodel-components-models-common-utils-compute_lm_head_logits) | Project hidden states through `lm_head` and wrap the result. | | [`generation_config_from_model_config`](#nemo_automodel-components-models-common-utils-generation_config_from_model_config) | Generation config a custom causal LM starts with, seeded the way transformers does it. | | [`get_is_first_microbatch`](#nemo_automodel-components-models-common-utils-get_is_first_microbatch) | Get the global IS\_FIRST\_MICROBATCH flag. | | [`get_is_optim_step`](#nemo_automodel-components-models-common-utils-get_is_optim_step) | Get the global IS\_OPTIM\_STEP flag. | | [`get_rope_config`](#nemo_automodel-components-models-common-utils-get_rope_config) | Extract rope configuration from `config.rope_parameters`. | | [`initialize_linear_module`](#nemo_automodel-components-models-common-utils-initialize_linear_module) | Initialize Linear module with the specified backend. | | [`initialize_rms_norm_module`](#nemo_automodel-components-models-common-utils-initialize_rms_norm_module) | Initialize RMSNorm module with the specified backend. | | [`is_tensor_unallocated`](#nemo_automodel-components-models-common-utils-is_tensor_unallocated) | Check if tensor is unallocated (meta tensor, fake tensor, etc.). | | [`load_pretrained_generation_config`](#nemo_automodel-components-models-common-utils-load_pretrained_generation_config) | The generation config a checkpoint carries, or `None` when it carries none. | | [`restore_pretrained_generation_config`](#nemo_automodel-components-models-common-utils-restore_pretrained_generation_config) | Replace a model's `generation_config` with the one its checkpoint carries, if any. | | [`set_is_first_microbatch`](#nemo_automodel-components-models-common-utils-set_is_first_microbatch) | Set the global IS\_FIRST\_MICROBATCH flag for FP8 weight caching. | | [`set_is_optim_step`](#nemo_automodel-components-models-common-utils-set_is_optim_step) | Set the global IS\_OPTIM\_STEP flag. | | [`yield_fp32_model`](#nemo_automodel-components-models-common-utils-yield_fp32_model) | Run a block with the model temporarily in fp32, then cast it to `restore_dtype`. | ### Data [`AttentionBackend`](#nemo_automodel-components-models-common-utils-AttentionBackend) [`HAVE_DEEP_EP`](#nemo_automodel-components-models-common-utils-HAVE_DEEP_EP) [`HAVE_TE`](#nemo_automodel-components-models-common-utils-HAVE_TE) [`HAVE_UCCL_EP`](#nemo_automodel-components-models-common-utils-HAVE_UCCL_EP) [`HUB_LOADING_KWARGS`](#nemo_automodel-components-models-common-utils-HUB_LOADING_KWARGS) [`IS_FIRST_MICROBATCH`](#nemo_automodel-components-models-common-utils-IS_FIRST_MICROBATCH) [`IS_OPTIM_STEP`](#nemo_automodel-components-models-common-utils-IS_OPTIM_STEP) [`_patch_te_modules`](#nemo_automodel-components-models-common-utils-_patch_te_modules) [`logger`](#nemo_automodel-components-models-common-utils-logger) ### API ```python class nemo_automodel.components.models.common.utils.BackendConfig( attn: nemo_automodel.components.models.common.utils.AttentionBackend = 'te' if HAVE_TE and torch.c..., sparse_attn: typing.Literal['generic', 'msa'] = 'generic', linear: typing.Literal['torch', 'te', 'quack'] = 'te' if HAVE_TE and torch.c..., rms_norm: typing.Literal['torch', 'torch_fp32', 'te', 'quack'] = 'torch_fp32', rope: typing.Literal['torch', 'quack', 'triton'] = 'torch', rope_fusion: bool = HAVE_TE and torch.cuda.is_a..., experts: typing.Literal['torch', 'te', 'gmm', 'torch_mm', 'torch_mm_mxfp8'] = 'torch_mm' if torch.cuda.is..., dispatcher: typing.Literal['torch', 'deepep', 'hybridep', 'uccl_ep', 'mok'] = 'deepep' if HAVE_DEEP_EP an..., dispatcher_num_sms: int = 20, dispatcher_share_token_dispatcher: bool = True, dispatcher_async_dispatch: bool = False, mok: nemo_automodel.components.models.common.utils.MoKBackendConfig = MoKBackendConfig(), dispatcher_hybridep_permute_fusion: bool = False, dispatcher_hybridep_compact_routing: bool = False, dispatcher_hybridep_num_sms_preprocessing: int | None = None, dispatcher_hybridep_num_blocks_permute: int | None = None, dispatcher_hybridep_num_blocks_unpermute: int | None = None, enable_deepep: bool | None = None, fake_balanced_gate: bool = False, fake_gate_noise: float = 0.0, enable_hf_state_dict_adapter: bool = True, enable_fsdp_optimizations: bool = False, te_fp8: nemo_automodel.components.models.common.utils.TEFp8Config | None = None, gate_precision: str | torch.dtype | None = None, compile_attn: bool = False, compile_situ: bool = False, compile_norm: bool = False, compile_hc: bool = False, shared_expert_overlap: bool = False, compile_router_weight: bool = False, benchmark_static_routing: bool = False, cuda_graph: nemo_automodel.components.models.common.utils.CudaGraphConfig = CudaGraphConfig() ) ``` Dataclass Backend configuration for model components. **`attn`** `AttentionBackend` 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_static_routing`** `bool = False` 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. --- **`compile_attn`** `bool = False` 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. --- **`compile_hc`** `bool = 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. --- **`compile_norm`** `bool = False` 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. --- **`compile_router_weight`** `bool = False` 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. --- **`compile_situ`** `bool = False` 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. --- **`cuda_graph`** `CudaGraphConfig = field(default_factory=CudaGraphConfig)` Scoped partial CUDA-graph configuration. --- **`dispatcher`** `Literal['torch', 'deepep', 'hybridep', 'uccl_ep', 'mok']` 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. --- **`dispatcher_async_dispatch`** `bool = False` Whether DeepEP/UCCL-EP dispatch and combine should return asynchronously and allocate their outputs on the communication stream. --- **`dispatcher_hybridep_compact_routing`** `bool = False` 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. --- **`dispatcher_hybridep_num_blocks_permute`** `int | None = None` Optional number of dispatch permutation blocks. `None` preserves HybridEP's default. --- **`dispatcher_hybridep_num_blocks_unpermute`** `int | None = None` Optional number of combine unpermutation blocks. `None` preserves HybridEP's default. --- **`dispatcher_hybridep_num_sms_preprocessing`** `int | None = None` Optional number of SMs used by HybridEP routing-metadata preprocessing. --- **`dispatcher_hybridep_permute_fusion`** `bool = False` Fuse HybridEP token permutation into dispatch and unpermutation into combine. Has no effect for other dispatcher backends. --- **`dispatcher_num_sms`** `int = 20` --- **`dispatcher_share_token_dispatcher`** `bool = True` Whether flex token dispatchers share a communication manager instance across MoE layers. --- **`enable_deepep`** `bool | None = None` Removed and ignored. Logs a warning if set; configure "dispatcher" and "experts" explicitly instead. --- **`enable_fsdp_optimizations`** `bool = False` Whether to enable FSDP2 optimizations. --- **`enable_hf_state_dict_adapter`** `bool = True` Whether to enable HuggingFace state dict adapter. --- **`experts`** `Literal['torch', 'te', 'gmm', 'torch_mm', 'torch_mm_mxfp8']` 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). --- **`fake_balanced_gate`** `bool = False` If True, replace the learned Gate with FakeBalancedGate that assigns tokens to experts without learned routing weights. --- **`fake_gate_noise`** `float = 0.0` 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. --- **`gate_precision`** `str | dtype | None = None` Optional dtype override for the gate computation. Accepts torch.dtype or string (e.g., "torch.float32", "float32"). --- **`linear`** `Literal['torch', 'te', 'quack']` Linear layer backend ("torch", "te", or "quack"). --- **`mok`** `MoKBackendConfig = field(default_factory=MoKBackendConfig)` Mixture-of-Kittens tuning settings used when `dispatcher="mok"`. --- **`rms_norm`** `Literal['torch', 'torch_fp32', 'te', 'quack'] = 'torch_fp32'` RMSNorm backend ("torch", "torch\_fp32", "te", or "quack"). --- **`rope`** `Literal['torch', 'quack', 'triton'] = 'torch'` 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). --- **`rope_fusion`** `bool = HAVE_TE and torch.cuda.is_available()` Whether to use fused RoPE (requires TE). --- **`shared_expert_overlap`** `bool = False` 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_attn`** `Literal['generic', 'msa'] = 'generic'` 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. --- **`te_fp8`** `TEFp8Config | None = None` --- ```python nemo_automodel.components.models.common.utils.BackendConfig.__post_init__() -> None ``` ```python class nemo_automodel.components.models.common.utils.CudaGraphConfig( modules: list[typing.Literal['attn', 'te_dpa', 'moe_router', 'moe_preprocess']] = list() ) ``` Dataclass Configuration for scoped partial CUDA graphs. **`modules`** `list[Literal['attn', 'te_dpa', 'moe_router', 'moe_preprocess']] = field(default_factory=list)` 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. --- ```python nemo_automodel.components.models.common.utils.CudaGraphConfig.__post_init__() -> None ``` Validate the declarative CUDA-graph configuration. ```python class nemo_automodel.components.models.common.utils.Float32RMSNorm( dim, eps = 1e-05, device = None, dtype = torch.bfloat16 ) ``` **Bases:** `Module` RMSNorm with explicit fp32 computation for training stability. Weights stay in the model's dtype (e.g. bf16) for FSDP2 compatibility. Inputs are upcast to fp32, norm is computed in fp32, and the output is cast back to the original input dtype. **`weight`** --- ```python nemo_automodel.components.models.common.utils.Float32RMSNorm.forward( x ) ``` ```python nemo_automodel.components.models.common.utils.Float32RMSNorm.reset_parameters() ``` ```python class nemo_automodel.components.models.common.utils.MoKBackendConfig( precision: typing.Literal['bf16', 'mxfp8'] = 'bf16', fwd_num_comm_sms: int = 40, bwd_num_comm_sms: int = 28, minibatch_size: int = 4096, macrobatch_size: int = 131072, schedule_capacity_multiplier: float = 0.5, all_gather_top_experts_chunk_bytes: int = 2048 ) ``` Dataclass Configuration for the optional Mixture-of-Kittens MoE backend. **`all_gather_top_experts_chunk_bytes`** `int = 2048` Routing all-gather chunk size in bytes. --- **`bwd_num_comm_sms`** `int = 28` Communication SMs reserved during the backward megakernel. --- **`fwd_num_comm_sms`** `int = 40` Communication SMs reserved during the forward megakernel. --- **`macrobatch_size`** `int = 131072` Routed-token ring-buffer capacity and a multiple of `minibatch_size`. --- **`minibatch_size`** `int = 4096` Routed-token overlap granularity, divisible by 256. --- **`precision`** `Literal['bf16', 'mxfp8'] = 'bf16'` 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. --- **`schedule_capacity_multiplier`** `float = 0.5` Worst-case fraction of routes sent to one EP rank. --- ```python nemo_automodel.components.models.common.utils.MoKBackendConfig.__post_init__() -> None ``` Validate settings that do not depend on a CUDA device or EP group. ```python nemo_automodel.components.models.common.utils.MoKBackendConfig.build() -> nemo_automodel.components.models.common.utils._MoKFunctionalConfig ``` Build the optional MoK runtime configuration. **Returns:** `_MoKFunctionalConfig` A `mok.functional.MoKConfig` with these declarative settings. **Raises:** * `UnavailableError`: If Mixture-of-Kittens is not installed. ```python class nemo_automodel.components.models.common.utils.TEFp8Config( recipe: typing.Literal['current', 'block', 'mxfp8'] | typing.Any = 'current' ) ``` Dataclass Configuration for Transformer Engine FP8 quantization. When present (not None) in BackendConfig, FP8 is enabled. The `recipe` field accepts either a string shorthand (`"current"`, `"block"`, or `"mxfp8"`) or a pre-built TE recipe object (e.g. `Float8CurrentScaling(fp8_dpa=True)`). `"mxfp8"` selects TE's `MXFP8BlockScaling` recipe (e4m3 data + e8m0 block scales). Unlike torchao's MXFP8 grouped GEMM, TE's MXFP8 backward is mature (no e8m0-overflow NaN), which is why GPT-OSS experts (grouped + bias) use the `experts="te"` path with this recipe instead of `experts="torch_mm_mxfp8"`. **`recipe`** `Literal['current', 'block', 'mxfp8'] | Any = 'current'` --- ```python nemo_automodel.components.models.common.utils.TEFp8Config.build_recipe() ``` Build and return the TE FP8 recipe object. If `recipe` is already a TE recipe object (e.g. `Float8CurrentScaling(...)`), it is returned directly. String values `"current"`, `"block"`, and `"mxfp8"` are mapped to the corresponding TE recipe class. ```python nemo_automodel.components.models.common.utils.TEFp8Config.maybe_te_autocast() ``` Return te\_autocast context manager for FP8. ```python class nemo_automodel.components.models.common.utils._MoKFunctionalConfig() ``` Protocol Structural type for the optional `mok.functional.MoKConfig`. **`all_gather_top_experts_chunk_bytes`** `int` --- **`bwd_num_comm_sms`** `int` --- **`fwd_num_comm_sms`** `int` --- **`macrobatch_size`** `int` --- **`minibatch_size`** `int` --- **`schedule_capacity_multiplier`** `float` --- ```python nemo_automodel.components.models.common.utils._float32_rms_norm_backward( ctx, grad_output ) ``` Differentiate the fp32 RMSNorm computation. **Parameters:** **`ctx`** Autograd context containing x \[..., hidden] and weight \[hidden]. --- **`grad_output`** Output gradient of shape \[..., hidden]. --- **Returns:** Gradients for x \[..., hidden] and weight \[hidden] in their input dtypes, ```python nemo_automodel.components.models.common.utils._float32_rms_norm_fwd( x: torch.Tensor, weight: torch.Tensor, eps: float ) -> torch.Tensor ``` Compiled fp32 RMSNorm forward — standalone function to minimize dynamo guards. The opaque forward keeps the same computation in grad and no\_grad contexts. **Parameters:** **`x`** `torch.Tensor` Tensor of shape \[..., hidden], with arbitrary leading dimensions. --- **`weight`** `torch.Tensor` Tensor of shape \[hidden] on the same device as x. --- **`eps`** `float` Epsilon added to the mean square. --- **Returns:** `torch.Tensor` Tensor of shape \[..., hidden] in x's dtype. ```python nemo_automodel.components.models.common.utils._float32_rms_norm_impl( x: torch.Tensor, weight: torch.Tensor, eps: float ) -> torch.Tensor ``` Compute RMSNorm in fp32 with an opaque, device-independent forward. **Parameters:** **`x`** `torch.Tensor` Tensor of shape \[..., hidden], with arbitrary leading dimensions. --- **`weight`** `torch.Tensor` Tensor of shape \[hidden] on the same device as x. --- **`eps`** `float` Epsilon added to the mean square. --- **Returns:** `torch.Tensor` Tensor of shape \[..., hidden] in x's dtype, without aliasing either input. ```python nemo_automodel.components.models.common.utils._float32_rms_norm_meta( x: torch.Tensor, weight: torch.Tensor, eps: float ) -> torch.Tensor ``` Describe the output metadata for inputs x \[..., hidden] and weight \[hidden]. **Parameters:** **`x`** `torch.Tensor` Tensor of shape \[..., hidden], with arbitrary leading dimensions. --- **`weight`** `torch.Tensor` Tensor of shape \[hidden] on the same device as x. --- **`eps`** `float` Epsilon added to the mean square. --- **Returns:** `torch.Tensor` Tensor of shape \[..., hidden] with x's dtype and device. ```python nemo_automodel.components.models.common.utils._float32_rms_norm_setup_context( ctx, inputs, output ) ``` Save the inputs required for the RMSNorm gradient. **Parameters:** **`ctx`** Autograd context that owns the saved tensors. --- **`inputs`** Tuple containing x \[..., hidden], weight \[hidden], and eps. --- **`output`** Tensor of shape \[..., hidden] in x's dtype. --- ```python nemo_automodel.components.models.common.utils._float32_rms_norm_sharding( x, weight, eps ) ``` Keep the normalized axis complete while allowing leading-axis sharding. **Parameters:** **`x`** Tensor metadata of global shape \[..., hidden]. --- **`weight`** Tensor metadata of global shape \[hidden]. --- **`eps`** Epsilon added to the mean square. --- **Returns:** Supported output/input placements on each mesh axis. Leading dimensions ```python nemo_automodel.components.models.common.utils._get_fp32_module_keywords( model: torch.nn.Module ) -> list[str] ``` Collect module name patterns that must remain in fp32. Reads `_keep_in_fp32_modules` and `_keep_in_fp32_modules_strict` from the model (the same attributes HuggingFace transformers uses). **Parameters:** **`model`** `nn.Module` The model to inspect. --- **Returns:** `list[str]` De-duplicated list of module-name keywords to keep in fp32. ```python nemo_automodel.components.models.common.utils._get_strict_fp32_module_keywords( model: torch.nn.Module ) -> list[str] ``` ```python nemo_automodel.components.models.common.utils._has_dtensor_params( model: torch.nn.Module ) -> bool ``` Check if any model parameter is a DTensor (FSDP2 sharded). ```python nemo_automodel.components.models.common.utils._make_lazy_te_patcher() ``` Return a callable that patches TE modules exactly once. Uses a closure instead of module-level global state to track whether the patch has already been applied. The actual `transformer_engine` import is deferred until the first call so that importing this module never triggers heavy native-library loads (flash-attn, CUDA kernels, etc.). Two patches are applied: 1. Unallocated tensor handling: TE kernels don't support meta/fake tensors, so we short-circuit with empty tensors for PP shape inference. 2. is\_first\_microbatch injection: Reads the global IS\_FIRST\_MICROBATCH flag and passes it to TE Linear/GroupedLinear for FP8 weight caching during gradient accumulation (quantize on first microbatch, reuse cached on rest). ```python nemo_automodel.components.models.common.utils._restore_fp32_buffers( model: torch.nn.Module, fp32_keywords: list[str] ) -> None ``` Cast only matching buffers (not parameters) back to float32. Safe for FSDP2-sharded models because buffers are plain tensors, not DTensors managed by FSDP2. **Parameters:** **`model`** `nn.Module` The model (already cast to the target dtype). --- **`fp32_keywords`** `list[str]` Substrings matched against dot-separated module names. --- ```python nemo_automodel.components.models.common.utils._restore_fp32_modules( model: torch.nn.Module, fp32_keywords: list[str] ) -> None ``` Cast modules or individual tensors matching *fp32\_keywords* back to float32. Only safe for unsharded models (plain tensors). FSDP2 requires uniform dtype within each parameter group, so this must not be called on DTensor-sharded models. Keywords may name modules (for example `norm`) or individual parameters (for example `attn_hc.fn`), matching HuggingFace's strict fp32 module declarations. **Parameters:** **`model`** `nn.Module` The model (already cast to the target dtype). --- **`fp32_keywords`** `list[str]` Substrings matched against dot-separated module names. --- ```python nemo_automodel.components.models.common.utils._restore_fp32_tensor_snapshots( model: torch.nn.Module, parameter_snapshots: dict[str, torch.Tensor], buffer_snapshots: dict[str, torch.Tensor] ) -> None ``` Restore fp32-preserved tensors from pre-cast snapshots. ```python nemo_automodel.components.models.common.utils._snapshot_fp32_tensors( model: torch.nn.Module, parameter_keywords: list[str], buffer_keywords: list[str] ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]] ``` Clone fp32-preserved tensors before a broad dtype cast. Casting `fp32 -> bf16 -> fp32` restores the dtype but not the original values. Snapshot the matching tensors first so strict fp32 state such as router correction bias or recurrent-decay parameters is restored exactly. ```python nemo_automodel.components.models.common.utils.cast_frozen_modules_to_compute_dtype( model: torch.nn.Module, compute_dtype: torch.dtype | None ) -> None ``` Cast frozen floating-point tensors to `compute_dtype`. When parameters are stored in fp32 (the fp32-master-weights pattern) while compute runs in bf16, a fully frozen submodule -- such as a frozen vision tower -- can still produce fp32 values that flow into bf16 trainable modules and raise a dtype mismatch in the next matmul. This walks each tensor independently so frozen base weights are still cast when a module also contains trainable adapter weights. In that mixed-module case, its trainable plain-tensor descendants are cast as well so unsharded execution does not introduce a new mismatch in the adapter path. Parameters and buffers are handled differently: * **Parameters** are cast only when they are plain (unsharded) tensors. Sharded (DTensor) params are left as-is: FSDP all-gathers them to the compute dtype during forward, and changing a sharded param's dtype in place would desync FSDP's flat-parameter and `orig_dtype` bookkeeping. * **Buffers** are always cast. Buffers are never sharded, so they stay in their stored dtype regardless of the wrapper; an fp32 buffer (for example a standardization constant) used in a forward op promotes the surrounding bf16 activations to fp32. Tensors whose qualified name matches `_keep_in_fp32_modules` or `_keep_in_fp32_modules_strict` are left in fp32. The function is a no-op when `compute_dtype` is None and for tensors already in `compute_dtype`. Frozen parameters are never updated during training. Trainable plain tensors are cast only in mixed subtrees, where unsharded execution must match the requested compute dtype. **Parameters:** **`model`** `nn.Module` The model, already materialized, checkpoint-loaded, and sharded. --- **`compute_dtype`** `torch.dtype | None` The compute dtype (`mp_policy.param_dtype`); None disables the cast. --- ```python nemo_automodel.components.models.common.utils.cast_model_to_dtype( model: torch.nn.Module, dtype: torch.dtype = torch.bfloat16, skip_modules: tuple[str, ...] = () ) -> None ``` Cast model parameters to the target dtype, keeping fp32 modules in full precision. Respects `_keep_in_fp32_modules` / `_keep_in_fp32_modules_strict` on the model (the same attributes HuggingFace transformers uses). Uses `nn.Module.to()` which is safe for both plain tensors and DTensors (FSDP2 sharded parameters). When the model is already FSDP2-sharded (parameters are DTensors), strict fp32 modules are restored to fp32 because they are expected to be isolated as uniform fp32 FSDP units. Non-strict fp32 hints only restore matching buffers, since their parameters may share an FSDP unit with lower-precision parameters. **Parameters:** **`model`** `nn.Module` The model whose parameters should be cast. --- **`dtype`** `torch.dtype` — default: torch.bfloat16 Target dtype (e.g. `torch.bfloat16`). --- **`skip_modules`** `tuple[str, ...]` — default: () Names of immediate submodules to leave entirely untouched (kept at their current dtype). Unlike the `_keep_in_fp32_modules` restore path, these are *detached* during the cast so `model.to()` never visits them — the only reliable way to preserve an fp32 parameter once it is FSDP2-sharded (post-shard `.data` reassignment does not stick). The caller must guarantee each skipped submodule is its own dtype-uniform FSDP group (e.g. Qwen3.5's `_fp32_params` holder, sharded separately in fp32), so leaving it fp32 cannot break FSDP's uniform-dtype rule. --- ```python nemo_automodel.components.models.common.utils.compute_lm_head_logits( lm_head: torch.nn.Module | None, hidden_states: torch.Tensor, logits_to_keep: int | torch.Tensor = 0, is_thd: bool = False, fp32_lm_head: bool = False, output_hidden_states: bool = False ) -> transformers.modeling_outputs.CausalLMOutputWithPast ``` Project hidden states through `lm_head` and wrap the result. Centralizes the lm\_head projection and output packaging shared by every custom `*ForCausalLM` / `*ForConditionalGeneration` `forward()`. The returned `CausalLMOutputWithPast` carries the projected `logits` and, when requested, the final `hidden_states`; callers that also need `loss`, `past_key_values`, etc. read `.logits` and build their own output. * `lm_head is None` (e.g. a non-final pipeline-parallel stage that does not own the head): `hidden_states` is passed through as `logits` so the next stage receives it. * `logits_to_keep == 0` (training default): every position is projected. The full range is deliberately *not* sliced, because `slice(0, None)` on a DTensor is unsupported (it raises on the `aten.alias` op under tensor parallel with sequence parallelism). * `logits_to_keep` as a positive int or a tensor of indices: only the requested positions are projected. Both 2D `[T, H]` (THD/packed) and 3D `[B, S, H]` (BSHD) hidden states are handled. * `is_thd`: THD/packed inputs yield 2D `[T, V]` logits; the leading batch dim is restored (`unsqueeze(0)` -> `[1, T, V]`) so downstream code sees a uniform `[B, S, V]` layout. The same restoration is applied to the `hidden_states` field. Only applied while the tensor is still 2D, so an `inputs_embeds` path that already produced `[1, T, *]` is left untouched. * `fp32_lm_head`: run the projection in fp32 and cast the logits back to the input dtype. Used by models whose `lm_head.weight` has been promoted to fp32 (e.g. via the MoE `lm_head_precision` setting). The matmul goes through `lm_head` (`nn.Linear`, DTensor-aware under FSDP2) rather than `F.linear` so DTensor redistribution is preserved. * Otherwise, cast the sliced input to the head's effective compute dtype. For a standalone FSDP2 head this comes from its mixed-precision policy, since its sharded weight may still expose the master-copy dtype before the forward pre-hook runs. Other heads use their materialized weight dtype. * `output_hidden_states`: when set, the (full-sequence, THD-restored) `hidden_states` are attached to the output so the fused cross-entropy path can recompute logits over every position; otherwise the field is `None`. **Parameters:** **`lm_head`** `nn.Module | None` The language-model head module, or `None` on a pipeline stage that does not own it. --- **`hidden_states`** `torch.Tensor` Final hidden states, shaped `[T, H]` or `[B, S, H]`. --- **`logits_to_keep`** `int | torch.Tensor` — default: 0 `0` to project every position; a positive int to keep the last `N` positions; or a tensor of position indices. --- **`is_thd`** `bool` — default: False 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. --- **`fp32_lm_head`** `bool` — default: False Project in fp32 and cast the result back to the input dtype. Ignored when `lm_head` is `None`. --- **`output_hidden_states`** `bool` — default: False Attach the final hidden states to the output. --- **Returns:** `CausalLMOutputWithPast` A `CausalLMOutputWithPast` whose `logits` are the projected logits ```python nemo_automodel.components.models.common.utils.generation_config_from_model_config( config: transformers.PretrainedConfig ) -> transformers.generation.GenerationConfig ``` Generation config a custom causal LM starts with, seeded the way transformers does it. `PreTrainedModel.__init__` derives `generation_config` from the model config so the bos/eos/pad ids are set from the start. A custom model that builds a bare `GenerationConfig()` instead has no stop token, and because the consolidated export writes `model.generation_config` out as `generation_config.json` (which beats `config.json` on reload), the exported model never stops generating either. Objects that are not `PretrainedConfig` instances (test doubles) get the defaults. ```python nemo_automodel.components.models.common.utils.get_is_first_microbatch() -> bool | None ``` Get the global IS\_FIRST\_MICROBATCH flag. **Returns:** `bool | None` True/False/None indicating microbatch position for FP8 weight caching. ```python nemo_automodel.components.models.common.utils.get_is_optim_step() -> bool ``` Get the global IS\_OPTIM\_STEP flag. **Returns:** `bool` Whether we are in an optimization step. ```python nemo_automodel.components.models.common.utils.get_rope_config( config ) -> tuple[float, dict, float] ``` Extract rope configuration from `config.rope_parameters`. **Parameters:** **`config`** A HuggingFace model config object. --- **Returns:** `tuple[float, dict, float]` Tuple of (rope\_theta, rope\_parameters, partial\_rotary\_factor). ```python nemo_automodel.components.models.common.utils.initialize_linear_module( linear_impl: str, in_features: int, out_features: int, bias: bool = False, device: torch.device | str | None = None, dtype: torch.dtype = torch.bfloat16 ) -> torch.nn.Module ``` Initialize Linear module with the specified backend. For TE backend, creates TE module directly on specified device. Call reset\_parameters() to materialize weights if created on meta device. **Parameters:** **`linear_impl`** `str` Backend implementation ("te", "torch", or "quack") --- **`in_features`** `int` Input features --- **`out_features`** `int` Output features --- **`bias`** `bool` — default: False Whether to use bias --- **`device`** `torch.device | str | None` — default: None Device to create module on (None uses PyTorch default, typically CPU) --- **`dtype`** `torch.dtype` — default: torch.bfloat16 Parameter dtype --- **Returns:** `nn.Module` Linear module ```python nemo_automodel.components.models.common.utils.initialize_rms_norm_module( rms_norm_impl: str, dim: int, eps: float = 1e-05, device: torch.device | str | None = None, dtype: torch.dtype = torch.bfloat16 ) -> torch.nn.Module ``` Initialize RMSNorm module with the specified backend. For TE backend, creates TE module directly on specified device. Call reset\_parameters() to materialize weights if created on meta device. **Parameters:** **`rms_norm_impl`** `str` 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 --- **`dim`** `int` Normalized dimension --- **`eps`** `float` — default: 1e-05 Epsilon for numerical stability --- **`device`** `torch.device | str | None` — default: None Device to create module on (None uses PyTorch default, typically CPU) --- **`dtype`** `torch.dtype` — default: torch.bfloat16 Parameter dtype --- **Returns:** `nn.Module` RMSNorm module ```python nemo_automodel.components.models.common.utils.is_tensor_unallocated( tensor: torch.Tensor ) -> bool ``` Check if tensor is unallocated (meta tensor, fake tensor, etc.). TE kernels don't support meta tensors, fake tensors, or unallocated tensors. This helper detects such cases for fallback handling. **Parameters:** **`tensor`** `torch.Tensor` Tensor to check --- **Returns:** `bool` True if tensor is unallocated or cannot be accessed ```python nemo_automodel.components.models.common.utils.load_pretrained_generation_config( pretrained_model_name_or_path: str | os.PathLike[str], config: transformers.PretrainedConfig | None = None, config_overrides: collections.abc.Collection[str] | None = None, loading_kwargs: typing.Any = {} ) -> transformers.generation.GenerationConfig | None ``` The generation config a checkpoint carries, or `None` when it carries none. Mirrors what `PreTrainedModel.from_pretrained` does once the weights are in: `generation_config.json` wins when present, because that is where hub checkpoints keep extra stop tokens (Nemotron-3 lists `[2, 11]` there but only `2` in `config.json`) and their sampling defaults. Without it, the generation fields of the raw `config.json` are used, as HF does for legacy checkpoints that still keep `do_sample`/`temperature` there (the in-memory config has already dropped them). **Parameters:** **`pretrained_model_name_or_path`** `str | os.PathLike[str]` Hub id or local directory of the checkpoint. --- **`config`** `PretrainedConfig | None` — default: None The model's in-memory config. Its explicit values (an `eos_token_id` passed to `from_pretrained` lands here) win over the `config.json` fallback, which cannot know about them. --- **`config_overrides`** `Collection[str] | None` — default: None Names the caller explicitly set on *config*. Needed because a value cannot be recognized as deliberate once it equals a default. --- **`**loading_kwargs`** `Any` — default: \{} The `HUB_LOADING_KWARGS` the caller was loading with, so the generation config is read from the same place as the weights. --- ```python nemo_automodel.components.models.common.utils.restore_pretrained_generation_config( model: torch.nn.Module, pretrained_model_name_or_path: str | os.PathLike[str], config_overrides: collections.abc.Collection[str] | None = None, loading_kwargs: typing.Any = {} ) -> None ``` Replace a model's `generation_config` with the one its checkpoint carries, if any. Applies only to models that expose a real `GenerationConfig` (the ones that can generate); a model without one, or a checkpoint without generation settings, is left as it is. `config_overrides` names what the caller explicitly set on the model's config, and `loading_kwargs` are the `HUB_LOADING_KWARGS` the model was loaded with. ```python nemo_automodel.components.models.common.utils.set_is_first_microbatch( value: bool | None ) -> None ``` Set the global IS\_FIRST\_MICROBATCH flag for FP8 weight caching. **Parameters:** **`value`** `bool | None` True for first microbatch (quantize+cache), False for subsequent (use cached), None to disable caching. --- ```python nemo_automodel.components.models.common.utils.set_is_optim_step( value: bool ) -> None ``` Set the global IS\_OPTIM\_STEP flag. **Parameters:** **`value`** `bool` Whether we are in an optimization step. --- ```python nemo_automodel.components.models.common.utils.yield_fp32_model( model: torch.nn.Module, restore_dtype: torch.dtype | None = None ) ``` Run a block with the model temporarily in fp32, then cast it to `restore_dtype`. On entry the whole model is cast to fp32; on exit it is cast to `restore_dtype` (which defaults to the model's pre-context floating-point dtype, so by default the original dtype is restored). The exit cast is a no-op when the target is already fp32. The motivating use is from-scratch weight initialization. Sampling a random init directly in a reduced-precision dtype (e.g. bf16) distorts the init's variance/mean schedule: bf16's 8-bit mantissa quantizes the small init magnitudes and biases the truncation/scaling arithmetic used by `normal_` / `trunc_normal_`. In a deep residual stack this compounds and produces genuinely huge gradients on the first optimization steps of from-scratch pretraining (flat / diverging loss). Sampling in fp32 and then casting back avoids this while keeping reduced-precision storage: the round-to-bf16 of a correct fp32 sample is an unbiased per-element perturbation that preserves the init statistics. Wrap the body of a model's `initialize_weights` to keep that round-trip in one place. Works whether or not the model is already FSDP2-sharded: both casts are *uniform* whole-model casts, so FSDP2's invariant that every parameter in a group shares one dtype is preserved. In the AutoModel pipeline `initialize_weights` actually runs after sharding (via `checkpointer.initialize_model_weights`), i.e. on DTensor params, which is supported. `_keep_in_fp32_modules` / `_keep_in_fp32_modules_strict` handling is delegated to `cast_model_to_dtype`: on an unsharded model those modules' params and buffers are restored to fp32 on exit; on a sharded model, strict fp32 modules are restored while non-strict modules only have their buffers restored. **Example** \>>> with yield\_fp32\_model(self, dtype): ... self.model.init\_weights(buffer\_device=buffer\_device) **Parameters:** **`model`** `nn.Module` The model to run in fp32 within the context. --- **`restore_dtype`** `torch.dtype | None` — default: None The dtype to cast the model to on exit. Defaults to the model's current floating-point dtype (captured before the fp32 cast), i.e. the original dtype. --- ```python nemo_automodel.components.models.common.utils.AttentionBackend = Literal['torch', 'te', 'sdpa', 'flex', 'eager', 'tilelang', 'cudnn', 'fa4', 'mag... ``` ```python nemo_automodel.components.models.common.utils.HAVE_DEEP_EP = importlib.util.find_spec('deep_ep') is not None ``` ```python nemo_automodel.components.models.common.utils.HAVE_TE = importlib.util.find_spec('transformer_engine') is not None ``` ```python nemo_automodel.components.models.common.utils.HAVE_UCCL_EP = importlib.util.find_spec('uccl') is not None or importlib.util.find_spec('ep') i... ``` ```python nemo_automodel.components.models.common.utils.HUB_LOADING_KWARGS = ('cache_dir', 'force_download', 'proxies', 'local_files_only', 'token', 'revisio... ``` ```python nemo_automodel.components.models.common.utils.IS_FIRST_MICROBATCH: bool | None = None ``` ```python nemo_automodel.components.models.common.utils.IS_OPTIM_STEP = False ``` ```python nemo_automodel.components.models.common.utils._patch_te_modules = _make_lazy_te_patcher() ``` ```python nemo_automodel.components.models.common.utils.logger = logging.getLogger(__name__) ```