transformer_config.py

# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

import logging
import math
import os
import warnings
from copy import deepcopy
from dataclasses import dataclass, field
from typing import Callable, List, Literal, Optional, Self, Tuple, Union

import torch
import torch.nn.functional as F

from megatron.core._rank_utils import warn_single_rank
from megatron.core.activations import squared_relu
from megatron.core.context_parallel import CPLayout
from megatron.core.enums import Fp4Recipe, Fp8Recipe
from megatron.core.inference.moe import InferenceGroupedGemmBackend
from megatron.core.quantization.quant_config import RecipeConfig
from megatron.core.transformer.cuda_graph_config import (
    ALLOWED_INFERENCE_SCOPES,
    get_deprecated_cuda_graph_modules_migration,
    normalize_cuda_graph_modules,
    normalize_inference_cuda_graph_scope,
    validate_deprecated_cuda_graph_modules_migration_inputs,
)
from megatron.core.transformer.enums import (
    AttnBackend,
    CudaGraphModule,
    CudaGraphScope,
    InferenceCudaGraphScope,
    LayerType,
)
from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
from megatron.core.transformer.wide_residual_config import WideResidualConfig

from .._rank_utils import log_single_rank
from ..fusions.fused_bias_geglu import quick_gelu
from ..model_parallel_config import ModelParallelConfig
from ..utils import (
    _validate_dsa_kernel_backend_dependencies,
    get_te_version,
    init_method_normal,
    is_te_min_version,
    is_torch_min_version,
    mup_scaled_init_method_normal,
    scaled_init_method_normal,
)

logger = logging.getLogger(__name__)

try:
    from packaging.version import Version as PkgVersion

    HAVE_PACKAGING = True
except ImportError:
    HAVE_PACKAGING = False


@dataclass
class TransformerConfig(ModelParallelConfig):
    """Configuration object for megatron-core transformers.

    The initialization function has an argument for each parameter,
    including those in ModelParallelConfig.
    """

    ####################
    # model architecture
    ####################

    num_layers: int = field(default=0, metadata={"argparse_meta": {"default": None}})
    """Number of transformer layers in a transformer block."""

    mtp_num_layers: Optional[int] = None
    """Number of Multi-Token Prediction (MTP) Layers.
    MTP extends the prediction scope to multiple future tokens at each position.
    This MTP implementation sequentially predict additional tokens
    by using D sequential modules to predict D additional tokens.
    """

    mtp_loss_scaling_factor: Optional[float] = 0.1
    """Weighting factor of Multi-Token Prediction (MTP) loss.
    We compute the average of the MTP losses across all depths, 
    and multiply it the scaling factor to obtain the overall MTP loss, 
    which serves as an additional training objective.
    """

    mtp_use_repeated_layer: bool = False
    """Use a single MTP layer repeatedly instead of multiple separate layers."""

    freeze_base_model_for_mtp: bool = False
    """Freeze every non-MTP parameter and avoid recording backbone activations."""

    keep_mtp_in_bf16: bool = False
    """Keep MTP layers out of FP8 and FP4 quantization contexts."""

    mtp_detach_heads: bool = False
    """If True, detach MTP head inputs from the main model graph.
    This prevents MTP loss gradients from flowing back to the main model,
    only training the MTP heads themselves."""

    mtp_hsm: bool = False
    """Enable uniform per-token Hidden State Mixing (HSM) for MTP layers.
    At every MTP depth, each token independently draws its input from the main model
    hidden state and the outputs of the earlier depths, all aligned on the same target
    token. Only takes effect during training and requires at least two MTP layers,
    since a single depth has nothing to mix. Model constructors validate this
    against the resolved architecture."""

    mtp_hybrid_override_pattern: Optional[str] = None
    """DEPRECATED: Use unified hybrid_layer_pattern instead.
    Legacy argument for loading old checkpoints.
    Force a specific hybrid layer pattern for MTP layers.
    """

    num_layers_in_first_pipeline_stage: Optional[int] = None
    """Number of transformer layers on first pipeline stage.
    None implies equal layer division across PP ranks."""

    num_layers_in_last_pipeline_stage: Optional[int] = None
    """Number of transformer layers on last pipeline stage.
    None implies equal layer division across PP ranks."""

    pipeline_model_parallel_layout: Optional[Union[str, list, PipelineParallelLayerLayout]] = None
    """Custom definition of the pipeline parallel partitioning.
    Support type:
    - str: e.g., 'Et*3|(tt|)*29,m|L'. Stages are split by '|', replicated stages or layers
    can be described with multiplication. Commas can be used cosmetically.
    - list: e.g., [['embedding', 'decoder'], ['decoder', 'decoder', 'decoder', 'loss']].
    - PipelineParallelLayerLayout: a PipelineParallelLayerLayout object.
    If given either a string or a list, it will be transferred into a PipelineParallelLayerLayout
    in post init. Let i = a * pp_size + b, then layout[i] gives a list of the layers 
    in the a-th vpp stage and the b-th pp stage, i.e., vpp(0)pp(0), vpp(0)pp(1), ..., 
    vpp(i)pp(j), vpp(i)pp(j+1), ..., vpp(-1)pp(-2), vpp(-1)pp(-1).
    In the inner lists of layers, 'embedding' or 'E' denotes the embedding layer, 'loss' or 'L'
    denotes the loss function, and 'decoder' or 't' denotes the transformer decoder layer.
    Examples:
        [['embedding', 'decoder'], ['decoder', 'decoder', 'decoder', 'loss']]:
        pp = 2, vpp = None
        pp rank 0 holds: embedding, decoder
        pp rank 1 holds: decoder*3, loss
        'E|(tt|)*2,(t|)*4,mL':
        pp = 2, vpp = 4
        vpp rank 0 pp rank 0 holds: embedding
        vpp rank 0 pp rank 1~2 holds: decoder*2
        vpp rank 0 pp rank 3 holds: decoder
        vpp rank 1 pp rank 0~2 holds: decoder
        vpp rank 1 pp rank 3 holds: mtp, loss"""

    account_for_embedding_in_pipeline_split: bool = False
    """If set, the embedding layer will be treated as a standard transformer
    layer in the context of partition and placement for pipeline parallelism."""

    account_for_loss_in_pipeline_split: bool = False
    """If set, the loss layer will be treated as a standard transformer
    layer in the context of partition and placement for pipeline parallelism."""

    hidden_size: int = field(default=0, metadata={"argparse_meta": {"default": None}})
    """Transformer hidden size."""

    num_attention_heads: int = field(default=0, metadata={"argparse_meta": {"default": None}})
    """Number of transformer attention heads."""

    attention_backend: AttnBackend = AttnBackend.auto
    """Attention backend to run. By default we let transformer engine
    decide the best backend to run (except in the case of local).
    If attention backend is local we use the local pytorch implementation in mcore.
    Users can specify exact backend by changing this config. """

    flash_attention_version: Optional[Literal[2, 3, 4]] = None
    """Pin the FlashAttention generation (2, 3, or 4) used by both the training
    (TransformerEngine) and inference (mcore dynamic-batching) attention paths. When
    None, each path selects a version automatically based on what is installed. Pinning
    is required for batch-invariant mode: the training-side logprob recompute and the
    inference engine must run the same kernel, since different FlashAttention
    generations use different tile sizes and softmax accumulation orders and therefore
    differ bitwise. On the training side this is enforced via TransformerEngine's
    NVTE_FLASH_ATTN_V2/V3/V4 selection environment variables."""

    softmax_scale: Optional[float] = None
    """Softmax scale for attention scaling."""

    softmax_type: Literal['vanilla', 'off-by-one', 'learnable'] = 'vanilla'
    """Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. 
       Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and 
       and learnable offset."""

    attn_logit_softcapping: Optional[float] = None
    """If not None, cap the attention logits at this value using cap * tanh(logits / cap) before
    softmax. Must be positive; use None to disable softcapping. Note that TransformerEngine
    spells the disabled state as 0.0 rather than None, so 0.0 is rejected here to keep the two
    from meaning different things."""

    num_query_groups: Optional[int] = field(
        default=None, metadata={"argparse_meta": {"default": 1}}
    )
    """Number of query groups for group query attention. If None, normal attention is used."""

    ffn_hidden_size: Optional[int] = None
    """Transformer Feed-Forward Network hidden size. This is set to 4*hidden_size
    if not provided."""

    kv_channels: Optional[int] = None
    """Projection weights dimension in multi-head attention. This is set to hidden_size //
    num_attention_heads if not provided."""

    hidden_dropout: float = 0.1
    """Dropout probability for transformer hidden state."""

    attention_dropout: float = 0.1
    """Post attention dropout probability."""

    fp32_residual_connection: bool = False
    """If true, move residual connections to fp32."""

    # @jcasper should we keep this option?
    apply_residual_connection_post_layernorm: bool = False
    """If True, uses the original BERT residule connection ordering."""

    layernorm_epsilon: float = field(
        default=1e-5, metadata={"argparse_meta": {"arg_names": ["--norm-epsilon"]}}
    )
    """Epsilon value for any LayerNorm/RMSNorm operations."""

    layernorm_zero_centered_gamma: bool = field(
        default=False, metadata={"argparse_meta": {"arg_names": ["--apply-layernorm-1p"]}}
    )
    """If set to True, the LayerNorm is adjusted to center the gamma values around 0. This improves
    numerical stability."""

    add_bias_linear: bool = field(
        default=True, metadata={"argparse_meta": {"arg_names": ["--disable-bias-linear"]}}
    )
    """Include/exclude a bias term in all linear layers (QKV projections, after core attention,
    and two in MLP layer)."""

    add_qkv_bias: bool = False
    """Add a bias term only for QKV projections."""

    gated_linear_unit: bool = False
    """Use a gated linear unit for the first linear layer in the MLP."""

    activation_func: Callable[[torch.Tensor], torch.Tensor] = F.gelu
    """Activation function to use for the non-linearity in the MLP."""

    activation_func_fp8_input_store: bool = False
    """Store the input of MLP activation function in FP8 for backprop to save memory.
    The stored input is casted back to the original precision before backprop compuatation."""

    glu_linear_offset: float = 0.0
    """Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). Only 
    used when gated_linear_unit is True"""

    activation_func_clamp_value: Optional[float] = None
    """Clamp the output of the linear_fc1 in the activation function. Only used when activation_func
    is quick_gelu or SwiGLU (MoE only)."""

    activation_func_tanh_clamp_scale: Optional[float] = None
    """If set, precondition the input of the activation function with `s * tanh(x / s)`, where `s`
    is this value. For a gated activation (silu only) this instead selects SiTU-GLU. The fused MoE
    path (use_transformer_engine_op_fuser) requires a Transformer Engine with ScaledTanhSReLU."""

    activation_func_tanh_clamp_scale_linear: Optional[float] = None
    """Soft clamp scale for the linear (up) half of a gated activation, decoupled from the gate
    scale in activation_func_tanh_clamp_scale. Requires activation_func_tanh_clamp_scale."""

    num_moe_experts: Optional[int] = None
    """Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None
    for no MoE."""

    rotary_interleaved: bool = False
    """True is rotate pairs of even and odd dimensions (RoFormer style), False is rotate pairs of
    first half and second half (LLaMa style). Default to False."""

    window_size: Optional[Tuple[int, int]] = None
    """If not None, then will use sliding window attention. The size of the window is specified by
    the numbers inside the tuple; -1 is special value meaning "infinite window size"."""

    window_attn_skip_freq: Optional[Union[int, List[int]]] = None
    """Frequency of full attention layers among sliding window attention layers. Accepts either:
    - An integer N: Represents a (N-1):1 ratio, one full attention layer after (N-1) SWA layers.
    - A list that defines a custom pattern, e.g.: [1,1,1,1,0,0,0,0], where 1 represents SWA. """

    normalization: Literal['LayerNorm', 'RMSNorm'] = "LayerNorm"
    """Which norm to use for normalization layers, valid options are `LayerNorm` and `RMSNorm`."""

    qk_layernorm: bool = False
    """Whether to apply `normalization` type of normalization to the query and key embeddings."""

    qk_l2_norm: bool = False
    """Whether to apply llama 4-style qk L2 norm."""

    qk_clip: bool = False
    """Whether to clip the query and key weights. Needed for Muon MLA Model training."""

    qk_clip_alpha: float = 0.5
    """The balancing alpha for qk-clip. Q = Q * (eta ** alpha)"""

    qk_clip_threshold: float = 100
    """The balancing threshold for qk-clip. eta = min(threshold / max_attention_logits, 1.0)"""

    log_max_attention_logit: bool = False
    """Whether to log the max attention logit across whole model. Decoupled from qk_clip,
    defualts to False. Setting qk_clip will automatically log the max logit"""

    attention_output_gate: bool = False
    """Whether to apply output gate to the attention layers."""

    test_mode: bool = False
    """Whether to run real-time tests."""

    calculate_per_token_loss: bool = False
    """Whether cross entropy loss is calculated over the actual number of non-padded tokens in the
    global batch, versus the default behavior of assuming all tokens are non-padded."""

    multi_latent_attention: bool = False
    """Whether to use multi-latent attention."""

    no_rope_freq: Optional[Union[int, List[int]]] = None
    """Controls which layers perform Rotary Position Embedding (RoPE). Accepts either:
    An integer N: Creates a pattern where RoPE is skipped every N-1 layers. For example,
    no_rope=4 means RoPE is applied for 3 layers, then skipped for 1 layer, repeating this pattern.
    A list of integers: Defines a custom pattern where 1 means skip RoPE and 0 means apply RoPE.
    For example, [0,1,1,0] means: apply RoPE, skip RoPE, skip RoPE, apply RoPE."""

    ####################
    # attention variant
    ####################
    experimental_attention_variant: Optional[
        Literal['gdn', 'gdn2', 'dsa', 'dsv4_hybrid', 'gated_delta_net']
    ] = None
    """Type of attention variant to use. Supports gdn, gdn2, dsa, and dsv4_hybrid.
    gdn2 selects the GDN2 (Gated DeltaNet-2) variant of the gated delta net layer, with
    channel-wise decay, erase and write gates; it requires flash-linear-attention >= 0.5.1.
    Both gdn and gdn2 also select the layer built for the hybrid layer pattern symbol 'G'.
    'gated_delta_net' is a deprecated alias of 'gdn': it is normalized to 'gdn' in
    __post_init__ and emits a DeprecationWarning."""

    experimental_attention_variant_loss_scale_func: Optional[Callable[[torch.Tensor], None]] = None
    """Optional hook for experimental attention variants to receive the main loss scale."""

    ####################
    # DSA
    ####################
    dsa_indexer_n_heads: Optional[int] = None
    """Number of DSA indexer heads."""

    dsa_indexer_head_dim: Optional[int] = None
    """Dimension per DSA indexer head."""

    dsa_indexer_topk: Optional[int] = None
    """Number of top-k tokens to select in DSA indexer."""

    dsa_indexer_topk_freq: int = 1
    """Frequency of DSA indexer top-k computation across layers.
    A value greater than 1 enables cross-layer top-k sharing."""

    dsa_indexer_skip_topk_offset: int = 0
    """Layer offset for DSA cross-layer top-k sharing."""

    dsa_indexer_loss_coeff: Optional[float] = None
    """Coefficient for the DSA indexer KL divergence loss. Set to 0 to disable indexer loss."""

    dsa_indexer_use_sparse_loss: bool = False
    """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the
    top-k indices."""

    dsa_kernel_backend: Literal["none", "tilelang", "cudnn"] | None = None
    """Optional fused DSA kernel backend.
    When unset, DSv4 hybrid uses ``cudnn`` and other attention variants use ``none``.
    ``none`` disables fused DSA kernels. Explicit ``tilelang`` or ``cudnn`` enables only that
    backend. Unsupported DSA layouts continue to use the PyTorch fallback."""

    dsa_indexer_rope_interleaved: bool = False
    """Whether DSA indexer RoPE should use MLA-style interleaving."""

    dsa_indexer_rotate_activation: bool = True
    """Whether DSA indexer should apply Hadamard rotate_activation to q/k before scoring."""

    dsa_indexer_scoring_relu: bool = True
    """Whether DSA indexer should apply ReLU to q@k^T scores before weighting."""

    dsa_indexer_k_norm_epsilon: Optional[float] = None
    """Optional epsilon override for the DSA indexer key LayerNorm."""

    dsa_indexer_k_norm_fp32: bool = False
    """Whether DSA indexer key LayerNorm should run on fp32 inputs."""

    ####################
    # Compressed sparse attention
    ####################
    csa_window_size: int = 128
    """Sliding window size for compressed sparse attention."""

    csa_compress_ratios: Optional[List[int]] = None
    """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...].
    The decoder occupies the first ``num_layers`` entries. MTP attention layers use the tail;
    HybridModel patterns with multiple inner layers therefore require one entry per inner layer."""

    csa_compress_rotary_base: float = 40000.0
    """RoPE base for compressed KV positions in compressed sparse attention."""

    csa_dense_mode: bool = False
    """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be
    disabled."""

    ####################
    # linear attention
    ####################
    linear_attention_freq: Optional[Union[int, List[int]]] = None
    """Frequency between LA (linear attention) layers 
    and SDPA (scaled dot-product attention) layers.
    Accepts either:
    - An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer
    - A list that defines a custom pattern, e.g.: [1,1,1,0,1,1,1,0,1,1,1,0]"""

    linear_conv_kernel_dim: Optional[int] = 4
    """Conv kernel dimension for the gated delta net."""

    linear_key_head_dim: Optional[int] = 128
    """Query and key head dimension for the gated delta net."""

    linear_value_head_dim: Optional[int] = 128
    """Value and gate head dimension for the gated delta net."""

    linear_num_key_heads: Optional[int] = 16
    """Number of query and key heads for the gated delta net."""

    linear_num_value_heads: Optional[int] = 32
    """Number of value and gate heads for the gated delta net."""

    gdn_pre_gated_delta_rule_fusion: bool = False
    """Whether to use the streamed Triton fusion for GatedDeltaNet pre-GDR preprocessing."""

    gdn_gated_output_norm_fusion: bool = False
    """Fuse GatedDeltaNet output RMSNorm and SiLU gating. Unsupported configurations and
    layouts raise on every forward; see docs/developer/gdn_ew_fusion.md for requirements."""

    ####################
    # initialization
    ####################
    init_method: Optional[Callable] = None
    """Method to initialize weights. Note that bias is always set to zero. Should be a function that
    takes a single Tensor and initializes it. If None, will be set to
    megatron.core.utils.init_method_normal(init_method_std) which is torch nn init normal with
    mean=0.0 and std=init_method_std."""

    output_layer_init_method: Optional[Callable] = None
    """Method to initialize weights of the output layer of both attention and MLP blocks. If None,
    will be set to megatron.core.utils.scaled_init_method_normal(init_method_std) which is torch nn
    init normal with mean=0.0 and std=init_method_std / math.sqrt(2.0 * num_layers).
    Note: this does not control vocab readout/unembedding initialization."""

    init_method_std: float = 0.02
    """Standard deviation of the zero mean normal for the default initialization method, not used if
    init_method and output_layer_init_method are provided."""

    embedding_init_method: Optional[Callable] = None
    """
    Method to initialize weights of the embedding layer. If None, will be set as described 
    in init_method above.
    """

    embedding_init_method_std: Optional[float] = None
    """
    Standard deviation of the zero mean normal for the default initialization method for the 
    embedding layer. If None, will be set to init_method_std. Setting this to a value around
    1.0 may avoid loss spikes in training. Setting this to any value will also skip applying
    weight decay on embedding weights to avoid shrinkage towards zero.
    See https://arxiv.org/abs/2312.16903 for more details.
    """

    init_model_with_meta_device: bool = False
    """
    If True, initializes the model with the meta device. This is helpful for
    training of very large models. This feature is only works when megatron fsdp is turned on.
    """

    ####################
    # MuP (Maximal Update Parameterization)
    ####################
    use_mup: bool = False
    """
    Enable Maximal Update Parameterization (MuP) for hyperparameter transfer across
    model widths. When enabled, learning rates and initialization are scaled according
    to the width multiplier to ensure consistent training dynamics.
    """

    mup_width_mult: float = 1.0
    """
    Width multiplier for MuP scaling, computed as hidden_size / mup_base_hidden_size.
    This value is automatically computed in __post_init__ when use_mup is enabled.
    """

    mup_base_hidden_size: Optional[int] = None
    """
    Base hidden size for MuP width scaling. This is the reference width from which
    scaling factors are computed. Defaults to hidden_size if not specified (base model
    case where width_mult=1.0). Set this to your base/proxy model's hidden size when
    scaling up.
    """

    mup_embedding_mult: float = 1.0
    """
    Multiplier for embedding layer output. Applied after the embedding lookup.
    Default: 1.0 (no scaling).
    """

    mup_output_mult: float = 1.0
    """
    Multiplier for output logits before softmax. When MuP is enabled and this is left
    at 1.0, it is auto-set to 1/mup_width_mult to keep output variance stable across
    widths. Override to customize output scaling.
    Default: 1.0.
    """

    mup_base_head_dim: Optional[float] = None
    """
    Base head dimension for MuP attention scaling. When set,
    softmax_scale = sqrt(mup_base_head_dim) / (kv_channels ** mup_attn_scale_power).
    Set to base model's d_head (e.g., 64) to match standard 1/sqrt(d_head) scaling
    at the base width, ensuring non-MuP compatibility for that specific value.
    """

    mup_attn_scale_power: float = 1.0
    """
    Power for attention scaling: softmax_scale = 1 / (kv_channels ** mup_attn_scale_power).
    0.5 = standard attention (1/sqrt(d_head)), 1.0 = MuP attention (1/d_head).
    Default: 1.0 (MuP scaling when use_mup is True). Set to 0.5 for standard scaling.
    """

    ####################
    # mixed-precision
    ####################
    apply_query_key_layer_scaling: bool = False
    """If true, scale Q * K^T by 1 / layer-number. This improve numeric stability when training with
    fp16. Also sets `attention_softmax_in_fp32` to True."""

    attention_softmax_in_fp32: bool = True
    """If True, run attention masking and softmax in fp32. This should be True if
    apply_query_key_layer_scaling is True."""

    disable_bf16_reduced_precision_matmul: bool = False
    """If True, sets torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction=False to
    prevent matmul from using reduced precision accumulation when using BF16."""

    ####################
    # fusion
    ####################
    bias_activation_fusion: bool = False
    """If True, fuses bias addition and the activation function when possible."""

    masked_softmax_fusion: bool = False
    """If True, uses softmax fusion."""

    persist_layer_norm: bool = False
    """If True, uses the persistent fused layer norm kernel. This kernel only supports a fixed set
    of hidden sizes."""

    memory_efficient_layer_norm: bool = False
    """If True, and using local layers (not from TransformerEngine), tells Apex to use the memory
    efficient fused LayerNorm kernel. Ignored if not using LayerNorm."""

    bias_dropout_fusion: bool = False  # TODO: this should be bias_dropout_add_fusion?
    """If True, uses bias dropout fusion."""

    apply_rope_fusion: bool = False
    """If True, use fused RoPE kernel."""

    use_fused_weighted_squared_relu: bool = False
    """If True, uses the fused squared relu kernel: for MoE experts, the per-token
    weighted variant; for the dense MLP, the tanh soft-clamped variant when
    activation_func_tanh_clamp_scale is set."""

    fused_single_qkv_rope: bool = False
    """If set, avoid splitting QKV before ROPE forward and avoid concatenating ROPE dgrads."""

    fused_residual_rmsnorm: bool = False
    """If True, fuses residual connection and RMSNorm backward pass when TE is used."""

    use_transformer_engine_op_fuser: bool = False
    """If True, submodules may use Transformer Engine's operation fuser
    API to enable advanced fusions."""

    ####################
    # activation recomputation
    ####################
    recompute_granularity: Optional[Literal['full', 'selective']] = None
    """Determines which type of activation recompute to use.  Megatron-core supports 'selective'
    activation checkpointing where the submodules set in --recompute-modules is checkpointed.
    The default is "core_attn" which is the memory intensive part of attention.
    These memory intensive activations are also less compute intensive which makes activation
    checkpointing more efficient for LLMs (20B+).  See Reducing Activation Recomputation in Large
    Transformer Models (https://arxiv.org/abs/2205.05198) for more details.  'full' will checkpoint
    the entire transformer layer.  If None, no recompute is performed and all activations are saved.
    If set, must be 'selective' or 'full'. 'selective' always uses all layers.
    """

    recompute_method: Optional[Literal['uniform', 'block']] = None
    """Determines which transformer layers will be recomputed. uniform will uniformly divide the
    total number of transformer layers in a transformer block and recompute the input activation of
    each divided chunk at the specified granularity.  block will recompute the input activations for
    only a set number of transformer layers per pipeline stage.  The rest of the layers in the
    pipeline stage will not have any activations recomputed.  If None, and recompute is enabled, all
    layers will do recomputation. If set, must be 'uniform' or 'block'."""

    recompute_num_layers: Optional[int] = None
    """When recompute_method is uniform, recompute_num_layers is the number of transformer layers in
    each uniformly divided recompute unit.  When recompute_method is block, recompute_num_layers is
    the number of transformer layers to recompute within each pipeline stage.  Must be None for
    'selective' activation checkpointing."""

    distribute_saved_activations: Optional[bool] = False
    """If True, distribute recomputed activations across the model parallel group."""

    recompute_modules: Optional[List[str]] = None
    """The submodules to recompute.
    choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe",
    "shared_experts", "gdn_norm_out", "gdp_in_proj", "gdp_qkv", "mhc",
    "shortcut_pre_mlp_layernorm", "residual_stream".
    default: ["core_attn"].
    "core_attn": recompute the core attention part of the transformer layer.
    "moe_act": recompute the MoE MLP activation function.
    "layernorm": recompute the input_layernorm and pre_mlp_layernorm.
    "mla_up_proj": recompute the MLA up projection and RoPE applying parts.
    "mlp": recompute the dense MLP submodule.
    "moe": recompute the MoE layer.
    "shared_experts": recompute the shared experts in the MoE layer.
    "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all.
    "gdp_in_proj": recompute the GatedDeltaProduct input projection and its CP gather/split
    preprocessing.
    "gdp_qkv": recompute the GatedDeltaProduct causal conv and QKV preparation.
    "mhc": recompute HyperConnection intermediate activations via
            CheckpointWithoutOutput + CheckpointWithoutOutputManager. Requires
            enable_mhc_connections=True. Cannot be used with "mlp".
    "shortcut_pre_mlp_layernorm": recompute the shortcut router's input normalization.
            Requires moe_shortcut_connection=True and selective recomputation.
    "residual_stream": replay wide-residual reads, connected norms, and writes via
            CheckpointWithoutOutput + CheckpointWithoutOutputManager.
    "moe_act", "layernorm", "mla_up_proj", "gdn_norm_out", "gdp_in_proj", "gdp_qkv", "mhc",
    "shortcut_pre_mlp_layernorm", and "residual_stream" use output-discarding checkpointing;
    "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing.
    """

    ####################
    # fp8 related
    ####################
    fp8: Optional[Literal['e4m3', 'hybrid']] = field(
        default=None, metadata={"argparse_meta": {"arg_names": ["--fp8-format"]}}
    )
    """If set, enables the use of FP8 precision through Transformer Engine. There are 2 predefined
    choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses e4m3 for all FP8
    activation and weight tensors and e5m2 for all FP8 output activation gradient tensors."""

    fp8_recipe: Optional[Literal['tensorwise', 'delayed', 'mxfp8', 'blockwise', 'custom']] = (
        "delayed"
    )
    """If set, enables the use of FP8 precision through Transformer Engine. There are 5 predefined
    choices (1) 'tensorwise' uses per tensor current scaling recipe, (2) 'delayed'
    uses delayed scaling recipe, 3) 'mxfp8' for Blackwell architecture only,
    4) 'blockwise' for blockwise scaling recipe, 5) 'custom' for custom quantization recipe."""

    fp8_param: bool = False
    """If set, keep the parameters in fp8 precision to save memory. This option must be used
    together with fp8 mode (i.e., TransformerConfig.fp8 is not None). Note that not all parameters
    will be converted to fp8; for example, biases will remain unchanged. The parameters affected are
    primarily the weights of GEMMs. The specific parameters that will be converted to fp8 are
    determined by TE."""

    fp8_quantizer_factory: Optional[str] = None
    """Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory.
    Required when fp8_recipe is custom."""

    fp8_margin: int = 0
    """Margin for the scaling factor computation."""

    fp8_interval: int = 1
    """DEPRECATED from TransformerEngine v1.8.0. This flag is ignored.
    Controls how often the scaling factor is recomputed.
    """

    fp8_amax_history_len: int = 1
    """The length of the amax history window used for scaling factor computation."""

    fp8_amax_compute_algo: Literal['most_recent', 'max'] = "most_recent"
    """Algorithm used for choosing the `amax` value for the scaling factor computation. There are 2
    predefined choices: `max` chooses the largest `amax` in the history window, while `most_recent`
    always chooses the most recently seen value.

    """

    fp8_wgrad: bool = True
    """When set to False, override FP8 config options and do the wgrad computation
    in higher precision."""

    fp8_output_proj: bool = False
    """If True, run the LM-head output projection with a TE ColumnParallelLinear
    under the MXFP8 autocast context. Only active when fp8=True and
    fp8_recipe='mxfp8'."""

    fp8_dot_product_attention: bool = False
    """When set to True, use the FP8 implementation of Dot Product Attention."""

    fp8_multi_head_attention: bool = False
    """When set to True, use the FP8 implementation of Multi Head Attention."""

    tp_only_amax_red: bool = False
    """When set to True, reduce the FP8 AMAX only in the TP or TP-CP domain"""

    first_last_layers_bf16: bool = False
    """If True, retains first and last N TransformerBlocks in BF16 as opposed to FP8."""

    num_layers_at_start_in_bf16: int = 1
    """Number of layers at the start of the model to keep in BF16 precision when
    first_last_layers_bf16 is True."""

    num_layers_at_end_in_bf16: int = 1
    """Number of layers at the end of the model to keep in BF16 precision when
    first_last_layers_bf16 is True."""

    use_kitchen: bool = False
    """Use the kitchen extension for transformer quantization."""

    use_kitchen_attention: bool = False
    """Use the kitchen extension for attention (instead of TE's attention)."""

    kitchen_attention_backend: Literal["sdpa", "fa"] = "sdpa"
    """Which kitchen attention backend to use when use_kitchen_attention=True.
    "sdpa" for KitchenDotProductAttention, "fa" for KitchenFlashAttention."""

    ####################
    # fp4 related
    ####################
    fp4: Optional[Literal['e2m1']] = field(
        default=None, metadata={"argparse_meta": {"arg_names": ["--fp4-format"]}}
    )
    """If set, enables the use of FP4 precision through Transformer Engine. Currently only 
    supports 'nvfp4' which uses NVFP4BlockScaling recipe (requires TE >= 2.7.0.dev0)."""

    fp4_recipe: Optional[Literal['nvfp4', 'custom']] = "nvfp4"
    """If set, enables the use of FP4 precision through Transformer Engine. Currently only
    'nvfp4' is supported which uses NVFP4BlockScaling recipe for Blackwell+ architecture."""

    fp4_param: bool = field(
        default=False, metadata={"argparse_meta": {"arg_names": ["--fp4-param-gather"]}}
    )
    """If set, keep the parameters in fp4 precision to save memory. This option must be used
    together with fp4 mode (i.e., TransformerConfig.fp4 is not None). Note that not all parameters
    will be converted to fp4; for example, biases will remain unchanged."""

    fp4_quantizer_factory: Optional[str] = None
    """Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory.
    Required when fp4_recipe is custom."""

    high_priority_a2a_comm_stream: bool = False
    """If True, the communication stream created by set_streams for combined 1f1b
    a2a overlap is created with CUDA high priority."""

    ####################
    # MoE related
    ####################
    moe_shared_expert_intermediate_size: Optional[int] = None
    """Shared expert total ffn hidden size.
    It should be equal to 'num_shared_experts * ffn_size_of_each_shared_expert' if
    there are multiple shared experts.
    None means no shared expert.
    By default, the shared experts execute before the router. However, when
    moe_shared_expert_overlap or overlap_moe_expert_parallel_comm is set,
    the shared experts execute after the router, before the routed experts.
    This makes the gradients from the router and the shared experts added in
    different orders to the hidden_states, causing minor numerical differences
    in the hidden_states gradient."""

    moe_shared_expert_gate: bool = False
    """Enable gate for shared expert. Only effective when 
    moe-shared-expert-intermediate-size is set."""

    moe_shared_expert_overlap: bool = False
    """Enable overlapping between shared expert computations and dispatcher communications.
    Without this, the shared experts execute before the router. 
    Only effective when moe-shared-expert-intermediate-size is set.
    """

    use_grouped_gemm_for_shared_expert: bool = False
    """Use GroupedLinear(num_groups=1) for the shared expert MLP to trigger the
    Transformer Engine grouped SwiGLU fusion path. Only effective when
    moe-shared-expert-intermediate-size is set.
    """

    moe_shared_expert_glu_interleave_size: Optional[int] = None
    """When set, GLU activations in the shared expert MLP will use a block
    interleaved format. This is only effective when
    use_grouped_gemm_for_shared_expert is set.
    """
    moe_shortcut_connection: bool = False
    """Enable ScMoE shortcut-connected routing. When enabled, the MoE router and routed experts
    process the preceding layer's output (via a shortcut connection) instead of the current layer's
    post-attention representation, allowing the two layers to be run in parallel and hiding the MoE
    layer's A2A coommuunication. Supported only by HybridStack and requires num_moe_experts > 0. 
    CUDA graphs are not supported. For the first MoE layer (no preceding layer), falls back to 
    standard routing."""

    moe_shortcut_post_norm: bool = False
    """Apply the configured normalization to the combined routed and shared expert output.
    Requires moe_shortcut_connection = True."""

    moe_shortcut_parallel: bool = False
    """Overlap shortcut MoE All-to-All communication with paired Attention/Mamba compute.
    Dispatch and combine collectives run on a side CUDA stream; routing, experts, and paired
    compute remain on the main stream. Requires moe_shortcut_connection = True and
    num_moe_experts > 0. Mutually exclusive with moe_shared_expert_overlap and unsupported with
    full activation recomputation."""

    moe_layer_freq: Union[int, List[int]] = 1
    """Frequency between MoE layers and Dense layers. Accepts either:
    - An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers.
    - A list that defines a custom pattern, e.g.: [1,1,1,0,1,1,1,0,1,1,1,0]"""

    moe_ffn_hidden_size: Optional[int] = None
    """MoE Feed-Forward Network hidden size. If not specified, defaults to the ffn_hidden_size."""

    moe_router_load_balancing_type: Union[str, List[str]] = "aux_loss"
    """The load balancing strategy for the router.
    Options:
    - "aux_loss": Load balancing loss used in GShard and SwitchTransformer, calculated at
    micro-batch level.
    - "seq_aux_loss": Load balancing loss used in DeepSeekV2 and DeepSeekV3, computes loss
    for each individual sample.
    - "global_aux_loss": Load balancing loss calculated at global batch level.
    - "sinkhorn": Balancing algorithm used in S-BASE.
    - "quantile_balancing": Dual coordinate-descent quantile balancing (QB). Load balance is
    handled entirely by an internal per-expert bias update; auxiliary losses must be disabled
    (`moe_aux_loss_coeff` = 0) when QB is selected.
    - "none": No load balancing.
    A list of strings can be provided to combine multiple aux-loss load balancing types.
    The default is "aux_loss".
    """

    moe_router_topk: int = 2
    """Number of experts to route to for each token."""

    moe_enable_routing_replay: bool = False
    """If True, enable the routing replay feature for MoE layers."""

    moe_router_topk_limited_devices: Optional[int] = None
    """Number of EP ranks to consider for each token in group-limited routing,
    DEPRECATED and replaced by moe_router_num_groups and moe_router_group_topk.
    """

    moe_router_padding_for_quantization: Optional[bool] = False
    """Whether to pad the routing_map to make sure the number of tokens each expert receives
    is a multiple of 16/32 for quantized precision (e.g., FP8, FP4). This can remove the explicit
    padding in the GroupedMLP layer."""

    moe_router_padding_for_fp8: Optional[bool] = False
    """[Compatibility alias for moe_router_padding_for_quantization]
    Enabling this will also enable moe_router_padding_for_quantization."""

    moe_router_num_groups: Optional[int] = None
    """Number of groups to divide experts into for group-limited routing.
    When using group-limited routing:
    1. Experts are divided into 'moe_router_num_groups' equal-sized groups
    2. For each token, 'moe_router_group_topk' groups are selected based on sum of
    top-('moe_router_topk'/'moe_router_group_topk') routing scores within each group
    3. From these selected groups, 'moe_router_topk' individual experts are chosen
    Two common use cases:
    - Device-limited routing: Set 'moe_router_num_groups' equal to expert parallel size (EP)
    to limit each token to experts on a subset of devices
    (See DeepSeek-V2: https://arxiv.org/pdf/2405.04434)
    - Node-limited routing: Set 'moe_router_num_groups' equal to number of nodes in EP group
    to limit each token to experts on a subset of nodes
    (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437)
    """

    moe_router_group_topk: Optional[int] = None
    """Number of selected groups for group-limited routing."""

    moe_router_pre_softmax: bool = False
    """Enable pre-softmax(pre-sigmoid) routing for MoE, which means softmax is before the
    top-k selection.
    By default, softmax is done after top-k."""

    moe_router_topk_scaling_factor: Optional[float] = None
    """Scaling factor for routing score in top-k selection. Defaults to None, which means no
    scaling."""

    moe_router_score_function: Literal['softmax', 'sigmoid', 'sqrtsoftplus'] = "softmax"
    """Score function for MoE routing. Can be "softmax", "sigmoid" or "sqrtsoftplus"."""

    moe_router_dtype: Optional[Literal['fp32', 'fp64']] = None
    """Data type for routing and expert output weighted averaging. Using fp32 or fp64 can
    improve stability especially when the number of experts is large (e.g. finegrained-moe).
    None means no changes for dtype."""

    moe_router_skip_muon: bool = False
    """Use the scalar optimizer instead of Muon for MoE router parameters."""

    moe_router_enable_expert_bias: bool = False
    """TopK routing with dynamic per-expert bias in the aux-loss-free load balancing strategy.
    The routing decision is based on the sum of the routing scores and the expert bias.
    See https://arxiv.org/abs/2408.15664 for details."""

    moe_router_bias_update_rate: float = 1e-3
    """The expert bias is updated based on the number of assigned tokens to each expert
    in a global batch, where the bias is increased for the experts with less assigned tokens
    and decreased for the experts with more assigned tokens.
    The default value 1e-3 is same as that used in DeepSeekV3."""

    moe_router_quantile_balancing_ema: float = 0.0
    """EMA coefficient for the quantile-balancing per-expert bias (`qb_beta`), used only when
    `moe_router_load_balancing_type` is "quantile_balancing". At each global batch the bias is
    updated as `qb_beta = ema * qb_beta + (1 - ema) * local_quantile`. The default 0.0 means
    no memory: the bias is replaced by the latest global-batch quantile estimate each step."""

    moe_router_force_load_balancing: bool = False
    """[Experimental] Force load balancing with random logits for MoE router, supports naive topk 
    and group-limited topk. This is an experimental feature and only for benchmark."""

    moe_router_force_biased: Optional[float] = None
    """Apply random expert bias in normal distribution with specified std
    to router logits. Shared seed across all ranks ensures identical bias.
    If positive, generates new random bias each forward pass.
    If negative, generates bias once per layer and reuses it (abs value is std).
    This is an experimental feature for benchmarking purposes."""

    moe_num_hash_layers: int = 0
    """Number of leading MoE layers that use hash-based routing.
    In HybridModel this counts MoE positions in the layer pattern rather than
    all hybrid symbols. Other transformer stacks use the layer number directly."""

    hash_moe_vocab_size: Optional[int] = None
    """TP-independent vocabulary size of the token-to-expert lookup table.
    Required when ``moe_num_hash_layers > 0``."""

    use_grouped_gemm_for_dense_mlp: bool = False
    """Use GroupedLinear(num_groups=1) for dense MLP to trigger the
    ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 fusion on SM100+ with MXFP8 recipe.
    Requires ``use_te_op_fuser=True`` and SwiGLU activation.
    """

    moe_grouped_gemm: bool = False
    """Use grouped GEMM to execute multiple local MoE experts together.

    The concrete implementation is selected by Transformer Engine. Set
    ``moe_use_grouped_tensor=True`` to use its CUDA-graph-safe GroupedTensor path.
    """

    moe_use_grouped_tensor: bool = False
    """Use Transformer Engine's native GroupedTensor path for grouped MoE GEMMs.

    This path uses padded expert segments and CUDA split metadata so it can be captured in CUDA
    graphs. Enabling the Transformer Engine operation fuser also enables this option.
    """

    moe_single_grouped_weight: bool = False
    """When using TE GroupedLinear for MoE experts, store expert weights as a single grouped
    parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` and
    ``moe_use_grouped_tensor=True``.
    """

    moe_single_grouped_bias: bool = False
    """When using TE GroupedLinear for MoE experts, store expert biases as a single grouped
    parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``,
    ``moe_use_grouped_tensor=True``, and ``add_bias_linear=True``."""

    moe_aux_loss_coeff: Union[float, List[float]] = 0.0
    """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended.
    If a list of load balancing types is provided for `moe_router_load_balancing_type`,
    a corresponding list of coefficients should be provided here."""

    moe_z_loss_coeff: Optional[float] = None  # 1e-3 would be a good start value for z-loss
    """Scaling coefficient for the z-loss. A starting value of 1e-3 is recommended."""

    moe_input_jitter_eps: Optional[float] = None
    """Add noise to the input tensor by applying jitter with a specified epsilon value."""

    moe_token_dropping: bool = False
    """This feature involves selectively dropping and padding tokens for each expert to achieve a
    specified capacity, similar to GShard, Switch-Transformer, and DeepSpeed-MoE. Note that this is
    currently unsupported so should remain False."""

    moe_token_dispatcher_type: Literal['allgather', 'alltoall', 'flex'] = "allgather"
    """The type of token dispatcher to use. The default is 'allgather'.
    Options are 'allgather','alltoall' and 'flex'."""

    moe_enable_deepep: bool = False
    """[Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models."""

    moe_flex_dispatcher_backend: Literal['deepep', 'hybridep', 'ncclep'] = "deepep"
    """[Experimental] The backend to use for flex token dispatcher. The default is "deepep".
    Options are "deepep", "hybridep", and "ncclep". Currently only "hybridep" backend supports
    the MNNVL case. "ncclep" uses NVIDIA NCCL Expert Parallelism via TransformerEngine's
    transformer_engine.pytorch.ep API."""

    moe_permute_fusion_into_hybridep: bool = False
    """Fuse token rearrangement ops during token dispatching for HybridEP."""

    moe_hybridep_pad_uneven_dispatch_inputs: bool = False
    """Pad uneven HybridEP dispatch inputs to the group maximum before dispatch.
    Enable when local HybridEP input token counts can differ across ranks, for example
    with dynamically packed THD inputs. Leave disabled when dispatcher inputs are
    already padded to equal token counts.
    """

    moe_per_layer_logging: bool = False
    """Enable per-layer logging for MoE, currently supports auxiliary loss and z loss."""

    moe_expert_capacity_factor: Optional[float] = None
    """moe_expert_capacity_factor (float): The capacity factor for each expert, None means no token
    will be dropped. The default is None."""

    moe_pad_expert_input_to_capacity: bool = False
    """moe_pad_expert_input_to_capacity (bool): If True, pads the input for each expert to match
    the expert capacity length, effective only after the moe_expert_capacity_factor is set. The
    default setting is False."""

    moe_pad_experts_for_cuda_graph_inference: bool = False
    """moe_pad_experts_for_cuda_graph_inference (bool): If True, the router will switch to dropping
    and padding during decode time which does not have a D2H sync. The capacity factor is set to the
    max that an expert could see during inference so no tokens are actually dropped. The default
    setting is False."""

    moe_token_drop_policy: Literal['probs', 'position'] = "probs"
    """The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with
    the lowest probabilities will be dropped. If "position", tokens at the end of each batch will
    be dropped.
    """

    moe_layer_recompute: bool = False
    """Memory optimization: checkpointing moe_layer to save actiavtion memory."""

    moe_permute_fusion: bool = False
    """Fuse token rearrangement ops during token dispatching."""

    moe_router_fusion: bool = False
    """Enable fusion for MoE TopK routing and aux-loss computation. This is only
    supported in TransformerEngine 2.7.0 and above.
    """

    moe_router_aux_loss_fusion: Optional[bool] = None
    """Enable fusion for the MoE aux loss only, independently of the fused TopK routing.
    ``None`` follows ``moe_router_fusion`` and is resolved to a concrete bool in
    ``__post_init__``.
    """

    moe_apply_probs_on_input: bool = False
    """Apply probs on input of experts instead of applying after activation and glu."""

    moe_latent_size: Optional[int] = None
    """Latent projection dimension for MoE. If None, MoE latent projections are not used."""

    moe_use_norm_before_up_proj: bool = False
    """Apply normalization before the latent-to-hidden MoE projection. Requires
    ``moe_latent_size`` to be set."""

    gtp_remat_opt_in_modules: list[str] = field(default_factory=list)
    """Extra modules to apply GTP_remat weight sharding to, beyond the default set (attention,
    Mamba, MLP, expert linears, embeddings). Allowed values:

      - ``"moe_latent_proj"`` — shard ``fc1_latent_proj`` / ``fc2_latent_proj`` (MoE latent
        projections, ``parallel_mode="duplicated"``). Only beneficial when ``moe_latent_size``
        is large enough for the all-gather to amortize.
    """

    moe_flex_dispatcher_num_sms: Optional[int] = None
    """Number of SMs for the flex token dispatcher's dispatch/combine communication, for all
    backends (deepep, hybridep, ncclep). None lets each backend use its own default. Unifies the
    deprecated per-backend moe_{deepep,hybridep}_num_sms knobs (routed in __post_init__)."""

    moe_deepep_num_sms: Optional[int] = None
    """DEPRECATED: use moe_flex_dispatcher_num_sms. Number of SMs to use for DeepEP (historical
    default 20). If set, routed to moe_flex_dispatcher_num_sms in __post_init__."""

    moe_hybridep_num_sms: Optional[int] = None
    """DEPRECATED: use moe_flex_dispatcher_num_sms. Number of SMs to use for HybridEP (None uses the
    default from DeepEP). If set, routed to moe_flex_dispatcher_num_sms in __post_init__."""

    moe_hybridep_num_blocks_permute: Optional[int] = None
    """Number of CUDA thread blocks for the permute part in HybridEP.
    When permute_fusion_into_hybridep is True, this sets the number
    of SMs for the permute part (only 1 block per SM)."""

    moe_hybridep_num_blocks_unpermute: Optional[int] = None
    """Number of CUDA thread blocks for the unpermute part in HybridEP.
    When permute_fusion_into_hybridep is True, this sets the number
    of SMs for the unpermute part (only 1 block per SM)."""

    moe_hybridep_num_sms_preprocessing: int = 108
    """Number of SMs to use for HybridEP preprocessing (metadata scan kernel)."""

    moe_hybridep_routing_map_mode: Literal['indices', 'bool'] = 'indices'
    """Routing-map format for HybridEP. ``indices`` requests int16 top-k indices and is the
    default, while ``bool`` forces the bool token-to-expert map. Index routing remains gated on
    Transformer Engine and HybridEP support and the int16 expert limit; unsupported configurations
    fall back to the bool routing-map path. ``moe_pad_expert_input_to_capacity`` also forces the
    bool path (with a warning): a pad-to-capacity routing map can hold more than topk assignments
    per token, which dense top-k indices cannot represent."""

    moe_ncclep_zero_copy: bool = False
    """For the 'ncclep' flex dispatcher: use the NCCL symmetric-memory zero-copy IO path
    (ep_bootstrap zero_copy + symm-mem-backed receive/combine buffers) instead of the default HBM
    staged-copy path, saving one copy on the wire. Requires moe_expert_rank_capacity_factor (the
    symm-mem buffers are a fixed [recv_capacity, hidden] and cannot be resized per step) and the
    fused op (use_transformer_engine_op_fuser). Defaults to False."""

    moe_dispatch_fwd_dtype: Literal['bf16', 'mxfp8'] = 'bf16'
    """Wire dtype of the MoE dispatch forward payload ('ncclep' flex dispatcher only). With
    'mxfp8', TransformerEngine quantizes the payload before the all-to-all and the receive
    buffer comes back as a per-expert MXFP8 GroupedTensor that the grouped GEMM consumes
    directly. Requires moe_grouped_gemm and use_transformer_engine_op_fuser. Defaults to
    'bf16' (no quantization on the wire)."""

    moe_combine_bwd_dtype: Literal['bf16', 'mxfp8'] = 'bf16'
    """Wire dtype of the MoE combine backward gradient ('ncclep' flex dispatcher only). With
    'mxfp8', TransformerEngine quantizes the gradient before the all-to-all and the
    expert-output gradient comes back as a per-expert MXFP8 GroupedTensor that the grouped
    GEMM backward consumes directly. Same requirements as moe_dispatch_fwd_dtype. Defaults to
    'bf16' (no quantization on the wire)."""

    moe_mlp_glu_interleave_size: Optional[int] = None
    """When set, GLU activations in the MoE grouped MLP layer will use a
    block interleaved format. Instead of interpreting the input tensor
    as a concatenation of gates and linear units, it will be
    interpreted as alternating blocks of gates and linear units.
    This data format is experimental and primarily intended to enable
    advanced fused kernels."""

    moe_expert_rank_capacity_factor: Optional[float] = None
    """moe_expert_rank_capacity_factor (float): The capacity factor for each expert rank, i.e. the
    per-rank token budget. None means no token will be dropped. The default is None.
    With the 'hybridep' backend, tokens exceeding this budget are dropped. With the 'ncclep'
    backend, setting it selects static shapes (fixed receive buffer, CUDA-graph capturable) and
    leaving it None selects eager mode (receive buffer sized per step from the received-token
    count)."""

    ##################
    # Context Parallel
    ##################
    cp_comm_type: Optional[Union[str, List[str]]] = None
    """Inter-gpu communication type for context parallelism.
    str: all layers share same communication type.
    List[str]: each layer has its separate communication type.
    cp_comm_type of each layer can be "p2p" or "all_gather" or "a2a" or "a2a+p2p".
    "p2p": Exchange KV chunks with P2P communications in ring topology. P2P is async and can be
    overlapped with attention compute.
    "all_gather": All-gather to get full sequence of KV before attention. The all-gather is not
    async, and cannot be overlapped.
    "a2a": Like DeepSpeed Ulysses, scatter attention heads across the CP group, and gather to get
    full sequence of QKV.
    "a2a+p2p": A hierarchical implementation of context parallelism to attention.
    It uses A2A communications in low-level CP groups (e.g., via NVLink),
    and P2P communications in high-level CP groups (e.g., via IBLink).
    """

    linear_cp_mode: Literal["headwise", "chunkwise"] = "headwise"
    """Context-parallel algorithm for recurrent and linear-attention layers."""

    linear_cp_layout: CPLayout = "zigzag"
    """CP layout for linear-attention layers."""

    attention_cp_layout: CPLayout = "zigzag"
    """CP layout for softmax-attention layers."""

    ##################
    # Cuda Graphs
    ##################
    enable_cuda_graph: bool = False
    """DEPRECATED and replaced by cuda_graph_impl.
    When set to true, either partial CUDA graph (1/many CUDA graph per layer) or full iteration
    CUDA graph (1 CUDA graph for whole iteration excluding optimizer) is enabled. cuda_graph_modules
    determines the scope of graph capture."""

    cuda_graph_use_single_mempool: bool = True
    """For cuda_graph_impl "local" with cuda_graph_scope "full_iteration" only.

    When True, full-iteration graph replay (training and evaluation) and optimizer graph
    capture/replay share the same CUDA graph memory pool."""

    cuda_graph_retain_backward_graph: bool = False
    """When set to true, cudagraph backward passes will be graph captured with 'retain_grad=True'
    This may enable cudagraphs for certain modules that are not completely cudagraph safe. For
    more details, see: https://pytorch.org/docs/stable/generated/torch.Tensor.backward.html."""

    cuda_graph_warmup_steps: int = 3
    """Number of warmup steps for CUDA graphs. Note: GTP (``gtp_weight_remat_size > 1``) forces a
    minimum of 2 per-graph warmup steps regardless of this value, because the first warmup builds
    the weight-prefetch chain and the second exercises the prefetch path before capture."""

    external_cuda_graph: bool = False
    """DEPRECATED and replaced by cuda_graph_impl.
    When set to true, TransformerLayer layers are swapped with user provided CUDA graphs."""

    cuda_graph_impl: Literal['none', 'local', 'transformer_engine', 'full_iteration'] = "none"
    """Determines the CUDA graph capture implementation.
    "none": no CUDA graph.
    "local": MCore CUDA graph implementation. During training, graphable modules own per-layer
    CUDA graphs controlled by cuda_graph_modules. During inference, graph ownership is controlled
    separately by inference_cuda_graph_scope.
    "transformer_engine": Transformer Engine CUDA graph implementation. During training, TE
    make_graphed_callables() creates per-layer CUDA graphs controlled by cuda_graph_modules.
    Inference CUDA graphs are not supported; inference_cuda_graph_scope must be "none".
    "full_iteration": full-iteration CUDA graph implementation for the training iteration
    (1 CUDA graph for the whole forward-backward path excluding the optimizer step). Inference
    CUDA graphs are not supported; inference_cuda_graph_scope must be "none".
    cuda_graph_modules has no effect when cuda_graph_impl="none" and must be empty when
    cuda_graph_impl="full_iteration"."""

    cuda_graph_modules: Union[str, CudaGraphModule, List[str], List[CudaGraphModule]] = "full"
    """Selects training capture coverage within per-layer CUDA graphs (local and
    transformer_engine implementations).
    Valid values are "attn", "mlp", "moe", "moe_router", "moe_preprocess", and "mamba":
    "attn": captures operations in TransformerLayer._forward_attention().
    "mlp": captures operations in TransformerLayer._forward_mlp() for a dense layer.
    "moe": captures operations in TransformerLayer._forward_mlp() for a MoE layer.
    "moe_router": captures operations in TransformerLayer._forward_mlp() up to MoELayer.router(),
    including the shared experts if they are not overlapped with EP comm.
    "moe_preprocess": captures operations in MoELayer.preprocess(). Must be used together with
    "moe_router".
    "mamba": captures the mamba layer.
    An empty list means capturing the whole Transformer layer.
    This field is meaningless when cuda_graph_impl="full_iteration" and must be empty.
    Backward compatibility: "full" is deprecated but kept for backward compatibility; it is
    transformed to an empty list in __post_init__. The deprecated values "full_iteration" and
    "full_iteration_inference" are also accepted and migrated to the new API in __post_init__."""

    inference_cuda_graph_scope: Optional[InferenceCudaGraphScope] = field(
        default=None,
        metadata={
            "argparse_meta": {
                "type": str,
                "choices": [scope.name for scope in InferenceCudaGraphScope],
            }
        },
    )
    """Controls the CUDA graph scope during inference.
    When unset, the effective default is derived from cuda_graph_impl:
    "local" -> "layer", all other impls -> "none".
    "none": inference runs in eager mode (no CUDA graphs).
    "layer": inference graphs are owned at the module/layer boundary, e.g. TransformerLayer or
    MambaLayer.
    "block": inference graphs are owned by the enclosing block, e.g. TransformerBlock or
    HybridBlock.
    Currently supported combinations are:
    cuda_graph_impl="local" -> "layer" or "block";
    all other cuda_graph_impl values -> "none"."""

    cuda_graph_scope: Optional[
        Union[
            str, CudaGraphModule, CudaGraphScope, List[Union[str, CudaGraphModule, CudaGraphScope]]
        ]
    ] = None
    """Deprecated: renamed to cuda_graph_modules. Accepted for backward compatibility and
    migrated to cuda_graph_modules in __post_init__. Will be removed in a future release.
    CudaGraphScope instances deserialized from pre-refactor checkpoints are converted to their
    string names before normalization so existing CUDA_GRAPH_MODULES_DEPRECATIONS handles them."""

    ####################
    # Hyper-Connection Configuration
    ####################
    enable_mhc_connections: bool = False
    """Enable mHC residual connections."""

    mhc_num_residual_streams: int = 4
    """Number of residual streams (n in paper)."""

    mhc_sinkhorn_iterations: int = 20
    """Number of Sinkhorn-Knopp iterations for doubly stochastic projection."""

    mhc_init_gating_factor: float = 0.01
    """Initial value of Gating Factor (alpha in paper)."""

    use_fused_mhc: bool = False
    """Use fused kernels for mHC operations when supported.

    With the default ``auto`` backend policy, selection is operation-specific
    and unavailable accelerated implementations fall back to native torch.
    Set ``mhc_fused_backend`` to request an explicit backend policy.
    """

    mhc_fused_backend: Literal["auto", "native", "triton", "cutile"] = "auto"
    """Backend policy for fused mHC operations.

    ``auto`` selects the fastest available implementation for each operation.
    Explicit policies require the requested dependency and device support, and
    never select a different accelerated backend. Operations without an
    implementation in the selected backend retain their native implementation.
    """

    mhc_recompute_layer_num: Optional[int] = None
    """Number of layers per MHC recompute block.

    When set, every `mhc_recompute_layer_num` layers form a recompute block. The last layer
    in each recompute block (i.e., layer_number % mhc_recompute_layer_num == 0 or the final
    layer in the transformer block) will:
    - NOT checkpoint its final MLP BDA
    - Register the unified recompute hook on its MLP BDA output
    - A new CheckpointWithoutOutputManager is created for subsequent layers

    If None, all layers in the transformer block share a single recompute block.

    Must be a positive integer when set."""

    ####################
    # Wide Residual Configuration
    ####################
    wide_residual: Optional[WideResidualConfig] = None
    """Optional streamwise wide-residual architecture configuration.

    When set, the model carries ``num_streams * hidden_size`` features between
    layers while attention and MLP branches continue to operate at ``hidden_size``.
    """

    residual_stream_recompute_num_layers: Optional[int] = None
    """Number of local layers per ordered residual-stream replay block.

    ``None`` places all local layers in one block. This setting requires selective
    recomputation with ``"residual_stream"`` in ``recompute_modules``.
    """

    ####################
    # miscellaneous
    ####################
    clone_scatter_output_in_embedding: bool = True
    """When set to True, clone the output of scatter_to_sequence_parallel_region in embedding layer
    to facilitate garbage collection of input."""

    disable_parameter_transpose_cache: bool = False
    """When set to true, the parameter transposes are not cached for subsequent iterations."""

    config_logger_dir: str = ""
    """When non-empty, dumps entry-point configs to config_logger_dir"""

    flash_decode: bool = False
    """ Use the optimized flash decoding kernel during inference. """

    batch_invariant_mode: bool = False
    """If true, uses batch-invariant kernels that provide deterministic forward execution regardless
       of batch size. This ensures bitwise identical results when the same inputs are processed
       in different batch configurations. This will significantly affect speed of 
       training and inference as the kernels are not full optimized.
       Defaults to False."""

    batch_invariant_backend: Literal["te_native", "deepgemm", "triton"] = "te_native"
    """Which batch-invariant GEMM backend to use when batch_invariant_mode is
    enabled: "te_native" (default: keep the native cuBLASLt kernels and obtain
    invariance via workspace starvation — lowest overhead, no extra
    dependencies, and the configuration verified bitwise-identical to the TE
    training forward), "deepgemm" (DeepGEMM bf16 kernels), or "triton"
    (persistent Triton matmul; any dtype)."""

    batch_invariant_collective: Literal["ordered", "multimem"] = "ordered"
    """Cross-rank EP combine collective under batch_invariant_mode. "ordered"
    (default) reduces with an explicit fixed rank-order fp32 Triton kernel —
    deterministic by construction on any hardware. "multimem" keeps the native
    NVLS in-switch reduce-scatter: measured to return the correctly-rounded
    exact fp32 sum (bitwise-equal to an fp64 reference over 16.7M adversarial
    channels on B200), deterministic and batch-invariant, with better scaling
    at large NVLink domains; software paths that must match it bitwise should
    accumulate in fp64."""

    use_te_activation_func: bool = False
    """Whether to use ffn activation functions implemented by TransformerEngine"""

    use_te_rng_tracker: bool = False
    """ Whether to use the TE or MCore version of the RNG tracker. """

    inference_rng_tracker: bool = False
    """ Whether we should instantiate a separate RNG tracker for inference. """

    inference_sampling_seed: int = 42
    """ Random seed to use for sampling during inference. """

    symmetric_ar_type: Optional[Literal['two_shot', "one_shot", "multimem_all_reduce"]] = None
    """What type of symmetric all reduce to use. The default is None
    which is no use of symmetric memory.
    """

    nccl_all_reduce_for_prefill: bool = False
    """If True, use NCCL all-reduce kernels when symmetric all-reduce is enabled."""

    use_inference_optimized_layers: bool = False
    """If True, use inference optimized transformer layers during inference."""

    inference_fuse_tp_communication: bool = False
    """ If true, uses a fused reduce-scatter-residual-norm-allgather kernel during inference. """

    inference_disable_triton_nvls_kernels: bool = False
    """ If true, disables the use of Triton NVLS kernels during inference. """

    inference_grouped_gemm_backend: Literal['flashinfer', 'torch', 'vllm'] = "vllm"
    """Specifies the backend to use for grouped GEMM operations during inference.
    Options:
    - 'flashinfer': Uses FlashInfer cutlass_fused_moe for BF16 and TRT-LLM routed
      block-scale MoE for MXFP8. The MXFP8 path retains canonical expert weights
      for refit and also stores a padded TRT-LLM Major-K copy, increasing
      expert-weight memory relative to the torch backend.
    - 'torch': Uses torch.nn.functional.grouped_mm (mcore_fused_moe with Triton kernels).
      Supports both BF16 and MXFP8.
    - 'vllm': Uses vLLM's Triton fused MoE kernel for BF16. Avoids physical token
      permutation via indirect addressing. MXFP8 expert layers use MCore's scaled
      grouped-GEMM path, allowing per-layer mixed BF16/MXFP8 policies.
    """

    inference_moe_disable_fused_quant_kernels: bool = False
    """When False (default), use fused kernels that combine permute/activation with
    MXFP8 quantization + swizzle into a single kernel launch. Only applies when
    fp8_recipe='mxfp8' with inference_grouped_gemm_backend='torch' or 'vllm'. Set to
    True to disable fusion and use separate kernel launches (useful for debugging)."""

    inference_flashinfer_mxfp8_token_capacity: int | None = None
    """Optional fixed token-row capacity for FlashInfer routed MXFP8 MoE.

    Decode-only dynamic-inference graphs use this fixed prefix when their
    host-known EP-wide token ceiling fits. Prefill, mixed, static-inference,
    and oversized decode graphs retain the full dispatcher buffer. Requires
    the NVLS inference dispatcher and EP > 1.
    """

    inference_moe_token_dispatcher_type: Literal['nccl', 'nvls'] = 'nvls'
    """Token dispatcher to use for MoE expert parallelism during inference.
    - 'nccl': AllGather/ReduceScatter via NCCL. Fixed token counts per rank; requires
      decode-only CUDA graphs (forced automatically).
    - 'nvls': Variable-count AllGather-V/ReduceScatter-V via NVLS multimem kernels.
      Requires Hopper+ GPUs with NVLink and symmetric memory. Default.
    Only applies when transformer_impl='inference_optimized' and EP > 1."""

    mrope_section: Optional[List[int]] = None
    """ Multimodal rope section is for channel dimension of temporal, height and width
    in rope calculation. """

    is_hybrid_model: bool = False
    """ Indicates whether this is a hybrid model. """

    mamba_state_dim: int = 128
    """The dimensionality of the state representation in Mamba layers."""

    mamba_head_dim: int = 64
    """The dimensionality of the heads in the Mamba layers."""

    mamba_num_groups: int = 8
    """The number of groups used in Mamba layers."""

    mamba_num_heads: Optional[int] = None
    """The number of heads used in Mamba layers.
    If None, the number of heads will be hidden_size * expand // mamba_head_dim."""

    gdp_num_householder: int = 3
    """The number of Householder reflections used in Gated Delta Product layers."""

    mamba_training_ssm_states_dtype: Optional[torch.dtype] = None
    """dtype of the materialized inter-chunk SSM states in Mamba training forwards and backwards.
    None causes the states to follow the activation dtype."""

    use_mamba_mem_eff_path: bool = field(
        default=True, metadata={"argparse_meta": {"arg_names": ["--disable-mamba-mem-eff-path"]}}
    )
    """Controls usage of the memory efficient path for Mamba layers."""

    gdp_cutedsl_kernel: bool = False
    """Whether to use the CuTeDSL kernel for the GatedDeltaProduct mixer."""

    gdp_num_chunk_states_to_recompute: int = 2
    """Checkpoint-coarsening ratio N in [0, 64] for the CuTeDSL GatedDeltaProduct kernel.
    N=0 checkpoints every chunk state for the backward pass (dense, no recompute). N>=1
    stores one checkpoint per group of N+1 chunks (1/(N+1) the checkpoint memory) and the
    backward recomputes each group's N missing chunk states, trading recompute time for
    activation memory monotonically in N (sweet spot N=2..3). Only honored by kernel
    builds that expose ``recompute_chunk_num`` or the legacy
    ``num_chunk_states_to_recompute`` argument."""

    mlp_chunks_for_prefill: int = 1
    """The number of chunks along the sequence dimension to use for MLP computation
    during prefill."""
    mlp_chunks_for_training: int = 1
    """The number of chunks along the sequence dimension to use for MLP computation
    during training."""

    heterogeneous_block_specs: bool = False
    """Whether to use heterogeneous block specs (nemotron-nas architecture)."""

    hetereogenous_dist_checkpoint: bool = False
    """Whether to use heterogenous layers in distributed checkpoint."""

    ####################
    # Quantization
    ####################
    quant_recipe: Optional[RecipeConfig] = None
    """Configuration of any per-module quantization settings to be applied to the model"""

    transformer_impl: Literal['local', 'transformer_engine', 'inference_optimized'] = (
        "transformer_engine"
    )
    """Transformer implementation to use.
    Options are 'transformer_engine' for Transformer Engine and 'local' for MCore."""

    #####################################
    # Fine-grained Activation Offloading
    #####################################
    fine_grained_activation_offloading: bool = False
    """If True, offload the input of the specified modules to the CPU.
    Fine-grained activation offloading is a module-level offloading method
    instead of a layer-level offloading method like cpu_offloading."""

    offload_modules: Optional[list[str]] = field(default_factory=list)
    """The submodules to offload its input.
    choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj",
             "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp", "gdp_qkv",
             "shortcut_post_norm".
    "attn_norm": offload the input of the normalization in the attention part.
    "qkv_linear": offload the input of the qkv linear part.
    "core_attn": offload the input of the core attention part.
    "attn_proj": offload the input of the attn linear projection part.
    "mlp_norm": offload the input of the normalization in the mlp part.
    "expert_fc1": offload the input of the expert fc1 part.
    "moe_act": offload the input of the moe act part.
    "fused_group_mlp": offload the input of the whole fused grouped MLP.
    "gdp_qkv": offload the input of the causal conv and QKV preparation in the
               GatedDeltaProduct mixer.
    "shortcut_post_norm": offload the input of the shortcut output normalization.
            Requires moe_shortcut_connection=True.
    """
    min_offloaded_tensor_size: int = 1024 * 1024
    """The minimum size of the tensor to be offloaded."""

    delay_offload_until_cuda_graph: bool = False
    """If True, delay the offload until the CUDA graph is executed for minimal CPU overhead.
    For more details, see the documentation:
    https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#cuda-graph-integration.
    """

    delta_offload_bytes_across_pp_ranks: int = 0
    """Difference of offload bytes across PP ranks to balance the offload load.
    For more details, see the documentation:
    https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#tuning-parameters.
    """

    activation_offload_fraction: float = 1.0
    """Fraction of eligible activation offload groups to offload across configured modules.
    For details, see:
    https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#activation-offload-fraction.
    """

    moe_paged_stash: bool = False
    """If True, enable paged stash for all routed-expert activations needed for backward"""

    moe_paged_stash_page_size: int = 64
    """Number of tokens per page for paged stash memory management."""

    moe_paged_stash_buffer_size_factor_cuda: float = 1.10
    """Scale factor for paged stash CUDA buffer allocation.

    Sign selects sizing: positive = avg-based, negative = actual-max. Magnitude is headroom
    (e.g. 1.10 = 10%%)."""

    moe_paged_stash_buffer_size_factor_cpu: float = 0.0
    """Scale factor for paged stash host buffer. 0 disables host buffer.

    Same sign convention as moe_paged_stash_buffer_size_factor_cuda: positive = avg-based,
    negative = actual-max; scale = abs(factor)."""

    fine_grained_offloading_max_inflight_offloads: Optional[int] = None
    """Per fine-grained offloading group name, max number of inflight offloads for that name not
    yet joined on the main stream (wait_event on D2H). The same cap applies to every name (e.g.,
    ``moe_act`` and ``qkv_linear`` each have their own pending queue). 0 = wait after every
    offload for that name. 1 = at most one not-yet-waited offload per name, etc. None = do not
    insert these joins. This feature is particularly useful when using with full-iteration CUDA
    graphs"""

    @classmethod
    def from_config(cls, config: "TransformerConfig") -> Self:
        """Create this config type from an existing normalized transformer config.

        The source config's complete instance state is deep-copied without invoking
        the target class's initializer or ``__post_init__``. This preserves normalized
        values and dynamically added attributes while producing an independent config.

        Args:
            config: The transformer config to copy.

        Returns:
            An independent copy of ``config`` whose type is ``cls``.
        """
        new_config = cls.__new__(cls)
        new_config.__dict__ = deepcopy(config.__dict__, {id(config): new_config})
        return new_config

    def _validate_cp_layouts(self) -> None:
        """Validate context-parallel layout settings."""
        if self.linear_cp_mode == "chunkwise" and self.linear_cp_layout != "contiguous":
            raise ValueError("linear_cp_mode='chunkwise' requires linear_cp_layout='contiguous'.")
        if self.context_parallel_size > 1 and self.attention_cp_layout == "contiguous":
            raise ValueError(
                "attention_cp_layout='contiguous' is not yet supported with context parallelism."
            )
        if self.linear_cp_layout == "contiguous" and self.hybrid_context_parallel:
            raise ValueError(
                "hybrid_context_parallel is not supported with linear_cp_layout='contiguous'."
            )
        if (
            self.sequence_packing_scheduler is not None
            and self.context_parallel_size > 1
            and self.linear_cp_layout != self.attention_cp_layout
        ):
            raise ValueError(
                "The sequence-packing scheduler does not support CP layout conversion."
            )
        if (
            self.context_parallel_size > 1
            and self.linear_cp_layout != self.attention_cp_layout
            and self.sequence_parallel
            and self.tensor_model_parallel_size > 1
            and self.tensor_model_parallel_size % 2 != 0
        ):
            raise ValueError(
                "Sequence-parallel CP layout conversion requires an even "
                f"tensor-parallel size, got {self.tensor_model_parallel_size}."
            )

    def __post_init__(self):
        """Python dataclass method that is used to modify attributes after initialization.
        See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more
        details.
        """
        super().__post_init__()
        self._validate_cp_layouts()

        if self.attn_logit_softcapping is not None and not (
            math.isfinite(self.attn_logit_softcapping) and self.attn_logit_softcapping > 0
        ):
            raise ValueError(
                "attn_logit_softcapping must be a positive finite value, got "
                f"{self.attn_logit_softcapping}. Use None to disable softcapping. A cap of 0.0 "
                "disables softcapping in TransformerEngine but collapses every logit to zero in "
                "the local attention path, a negative cap is silently applied as its absolute "
                "value there while FlashAttention ignores it entirely, and a non-finite cap "
                "produces NaN logits."
            )

        # Unset means "follow moe_router_fusion". Resolve it here so every consumer
        # downstream reads a plain bool.
        if self.moe_router_aux_loss_fusion is None:
            self.moe_router_aux_loss_fusion = self.moe_router_fusion

        if self.wide_residual is not None:
            if self.enable_mhc_connections:
                raise ValueError("wide_residual and enable_mhc_connections are mutually exclusive.")
            if self.cuda_graph_impl != "none" or self.enable_cuda_graph or self.external_cuda_graph:
                raise NotImplementedError(
                    "wide_residual does not yet support CUDA graphs; use cuda_graph_impl='none'."
                )
            if self.pipeline_model_parallel_size > 1:
                raise NotImplementedError(
                    "wide_residual does not yet support pipeline_model_parallel_size > 1. "
                    "Inter-stage communication buffers are still sized from hidden_size."
                )
            if self.inference_fuse_tp_communication:
                raise NotImplementedError(
                    "wide_residual is not compatible with inference_fuse_tp_communication. "
                    "The fused inference path assumes an ordinary-width residual tensor."
                )
            if self.heterogeneous_block_specs:
                raise NotImplementedError(
                    "wide_residual does not yet support heterogeneous_block_specs. "
                    "Residual-stream width is currently owned by the enclosing block."
                )
            if self.overlap_moe_expert_parallel_comm:
                raise NotImplementedError(
                    "wide_residual does not yet support overlap_moe_expert_parallel_comm. "
                    "The fine-grained EP-overlap schedule invokes the pre-MLP norm and MLP BDA "
                    "outside TransformerLayer._forward_mlp, bypassing the wide-residual MLP "
                    "read and write connection."
                )

        # Resolve deprecated attention variant spellings up front so that every consumer
        # downstream only has to handle the canonical names. Imported lazily because the
        # spec module imports this one.
        from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
            is_gated_delta_net_variant,
            normalize_experimental_attention_variant,
        )

        if self.experimental_attention_variant is not None:
            self.experimental_attention_variant = normalize_experimental_attention_variant(
                self.experimental_attention_variant
            )

        if self.dsa_kernel_backend is None:
            self.dsa_kernel_backend = (
                "cudnn" if self.experimental_attention_variant == "dsv4_hybrid" else "none"
            )

        if self.use_transformer_engine_op_fuser and self.moe_grouped_gemm:
            self.moe_use_grouped_tensor = True

        if self.moe_use_grouped_tensor and not self.moe_grouped_gemm:
            raise ValueError("moe_use_grouped_tensor=True requires moe_grouped_gemm=True.")

        # When fp32 residual connections are enabled, pipeline parallel communication must
        # use fp32 to match the dtype of the residual stream between pipeline stages.
        if self.fp32_residual_connection and self.pipeline_dtype is not None:
            if self.pipeline_dtype != torch.float:
                log_single_rank(
                    logger,
                    logging.WARNING,
                    f"fp32_residual_connection is enabled, overriding pipeline_dtype "
                    f"from {self.pipeline_dtype} to torch.float to match the "
                    f"residual stream dtype between pipeline stages.",
                )
            self.pipeline_dtype = torch.float

        if self.fp16 and self.bf16:
            raise ValueError(
                f"Only one of self.fp16: {self.fp16} and self.bf16 {self.bf16} should be True."
            )

        if self.gdp_num_householder < 1:
            raise ValueError(
                f"gdp_num_householder must be positive, got {self.gdp_num_householder}."
            )

        # Apply BF16 matmul precision setting if needed
        if self.bf16 and self.disable_bf16_reduced_precision_matmul:
            torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False

        if self.num_attention_heads % self.tensor_model_parallel_size != 0:
            raise ValueError(
                f"num_attention_heads ({self.num_attention_heads}) must be a multiple of "
                f"tensor_model_parallel_size ({self.tensor_model_parallel_size})."
            )

        if self.ffn_hidden_size is None:
            self.ffn_hidden_size = 4 * self.hidden_size

        if self.kv_channels is None:
            self.kv_channels = self.hidden_size // self.num_attention_heads

        if self.num_query_groups is None:
            self.num_query_groups = self.num_attention_heads

        if self.num_query_groups > 0 and self.num_attention_heads % self.num_query_groups != 0:
            raise ValueError(
                f"num_query_groups ({self.num_query_groups}) must be a divisor of "
                f"num_attention_heads ({self.num_attention_heads})."
            )

        if (
            self.num_query_groups % self.tensor_model_parallel_size != 0
            and self.tensor_model_parallel_size % self.num_query_groups != 0
        ):
            raise ValueError(
                f"num_query_groups ({self.num_query_groups}) must be a multiple or divisor of "
                f"tensor_model_parallel_size ({self.tensor_model_parallel_size})."
            )

        if is_gated_delta_net_variant(self.experimental_attention_variant):
            # gdn2 may also be enabled for GDN layers built via the hybrid layer pattern
            # symbol 'G', where linear_attention_freq is unused; the GPT experimental
            # attention route raises a clear error downstream if it is missing.
            if self.experimental_attention_variant == "gdn":
                assert (
                    self.linear_attention_freq is not None
                ), "linear_attention_freq must be set for linear gdn."

            # Check required parameters
            assert (
                self.linear_conv_kernel_dim is not None
            ), "linear_conv_kernel_dim must be set for gated delta net."
            assert (
                self.linear_key_head_dim is not None
            ), "linear_key_head_dim must be set for gated delta net."
            assert (
                self.linear_value_head_dim is not None
            ), "linear_value_head_dim must be set for gated delta net."
            assert (
                self.linear_num_key_heads is not None
            ), "linear_num_key_heads must be set for gated delta net."
            assert (
                self.linear_num_value_heads is not None
            ), "linear_num_value_heads must be set for gated delta net."
            assert self.linear_num_value_heads % self.linear_num_key_heads == 0, (
                f"linear_num_value_heads ({self.linear_num_value_heads}) must be a multiple of "
                f"linear_num_key_heads ({self.linear_num_key_heads})."
            )

            # Check tensor parallelism compatibility
            tp_cp_size = self.tensor_model_parallel_size * self.context_parallel_size
            assert self.linear_num_key_heads % tp_cp_size == 0, (
                f"{self.linear_num_key_heads=} must be a multiple of "
                f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})."
            )
            assert self.linear_num_value_heads % tp_cp_size == 0, (
                f"{self.linear_num_value_heads=} must be a multiple of "
                f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})."
            )
        elif self.experimental_attention_variant == "dsa":
            _validate_dsa_kernel_backend_dependencies(self.dsa_kernel_backend)
            if self.add_bias_linear:
                raise ValueError(
                    "DSA uses AbsorbedMLASelfAttention, which requires add_bias_linear=False. "
                    "Disable linear bias for DSA configs."
                )
            if self.dsa_indexer_topk_freq < 1:
                raise ValueError(
                    f"dsa_indexer_topk_freq must be positive, got {self.dsa_indexer_topk_freq}."
                )
            if self.dsa_indexer_skip_topk_offset < 0:
                raise ValueError(
                    "dsa_indexer_skip_topk_offset must be non-negative, got "
                    f"{self.dsa_indexer_skip_topk_offset}."
                )
        elif self.experimental_attention_variant == "dsv4_hybrid":
            assert self.multi_latent_attention, "DSv4 Hybrid requires multi_latent_attention."
            assert self.csa_compress_ratios is not None, "csa_compress_ratios must be set"
            mtp_layers = self.mtp_num_layers or 0
            minimum_len = self.num_layers + mtp_layers
            assert len(self.csa_compress_ratios) >= minimum_len, (
                f"csa_compress_ratios length ({len(self.csa_compress_ratios)}) must be at least "
                f"num_layers + mtp_num_layers "
                f"({self.num_layers} + {mtp_layers} = {minimum_len})"
            )
            assert all(
                ratio in [0, 4, 128] for ratio in self.csa_compress_ratios
            ), "csa_compress_ratios must be 0, 4, or 128"
            assert (
                self.tensor_model_parallel_size == 1
            ), "DSv4 Hybrid Attention only supports TP size 1."
            assert (
                self.context_parallel_size == 1
            ), "DSv4 Hybrid Attention does not support context parallelism yet."
            assert not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention."
            if self.dsa_kernel_backend == "tilelang":
                raise ValueError(
                    "dsv4_hybrid does not support dsa_kernel_backend='tilelang'; use 'cudnn' "
                    "for fused CSA kernels or 'none' for the PyTorch fallback."
                )
            _validate_dsa_kernel_backend_dependencies(self.dsa_kernel_backend)
            if self.dsa_kernel_backend == "cudnn":
                sm = torch.cuda.get_device_capability()
                assert sm[0] >= 9, (
                    "dsa_kernel_backend='cudnn' requires SM90+ (Hopper or later), "
                    f"but current device has compute capability {sm[0]}.{sm[1]}."
                )
                uses_ratio4_indexer = 4 in self.csa_compress_ratios and not self.csa_dense_mode
                indexer_loss_enabled = (self.dsa_indexer_loss_coeff or 0.0) > 0
                if (
                    sm[0] == 9
                    and uses_ratio4_indexer
                    and indexer_loss_enabled
                    and not self.dsa_indexer_use_sparse_loss
                ):
                    raise ValueError(
                        "DSv4 with fused DSA and dense indexer loss is not supported on SM90 "
                        "because the cuDNN Frontend SM90 dense DSA kernels are not reliable for "
                        "this path. Use sparse indexer loss or set dsa_kernel_backend='none'."
                    )
            self.hetereogenous_dist_checkpoint = True

        if self.gdn_pre_gated_delta_rule_fusion and self.experimental_attention_variant != "gdn":
            raise ValueError(
                "gdn_pre_gated_delta_rule_fusion is only supported with "
                "experimental_attention_variant='gdn' "
                "or deprecated alias experimental_attention_variant='gated_delta_net'."
            )

        if self.gdn_gated_output_norm_fusion and self.experimental_attention_variant != "gdn":
            raise ValueError(
                "gdn_gated_output_norm_fusion is only supported with "
                "experimental_attention_variant='gdn' "
                "or deprecated alias experimental_attention_variant='gated_delta_net'."
            )

        if self.fp8:
            # cannot support first last layer bf16 with delayed scaling
            if self.first_last_layers_bf16 and self.fp8_recipe == Fp8Recipe.delayed:
                raise ValueError("Delayed scaling does not support first / last layer in BF16.")

            # max bf16 layers per pipeline stage
            max_bf16_layers_per_pipeline_stage = (
                self.num_layers // self.pipeline_model_parallel_size
            )

            # check start/end bf16 layer counts are valid
            if self.first_last_layers_bf16:
                if (
                    self.num_layers_at_start_in_bf16 < 0
                    or self.num_layers_at_start_in_bf16 > max_bf16_layers_per_pipeline_stage
                ):
                    raise ValueError(
                        f"num_layers_at_start_in_bf16 ({self.num_layers_at_start_in_bf16}) must be "
                        f"between 0 and number of layers per pipeline stage "
                        f"({max_bf16_layers_per_pipeline_stage})."
                    )
                if (
                    self.num_layers_at_end_in_bf16 < 0
                    or self.num_layers_at_end_in_bf16 > max_bf16_layers_per_pipeline_stage
                ):
                    raise ValueError(
                        f"num_layers_at_end_in_bf16 ({self.num_layers_at_end_in_bf16}) must be "
                        f"between 0 and number of layers per pipeline stage "
                        f"({max_bf16_layers_per_pipeline_stage})."
                    )

            if self.fp8_recipe == Fp8Recipe.custom:
                if not self.fp8_quantizer_factory:
                    raise ValueError(
                        "fp8_quantizer_factory must be provided when fp8_recipe is 'custom'. "
                        "Specify a Python import path (e.g., package.module.quantizer_factory) "
                        "via --fp8-quantizer-factory."
                    )

        if self.fp8_param and not self.fp8:
            raise ValueError("fp8_param must be used together with fp8 mode.")

        if self.fp8_output_proj:
            if not self.fp8:
                raise ValueError("fp8_output_proj must be used together with fp8 mode.")
            if self.fp8_recipe != Fp8Recipe.mxfp8:
                raise ValueError(
                    f"fp8_output_proj requires fp8_recipe='mxfp8', got " f"'{self.fp8_recipe}'."
                )

        # FP4 validation
        if self.fp4_param and not self.fp4:
            raise ValueError("fp4_param must be used together with fp4 mode.")

        if self.fp4 and self.fp8:
            raise ValueError("fp4 and fp8 cannot be used simultaneously. Please choose one.")

        if self.fp4 and self.fp4_recipe == Fp4Recipe.custom:
            if not self.fp4_quantizer_factory:
                raise ValueError(
                    "fp4_quantizer_factory must be provided when fp4_recipe is 'custom'. "
                    "Specify a Python import path (e.g., package.module.quantizer_factory) "
                    "via --fp4-quantizer-factory."
                )

        if self.apply_query_key_layer_scaling:
            self.attention_softmax_in_fp32 = True

        if self.expert_model_parallel_size > 1 and self.num_moe_experts is None:
            raise ValueError("num_moe_experts must be non None to use expert-parallel.")

        if self.transformer_impl == "inference_optimized" and self.num_moe_experts is not None:
            self.inference_grouped_gemm_backend = InferenceGroupedGemmBackend.from_config(
                self.inference_grouped_gemm_backend
            )

            mxfp8_enabled = bool(self.fp8) and self.fp8_recipe == Fp8Recipe.mxfp8
            if self.expert_tensor_parallel_size > 1:
                raise ValueError(
                    "Inference-optimized MoE layers does not support expert tensor parallelism."
                )
            if self.moe_expert_capacity_factor is not None:
                raise ValueError("Inference-optimized MoE layers only support dropless MoE ")
            if self.moe_router_padding_for_quantization:
                raise ValueError(
                    "Inference-optimized MoE layers do not support padded "
                    "routing map for quantization."
                )
            if self.moe_router_dtype != "fp32":
                raise ValueError(
                    "--transformer-impl='inference_optimized' requires --moe-router-dtype=fp32 "
                    "to avoid costly dtype conversions during decode."
                )

            # Gated linear units (SwiGLU/GeGLU) are supported by the torch and vLLM
            # grouped-GEMM backends.
            if self.gated_linear_unit and self.inference_grouped_gemm_backend not in (
                InferenceGroupedGemmBackend.TORCH,
                InferenceGroupedGemmBackend.VLLM,
            ):
                raise ValueError(
                    "--transformer-impl='inference_optimized' supports gated linear units "
                    "(SwiGLU/GeGLU) only with --inference-grouped-gemm-backend torch or vllm, "
                    f"got '{self.inference_grouped_gemm_backend}'."
                )

            if mxfp8_enabled:
                if not self.fp8_param:
                    raise ValueError(
                        "fp8_param must be enabled when using "
                        "--transformer-impl='inference_optimized' with --fp8-recipe='mxfp8'. "
                        "Please set --fp8-param-gather."
                    )

            if (
                self.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.FLASHINFER
                and mxfp8_enabled
                and (self.gated_linear_unit or self.activation_func != squared_relu)
            ):
                raise ValueError(
                    "FlashInfer routed MXFP8 MoE currently supports only non-gated "
                    "squared-ReLU experts. Set activation_func=squared_relu and "
                    "gated_linear_unit=False, or select inference_grouped_gemm_backend "
                    "'torch' or 'vllm'."
                )

            if self.inference_flashinfer_mxfp8_token_capacity is not None:
                if self.inference_flashinfer_mxfp8_token_capacity <= 0:
                    raise ValueError(
                        "inference_flashinfer_mxfp8_token_capacity must be > 0, got "
                        f"{self.inference_flashinfer_mxfp8_token_capacity}"
                    )
                if (
                    self.inference_grouped_gemm_backend != InferenceGroupedGemmBackend.FLASHINFER
                    or not mxfp8_enabled
                    or self.inference_moe_token_dispatcher_type != "nvls"
                    or self.expert_model_parallel_size <= 1
                ):
                    raise ValueError(
                        "inference_flashinfer_mxfp8_token_capacity requires "
                        "inference_grouped_gemm_backend='flashinfer', FP8 enabled with "
                        "fp8_recipe='mxfp8', "
                        "inference_moe_token_dispatcher_type='nvls' and "
                        "expert_model_parallel_size > 1"
                    )

            if self.batch_invariant_mode:
                if self.inference_grouped_gemm_backend not in (
                    InferenceGroupedGemmBackend.FLASHINFER,
                    InferenceGroupedGemmBackend.TORCH,
                    InferenceGroupedGemmBackend.VLLM,
                ):
                    raise ValueError(
                        "batch_invariant_mode requires inference_grouped_gemm_backend "
                        "'flashinfer', 'torch', or 'vllm'."
                    )
                if (
                    self.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.FLASHINFER
                    and not mxfp8_enabled
                ):
                    raise ValueError(
                        "batch_invariant_mode currently supports the FlashInfer grouped-GEMM "
                        "backend only for an MXFP8 model configuration. Selectively BF16 "
                        "expert layers within that configuration remain supported."
                    )
                if (
                    self.expert_model_parallel_size > 1
                    and self.inference_moe_token_dispatcher_type != "nvls"
                ):
                    raise ValueError(
                        "batch_invariant_mode with inference-optimized MoE and expert "
                        "parallelism requires inference_moe_token_dispatcher_type='nvls'."
                    )

        if self.num_moe_experts is not None and self.num_moe_experts <= 0:
            raise ValueError("num_moe_experts must be non-negative.")

        if self.num_moe_experts is not None and self.moe_ffn_hidden_size is None:
            self.moe_ffn_hidden_size = self.ffn_hidden_size
            warnings.warn("moe_ffn_hidden_size is not set, using ffn_hidden_size instead.")

        if self.num_moe_experts is None and self.moe_ffn_hidden_size is not None:
            is_mixed_model = (
                isinstance(self.moe_layer_freq, list) and 0 in self.moe_layer_freq
            ) or (isinstance(self.moe_layer_freq, int) and self.moe_layer_freq > 1)
            if is_mixed_model:
                warnings.warn(
                    "moe_ffn_hidden_size is set but num_moe_experts is None. "
                    "This is expected for dense layers in a mixed dense/MoE model "
                    "(moe_layer_freq indicates some layers remain dense). "
                    "moe_ffn_hidden_size will be ignored for this layer."
                )
                self.moe_ffn_hidden_size = None
            else:
                raise ValueError(
                    "moe_ffn_hidden_size is set but num_moe_experts is None. "
                    "Please set num_moe_experts or remove moe_ffn_hidden_size."
                )

        if self.moe_single_grouped_weight or self.moe_single_grouped_bias:
            if not self.moe_grouped_gemm:
                raise ValueError(
                    "moe_single_grouped_weight and moe_single_grouped_bias require "
                    "moe_grouped_gemm=True."
                )
            if not is_te_min_version("2.14.0"):
                raise ValueError(
                    "moe_single_grouped_weight and moe_single_grouped_bias require "
                    f"transformer-engine>=2.14.0, but your version is {get_te_version()}."
                )
        if self.moe_single_grouped_weight:
            # Single grouped weights are supported for high-precision primary weights
            # (BF16/FP16), MXFP8 primary weights, and NVFP4 primary weights.
            # Other quantized primary-weight paths need grouped partial-cast support
            # before they are safe to enable.
            if (self.fp8 and self.fp8_recipe != Fp8Recipe.mxfp8) or (
                self.fp4 and self.fp4_recipe != Fp4Recipe.nvfp4
            ):
                raise ValueError(
                    "moe_single_grouped_weight is currently supported with high-precision "
                    "primary weights, fp8_recipe='mxfp8', or fp4_recipe='nvfp4'."
                )
            if self.fp4 and not self.fp4_param:
                raise ValueError(
                    "moe_single_grouped_weight with FP4 compute requires fp4_param=True "
                    "(--fp4-param-gather). Without FP4 parameter gather, Transformer Engine "
                    "uses a split-quantize fallback that is being deprecated."
                )
            if not self.moe_use_grouped_tensor:
                raise ValueError("moe_single_grouped_weight requires moe_use_grouped_tensor=True.")
        if self.moe_single_grouped_bias and not self.add_bias_linear:
            raise ValueError("moe_single_grouped_bias requires add_bias_linear=True.")
        if self.moe_single_grouped_bias and not self.moe_use_grouped_tensor:
            raise ValueError("moe_single_grouped_bias requires moe_use_grouped_tensor=True.")

        if self.moe_enable_deepep:
            if self.moe_token_dispatcher_type != "flex":
                raise ValueError("DeepEP backend is only supported with flex token dispatcher.")
            if self.moe_flex_dispatcher_backend == "hybridep":
                raise ValueError("Only one backend is supported for flex token dispatcher.")
            self.moe_flex_dispatcher_backend = "deepep"
            warnings.warn(
                "moe_enable_deepep is deprecated."
                "Please use --moe-flex-dispatcher-backend=deepep instead."
            )

        if self.moe_token_dispatcher_type == "flex":
            if self.moe_pad_expert_input_to_capacity and (
                self.moe_enable_deepep or self.moe_flex_dispatcher_backend == "deepep"
            ):
                raise ValueError(
                    "Flex token dispatcher with deepep backend does not support "
                    "moe_pad_expert_input_to_capacity"
                )

        if self.moe_hybridep_routing_map_mode not in ("indices", "bool"):
            raise ValueError("moe_hybridep_routing_map_mode must be one of 'indices' or 'bool'.")
        if (
            self.moe_hybridep_routing_map_mode == "indices"
            and self.moe_pad_expert_input_to_capacity
            and self.moe_token_dispatcher_type == "flex"
            and self.moe_flex_dispatcher_backend == "hybridep"
        ):
            warnings.warn(
                "moe_hybridep_routing_map_mode='indices' is disabled by "
                "moe_pad_expert_input_to_capacity: a pad-to-capacity routing map can hold more "
                "than topk assignments per token, which dense top-k indices cannot represent. "
                "HybridEP will use the bool routing-map path."
            )

        if self.moe_flex_dispatcher_backend == "ncclep":
            if self.moe_token_dispatcher_type != "flex":
                raise ValueError(
                    "moe_flex_dispatcher_backend='ncclep' requires "
                    "moe_token_dispatcher_type='flex'."
                )

        if self.moe_dispatch_fwd_dtype != 'bf16' or self.moe_combine_bwd_dtype != 'bf16':
            if (
                self.moe_token_dispatcher_type != "flex"
                or self.moe_flex_dispatcher_backend != "ncclep"
            ):
                raise ValueError(
                    "moe_dispatch_fwd_dtype / moe_combine_bwd_dtype require the 'ncclep' flex "
                    "dispatcher backend."
                )
            if not (self.use_transformer_engine_op_fuser and self.moe_grouped_gemm):
                raise ValueError(
                    "moe_dispatch_fwd_dtype / moe_combine_bwd_dtype = 'mxfp8' require BOTH "
                    "use_transformer_engine_op_fuser and moe_grouped_gemm: only the fused "
                    "grouped GEMM path consumes the pre-quantized MXFP8 GroupedTensor payload."
                )

        if self.moe_use_norm_before_up_proj and self.moe_latent_size is None:
            raise ValueError("moe_use_norm_before_up_proj requires moe_latent_size to be set.")

        # moe_deepep_num_sms / moe_hybridep_num_sms are deprecated and unified into
        # moe_flex_dispatcher_num_sms. If either is set, route it (an explicit
        # moe_flex_dispatcher_num_sms takes precedence) and warn.
        _deprecated_num_sms = {
            name: getattr(self, name)
            for name in ("moe_deepep_num_sms", "moe_hybridep_num_sms")
            if getattr(self, name) is not None
        }
        if _deprecated_num_sms:
            warn_single_rank(
                f"{', '.join(_deprecated_num_sms)} is deprecated. "
                "Use moe_flex_dispatcher_num_sms instead."
            )
            if self.moe_flex_dispatcher_num_sms is None:
                if len(set(_deprecated_num_sms.values())) > 1:
                    raise ValueError(
                        f"Conflicting deprecated SM-count knobs {_deprecated_num_sms}; set a "
                        "single moe_flex_dispatcher_num_sms instead."
                    )
                self.moe_flex_dispatcher_num_sms = next(iter(_deprecated_num_sms.values()))
        shortcut_pre_norm_recompute = "shortcut_pre_mlp_layernorm" in (self.recompute_modules or [])
        shortcut_post_norm_offload = "shortcut_post_norm" in (self.offload_modules or [])
        if (shortcut_pre_norm_recompute or shortcut_post_norm_offload) and not (
            self.moe_shortcut_connection
        ):
            raise ValueError(
                "shortcut_pre_mlp_layernorm recompute and shortcut_post_norm offload require "
                "moe_shortcut_connection=True."
            )
        if shortcut_pre_norm_recompute and self.recompute_granularity != "selective":
            raise ValueError(
                "shortcut_pre_mlp_layernorm in recompute_modules requires "
                "recompute_granularity='selective'."
            )

        if self.moe_shortcut_connection:
            assert (
                self.num_moe_experts is not None and self.num_moe_experts > 0
            ), "moe_shortcut_connection requires MoE to be enabled (num_moe_experts > 0)"
            if self.recompute_granularity == 'full':
                raise ValueError(
                    "moe_shortcut_connection is not supported with full activation recomputation"
                )
            if self.moe_shared_expert_overlap:
                raise ValueError(
                    "moe_shortcut_connection is mutually exclusive with "
                    "moe_shared_expert_overlap. ScMoE computes shared experts inline."
                )

        if self.moe_shortcut_post_norm and not self.moe_shortcut_connection:
            raise ValueError("moe_shortcut_post_norm requires moe_shortcut_connection = True.")
        if shortcut_post_norm_offload and not self.moe_shortcut_post_norm:
            raise ValueError(
                "shortcut_post_norm in offload_modules requires " "moe_shortcut_post_norm = True."
            )

        if self.moe_shortcut_parallel:
            assert (
                self.moe_shortcut_connection
            ), "moe_shortcut_parallel requires moe_shortcut_connection = True"
            assert (
                self.num_moe_experts is not None and self.num_moe_experts > 0
            ), "moe_shortcut_parallel requires MoE to be enabled (num_moe_experts > 0)"

        if self.moe_shared_expert_intermediate_size is not None:
            if self.moe_shared_expert_intermediate_size <= 0:
                raise ValueError(
                    f"moe_shared_expert_intermediate_size must be "
                    f"num_shared_experts * ffn_size_of_each_shared_expert, "
                    f"but got {self.moe_shared_expert_intermediate_size}"
                )
            if self.moe_shared_expert_overlap and self.moe_token_dispatcher_type not in [
                "alltoall",
                "flex",
            ]:
                raise ValueError(
                    f"moe_shared_expert_overlap only works with alltoall or flex token dispatcher."
                )

        if isinstance(self.moe_router_load_balancing_type, list):
            assert isinstance(self.moe_aux_loss_coeff, list) and len(
                self.moe_aux_loss_coeff
            ) == len(self.moe_router_load_balancing_type), (
                "moe_aux_loss_coeff must be a list of the same length as "
                "moe_router_load_balancing_type"
            )

        if self.moe_expert_capacity_factor is not None:
            if self.moe_expert_capacity_factor < 0:
                self.moe_expert_capacity_factor = None
            if isinstance(self.moe_router_load_balancing_type, list):
                for load_balancing_type in self.moe_router_load_balancing_type:
                    if load_balancing_type not in [
                        "aux_loss",
                        "seq_aux_loss",
                        "global_aux_loss",
                        "none",
                    ]:
                        raise ValueError(
                            "moe_expert_capacity_factor only works with aux_loss, "
                            "seq_aux_loss, global_aux_loss or none load balancing"
                        )
            elif self.moe_router_load_balancing_type not in [
                "aux_loss",
                "seq_aux_loss",
                "global_aux_loss",
                "none",
            ]:
                raise ValueError(
                    "moe_expert_capacity_factor only works with aux_loss, "
                    "seq_aux_loss, global_aux_loss or none load balancing"
                )

        if self.moe_pad_expert_input_to_capacity:
            if self.moe_expert_capacity_factor is None:
                raise ValueError(
                    "moe_expert_capacity_factor must be set to use moe_pad_expert_input_to_capacity"
                )

        if self.moe_expert_rank_capacity_factor is not None:
            if self.moe_flex_dispatcher_backend not in ("hybridep", "ncclep"):
                raise ValueError(
                    "moe_expert_rank_capacity_factor requires moe_flex_dispatcher_backend to be "
                    "'hybridep' or 'ncclep'."
                )
            if (
                self.moe_flex_dispatcher_backend == "hybridep"
                and not self.use_transformer_engine_op_fuser
                and not self.moe_use_grouped_tensor
            ):
                raise ValueError(
                    "moe_expert_rank_capacity_factor with the 'hybridep' backend requires "
                    "use_transformer_engine_op_fuser=True or moe_use_grouped_tensor=True."
                )

        if self.cpu_offloading and (
            self.cpu_offloading_num_layers < 0 or self.cpu_offloading_num_layers >= self.num_layers
        ):
            raise ValueError(
                f"CPU offloading can be done only for layers less than {self.num_layers}"
            )

        if self.cpu_offloading and self.pipeline_model_parallel_size > 1:
            raise ValueError(
                "Currently there is no support for Pipeline parallelism with CPU offloading"
            )

        if self.cpu_offloading and self.recompute_granularity is not None:
            raise ValueError(
                "CPU offloading does not work when activation recomputation is enabled"
            )

        if self.recompute_granularity is not None:
            if self.recompute_granularity not in ["full", "selective"]:
                raise ValueError(
                    f'When using recompute_granuarlity: {self.recompute_granularity} must be "full"'
                    'or "selective".'
                )

            if self.recompute_method is not None:
                if self.recompute_method not in ["block", "uniform"]:
                    raise ValueError(
                        f'recompute_method: {self.recompute_method} must be "block" or "uniform".'
                    )
            elif self.recompute_granularity != "selective":
                raise ValueError(
                    f"Using recompute_granularity: {self.recompute_granularity} so "
                    'recompute_method must be "block" or "uniform"'
                )

            if self.recompute_granularity != "selective" and self.recompute_num_layers is None:
                raise ValueError(
                    f"When using recompute_granularity: {self.recompute_granularity} "
                    "recompute_num_layers must be between "
                    "1 and num_layers_per_pipeline_rank: "
                    f"{self.num_layers // self.pipeline_model_parallel_size}"
                )
            elif (
                self.recompute_granularity == "selective" and self.recompute_num_layers is not None
            ):
                raise ValueError(
                    f"When using recompute_granularity: {self.recompute_granularity} "
                    "recompute_num_layers must be None."
                )

            if self.distribute_saved_activations and self.sequence_parallel:
                raise ValueError(
                    f"distribute_saved_activations: {self.distribute_saved_activations} must be "
                    f"false when sequence parallel is enabled: {self.sequence_parallel}"
                )

        if self.recompute_modules is None:
            self.recompute_modules = ["core_attn"]

        if self.recompute_granularity == "selective":
            if len(self.recompute_modules) > 0:
                allowed_modules = {
                    "core_attn",
                    "moe_act",
                    "layernorm",
                    "mla_up_proj",
                    "mlp",
                    "moe",
                    "shared_experts",
                    "gdn_norm_out",
                    "gdp_in_proj",
                    "gdp_qkv",
                    "mhc",
                    "shortcut_pre_mlp_layernorm",
                    "residual_stream",
                }
                invalid_modules = set(self.recompute_modules) - allowed_modules
                assert not invalid_modules, (
                    f"Invalid choices for recompute_modules: {invalid_modules}. "
                    f"Allowed modules are: {allowed_modules}"
                )

            if "moe_act" in self.recompute_modules and not self.moe_grouped_gemm:
                raise ValueError(
                    "moe_act in recompute_modules is only supported with moe_grouped_gemm."
                )

            if "mla_up_proj" in self.recompute_modules and not self.multi_latent_attention:
                raise ValueError(
                    "mla_up_proj in recompute_modules is only supported with "
                    "multi_latent_attention."
                )

            if "gdn_norm_out" in self.recompute_modules and (
                not is_gated_delta_net_variant(self.experimental_attention_variant)
            ):
                raise ValueError(
                    "gdn_norm_out in recompute_modules is only supported with "
                    "experimental_attention_variant='gdn' or 'gdn2'."
                )

            if "core_attn" in self.recompute_modules:
                warnings.warn(
                    "If you are using transformer_engine as the transformer implementation, "
                    "the core_attn is from transformer_engine and may be the fused version. "
                    "For fused attention, you have no need to set 'core_attn' to recompute. "
                    "Please check that the core_attn recompute is really needed."
                )

            if "shared_experts" in self.recompute_modules:
                if (
                    self.moe_shared_expert_intermediate_size is not None
                    and self.moe_shared_expert_overlap
                ):
                    raise ValueError(
                        "shared_experts recompute cannot work with --moe-shared-expert-overlap."
                    )

            if self.fp8:
                fp8_output_discarding_modules = {
                    "moe_act",
                    "layernorm",
                    "shortcut_pre_mlp_layernorm",
                }
                if fp8_output_discarding_modules & set(self.recompute_modules):
                    if self.fp8_recipe == 'delayed':
                        raise ValueError(
                            "Delayed scaling does not support moe_act, layernorm, or "
                            "shortcut_pre_mlp_layernorm recompute for fp8."
                        )
                    if not is_te_min_version("2.6.0dev0"):
                        raise ValueError(
                            "moe_act, layernorm, and shortcut_pre_mlp_layernorm recompute for "
                            "fp8 need "
                            "transformer-engine>=2.6.0dev0, "
                            f"but your version is {get_te_version()}."
                        )

        if self.moe_layer_recompute:
            warnings.warn(
                "--moe-layer-recompute is deprecated. "
                "Use --recompute-granularity selective --recompute-modules moe_layer instead."
            )
            if self.recompute_granularity == "full":
                raise ValueError(
                    "Do not set --moe-layer-recompute with full recompute granularity. "
                )
            self.recompute_granularity = "selective"
            if "moe" not in self.recompute_modules:
                self.recompute_modules.append("moe")

        # Validation for "mhc" in recompute_modules
        if self.recompute_granularity == "selective" and "mhc" in self.recompute_modules:
            if not self.enable_mhc_connections:
                raise ValueError("'mhc' in recompute_modules requires enable_mhc_connections=True.")
            if "mlp" in self.recompute_modules:
                raise ValueError(
                    "'mhc' and 'mlp' in recompute_modules cannot be used together. "
                    "They use different checkpoint mechanisms that may conflict."
                )
            if self.mhc_recompute_layer_num is not None and (
                isinstance(self.mhc_recompute_layer_num, bool)
                or not isinstance(self.mhc_recompute_layer_num, int)
                or self.mhc_recompute_layer_num < 1
            ):
                raise ValueError(
                    "mhc_recompute_layer_num must be a positive integer when "
                    "'mhc' is in recompute_modules."
                )
            if self.fine_grained_activation_offloading:
                raise NotImplementedError(
                    "'mhc' in recompute_modules + fine_grained_activation_offloading is "
                    "not yet supported. The mHC recompute hook currently fires before "
                    "the offloading backward chunk is initialized, causing tensor_pop "
                    "on a None chunk. Disable one of them."
                )

        use_residual_stream_recompute = (
            self.recompute_granularity == "selective"
            and "residual_stream" in self.recompute_modules
        )
        if "residual_stream" in self.recompute_modules and not use_residual_stream_recompute:
            raise ValueError(
                "'residual_stream' in recompute_modules requires "
                "recompute_granularity='selective'."
            )
        if use_residual_stream_recompute:
            if self.wide_residual is None:
                raise ValueError(
                    "'residual_stream' recomputation requires a configured wide residual stream."
                )
            if self.residual_stream_recompute_num_layers is not None and (
                isinstance(self.residual_stream_recompute_num_layers, bool)
                or not isinstance(self.residual_stream_recompute_num_layers, int)
                or self.residual_stream_recompute_num_layers < 1
            ):
                raise ValueError(
                    "residual_stream_recompute_num_layers must be a positive integer or None."
                )
            if self.fine_grained_activation_offloading:
                replay_owned_norms = {"attn_norm", "mlp_norm"} & set(self.offload_modules or [])
                if replay_owned_norms:
                    warnings.warn(
                        "Residual-stream recomputation owns residual reads, connected-branch "
                        "norms, and writes. Fine-grained activation offloading will skip "
                        f"{sorted(replay_owned_norms)} only on connected branches."
                    )
            if self.cuda_graph_impl != "none":
                raise ValueError(
                    "'residual_stream' recomputation requires cuda_graph_impl='none' because "
                    "its Python checkpoint manager remains outside CUDA graph capture."
                )
        elif self.residual_stream_recompute_num_layers is not None:
            raise ValueError(
                "residual_stream_recompute_num_layers requires selective recomputation with "
                "'residual_stream' in recompute_modules."
            )

        if self.enable_mhc_connections and not (
            self.recompute_granularity == "selective" and "mhc" in self.recompute_modules
        ):
            warnings.warn(
                "HyperConnections are enabled but 'mhc' is not in "
                "recompute_modules with selective recompute. Consider adding 'mhc' to "
                "recompute_modules with selective recompute to reduce activation memory."
            )

        if self.use_fused_mhc and not self.enable_mhc_connections:
            raise ValueError("use_fused_mhc requires enable_mhc_connections=True.")

        valid_mhc_fused_backends = ("auto", "native", "triton", "cutile")
        if self.mhc_fused_backend not in valid_mhc_fused_backends:
            raise ValueError(
                f"Unknown mhc_fused_backend {self.mhc_fused_backend!r}; expected one of "
                f"{valid_mhc_fused_backends}."
            )
        if self.mhc_fused_backend != "auto" and not self.use_fused_mhc:
            raise ValueError("mhc_fused_backend requires use_fused_mhc=True when set explicitly.")

        if self.enable_mhc_connections and self.recompute_granularity == "full":
            raise NotImplementedError(
                "enable_mhc_connections is not yet compatible with full activation recompute. "
                "Use selective recompute with 'mhc' in recompute_modules, or disable "
                "activation recompute."
            )

        if self.enable_mhc_connections and self.inference_fuse_tp_communication:
            raise NotImplementedError(
                "enable_mhc_connections is not compatible with inference_fuse_tp_communication. "
                "The fused inference TP path assumes single-stream residual tensors."
            )

        # Note: mHC + MoE is deliberately NOT rejected here. HyperConnectionTransformerLayer
        # raises for a MoE MLP submodule at build time, which is the precise check; a
        # config-level `num_moe_experts` guard would also block the documented composition
        # (wrapping MoE as a HybridStack layer via HyperConnectionHybridLayer).

        if self.enable_mhc_connections:
            # TransformerBlock expands to n-stream at `pre_process` and contracts back at
            # the stage holding the final layernorm, so every intermediate pipeline stage
            # exchanges [s, b, n*C] while the p2p buffers are still sized from hidden_size.
            # Pipeline support must resize the p2p buffers before this guard can be lifted.
            if self.pipeline_model_parallel_size > 1:
                raise NotImplementedError(
                    "enable_mhc_connections does not support pipeline_model_parallel_size > 1 "
                    "yet. Inter-stage activations are n-stream ([s, b, n*C]) while pipeline "
                    "p2p buffers are sized from hidden_size, so the shapes disagree. Use "
                    "pipeline_model_parallel_size=1 until mHC pipeline support lands."
                )

            # The residual carried across an mHC layer is the n-stream tensor consumed by
            # `fused_h_res_h_post_bda` (a bmm against h_res), not the single-stream residual
            # the base layer casts. Upcasting it alone would mismatch the h_res dtype.
            if self.fp32_residual_connection:
                raise NotImplementedError(
                    "enable_mhc_connections is not compatible with fp32_residual_connection. "
                    "The mHC residual is the n-stream tensor consumed by the H_res batched "
                    "matmul, which must share the activation dtype."
                )

            # SinkhornKnopp.backward's analytic chain rule (grad = grad_M_init * M_init)
            # is exact only because the per-row max shift is annihilated by the first
            # row-normalization T_r. With zero iterations that invariance does not hold
            # and the dropped argmax term makes the gradient silently wrong.
            if self.mhc_sinkhorn_iterations < 1:
                raise ValueError(
                    "mhc_sinkhorn_iterations must be >= 1; the Sinkhorn-Knopp backward "
                    f"assumes at least one row-normalization pass, got "
                    f"{self.mhc_sinkhorn_iterations}."
                )

            if self.mhc_init_gating_factor < 0:
                raise ValueError(
                    "mhc_init_gating_factor must be non-negative, got "
                    f"{self.mhc_init_gating_factor}."
                )

        if self.fine_grained_activation_offloading:
            assert (
                not self.cpu_offloading
            ), "fine_grained_activation_offloading cannot be enabled with cpu_offloading."
            assert self.offload_modules is not None and len(self.offload_modules) > 0
            allowed_modules = {
                "core_attn",
                "attn_proj",
                "expert_fc1",
                "fused_group_mlp",
                "moe_act",
                "attn_norm",
                "mlp_norm",
                "qkv_linear",
                "gdp_qkv",
                "shortcut_post_norm",
            }
            invalid_modules = set(self.offload_modules) - allowed_modules
            assert not invalid_modules, (
                f'Invalid choices for offload_modules: {invalid_modules}. '
                f'Allowed modules are: {allowed_modules}'
            )
            if "attn_proj" in self.offload_modules and "core_attn" not in self.offload_modules:
                raise ValueError(
                    "attn_proj cannot be set to offload_modules alone without core_attn "
                    "because the input of attn_proj is the output of core_attn, "
                    "which is needed in core_attn.backward()."
                )
            if (
                "gdp_qkv" in self.offload_modules
                and self.recompute_granularity == "selective"
                and "gdp_in_proj" in self.recompute_modules
            ):
                raise ValueError(
                    "gdp_qkv cannot be set in offload_modules together with gdp_in_proj in "
                    "recompute_modules, because gdp_in_proj discards the input of the causal "
                    "conv and rematerializes it at the start of the mixer backward, leaving "
                    "nothing for the gdp_qkv offload group to keep on the host."
                )
            if self.recompute_granularity == "selective" and "moe" in self.recompute_modules:
                offload_inside_moe = {"moe_act", "expert_fc1", "fused_group_mlp"} & set(
                    self.offload_modules
                )
                assert not offload_inside_moe, (
                    f"Cannot offload {offload_inside_moe} while recomputing the entire MoE layer. "
                    f"'moe' in recompute_modules wraps the full MoE forward in a checkpoint, "
                    f"so offloading activations inside it is redundant and will cause errors. "
                    f"Either remove 'moe' from --recompute-modules or remove "
                    f"{offload_inside_moe} from --offload-modules."
                )
            assert (
                self.min_offloaded_tensor_size >= 0
            ), "min_offloaded_tensor_size must be non-negative."
            assert (
                self.activation_offload_fraction >= 0 and self.activation_offload_fraction <= 1
            ), "activation_offload_fraction must be in range [0, 1]."
            assert (
                self.delta_offload_bytes_across_pp_ranks >= 0
            ), "delta_offload_bytes_across_pp_ranks must be non-negative."

            if "fused_group_mlp" in self.offload_modules:
                if not self.use_transformer_engine_op_fuser:
                    raise ValueError("fused_group_mlp requires use_transformer_engine_op_fuser.")
                moe_partial_offload = {"expert_fc1", "moe_act"} & set(self.offload_modules)
                if moe_partial_offload:
                    raise ValueError(
                        "fused_group_mlp offloads the whole fused grouped MLP and cannot be "
                        f"combined with expert_fc1 or moe_act. Remove: {moe_partial_offload}"
                    )

        if self.gdp_cutedsl_kernel:
            assert 0 <= self.gdp_num_chunk_states_to_recompute <= 64, (
                "gdp_num_chunk_states_to_recompute must be in range [0, 64], got "
                f"{self.gdp_num_chunk_states_to_recompute}."
            )

        if self.gtp_remat_opt_in_modules:
            _allowed_gtp_remat_opt_in_modules = {"moe_latent_proj"}
            invalid = set(self.gtp_remat_opt_in_modules) - _allowed_gtp_remat_opt_in_modules
            assert not invalid, (
                f"Invalid choices for gtp_remat_opt_in_modules: {invalid}. "
                f"Allowed modules are: {_allowed_gtp_remat_opt_in_modules}"
            )

        if self.moe_paged_stash:
            if self.cpu_offloading:
                raise ValueError("moe_paged_stash cannot be enabled with cpu_offloading.")
            if self.moe_expert_rank_capacity_factor is None:
                raise ValueError(
                    "moe_paged_stash requires moe_expert_rank_capacity_factor to be set; "
                    "there is no need to use paged stashing without it."
                )
            moe_offload_conflict = {"expert_fc1", "moe_act", "fused_group_mlp"} & set(
                self.offload_modules
            )
            if moe_offload_conflict:
                raise ValueError(
                    "When moe_paged_stash is enabled, offload_modules must not include "
                    f"expert_fc1, moe_act, or fused_group_mlp "
                    f"(paged stash covers those activations). "
                    f"Remove: {moe_offload_conflict}"
                )

        if (
            self.num_layers_in_first_pipeline_stage is not None
            or self.num_layers_in_last_pipeline_stage is not None
        ) and (
            self.account_for_embedding_in_pipeline_split or self.account_for_loss_in_pipeline_split
        ):
            raise ValueError(
                "num_layers_in_first_pipeline_stage and num_layers_in_last_pipeline_stage cannot be"
                "set at the same time with account_for_embedding_in_pipeline_split"
                "and account_for_loss_in_pipeline_split"
            )

        # PP layout
        if self.pipeline_model_parallel_layout is not None:
            # If pipeline layout is set, we will check the conflicts
            # with other pipeline layout arguments.
            any_conflict = (
                self.num_layers_in_first_pipeline_stage is not None
                or self.num_layers_in_last_pipeline_stage is not None
                or self.account_for_embedding_in_pipeline_split
                or self.account_for_loss_in_pipeline_split
            )
            if any_conflict:
                raise ValueError(
                    "pipeline_model_parallel_layout cannot be set"
                    " with other pipeline layout arguments."
                    f" {self.num_layers_in_first_pipeline_stage=},"
                    f" {self.num_layers_in_last_pipeline_stage=},"
                    f" {self.account_for_embedding_in_pipeline_split=},"
                    f" {self.account_for_loss_in_pipeline_split=}."
                )

            # Transfer pipeline_model_parallel_layout from str or list to
            # PipelineParallelLayerLayout
            if isinstance(self.pipeline_model_parallel_layout, str):
                self.pipeline_model_parallel_layout = PipelineParallelLayerLayout.from_str(
                    layout=self.pipeline_model_parallel_layout,
                    pipeline_model_parallel_size=self.pipeline_model_parallel_size,
                )
            elif isinstance(self.pipeline_model_parallel_layout, list):
                # Since list is not hashable, the initialization will not be cached.
                self.pipeline_model_parallel_layout = PipelineParallelLayerLayout(
                    layout=self.pipeline_model_parallel_layout,
                    pipeline_model_parallel_size=self.pipeline_model_parallel_size,
                )

            # Check whether the input VPP size conflicts with the PP layout
            detected_vpp_size = (
                self.pipeline_model_parallel_layout.virtual_pipeline_model_parallel_size
            )
            if self.virtual_pipeline_model_parallel_size is not None:
                assert self.virtual_pipeline_model_parallel_size == detected_vpp_size, (
                    f"virtual_pipeline_model_parallel_size conflicts with"
                    f" pipeline_model_parallel_layout,"
                    f" ({self.virtual_pipeline_model_parallel_size=}, "
                    f" {detected_vpp_size=})"
                )
            elif detected_vpp_size > 1:
                self.virtual_pipeline_model_parallel_size = detected_vpp_size

            # Check whether the layout is valid.
            self.mtp_standalone = self.pipeline_model_parallel_layout.validate_layer_layout(
                num_layers=self.num_layers, mtp_num_layers=self.mtp_num_layers
            )

        # Uneven PP
        elif (
            self.num_layers_in_first_pipeline_stage is not None
            or self.num_layers_in_last_pipeline_stage is not None
        ):
            pipeline_parallel_size = self.pipeline_model_parallel_size
            num_layers = self.num_layers

            if self.num_layers_in_first_pipeline_stage is not None:
                if self.num_layers_in_first_pipeline_stage <= 0:
                    raise ValueError("num_layers_in_first_pipeline_stage must be larger than 0")

                if self.virtual_pipeline_model_parallel_size is not None:
                    if (
                        self.num_layers_in_first_pipeline_stage
                        % self.virtual_pipeline_model_parallel_size
                        != 0
                    ):
                        raise ValueError(
                            f"number of layers at first stage: "
                            f"{self.num_layers_in_first_pipeline_stage}"
                            f"must be divisible by virtual pipeline"
                            f"parallel degree {self.virtual_pipeline_model_parallel_size}"
                        )
                num_layers -= self.num_layers_in_first_pipeline_stage
                pipeline_parallel_size -= 1

            if self.num_layers_in_last_pipeline_stage is not None:
                if self.num_layers_in_last_pipeline_stage <= 0:
                    raise ValueError("num_layers_in_last_pipeline_stage must be larger than 0")

                if self.virtual_pipeline_model_parallel_size is not None:
                    if (
                        self.num_layers_in_last_pipeline_stage
                        % self.virtual_pipeline_model_parallel_size
                        != 0
                    ):
                        raise ValueError(
                            f"number of layers at last stage: "
                            f"{self.num_layers_in_last_pipeline_stage}"
                            f"must be divisible by virtual pipeline"
                            f"parallel degree {self.virtual_pipeline_model_parallel_size}"
                        )
                num_layers -= self.num_layers_in_last_pipeline_stage
                pipeline_parallel_size -= 1

            # Ensure you either have middle pp stages and layers or none of them.
            if bool(num_layers) != bool(pipeline_parallel_size):
                raise ValueError(
                    f"Mismatch: {num_layers} middle layers remaining but {pipeline_parallel_size} "
                    f"middle PP stages available."
                )

            # Here pipeline_parallel_size is the number of middle PP stages. If there are middle
            # PP stages, check number of layers at middle stage is divisible by middle PP size.
            if pipeline_parallel_size and not num_layers % pipeline_parallel_size == 0:
                raise ValueError(
                    f"number of layers at middle stage: {num_layers} must be divisible by"
                    f"the middle pipeline model parallel size {pipeline_parallel_size}"
                )

            # If there are middle PP stages, check number of layers
            # on each middle PP rank is divisible by VPP size.
            if pipeline_parallel_size and self.virtual_pipeline_model_parallel_size is not None:
                num_layers_per_middle_pipeline_rank = num_layers // pipeline_parallel_size
                if (
                    not num_layers_per_middle_pipeline_rank
                    % self.virtual_pipeline_model_parallel_size
                    == 0
                ):
                    raise ValueError(
                        f"number of layers on each middle pipeline rank:"
                        f"{num_layers_per_middle_pipeline_rank} must be divisible by virtual"
                        f"pipeline parallel degree {self.virtual_pipeline_model_parallel_size}"
                    )

        elif (
            self.account_for_embedding_in_pipeline_split or self.account_for_loss_in_pipeline_split
        ):
            if self.virtual_pipeline_model_parallel_size is None:
                num_layers = self.num_layers

                if self.account_for_embedding_in_pipeline_split:
                    num_layers += 1

                if self.account_for_loss_in_pipeline_split:
                    num_layers += 1

                if not num_layers % self.pipeline_model_parallel_size == 0:
                    raise ValueError(
                        f"number of middle layers: {num_layers} must be divisible by "
                        f"middle pipeline_model_parallel_size {self.pipeline_model_parallel_size}"
                    )
            else:
                num_layers = self.num_layers
                if self.account_for_embedding_in_pipeline_split:
                    num_layers += 1

                if self.account_for_loss_in_pipeline_split:
                    num_layers += 1

                if not num_layers % self.pipeline_model_parallel_size == 0:
                    raise ValueError(
                        f"num_layers: {num_layers} after enable"
                        f"account_for_embedding_in_pipeline_split or "
                        f"account_for_loss_in_pipeline_split must be divisible"
                        f"by pipeline_model_parallel_size "
                        f"{self.pipeline_model_parallel_size}"
                    )

                num_layers_per_pipeline_rank = num_layers // self.pipeline_model_parallel_size
                if (
                    not num_layers_per_pipeline_rank % self.virtual_pipeline_model_parallel_size
                    == 0
                ):
                    raise ValueError(
                        f"number of layers on each pipeline rank: {num_layers_per_pipeline_rank}"
                        f"(after enable account_for_embedding_in_pipeline_split or "
                        f"account_for_loss_in_pipeline_split) must be divisible by"
                        f"virtual_pipeline_model_parallel_size"
                        f"{self.virtual_pipeline_model_parallel_size}"
                    )

        if self.apply_query_key_layer_scaling:
            self.attention_softmax_in_fp32 = True

        if self.bias_activation_fusion:
            if self.activation_func not in [F.gelu, F.silu, quick_gelu]:
                raise ValueError(
                    "When bias_activation_fusion is True, activation function should be either "
                    "gelu, swiglu, or quick_geglu"
                )
            if (
                self.activation_func == F.gelu
                and not self.gated_linear_unit
                and not self.add_bias_linear
            ):
                raise ValueError(
                    "When bias_activation_fusion is True, gated_linear_unit is False "
                    "and activation function is gelu, add_bias_linear must also be True."
                )
            if self.activation_func == quick_gelu and not self.gated_linear_unit:
                raise ValueError(
                    "When bias_activation_fusion is True and activation function is quick_gelu, "
                    "gated_linear_unit must be True."
                )
            if self.glu_linear_offset != 0.0 and self.activation_func != quick_gelu:
                raise ValueError(
                    "When bias_activation_fusion is True and glu_linear_offset is non-zero, "
                    "activation function must be quick_gelu."
                )

            if self.use_te_activation_func:
                raise ValueError(
                    "bias_activation_fusion and use_te_activation_func cannot be both true. "
                    "If you use bias in MLP FC1, we recommend setting bias_activation_fusion "
                    "to True and use_te_activation_func to False."
                )

        if self.fused_residual_rmsnorm:
            if self.normalization != "RMSNorm":
                raise ValueError(
                    "fused_residual_rmsnorm is only supported when normalization is RMSNorm."
                )

        if self.use_te_activation_func:
            if self.activation_func not in (F.gelu, F.silu, F.relu):
                raise ValueError(
                    "TransformerEngine only support gelu, geglu, silu, swiglu, relu, reglu. "
                    "If you don't want to use TransformerEngine activation function, set "
                    "use_te_activation_func to False"
                )

        if self.activation_func_tanh_clamp_scale is not None:
            if self.activation_func_tanh_clamp_scale <= 0.0:
                raise ValueError(
                    "activation_func_tanh_clamp_scale must be positive, got "
                    f"{self.activation_func_tanh_clamp_scale}."
                )
            if self.use_te_activation_func:
                raise ValueError(
                    "activation_func_tanh_clamp_scale is not supported with "
                    "use_te_activation_func, since the TE activation modules cannot clamp."
                )
            if self.gated_linear_unit and self.activation_func != F.silu:
                raise ValueError(
                    "activation_func_tanh_clamp_scale with a gated activation is implemented as "
                    "SiTU-GLU, which replaces the swish gate, so it requires silu."
                )
            if self.bias_activation_fusion and not (
                self.activation_func == F.silu and self.gated_linear_unit
            ):
                raise ValueError(
                    "activation_func_tanh_clamp_scale with bias_activation_fusion is only "
                    "implemented for SwiGLU. The fused gelu, geglu and quick_geglu kernels do not "
                    "apply the clamp, so set bias_activation_fusion to False."
                )
            if self.activation_func_clamp_value is not None:
                raise ValueError(
                    "activation_func_tanh_clamp_scale and activation_func_clamp_value both clamp "
                    "the activation input; set only one of them."
                )

        if self.activation_func_tanh_clamp_scale_linear is not None:
            if self.activation_func_tanh_clamp_scale_linear <= 0.0:
                raise ValueError(
                    "activation_func_tanh_clamp_scale_linear must be positive, got "
                    f"{self.activation_func_tanh_clamp_scale_linear}."
                )
            if self.activation_func_tanh_clamp_scale is None:
                raise ValueError(
                    "activation_func_tanh_clamp_scale_linear only clamps the linear half of "
                    "SiTU-GLU, so it requires activation_func_tanh_clamp_scale for the gate half. "
                    "Clamping the linear half alone would leave the gate unbounded."
                )
            if not self.gated_linear_unit:
                raise ValueError(
                    "activation_func_tanh_clamp_scale_linear requires gated_linear_unit."
                )

        if self.activation_func_fp8_input_store:
            if self.activation_func != F.silu or not self.gated_linear_unit:
                raise ValueError("Storing activation input in FP8 is supported only for SwiGLU.")

        if (
            self.activation_func_clamp_value is not None
            and self.activation_func == F.silu
            and self.gated_linear_unit
        ):
            if (
                not math.isfinite(self.activation_func_clamp_value)
                or self.activation_func_clamp_value <= 0
            ):
                raise ValueError(
                    "activation_func_clamp_value for SwiGLU must be finite and greater than zero."
                )
            if self.num_moe_experts is None:
                raise ValueError(
                    "activation_func_clamp_value for SwiGLU is only supported with MoE."
                )
            if self.glu_linear_offset != 0.0:
                raise ValueError(
                    "glu_linear_offset must be zero when activation_func_clamp_value "
                    "is set for SwiGLU."
                )
            if self.use_te_activation_func:
                raise ValueError(
                    "use_te_activation_func must be False "
                    "when activation_func_clamp_value is not None for SwiGLU"
                )

        if self.apply_rope_fusion:
            if self.multi_latent_attention:
                warnings.warn(
                    "apply_rope_fusion for multi-latent attention only supports training. "
                    "It is experimental and may change in future versions."
                )
            else:
                if self.rotary_interleaved:
                    if not is_te_min_version("2.3.0"):
                        raise ValueError(
                            "rotary_interleaved does not work with apply_rope_fusion for "
                            "TE < 2.3.0. Please install TE >= 2.3.0"
                        )

                from megatron.core.models.common.embeddings.rope_utils import (
                    fused_apply_rotary_pos_emb,
                    fused_apply_rotary_pos_emb_thd,
                )

                if fused_apply_rotary_pos_emb is None and fused_apply_rotary_pos_emb_thd is None:
                    raise ValueError(
                        "apply_rope_fusion is not available. Please install TE >= 1.4."
                    )

        if self.fused_single_qkv_rope:
            if self.attention_output_gate:
                raise ValueError("fused_single_qkv_rope does not support gated attention for now.")

        if self.multi_latent_attention and self.rotary_interleaved:
            raise ValueError("rotary_interleaved does not work with multi_latent_attention.")

        # MuP (Maximal Update Parameterization) configuration
        if self.use_mup:
            # Default base_hidden_size to hidden_size (base model case, width_mult=1.0)
            if self.mup_base_hidden_size is None:
                self.mup_base_hidden_size = self.hidden_size
            assert self.mup_base_hidden_size > 0, "--mup-base-hidden-size must be positive."
            # Compute width multiplier
            self.mup_width_mult = self.hidden_size / self.mup_base_hidden_size

            # MuP attention scaling: 1/d_head instead of 1/sqrt(d_head).
            if self.softmax_scale is None:
                base_head_scale = (
                    1.0 if self.mup_base_head_dim is None else self.mup_base_head_dim**0.5
                )
                self.softmax_scale = base_head_scale / (self.kv_channels**self.mup_attn_scale_power)

            # MuP output scaling: scale logits by 1/width_mult to keep outputs O(1).
            # Only auto-set if user hasn't explicitly configured it.
            if self.mup_output_mult == 1.0 and self.mup_width_mult != 1.0:
                self.mup_output_mult = 1.0 / self.mup_width_mult

            overridden_init_methods = []
            if self.init_method is not None:
                overridden_init_methods.append("init_method")
            if self.output_layer_init_method is not None:
                overridden_init_methods.append("output_layer_init_method")
            if overridden_init_methods:
                overridden_init_methods_text = " and ".join(overridden_init_methods)
                verb = "is" if len(overridden_init_methods) == 1 else "are"
                warnings.warn(
                    "use_mup is enabled, but custom "
                    + overridden_init_methods_text
                    + f" {verb} set. This may break MuP initialization assumptions.",
                    UserWarning,
                )

        # Set the embedding init method.
        # NOTE: This block must run AFTER the MuP block above but BEFORE the init_method
        # block below. When MuP is enabled and init_method is None (the common case),
        # embedding_init_method is set here using the unscaled init_method_std, while
        # init_method (set below) gets MuP width-scaling. This ordering ensures embeddings
        # use the base (unscaled) initialization as required by MuP.
        if self.embedding_init_method_std is None:
            # By default, use the same init std as you use for every other non-output layer.
            self.embedding_init_method_std = self.init_method_std

        if self.embedding_init_method is None:
            if self.init_method is None or (self.embedding_init_method_std != self.init_method_std):
                # In this case, we set both the init method and the embedding init method to
                #  whatever std value requested (or defaulted) for the embedding_init_layer
                self.embedding_init_method = init_method_normal(self.embedding_init_method_std)
            else:
                # Replicate the current behavior where if you are not changing the std of the
                #  embedding init differently and the init method is set, we fallback to the
                #  init method for this layer. Since we are here after an OR we know that
                #  init_method is not None
                self.embedding_init_method = self.init_method

        if self.init_method is None:
            if self.use_mup:
                # MuP: scale std by 1/sqrt(width_mult).
                self.init_method = init_method_normal(
                    self.init_method_std / math.sqrt(self.mup_width_mult)
                )
            else:
                self.init_method = init_method_normal(self.init_method_std)

        if self.output_layer_init_method is None:
            if self.use_mup:
                # MuP: depth and width scaling for output layers.
                self.output_layer_init_method = mup_scaled_init_method_normal(
                    self.init_method_std,
                    self.num_layers,
                    self.mup_width_mult,
                    multiplier=2.0 if not self.is_hybrid_model else 1.0,
                )
            else:
                self.output_layer_init_method = scaled_init_method_normal(
                    self.init_method_std,
                    self.num_layers,
                    multiplier=2.0 if not self.is_hybrid_model else 1.0,
                )

        if self.num_moe_experts is not None and self.add_bias_linear:
            assert (
                self.expert_tensor_parallel_size == 1
            ), "Bias in Moe is only supported when ETP==1"

        if self.moe_router_enable_expert_bias and self.moe_router_score_function not in (
            "sigmoid",
            "sqrtsoftplus",
        ):
            raise ValueError(
                "Expert bias for aux-loss-free routing only supports 'sigmoid' and 'sqrtsoftplus' "
                "score functions. Please set --moe-router-score-function to 'sigmoid' or "
                "'sqrtsoftplus', or unset --moe-router-enable-expert-bias."
            )

        if self.moe_num_hash_layers > 0:
            if self.moe_shortcut_connection:
                raise ValueError(
                    "ShortcutMoE does not yet forward the token IDs required for hash MoE routing."
                )
            assert (
                self.hash_moe_vocab_size is not None and self.hash_moe_vocab_size > 0
            ), "hash_moe_vocab_size must be positive when moe_num_hash_layers > 0."
            assert (
                self.num_moe_experts is not None
            ), "num_moe_experts must be set when moe_num_hash_layers > 0."
            if not 1 <= self.moe_router_topk <= self.num_moe_experts:
                raise ValueError("Hash MoE requires 1 <= moe_router_topk <= num_moe_experts.")
            if self.pipeline_model_parallel_size > 1 and not self.is_hybrid_model:
                assert self.pipeline_model_parallel_layout is not None, (
                    "pipeline_model_parallel_layout must be set when using hash MoE "
                    "layers with pipeline parallelism (PP > 1)."
                )
                embedding_stage = self.pipeline_model_parallel_layout.layout[0][0]
                n_decoders_with_embedding = embedding_stage.count(LayerType.decoder)
                assert self.moe_num_hash_layers <= n_decoders_with_embedding, (
                    "All hash MoE layers must currently share the virtual pipeline stage "
                    "that owns the embedding. The embedding stage has "
                    f"{n_decoders_with_embedding} decoder layers, but "
                    f"moe_num_hash_layers={self.moe_num_hash_layers}."
                )
            assert (
                not self.overlap_moe_expert_parallel_comm
            ), "overlap_moe_expert_parallel_comm does not support hash MoE layers yet."

        if self.num_moe_experts and self.fp8:
            # TE version below 1.7.0 will raise Error when handle zeros tokens for expert
            if not is_te_min_version("1.7.0.dev0"):
                raise ValueError(
                    "Only transformer-engine>=1.7.0 supports MoE FP8 training, "
                    f"but your version is {get_te_version()}."
                )

            if self.moe_grouped_gemm and not is_te_min_version("1.11.0"):
                raise ValueError(
                    "Only transformer-engine>=1.11.0 supports FP8 grouped gemm, "
                    f"but your version is {get_te_version()}."
                )

        if self.moe_router_padding_for_fp8:
            # enable moe_router_padding_for_quantization
            warnings.warn(
                "--moe-router-padding-for-fp8 is going to be deprecated. "
                "Use --moe-router-padding-for-quantization instead."
            )
            self.moe_router_padding_for_quantization = True

        if self.moe_router_padding_for_quantization:
            if self.fp8 is None and self.fp4 is None:
                raise ValueError(
                    "fp8/fp4 must be specified when moe_router_padding_for_quantization is True."
                )

            if self.moe_token_dispatcher_type in ["allgather", "alltoall_seq"]:
                raise ValueError(
                    "allgather and alltoall_seq dispatcher does not support "
                    "moe_router_padding_for_quantization."
                )

        if (
            self.moe_router_topk == 1
            and self.moe_router_score_function == "softmax"
            and not self.moe_router_pre_softmax
            and self.moe_router_load_balancing_type != "sinkhorn"
        ):
            # Requires applying softmax before selecting the top-k when k is 1,
            # since softmax on a [num_tokens, 1] would yield a zero gradient.
            raise ValueError("Please use --moe-router-pre-softmax when topk is 1.")

        if self.moe_router_group_topk:
            if self.moe_router_topk_limited_devices:
                raise ValueError(
                    "moe_router_topk_limited_devices is deprecated and replaced by "
                    "moe_router_group_topk and moe_router_num_groups."
                )
            if not self.moe_router_num_groups:
                raise ValueError(
                    "When using group limited routing, moe_router_num_groups must be specified."
                )
            else:
                assert self.num_moe_experts % self.moe_router_num_groups == 0, (
                    f"num_moe_experts ({self.num_moe_experts}) should be divisible by "
                    f"moe_router_num_groups ({self.moe_router_num_groups})."
                )
                assert self.moe_router_group_topk <= self.moe_router_num_groups, (
                    f"moe_router_group_topk ({self.moe_router_group_topk}) should be smaller than "
                    f"moe_router_num_groups ({self.moe_router_num_groups})."
                )
        elif self.moe_router_topk_limited_devices:
            warnings.warn(
                "moe_router_topk_limited_devices is deprecated. Use moe_router_group_topk and "
                "moe_router_num_groups instead."
            )
            self.moe_router_group_topk = self.moe_router_topk_limited_devices
            self.moe_router_num_groups = self.expert_model_parallel_size

        if self.enable_cuda_graph or self.external_cuda_graph:
            assert (
                self.cuda_graph_impl == "none"
            ), "Do not use enable_cuda_graph or external_cuda_graph with cuda_graph_impl."
            assert (
                not self.enable_cuda_graph or not self.external_cuda_graph
            ), "enable_cuda_graph and external_cuda_graph cannot be enabled at the same time."

            if self.enable_cuda_graph:
                warnings.warn('enable_cuda_graph is deprecated, use cuda_graph_impl=local instead.')
                self.cuda_graph_impl = "local"
            if self.external_cuda_graph:
                warnings.warn(
                    'external_cuda_graph is deprecated, '
                    'use cuda_graph_impl=transformer_engine instead.'
                )
                self.cuda_graph_impl = "transformer_engine"

        if self.cuda_graph_scope is not None:
            assert self.cuda_graph_modules in (
                "full",
                None,
                [],
            ), "cuda_graph_scope and cuda_graph_modules cannot be set together."
            warnings.warn(
                "cuda_graph_scope is deprecated, use cuda_graph_modules instead.",
                DeprecationWarning,
                stacklevel=2,
            )

            # CudaGraphScope is preserved as a standalone class (not an alias of CudaGraphModule)
            # so pre-refactor checkpoint values deserialize with the correct member identity.
            # Convert to string names here so normalize_cuda_graph_modules handles them uniformly.
            def _scope_to_str(s):
                return s.name if isinstance(s, CudaGraphScope) else s

            scope = self.cuda_graph_scope
            if isinstance(scope, list):
                self.cuda_graph_modules = [_scope_to_str(s) for s in scope]
            else:
                self.cuda_graph_modules = _scope_to_str(scope)
            self.cuda_graph_scope = None

        normalized_scopes, deprecated_scopes, used_full_scope = normalize_cuda_graph_modules(
            self.cuda_graph_modules
        )
        validate_deprecated_cuda_graph_modules_migration_inputs(
            deprecated_scopes, self.cuda_graph_impl, self.inference_cuda_graph_scope
        )
        if used_full_scope:
            warnings.warn(
                "full scope is deprecated. "
                "Use empty cuda_graph_modules to capture the whole layer."
            )
        for scope, attr, value in deprecated_scopes:
            migration = get_deprecated_cuda_graph_modules_migration(
                scope, attr, value, self.cuda_graph_impl
            )
            if migration is None:
                warnings.warn(
                    f"cuda_graph_modules '{scope}' is deprecated and has no effect when "
                    "cuda_graph_impl='none'. Use cuda_graph_impl='local' with "
                    "inference_cuda_graph_scope='block' to enable inference CUDA graphs.",
                    DeprecationWarning,
                    stacklevel=2,
                )
                continue
            migration_attr, migration_value = migration
            warnings.warn(
                f"cuda_graph_modules '{scope}' is deprecated. "
                f"Use {migration_attr}={migration_value!r} instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            setattr(self, migration_attr, migration_value)
        self.cuda_graph_modules = normalized_scopes
        assert all(
            isinstance(scope, CudaGraphModule) for scope in self.cuda_graph_modules
        ), f"cuda_graph_modules must be a list of CudaGraphModule, got {self.cuda_graph_modules}."

        assert self.cuda_graph_impl in [
            "none",
            "transformer_engine",
            "local",
            "full_iteration",
        ], f"Invalid cuda graph implementation: {self.cuda_graph_impl}"

        self.inference_cuda_graph_scope = normalize_inference_cuda_graph_scope(
            self.inference_cuda_graph_scope, self.cuda_graph_impl
        )

        assert self.inference_cuda_graph_scope in ALLOWED_INFERENCE_SCOPES[self.cuda_graph_impl], (
            "Invalid inference CUDA graph scope "
            f"{self.inference_cuda_graph_scope.name!r} for cuda_graph_impl="
            f"{self.cuda_graph_impl!r}."
        )
        assert not (
            self.cuda_graph_impl == "full_iteration" and self.cuda_graph_modules
        ), 'cuda_graph_modules must be empty when cuda_graph_impl="full_iteration".'

        assert not (
            self.moe_shortcut_connection and self.cuda_graph_impl != "none"
        ), "CUDA graphs are not supported with moe_shortcut_connection."

        if self.cuda_graph_impl != "none":

            if self.cpu_offloading and self.cuda_graph_impl != "full_iteration":
                raise ValueError("CUDA graphs not supported with CPU offloading.")

            # Check cuda graph scopes for per-layer implementations.
            if self.cuda_graph_impl in ("local", "transformer_engine"):
                if self.cuda_graph_impl == "local":
                    # local impl doesn't currently distinguish between moe_preprocess or moe_router
                    # so just set both if either is specified.
                    if (
                        CudaGraphModule.moe_router in self.cuda_graph_modules
                        or CudaGraphModule.moe_preprocess in self.cuda_graph_modules
                    ):
                        if CudaGraphModule.moe_router not in self.cuda_graph_modules:
                            self.cuda_graph_modules.append(CudaGraphModule.moe_router)
                        if CudaGraphModule.moe_preprocess not in self.cuda_graph_modules:
                            self.cuda_graph_modules.append(CudaGraphModule.moe_preprocess)

                assert (
                    CudaGraphModule.moe not in self.cuda_graph_modules
                    or CudaGraphModule.moe_router not in self.cuda_graph_modules
                ), 'cuda_graph_modules must not contain both moe and moe_router.'
                if CudaGraphModule.moe_preprocess in self.cuda_graph_modules:
                    assert (
                        CudaGraphModule.moe_router in self.cuda_graph_modules
                    ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.'
                if self.num_moe_experts is None or self.num_moe_experts <= 1:
                    assert (
                        CudaGraphModule.moe not in self.cuda_graph_modules
                        and CudaGraphModule.moe_router not in self.cuda_graph_modules
                    ), 'moe cuda graph is only supported for MoE.'
                else:
                    if self.moe_layer_freq == 1 or (
                        isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq
                    ):
                        assert CudaGraphModule.mlp not in self.cuda_graph_modules, (
                            'mlp cuda graph is only supported for dense layers, '
                            'but not found in the model.'
                        )
                    if (
                        self.moe_expert_capacity_factor is None
                        or not self.moe_pad_expert_input_to_capacity
                    ):
                        assert (
                            CudaGraphModule.moe not in self.cuda_graph_modules
                        ), 'moe cuda graph is only supported with drop-padding MoE.'
                        if self.moe_token_dispatcher_type == 'alltoall' and (
                            self.moe_expert_capacity_factor is not None
                            or self.moe_router_padding_for_fp8
                        ):
                            assert CudaGraphModule.moe_preprocess not in self.cuda_graph_modules, (
                                'moe_preprocess cuda graph is not supported when there are '
                                'DtoH copies and synchronizations in the preprocess step.'
                            )

            if self.recompute_granularity:
                if self.recompute_granularity != "selective":
                    assert (
                        self.cuda_graph_impl == "full_iteration"
                    ), "full recompute is only supported with full iteration CUDA graph."
                else:
                    # The recompute module should be inside or outside of the graph scope.
                    # Recompute module coverring graph scope is not allowed.
                    if (
                        self.cuda_graph_impl == "transformer_engine"
                        and "moe" in self.recompute_modules
                    ):
                        assert (
                            CudaGraphModule.moe_router not in self.cuda_graph_modules
                        ), "moe recompute is not supported with moe_router CUDA graph with: "
                        "--cuda-graph-impl transformer_engine."

                    # Graphed recompute module doesn't accept random number.
                    # full_cudagraph means either full_iteration impl or an empty per-layer scope
                    # (which captures the whole layer).
                    if self.cuda_graph_impl == "full_iteration" or not self.cuda_graph_modules:
                        full_cudagraph = True
                    else:
                        full_cudagraph = False
                    if self.attention_dropout != 0.0:
                        assert (
                            not full_cudagraph
                            and CudaGraphModule.attn not in self.cuda_graph_modules
                        ) or "core_attn" not in self.recompute_modules, (
                            "attention dropout is not supported with graphed attention "
                            "recomputation."
                        )
                    if self.hidden_dropout != 0.0:
                        assert (
                            (
                                not full_cudagraph
                                and CudaGraphModule.mlp not in self.cuda_graph_modules
                            )
                            or "mlp" not in self.recompute_modules
                        ) and (
                            (
                                not full_cudagraph
                                and CudaGraphModule.moe not in self.cuda_graph_modules
                            )
                            or "moe" not in self.recompute_modules
                        ), "hidden dropout is not supported with graphed MLP/MoE recomputation."
                    if self.moe_input_jitter_eps is not None:
                        assert (
                            not full_cudagraph
                            and CudaGraphModule.moe not in self.cuda_graph_modules
                        ) or "moe" not in self.recompute_modules, (
                            "moe_input_jitter_eps is not supported with graphed moe recomputation."
                        )

                    if (
                        self.gtp_weight_remat_size > 1
                        and self.cuda_graph_impl == "local"
                        and (self.fp8 is not None or self.fp4 is not None)
                        and self.moe_shared_expert_intermediate_size is not None
                        and not self.moe_shared_expert_overlap
                        and (
                            full_cudagraph
                            or CudaGraphModule.moe in self.cuda_graph_modules
                            or CudaGraphModule.moe_router in self.cuda_graph_modules
                        )
                    ):
                        assert "shared_experts" not in self.recompute_modules, (
                            "GTP + local CUDA graphs that capture shared_experts "
                            "(moe_router/moe scope) cannot recompute it under fp8/fp4: "
                            "te_checkpoint requires .backward(), but the local fwd-graph "
                            "warmup uses .grad(). Drop 'shared_experts' from "
                            "--recompute-modules (GTP-shard + offload instead), or use "
                            "--cuda-graph-impl full_iteration."
                        )

            if self.fine_grained_activation_offloading:
                offload_modules = set(self.offload_modules or [])
                if self.cuda_graph_impl == "local":
                    local_supported_offload_modules = {"expert_fc1", "moe_act", "fused_group_mlp"}
                    unsupported_offload_modules = offload_modules - local_supported_offload_modules
                    assert not unsupported_offload_modules, (
                        "fine-grained activation offloading with cuda_graph_impl='local' "
                        "only supports offload_modules 'expert_fc1', 'moe_act', and "
                        "'fused_group_mlp'. "
                        f"Unsupported offload_modules: {sorted(unsupported_offload_modules)}."
                    )
                    assert self.cuda_graph_modules, (
                        "fine-grained activation offloading with cuda_graph_impl='local' "
                        "is not supported with whole-layer CUDA graph capture."
                    )
                local_partial_moe_offload = (
                    self.cuda_graph_impl == "local"
                    and bool(offload_modules)
                    and offload_modules <= {"expert_fc1", "moe_act", "fused_group_mlp"}
                    and CudaGraphModule.moe not in self.cuda_graph_modules
                )
                assert (
                    self.cuda_graph_impl in ("transformer_engine", "full_iteration")
                    or local_partial_moe_offload
                ), (
                    "fine-grained activation offloading is only supported with "
                    "transformer_engine CUDA graph implementation or local CUDA graph "
                    "implementation with partial MoE offload. Local partial CUDA graphs "
                    "are supported only for expert_fc1, moe_act, or fused_group_mlp "
                    "offload when the full MoE module is not captured."
                )
                assert (
                    CudaGraphModule.moe not in self.cuda_graph_modules
                ), "Token-drop MoE is temporarily not supported with activation offloading."
                assert self.cuda_graph_warmup_steps > 0, (
                    "cuda_graph_warmup_steps must be greater than 0 when enabling "
                    "fine-grained activation offloading."
                )
                if self.cuda_graph_impl == "full_iteration":
                    assert (
                        self.fine_grained_offloading_max_inflight_offloads is not None
                        and self.fine_grained_offloading_max_inflight_offloads >= 0
                    ), (
                        "fine_grained_offloading_max_inflight_offloads must be a non-negative "
                        "integer when using fine-grained activation offloading with "
                        "full-iteration CUDA graphs"
                    )

        # Only meaningful for MoE models; dense models never dispatch tokens,
        # so the (unused) dispatcher default must not fail validation.
        if self.num_moe_experts is not None and self.moe_token_dispatcher_type in ["allgather"]:
            if self.variable_seq_lengths is True:
                raise ValueError(
                    f"Token dispatcher type: {self.moe_token_dispatcher_type} does not support "
                    f"variable sequence length, please use alltoall dispatcher instead."
                )

        if self.moe_permute_fusion:
            from megatron.core.transformer.moe.moe_utils import (
                fused_permute,
                fused_permute_with_probs,
                fused_sort_chunks_by_index,
                fused_sort_chunks_by_index_with_probs,
                fused_unpermute,
            )

            if (
                fused_permute is None
                or fused_permute_with_probs is None
                or fused_sort_chunks_by_index is None
                or fused_sort_chunks_by_index_with_probs is None
                or fused_unpermute is None
            ):
                raise ValueError("fused permutation is not available. Please install TE >= 2.1.0.")

        if self.overlap_moe_expert_parallel_comm:
            # TODO: remove this after we fix the hang issue with torch version < 2.6.0
            assert is_torch_min_version(
                "2.6.0"
            ), "A2A Overlap encounters hang issue with torch version < 2.6.0"
            if self.pipeline_model_parallel_size > 1:
                assert self.virtual_pipeline_model_parallel_size is not None, (
                    "If enabling EP A2A overlap, virtual_pipeline_model_parallel_size "
                    "must be specified when pipeline_model_parallel_size > 1"
                )
            # Expert model parallelism requirements
            assert (
                self.expert_model_parallel_size > 1
            ), 'overlap_moe_expert_parallel_comm is only supported with expert model parallelism'
            assert self.moe_token_dispatcher_type in [
                'alltoall',
                'flex',
            ], 'overlap_moe_expert_parallel_comm is supported with alltoall/flex token dispatcher'

            assert (
                self.recompute_granularity != 'full'
            ), 'disable full recomputation when enabling overlap_moe_expert_parallel_comm'
            assert (
                self.recompute_method is None
            ), 'disable recomputation method when enabling overlap_moe_expert_parallel_comm'
            assert (
                self.recompute_num_layers is None
            ), 'recompute_num_layers must be None when enabling overlap_moe_expert_parallel_comm'
            assert (
                "moe" not in self.recompute_modules
            ), 'disable moe in recompute_modules when enabling overlap_moe_expert_parallel_comm'

            # Check if bf16 or fp16 is used
            assert (
                self.bf16 or self.fp16
            ), 'overlap_moe_expert_parallel_comm is only supported with bf16 or fp16 model'

            assert (
                not self.moe_shared_expert_overlap
            ), 'disable moe_shared_expert_overlap when enabling overlap_moe_expert_parallel_comm'
            assert self.mtp_num_layers in (
                None,
                0,
                1,
            ), 'MTP supports at most one layer when enabling overlap_moe_expert_parallel_comm.'

            # NCCL EP (ncclep flex backend) mirrors hybridep's comm/compute overlap, but a few
            # configs are not yet safe under the 1F1B split and are gated here.
            if (
                self.moe_token_dispatcher_type == 'flex'
                and self.moe_flex_dispatcher_backend == 'ncclep'
            ):
                assert not (
                    self.fine_grained_activation_offloading
                    and 'expert_fc1' in (self.offload_modules or [])
                ), (
                    "overlap_moe_expert_parallel_comm with ncclep does not support offloading "
                    "'expert_fc1': it forces expert FC1 to save the raw bf16 input, which the "
                    'overlap path eagerly frees.'
                )

            if self.cuda_graph_impl != "none":
                if self.cuda_graph_impl == "transformer_engine":
                    assert (
                        CudaGraphModule.moe not in self.cuda_graph_modules
                        and CudaGraphModule.mlp not in self.cuda_graph_modules
                    ), (
                        'CUDA graph scope on moe and mlp is not '
                        'supported with overlap_moe_expert_parallel_comm'
                    )

        # Check delay_wgrad_compute compatibility
        if self.delay_wgrad_compute:
            assert (
                self.overlap_moe_expert_parallel_comm
            ), 'overlap_moe_expert_parallel_comm must be enabled when enabling delay_wgrad_compute'
            if self.cuda_graph_impl == "transformer_engine":
                assert is_te_min_version("2.10.0"), (
                    'TE version >= 2.10.0 is required for delay_wgrad_compute with '
                    'partial cuda graph'
                )

        if self.overlap_dispatch_backward_with_experts_wgrad:
            assert not self.overlap_moe_expert_parallel_comm, (
                'overlap_moe_expert_parallel_comm must be disabled when enabling '
                'overlap_dispatch_backward_with_experts_wgrad.'
            )
            assert is_te_min_version(
                "2.3.0"
            ), 'TE version >= 2.3.0 is required for overlap_dispatch_backward_with_experts_wgrad'
            assert not self.delay_wgrad_compute, (
                'delay_wgrad_compute and overlap_dispatch_backward_with_experts_wgrad '
                'are mutually exclusive; use only one'
            )

        if self.ep_overlap_early_attn_memory_release:
            assert self.overlap_moe_expert_parallel_comm, (
                'overlap_moe_expert_parallel_comm must be enabled when enabling '
                'ep_overlap_early_attn_memory_release'
            )

        if self.context_parallel_size > 1 and self.cp_comm_type is not None:
            if isinstance(self.cp_comm_type, list):
                assert len(self.cp_comm_type) == self.num_layers, (
                    f"Length of cp_comm_type ({len(self.cp_comm_type)}) should equal to "
                    f"the total number of transformer layers ({self.num_layers})!"
                )
            else:
                assert isinstance(
                    self.cp_comm_type, str
                ), "Unsupported communication type for context parallelism!"

        assert (
            self.pipeline_model_parallel_size > 0
        ), f"Pipeline model parallel size must be larger than 0 \
            when enable --standalone-embedding-stage and --standalone-loss-stage"

        if (
            self.num_moe_experts is not None
            and self.num_moe_experts >= 32
            and not self.moe_router_dtype
        ):
            warnings.warn(
                "Using a large number of experts (e.g. >=32) without fp32 routing. "
                "Consider enabling moe_router_dtype for better numerical stability."
            )
        if self.symmetric_ar_type is not None:
            if not HAVE_PACKAGING:
                raise ImportError(
                    "packaging is not installed. Please install it with `pip install packaging`."
                )
            assert is_torch_min_version("2.7.0a0"), "Must have at least torch version 2.7 or higher"
            assert is_te_min_version("2.3.0") or get_te_version() == PkgVersion(
                "2.3.0.dev0+39c0e70"
            ), "Must have at least TE version 2.3 or higher to use symmetric memory all reduce"

        if self.no_rope_freq:
            assert not self.flash_decode, "flash_decode cannot be used with no_rope."
            if isinstance(self.no_rope_freq, int):
                assert self.num_layers % self.no_rope_freq == 0, (
                    f"no_rope_freq={self.no_rope_freq} should be "
                    f"divisible by num_layers={self.num_layers}."
                )
                # Convert integer pattern to list pattern
                # e.g. no_rope=4 with num_layers=8 becomes [0,0,0,1,0,0,0,1]
                pattern = [0] * (self.no_rope_freq - 1) + [1]
                self.no_rope_freq = pattern * (self.num_layers // self.no_rope_freq)
            else:
                assert len(self.no_rope_freq) == self.num_layers, (
                    f"Length of no_rope list ({len(self.no_rope_freq)}) must match "
                    f"the number of layers ({self.num_layers})"
                )

        if self.transformer_impl == "inference_optimized":
            assert self.normalization == "RMSNorm"
            assert not self.layernorm_zero_centered_gamma
            assert not self.add_bias_linear
            assert not self.add_qkv_bias
            assert not self.use_kitchen

        if self.experimental_attention_variant == "dsa":
            assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention"
            if self.context_parallel_size > 1:
                cp_comm_types = (
                    self.cp_comm_type
                    if isinstance(self.cp_comm_type, list)
                    else [self.cp_comm_type]
                )
                assert all(
                    cp_comm_type is not None
                    and cp_comm_type.replace("_", "").lower() == "allgather"
                    for cp_comm_type in cp_comm_types
                ), (
                    "DSAttention context parallelism currently supports "
                    "cp_comm_type=allgather only."
                )

        if self.inference_fuse_tp_communication:
            assert self.transformer_impl == "inference_optimized", (
                "inference_fuse_tp_communication is only supported "
                "for inference_optimized transformer implementation."
            )
            assert (
                self.num_moe_experts is None
            ), "--inference-fuse-tp-communication is not supported for MoE models."

        if self.inference_disable_triton_nvls_kernels:
            assert self.transformer_impl == "inference_optimized", (
                "inference_disable_triton_nvls_kernels is only supported "
                "for inference_optimized transformer implementation."
            )

        if self.flash_attention_version is not None:
            assert self.flash_attention_version in (2, 3, 4), (
                "flash_attention_version must be one of 2, 3, or 4, got "
                f"{self.flash_attention_version}"
            )

        if self.batch_invariant_mode:
            from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
                _BATCH_INVARIANT_BACKENDS,
            )

            # argparse validates via the Literal annotation; guard here too so
            # programmatic TransformerConfig construction fails at build time
            # rather than inside enable_batch_invariant_mode() after model init.
            assert self.batch_invariant_backend in _BATCH_INVARIANT_BACKENDS, (
                f"Unknown batch_invariant_backend {self.batch_invariant_backend!r}; "
                f"expected one of {_BATCH_INVARIANT_BACKENDS}."
            )
            assert self.params_dtype == torch.bfloat16, (
                "Batch invariant mode supports BF16 model parameters only; "
                f"got {self.params_dtype}."
            )
            assert (
                self.attention_backend == AttnBackend.flash
            ), "Batch invariant mode only supports FlashAttention (--attention-backend flash)"
            # The training (TransformerEngine) and inference attention paths must run
            # the same FlashAttention kernel, so the version cannot be left to each
            # path's autodetection. FlashAttention-2 is excluded because it does not
            # expose the fixed num_splits schedule the batch-invariant kernels require.
            assert self.flash_attention_version in (3, 4), (
                "Batch invariant mode requires --flash-attention-version 3 or 4 so the "
                "training and inference attention paths run the same batch-invariant "
                f"FlashAttention kernel (got {self.flash_attention_version})."
            )
            # Context parallelism routes through TE's FA2 fwd/bwd kernels directly, which
            # cannot be pinned to another version; dropout is not batch-invariant.
            assert (
                self.context_parallel_size == 1
            ), "Batch invariant mode does not support context parallelism"
            assert (
                self.attention_dropout == 0.0
            ), "Batch invariant mode does not support attention dropout"
            if (self.num_moe_experts or 0) > 0:
                from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
                    HAVE_DEEPGEMM_BF16,
                )

                if self.transformer_impl != "inference_optimized":
                    assert self.moe_token_dispatcher_type == "alltoall", (
                        "Batch-invariant MoE training requires "
                        "moe_token_dispatcher_type='alltoall'."
                    )
                    if self.batch_invariant_backend == "te_native":
                        assert (
                            not self.use_transformer_engine_op_fuser
                        ), "Batch-invariant training does not support TE op fuser."
                        assert self.moe_use_grouped_tensor or not bool(
                            int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))
                        ), (
                            "Enable device-metadata GEMM with moe_use_grouped_tensor=True, "
                            "not just NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM, so expert "
                            "padding uses the required 256-row alignment."
                        )
                mxfp8_params_enabled = (
                    bool(self.fp8)
                    and self.fp8_recipe == Fp8Recipe.mxfp8
                    and self.fp8_param
                    and not self.fp4
                )
                # DeepGEMM is used by the "deepgemm"/"triton" backends, and by
                # the torch inference path for BF16 experts. MXFP8 experts use
                # torch scaled_grouped_mm directly and do not need DeepGEMM.
                needs_deepgemm = self.batch_invariant_backend in ("deepgemm", "triton") or (
                    self.transformer_impl == "inference_optimized"
                    and self.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.TORCH
                    and not mxfp8_params_enabled
                )
                assert not needs_deepgemm or HAVE_DEEPGEMM_BF16, (
                    "batch_invariant_mode=True with MoE requires DeepGEMM with bf16 "
                    "grouped-GEMM bindings (m_grouped_bf16_gemm_nt_contiguous) for "
                    "this backend combination. "
                    "Install via `uv pip install -e .[batch_invariant]`."
                )
                squared_relu_or_swiglu = (
                    not self.gated_linear_unit and self.activation_func == squared_relu
                ) or (self.gated_linear_unit and self.activation_func == F.silu)
                if self.transformer_impl == "inference_optimized":
                    mxfp8_supported = mxfp8_params_enabled and (
                        (
                            self.inference_grouped_gemm_backend
                            in (InferenceGroupedGemmBackend.TORCH, InferenceGroupedGemmBackend.VLLM)
                            and squared_relu_or_swiglu
                        )
                        or (
                            self.inference_grouped_gemm_backend
                            == InferenceGroupedGemmBackend.FLASHINFER
                            and not self.gated_linear_unit
                            and self.activation_func == squared_relu
                        )
                    )
                else:
                    # The training policy uses TE GroupedLinear directly; the inference
                    # backend selector is generation-only and therefore irrelevant here.
                    mxfp8_supported = (
                        mxfp8_params_enabled
                        and self.moe_grouped_gemm
                        and self.batch_invariant_backend == "te_native"
                        and squared_relu_or_swiglu
                    )
                assert mxfp8_supported or not (self.fp8 or self.fp4), (
                    "Batch-invariant MoE supports BF16; TE MXFP8 squared-ReLU/SwiGLU "
                    "training experts; and Torch/vLLM MXFP8 squared-ReLU/SwiGLU "
                    "or FlashInfer MXFP8 squared-ReLU inference experts."
                )
                assert not (self.moe_permute_fusion or self.moe_permute_fusion_into_hybridep), (
                    "Batch-invariant MoE requires the unfused permute/unpermute path so "
                    "top-k reductions use the fixed batch-invariant add tree."
                )
                assert not (
                    self.moe_pad_expert_input_to_capacity
                    or self.moe_pad_experts_for_cuda_graph_inference
                ), (
                    "Batch-invariant MoE supports dynamic dropless routing only. "
                    "Disable MoE capacity/expert padding."
                )

        # Scheduler-value, max-seqlen, and variable_seq_lengths handling live in
        # ModelParallelConfig.__post_init__ next to the field definitions; only the
        # transformer-stack requirements are validated here.
        if self.sequence_packing_scheduler is not None:
            # Check TE version.
            if not HAVE_PACKAGING:
                raise ImportError(
                    "packaging is not installed. Please install it with `pip install packaging`."
                )
            # TODO: remove this after we fix the convergence issue with TE < 2.9.
            if not (
                is_te_min_version("2.9.0") or get_te_version() == PkgVersion("2.9.0.dev0+5b3092a")
            ):
                raise ValueError(
                    "SFT sequence packing requires Transformer Engine >= 2.9.0 "
                    f"but got {get_te_version()} (TE < 2.9.0 may have convergence issues)."
                )

            # TODO(tailaim): add support for other dispatcher types
            # Only relevant for MoE models; dense models never dispatch tokens,
            # so the (unused) dispatcher default must not fail validation. For
            # allgather specifically, the general variable_seq_lengths check
            # above raises first (packing derives variable_seq_lengths=True).
            if self.num_moe_experts is not None:
                assert self.moe_token_dispatcher_type == "alltoall", (
                    f"sequence_packing only supports moe_token_dispatcher_type='alltoall', "
                    f"got '{self.moe_token_dispatcher_type}'"
                )


@dataclass
class MLATransformerConfig(TransformerConfig):
    """Configuration object for megatron-core Multi-Latent Attention (MLA) transformers.

    The initialization function has an argument for each parameter, including those in
    ModelParallelConfig. Included YaRN RoPE parameters that is fused in MLA.
    """

    multi_latent_attention: bool = True
    """Whether to use Multi-Latent Attention."""

    use_fused_mla_q_uproj: bool = False
    """Use the cuDNN fused MLA Q up-proj + per-head RoPE + MXFP8-quant kernel (SM100 only).
    Requires apply_rope_fusion=True, q_lora_rank set, TP=1, SBHD, MXFP8 DPA, and zero
    attention dropout."""

    q_lora_rank: int = 512
    """Rank of Query tensor's low rank representation."""

    kv_lora_rank: int = 512
    """Rank of Key and Value tensors' low rank representation.
       This is not used for DSv4 Hybrid Attention and will be overridden automatically."""

    attention_latent_norm_epsilon: float | None = None
    """Epsilon for the primary query and key-value latent norms in attention.
       If unset, inherit ``layernorm_epsilon`` for backward compatibility."""

    qk_head_dim: int = 128
    """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim
       This is not used for DSv4 Hybrid Attention and will be overridden automatically."""

    qk_pos_emb_head_dim: int = 64
    """Dimension of the position embedding in the QK projection."""

    v_head_dim: int = 128
    """Dimension of the head in the V projection."""

    normalization: str = "RMSNorm"
    """Default normalization layer for MLA models is RMSNorm."""

    rope_type: str = "yarn"
    """Type of RoPE to use. Default to yarn, options are rope and yarn."""

    rotary_base: float = 10000
    """Rotary base for the rotary embeddings, used by rope and yarn."""

    rotary_percent: float = 1.0
    """Rotary percent for the rotary embeddings, used by rope."""

    rotary_scaling_factor: float = 40
    """Rotary scaling factor for the rotary embeddings, used by yarn."""

    original_max_position_embeddings: int = 4096
    """Original maximum position embeddings for the original model, used by yarn."""

    beta_fast: float = 32
    """Beta fast for YaRN RoPE, used by yarn."""

    beta_slow: float = 1
    """Beta slow for YaRN RoPE, used by yarn."""

    mscale: float = 1.0
    """Mscale for YaRN RoPE in Multi-Latent Attention, used by yarn."""

    mscale_all_dim: float = 0.0
    """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn."""

    output_projection_groups: int = 8
    """Number of groups for the grouped low-rank output projection (wo_a)."""

    output_projection_lora_rank: int = 1024
    """Low-rank dimension per group for the grouped output projection (wo_a)."""

    cache_mla_latents: bool = False
    """Cache the low dimensional tensors for MLA rather than full KV cache.
       This is only for the dynamic inference backend and requires that 
       Flash MLA is installed."""

    mla_down_proj_fusion: bool = False
    """Enable fused q/kv down-projection and fused input layernorm when backend supports.
       Otherwise fall back to the unfused MLA.
    """

    def __post_init__(self):
        super().__post_init__()
        if self.attention_latent_norm_epsilon is None:
            self.attention_latent_norm_epsilon = self.layernorm_epsilon

        if (
            self.multi_latent_attention
            and self.apply_rope_fusion
            and self.rope_type != "yarn"
            and self.experimental_attention_variant != "dsv4_hybrid"
        ):
            raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.")

        if self.use_fused_mla_q_uproj and (
            self.fp8 is None
            or self.fp8_recipe != Fp8Recipe.mxfp8
            or not self.fp8_dot_product_attention
            or self.attention_dropout != 0.0
        ):
            raise ValueError(
                "use_fused_mla_q_uproj requires FP8 with fp8_recipe='mxfp8' and "
                "fp8_dot_product_attention=True so TE interprets the pre-quantized Q/K/V "
                "scale layout correctly, plus attention_dropout=0.0 because MXFP8 attention "
                "backward does not support dropout."
            )

        if self.attention_output_gate:
            raise NotImplementedError("Output gate is not supported for MLA yet.")

        # DSv4 hybrid: derive qk_head_dim and kv_lora_rank from v_head_dim and qk_pos_emb_head_dim.
        if self.experimental_attention_variant == "dsv4_hybrid":
            assert (
                not self.mla_down_proj_fusion
            ), "MLA down projection fusion must be disabled for DSv4 hybrid mode."
            assert self.q_lora_rank is not None, "DSv4 hybrid mode requires q_lora_rank."
            assert (
                self.output_projection_groups > 0
            ), "DSv4 hybrid mode requires output_projection_groups to be positive."
            assert (
                self.output_projection_lora_rank > 0
            ), "DSv4 hybrid mode requires output_projection_lora_rank to be positive."
            assert (
                self.num_attention_heads * self.v_head_dim
            ) % self.output_projection_groups == 0, (
                "num_attention_heads * v_head_dim must be divisible by " "output_projection_groups."
            )
            log_single_rank(
                logger,
                logging.WARNING,
                "DSv4 hybrid mode is enabled, deriving qk_head_dim and kv_lora_rank from "
                "v_head_dim and qk_pos_emb_head_dim",
            )
            derived = self.v_head_dim - self.qk_pos_emb_head_dim
            assert derived > 0, "v_head_dim must be greater than qk_pos_emb_head_dim."
            self.qk_head_dim = derived
            self.kv_lora_rank = derived

        if self.cache_mla_latents:
            assert (
                self.apply_rope_fusion is False
            ), "Rope Fusion is not compatible with caching latents"