diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 7b6cda9975d..caf790ad51e 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -458,6 +458,8 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC attention_layer_map, dsa_layer_map = operator.itemgetter( Symbols.ATTENTION, Symbols.DS_ATTENTION )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)) + if Symbols.KDA in mamba_inference_state_config.layer_type_list: + raise NotImplementedError("KDA layers are not supported for inference.") recurrent_layer_map = {} for global_layer_idx, layer_type in enumerate( mamba_inference_state_config.layer_type_list diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index c22d8784495..c3e32cca772 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -7,6 +7,21 @@ import torch +def _parse_pad_packed_seq_alignment(value): + """Parse THD packed-sequence padding alignment.""" + if value == "max": + return value + try: + alignment = int(value) + except (TypeError, ValueError) as exc: + raise ValueError( + "pad_packed_seq_alignment must be 'max' or a positive integer alignment." + ) from exc + if alignment <= 0: + raise ValueError("pad_packed_seq_alignment must be 'max' or a positive integer alignment.") + return alignment + + def resolve_tensor_parallel_weight_shards( tensor_model_parallel_size: int, tensor_parallel_num_weight_shards: Optional[int], diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index 87c63a126bb..96592b42028 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -1,11 +1,16 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import warnings from typing import List, Optional from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.models.backends import BackendSpecProvider -from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNet2, GatedDeltaNetSubmodules +from megatron.core.ssm.gated_delta_net import ( + GatedDeltaNet, + GatedDeltaNetSubmodules, + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, +) from megatron.core.transformer.enums import AttnMaskType, LayerType from megatron.core.transformer.experimental_attention_variant.absorbed_mla import ( AbsorbedMLASelfAttention, @@ -19,6 +24,7 @@ is_dsa_skip_topk_layer, source_dsa_compute_layer, ) +from megatron.core.transformer.hyper_connection import HyperConnectionModule from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_block import ( @@ -27,6 +33,7 @@ ) from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, MlpBuilder, TransformerLayer, TransformerLayerSubmodules, @@ -58,7 +65,7 @@ ########## # Canonical ``experimental_attention_variant`` names served by the gated delta net family. -GDN_ATTENTION_VARIANTS = ("gdn", "gdn2") +GDN_ATTENTION_VARIANTS = ("gdn", "kda") # Deprecated ``experimental_attention_variant`` spellings mapped to their canonical name. _DEPRECATED_ATTENTION_VARIANT_ALIASES = {"gated_delta_net": "gdn"} @@ -72,25 +79,33 @@ def get_gated_delta_net_module_spec( config: TransformerConfig, backend: BackendSpecProvider = None ) -> ModuleSpec: - """Build module spec for GatedDeltaNet attention.""" + """Build a module spec for a GDN-family attention variant.""" if backend is None: backend = _get_backend_spec_provider(config=config) rms_norm = config.normalization == "RMSNorm" - # gdn2 reuses the GDN submodules and spec structure with the GatedDeltaNet2 module. - gdn_module = ( - GatedDeltaNet2 if config.experimental_attention_variant == "gdn2" else GatedDeltaNet - ) - attention = ModuleSpec( - module=gdn_module, - submodules=GatedDeltaNetSubmodules( - in_proj=backend.column_parallel_layer_norm_linear(), - out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False), - out_proj=backend.row_parallel_linear(), - ), - metainfo={"fuse_input_layernorm": True}, - ) + if config.experimental_attention_variant == "kda": + attention = ModuleSpec( + module=KimiDeltaAttention, + submodules=KimiDeltaAttentionSubmodules( + in_proj=backend.column_parallel_linear(), + beta_proj=backend.column_parallel_linear(), + out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False), + out_proj=backend.row_parallel_linear(), + ), + metainfo={"fuse_input_layernorm": False}, + ) + else: + attention = ModuleSpec( + module=GatedDeltaNet, + submodules=GatedDeltaNetSubmodules( + in_proj=backend.column_parallel_layer_norm_linear(), + out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False), + out_proj=backend.row_parallel_linear(), + ), + metainfo={"fuse_input_layernorm": True}, + ) return attention @@ -245,6 +260,10 @@ def get_transformer_layer_with_experimental_attention_variant_spec( # Get GPT decoder block layer specs rms_norm = config.normalization == "RMSNorm" + enable_hc = config.enable_hyper_connections + hc_module = HyperConnectionModule if enable_hc else IdentityOp + layer_module = HyperConnectionTransformerLayer if enable_hc else TransformerLayer + layer_specs = [] for layer_number in range(config.num_layers): attention = ( @@ -271,14 +290,16 @@ def get_transformer_layer_with_experimental_attention_variant_spec( layer_specs.append( ModuleSpec( - module=TransformerLayer, + module=layer_module, submodules=TransformerLayerSubmodules( input_layernorm=input_layernorm, self_attention=attention, self_attn_bda=get_bias_dropout_add, + self_attention_hyper_connection=hc_module, pre_mlp_layernorm=pre_mlp_layernorm, mlp=not_none(mlp), mlp_bda=get_bias_dropout_add, + mlp_hyper_connection=hc_module, ), ) ) @@ -349,19 +370,7 @@ def get_transformer_block_with_experimental_attention_variant_spec( def normalize_experimental_attention_variant( experimental_attention_variant: Optional[str], ) -> Optional[str]: - """Resolve a deprecated ``experimental_attention_variant`` spelling to its canonical name. - - ``gated_delta_net`` is the deprecated spelling of ``gdn``. Passing it emits a - ``DeprecationWarning`` and returns the canonical name so that every downstream - consumer only has to handle ``gdn``. - - Args: - experimental_attention_variant: The configured variant name, possibly a - deprecated alias. - - Returns: - The canonical variant name, or the argument unchanged when it is not an alias. - """ + """Resolve a deprecated attention-variant spelling to its canonical name.""" canonical = _DEPRECATED_ATTENTION_VARIANT_ALIASES.get(experimental_attention_variant) if canonical is None: return experimental_attention_variant @@ -376,11 +385,7 @@ def normalize_experimental_attention_variant( def is_gated_delta_net_variant(experimental_attention_variant: Optional[str]) -> bool: - """Check if the experimental attention variant is served by a gated delta net layer. - - Accepts the deprecated ``gated_delta_net`` spelling without warning; use - :func:`normalize_experimental_attention_variant` to emit the deprecation notice. - """ + """Return whether a name selects a GDN-family attention implementation.""" canonical = _DEPRECATED_ATTENTION_VARIANT_ALIASES.get( experimental_attention_variant, experimental_attention_variant ) @@ -501,7 +506,7 @@ def _get_backend_spec_provider(config: TransformerConfig) -> BackendSpecProvider ) backend: BackendSpecProvider = ( KitchenSpecProvider( - fallback=TESpecProvider(), + fallback=TESpecProvider(fallback_to_eager_attn=config.fallback_to_eager_attn), use_kitchen_attention=config.use_kitchen_attention, kitchen_attention_backend=config.kitchen_attention_backend, ) @@ -537,6 +542,7 @@ def _get_self_attention_module_spec( qk_l2_norm=config.qk_l2_norm, use_kitchen=config.use_kitchen, use_te_activation_func=config.use_te_activation_func, + fallback_to_eager_attn=config.fallback_to_eager_attn, use_kitchen_attention=config.use_kitchen_attention, kitchen_attention_backend=config.kitchen_attention_backend, mla_down_proj_fusion=getattr(config, "mla_down_proj_fusion", False), diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index 20f62f24591..910166d049d 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/megatron/core/models/gpt/gpt_layer_specs.py @@ -125,6 +125,7 @@ def get_gpt_layer_with_inference_submodules( linear_kv_down_proj=backend.linear(), linear_kv_up_proj=linear_kv_up_proj, core_attention=backend.core_attention(), + linear_gate=backend.column_parallel_linear(), linear_proj=backend.row_parallel_linear(), q_layernorm=IdentityOp, kv_layernorm=IdentityOp, @@ -304,6 +305,7 @@ def get_gpt_layer_with_transformer_engine_submodules( linear_kv_down_proj=backend.linear(), linear_kv_up_proj=linear_kv_up_proj, core_attention=backend.core_attention(), + linear_gate=backend.column_parallel_linear(), linear_proj=backend.row_parallel_linear(), q_layernorm=IdentityOp, kv_layernorm=IdentityOp, @@ -424,6 +426,7 @@ def get_gpt_layer_local_submodules( linear_kv_down_proj=backend.column_parallel_linear(), linear_kv_up_proj=backend.column_parallel_linear(), core_attention=backend.core_attention(), + linear_gate=backend.column_parallel_linear(), linear_proj=backend.row_parallel_linear(), q_layernorm=qk_norm if qk_layernorm else IdentityOp, kv_layernorm=qk_norm if qk_layernorm else IdentityOp, diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 20a85839658..84637304d6f 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -1,11 +1,10 @@ -# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2024, Tri Dao, Albert Gu. # Some of this code was adopted from https://github.com/state-spaces/mamba/ # This source code is licensed under the Apache license found in the # LICENSE file in the root directory of this source tree. -import copy from contextlib import nullcontext from dataclasses import dataclass from typing import List, Optional, Tuple, Union @@ -13,30 +12,33 @@ import torch from torch import Tensor, nn -from megatron.core.context_parallel import ContextParallelLayoutManager from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding from megatron.core.enums import Fp8Recipe -from megatron.core.extensions.transformer_engine import TELayerNormColumnParallelLinear, TENorm +from megatron.core.extensions.transformer_engine import TENorm from megatron.core.fp4_utils import get_fp4_context from megatron.core.fp8_utils import get_fp8_context from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.inference.utils import InferenceMode from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols as LayerSymbols -from megatron.core.models.hybrid.layers.hybrid_hyper_connection import HyperConnectionHybridLayer from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager from megatron.core.transformer import TransformerConfig from megatron.core.transformer.cuda_graphs import annotate_first_last_layer +from megatron.core.transformer.enums import CudaGraphModule from megatron.core.transformer.hyper_connection import ( HyperConnectionModule, learned_output_contract, ) from megatron.core.transformer.identity_op import IdentityOp -from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 -from megatron.core.transformer.multi_latent_attention import FusedMLASelfAttention +from megatron.core.transformer.module import ( + GraphableMegatronModule, + MegatronModule, + convert_module_to_dtype_except_fp32_marked, + mark_keep_in_fp32, +) from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_layer import TransformerLayer from megatron.core.transformer.utils import ( @@ -55,14 +57,491 @@ class HybridStackSubmodules: mamba_layer: Union[ModuleSpec, type] = IdentityOp gdn_layer: Union[ModuleSpec, type] = IdentityOp + kda_layer: Union[ModuleSpec, type] = IdentityOp attention_layer: Union[ModuleSpec, type] = IdentityOp dsa_layer: Union[ModuleSpec, type] = IdentityOp mla_layer: Union[ModuleSpec, type] = IdentityOp + csa_layer: Union[ModuleSpec, type] = IdentityOp + hca_layer: Union[ModuleSpec, type] = IdentityOp + window_layer: Union[ModuleSpec, type] = IdentityOp mlp_layer: Union[ModuleSpec, type] = IdentityOp moe_layer: Union[ModuleSpec, type] = IdentityOp mtp_block_spec: Optional[ModuleSpec] = None +class HyperConnectionHybridLayer(GraphableMegatronModule): + """Layer-boundary mHC wrapper for HybridStack layers. + + Hybrid layers already own their local residual paths. For this initial + integration we treat each hybrid layer as a single function by aggregating + n streams to the layer input, running the existing layer, and feeding only + the layer delta back through mHC expansion. The expansion path intentionally + uses zero additional dropout because the wrapped hybrid layer has already + applied its local dropout/residual update before the delta is computed. + + Checkpoint compatibility: this is a *wrapper* (the inner layer is held as + `self.inner_layer`), so wrapped-layer state_dict keys are nested under + `inner_layer.` (e.g. `layers.0.inner_layer.input_layernorm.weight` instead + of `layers.0.input_layernorm.weight`). HybridStack checkpoints saved with + `enable_hyper_connections=False` cannot be loaded into a model with + `enable_hyper_connections=True` (and vice versa) without a key-mapping + migration. Note: this differs from `HyperConnectionTransformerLayer`, + which subclasses `TransformerLayer` and only adds new sibling fields, + keeping all base keys stable. + + CUDA graphs: this wrapper subclasses ``GraphableMegatronModule`` so that, with + ``cuda_graph_impl="transformer_engine"``, wrapped layers are captured per-layer — + mirroring ``HyperConnectionTransformerLayer`` on the GPT path. Without this, the TE + graph discovery (``_layer_is_graphable``) only inspects the top-level layer type and + silently skips every wrapped layer, so an mHC-enabled HybridStack would run entirely + eager. Two capture modes: + + * Non-MoE inner layers (attention variants, Mamba): the whole wrapper forward + (mHC aggregate + inner layer + n-stream BDA) is captured as one graph. The inner + layer's own ``__call__`` graph routing is bypassed during capture (see + ``_call_inner_layer``) to avoid nested capture. + * MoE inner layers, when ``moe_router`` is in ``cuda_graph_modules``: the expert + all-to-all is not graph-safe, so only the deterministic prefix is graphed (mHC + ``compute_mappings``/``aggregate`` + the inner layer's router/preprocess). The graph + outputs the router intermediates, the mHC state (``h_post``, ``h_res``) and the + n-stream residual; on replay the experts run eagerly and the n-stream BDA (eager) + consumes the inner's raw ``mlp_output_with_bias`` as the layer delta. Routing the + residual *through the graph* (not reusing the layer input directly in the eager BDA) + keeps the backward gradient flowing into the captured graph, which is required for + bit-identical training — again mirroring ``HyperConnectionTransformerLayer``. + + ``_get_submodules_under_cudagraphs`` returns the submodules whose params the wrapper + graph's manual hooks must drive: ``[self]`` for whole-wrapper capture, or the mHC module + + the inner router/preprocess submodules for partial MoE capture (experts stay eager). + """ + + def __init__(self, config: TransformerConfig, layer: MegatronModule) -> None: + super().__init__(config=config) + if ( + config.cuda_graph_impl in ("transformer_engine", "full_iteration") + and config.recompute_granularity == "selective" + and "mhc" in (config.recompute_modules or []) + ): + raise ValueError( + "mHC selective recompute with CUDA Graphs (cuda_graph_impl=" + f"{config.cuda_graph_impl!r}) is not supported for HybridStack mHC " + "layers. For per-layer Transformer Engine capture, the hybrid wrapper " + "captures the mHC producer inside the graph, so per-microbatch " + "checkpoint registration cannot run; the guarded attention-only split " + "exists only on the GPT HyperConnectionTransformerLayer path. " + "Full-iteration capture is rejected here because it has only been " + "validated on that same GPT path -- the config-level gate that admits " + "it is not model-family aware. Disable CUDA graphs or remove 'mhc' " + "from recompute_modules." + ) + self.inner_layer = layer + self.layer_number = layer.layer_number + self.hyper_connection = HyperConnectionModule(config=config, layer_number=self.layer_number) + if config.params_dtype is not None: + convert_module_to_dtype_except_fp32_marked(self.hyper_connection, config.params_dtype) + if hasattr(layer, 'tp_group'): + self.tp_group = layer.tp_group + + def get_layer_static_inputs(self, seq_length, micro_batch_size): + """Override to produce n-stream hidden_states of shape [s, b, n*C]. + + CUDA graph capture allocates static buffers sized by this method. The base + returns [s, b, C], but mHC layers carry n-stream hidden states [s, b, n*C]. + Mirrors ``HyperConnectionTransformerLayer.get_layer_static_inputs``. + """ + if hasattr(self.inner_layer, "get_layer_static_inputs"): + static_inputs = self.inner_layer.get_layer_static_inputs(seq_length, micro_batch_size) + else: + static_inputs = super().get_layer_static_inputs(seq_length, micro_batch_size) + hs = static_inputs["hidden_states"] + n = self.config.num_residual_streams + static_inputs["hidden_states"] = torch.ones( + (hs.shape[0], hs.shape[1], n * self.config.hidden_size), + dtype=hs.dtype, + requires_grad=hs.requires_grad, + device=hs.device, + ) + return static_inputs + + @staticmethod + def _decompose_packed_seq_params_to_kwargs(kwargs): + """Decompose PackedSeqParams into tensor kwargs for TE CUDA graphs.""" + packed_seq_params = kwargs.pop('packed_seq_params', None) + if packed_seq_params is None: + return + kwargs['cu_seqlens_q'] = packed_seq_params.cu_seqlens_q + kwargs['cu_seqlens_kv'] = packed_seq_params.cu_seqlens_kv + kwargs['cu_seqlens_q_padded'] = packed_seq_params.cu_seqlens_q_padded + kwargs['cu_seqlens_kv_padded'] = packed_seq_params.cu_seqlens_kv_padded + + def _reconstruct_packed_seq_params_from_kwargs(self, kwargs): + """Reconstruct THD PackedSeqParams from tensor kwargs in the graph capture path.""" + if 'cu_seqlens_q' not in kwargs: + return + max_seqlen = self.config.max_seqlen_per_dp_cp_rank * self.config.context_parallel_size + packed_seq_params = PackedSeqParams( + qkv_format='thd', + cp_partition_mode=self.config.cp_partition_mode, + cu_seqlens_q=kwargs.pop('cu_seqlens_q'), + cu_seqlens_kv=kwargs.pop('cu_seqlens_kv'), + cu_seqlens_q_padded=kwargs.pop('cu_seqlens_q_padded'), + cu_seqlens_kv_padded=kwargs.pop('cu_seqlens_kv_padded'), + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + # This Python flag is baked into the captured graph and cannot vary + # between replay batches. Use the conservative THD-safe branch. + pad_between_seqs=True, + ) + kwargs['packed_seq_params'] = packed_seq_params + + def __call__(self, *args, **kwargs): + # Keep non-tensor recompute state out of TE CUDA graph inputs; the GPT + # hyper-connection path follows the same pattern. + self._mhc_recompute_manager = kwargs.pop("mhc_recompute_manager", None) + return super().__call__(*args, **kwargs) + + def _inner_is_moe(self) -> bool: + """True when the inner layer is an MoE ``TransformerLayer``. Such layers use the + GPT-style raw-delta path (feed the inner's ``mlp_output_with_bias`` straight to the + n-stream BDA) in both the eager forward and the CUDA-graph replay.""" + from megatron.core.transformer.moe.moe_layer import MoELayer + + return isinstance(self.inner_layer, TransformerLayer) and isinstance( + getattr(self.inner_layer, 'mlp', None), MoELayer + ) + + def _inner_is_partial_moe_capture(self) -> bool: + """True when the inner layer is MoE and the configured ``cuda_graph_modules`` request + partial MoE capture (``moe_router``). + + In that case the wrapper does NOT capture the whole forward as one graph (the expert + all-to-all is not graph-safe). Instead it graphs the deterministic prefix (mHC aggregate + + the inner layer's router/preprocess) and runs the experts + mHC BDA eagerly — mirroring + how ``HyperConnectionTransformerLayer`` graphs MoE layers on the GPT path. Whole-wrapper + capture is still used for non-MoE inner layers (attention variants, Mamba). + """ + return ( + self._inner_is_moe() + and bool(self.config.cuda_graph_modules) + and CudaGraphModule.moe_router in self.config.cuda_graph_modules + ) + + def _te_cuda_graph_capture(self, *args, **kwargs): + """Capture the graph-safe portion of the wrapper forward. + + For non-MoE inner layers the whole wrapper forward (mHC aggregate + inner layer + + n-stream BDA) is captured as one graph. For MoE inner layers under ``moe_router`` + partial capture, only the deterministic prefix is graphed: the mHC + ``compute_mappings``/``aggregate`` followed by the inner layer's router/preprocess. + The captured outputs are the inner router/preprocess intermediates plus the mHC + state (``h_post``, ``h_res``) and the aggregated single-stream input needed to + reconstruct the layer delta on replay. ``context`` is ``None`` for the graphed + hybrid layer types, so it is dropped (a tuple containing ``None`` cannot be a + CUDA-graph output). + """ + self._reconstruct_packed_seq_params_from_kwargs(kwargs) + + if self._inner_is_partial_moe_capture(): + hidden_states = args[0] if args else kwargs["hidden_states"] + aggregated, h_res, h_post, residual = self.hyper_connection(hidden_states) + inner_kwargs = dict(kwargs) + inner_kwargs.pop("hidden_states", None) + inner_out = list(self.inner_layer._te_cuda_graph_capture(aggregated, **inner_kwargs)) + # inner_out = router/preprocess intermediates ending in the inner residual; + # append the mHC state AND the n-stream `residual` returned by the (graphed) + # hyper_connection. Routing `residual` through the graph as an output keeps its + # backward grad flowing into the graph's backward (mirrors + # HyperConnectionTransformerLayer), instead of a second autograd path the captured + # backward does not account for. The experts' raw mlp_output_with_bias (produced + # on replay) is the layer delta, so `aggregated` need not be captured. + return tuple(inner_out) + (h_post, h_res, residual) + + hidden_states, context = self.forward(*args, **kwargs) + cuda_graph_outputs = [hidden_states] + if context is not None: + cuda_graph_outputs.append(context) + return tuple(cuda_graph_outputs) + + def _te_cuda_graph_replay(self, *args, **kwargs): + """Replay the captured graph and restore the (hidden_states, context) contract. + + Non-MoE inner layers: the whole wrapper forward was captured, so the only graph + output is the layer's n-stream hidden_states; re-append ``None`` for context. + + MoE inner layers (partial capture): replay the graphed prefix, then run the + experts eagerly and apply the mHC n-stream BDA — reproducing exactly the eager + wrapper tail (``layer_delta = layer_output - aggregated`` then + ``fused_h_res_h_post_bda``), just with the deterministic prefix graphed. + """ + self._decompose_packed_seq_params_to_kwargs(kwargs) + + if self._inner_is_partial_moe_capture(): + out = list(super()._te_cuda_graph_replay(*args, **kwargs)) + residual = out.pop() # n-stream [s, b, n*C] — graph output (see capture) + h_res = out.pop() + h_post = out.pop() + # Resume the inner MoE experts eagerly to the raw delta (mlp_output_with_bias), + # then let the n-stream BDA own the residual — identical to the eager forward + # (`_call_inner_transformer_layer_without_local_bda` fast path → + # fused_h_res_h_post_bda), just with the router/preprocess prefix graphed. + # Mirror the eager fast path's BDA args (it feeds the inner layer's + # `hidden_dropout` / `bias_dropout_fusion`) so replay == eager bit-for-bit. + mlp_output_with_bias = self.inner_layer.resume_moe_experts_after_partial_cudagraph(out) + hidden_states = self.hyper_connection.fused_h_res_h_post_bda( + h_res, + residual, + h_post, + mlp_output_with_bias, + dropout_prob=self.inner_layer.hidden_dropout, + training=self.training, + fused=self.inner_layer.config.bias_dropout_fusion, + manager=None, + ) + if ( + self.config.fp32_residual_connection + and self.config.params_dtype is not None + and hidden_states.dtype != self.config.params_dtype + ): + hidden_states = hidden_states.to(self.config.params_dtype) + return hidden_states, None + + cuda_graph_output = list(super()._te_cuda_graph_replay(*args, **kwargs)) + return cuda_graph_output[0], None + + def _get_submodules_under_cudagraphs(self): + """Submodules whose params are driven by the wrapper graph's manual hooks. + + Whole-wrapper capture covers the entire wrapper (``[self]``, the base default). + For partial MoE capture only the graphed prefix is covered — the mHC module plus + the inner layer's router/preprocess submodules — so the experts (run eagerly) + keep their normal forward hooks. + """ + if self._inner_is_partial_moe_capture(): + return [self.hyper_connection] + self.inner_layer._get_submodules_under_cudagraphs() + return super()._get_submodules_under_cudagraphs() + + def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]: + """Delegate Mamba inference state shape requests to the wrapped layer.""" + if not hasattr(self.inner_layer, 'mamba_state_shapes_per_request'): + return None + return self.inner_layer.mamba_state_shapes_per_request() + + def _call_inner_layer( + self, + hidden_states: Tensor, + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext], + rotary_pos_emb: Optional[Tensor], + sequence_len_offset: Optional[Tensor], + packed_seq_params: Optional[PackedSeqParams], + padding_mask: Optional[Tensor], + input_ids: Optional[Tensor] = None, + ) -> Tuple[Tensor, Optional[Tensor]]: + # When this wrapper is itself being CUDA-graph captured, the inner layer + # must run as a plain forward: routing through its ``__call__`` would + # trigger nested TE graph capture (the inner layer is also a + # GraphableMegatronModule). During eager steps we keep ``__call__`` so the + # inner layer's forward pre-hooks (e.g. param all-gather) fire normally; + # under graph replay these are driven by the wrapper's manual hooks. + from megatron.core.transformer.cuda_graphs import is_graph_capturing + + inner = self.inner_layer.forward if is_graph_capturing() else self.inner_layer + + if isinstance(self.inner_layer, TransformerLayer): + output = inner( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + sequence_len_offset=sequence_len_offset, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + input_ids=input_ids, + _called_from_hybrid_mhc_wrapper=True, + ) + else: + # Non-transformer layers (e.g. MambaLayer; GatedDeltaNet which does + # accept `sequence_len_offset` is currently always wrapped inside a + # TransformerLayer spec, so it takes the branch above) do not accept + # rotary_pos_emb / sequence_len_offset / padding_mask — pass only + # the common arguments. New layer types that consume any of these + # must add explicit handling here. + output = inner( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + ) + + if isinstance(output, tuple): + context = output[1] if len(output) > 1 else None + return output[0], context + return output, None + + def _call_inner_transformer_layer_without_local_bda( + self, + hidden_states: Tensor, + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext], + rotary_pos_emb: Optional[Tensor], + sequence_len_offset: Optional[Tensor], + packed_seq_params: Optional[PackedSeqParams], + padding_mask: Optional[Tensor], + input_ids: Optional[Tensor] = None, + ) -> Optional[Tuple[Tuple[Tensor, Optional[Tensor]], Optional[Tensor], float, bool]]: + """Return a raw TransformerLayer branch output when the wrapped layer is split. + + Hybrid DSv4 layers are usually attention-only (`W/C/H/D`) or MLP/MoE-only (`-/E`) + TransformerLayer instances. For those layers, skip the inner layer's local + residual+BDA and feed the raw branch output directly into the mHC BDA, matching the + GPT mHC path and avoiding a residual add followed by `layer_output - aggregated`. + """ + if not isinstance(self.inner_layer, TransformerLayer): + return None + + layer = self.inner_layer + if (not layer.training) and layer.config.inference_fuse_tp_communication: + return None + + has_attention = not isinstance(layer.self_attention, IdentityOp) + has_cross_attention = not isinstance(layer.cross_attention, IdentityOp) + has_mlp = not isinstance(layer.mlp, IdentityOp) + + if has_cross_attention or has_attention == has_mlp: + return None + + if has_attention: + output_with_bias, attn_norm_manager, residual = ( + layer._forward_self_attention_output_with_bias( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + ) + ) + output_with_bias = layer._group_offload_output_with_bias( + output_with_bias, attn_norm_manager, forced_released_tensors=[residual] + ) + return output_with_bias, None, layer.hidden_dropout, layer.config.bias_dropout_fusion + + output_with_bias, residual = layer._forward_mlp_output_with_bias( + hidden_states, + inference_context=inference_context, + padding_mask=padding_mask, + input_ids=input_ids, + packed_seq_params=packed_seq_params, + ) + if layer.mlp_norm_manager is not None: + output_with_bias = layer._group_offload_output_with_bias( + output_with_bias, layer.mlp_norm_manager, forced_released_tensors=[residual] + ) + layer.mlp_norm_manager = None + return output_with_bias, None, layer.hidden_dropout, layer.config.bias_dropout_fusion + + def forward( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor] = None, + inference_context: Optional[BaseInferenceContext] = None, + rotary_pos_emb: Optional[Tensor] = None, + sequence_len_offset: Optional[Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + padding_mask: Optional[Tensor] = None, + input_ids: Optional[Tensor] = None, + mhc_recompute_manager=None, + ) -> Tuple[Tensor, Optional[Tensor]]: + """Run the wrapped hybrid layer through one layer-boundary mHC update. + + ``attention_mask`` defaults to ``None`` so that CUDA-graph capture, which + calls this forward with only the static ``hidden_states`` input, does not + fail on a missing positional argument (causal masking is inferred by the + attention backend when the mask is ``None``). + """ + + if mhc_recompute_manager is None: + mhc_recompute_manager = getattr(self, '_mhc_recompute_manager', None) + + aggregated, h_res, h_post, residual = self.hyper_connection( + hidden_states, mhc_recompute_manager=mhc_recompute_manager + ) + fast_path_result = self._call_inner_transformer_layer_without_local_bda( + aggregated, + attention_mask, + inference_context, + rotary_pos_emb, + sequence_len_offset, + packed_seq_params, + padding_mask, + input_ids, + ) + + if fast_path_result is None: + layer_output, context = self._call_inner_layer( + aggregated, + attention_mask, + inference_context, + rotary_pos_emb, + sequence_len_offset, + packed_seq_params, + padding_mask, + input_ids, + ) + # The inner hybrid layer already applied its own local residual/dropout, so + # it returns `aggregated + f(aggregated)`. We feed only the function + # delta `f(aggregated)` into the n-stream BDA so it does not double-count + # the residual that mHC owns. + if self.config.fp32_residual_connection and aggregated.dtype != layer_output.dtype: + aggregated = aggregated.to(layer_output.dtype) + layer_output_with_bias = (layer_output - aggregated, None) + dropout_prob = 0.0 + bias_dropout_fusion = False + else: + layer_output_with_bias, context, dropout_prob, bias_dropout_fusion = fast_path_result + + layer_output = layer_output_with_bias[0] + # Sanity check: this contract requires the branch output to preserve shape; + # any mismatch indicates a future layer type is breaking the residual assumption + # and would silently corrupt the n-stream state. + if layer_output.shape != aggregated.shape: + raise RuntimeError( + "HyperConnectionHybridLayer requires wrapped branches to preserve " + f"hidden-state shape. Got {tuple(layer_output.shape)} from wrapped branch " + f"vs {tuple(aggregated.shape)} input." + ) + is_last_in_recompute_block = bool( + mhc_recompute_manager is not None + and getattr(mhc_recompute_manager, "is_last_layer_in_recompute_block", False) + ) + mhc_bda_manager = None if is_last_in_recompute_block else mhc_recompute_manager + + hidden_states = self.hyper_connection.fused_h_res_h_post_bda( + h_res, + residual, + h_post, + layer_output_with_bias, + dropout_prob=dropout_prob, + training=self.training, + fused=bias_dropout_fusion, + manager=mhc_bda_manager, + ) + # In `HyperConnectionTransformerLayer` the n-stream output stays in compute + # dtype because the post-attention `x` is in compute dtype. In the hybrid + # wrapper, `layer_delta` may be fp32 (when `fp32_residual_connection=True` + # or an inner layer upcasts), so `fused_h_res_h_post_bda`'s `output.to(x.dtype)` + # would leave the result in fp32 and silently propagate fp32 n-stream + # hidden states to every subsequent layer (~2x activation memory). Restore + # the compute-dtype contract here. + if ( + self.config.fp32_residual_connection + and self.config.params_dtype is not None + and hidden_states.dtype != self.config.params_dtype + ): + hidden_states = hidden_states.to(self.config.params_dtype) + return hidden_states, context + + class HybridStack(MegatronModule): """ Constructor for the HybridStack class. @@ -86,6 +565,7 @@ class HybridStack(MegatronModule): pg_collection (ProcessGroupCollection): the required model communication process groups to use. is_mtp_layer (bool, optional): whether this is an MTP layer. Defaults to False. + mtp_layer_number (int, optional): enclosing MTP depth for logging nested MTP metrics. """ def __init__( @@ -101,6 +581,7 @@ def __init__( dtype=None, pg_collection: ProcessGroupCollection = None, is_mtp_layer: bool = False, + mtp_layer_number: Optional[int] = None, name: str | None = None, ) -> None: """ @@ -112,18 +593,19 @@ def __init__( self.post_layer_norm = post_layer_norm self.post_process = post_process self.is_mtp_layer = is_mtp_layer + self.mtp_layer_number = mtp_layer_number assert pg_collection is not None, "pg_collection must be provided for HybridStack" self.pp_group = pg_collection.pp self.tp_group = pg_collection.tp - self.cp_group = pg_collection.cp - self.tp_cp_group = pg_collection.tp_cp # Required for pipeline parallel schedules self.input_tensor = None self.pg_collection = pg_collection + # Lazily populated mHC recompute layout cache (deterministic from config + # and num_layers); see `_build_mhc_recompute_layer_plan`. self._mhc_block_end_plan: Optional[List[bool]] = None assert layer_type_list is not None, ( @@ -131,26 +613,6 @@ def __init__( "--hybrid-layer-pattern by HybridModel." ) self.layer_type_list = layer_type_list - self._cp_layout_manager = None - if self.cp_group.size() > 1: - layer_layouts = tuple( - ( - self.config.attention_cp_layout - if layer_type in LayerSymbols.ATTENTION_LAYERS - else self.config.linear_cp_layout - ) - for layer_type in self.layer_type_list - ) - self._cp_layout_manager = ContextParallelLayoutManager( - layer_layouts=layer_layouts, - boundary_layout=self.config.linear_cp_layout, - sequence_parallel=self.config.sequence_parallel, - cp_group=self.cp_group, - tp_group=self.tp_group, - tp_cp_group=self.tp_cp_group, - ) - if getattr(self.config, "mla_down_proj_fusion", False): - submodules = self._fuse_mla_down_proj(submodules) # Build layers from the pre-selected segment self.layers = nn.ModuleList() @@ -203,6 +665,41 @@ def __init__( is_mtp_layer=is_mtp_layer, add_layer_offset=False, pp_layer_offset=pp_layer_offset, + name=(name + f".layers.{i}") if name is not None else None, + ) + elif layer_type == LayerSymbols.CSA: + # DSv4 Compressed Sparse Attention (compress_ratio fixed by the spec). + layer = build_module( + submodules.csa_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + ) + elif layer_type == LayerSymbols.HCA: + # DSv4 Heavily Compressed Attention (compress_ratio fixed by the spec). + layer = build_module( + submodules.hca_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + ) + elif layer_type == LayerSymbols.WINDOW: + # DSv4 sliding-window-only attention (compress_ratio=0 fixed by the spec; + # no compressor / no top-k indexer — attends only within csa_window_size). + layer = build_module( + submodules.window_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, ) elif layer_type == LayerSymbols.MLP: layer = build_module( @@ -224,27 +721,29 @@ def __init__( name=(name + f".layers.{i}") if name is not None else None, ) elif layer_type == LayerSymbols.GDN: - gdn_layer_spec = submodules.gdn_layer - if self.config.experimental_attention_variant == "gdn2": - # 'G' layers build the GDN2 variant when the gdn2 experimental - # attention variant is selected. - from megatron.core.ssm.gated_delta_net import GatedDeltaNet2 - - gdn_layer_spec = copy.deepcopy(gdn_layer_spec) - gdn_layer_spec.submodules.self_attention.module = GatedDeltaNet2 layer = build_module( - gdn_layer_spec, + submodules.gdn_layer, config=self.config, layer_number=layer_number, pg_collection=pg_collection, # Set to False as we do not want to change offset. add_layer_offset=False, - pp_layer_offset=pp_layer_offset, + name=(name + f".layers.{i}") if name is not None else None, + ) + elif layer_type == LayerSymbols.KDA: + layer = build_module( + submodules.kda_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + add_layer_offset=False, name=(name + f".layers.{i}") if name is not None else None, ) else: raise ValueError("unexpected layer_type") - if self.config.enable_mhc_connections: + if self.is_mtp_layer and self.mtp_layer_number is not None: + self._set_mtp_layer_number_for_moe_metrics(layer, self.mtp_layer_number) + if self.config.enable_hyper_connections: layer = HyperConnectionHybridLayer(config=self.config, layer=layer) self.layers.append(layer) @@ -262,8 +761,12 @@ def __init__( eps=self.config.layernorm_epsilon, ) - if self.config.enable_mhc_connections and self.post_process and not self.is_mtp_layer: - hc_mult = self.config.mhc_num_residual_streams + # Skip hc_head_* params inside the nested MTP HybridStack — `forward()` + # no longer calls `learned_output_contract` there (MTP owns that), so these + # params would be orphaned and break DDP's per-param grad-ready accounting + # with a `len(per_param_grad_ready_counts) != len(params)` AssertionError. + if self.config.enable_hyper_connections and self.post_process and not self.is_mtp_layer: + hc_mult = self.config.num_residual_streams hc_dim = self.config.hidden_size * hc_mult self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim))) self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult))) @@ -274,25 +777,15 @@ def __init__( setattr(self.hc_head_base, 'sequence_parallel', True) setattr(self.hc_head_scale, 'sequence_parallel', True) - def _fuse_mla_down_proj(self, submodules: HybridStackSubmodules) -> HybridStackSubmodules: - # Avoid modifying the original object so users don't get surprised about their `submodules` - # being modified underneath them. - submodules = copy.deepcopy(submodules) - mla_spec = submodules.mla_layer - # We always fuse the input layernorm because Hybrid always uses TransformerEngine. - mla_spec.submodules.input_layernorm = IdentityOp - mla_spec.submodules.self_attention.module = FusedMLASelfAttention - mla_spec.submodules.self_attention.submodules.linear_qkv_down_proj = ( - TELayerNormColumnParallelLinear - ) - mla_spec.submodules.self_attention.submodules.linear_q_down_proj = None - mla_spec.submodules.self_attention.submodules.linear_kv_down_proj = None - mla_spec.submodules.sharded_state_dict_keys_map = { - "self_attention.linear_q_down_proj.layer_norm_": "input_layernorm.", - "self_attention.linear_kv_down_proj.layer_norm_": "input_layernorm.", - "self_attention.linear_qkv_down_proj.layer_norm_": "input_layernorm.", - } - return submodules + @staticmethod + def _set_mtp_layer_number_for_moe_metrics( + layer: torch.nn.Module, mtp_layer_number: int + ) -> None: + """Tell nested MTP MoE routers which MTP depth they belong to for logging.""" + for module in layer.modules(): + router = getattr(module, "router", None) + if router is not None and getattr(router, "is_mtp_layer", False): + router.mtp_layer_number = mtp_layer_number def set_input_tensor(self, input_tensor: Tensor): """Set input tensor to be used instead of forward()'s input. @@ -306,62 +799,66 @@ def set_input_tensor(self, input_tensor: Tensor): def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]: """ - Returns the recurrent mixer's conv and SSM state shapes per input sequence - if this block contains Mamba or GDN layers (this may not be the case with PP > 1). + Returns the Mamba conv and ssm states shapes per input sequence + if this block contains Mamba layers (this may not be the case with PP > 1). """ for layer_type, layer in zip(self.layer_type_list, self.layers): if layer_type == LayerSymbols.MAMBA: return layer.mamba_state_shapes_per_request() - if layer_type == LayerSymbols.GDN: - if hasattr(layer, 'mamba_state_shapes_per_request'): - state_shapes = layer.mamba_state_shapes_per_request() - if state_shapes is not None: - return state_shapes - return layer.self_attention.mamba_state_shapes_per_request() return None def _compute_mhc_block_end_plan(self) -> List[bool]: - """Compute deterministic per-layer mHC recompute block boundaries.""" + """Compute per-layer block-end markers (deterministic from config).""" num_layers = len(self.layers) - block_ends: List[bool] = [False] * num_layers + is_recompute_block_end: List[bool] = [False] * num_layers if num_layers == 0: - return block_ends - - layers_per_block = self.config.mhc_recompute_layer_num - for layer_idx in range(num_layers): - is_last_in_stack = layer_idx == num_layers - 1 - block_ends[layer_idx] = is_last_in_stack or ( - layers_per_block is not None and (layer_idx + 1) % layers_per_block == 0 - ) - return block_ends + return is_recompute_block_end + mhc_recompute_layer_num = self.config.mhc_recompute_layer_num + for l_no in range(num_layers): + is_last_in_stack = l_no == num_layers - 1 + is_last_in_recompute_block = is_last_in_stack + if mhc_recompute_layer_num is not None: + is_last_in_recompute_block = is_last_in_stack or ( + (l_no + 1) % mhc_recompute_layer_num == 0 + ) + is_recompute_block_end[l_no] = is_last_in_recompute_block + return is_recompute_block_end def _build_mhc_recompute_layer_plan( self, use_mhc_recompute: bool ) -> Tuple[List[Optional[CheckpointWithoutOutputManager]], List[bool]]: - """Build single-use recompute managers for this forward pass.""" + """Pre-build per-layer MHC recompute managers and block-end markers. + + The block-end plan is deterministic from config and cached on the + instance; only the per-block ``CheckpointWithoutOutputManager`` + instances are allocated fresh per forward pass (managers are single-use). + This mirrors the caching scheme used by ``TransformerBlock``. + """ num_layers = len(self.layers) if not use_mhc_recompute or num_layers == 0: return [None] * num_layers, [False] * num_layers if self._mhc_block_end_plan is None: self._mhc_block_end_plan = self._compute_mhc_block_end_plan() - block_ends = self._mhc_block_end_plan + is_recompute_block_end = self._mhc_block_end_plan layer_managers: List[Optional[CheckpointWithoutOutputManager]] = [None] * num_layers - manager = CheckpointWithoutOutputManager() - for layer_idx in range(num_layers): - layer_managers[layer_idx] = manager - if block_ends[layer_idx] and layer_idx != num_layers - 1: - manager = CheckpointWithoutOutputManager() - return layer_managers, block_ends + mhc_manager = CheckpointWithoutOutputManager() + for l_no in range(num_layers): + layer_managers[l_no] = mhc_manager + if is_recompute_block_end[l_no] and l_no != num_layers - 1: + mhc_manager = CheckpointWithoutOutputManager() + return layer_managers, is_recompute_block_end @staticmethod def _finalize_mhc_recompute_layer( - manager: Optional[CheckpointWithoutOutputManager], hidden_states: Tensor, is_block_end: bool + mhc_manager: Optional[CheckpointWithoutOutputManager], + hidden_states: Tensor, + is_last_in_recompute_block: bool, ) -> None: - """Finalize the current mHC recompute block when its last layer finishes.""" - if manager is not None and is_block_end: - manager.discard_all_outputs_and_register_unified_recompute(hidden_states) + """Finalize MHC recompute state for the current layer when a block ends.""" + if mhc_manager is not None and is_last_in_recompute_block: + mhc_manager.discard_all_outputs_and_register_unified_recompute(hidden_states) def forward( self, @@ -373,7 +870,8 @@ def forward( inference_params: Optional[BaseInferenceContext] = None, packed_seq_params: Optional[PackedSeqParams] = None, padding_mask=None, - ): + input_ids: Optional[Tensor] = None, + ) -> Union[Tensor, Tuple[Tensor, Tensor]]: """ Forward function of the HybridStack class. @@ -389,15 +887,15 @@ def forward( rotary_pos_emb (Tensor, optional): the rotary positional embeddings. Defaults to None. Returns: - Tensor: the output tensor. + Tensor in the common case. A 2-tuple ``(hidden_states, mhc_multistream)`` ONLY when + ``enable_hyper_connections and post_process and mtp_num_layers > 0 and not + is_mtp_layer`` — the extra element is the pre-contraction multi-stream tensor that + MTP's ``_concat_embeddings`` consumes. Callers (e.g. ``HybridModel.forward``) must + handle both; pipeline send/recv only ever transfers the contracted ``hidden_states``. """ inference_context = deprecate_inference_params(inference_context, inference_params) - cp_layout_state = None - if self._cp_layout_manager is not None: - cp_layout_state = self._cp_layout_manager.build_forward_state(packed_seq_params) - if not self.pre_process: # See set_input_tensor() hidden_states = self.input_tensor @@ -406,9 +904,13 @@ def forward( if isinstance(hidden_states, WrappedTensor): hidden_states = hidden_states.unwrap() - if self.config.enable_mhc_connections and self.pre_process and not self.is_mtp_layer: + # Skip input_expand inside MTP nested HybridStack: when mHC + MTP, the outer + # decoder hands in already-multi-stream hidden_states via mhc_multistream + # (see multi_token_prediction.py _concat_embeddings), so expanding again would + # produce [s, b, n*(n*h)] instead of [s, b, n*h] and break HC mapping_proj. + if self.config.enable_hyper_connections and self.pre_process and not self.is_mtp_layer: hidden_states = HyperConnectionModule.input_expand( - hidden_states, self.config.mhc_num_residual_streams + hidden_states, self.config.num_residual_streams ) if inference_context and inference_context.is_static_batching(): @@ -460,11 +962,13 @@ def get_inner_quant_context(config, layer_number): use_mhc_recompute = ( self.training - and self.config.enable_mhc_connections + and self.config.enable_hyper_connections and self.config.recompute_granularity == 'selective' and "mhc" in self.config.recompute_modules ) - mhc_layer_managers, mhc_block_ends = self._build_mhc_recompute_layer_plan(use_mhc_recompute) + mhc_layer_managers, mhc_is_last_in_recompute_block = self._build_mhc_recompute_layer_plan( + use_mhc_recompute + ) with outer_fp8_context: if self.config.recompute_granularity == 'full' and self.training: @@ -478,23 +982,21 @@ def get_inner_quant_context(config, layer_number): attention_bias=None, packed_seq_params=packed_seq_params, padding_mask=padding_mask, + input_ids=input_ids, use_inner_quantization_context=(use_inner_fp8_context or use_fp4_context), - cp_layout_state=cp_layout_state, ) else: - for layer_idx, layer in enumerate(self.layers): - layer_packed_seq_params = packed_seq_params - if cp_layout_state is not None: - hidden_states, layer_packed_seq_params = cp_layout_state.prepare_layer( - layer_idx, hidden_states - ) + for l_no, layer in enumerate(self.layers): # Layers have 1-indexed layer numbers attribute. inner_quant_context = get_inner_quant_context( self.config, layer.layer_number - 1 ) - mhc_manager = mhc_layer_managers[layer_idx] + + mhc_manager = mhc_layer_managers[l_no] if mhc_manager is not None: - mhc_manager.is_last_layer_in_recompute_block = mhc_block_ends[layer_idx] + mhc_manager.is_last_layer_in_recompute_block = ( + mhc_is_last_in_recompute_block[l_no] + ) with inner_quant_context: if isinstance(layer, (TransformerLayer, HyperConnectionHybridLayer)): @@ -504,9 +1006,11 @@ def get_inner_quant_context(config, layer_number): inference_context=inference_context, rotary_pos_emb=rotary_pos_emb, sequence_len_offset=sequence_len_offset, - packed_seq_params=layer_packed_seq_params, + packed_seq_params=packed_seq_params, padding_mask=padding_mask, ) + if input_ids is not None: + layer_kwargs["input_ids"] = input_ids if mhc_manager is not None and isinstance( layer, HyperConnectionHybridLayer ): @@ -517,7 +1021,7 @@ def get_inner_quant_context(config, layer_number): hidden_states=hidden_states, attention_mask=attention_mask, inference_context=inference_context, - packed_seq_params=layer_packed_seq_params, + packed_seq_params=packed_seq_params, ) # The attention layer (currently a simplified transformer layer) @@ -525,17 +1029,27 @@ def get_inner_quant_context(config, layer_number): # for cross-attention, and is not needed in our model. if isinstance(hidden_states, tuple): hidden_states = hidden_states[0] - if cp_layout_state is not None: - hidden_states = cp_layout_state.finalize_layer(layer_idx, hidden_states) self._finalize_mhc_recompute_layer( - manager=mhc_manager, + mhc_manager=mhc_manager, hidden_states=hidden_states, - is_block_end=mhc_block_ends[layer_idx], + is_last_in_recompute_block=mhc_is_last_in_recompute_block[l_no], ) + # When mHC + MTP, save the pre-contraction multi-stream tensor for MTP input. + # MTP's _concat_embeddings mHC branch expects [s, b, n*h] (multi-stream), while + # the contracted hidden_states is [s, b, h]. Mirrors transformer_block.py:948-988. + # Only the OUTER decoder stack does this; nested MTP stacks (is_mtp_layer=True) + # must keep returning a single Tensor so MTP's _postprocess receives the right + # type for learned_output_contract. + # On the final stage of a (non-MTP) stack with mHC active, capture the pre-contraction + # multi-stream tensor for MTP's `_concat_embeddings` (only meaningful when MTP layers + # exist, i.e. mtp_num_layers > 0), THEN contract the streams. Combining capture and + # contraction avoids repeating the condition. Nested MTP HybridStacks (is_mtp_layer=True) + # must NOT contract here — MTP's own `_postprocess` calls learned_output_contract + + # final_layernorm itself, so doing it here would double-collapse the multi-stream tensor. mhc_multistream = None - if self.config.enable_mhc_connections and self.post_process and not self.is_mtp_layer: + if self.config.enable_hyper_connections and self.post_process and not self.is_mtp_layer: if (self.config.mtp_num_layers or 0) > 0: mhc_multistream = hidden_states hidden_states = learned_output_contract( @@ -543,7 +1057,7 @@ def get_inner_quant_context(config, layer_number): self.hc_head_fn, self.hc_head_base, self.hc_head_scale, - self.config.mhc_num_residual_streams, + self.config.num_residual_streams, self.config.layernorm_epsilon, ) @@ -626,7 +1140,7 @@ def sharded_state_dict( make_sharded_tensors_for_checkpoint( local_state_dict, prefix, - sharded_offsets=sharded_offsets, + sharded_offsets=sharded_offsets or (), tp_group=self.tp_group, dp_cp_group=metadata['dp_cp_group'], ) diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index bcd429dbfd7..6aa866ba158 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -16,15 +16,20 @@ class Symbols: MAMBA = "M" GDN = 'G' + KDA = 'K' ATTENTION = "*" DS_ATTENTION = "D" MLA = "+" + CSA = "C" # DSv4 Compressed Sparse Attention (compress_ratio=4) + HCA = "H" # DSv4 Heavily Compressed Attention (compress_ratio=128) + WINDOW = "W" # DSv4 sliding-window-only attention (compress_ratio=0; no compressor/indexer) MLP = "-" MOE = 'E' PIPE = '|' MTP_SEPARATOR = "/" - VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, MLA, MLP, MOE} - ATTENTION_LAYERS = {ATTENTION, DS_ATTENTION, MLA} + VALID_LAYERS = {MAMBA, GDN, KDA, ATTENTION, DS_ATTENTION, MLA, CSA, HCA, WINDOW, MLP, MOE} + # MLA-based attention layers (incompatible with standard '*' attention in one model). + MLA_ATTENTION = {MLA, DS_ATTENTION, CSA, HCA, WINDOW} @classmethod def name_sorted_valid_layer_symbols(cls) -> list[str]: @@ -175,10 +180,10 @@ def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]: Examples: >>> get_hybrid_layer_counts("M*M*") - {'*': 2, 'G': 0, 'D': 0, 'M': 2, '-': 0, 'E': 0} + {'*': 2, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'K': 0, 'M': 2, '+': 0, '-': 0, 'E': 0, 'W': 0} >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") - {'*': 1, 'G': 0, 'D': 0, 'M': 8, '-': 4, 'E': 0} + {'*': 1, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'K': 0, 'M': 8, '+': 0, '-': 4, 'E': 0, 'W': 0} """ parsed = parse_hybrid_pattern(pattern) counts = {symbol: 0 for symbol in Symbols.name_sorted_valid_layer_symbols()} @@ -294,9 +299,12 @@ def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False) f"Valid symbols are: {valid_chars}" ) - # Disallow Attention + MLA/DSA hybridity. - if Symbols.ATTENTION in pattern and (Symbols.DS_ATTENTION in pattern or Symbols.MLA in pattern): - raise ValueError("Not supported to have both Attention and MLA/DSA in one model") + # Disallow standard Attention ('*') mixed with any MLA-based attention (+/D/C/H/W). + # MLA variants (DSA / CSA / HCA / Window) may freely coexist with each other. + if Symbols.ATTENTION in pattern and any(s in pattern for s in Symbols.MLA_ATTENTION): + raise ValueError( + "Not supported to have both Attention and MLA/DSA/CSA/HCA/Window in one model" + ) def validate_segment_layers(segment: str) -> List[str]: @@ -322,9 +330,11 @@ def validate_segment_layers(segment: str) -> List[str]: f"one of {Symbols.VALID_LAYERS}" ) - # Disallow Attention + MLA/DSA hybridity. - if Symbols.ATTENTION in segment and (Symbols.DS_ATTENTION in segment or Symbols.MLA in segment): - raise ValueError("Not supported to have both Attention and MLA/DSA in one model") + # Disallow standard Attention ('*') mixed with any MLA-based attention (+/D/C/H/W). + if Symbols.ATTENTION in segment and any(s in segment for s in Symbols.MLA_ATTENTION): + raise ValueError( + "Not supported to have both Attention and MLA/DSA/CSA/HCA/Window in one model" + ) return layer_type_list diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index f0be7ef891b..c09d6887b33 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from functools import partial from megatron.core.extensions.transformer_engine import ( @@ -15,10 +15,11 @@ get_moe_module_spec, ) from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules -from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules -from megatron.core.ssm.gated_delta_product import ( - GatedDeltaProductMixer, - GatedDeltaProductMixerSubmodules, +from megatron.core.ssm.gated_delta_net import ( + GatedDeltaNet, + GatedDeltaNetSubmodules, + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, ) from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules @@ -81,8 +82,9 @@ submodules=MultiTokenPredictionLayerSubmodules( enorm=TENorm, hnorm=TENorm, - # Hybrid MTP selects the combined projection normally and - # per-stream projections when mHC is enabled. + # Both projection forms are populated so the layer can switch on + # `config.enable_hyper_connections` at runtime: mHC=True uses + # `e_proj`+`h_proj` per-stream, mHC=False uses fused `eh_proj`. eh_proj=TEColumnParallelLinear, e_proj=TEColumnParallelLinear, h_proj=TEColumnParallelLinear, @@ -95,19 +97,6 @@ ) -def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): - return ModuleSpec( - module=MambaLayer, - submodules=MambaLayerSubmodules( - mixer=ModuleSpec( - module=GatedDeltaProductMixer, - submodules=GatedDeltaProductMixerSubmodules(in_proj=in_proj, out_proj=out_proj), - ), - mamba_bda=get_bias_dropout_add, - ), - ) - - hybrid_stack_spec = ModuleSpec( module=HybridStack, submodules=HybridStackSubmodules( @@ -137,6 +126,22 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): self_attn_bda=get_bias_dropout_add, ), ), + kda_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, + self_attention=ModuleSpec( + module=KimiDeltaAttention, + submodules=KimiDeltaAttentionSubmodules( + in_proj=TEColumnParallelLinear, + beta_proj=TEColumnParallelLinear, + out_norm=TENorm, + out_proj=TERowParallelLinear, + ), + ), + self_attn_bda=get_bias_dropout_add, + ), + ), # Started with spec from gpt_layer_specs.py (with MLP removed) # Using the TE spec because we had problems getting the non-TE spec # working @@ -204,6 +209,7 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): linear_kv_down_proj=TELinear, linear_kv_up_proj=TEColumnParallelLinear, core_attention=TEDotProductAttention, + linear_gate=TEColumnParallelLinear, linear_proj=TERowParallelLinear, q_layernorm=IdentityOp, kv_layernorm=IdentityOp, @@ -238,22 +244,6 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): ) -gated_delta_product_stack_spec = ModuleSpec( - module=HybridStack, - submodules=HybridStackSubmodules( - mamba_layer=_get_gated_delta_product_mamba_layer_spec( - TELayerNormColumnParallelLinear, TERowParallelLinear - ), - gdn_layer=hybrid_stack_spec.submodules.gdn_layer, - attention_layer=hybrid_stack_spec.submodules.attention_layer, - dsa_layer=hybrid_stack_spec.submodules.dsa_layer, - mlp_layer=hybrid_stack_spec.submodules.mlp_layer, - moe_layer=hybrid_stack_spec.submodules.moe_layer, - mtp_block_spec=hybrid_stack_spec.submodules.mtp_block_spec, - ), -) - - hybrid_inference_stack_spec = ModuleSpec( module=HybridStack, submodules=HybridStackSubmodules( @@ -270,20 +260,6 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): mamba_bda=get_bias_dropout_add, ), ), - gdn_layer=ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - self_attention=ModuleSpec( - module=GatedDeltaNet, - submodules=GatedDeltaNetSubmodules( - in_proj=InferenceLayerNormColumnParallelLinear, - out_norm=TENorm, - out_proj=InferenceRowParallelLinear, - ), - ), - self_attn_bda=get_bias_dropout_add, - ), - ), # Started with spec from gpt_layer_specs.py (with MLP removed) # Using the TE spec because we had problems getting the non-TE spec # working @@ -337,28 +313,6 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): self_attn_bda=get_bias_dropout_add, ), ), - mla_layer=ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - input_layernorm=TENorm, - self_attention=ModuleSpec( - module=MLASelfAttention, - params={"attn_mask_type": AttnMaskType.causal}, - submodules=MLASelfAttentionSubmodules( - linear_q_proj=TEColumnParallelLinear, - linear_q_down_proj=TELinear, - linear_q_up_proj=TEColumnParallelLinear, - linear_kv_down_proj=TELinear, - linear_kv_up_proj=TEColumnParallelLinear, - core_attention=TEDotProductAttention, - linear_proj=InferenceRowParallelLinear, - q_layernorm=IdentityOp, - kv_layernorm=IdentityOp, - ), - ), - self_attn_bda=get_bias_dropout_add, - ), - ), # Started with spec from gpt_layer_specs.py # Using the TE spec because we had problems getting the non-TE spec # working @@ -391,7 +345,8 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): submodules=MultiTokenPredictionLayerSubmodules( enorm=TENorm, hnorm=TENorm, - # Keep both projection forms available for Hybrid MTP. + # Populate both projection forms so the layer can switch on + # `config.enable_hyper_connections` at runtime. eh_proj=InferenceColumnParallelLinear, e_proj=InferenceColumnParallelLinear, h_proj=InferenceColumnParallelLinear, @@ -406,24 +361,59 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): ) -gated_delta_product_inference_stack_spec = ModuleSpec( - module=HybridStack, - submodules=HybridStackSubmodules( - mamba_layer=_get_gated_delta_product_mamba_layer_spec( - InferenceLayerNormColumnParallelLinear, InferenceRowParallelLinear - ), - gdn_layer=hybrid_inference_stack_spec.submodules.gdn_layer, - attention_layer=hybrid_inference_stack_spec.submodules.attention_layer, - dsa_layer=hybrid_inference_stack_spec.submodules.dsa_layer, - mlp_layer=hybrid_inference_stack_spec.submodules.mlp_layer, - moe_layer=hybrid_inference_stack_spec.submodules.moe_layer, - mtp_block_spec=hybrid_inference_stack_spec.submodules.mtp_block_spec, - ), -) - - # Backward-compatible aliases mamba_stack_spec = hybrid_stack_spec mamba_inference_stack_spec = hybrid_inference_stack_spec -gdp_stack_spec = gated_delta_product_stack_spec -gdp_inference_stack_spec = gated_delta_product_inference_stack_spec + + +def hybrid_dsv4_stack_spec(config): + """Config-aware hybrid stack spec whose ``D`` (DS_ATTENTION) layer runs the DSv4 + ``CompressedSparseAttention`` (CSA/HCA + ``CSAIndexer``), identical to the GPT + ``dsv4_hybrid`` path, instead of the legacy ``DSAttention``. + + The default ``hybrid_stack_spec`` wires the ``D`` layer to ``AbsorbedMLASelfAttention + + DSAttention`` (DS3.2-style sparse attention with no CSA/HCA compression). To run real + DSv4 on HybridModel — and to be numerically equivalent to a GPT ``dsv4_hybrid`` + attention layer — we reuse GPT's own ``get_dsv4_hybrid_module_spec_for_backend`` + (which is config-aware, e.g. picks the qk-layernorm form from ``config``) so the two + model paths build the *same* attention module. Selected via + ``--spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_dsv4_stack_spec``; + ``hybrid_builder`` invokes this function with ``config`` when the spec is callable. + """ + import dataclasses + + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + _get_backend_spec_provider, + get_dsv4_hybrid_module_spec_for_backend, + ) + + backend = _get_backend_spec_provider(config) + dsv4_attention = get_dsv4_hybrid_module_spec_for_backend(config, backend) + + def _wrap_dsv4_layer(compress_ratio=None): + # Wrap the DSv4 attention in a hybrid TransformerLayer. When compress_ratio is given + # (the 'C'/'H'/'W' layer symbols), bake it into the attention params so the layer uses a + # fixed CSA(4)/HCA(128)/window-only(0) ratio regardless of csa_compress_ratios; otherwise + # the layer reads its ratio from config.csa_compress_ratios (array-driven 'D' / GPT-parity + # path). compress_ratio=0 builds neither the compressor nor the top-k indexer, so the + # 'W' layer reuses the entire CSA/HCA code path as pure sliding-window attention. + attn = dsv4_attention + if compress_ratio is not None: + attn = dataclasses.replace( + dsv4_attention, params={**dsv4_attention.params, "compress_ratio": compress_ratio} + ) + return ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, self_attention=attn, self_attn_bda=get_bias_dropout_add + ), + ) + + submodules = dataclasses.replace( + hybrid_stack_spec.submodules, + dsa_layer=_wrap_dsv4_layer(), # 'D': array-driven (or window) DSv4 attention + csa_layer=_wrap_dsv4_layer(compress_ratio=4), # 'C': CSA + hca_layer=_wrap_dsv4_layer(compress_ratio=128), # 'H': HCA + window_layer=_wrap_dsv4_layer(compress_ratio=0), # 'W': sliding-window-only + ) + return ModuleSpec(module=HybridStack, submodules=submodules) diff --git a/megatron/core/msc_utils.py b/megatron/core/msc_utils.py index c7dd7dc07d6..1764ed756d9 100644 --- a/megatron/core/msc_utils.py +++ b/megatron/core/msc_utils.py @@ -55,6 +55,13 @@ def __setstate__(self, state): MultiStorageClientFeature = _FeatureFlag(default=False) +def open_file(*args, **kwargs): + """Open a path through MSC when enabled, otherwise use the builtin open.""" + if MultiStorageClientFeature.is_enabled(): + return MultiStorageClientFeature.import_package().open(*args, **kwargs) + return open(*args, **kwargs) + + class MaybeMultiStorageClient: """ Helper class to use MultiStorageClient @@ -108,4 +115,4 @@ def __dir__(self): maybe_msc = MaybeMultiStorageClient() -__all__ = ['MultiStorageClientFeature', 'maybe_msc'] +__all__ = ['MultiStorageClientFeature', 'maybe_msc', 'open_file'] diff --git a/megatron/core/optimizer/qk_clip.py b/megatron/core/optimizer/qk_clip.py index e1c9133261a..e50d2d86db5 100644 --- a/megatron/core/optimizer/qk_clip.py +++ b/megatron/core/optimizer/qk_clip.py @@ -3,6 +3,29 @@ import torch from megatron.core import parallel_state +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.transformer.attention import Attention + + +def _iter_qk_clip_modules(model_chunk): + """Yield attention modules that need QK clipping. + + Preserve the legacy traversal for other model types. HybridModel may own + nested MTP stacks, so discover its attention modules recursively. + """ + + model_module = model_chunk.module.module + if isinstance(model_module, HybridModel): + yield from ( + module + for module in model_module.modules() + if isinstance(module, Attention) and hasattr(module, 'clip_qk') + ) + return + + for transformer_layer in model_module.decoder.layers: + if hasattr(transformer_layer.self_attention, 'clip_qk'): + yield transformer_layer.self_attention def clip_qk(model, log_max_only=False) -> float: @@ -20,33 +43,25 @@ def clip_qk(model, log_max_only=False) -> float: with torch.no_grad(): log_max_attention_logit = 0 for model_chunk in model: - for transformer_layer in model_chunk.module.module.decoder.layers: - if hasattr(transformer_layer.self_attention, 'clip_qk'): - if ( - transformer_layer.self_attention.core_attention.current_max_attn_logits - is None - ): - continue - torch.distributed.all_reduce( - transformer_layer.self_attention.core_attention.current_max_attn_logits, - op=torch.distributed.ReduceOp.MAX, - group=parallel_state.get_data_parallel_group(with_context_parallel=True), - ) - log_max_attention_logit = max( - log_max_attention_logit, - torch.max( - transformer_layer.self_attention.core_attention.current_max_attn_logits - ).item(), - ) - if not log_max_only: - transformer_layer.self_attention.clip_qk() - else: - # When qk-clip is disabled, clip_qk() is not called and - # would otherwise never reset current_max_attn_logits. - # Reset it here so the logged value reflects the current - # step and stale references are not retained. - transformer_layer.self_attention.core_attention.current_max_attn_logits = ( - None - ) + for attention in _iter_qk_clip_modules(model_chunk): + current_max_attn_logits = attention.core_attention.current_max_attn_logits + if current_max_attn_logits is None: + continue + torch.distributed.all_reduce( + current_max_attn_logits, + op=torch.distributed.ReduceOp.MAX, + group=parallel_state.get_data_parallel_group(with_context_parallel=True), + ) + log_max_attention_logit = max( + log_max_attention_logit, torch.max(current_max_attn_logits).item() + ) + if not log_max_only: + attention.clip_qk() + else: + # When qk-clip is disabled, clip_qk() is not called and + # would otherwise never reset current_max_attn_logits. + # Reset it here so the logged value reflects the current + # step and stale references are not retained. + attention.core_attention.current_max_attn_logits = None return log_max_attention_logit diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 52e7dbaa516..63c68d9881f 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -26,6 +26,7 @@ class PackedSeqParams: seq_idx: Tensor = None tokens_per_sample: int = None pad_between_seqs: bool = None + cp_partition_mode: str = "zigzag" cp_scatter_cache: object = None def __post_init__(self): @@ -67,3 +68,12 @@ def __post_init__(self): .to(torch.int32) .unsqueeze(0) # Add a batch dimension ) + + +def resolve_cp_group( + static_cp_group: dist.ProcessGroup, packed_seq_params: PackedSeqParams = None +) -> dist.ProcessGroup: + """Return the packed-sequence CP group when present, otherwise the static group.""" + if packed_seq_params is not None and packed_seq_params.cp_group is not None: + return packed_seq_params.cp_group + return static_cp_group diff --git a/megatron/core/ssm/gated_delta_net/__init__.py b/megatron/core/ssm/gated_delta_net/__init__.py index e5a5b4e6a9f..ba139891c36 100644 --- a/megatron/core/ssm/gated_delta_net/__init__.py +++ b/megatron/core/ssm/gated_delta_net/__init__.py @@ -1,6 +1,6 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Gated Delta Net (GDN) family of layers. +"""Gated Delta Net (GDN) layer exports. This package replaces the former ``megatron/core/ssm/gated_delta_net.py`` module at the same import path; the names below preserve that module's public surface. @@ -9,6 +9,9 @@ from megatron.core.ssm.gated_delta_net.common import ( HAVE_FLA, GatedDeltaNetSubmodules, + _build_head_perm_for_split_sections, + _build_thd_cp_a2a_perm, + _split_tensor_factory, causal_conv1d, chunk_gated_delta_rule, get_parameter_local_cp, @@ -17,26 +20,25 @@ tensor_a2a_hp2cp, ) from megatron.core.ssm.gated_delta_net.gdn import GatedDeltaNet, torch_chunk_gated_delta_rule -from megatron.core.ssm.gated_delta_net.gdn2 import ( - HAVE_FLA_GDN2, - GatedDeltaNet2, - chunk_gdn2, - torch_chunk_gdn2, +from megatron.core.ssm.gated_delta_net.kda import ( + HAVE_FLA_KDA, + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, ) __all__ = [ "HAVE_FLA", - "HAVE_FLA_GDN2", + "HAVE_FLA_KDA", "GatedDeltaNet", - "GatedDeltaNet2", "GatedDeltaNetSubmodules", + "KimiDeltaAttention", + "KimiDeltaAttentionSubmodules", + "_split_tensor_factory", "causal_conv1d", "chunk_gated_delta_rule", - "chunk_gdn2", "get_parameter_local_cp", "l2norm", "tensor_a2a_cp2hp", "tensor_a2a_hp2cp", "torch_chunk_gated_delta_rule", - "torch_chunk_gdn2", ] diff --git a/megatron/core/ssm/gated_delta_net/common.py b/megatron/core/ssm/gated_delta_net/common.py index 63f2d244840..0483fec8fe4 100644 --- a/megatron/core/ssm/gated_delta_net/common.py +++ b/megatron/core/ssm/gated_delta_net/common.py @@ -1,20 +1,16 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh. # Some of this code was adopted from https://github.com/huggingface/transformers # This source code is licensed under the Apache license found in the # LICENSE file in the root directory of this source tree. -# pylint: disable=unused-import - -import logging from dataclasses import dataclass from functools import lru_cache -from typing import Callable, Optional, Protocol, Union +from typing import Optional, Protocol, Union import torch import torch.nn as nn -import torch.nn.functional as F from megatron.core.fp8_utils import get_fp8_align_size from megatron.core.inference.contexts import BaseInferenceContext @@ -38,29 +34,44 @@ make_sharded_tensors_for_checkpoint, sharded_state_dict_default, ) -from megatron.core.utils import nvtx_range_pop, nvtx_range_push try: from fla.modules.convolution import causal_conv1d from fla.modules.l2norm import l2norm + from fla.ops.cp import build_cp_context from fla.ops.gated_delta_rule import chunk_gated_delta_rule HAVE_FLA = True except ImportError: + build_cp_context = None causal_conv1d = None l2norm = None chunk_gated_delta_rule = None HAVE_FLA = False -logger = logging.getLogger(__name__) +__all__ = [ + "HAVE_FLA", + "GatedDeltaNetSubmodules", + "_GDNBase", + "_build_head_perm_for_split_sections", + "_build_thd_cp_a2a_perm", + "_split_tensor_factory", + "a2a_cp_to_hp", + "a2a_hp_to_cp", + "build_cp_context", + "causal_conv1d", + "chunk_gated_delta_rule", + "get_parameter_local_cp", + "l2norm", + "tensor_a2a_cp2hp", + "tensor_a2a_hp2cp", +] @dataclass class GatedDeltaNetSubmodules: - """ - Contains the module specs for the input linear, output norm, and output linear layers. - """ + """Module specs shared by GDN-family layers.""" in_proj: Union[ModuleSpec, type] = IdentityOp out_norm: Union[ModuleSpec, type] = IdentityOp @@ -68,12 +79,7 @@ class GatedDeltaNetSubmodules: class GatedDeltaRuleInterface(Protocol): - """ - Unified typing protocol for linear attention interfaces, compliant to upstream FLA interfaces. - - Only ``q``/``k``/``v``/``g`` are common to every kernel, and only as keywords: each - variant inserts its own gates after ``g`` (e.g., ``beta`` for GDN, ``b``/``w`` for GDN2). - """ + """Callable interface shared by GDN-family kernels.""" def __call__( self, @@ -81,6 +87,7 @@ def __call__( k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, + beta: torch.Tensor, *, scale: float | None = None, initial_state: torch.Tensor | None = None, @@ -92,16 +99,16 @@ def __call__( class _GDNBase(MegatronModule): - """Common base class for the Gated Delta Net (GDN) family of layers. + """Shared implementation for the GDN-family layers. - Hosts everything the GDN variants share: the fused input projection, causal - convolution on q/k/v, the CP all-to-all plumbing, the kernel-input preparation - skeleton, the gated output norm + projection, and sharded checkpointing. + Provides the projection, Q/K/V causal convolution, gated delta-rule parameters, + context-parallel layout handling, gated output normalization, output projection, + and sharded-checkpoint plumbing shared by head-wise GatedDeltaNet and channel-wise + Kimi Delta Attention. """ dt_bias_dim: int a_log_dim: int - in_proj_qkvg_dim: int in_proj_extra_dim: int in_proj_dim: int @@ -124,13 +131,14 @@ def __init__( *, name: str | None = None, cp_comm_type: str | None = None, - pp_layer_offset: int = 0, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, ): """ Args: config: The config of the model. submodules: Contains the module specs for the input and output linear layers. - layer_number: The layer number of this GDN layer. + layer_number: The layer number of this GDN-family layer. bias: Whether to use bias in the linear layers. conv_bias: Whether to use bias in the causal convolution. conv_init: The initialization range for the causal convolution weights. @@ -138,11 +146,13 @@ def __init__( A_init_range: The initialization range for the attention weights. pg_collection: The required process groups to use for tensor model parallel and context parallel. - name (str | None): module instance name passed top-down from its paranet module + name (str | None): Optional module path prefix used for child module names. cp_comm_type (Optional[str]): Accepted for TransformerLayer compatibility and ignored; GDN implements context parallelism with its own all-to-alls rather than the attention CP communication schemes. - pp_layer_offset: Offset of this pipeline stage's first global layer. + pp_layer_offset (Optional[int]): Pipeline layer offset forwarded by + TransformerLayer. Stored for MTP/TransformerLayer API compatibility. + is_mtp_layer (bool): Whether this module is inside an MTP prediction depth. """ if not HAVE_FLA: raise ImportError( @@ -154,16 +164,19 @@ def __init__( # Attributes from arguments self.layer_number = layer_number - self.pp_layer_offset = pp_layer_offset + self._pp_layer_offset = pp_layer_offset + self.is_mtp_layer = is_mtp_layer self.bias = bias self.conv_bias = conv_bias self.conv_init = conv_init assert A_init_range[0] >= 0 and A_init_range[1] >= A_init_range[0] self.A_init_range = A_init_range self.use_qk_l2norm = use_qk_l2norm - assert pg_collection is not None, "pg_collection must be provided for GatedDeltaNet" + assert pg_collection is not None, "pg_collection must be provided for a GDN-family layer" self.pg_collection = pg_collection self.tp_group = pg_collection.tp + # Static/max CP size from model construction. Runtime dynamic CP paths must resolve + # the effective group from packed_seq_params instead of using this value. self.cp_size = self.pg_collection.cp.size() self.tp_size = self.pg_collection.tp.size() self.sp_size = self.tp_size if config.sequence_parallel else 1 @@ -183,6 +196,21 @@ def __init__( self.qk_dim_local_tp = self.qk_dim // self.tp_size self.v_dim_local_tp = self.v_dim // self.tp_size + # Headwise CP shards heads over the CP group; chunkwise CP keeps heads local. + if self.config.linear_cp_mode == "headwise": + num_key_heads_per_tp = self.num_key_heads // self.tp_size + num_value_heads_per_tp = self.num_value_heads // self.tp_size + assert num_key_heads_per_tp % self.cp_size == 0, ( + f"GDN-family head-parallel CP requires the static (max) cp_size ({self.cp_size}) " + f"to evenly divide num_key_heads per TP rank ({num_key_heads_per_tp}); " + f"all runtime dynamic cp_size values divide the static one and so will also divide." + ) + assert num_value_heads_per_tp % self.cp_size == 0, ( + f"GDN-family head-parallel CP requires the static (max) cp_size ({self.cp_size}) " + f"to evenly divide num_value_heads per TP rank ({num_value_heads_per_tp}); " + f"all runtime dynamic cp_size values divide the static one and so will also divide." + ) + self.num_v_heads_local_tp = self.num_value_heads // self.tp_size self.num_k_heads_local_tp = self.num_key_heads // self.tp_size @@ -192,20 +220,21 @@ def __init__( "in_proj_extra_dim", "in_proj_split_names", "in_proj_split_sections", - "feat_dim_split", "gated_delta_rule", ) self._setup_variant_attrs() for attr in attrs_to_check: - assert getattr(self, attr, None) is not None, f"Attribute {attr} for GDN is not set" - # QK, V, gate, shared across all variants - self.in_proj_qkvg_dim = self.qk_dim * 2 + self.v_dim * 2 - self.in_proj_dim = self.in_proj_qkvg_dim + self.in_proj_extra_dim + assert hasattr(self, attr), f"Attribute {attr} for the GDN-family variant is not set" + assert ( + getattr(self, attr) is not None + ), f"Attribute {attr} for the GDN-family variant is not set" + # Full input projection width: q, k, v, output gate, and variant-specific gate features. + self.in_proj_dim = self.qk_dim * 2 + self.v_dim * 2 + self.in_proj_extra_dim if self.config.fp8: fp8_align_size = get_fp8_align_size(self.config.fp8_recipe) assert self.in_proj_dim % fp8_align_size == 0, ( - "For FP8, the innermost dimension of the GDN layer " + "For FP8, the innermost dimension of the GDN-family layer " "input projection output tensor must be a multiple of 16." ) self.in_proj = build_module( @@ -245,7 +274,9 @@ def __init__( self.dt_bias = nn.Parameter( torch.empty( - self.dt_bias_dim, dtype=self.config.params_dtype, device=torch.cuda.current_device() + self.dt_bias_dim, + dtype=getattr(self, "gate_params_dtype", self.config.params_dtype), + device=torch.cuda.current_device(), ) ) setattr(self.dt_bias, "tensor_model_parallel", True) @@ -253,7 +284,9 @@ def __init__( self.A_log = nn.Parameter( torch.empty( - self.a_log_dim, dtype=self.config.params_dtype, device=torch.cuda.current_device() + self.a_log_dim, + dtype=getattr(self, "gate_params_dtype", self.config.params_dtype), + device=torch.cuda.current_device(), ) ) setattr(self.A_log, "tensor_model_parallel", True) @@ -268,8 +301,10 @@ def __init__( ) self.recompute_norm_out = False self.norm_out_checkpoint = None - if self.config.recompute_granularity == "selective": + self.recompute_gdn = False + if self.config.recompute_granularity == "selective" and self.config.recompute_modules: self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules + self.recompute_gdn = "gdn" in self.config.recompute_modules self.out_proj = build_module( submodules.out_proj, @@ -285,17 +320,20 @@ def __init__( tp_group=self.pg_collection.tp, name=(name + ".out_proj") if name is not None else None, ) + # TODO: Packed sequence cu_seqlens can vary per batch; cache only static SBHD + # cp_context entries here and revisit routing metadata lifetime in the CP layout refactor. + self._chunkwise_cp_context_cache: dict[tuple[int, int], tuple[torch.Tensor, object]] = {} self.reset_parameters() def _setup_variant_attrs(self): - """Set variant specifics on the module. Called once from ``__init__``. + """Set variant projection sections, gate parameter sizes, and kernel callable. Must set: - - ``in_proj_dim`` + - ``in_proj_extra_dim`` (the in_proj sections beyond q/k/v/z; the base + class derives ``in_proj_dim`` from it) - ``in_proj_split_names`` - ``in_proj_split_sections`` - - ``feat_dim_split`` - ``dt_bias_dim`` / ``a_log_dim`` (sizes of the gate parameters, which the base class creates after the conv1d module to preserve the original parameter registration order) @@ -306,8 +344,8 @@ def _setup_variant_attrs(self): def _reset_dt_bias(self): """Initialize ``dt_bias``. Called from ``reset_parameters`` under the RNG tracker. - Defaults to ones; variants whose kernel expects a different step-size - parametrization override this. + Defaults to ones; subclasses can override this if their kernel expects a + different step-size parametrization. """ torch.ones( self.dt_bias_dim, @@ -324,9 +362,7 @@ def reset_parameters(self): nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init) self._reset_dt_bias() A = torch.empty( - self.A_log.shape[0], - dtype=self.config.params_dtype, - device=torch.cuda.current_device(), + self.A_log.shape[0], dtype=self.A_log.dtype, device=self.A_log.device ).uniform_(*self.A_init_range) self.A_log.data.copy_(torch.log(A)) @@ -344,29 +380,6 @@ def forward( # pylint: disable=missing-function-docstring raise NotImplementedError - def _gated_norm_and_a2a( - self, - core_attn_out: torch.Tensor, - gate: torch.Tensor, - thd_cp_a2a_inv: torch.Tensor | None, - batch: int, - seq_len: int, - packed_seq_params: PackedSeqParams | None = None, - ) -> torch.Tensor: - # RMSNorm - nvtx_range_push(suffix="gated_norm") - norm_out_hp = self._apply_gated_norm(core_attn_out, gate) - nvtx_range_pop(suffix="gated_norm") - - # Transpose: b s x --> s b x - # From bshd back to sbhd format - norm_out_hp = norm_out_hp.reshape(batch, seq_len, -1) - norm_out_hp = norm_out_hp.transpose(0, 1).contiguous() - - return a2a_hp_to_cp( - norm_out_hp, self.cp_size, self.pg_collection.cp, packed_seq_params, thd_cp_a2a_inv - ) - @jit_fuser def _apply_gated_norm(self, x, gate): # Output Norm @@ -389,24 +402,25 @@ def _prepare_input_for_gated_delta_rule( batch: int, seq_len: int, *gate_feats: tuple[torch.Tensor], + cp_size_headwise: int | None = None, ) -> dict[str, torch.Tensor]: """ Prepare all gated delta rule kernel inputs. Fuses split, reshape, L2 norm, decay/gate activations, repeat_interleave, and - contiguous operations. ``gate_feats`` holds the variant-specific in_proj - sections, which ``_compute_gates`` turns into the decay and gating tensors. + contiguous operations. ``gate_feats`` holds the in_proj sections after qkv + and gate, which ``_compute_gates`` turns into the decay and gating tensors. Returns: (dict[str, Tensor]): Kernel inputs keyed by kernel argument name (``q``, - ``k``, ``v``, ``g``, plus the variant-specific gates), and the output + ``k``, ``v``, ``g``, and ``beta``), and the output gate (z) tensor under the ``gate`` key, which is not a kernel input. """ + cp_size = 1 if cp_size_headwise is None else cp_size_headwise + # Split qkv into query_key and value query_key, value = torch.split( - qkv, - [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size], - dim=-1, + qkv, [2 * self.qk_dim_local_tp // cp_size, self.v_dim_local_tp // cp_size], dim=-1 ) # Reshape query_key and value @@ -418,7 +432,7 @@ def _prepare_input_for_gated_delta_rule( query_key = l2norm(query_key.contiguous()) # Split query and key - split_size = self.qk_dim_local_tp // self.key_head_dim // self.cp_size + split_size = self.qk_dim_local_tp // self.key_head_dim // cp_size query, key = torch.split(query_key, [split_size, split_size], dim=2) # Expand query and key if needed (grouped query attention) @@ -450,19 +464,18 @@ def _compute_gates( *gate_feats: tuple[torch.Tensor], ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Compute the log-decay ``g`` and the variant-specific kernel inputs. + Compute the log-decay ``g`` and remaining kernel inputs. Args: A_log_local_cp: CP-local slice of ``A_log``. dt_bias_local_cp: CP-local slice of ``dt_bias``. batch: Batch size. seq_len: Sequence length. - gate_feats: The variant-specific in_proj output sections (everything after - the qkv and output-gate sections, in ``feat_dim_split`` order). + gate_feats: The in_proj output sections after the qkv and output-gate sections. Returns: (tuple[Tensor, dict[str, Tensor]]): The log-decay ``g`` and a dict of the - remaining variant-specific kernel inputs keyed by kernel argument name. + remaining kernel inputs keyed by kernel argument name. """ raise NotImplementedError @@ -647,7 +660,7 @@ def _build_head_perm_for_split_sections( def get_parameter_local_cp( param: torch.Tensor, dim: int, - cp_group: torch.distributed.ProcessGroup, + cp_group: torch.distributed.ProcessGroup | None, split_sections: Optional[list[int]] = None, ) -> torch.Tensor: """Get the local parameter for the current context parallel rank. @@ -665,13 +678,15 @@ def get_parameter_local_cp( torch.Tensor: The local parameter for the current context parallel rank. """ - cp_size = cp_group.size() - cp_rank = cp_group.rank() + cp_size = cp_group.size() if cp_group is not None else 1 # No need to split if CP size is 1. if cp_size == 1: return param + assert cp_group is not None + cp_rank = cp_group.rank() + # Split first if needed. if split_sections is not None: inputs = torch.split(param, split_sections, dim=dim) @@ -693,7 +708,7 @@ def tensor_a2a_cp2hp( tensor: torch.Tensor, seq_dim: int, head_dim: int, - cp_group: torch.distributed.ProcessGroup, + cp_group: torch.distributed.ProcessGroup | None, split_sections: Optional[list[int]] = None, undo_attention_load_balancing: bool = True, ): @@ -714,12 +729,14 @@ def tensor_a2a_cp2hp( torch.Tensor: The all-to-all tensor. """ - cp_size = cp_group.size() + cp_size = cp_group.size() if cp_group is not None else 1 # No need to all-to-all if CP size is 1. if cp_size == 1: return tensor + assert cp_group is not None + # Limitations of mamba_context_parallel._all_to_all_cp2hp. assert seq_dim == 0, f"tensor_a2a_cp2hp only supports seq_dim == 0 for now, but got {seq_dim=}" assert ( @@ -756,7 +773,7 @@ def tensor_a2a_hp2cp( tensor: torch.Tensor, seq_dim: int, head_dim: int, - cp_group: torch.distributed.ProcessGroup, + cp_group: torch.distributed.ProcessGroup | None, split_sections: Optional[list[int]] = None, redo_attention_load_balancing: bool = True, ): @@ -777,12 +794,14 @@ def tensor_a2a_hp2cp( torch.Tensor: The all-to-all tensor. """ - cp_size = cp_group.size() + cp_size = cp_group.size() if cp_group is not None else 1 # No need to all-to-all if CP size is 1. if cp_size == 1: return tensor + assert cp_group is not None + # Limitations of mamba_context_parallel._all_to_all_hp2cp. assert seq_dim == 0, f"tensor_a2a_hp2cp only supports seq_dim == 0 for now, but got {seq_dim=}" assert ( diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py index c7a983841cb..e56d3a95d72 100644 --- a/megatron/core/ssm/gated_delta_net/gdn.py +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh. # Some of this code was adopted from https://github.com/huggingface/transformers @@ -12,35 +12,36 @@ import torch.nn.functional as F from megatron.core import tensor_parallel -from megatron.core.inference.contexts import BaseInferenceContext, DynamicInferenceContext -from megatron.core.inference.contexts.attention_context.triton.tensor_ops import ( - tensor_masked_update, -) +from megatron.core.context_parallel import contiguous_to_zigzag, zigzag_to_contiguous +from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.jit import jit_fuser -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group +from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gated_delta_net.common import ( _GDNBase, a2a_cp_to_hp, + a2a_hp_to_cp, + build_cp_context, causal_conv1d, chunk_gated_delta_rule, get_parameter_local_cp, l2norm, ) -from megatron.core.ssm.ssm_inference import SSMDynamicInferenceMixin from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push -try: - from fla.modules.convolution import causal_conv1d_update - from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule -except ImportError: - causal_conv1d_update = None - fused_recurrent_gated_delta_rule = None +class GatedDeltaNet(_GDNBase): + """Gated DeltaNet with a head-wise scalar memory-decay gate.""" -class GatedDeltaNet(SSMDynamicInferenceMixin, _GDNBase): - # pylint: disable=missing-class-docstring def _setup_variant_attrs(self): """Set the GDN in_proj sizing, split tables, gate parameter dims, and kernel.""" + self.gdn_pre_gated_delta_rule_fusion = self.config.gdn_pre_gated_delta_rule_fusion + if self.config.deterministic_mode and self.gdn_pre_gated_delta_rule_fusion: + raise ValueError( + "Pre-GDR fusion is non-deterministic, but deterministic_mode=True. " + "Disable gdn_pre_gated_delta_rule_fusion or deterministic_mode." + ) + # alpha, beta self.in_proj_extra_dim = self.num_value_heads * 2 @@ -56,13 +57,6 @@ def _setup_variant_attrs(self): self.num_value_heads // self.tp_size, # beta self.num_value_heads // self.tp_size, # alpha ) - self.feat_dim_split = ( - (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // self.cp_size, # qkv - self.v_dim_local_tp // self.cp_size, # gate (z) - self.num_value_heads // self.tp_size // self.cp_size, # beta - self.num_value_heads // self.tp_size // self.cp_size, # alpha - ) - self.dt_bias_dim = self.num_v_heads_local_tp self.a_log_dim = self.num_v_heads_local_tp @@ -70,7 +64,15 @@ def _setup_variant_attrs(self): self.gated_delta_rule = torch_chunk_gated_delta_rule else: self.gated_delta_rule = chunk_gated_delta_rule - self.chunk_size = 64 + + def _get_feat_dim_split(self, cp_size_headwise: int) -> tuple[int, int, int, int]: + """Return GDN1 qkv/z/beta/alpha split sizes for a runtime headwise CP size.""" + return ( + (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size_headwise, + self.v_dim_local_tp // cp_size_headwise, + self.num_value_heads // self.tp_size // cp_size_headwise, + self.num_value_heads // self.tp_size // cp_size_headwise, + ) @jit_fuser def _compute_gates( @@ -96,9 +98,10 @@ def forward( packed_seq_params: Optional[PackedSeqParams] = None, sequence_len_offset: Optional[int] = None, *, + pg_collection: Optional[ProcessGroupCollection] = None, inference_params: Optional[BaseInferenceContext] = None, **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor | None]: + ): """ Perform a forward pass through the GDN module. @@ -108,30 +111,68 @@ def forward( inference_context = deprecate_inference_params(inference_context, inference_params) - seq_len, batch, _ = hidden_states.shape - seq_len = seq_len * self.sp_size * self.cp_size + active_pg_collection = pg_collection if pg_collection is not None else self.pg_collection + base_cp_group = active_pg_collection.cp + cp_group = resolve_cp_group(base_cp_group, packed_seq_params) + if self.config.linear_cp_mode == "chunkwise": + cp_group_chunkwise = cp_group + cp_group_headwise = None + elif self.config.linear_cp_mode == "headwise": + cp_group_chunkwise = None + cp_group_headwise = cp_group + elif cp_group.size() == 1: + cp_group_chunkwise = None + cp_group_headwise = None + else: + raise ValueError( + f"Unsupported linear_cp_mode {self.config.linear_cp_mode!r}; " + "expected 'headwise' or 'chunkwise'." + ) + cp_size_chunkwise = cp_group_chunkwise.size() if cp_group_chunkwise is not None else 1 + cp_size_headwise = cp_group_headwise.size() if cp_group_headwise is not None else 1 + cp_size_runtime = cp_group.size() + convert_from_zigzag = ( + self.config.linear_cp_mode == "chunkwise" + and cp_size_chunkwise > 1 + and self.config.cp_partition_mode == "zigzag" + ) + if convert_from_zigzag: + hidden_states = zigzag_to_contiguous( + hidden_states, + cp_group_chunkwise, + self.config.sequence_parallel, + self.tp_group, + getattr(active_pg_collection, "tp_cp", None), + ) + + seq_len_local, batch, _ = hidden_states.shape + seq_len_post_headwise = seq_len_local * self.sp_size * cp_size_headwise + seq_len_global = seq_len_post_headwise * cp_size_chunkwise if inference_context is not None: - if inference_context.is_dynamic_batching(): - assert ( - not self.config.deterministic_mode - ), "GDN dynamic inference requires the FLA recurrent kernels." - assert ( - not self.config.batch_invariant_mode - ), "GDN dynamic inference does not support batch-invariant mode." - assert ( - self.cp_size == 1 - ), "Context parallelism is not supported for GDN dynamic inference." - assert ( - inference_context.num_speculative_tokens == 0 - ), "GDN dynamic inference does not support speculative decoding." - assert ( - not inference_context.enable_prefix_caching - ), "GDN dynamic inference does not support prefix caching." - return self.ssm_dynamic_inference(hidden_states, inference_context) - assert inference_context.is_static_batching() + assert ( + inference_context.is_static_batching() + ), "GDN does not currently support dynamic inference batching." assert not self.config.sequence_parallel - raise NotImplementedError("GDN static-batching inference is not supported.") + # TODO: support inference + raise NotImplementedError("GDN does not support inference for now.") + + if cp_size_headwise > 1 and ( + ( + packed_seq_params is not None + and packed_seq_params.qkv_format == "thd" + and getattr(packed_seq_params, "cp_partition_mode", self.config.cp_partition_mode) + != "zigzag" + ) + or ( + (packed_seq_params is None or packed_seq_params.qkv_format != "thd") + and self.config.cp_partition_mode != "zigzag" + ) + ): + raise ValueError( + "GatedDeltaNet with headwise CP requires zigzag layout. CP partition " + "conversion must be handled before calling GatedDeltaNet." + ) if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': assert batch == 1, "Packed sequence expects batch dimension to be 1" @@ -143,16 +184,16 @@ def forward( cu_seqlens_q = self._resolve_cu_seqlens( packed_seq_params.cu_seqlens_q_padded, packed_seq_params.cu_seqlens_q, - seq_len, + seq_len_global, "cu_seqlens_q", - cp_size=self.cp_size, + cp_size=cp_size_runtime, ) cu_seqlens_kv = self._resolve_cu_seqlens( packed_seq_params.cu_seqlens_kv_padded, packed_seq_params.cu_seqlens_kv, - seq_len, + seq_len_global, "cu_seqlens_kv", - cp_size=self.cp_size, + cp_size=cp_size_runtime, ) assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " @@ -167,6 +208,90 @@ def forward( cu_seqlens_q = None cu_seqlens_kv = None + if cp_size_chunkwise > 1: + if cu_seqlens_q is None: + cache_key = (seq_len_global, batch) + cached = self._chunkwise_cp_context_cache.get(cache_key) + if cached is None: + cached_cu_seqlens = ( + torch.arange( + batch + 1, device=torch.cuda.current_device(), dtype=torch.long + ) + * seq_len_global + ) + cached_ctx = build_cp_context( + cu_seqlens=cached_cu_seqlens, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + cached = (cached_cu_seqlens, cached_ctx) + self._chunkwise_cp_context_cache[cache_key] = cached + cu_seqlens_q, chunkwise_cp_context = cached + else: + chunkwise_cp_context = build_cp_context( + cu_seqlens=cu_seqlens_q, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + else: + chunkwise_cp_context = None + + if self.recompute_gdn and self.training: + + def _checkpointed_compute(hidden_states): + return self._forward_compute( + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ) + + out, out_bias = tensor_parallel.checkpoint(_checkpointed_compute, False, hidden_states) + else: + out, out_bias = self._forward_compute( + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ) + + if convert_from_zigzag: + out = contiguous_to_zigzag( + out, + cp_group_chunkwise, + self.config.sequence_parallel, + self.tp_group, + getattr(active_pg_collection, "tp_cp", None), + ) + + return out, out_bias + + def _forward_compute( + self, + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ): + """Core GDN computation (in_proj -> conv1d -> gated_delta_rule -> norm -> out_proj).""" # Input projection nvtx_range_push(suffix="in_proj") qkvzba, _ = self.in_proj(hidden_states) @@ -175,24 +300,171 @@ def forward( qkvzba, thd_cp_a2a_inv = a2a_cp_to_hp( qkvzba, self.in_proj_split_sections, - self.cp_size, - self.pg_collection.cp, + cp_size_headwise, + cp_group_headwise, cu_seqlens_q, - seq_len, + seq_len_post_headwise, packed_seq_params, ) - # Transpose: s b x --> b s x - # From sbhd to bshd format - qkvzba = qkvzba.transpose(0, 1) + if self.gdn_pre_gated_delta_rule_fusion: + if cp_size_chunkwise > 1 and batch > 1: + raise ValueError( + "GDN chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " + "when cp_context is used. Use packed THD input or micro_batch_size=1." + ) + if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + nvtx_range_push(suffix="fused_streamed_pre_gated_delta_rule") + seq_idx = ( + packed_seq_params.seq_idx + if packed_seq_params is not None + and packed_seq_params.qkv_format == 'thd' + and cp_size_chunkwise == 1 + else None + ) + fused_cu_seqlens_q = ( + cu_seqlens_q + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + else None + ) + query, key, value, gate, beta, g = self._fused_streamed_pre_gated_delta_rule( + qkvzba, + cu_seqlens_q=fused_cu_seqlens_q, + seq_idx=seq_idx, + cp_group=cp_group_chunkwise if cp_size_chunkwise > 1 else None, + cp_group_headwise=cp_group_headwise, + ) + kernel_inputs = {"q": query, "k": key, "v": value, "g": g, "beta": beta} + nvtx_range_pop(suffix="fused_streamed_pre_gated_delta_rule") + else: + nvtx_range_push(suffix="pre_gated_delta_rule") + if cp_size_chunkwise > 1 and packed_seq_params is None and batch > 1: + # TODO: If additional gated delta rule backends are added, handle this + # SBHD + chunkwise CP + batch>1 case per backend instead of + # unconditionally rejecting it. + raise ValueError( + "GDN chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " + "when cp_context is used. Use packed THD input or micro_batch_size=1." + ) + if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + query, key, value, gate, beta, g = self.pre_gated_delta_rule( + qkvzba, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + chunkwise_cp_context, + packed_seq_params=packed_seq_params, + ) + kernel_inputs = {"q": query, "k": key, "v": value, "g": g, "beta": beta} + nvtx_range_pop(suffix="pre_gated_delta_rule") + + nvtx_range_push(suffix="gated_delta_rule") + core_attn_out, _ = self.gated_delta_rule( + **kernel_inputs, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens_q, + cp_context=chunkwise_cp_context, + ) + nvtx_range_pop(suffix="gated_delta_rule") + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() + norm_func = partial( + self._gated_norm_and_layout_restore, + thd_cp_a2a_inv=thd_cp_a2a_inv, + batch=batch, + seq_len=seq_len_post_headwise, + packed_seq_params=packed_seq_params, + cp_size_headwise=cp_size_headwise, + cp_group_headwise=cp_group_headwise, + cp_size_chunkwise=cp_size_chunkwise, + cp_group_chunkwise=cp_group_chunkwise, + cu_seqlens_q=cu_seqlens_q, + ) + norm_out = self.norm_out_checkpoint.checkpoint(norm_func, core_attn_out, gate) + else: + norm_out = self._gated_norm_and_layout_restore( + core_attn_out, + gate, + thd_cp_a2a_inv, + batch, + seq_len_post_headwise, + packed_seq_params, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + ) + + # Output projection + nvtx_range_push(suffix="out_proj") + out, out_bias = self.out_proj(norm_out) + nvtx_range_pop(suffix="out_proj") + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint.discard_output_and_register_recompute(out) + + return out, out_bias + + def _gated_norm_and_layout_restore( + self, + core_attn_out: torch.Tensor, + gate: torch.Tensor, + thd_cp_a2a_inv: torch.Tensor | None, + batch: int, + seq_len: int, + packed_seq_params: PackedSeqParams | None, + cp_size_headwise: int, + cp_group_headwise: torch.distributed.ProcessGroup | None, + cp_size_chunkwise: int, + cp_group_chunkwise: torch.distributed.ProcessGroup | None, + cu_seqlens_q: torch.Tensor | None, + ) -> torch.Tensor: + nvtx_range_push(suffix="gated_norm") + norm_out = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + norm_out = norm_out.reshape(batch, seq_len, -1) + norm_out = norm_out.transpose(0, 1).contiguous() + + return a2a_hp_to_cp( + norm_out, cp_size_headwise, cp_group_headwise, packed_seq_params, thd_cp_a2a_inv + ) - # Split the tensor into q, k, v, gate (z), and the variant-specific gate features - # (beta, alpha for GDN; f, b, w for GDN2) - qkv, gate, beta, alpha = self._split_projection(qkvzba, batch, seq_len) + def pre_gated_delta_rule( + self, + qkvzba, + batch, + seq_len, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q=None, + chunkwise_cp_context=None, + packed_seq_params=None, + ): + """Prepare QKV, gate, beta, and decay tensors before the gated delta rule.""" + + qkvzba = qkvzba.transpose(0, 1) + qkv, gate, beta, alpha = torch.split( + qkvzba, self._get_feat_dim_split(cp_size_headwise), dim=-1 + ) + gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) - # Convolution on qkv nvtx_range_push(suffix="conv1d") - seq_len = qkv.shape[1] + kernel_seq_len = qkv.shape[1] qkv_channels_split_sections = [ self.qk_dim_local_tp, self.qk_dim_local_tp, @@ -201,219 +473,155 @@ def forward( conv1d_weight = get_parameter_local_cp( self.conv1d.weight, dim=0, - cp_group=self.pg_collection.cp, + cp_group=cp_group_headwise, split_sections=qkv_channels_split_sections, ) conv1d_bias = ( get_parameter_local_cp( self.conv1d.bias, dim=0, - cp_group=self.pg_collection.cp, + cp_group=cp_group_headwise, split_sections=qkv_channels_split_sections, ) if self.conv_bias else None ) if self.config.deterministic_mode: - qkv = qkv.transpose(1, 2).contiguous() # b, s, d -> b, d, s + qkv = qkv.transpose(1, 2).contiguous() conv_out = F.conv1d( - input=qkv, # Torch-native only accept [b, d, s] format input + input=qkv, weight=conv1d_weight, bias=conv1d_bias, stride=self.conv1d.stride, padding=self.conv1d.padding, dilation=self.conv1d.dilation, - groups=self.conv_dim_local_tp // self.cp_size, + groups=self.conv_dim_local_tp // cp_size_headwise, ) - qkv = self.act_fn(conv_out[..., :seq_len]) - qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d + qkv = self.act_fn(conv_out[..., :kernel_seq_len]) + qkv = qkv.transpose(1, 2) else: assert self.activation in ["silu", "swish"] + orig_seq = qkv.shape[1] + pad_n = 0 + conv_input = qkv.contiguous() + conv_cu_seqlens = cu_seqlens_q + conv_cp_context = chunkwise_cp_context + if self.config.gdn_conv_pad_alignment is not None: + if packed_seq_params is None or cu_seqlens_q is None: + raise ValueError( + "gdn_conv_pad_alignment is only supported with packed sequence " + "parameters in THD format. SBHD inputs do not need causal-conv padding." + ) + if chunkwise_cp_context is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + pad_n = -orig_seq % self.config.gdn_conv_pad_alignment + if pad_n > 0: + conv_input = torch.nn.functional.pad(conv_input, (0, 0, 0, pad_n)) + conv_cu_seqlens = cu_seqlens_q.clone() + conv_cu_seqlens[-1] += pad_n qkv, _ = causal_conv1d( - x=qkv, # FLA conv1d accepts [b, s, d] format input - weight=conv1d_weight.squeeze(1), # d, 1, w -> d, w + x=conv_input, + weight=conv1d_weight.squeeze(1), bias=conv1d_bias, activation=self.activation, initial_state=None, output_final_state=False, - cu_seqlens=cu_seqlens_q, + cu_seqlens=conv_cu_seqlens, + cp_context=conv_cp_context, ) + if pad_n > 0: + qkv = qkv[:, :orig_seq, :] nvtx_range_pop(suffix="conv1d") - A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp) - dt_bias_local_cp = get_parameter_local_cp( - self.dt_bias, dim=0, cp_group=self.pg_collection.cp - ) + A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group_headwise) + dt_bias_local_cp = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group_headwise) - # Prepare all kernel inputs (split, reshape, L2 norm, gates, contiguous) nvtx_range_push(suffix="prepare_input_for_gated_delta_rule") kernel_inputs = self._prepare_input_for_gated_delta_rule( - qkv, gate, A_log_local_cp, dt_bias_local_cp, batch, seq_len, beta, alpha + qkv, + gate, + A_log_local_cp, + dt_bias_local_cp, + batch, + kernel_seq_len, + beta, + alpha, + cp_size_headwise=cp_size_headwise, ) - gate = kernel_inputs.pop("gate") nvtx_range_pop(suffix="prepare_input_for_gated_delta_rule") - nvtx_range_push(suffix="gated_delta_rule") - core_attn_out, _ = self.gated_delta_rule( - **kernel_inputs, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - cu_seqlens=cu_seqlens_q, - ) - nvtx_range_pop(suffix="gated_delta_rule") - - if self.recompute_norm_out: - self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() - norm_func = partial( - self._gated_norm_and_a2a, - thd_cp_a2a_inv=thd_cp_a2a_inv, - batch=batch, - seq_len=seq_len, - packed_seq_params=packed_seq_params, - ) - norm_out = self.norm_out_checkpoint.checkpoint(norm_func, core_attn_out, gate) - else: - norm_out = self._gated_norm_and_a2a( - core_attn_out, gate, thd_cp_a2a_inv, batch, seq_len, packed_seq_params - ) - - # Output projection - nvtx_range_push(suffix="out_proj") - out, out_bias = self.out_proj(norm_out) - nvtx_range_pop(suffix="out_proj") - - if self.recompute_norm_out: - self.norm_out_checkpoint.discard_output_and_register_recompute(out) - - return out, out_bias + gate = kernel_inputs.pop("gate") - def _split_projection( - self, projected: torch.Tensor, batch: int, seq_len: int - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Split the fused projection into qkv, output gate, beta, and alpha.""" - qkv, gate, beta, alpha = torch.split(projected, self.feat_dim_split, dim=-1) - gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) - return qkv, gate, beta, alpha - - def _prepare_inference_inputs( - self, qkv: torch.Tensor, beta: torch.Tensor, alpha: torch.Tensor, batch: int, seq_len: int - ) -> dict[str, torch.Tensor]: - """Prepare raw FLA inputs while leaving normalization and gates fused in-kernel.""" - query_key, value = torch.split(qkv, [2 * self.qk_dim_local_tp, self.v_dim_local_tp], dim=-1) - query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) - query, key = torch.chunk(query_key, 2, dim=2) - value = value.reshape(batch, seq_len, -1, self.value_head_dim) - return { - "q": query.contiguous(), - "k": key.contiguous(), - "v": value.contiguous(), - "g": alpha.contiguous(), - "beta": beta.contiguous(), - } - - def mamba_state_shapes_per_request(self) -> tuple[tuple[int, ...], tuple[int, ...]]: - """Return the TP-local convolution and delta-rule cache shapes.""" return ( - (self.conv_dim_local_tp, self.conv_kernel_dim), - (self.num_v_heads_local_tp, self.key_head_dim, self.value_head_dim), + kernel_inputs["q"], + kernel_inputs["k"], + kernel_inputs["v"], + gate, + kernel_inputs["beta"], + kernel_inputs["g"], ) - def ssm_decode( - self, - projected: torch.Tensor, - conv_state: torch.Tensor, - ssm_state: torch.Tensor, - batch_indices: torch.Tensor, - intermediate_conv_state: torch.Tensor | None = None, - intermediate_ssm_state: torch.Tensor | None = None, - ) -> torch.Tensor: - """Run one CUDA-graph-compatible GDN decode token per request.""" - batch, seq_len, _ = projected.shape - assert seq_len == 1, "GDN speculative decoding is not supported." - assert ( - intermediate_conv_state is None and intermediate_ssm_state is None - ), "GDN speculative decoding state capture is not supported." - assert causal_conv1d_update is not None and fused_recurrent_gated_delta_rule is not None - - qkv, gate, beta, alpha = self._split_projection(projected, batch, seq_len) - read_indices = batch_indices.clamp(min=0) - - active_conv_state = conv_state[read_indices].contiguous() - qkv_dtype = qkv.dtype - qkv, active_conv_state = causal_conv1d_update( - x=qkv.to(conv_state.dtype), - cache=active_conv_state, - weight=self.conv1d.weight.squeeze(1).to(conv_state.dtype), - bias=self.conv1d.bias.to(conv_state.dtype) if self.conv1d.bias is not None else None, - activation=self.activation, - ) - qkv = qkv.to(qkv_dtype) - tensor_masked_update(conv_state, batch_indices, active_conv_state) + def _fused_streamed_pre_gated_delta_rule( + self, qkvzba, cu_seqlens_q=None, seq_idx=None, cp_group=None, cp_group_headwise=None + ): + """Call the streamed fused pre-GDR wrapper.""" - kernel_inputs = self._prepare_inference_inputs(qkv, beta, alpha, batch, seq_len) - active_ssm_state = ssm_state[read_indices].contiguous() - core_attn_out, final_ssm_state = fused_recurrent_gated_delta_rule( - **kernel_inputs, - A_log=self.A_log, - dt_bias=self.dt_bias, - initial_state=active_ssm_state, - output_final_state=True, - use_qk_l2norm_in_kernel=self.use_qk_l2norm, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=True, - ) - tensor_masked_update(ssm_state, batch_indices, final_ssm_state) - return self._apply_gated_norm(core_attn_out, gate).reshape(batch, seq_len, -1) + try: + from megatron.core.fusions.fused_pre_gated_delta_rule import ( + fused_streamed_pre_gated_delta_rule, + ) + except ImportError as exc: + raise ImportError( + "gdn_pre_gated_delta_rule_fusion requires the streamed pre-GDR fusion " + "dependencies, including causal-conv1d." + ) from exc - def ssm_prefill( - self, - projected: torch.Tensor, - conv_state: torch.Tensor, - ssm_state: torch.Tensor, - context: DynamicInferenceContext, - ) -> torch.Tensor: - """Run packed variable-length GDN prefill and populate request states.""" - assert ( - not context.is_chunked_prefill_enabled() - ), "GDN dynamic inference does not support chunked prefill." - metadata = context.mamba_metadata - cu_seqlens = metadata.cu_seqlens - batch_indices = metadata.batch_indices_prefill - token_count = projected.shape[0] - - projected = projected.transpose(0, 1).contiguous() - qkv, gate, beta, alpha = self._split_projection(projected, 1, token_count) - read_indices = batch_indices.clamp(min=0) - - qkv_dtype = qkv.dtype - qkv, final_conv_state = causal_conv1d( - x=qkv.to(conv_state.dtype), - weight=self.conv1d.weight.squeeze(1).to(conv_state.dtype), - bias=self.conv1d.bias.to(conv_state.dtype) if self.conv1d.bias is not None else None, - activation=self.activation, - initial_state=conv_state[read_indices].contiguous(), - output_final_state=True, - cu_seqlens=cu_seqlens, + qkv_channels_split_sections = [ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ] + conv1d_weight = get_parameter_local_cp( + self.conv1d.weight, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + conv1d_bias = ( + get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + if self.conv_bias + else None + ) + A_log = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group_headwise) + dt_bias = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group_headwise) + num_value_heads = A_log.numel() + num_key_heads = (conv1d_weight.shape[0] - num_value_heads * self.value_head_dim) // ( + 2 * self.key_head_dim ) - qkv = qkv.to(qkv_dtype) - tensor_masked_update(conv_state, batch_indices, final_conv_state) - kernel_inputs = self._prepare_inference_inputs(qkv, beta, alpha, 1, token_count) - core_attn_out, final_ssm_state = chunk_gated_delta_rule( - **kernel_inputs, - A_log=self.A_log, - dt_bias=self.dt_bias, - initial_state=ssm_state[read_indices].contiguous(), - output_final_state=True, - use_qk_l2norm_in_kernel=self.use_qk_l2norm, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=True, - cu_seqlens=cu_seqlens, + return fused_streamed_pre_gated_delta_rule( + qkvzba, + conv1d_weight, + conv1d_bias, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=self.key_head_dim, + value_head_dim=self.value_head_dim, + use_qk_l2norm=self.use_qk_l2norm, + cu_seqlens=cu_seqlens_q, + seq_idx=seq_idx, + cp_group=cp_group, ) - tensor_masked_update(ssm_state, batch_indices, final_ssm_state) - y = self._apply_gated_norm(core_attn_out, gate) - return y.reshape(1, token_count, -1).transpose(0, 1).contiguous() #################### @@ -425,22 +633,20 @@ def torch_chunk_gated_delta_rule( v, g, beta, - scale=None, chunk_size=64, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False, cu_seqlens=None, - **kwargs, + cp_context=None, + scale=None, ) -> tuple[torch.Tensor, torch.Tensor | None]: # pylint: disable=line-too-long ''' Torch-native implementation of chunked gated delta rule for deterministic mode. Need this because FLA is not deterministic. - ``scale`` defaults to ``1 / sqrt(K)``, matching the FLA kernel. Extra keyword - arguments are accepted and ignored so this stays interchangeable with the FLA - kernel, which takes several options this implementation does not model. + ``scale`` defaults to ``1 / sqrt(K)``, matching the FLA kernel. Reference: https://github.com/huggingface/transformers/blob/144c8ce2809a2e21914017652700e1ecb450501e/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L470-L547 ''' @@ -448,65 +654,67 @@ def torch_chunk_gated_delta_rule( assert ( cu_seqlens is None ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now." + assert ( + cp_context is None + ), "cp_context is not supported for torch_chunk_gated_delta_rule for now." - query, key, value = q, k, v - initial_dtype = query.dtype + initial_dtype = q.dtype if use_qk_l2norm_in_kernel: - query = l2norm(query, dim=-1, eps=1e-6) - key = l2norm(key, dim=-1, eps=1e-6) - query, key, value, beta, g = [ - x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) + q = l2norm(q, dim=-1, eps=1e-6) + k = l2norm(k, dim=-1, eps=1e-6) + q, k, v, beta, g = [ + x.transpose(1, 2).contiguous().to(torch.float32) for x in (q, k, v, beta, g) ] - batch_size, num_heads, sequence_length, k_head_dim = key.shape - v_head_dim = value.shape[-1] + batch_size, num_heads, sequence_length, k_head_dim = k.shape + v_head_dim = v.shape[-1] pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size - query = F.pad(query, (0, 0, 0, pad_size)) - key = F.pad(key, (0, 0, 0, pad_size)) - value = F.pad(value, (0, 0, 0, pad_size)) + q = F.pad(q, (0, 0, 0, pad_size)) + k = F.pad(k, (0, 0, 0, pad_size)) + v = F.pad(v, (0, 0, 0, pad_size)) beta = F.pad(beta, (0, pad_size)) g = F.pad(g, (0, pad_size)) total_sequence_length = sequence_length + pad_size if scale is None: - scale = 1 / (query.shape[-1] ** 0.5) - query = query * scale + scale = 1 / (q.shape[-1] ** 0.5) + q = q * scale - v_beta = value * beta.unsqueeze(-1) - k_beta = key * beta.unsqueeze(-1) + v_beta = v * beta.unsqueeze(-1) + k_beta = k * beta.unsqueeze(-1) # reshape to chunks - query, key, value, k_beta, v_beta = [ + q, k, v, k_beta, v_beta = [ x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) - for x in (query, key, value, k_beta, v_beta) + for x in (q, k, v, k_beta, v_beta) ] g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) mask = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0 + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0 ) # chunk decay g = g.cumsum(dim=-1) decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() - attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) + attn = -((k_beta @ k.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) for i in range(1, chunk_size): row = attn[..., i, :i].clone() sub = attn[..., :i, :i].clone() attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) - value = attn @ v_beta + v = attn @ v_beta k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) last_recurrent_state = ( - torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value) + torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(v) if initial_state is None - else initial_state.to(value) + else initial_state.to(v) ) - core_attn_out = torch.zeros_like(value) + core_attn_out = torch.zeros_like(v) mask = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1 + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1 ) # for each chunk for i in range(0, total_sequence_length // chunk_size): - q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] + q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state v_new = v_i - v_prime diff --git a/megatron/core/ssm/gated_delta_net/kda.py b/megatron/core/ssm/gated_delta_net/kda.py new file mode 100644 index 00000000000..cad16568472 --- /dev/null +++ b/megatron/core/ssm/gated_delta_net/kda.py @@ -0,0 +1,670 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Kimi Delta Attention, a channel-wise Gated DeltaNet variant.""" + +import math +from dataclasses import dataclass +from functools import partial +from typing import Optional + +import torch +import torch.nn.functional as F + +from megatron.core import tensor_parallel +from megatron.core.context_parallel import contiguous_to_zigzag, zigzag_to_contiguous +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.jit import jit_fuser +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gated_delta_net.common import ( + HAVE_FLA, + GatedDeltaNetSubmodules, + _GDNBase, + a2a_cp_to_hp, + a2a_hp_to_cp, + build_cp_context, + causal_conv1d, + get_parameter_local_cp, +) +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.module import mark_keep_in_fp32 +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push + +try: + from fla.ops.kda import chunk_kda + + # KDA also relies on the shared FLA convolution, normalization, and CP helpers. + HAVE_FLA_KDA = HAVE_FLA +except ImportError: # pragma: no cover + chunk_kda = None + HAVE_FLA_KDA = False + + +@dataclass +class KimiDeltaAttentionSubmodules(GatedDeltaNetSubmodules): + """Submodules used by Kimi Delta Attention.""" + + beta_proj: ModuleSpec | type = IdentityOp + + +class KimiDeltaAttention(_GDNBase): + """Channel-wise Gated DeltaNet variant with direct Q/K/V/F/G projections. + + The initial implementation supports training with equal query/key and value + head layouts. Optional low-rank F/B/G projections and recurrent inference + are intentionally outside this implementation. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: KimiDeltaAttentionSubmodules, + layer_number: int | None = None, + bias: bool = False, + conv_bias: bool = False, + conv_init: float | None = None, + use_qk_l2norm: bool = True, + A_init_range: tuple[float, float] = (1, 16), + pg_collection: ProcessGroupCollection | None = None, + *, + name: str | None = None, + cp_comm_type: str | None = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + ) -> None: + if not HAVE_FLA or not HAVE_FLA_KDA: # pragma: no cover + raise ImportError( + "FLA KDA is not installed. Install flash-linear-attention with KDA support." + ) + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + bias=bias, + conv_bias=conv_bias, + conv_init=conv_init, + use_qk_l2norm=use_qk_l2norm, + A_init_range=A_init_range, + pg_collection=pg_collection, + name=name, + cp_comm_type=cp_comm_type, + pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, + ) + + # KDA keeps beta in a separate projection so its checkpoint layout remains + # independent from the direct Q/K/V/F/G projection. + self.beta_proj = build_module( + submodules.beta_proj, + self.hidden_size, + self.num_key_heads, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="beta_proj", + tp_group=self.pg_collection.tp, + name=(name + ".beta_proj") if name is not None else None, + ) + + # These kernel parameters participate in FP32 gate math and must remain + # FP32 through model casting, optimizer construction, and checkpointing. + mark_keep_in_fp32(self.dt_bias) + mark_keep_in_fp32(self.A_log) + + # Pass raw g to FLA and let the KDA kernel apply A_log and dt_bias. + self.use_gate_in_kernel = True + + def _setup_variant_attrs(self) -> None: + """Set KDA dimensions, projection checkpoint metadata, and kernel callable.""" + + self.gdn_pre_gated_delta_rule_fusion = self.config.gdn_pre_gated_delta_rule_fusion + + # Channel-wise raw memory-decay gate g. + self.in_proj_extra_dim = self.qk_dim + + # Per-section sizes (and names) of the in_proj output, local to this TP rank. + # Used for the CP head permutation, post-a2a split, and sharded checkpoint split. + self.in_proj_split_names = ["query", "key", "value", "g", "gate"] + self.in_proj_split_sections = ( + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ) + self.dt_bias_dim = self.qk_dim_local_tp + self.a_log_dim = self.num_k_heads_local_tp + self.gate_params_dtype = torch.float32 + self.gated_delta_rule = chunk_kda + + def _get_feat_dim_split(self, cp_size_headwise: int) -> tuple[int, int, int]: + """Return KDA qkv/raw-g/output-gate split sizes for runtime headwise CP.""" + + return ( + (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size_headwise, + self.qk_dim_local_tp // cp_size_headwise, + self.v_dim_local_tp // cp_size_headwise, + ) + + def _reset_dt_bias(self) -> None: + """Initialize the KDA channel-wise step-size bias in inverse-softplus space.""" + + dt_min, dt_max, dt_init_floor = 0.001, 0.1, 1e-4 + dt = torch.exp( + torch.rand(self.dt_bias_dim, device=self.dt_bias.device, dtype=self.dt_bias.dtype) + * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min) + ).clamp(min=dt_init_floor) + self.dt_bias.data.copy_(dt + torch.log(-torch.expm1(-dt))) + + @jit_fuser + def _compute_gates( + self, + A_log_local_cp: torch.Tensor, + dt_bias_local_cp: torch.Tensor, + batch: int, + seq_len: int, + *gate_feats: tuple[torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Shape the raw channel-wise decay gate and activate the write strength.""" + + # ``gate_feats`` follows the KDA pre-GDR order: raw g, then separate beta. + raw_g, beta = gate_feats + num_key_heads = A_log_local_cp.numel() + raw_g = raw_g.reshape(batch, seq_len, num_key_heads, self.key_head_dim) + beta = beta.reshape(batch, seq_len, num_key_heads).float().sigmoid() + return raw_g, {"beta": beta.contiguous()} + + @jit_fuser + def _apply_gated_norm(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Apply per-head RMSNorm followed by KDA's sigmoid output gate.""" + + x_dtype = x.dtype + x = x.reshape(-1, self.value_head_dim) + x = self.out_norm(x) + gate = gate.reshape(-1, self.value_head_dim) + return (x * torch.sigmoid(gate.float())).to(x_dtype) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + inference_context: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[int] = None, + *, + pg_collection: Optional[ProcessGroupCollection] = None, + inference_params: Optional[BaseInferenceContext] = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Run the direct-projection KDA training path.""" + + del attention_mask, sequence_len_offset, kwargs + inference_context = deprecate_inference_params(inference_context, inference_params) + + active_pg_collection = pg_collection if pg_collection is not None else self.pg_collection + base_cp_group = active_pg_collection.cp + cp_group = resolve_cp_group(base_cp_group, packed_seq_params) + if self.config.linear_cp_mode == "chunkwise": + cp_group_chunkwise = cp_group + cp_group_headwise = None + elif self.config.linear_cp_mode == "headwise": + cp_group_chunkwise = None + cp_group_headwise = cp_group + elif cp_group.size() == 1: + cp_group_chunkwise = None + cp_group_headwise = None + else: + raise ValueError( + f"Unsupported linear_cp_mode {self.config.linear_cp_mode!r}; " + "expected 'headwise' or 'chunkwise'." + ) + cp_size_chunkwise = cp_group_chunkwise.size() if cp_group_chunkwise is not None else 1 + cp_size_headwise = cp_group_headwise.size() if cp_group_headwise is not None else 1 + cp_size_runtime = cp_group.size() + convert_from_zigzag = ( + self.config.linear_cp_mode == "chunkwise" + and cp_size_chunkwise > 1 + and self.config.cp_partition_mode == "zigzag" + ) + if convert_from_zigzag: + hidden_states = zigzag_to_contiguous( + hidden_states, + cp_group_chunkwise, + self.config.sequence_parallel, + self.tp_group, + getattr(active_pg_collection, "tp_cp", None), + ) + + seq_len_local, batch, _ = hidden_states.shape + seq_len_post_headwise = seq_len_local * self.sp_size * cp_size_headwise + seq_len_global = seq_len_post_headwise * cp_size_chunkwise + + if inference_context is not None: + assert ( + inference_context.is_static_batching() + ), "KimiDeltaAttention does not currently support dynamic inference batching." + assert not self.config.sequence_parallel + raise NotImplementedError("KimiDeltaAttention does not support inference for now.") + + if cp_size_headwise > 1 and ( + ( + packed_seq_params is not None + and packed_seq_params.qkv_format == "thd" + and getattr(packed_seq_params, "cp_partition_mode", self.config.cp_partition_mode) + != "zigzag" + ) + or ( + (packed_seq_params is None or packed_seq_params.qkv_format != "thd") + and self.config.cp_partition_mode != "zigzag" + ) + ): + raise ValueError( + "KimiDeltaAttention with headwise CP requires zigzag layout. CP partition " + "conversion must be handled before calling KimiDeltaAttention." + ) + + if packed_seq_params is not None and packed_seq_params.qkv_format == "thd": + if batch != 1: + raise ValueError("Packed KDA expects batch dimension to be 1.") + cu_seqlens_q = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_q_padded, + packed_seq_params.cu_seqlens_q, + seq_len_global, + "cu_seqlens_q", + cp_size=cp_size_runtime, + ) + cu_seqlens_kv = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_kv_padded, + packed_seq_params.cu_seqlens_kv, + seq_len_global, + "cu_seqlens_kv", + cp_size=cp_size_runtime, + ) + self._validate_packed_cu_seqlens(cu_seqlens_q, cu_seqlens_kv) + else: + cu_seqlens_q = None + + if cp_size_chunkwise > 1: + if cu_seqlens_q is None: + cache_key = (seq_len_global, batch) + cached = self._chunkwise_cp_context_cache.get(cache_key) + if cached is None: + cached_cu_seqlens = ( + torch.arange( + batch + 1, device=torch.cuda.current_device(), dtype=torch.long + ) + * seq_len_global + ) + cached_ctx = build_cp_context( + cu_seqlens=cached_cu_seqlens, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + cached = (cached_cu_seqlens, cached_ctx) + self._chunkwise_cp_context_cache[cache_key] = cached + cu_seqlens_q, chunkwise_cp_context = cached + else: + chunkwise_cp_context = build_cp_context( + cu_seqlens=cu_seqlens_q, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + else: + chunkwise_cp_context = None + + if self.recompute_gdn and self.training: + + def _checkpointed_compute(hidden_states): + return self._forward_compute( + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ) + + out, out_bias = tensor_parallel.checkpoint(_checkpointed_compute, False, hidden_states) + else: + out, out_bias = self._forward_compute( + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ) + + if convert_from_zigzag: + out = contiguous_to_zigzag( + out, + cp_group_chunkwise, + self.config.sequence_parallel, + self.tp_group, + getattr(active_pg_collection, "tp_cp", None), + ) + + return out, out_bias + + def _forward_compute( + self, + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ): + """Core KDA computation (in_proj -> conv1d -> gated_delta_rule -> norm -> out_proj).""" + + # Input projections. Beta intentionally remains a separate matrix. + nvtx_range_push(suffix="in_proj") + qkvfg, _ = self.in_proj(hidden_states) + beta, _ = self.beta_proj(hidden_states) + nvtx_range_pop(suffix="in_proj") + + qkvfg, thd_cp_a2a_inv = a2a_cp_to_hp( + qkvfg, + self.in_proj_split_sections, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + seq_len_post_headwise, + packed_seq_params, + ) + beta, _ = a2a_cp_to_hp( + beta, + (self.num_k_heads_local_tp,), + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + seq_len_post_headwise, + packed_seq_params, + ) + + if self.gdn_pre_gated_delta_rule_fusion: + raise NotImplementedError( + "gdn_pre_gated_delta_rule_fusion is not implemented for KDA yet." + ) + + if cp_size_chunkwise > 1 and packed_seq_params is None and batch > 1: + raise ValueError( + "KDA chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " + "when cp_context is used. Use packed THD input or micro_batch_size=1." + ) + if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with KDA chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + + nvtx_range_push(suffix="pre_gated_delta_rule") + query, key, value, gate, beta, raw_g, A_log, dt_bias = self.pre_gated_delta_rule( + qkvfg, + beta, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + chunkwise_cp_context, + packed_seq_params=packed_seq_params, + ) + kernel_inputs = { + "q": query, + "k": key, + "v": value, + "g": raw_g, + "beta": beta, + "A_log": A_log, + "dt_bias": dt_bias, + } + nvtx_range_pop(suffix="pre_gated_delta_rule") + + nvtx_range_push(suffix="gated_delta_rule") + core_attn_out, _ = self.gated_delta_rule( + **kernel_inputs, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=self.use_gate_in_kernel, + safe_gate=self.config.kda_safe_gate, + lower_bound=self.config.kda_lower_bound, + cu_seqlens=cu_seqlens_q, + cp_context=chunkwise_cp_context, + ) + nvtx_range_pop(suffix="gated_delta_rule") + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() + norm_func = partial( + self._gated_norm_and_layout_restore, + thd_cp_a2a_inv=thd_cp_a2a_inv, + batch=batch, + seq_len=seq_len_post_headwise, + packed_seq_params=packed_seq_params, + cp_size_headwise=cp_size_headwise, + cp_group_headwise=cp_group_headwise, + cp_size_chunkwise=cp_size_chunkwise, + cp_group_chunkwise=cp_group_chunkwise, + cu_seqlens_q=cu_seqlens_q, + ) + norm_out = self.norm_out_checkpoint.checkpoint(norm_func, core_attn_out, gate) + else: + norm_out = self._gated_norm_and_layout_restore( + core_attn_out, + gate, + thd_cp_a2a_inv, + batch, + seq_len_post_headwise, + packed_seq_params, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + ) + + # Output projection. + nvtx_range_push(suffix="out_proj") + out, out_bias = self.out_proj(norm_out) + nvtx_range_pop(suffix="out_proj") + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint.discard_output_and_register_recompute(out) + + return out, out_bias + + def _gated_norm_and_layout_restore( + self, + core_attn_out: torch.Tensor, + gate: torch.Tensor, + thd_cp_a2a_inv: torch.Tensor | None, + batch: int, + seq_len: int, + packed_seq_params: PackedSeqParams | None, + cp_size_headwise: int, + cp_group_headwise: torch.distributed.ProcessGroup | None, + cp_size_chunkwise: int, + cp_group_chunkwise: torch.distributed.ProcessGroup | None, + cu_seqlens_q: torch.Tensor | None, + ) -> torch.Tensor: + """Apply KDA's gated output norm and restore its context-parallel layout.""" + + nvtx_range_push(suffix="gated_norm") + norm_out = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + norm_out = norm_out.reshape(batch, seq_len, -1) + norm_out = norm_out.transpose(0, 1).contiguous() + + return a2a_hp_to_cp( + norm_out, cp_size_headwise, cp_group_headwise, packed_seq_params, thd_cp_a2a_inv + ) + + def pre_gated_delta_rule( + self, + qkvfg, + beta, + batch, + seq_len, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q=None, + chunkwise_cp_context=None, + packed_seq_params=None, + ): + """Prepare QKV, output gate, beta, and raw decay tensors before KDA.""" + + qkvfg = qkvfg.transpose(0, 1) + beta = beta.transpose(0, 1) + qkv, raw_g, gate = torch.split(qkvfg, self._get_feat_dim_split(cp_size_headwise), dim=-1) + gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) + + nvtx_range_push(suffix="conv1d") + kernel_seq_len = qkv.shape[1] + qkv_channels_split_sections = [ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ] + conv1d_weight = get_parameter_local_cp( + self.conv1d.weight, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + conv1d_bias = ( + get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + if self.conv_bias + else None + ) + if self.config.deterministic_mode or causal_conv1d is None: + qkv = qkv.transpose(1, 2).contiguous() + conv_out = F.conv1d( + input=qkv, + weight=conv1d_weight, + bias=conv1d_bias, + stride=self.conv1d.stride, + padding=self.conv1d.padding, + dilation=self.conv1d.dilation, + groups=self.conv_dim_local_tp // cp_size_headwise, + ) + qkv = self.act_fn(conv_out[..., :kernel_seq_len]) + qkv = qkv.transpose(1, 2).contiguous() + else: + if self.activation not in ("silu", "swish"): + raise ValueError(f"FLA causal convolution requires SiLU, got {self.activation}.") + orig_seq = qkv.shape[1] + pad_n = 0 + conv_input = qkv.contiguous() + conv_cu_seqlens = cu_seqlens_q + conv_cp_context = chunkwise_cp_context + if self.config.gdn_conv_pad_alignment is not None: + if packed_seq_params is None or cu_seqlens_q is None: + raise ValueError( + "gdn_conv_pad_alignment is only supported with packed sequence " + "parameters in THD format. SBHD inputs do not need causal-conv padding." + ) + if chunkwise_cp_context is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with KDA chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + pad_n = -orig_seq % self.config.gdn_conv_pad_alignment + if pad_n > 0: + conv_input = torch.nn.functional.pad(conv_input, (0, 0, 0, pad_n)) + conv_cu_seqlens = cu_seqlens_q.clone() + conv_cu_seqlens[-1] += pad_n + qkv, _ = causal_conv1d( + x=conv_input, + weight=conv1d_weight.squeeze(1), + bias=conv1d_bias, + activation=self.activation, + initial_state=None, + output_final_state=False, + cu_seqlens=conv_cu_seqlens, + cp_context=conv_cp_context, + ) + if pad_n > 0: + qkv = qkv[:, :orig_seq, :] + nvtx_range_pop(suffix="conv1d") + + A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group_headwise) + dt_bias_local_cp = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group_headwise) + + nvtx_range_push(suffix="prepare_input_for_gated_delta_rule") + kernel_inputs = self._prepare_input_for_gated_delta_rule( + qkv, + gate, + A_log_local_cp, + dt_bias_local_cp, + batch, + kernel_seq_len, + raw_g, + beta, + cp_size_headwise=cp_size_headwise, + ) + nvtx_range_pop(suffix="prepare_input_for_gated_delta_rule") + + gate = kernel_inputs.pop("gate") + + return ( + kernel_inputs["q"], + kernel_inputs["k"], + kernel_inputs["v"], + gate, + kernel_inputs["beta"], + kernel_inputs["g"], + A_log_local_cp, + dt_bias_local_cp, + ) + + @staticmethod + def _validate_packed_cu_seqlens( + cu_seqlens_q: torch.Tensor, cu_seqlens_kv: torch.Tensor + ) -> None: + """Validate the self-attention boundary contract for packed KDA.""" + + if cu_seqlens_q.numel() < 2 or cu_seqlens_kv.numel() < 2: + raise ValueError( + "Packed KDA requires at least one sequence in both Q and KV boundaries." + ) + if cu_seqlens_q.shape != cu_seqlens_kv.shape or not torch.equal( + cu_seqlens_q, cu_seqlens_kv + ): + raise ValueError( + "Packed KDA requires cu_seqlens_q to equal cu_seqlens_kv, " + f"but got shapes {tuple(cu_seqlens_q.shape)} and " + f"{tuple(cu_seqlens_kv.shape)}." + ) + + def backward_dw(self) -> None: + """Execute weight-gradient computation for all KDA projections.""" + + super().backward_dw() + self.beta_proj.backward_dw() diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index ab7c9614cf5..62c16ff8b4c 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -1701,6 +1701,11 @@ def __init__( Args: name (str | None): module instance name passed top-down from its paranet module """ + if config.gated_attention_proj_granularity != 'elementwise': + raise ValueError( + "Regular attention does not support headwise " + "gated_attention_proj_granularity; use 'elementwise'." + ) super().__init__( config=config, submodules=submodules, diff --git a/megatron/core/transformer/fsdp_dtensor_checkpoint.py b/megatron/core/transformer/fsdp_dtensor_checkpoint.py index 327a4f4da05..2dde56c9c73 100644 --- a/megatron/core/transformer/fsdp_dtensor_checkpoint.py +++ b/megatron/core/transformer/fsdp_dtensor_checkpoint.py @@ -44,33 +44,11 @@ HAVE_MEGATRON_FSDP = False from megatron.core import parallel_state -from megatron.core.dist_checkpointing.core import CheckpointingException -from megatron.core.dist_checkpointing.validation import StrictHandling, parse_strict_flag from megatron.core.tensor_parallel.layers import copy_tensor_model_parallel_attributes from megatron.core.transformer.transformer_layer import TransformerLayer from megatron.core.utils import get_attr_wrapped_model -def _strip_wrapper_prefixes(path): - """Strip DDP/FSDP wrapper prefixes (module., model.) from a module or state dict path.""" - parts = path.split('.') - while parts and parts[0] in ('module', 'model'): - parts = parts[1:] - return '.'.join(parts) - - -def _intersect_slice(s1, s2): - """Intersection of two step-1 slices, or an empty slice when they do not overlap.""" - start = max(s1.start, s2.start) - stop = min(s1.stop, s2.stop) - return slice(0, 0) if start >= stop else slice(start, stop) - - -def _shift_slice(s, offset): - """Move a step-1 slice by ``offset``, e.g. to rebase it onto a shard's local storage.""" - return slice(s.start + offset, s.stop + offset) - - def get_ep_layer_offset(num_experts: int | None = None) -> int: """ Get the expert layer offset for the current model. @@ -237,16 +215,21 @@ def handle_swiglu_in_state_dict(model, model_state_dict, optimizer_state_dict): # layers have gated_linear_unit=False while language decoder layers # have gated_linear_unit=True. # ------------------------------------------------------------------ + def _strip_wrappers(path): + """Strip DDP/FSDP wrapper prefixes (module., model.) from a path.""" + parts = path.split('.') + while parts and parts[0] in ('module', 'model'): + parts = parts[1:] + return '.'.join(parts) + _layer_glu = {} for name, module in model.named_modules(): if isinstance(module, TransformerLayer): - _layer_glu[_strip_wrapper_prefixes(name)] = getattr( - module.config, 'gated_linear_unit', False - ) + _layer_glu[_strip_wrappers(name)] = getattr(module.config, 'gated_linear_unit', False) def _key_in_glu_layer(key): """Return True if *key* belongs to a TransformerLayer with gated_linear_unit=True.""" - norm_key = _strip_wrapper_prefixes(key) + norm_key = _strip_wrappers(key) best_glu, best_len = None, -1 for layer_path, uses_glu in _layer_glu.items(): if norm_key.startswith(layer_path + '.') and len(layer_path) > best_len: @@ -255,6 +238,17 @@ def _key_in_glu_layer(key): return True # no TransformerLayer found — assume GLU for backward compat return best_glu + def intersection(s1, s2): + # Only works for step=1 + start = max(s1.start, s2.start) + stop = min(s1.stop, s2.stop) + if start >= stop: + return slice(0, 0) # Empty slice if no intersection + return slice(start, stop) + + def offset_slice(s, offset): + return slice(s.start + offset, s.stop + offset) + def is_swiglu_key(key): """ Check if this key should be handled as SwiGLU linear_fc1 weight or bias. @@ -280,10 +274,21 @@ def is_swiglu_key(key): def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis, is_expert_param): """ Split the SWiGLU linear_fc1 parameter into two parts: weight_w and weight_v. + + Args: + data: The tensor to split. May be a DTensor (model state dict) or a + plain Tensor (optimizer states from FusedAdam). + dist_param: The corresponding model parameter (always a DTensor). + Used for global shape, numel, FSDP slice, and dist index metadata. + swiglu_shard_axis: Axis along which to split W and V gates. + is_expert_param: Whether this is an expert parameter (affects TP mesh). """ - assert data.shape[swiglu_shard_axis] % 2 == 0, ( - f"SWiGLU weights must have an even size along the shard axis {swiglu_shard_axis}, " - f"got {data.shape[swiglu_shard_axis]}" + # Use dist_param (always a DTensor) for global shape/numel, + # as data may be a regular Tensor (e.g., optimizer states). + global_shape = dist_param.shape + assert global_shape[swiglu_shard_axis] % 2 == 0, ( + f"SWiGLU FC1 must have even global size along axis {swiglu_shard_axis}, " + f"got {global_shape[swiglu_shard_axis]} (global_shape={list(global_shape)})" ) fsdp_slice = dist_param.megatron_fsdp_slice @@ -292,25 +297,39 @@ def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis, is_expert_param tp_mesh = megatron_fsdp_dist_index.get_submesh( [megatron_fsdp_dist_index.tp_dim], is_expert_parallel=is_expert_param ) - data_size = data.numel() // tp_mesh.mesh.numel() + data_size = dist_param.numel() // tp_mesh.mesh.numel() w_slice = slice(0, data_size // 2) v_slice = slice(data_size // 2, data_size) - view_shape = list(data.shape) + view_shape = list(global_shape) view_shape[swiglu_shard_axis] = -1 - local_tensor = data.to_local() + if isinstance(data, DTensor): + assert data.shape == global_shape, ( + f"DTensor shape mismatch: data.shape={data.shape} vs " + f"dist_param.shape={global_shape}" + ) + local_tensor = data.to_local() + else: + # Plain Tensor must already be the FSDP-local shard; other layouts + # (TP-local / global) would silently misalign in the slice below. + expected = fsdp_slice.stop - fsdp_slice.start + assert data.numel() == expected, ( + f"Plain Tensor must be FSDP-local shard " + f"(expected numel={expected}, got {data.numel()}; fsdp_slice={fsdp_slice})" + ) + local_tensor = data weight_w = local_tensor.view(-1)[ - _shift_slice(_intersect_slice(fsdp_slice, w_slice), -fsdp_slice.start) + offset_slice(intersection(fsdp_slice, w_slice), -fsdp_slice.start) ] weight_v = local_tensor.view(-1)[ - _shift_slice(_intersect_slice(fsdp_slice, v_slice), -fsdp_slice.start) + offset_slice(intersection(fsdp_slice, v_slice), -fsdp_slice.start) ] weight_w = weight_w.reshape(view_shape) weight_v = weight_v.reshape(view_shape) # Fake parameters w and v are used to provide the correct parameter # shape and Tensor-Parallelism information. - per_tp_rank_shape = list(data.shape) + per_tp_rank_shape = list(global_shape) if is_mcore_tensor_model_parallel(dist_param): tp_dim = get_mcore_tensor_parallel_partition_dim(dist_param) assert tp_dim is not None, "Tensor model parallel dimension not found" @@ -418,6 +437,13 @@ def handle_gdn_in_state_dict(model, model_state_dict, optimizer_state_dict): GDN_IN_PROJ_NAMES = ["query", "key", "value", "z", "beta", "alpha"] GDN_CONV1D_NAMES = ["query", "key", "value"] + def _strip_wrappers(path): + """Strip DDP/FSDP wrapper prefixes (module., model.) from a path.""" + parts = path.split('.') + while parts and parts[0] in ('module', 'model'): + parts = parts[1:] + return '.'.join(parts) + # ------------------------------------------------------------------ # Build per-GDN-module split-size map by walking the model tree. # GDN modules are identified by the presence of qk_dim / v_dim / @@ -431,7 +457,7 @@ def handle_gdn_in_state_dict(model, model_state_dict, optimizer_state_dict): qk = mod.qk_dim // tp v = mod.v_dim // tp nvh = mod.num_value_heads // tp - _gdn_info[_strip_wrapper_prefixes(name)] = { + _gdn_info[_strip_wrappers(name)] = { 'in_proj_sizes': [qk, qk, v, v, nvh, nvh], 'conv1d_sizes': [qk, qk, v], } @@ -442,7 +468,7 @@ def handle_gdn_in_state_dict(model, model_state_dict, optimizer_state_dict): def _match_gdn_key(key): """Return (split_sizes, sub_names, split_dim) if *key* is a GDN fused parameter that needs splitting, else ``None``.""" - norm = _strip_wrapper_prefixes(key) + norm = _strip_wrappers(key) for gdn_path, info in _gdn_info.items(): if not norm.startswith(gdn_path + '.'): continue @@ -453,27 +479,63 @@ def _match_gdn_key(key): return info['conv1d_sizes'], GDN_CONV1D_NAMES, 0 return None + def intersection(s1, s2): + start = max(s1.start, s2.start) + stop = min(s1.stop, s2.stop) + return slice(0, 0) if start >= stop else slice(start, stop) + + def offset_slice(s, offset): + return slice(s.start + offset, s.stop + offset) + def split_gdn_fused(data, dist_param, split_sizes, split_dim): - """Split a fused GDN projection DTensor into per-component DTensors.""" + """Split a fused GDN projection DTensor into per-component DTensors. + + Args: + data: The tensor to split. May be a DTensor (model state dict) or a + plain Tensor (optimizer states from FusedAdam). + dist_param: The corresponding model parameter (always a DTensor). + Used for global shape, numel, FSDP slice, and dist index metadata. + split_sizes: List of sizes for each component along split_dim. + split_dim: Dimension along which to split. + """ total_split = sum(split_sizes) if isinstance(data, DTensor) and data.shape[split_dim] == total_split: + # GDN tensors are already TP-local here (fast-path from #4799). return list( split_dtensor( data, split_sizes, dim=split_dim, update_uneven_dtensor_chunk_meta=True ) ) + # Use dist_param (always a DTensor) for global shape/numel, + # as data may be a regular Tensor (e.g., optimizer states). + global_shape = dist_param.shape + fsdp_slice = dist_param.megatron_fsdp_slice dist_index = dist_param.megatron_fsdp_dist_index tp_mesh = dist_index.get_submesh([dist_index.tp_dim], is_expert_parallel=False) - data_size = data.numel() // tp_mesh.mesh.numel() + data_size = dist_param.numel() // tp_mesh.mesh.numel() elems_per_unit = data_size // total_split - local_tensor = data.to_local() - view_shape = list(data.shape) + if isinstance(data, DTensor): + assert data.shape == global_shape, ( + f"DTensor shape mismatch: data.shape={data.shape} vs " + f"dist_param.shape={global_shape}" + ) + local_tensor = data.to_local() + else: + # Plain Tensor must already be the FSDP-local shard; other layouts + # (TP-local / global) would silently misalign in the slice below. + expected = fsdp_slice.stop - fsdp_slice.start + assert data.numel() == expected, ( + f"Plain Tensor must be FSDP-local shard " + f"(expected numel={expected}, got {data.numel()}; fsdp_slice={fsdp_slice})" + ) + local_tensor = data + view_shape = list(global_shape) - per_tp_rank_shape = list(data.shape) + per_tp_rank_shape = list(global_shape) if is_mcore_tensor_model_parallel(dist_param): tp_dim = get_mcore_tensor_parallel_partition_dim(dist_param) assert tp_dim is not None, "Tensor model parallel dimension not found" @@ -485,8 +547,8 @@ def split_gdn_fused(data, dist_param, split_sizes, split_dim): comp_flat = s * elems_per_unit comp_slice = slice(flat_offset, flat_offset + comp_flat) - shard = _intersect_slice(fsdp_slice, comp_slice) - comp_data = local_tensor.view(-1)[_shift_slice(shard, -fsdp_slice.start)] + shard = intersection(fsdp_slice, comp_slice) + comp_data = local_tensor.view(-1)[offset_slice(shard, -fsdp_slice.start)] comp_view = list(view_shape) comp_view[split_dim] = -1 @@ -555,296 +617,6 @@ def split_gdn_fused(data, dist_param, split_sizes, split_dim): return model_state_dict, optimizer_state_dict -def split_fused_fsdp_param(data, dist_param, split_sizes, is_expert_param=False, split_dim=0): - """Split a fused Megatron-FSDP parameter along ``split_dim`` into per-section DTensors. - - ``split_sizes`` are given in full, tensor-parallel-unsharded units along ``split_dim`` - (so they can be taken straight from the module config) and are scaled down to this - rank's TP shard internally. - - The returned DTensors alias the fused parameter's storage, so a save reads the fused - values and an in-place load writes back into the fused parameter. - - Same flat-slice arithmetic as :func:`handle_swiglu_in_state_dict`, generalized from an - even two-way split to arbitrary section sizes. - """ - assert HAVE_MEGATRON_FSDP, "This function requires Megatron-FSDP to be installed." - - fsdp_slice = dist_param.megatron_fsdp_slice - dist_index = dist_param.megatron_fsdp_dist_index - tp_mesh = dist_index.get_submesh([dist_index.tp_dim], is_expert_parallel=is_expert_param) - - per_tp_rank_shape = list(data.shape) - if is_mcore_tensor_model_parallel(dist_param): - tp_dim = get_mcore_tensor_parallel_partition_dim(dist_param) - assert tp_dim is not None, "Tensor model parallel dimension not found" - per_tp_rank_shape[tp_dim] //= tp_mesh.mesh.numel() - - total_full = sum(split_sizes) - assert data.shape[split_dim] == total_full, ( - f"Fused parameter is {data.shape[split_dim]} wide along dim {split_dim}, " - f"but the requested sections sum to {total_full}" - ) - - local_total = per_tp_rank_shape[split_dim] - local_sizes = [] - for size in split_sizes: - assert (size * local_total) % total_full == 0, ( - f"Section of size {size} does not divide evenly across the tensor-parallel " - f"group (fused dim {total_full} -> {local_total} on this rank)" - ) - local_sizes.append(size * local_total // total_full) - - data_size = 1 - for dim_size in per_tp_rank_shape: - data_size *= dim_size - elems_per_unit = data_size // local_total - - local_tensor = data.to_local() - view_shape = list(per_tp_rank_shape) - view_shape[split_dim] = -1 - - results = [] - flat_offset = 0 - for local_size in local_sizes: - section_numel = local_size * elems_per_unit - section_slice = slice(flat_offset, flat_offset + section_numel) - - shard = _intersect_slice(fsdp_slice, section_slice) - section_data = local_tensor.view(-1)[_shift_slice(shard, -fsdp_slice.start)] - section_data = section_data.reshape(view_shape) - - # A meta tensor carries the unsharded shape and TP attributes of the section. - meta_shape = list(per_tp_rank_shape) - meta_shape[split_dim] = local_size - section_meta = torch.empty(*meta_shape, device="meta") - copy_tensor_model_parallel_attributes(section_meta, dist_param) - - results.append( - make_fsdp_dtensor( - section_data.data, - section_meta, - dist_index=dist_index, - is_expert_param=is_expert_param, - run_check=True, - update_uneven_dtensor_chunk_meta=True, - ) - ) - flat_offset += section_numel - - return results - - -# The single down-projection registered by FusedMLASelfAttention, and the two separate -# projections registered by the unfused MLASelfAttention it replaces. -MLA_FUSED_DOWN_PROJ = 'linear_qkv_down_proj' -MLA_UNFUSED_DOWN_PROJS = ('linear_q_down_proj', 'linear_kv_down_proj') - - -def get_mla_fused_down_proj_splits(model): - """Map each FusedMLASelfAttention module path to its ``[q, kv]`` down-proj split sizes. - - Sizes match ``FusedMLASelfAttention.sharded_state_dict`` and are expressed in full, - tensor-parallel-unsharded units. Returns an empty dict for unfused models. - """ - splits = {} - for name, mod in model.named_modules(): - if not hasattr(mod, MLA_FUSED_DOWN_PROJ): - continue - config = mod.config - splits[_strip_wrapper_prefixes(name)] = [ - config.q_lora_rank, - config.kv_lora_rank + config.qk_pos_emb_head_dim, - ] - return splits - - -def match_mla_fused_down_proj_key(key, fused_splits): - """Match a state dict key against the fused MLA down-projections in ``fused_splits``. - - Returns ``(wrapper, attention_path, split_sizes, leaf)`` where ``wrapper`` is the - stripped ``module.``/``model.`` prefix and ``leaf`` is the part after the fused module - name, or ``None`` when the key is not a fused down-projection. - """ - norm = _strip_wrapper_prefixes(key) - for attention_path, split_sizes in fused_splits.items(): - fused_prefix = f'{attention_path}.{MLA_FUSED_DOWN_PROJ}.' - if norm.startswith(fused_prefix): - wrapper = key[: len(key) - len(norm)] - return wrapper, attention_path, split_sizes, norm[len(fused_prefix) :] - return None - - -def absorbed_input_layernorm_key(wrapper, attention_path, leaf): - """Rewrite a ``layer_norm_*`` leaf absorbed by the fused down-proj to ``input_layernorm.*``. - - With ``fuse_input_layernorm``, the fused module owns the layer's input layernorm as - ``layer_norm_weight``/``layer_norm_bias``; on disk it lives one level up under - ``input_layernorm``, per the layer spec's ``sharded_state_dict_keys_map``. - """ - assert '.' in attention_path, ( - f"Cannot locate the transformer layer owning {attention_path}.{leaf}; expected the " - f"attention module to be nested inside it." - ) - layer_path = attention_path.rsplit('.', 1)[0] - return f'{wrapper}{layer_path}.input_layernorm.{leaf[len("layer_norm_"):]}' - - -def handle_mla_down_proj_in_state_dict(model, model_state_dict, optimizer_state_dict): - """Rewrite a fused MLA down-projection into the unfused layout used on disk. - - ``mla_down_proj_fusion=True`` swaps ``MLASelfAttention`` for ``FusedMLASelfAttention``, - replacing ``linear_q_down_proj`` and ``linear_kv_down_proj`` with a single - ``linear_qkv_down_proj`` holding their row-concatenation, and on the Transformer Engine - backend also absorbing the layer's ``input_layernorm`` into it as ``layer_norm_*``. - - ``FusedMLASelfAttention.sharded_state_dict`` hides both changes from ``torch_dist`` - checkpoints so fused and unfused runs share one on-disk format. The Megatron-FSDP path - builds its keys from ``named_parameters()`` and never calls ``sharded_state_dict``, so - the same remapping is applied here. - - No-op for unfused models. - """ - assert HAVE_MEGATRON_FSDP, "This function requires Megatron-FSDP to be installed." - - fused_splits = get_mla_fused_down_proj_splits(model) - if not fused_splits: - return model_state_dict, optimizer_state_dict - - rewritten = 0 - - model_state_dict = model_state_dict.copy() - for key in list(model_state_dict.keys()): - match = match_mla_fused_down_proj_key(key, fused_splits) - if match is None: - continue - wrapper, attention_path, split_sizes, leaf = match - - if leaf.endswith('_extra_state'): - # Not a tensor; handle_fp8_extra_state_case drops these beforehand. - continue - - if leaf.startswith('layer_norm_'): - new_key = absorbed_input_layernorm_key(wrapper, attention_path, leaf) - model_state_dict[new_key] = model_state_dict.pop(key) - rewritten += 1 - continue - - sections = split_fused_fsdp_param( - model_state_dict[key], model.get_parameter(f'module.{key}'), split_sizes - ) - for proj_name, section in zip(MLA_UNFUSED_DOWN_PROJS, sections): - model_state_dict[f'{wrapper}{attention_path}.{proj_name}.{leaf}'] = section - del model_state_dict[key] - rewritten += 1 - - if rewritten: - logger.info( - f"[MLA] Rewrote {rewritten} fused {MLA_FUSED_DOWN_PROJ} key(s) across " - f"{len(fused_splits)} attention module(s) into the unfused layout." - ) - - if optimizer_state_dict is not None and len(optimizer_state_dict.get("state", {})) != 0: - optimizer_state_dict = optimizer_state_dict.copy() - optimizer_state = optimizer_state_dict["state"] - new_optimizer_state = {} - for key in list(optimizer_state.keys()): - match = match_mla_fused_down_proj_key(key, fused_splits) - if match is None: - new_optimizer_state[key] = optimizer_state[key] - continue - wrapper, attention_path, split_sizes, leaf = match - - if leaf.startswith('layer_norm_'): - # Absorbed input layernorm: moved, not split. - new_key = absorbed_input_layernorm_key(wrapper, attention_path, leaf) - new_optimizer_state[new_key] = optimizer_state[key] - continue - - new_keys = [ - f'{wrapper}{attention_path}.{proj_name}.{leaf}' - for proj_name in MLA_UNFUSED_DOWN_PROJS - ] - for new_key in new_keys: - new_optimizer_state[new_key] = optimizer_state[key].copy() - dist_param = model.get_parameter(key[len("module.") :]) - for subkey in ["exp_avg", "exp_avg_sq"]: - sections = split_fused_fsdp_param( - optimizer_state[key][subkey], dist_param, split_sizes - ) - for new_key, section in zip(new_keys, sections): - new_optimizer_state[new_key][subkey] = section - optimizer_state_dict["state"] = new_optimizer_state - - return model_state_dict, optimizer_state_dict - - -# MCore renamed MultiTokenPredictionLayer's inner transformer layer; checkpoints keep the -# original name. -MTP_INNER_LAYER = 'mtp_model_layer' -MTP_INNER_LAYER_CHECKPOINT_NAME = 'transformer_layer' - - -def get_mtp_inner_layer_paths(model): - """Paths of MTP layers whose inner layer is renamed on disk. - - Mamba MTP layers (``mtp_layer_pattern`` set) are excluded, matching - ``MultiTokenPredictionLayer.sharded_state_dict``: for them ``mtp_model_layer`` is - already the native checkpoint name. - """ - return [ - _strip_wrapper_prefixes(name) - for name, mod in model.named_modules() - if hasattr(mod, MTP_INNER_LAYER) and getattr(mod, 'mtp_layer_pattern', None) is None - ] - - -def rename_mtp_inner_layer_keys(state_dict, mtp_layer_paths): - """Rename ``mtp_model_layer.*`` keys to ``transformer_layer.*`` under ``mtp_layer_paths``.""" - state_dict = state_dict.copy() - for key in list(state_dict.keys()): - norm = _strip_wrapper_prefixes(key) - for mtp_path in mtp_layer_paths: - current_prefix = f'{mtp_path}.{MTP_INNER_LAYER}.' - if not norm.startswith(current_prefix): - continue - wrapper = key[: len(key) - len(norm)] - leaf = norm[len(current_prefix) :] - new_key = f'{wrapper}{mtp_path}.{MTP_INNER_LAYER_CHECKPOINT_NAME}.{leaf}' - state_dict[new_key] = state_dict.pop(key) - break - return state_dict - - -def handle_mtp_in_state_dict(model, model_state_dict, optimizer_state_dict): - """Rename MTP inner-layer keys to the ``transformer_layer`` name used on disk. - - ``MultiTokenPredictionLayer.sharded_state_dict`` remaps ``mtp_model_layer.*`` back to - ``transformer_layer.*`` so GPT MTP checkpoints written before the rename keep loading. - The Megatron-FSDP path never calls ``sharded_state_dict``, so the same mapping is - applied here. - - No-op for models without MTP layers. - """ - mtp_layer_paths = get_mtp_inner_layer_paths(model) - if not mtp_layer_paths: - return model_state_dict, optimizer_state_dict - - model_state_dict = rename_mtp_inner_layer_keys(model_state_dict, mtp_layer_paths) - logger.info( - f"[MTP] Renamed {MTP_INNER_LAYER} -> {MTP_INNER_LAYER_CHECKPOINT_NAME} for " - f"{len(mtp_layer_paths)} MTP layer(s)." - ) - - if optimizer_state_dict is not None and len(optimizer_state_dict.get("state", {})) != 0: - optimizer_state_dict = optimizer_state_dict.copy() - optimizer_state_dict["state"] = rename_mtp_inner_layer_keys( - optimizer_state_dict["state"], mtp_layer_paths - ) - - return model_state_dict, optimizer_state_dict - - def handle_fp8_extra_state_case(model_state_dict): """ Handle the case where FP8 extra state is present in the model state dict. @@ -907,108 +679,6 @@ def print_diff_in_state_dicts(state_dict_metadata, load_state_dict, limit=100): logger.info(f" {k}: meta shape={meta_shape}, load shape={load_shape}") -# Top-level sections holding model weights: "model" for a single chunk, "model0", "model1", -# ... with virtual pipeline parallelism. -_MODEL_SECTION_PATTERN = re.compile(r'^model\d*\.') - - -def get_unexpected_model_keys(state_dict_metadata, load_state_dict): - """Model weights this rank requests that the checkpoint does not contain. - - "Unexpected" follows :class:`~megatron.core.dist_checkpointing.validation.StrictHandling`: - keys accessed locally but absent from the checkpoint. The opposite direction (keys the - checkpoint holds that this rank does not request) is *not* reported, because that is - normal under pipeline, tensor and expert parallelism where every rank loads a subset of - the global checkpoint. - - Only the ``model`` sections are inspected, so optimizer and RNG state that is - intentionally absent does not count. Transformer Engine ``_extra_state`` entries are - skipped because they are dropped from both sides before saving and loading. - """ - metadata_keys = set(flatten_state_dict(state_dict_metadata).keys()) - return { - key - for key in flatten_state_dict(load_state_dict) - if _MODEL_SECTION_PATTERN.match(key) - and not key.endswith('._extra_state') - and key not in metadata_keys - } - - -def validate_fsdp_dtensor_model_load( - state_dict_metadata, - load_state_dict, - checkpoint_path, - strict=StrictHandling.RAISE_UNEXPECTED, - max_reported=20, -): - """Report model weights that an ``fsdp_dtensor`` load cannot supply. - - Torch DCP loads with ``allow_partial_load=True`` skip such keys silently, leaving those - parameters at their initialized values. That does not crash; it surfaces much later as - a convergence regression, so by default fail here instead. - - Only "unexpected" keys are determined (see :func:`get_unexpected_model_keys`), so the - ``*_ALL`` variants of ``strict`` behave like their ``*_UNEXPECTED`` counterparts; - identifying genuinely "missing" keys would require exchanging key sets across ranks. - - ``ASSUME_OK_UNEXPECTED`` is treated as ``RAISE_UNEXPECTED`` here. It means "rely on the - underlying strategy to raise", which a partial DCP load never does, and it exists to skip - the extra disk access an explicit check normally costs — but the caller already holds the - checkpoint metadata, so this check is free. Use ``IGNORE_ALL`` to actually skip it. - - Args: - state_dict_metadata: ``state_dict_metadata`` from the checkpoint's DCP metadata. - load_state_dict: state dict this rank is about to load into. - checkpoint_path: checkpoint location, for the error message. - strict (StrictHandling): how to handle a mismatch. Defaults to raising. - max_reported (int): cap on the number of keys named in the message. - - Returns: - Set[str]: the unexpected model keys, empty when ``strict`` disables the check. - - Raises: - CheckpointingException: if ``strict`` is a ``RAISE_*`` value and keys are missing. - """ - strict = parse_strict_flag(strict) - if strict is StrictHandling.IGNORE_ALL: - return set() - - unexpected_keys = get_unexpected_model_keys(state_dict_metadata, load_state_dict) - if not unexpected_keys or strict in ( - StrictHandling.RETURN_UNEXPECTED, - StrictHandling.RETURN_ALL, - ): - return unexpected_keys - - reported = sorted(unexpected_keys) - truncated = ( - f'\n ... and {len(reported) - max_reported} more' if len(reported) > max_reported else '' - ) - rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 - message = ( - f"Megatron-FSDP checkpoint load is incomplete: rank {rank} requests " - f"{len(reported)} model parameter(s) not present in the checkpoint at " - f"{checkpoint_path}, which would silently keep their initialized values.\n" - + '\n'.join(f' {key}' for key in reported[:max_reported]) - + truncated - + "\nThe model's parameter names have diverged from the checkpoint's, typically " - "because of a fused vs. unfused module variant or a renamed submodule. " - "fsdp_dtensor checkpoints bypass sharded_state_dict(), so key remapping " - "implemented there does not apply; add the mapping to " - "preprocess_fsdp_dtensor_state_dict instead." - ) - - if strict in ( - StrictHandling.ASSUME_OK_UNEXPECTED, - StrictHandling.RAISE_UNEXPECTED, - StrictHandling.RAISE_ALL, - ): - raise CheckpointingException(message) - logger.warning(message) - return unexpected_keys - - def validate_loaded_state_dict(state_dict, checkpoint_path): """ Validate the loaded state dict against the expected structure and types. diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 6ce93c76837..b26bf1a0d3e 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -1,15 +1,14 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from abc import ABC, abstractmethod -from typing import Optional, Union +from dataclasses import dataclass +from typing import Optional, Sequence, Union import torch from megatron.core.inference.utils import InferenceMode from megatron.core.jit import jit_fuser -from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( - is_batch_invariant_mode_enabled, -) +from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.moe.moe_logging import get_moe_metrics_tracker from megatron.core.transformer.moe.moe_utils import ( @@ -20,17 +19,30 @@ apply_router_token_dropping, compute_routing_scores_for_aux_loss, get_tokens_per_expert_and_token_count, - qb_dual_update, router_gating_linear, sinkhorn, switch_load_balancing_loss_func, topk_routing_with_score_function, - z_loss_func, ) from megatron.core.transformer.moe.router_replay import RouterReplay from megatron.core.transformer.transformer_config import TransformerConfig +@dataclass(frozen=True) +class _AuxLossGroupConfig: + """Process groups for local aux/seq-aux loss and its metric logging.""" + + loss_reduce_groups: Sequence[torch.distributed.ProcessGroup] + metric_reduce_group: Optional[torch.distributed.ProcessGroup] + metric_avg_group: Optional[torch.distributed.ProcessGroup] + metric_needs_dp_avg: bool + + @property + def metric_pre_reduce_groups(self) -> Optional[Sequence[torch.distributed.ProcessGroup]]: + """Groups to reduce eagerly before recording metrics, if tracker reduction is unsafe.""" + return self.loss_reduce_groups if self.metric_avg_group is not None else None + + class Router(ABC, MegatronModule): """Base Router class""" @@ -39,6 +51,7 @@ def __init__( config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None, is_mtp_layer: bool = False, + layer_number: Optional[int] = None, ) -> None: """ Initialize the Router module. @@ -52,7 +65,7 @@ def __init__( self.config = config self.num_experts = self.config.num_moe_experts self.moe_aux_loss_func = None - self.layer_number = None + self.layer_number = layer_number self.is_mtp_layer = is_mtp_layer self.tp_group = pg_collection.tp self.cp_group = pg_collection.cp @@ -70,10 +83,6 @@ def __init__( ) else: self.bias = None - if self.config.moe_router_skip_muon: - setattr(self.weight, 'use_muon', False) - if self.bias is not None: - setattr(self.bias, 'use_muon', False) # If calculate per token loss, we need to scale up moe aux loss by the number of tokens. # So we need to know if the model is configured to calculate per token loss. self.calculate_per_token_loss = self.config.calculate_per_token_loss @@ -166,6 +175,7 @@ def __init__( config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None, is_mtp_layer: bool = False, + layer_number: Optional[int] = None, ) -> None: """Initialize the zero token dropping router. @@ -174,14 +184,43 @@ def __init__( pg_collection (ProcessGroupCollection, optional): Process groups for MoE operations. is_mtp_layer (bool): Flag indicating if this router is part of an MTP layer. """ - super().__init__(config=config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer) + super().__init__( + config=config, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + layer_number=layer_number, + ) self.topk = self.config.moe_router_topk self.routing_type = self.config.moe_router_load_balancing_type self.score_function = self.config.moe_router_score_function self.input_jitter = None + self.mtp_layer_number: Optional[int] = None self.frozen_expert_bias = False - self.enable_expert_bias = self.config.moe_router_enable_expert_bias + if self.config.moe_n_hash_layers > 0: + assert layer_number is not None, "layer_number is required for the hash-based router." + self.is_hash_layer = ( + not self.is_mtp_layer + and self.config.moe_n_hash_layers > 0 + and layer_number <= self.config.moe_n_hash_layers + ) + if self.is_hash_layer: + # DSv4-Pro ships a pre-trained tid2eid table in its inference checkpoint; + # no public initialization recipe is documented. Round-robin is used here + # only as a placeholder so the layer is runnable from scratch. + vocab_size = self.config.actual_vocab_size + num_experts = self.config.num_moe_experts + ids = torch.arange(vocab_size, device=torch.cuda.current_device()) + tid2eid = torch.stack([(ids + k) % num_experts for k in range(self.topk)], dim=1).to( + torch.int32 + ) + self.register_buffer('tid2eid', tid2eid) + else: + self.tid2eid = None + + self.enable_expert_bias = ( + self.config.moe_router_enable_expert_bias and not self.is_hash_layer + ) if self.enable_expert_bias: self.register_buffer( 'local_tokens_per_expert', @@ -224,41 +263,6 @@ def __init__( self.global_tokens_per_expert = None self.ga_steps = None - # Quantile balancing replaces the aux loss with a per-expert bias `qb_beta`. - # `qb_beta_accum`/`qb_beta_count` collect the per-microbatch quantile, reduced - # and reset each global batch. - if self.routing_type == "quantile_balancing": - assert not self.is_aux_loss_enabled(), ( - "Quantile balancing handles load balance via the bias update; " - "aux losses must be disabled (set moe_aux_loss_coeff to 0)." - ) - self.register_buffer( - 'qb_beta', - torch.zeros( - self.config.num_moe_experts, - dtype=torch.float32, - device=torch.cuda.current_device(), - ), - ) - self.register_buffer( - 'qb_beta_accum', - torch.zeros( - self.config.num_moe_experts, - dtype=torch.float32, - device=torch.cuda.current_device(), - ), - persistent=False, - ) - self.register_buffer( - 'qb_beta_count', - torch.zeros((), dtype=torch.long, device=torch.cuda.current_device()), - persistent=False, - ) - else: - self.qb_beta = None - self.qb_beta_accum = None - self.qb_beta_count = None - self.router_replay = None if self.config.moe_enable_routing_replay: self.router_replay = RouterReplay() @@ -273,13 +277,6 @@ def _maintain_float32_expert_bias(self): if hasattr(self, 'expert_bias') and self.expert_bias is not None: if self.expert_bias.dtype != torch.float32: self.expert_bias.data = self.expert_bias.data.to(torch.float32) - # Keep the QB bias in fp32 for the same reason. - if hasattr(self, 'qb_beta') and self.qb_beta is not None: - if self.qb_beta.dtype != torch.float32: - self.qb_beta.data = self.qb_beta.data.to(torch.float32) - if hasattr(self, 'qb_beta_accum') and self.qb_beta_accum is not None: - if self.qb_beta_accum.dtype != torch.float32: - self.qb_beta_accum.data = self.qb_beta_accum.data.to(torch.float32) def sinkhorn_load_balancing(self, logits: torch.Tensor): """Apply sinkhorn routing to the logits tensor. @@ -314,81 +311,6 @@ def _sinkhorn_activation(logits): scores = logits * map return scores, map - def quantile_balancing(self, logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Apply quantile-balancing (QB) routing to the logits tensor. - - Selects top-k experts per token using a dual coordinate-descent update on - a per-expert bias ``qb_beta``. Load balance is handled entirely by the bias - update; auxiliary losses must be disabled when QB is active. - - Args: - logits (torch.Tensor): The logits tensor, shape ``[num_tokens, num_experts]``. - - Returns: - Tuple[torch.Tensor, torch.Tensor]: Sparse routing probs and boolean - routing map, each shaped ``[num_tokens, num_experts]``. - """ - assert ( - not self.config.moe_router_fusion - ), "Quantile balancing routing does not support moe_router_fusion." - assert ( - self.config.moe_router_num_groups is None and self.config.moe_router_group_topk is None - ), "Quantile balancing routing does not support group-limited routing." - - local_num_tokens = logits.shape[0] - # Gather logits across TP/CP so the quantile sees a whole sequence's tokens. - # The DP reduction and qb_beta update run at the global-batch boundary in - # finalize_model_grads._update_router_qb_beta. - gather_group = self.tp_cp_group - gather_size = gather_group.size() if gather_group is not None else 1 - - should_update_beta = self.training and torch.is_grad_enabled() - - with torch.no_grad(): - logits_fp32 = logits.detach().to(dtype=torch.float32) - - if gather_size > 1: - full_logits = torch.empty( - (local_num_tokens * gather_size, self.config.num_moe_experts), - dtype=logits_fp32.dtype, - device=logits_fp32.device, - ) - torch.distributed.all_gather_into_tensor( - full_logits, logits_fp32.contiguous(), group=gather_group - ) - gather_rank = torch.distributed.get_rank(group=gather_group) - else: - full_logits = logits_fp32 - gather_rank = 0 - - # Route with the previous batch's qb_beta; in training, accumulate this - # microbatch's quantile for the next update. - full_indices, beta_local = qb_dual_update( - full_logits, self.topk, self.qb_beta, update_beta=should_update_beta - ) - if should_update_beta: - self.qb_beta_accum.add_(beta_local) - self.qb_beta_count.add_(1) - - # Take this rank's rows (all_gather orders rows by rank). - if gather_size > 1: - indices = full_indices[ - gather_rank * local_num_tokens : (gather_rank + 1) * local_num_tokens - ].contiguous() - else: - indices = full_indices - - # QB only picks the experts; reuse the shared score function for the probs. - return topk_routing_with_score_function( - logits, - self.topk, - use_pre_softmax=self.config.moe_router_pre_softmax, - scaling_factor=self.config.moe_router_topk_scaling_factor, - score_function=self.score_function, - fused=self.config.moe_router_fusion, - precomputed_indices=indices, - ) - def get_aux_loss_coeff(self, aux_loss_type: str) -> float: """Return the aux loss coeff for the given auxiliary loss type. If the auxiliary loss type is not found, return 0.0. @@ -417,16 +339,19 @@ def _apply_aux_loss( scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor, with_padding_mask: bool = False, + packed_seq_params: Optional[PackedSeqParams] = None, ): """Apply the auxiliary loss for the given scores and routing map.""" aux_loss_coeff = self.get_aux_loss_coeff("aux_loss") if aux_loss_coeff == 0: return probs + aux_loss_groups = self._get_aux_loss_groups(packed_seq_params) global_tokens_per_expert, local_num_tokens, total_num_tokens = ( get_tokens_per_expert_and_token_count( routing_map=routing_map, - reduce_group=self.tp_cp_group, + reduce_group=aux_loss_groups.loss_reduce_groups[0], + reduce_groups=aux_loss_groups.loss_reduce_groups, topk=self.topk, with_padding_mask=with_padding_mask, ) @@ -441,13 +366,19 @@ def _apply_aux_loss( moe_aux_loss_coeff=aux_loss_coeff, fused=self.config.moe_router_fusion, ) + probs = self.attach_and_log_load_balancing_loss( probs, aux_loss_coeff, aux_loss, "load_balancing_loss", - self.tp_cp_group, + aux_loss_groups.metric_reduce_group, + avg_group=aux_loss_groups.metric_avg_group, + needs_dp_avg=aux_loss_groups.metric_needs_dp_avg, valid_token_count=local_num_tokens, + aux_loss_logging_reduce_groups=aux_loss_groups.metric_pre_reduce_groups, + aux_loss_scale_reduce_groups=aux_loss_groups.loss_reduce_groups, + aux_loss_scale_num_tokens=total_num_tokens, ) return probs @@ -459,6 +390,7 @@ def _apply_seq_aux_loss( seq_length: int, bsz: int, with_padding_mask: bool = False, + packed_seq_params: Optional[PackedSeqParams] = None, ): """Apply the sequence-level auxiliary loss for the given scores and routing map. @@ -474,10 +406,12 @@ def _apply_seq_aux_loss( scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1) routing_map = routing_map.reshape(seq_length, -1) + aux_loss_groups = self._get_aux_loss_groups(packed_seq_params) global_tokens_per_expert, local_num_tokens, total_num_tokens = ( get_tokens_per_expert_and_token_count( routing_map=routing_map, - reduce_group=self.tp_cp_group, + reduce_group=aux_loss_groups.loss_reduce_groups[0], + reduce_groups=aux_loss_groups.loss_reduce_groups, with_padding_mask=with_padding_mask, topk=self.topk * bsz, ) @@ -501,10 +435,13 @@ def _apply_seq_aux_loss( seq_aux_loss_coeff, aux_loss, "seq_load_balancing_loss", - self.tp_cp_group, - # local_num_tokens is per-sequence (bsz folded into the expert dim above); - # * bsz recovers the micro-batch total, else per-token-loss scaling keeps a 1/MBS. - valid_token_count=local_num_tokens * bsz, + aux_loss_groups.metric_reduce_group, + avg_group=aux_loss_groups.metric_avg_group, + needs_dp_avg=aux_loss_groups.metric_needs_dp_avg, + valid_token_count=local_num_tokens, + aux_loss_logging_reduce_groups=aux_loss_groups.metric_pre_reduce_groups, + aux_loss_scale_reduce_groups=aux_loss_groups.loss_reduce_groups, + aux_loss_scale_num_tokens=total_num_tokens, ) return probs @@ -520,7 +457,8 @@ def _apply_global_aux_loss( if global_aux_loss_coeff == 0: return probs - # Use unified function to compute tokens_per_expert and num_tokens + # Global aux loss intentionally uses the full static TP x DP x CP domain. + # Dynamic CP subgroups only affect local aux/seq-aux domains. global_tokens_per_expert, local_num_tokens, total_num_tokens = ( get_tokens_per_expert_and_token_count( routing_map=routing_map, @@ -529,7 +467,6 @@ def _apply_global_aux_loss( topk=self.topk, ) ) - self.global_tokens_per_expert += global_tokens_per_expert self.ga_steps += 1 averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps @@ -543,6 +480,7 @@ def _apply_global_aux_loss( moe_aux_loss_coeff=global_aux_loss_coeff, fused=self.config.moe_router_fusion, ) + probs = self.attach_and_log_load_balancing_loss( probs, global_aux_loss_coeff, @@ -551,18 +489,64 @@ def _apply_global_aux_loss( self.tp_dp_cp_group, needs_dp_avg=False, valid_token_count=local_num_tokens, + # The global aux-loss statistics/logging domain is TP x DP x CP, but + # per-token-loss gradient normalization already reduces the denominator + # across DP x CP in finalize_model_grads. Scale the aux-loss numerator + # over TP x CP only, matching the original static behavior while still + # using an exact valid-token count when padding is present. + aux_loss_scale_reduce_groups=(self.tp_cp_group,), ) return probs + def _get_aux_loss_groups( + self, packed_seq_params: Optional[PackedSeqParams] = None + ) -> _AuxLossGroupConfig: + """Return process groups for MoE aux-loss statistics and logging.""" + if ( + packed_seq_params is not None + and packed_seq_params.local_cp_size is not None + and packed_seq_params.cp_group is not None + ): + return _AuxLossGroupConfig( + loss_reduce_groups=(packed_seq_params.cp_group, self.tp_group), + metric_reduce_group=None, + metric_avg_group=self.tp_dp_cp_group, + metric_needs_dp_avg=False, + ) + + return _AuxLossGroupConfig( + loss_reduce_groups=(self.tp_cp_group,), + metric_reduce_group=self.tp_cp_group, + metric_avg_group=None, + metric_needs_dp_avg=True, + ) + + def _get_metric_layer_number(self) -> tuple[int, int]: + """Return the 1-based metric index and the total number of metric slots.""" + num_layers = self.config.num_layers + (self.config.mtp_num_layers or 0) + if not self.is_mtp_layer: + return self.layer_number, num_layers + + # Hybrid MTP depths can contain multiple internal sublayers (for example `/WE`). + # Metrics are allocated per MTP depth, not per internal hybrid sublayer. + mtp_layer_number = self.mtp_layer_number or self.layer_number + if self.config.mtp_num_layers is not None: + mtp_layer_number = min(mtp_layer_number, self.config.mtp_num_layers) + return self.config.num_layers + mtp_layer_number, num_layers + def attach_and_log_load_balancing_loss( self, activation: torch.Tensor, aux_loss_coeff: float, aux_loss: torch.Tensor, aux_loss_name: str, - reduce_group: torch.distributed.ProcessGroup, + reduce_group: Optional[torch.distributed.ProcessGroup], + avg_group: Optional[torch.distributed.ProcessGroup] = None, needs_dp_avg: bool = True, valid_token_count: Optional[Union[int, torch.Tensor]] = None, + aux_loss_logging_reduce_groups: Optional[Sequence[torch.distributed.ProcessGroup]] = None, + aux_loss_scale_reduce_groups: Optional[Sequence[torch.distributed.ProcessGroup]] = None, + aux_loss_scale_num_tokens: Optional[Union[int, torch.Tensor]] = None, ): """Attach aux loss function to activation and add to logging. @@ -571,7 +555,10 @@ def attach_and_log_load_balancing_loss( aux_loss_coeff (float): Coefficient for the aux loss. aux_loss (torch.Tensor): Computed aux loss. aux_loss_name (str): Name of the aux loss for logging. - reduce_group (torch.distributed.ProcessGroup): Process group for reduction. + reduce_group (torch.distributed.ProcessGroup, optional): Process group for deferred + logging reduction. + avg_group (torch.distributed.ProcessGroup, optional): Process group for deferred + logging average. needs_dp_avg (bool): Whether to average this metric across DP ranks after reduce_group. valid_token_count (int or torch.Tensor, optional): Number of valid tokens excluding padding tokens. Can be a Python int or a torch.Tensor (typically 0-d tensor). @@ -591,53 +578,49 @@ def attach_and_log_load_balancing_loss( # add the aux loss logging value to other layer's since it is difficult to get the # correct layer_number for MTP. It does not affect the correctness of the calculation # results and the reduced load_balancing_loss logging value. - num_layers = self.config.num_layers - if self.config.mtp_num_layers is not None: - num_layers += self.config.mtp_num_layers + layer_number, num_layers = self._get_metric_layer_number() - if self.is_mtp_layer: - layer_number = self.layer_number + self.config.num_layers - else: - layer_number = self.layer_number + metric_value = aux_loss / aux_loss_coeff + if aux_loss_logging_reduce_groups is not None: + metric_value = metric_value.detach().clone() + for group in aux_loss_logging_reduce_groups: + torch.distributed.all_reduce(metric_value, group=group) get_moe_metrics_tracker().record( aux_loss_name, - aux_loss / aux_loss_coeff, + metric_value, layer_number, num_layers, reduce_group=reduce_group, + avg_group=avg_group, needs_dp_avg=needs_dp_avg, ) if self.calculate_per_token_loss: - # Target final scaling on aux_loss gradients: 1 / (num_micro_batches * dp_size), - # matching the !calculate_per_token_loss path. - # - # --calculate-per-token-loss already divides every parameter gradient by - # total_global_tokens (the global non-padded token count summed in - # finalize_model_grads). The router's `num_local_tokens` (= activation.shape[0]) - # is sequence-parallel sharded — the router weight is marked - # `sequence_parallel=True` in Router.reset_parameters (see - # `setattr(self.weight, 'sequence_parallel', ...)` above), so each TP rank - # computes a partial gradient on the router weight from its local sequence - # shard, and `_allreduce_non_tensor_model_parallel_grads` SUMS those partial - # gradients across the TP group. Re-expressing total_global_tokens in terms of the - # router's `num_local_tokens`: - # total_global_tokens - # = num_micro_batches * dp_cp_size * loss_func_local_tokens - # = num_micro_batches * dp_cp_size * tp_size * num_local_tokens - # = num_micro_batches * dp_size * (num_local_tokens * tp_cp_group.size()) - # (using loss_func_local_tokens = tp_size * num_local_tokens, then regrouping - # dp_cp_size * tp_size as dp_size * tp_cp_group.size()). - # - # So pre-multiplying aux_loss by num_local_tokens * tp_cp_group.size() cancels - # that same factor in total_global_tokens above, leaving 1 / (num_micro_batches * - # dp_size) as the effective scaling on the aux_loss gradient — the target. - # Use valid_token_count (excluding padding) if provided, otherwise use total tokens. - num_local_tokens = ( - valid_token_count if valid_token_count is not None else activation.shape[0] - ) + # --calculate-per-token-loss divides all parameter gradients by the global + # non-padded token count in finalize_model_grads. Pre-multiplying by the + # valid-token count from this aux-loss domain makes the final objective a + # token-weighted average of per-domain aux losses. Use the reduced count + # directly: with THD padding or dynamic CP, valid token counts can differ + # by rank/group, so local_num_tokens * group_size is not generally correct. + if aux_loss_scale_num_tokens is None: + num_local_tokens = ( + valid_token_count if valid_token_count is not None else activation.shape[0] + ) + if torch.is_tensor(num_local_tokens): + aux_loss_scale_num_tokens = num_local_tokens.clone().to( + device=activation.device + ) + else: + aux_loss_scale_num_tokens = torch.tensor( + num_local_tokens, device=activation.device + ) + if aux_loss_scale_reduce_groups is None: + assert reduce_group is not None, "reduce_group is required for aux-loss scaling" + aux_loss_scale_reduce_groups = (reduce_group,) + for group in aux_loss_scale_reduce_groups: + torch.distributed.all_reduce(aux_loss_scale_num_tokens, group=group) activation = MoEAuxLossAutoScaler.apply( - activation, aux_loss * num_local_tokens * self.tp_cp_group.size() + activation, aux_loss * aux_loss_scale_num_tokens ) else: activation = MoEAuxLossAutoScaler.apply(activation, aux_loss) @@ -649,66 +632,52 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None): Args: logits (torch.Tensor): The logits of the router. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape in [num_tokens]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [num_tokens]. True = padding, + False = valid. Defaults to None. Returns: torch.Tensor: The logits after applying the z-loss. """ if self.config.moe_z_loss_coeff is not None and self.training and torch.is_grad_enabled(): # Skip Z loss calculations when using torch.no_grad() or checkpointing. - moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size() - z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask) - if self.calculate_per_token_loss: - # Same derivation as in attach_and_log_load_balancing_loss: - # - Target final scaling on z_loss gradients: 1 / (num_micro_batches * dp_size). - # - In terms of the router's `num_local_tokens`, the total_global_tokens - # divisor that finalize_model_grads applies factors as - # num_micro_batches * dp_size * (num_local_tokens * tp_cp_group.size()). - # - Pre-multiplying z_loss by num_local_tokens * tp_cp_group.size() cancels - # that same factor in total_global_tokens, leaving - # 1 / (num_micro_batches * dp_size) as the effective scaling — the target. - # The /tp_cp_group.size() on moe_z_loss_coeff above is a separate forward-side - # correction: z_loss is computed independently on each TP+CP rank's local - # logits and must be averaged across TP+CP rather than summed. - # Count valid tokens: sum of inverted mask (False -> True = valid) - num_local_tokens = ( - (~padding_mask).sum() if padding_mask is not None else logits.shape[0] - ) - logits = MoEAuxLossAutoScaler.apply( - logits, z_loss * num_local_tokens * self.tp_cp_group.size() - ) + logsum = torch.logsumexp(logits, dim=-1) + z_loss_values = torch.square(logsum) + if padding_mask is not None: + valid_mask = ~padding_mask + z_loss_values = z_loss_values * valid_mask + num_local_tokens = valid_mask.sum() + z_loss_sum = z_loss_values.sum() + z_loss_mean = z_loss_sum / torch.clamp(num_local_tokens, min=1) else: - logits = MoEAuxLossAutoScaler.apply(logits, z_loss) + z_loss_sum = z_loss_values.sum() + # Keep the token count as a Python scalar so CUDA graph capture does not + # record a CPU-to-CUDA tensor creation. + z_loss_mean = z_loss_sum / max(logits.shape[0], 1) - # TODO: repeated-MTP z_loss is scaled after MoEAuxLossAutoScaler.apply(), so this - # adjusts logging only; move the scaling above the attach point if z_loss is used. - # When using repeated MTP layers, the same MTP layer is called mtp_num_layers times. - # To avoid accumulating the z_loss multiple times, we scale it by 1/mtp_num_layers - # so the total loss is correct. + mtp_loss_scale = 1 if ( self.is_mtp_layer and self.config.mtp_use_repeated_layer and self.config.mtp_num_layers is not None ): - z_loss = z_loss / self.config.mtp_num_layers - - num_layers = self.config.num_layers - if self.config.mtp_num_layers is not None: - num_layers += self.config.mtp_num_layers + mtp_loss_scale = self.config.mtp_num_layers - if self.is_mtp_layer: - layer_number = self.layer_number + self.config.num_layers + if self.calculate_per_token_loss: + # --calculate-per-token-loss divides gradients by the global non-padded + # token count. Attach the local z-loss numerator directly so the final + # objective is a token-weighted z-loss over valid tokens. + z_loss = z_loss_sum * self.config.moe_z_loss_coeff / mtp_loss_scale + logits = MoEAuxLossAutoScaler.apply(logits, z_loss) else: - layer_number = self.layer_number + moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size() + z_loss = z_loss_mean * moe_z_loss_coeff / mtp_loss_scale + logits = MoEAuxLossAutoScaler.apply(logits, z_loss) + + layer_number, num_layers = self._get_metric_layer_number() get_moe_metrics_tracker().record( - "z_loss", - z_loss / moe_z_loss_coeff, - layer_number, - num_layers, - avg_group=self.tp_dp_cp_group, + "z_loss", z_loss_mean / mtp_loss_scale, layer_number, num_layers ) return logits @@ -744,17 +713,77 @@ def _apply_expert_bias( if self.enable_expert_bias and torch.is_grad_enabled(): with torch.no_grad(): if padding_mask is not None: - routing_map = routing_map & (~padding_mask).unsqueeze(-1) + flat_mask = padding_mask.reshape(-1) + assert ( + flat_mask.shape[0] == routing_map.shape[0] + ), f"padding_mask flat {flat_mask.shape} vs routing_map {routing_map.shape}" + routing_map = routing_map & (~flat_mask).unsqueeze(-1) self.local_tokens_per_expert += routing_map.sum(dim=0) - def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def _hash_routing(self, logits: torch.Tensor, input_ids: torch.Tensor): + """Hash-based routing: expert indices come from the tid2eid lookup table. + + Scores are still computed from the gating logits for weight computation, + but expert selection is determined by the pre-computed hash table. + + Args: + logits (torch.Tensor): Gating logits, shape [num_tokens, num_experts]. + input_ids (torch.Tensor): Token IDs, shape [seq_length, bsz]. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: routing_probs and routing_map. + """ + num_tokens, num_experts = logits.shape + + if self.score_function == "softmax": + scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) + elif self.score_function == "sigmoid": + scores = torch.sigmoid(logits.float()).type_as(logits) + elif self.score_function == "sqrtsoftplus": + scores = torch.nn.functional.softplus(logits.float()).sqrt().type_as(logits) + else: + raise ValueError(f"Invalid score_function: {self.score_function}") + + # input_ids is [b, s] from the model, but hidden_states are [s, b, h] + # and get flattened to [s*b, h]. Transpose to match. + flat_ids = input_ids.T.reshape(-1) + top_indices = self.tid2eid[flat_ids].long() # [num_tokens, topk] + if ( + self.config.moe_router_force_load_balancing + or self.config.moe_router_force_biased is not None + ): + # override top_indices with random topk indices + # logits in processed by apply_random_logits or apply_biased_logits + _, top_indices = torch.topk(logits, k=self.topk, dim=1) + + probs = scores.gather(1, top_indices) + if self.score_function != "softmax": + probs = probs / (probs.sum(dim=-1, keepdim=True) + 1e-20) + + if self.config.moe_router_topk_scaling_factor: + probs = probs * self.config.moe_router_topk_scaling_factor + + routing_probs = torch.zeros_like(logits).scatter(1, top_indices, probs) + routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool() + + return routing_probs, routing_map + + def routing( + self, + logits: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): """Top-k routing function Args: logits (torch.Tensor): Logits tensor after gating. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape [seq_length, bsz]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [seq_length, bsz]. True = padding, + False = valid. Defaults to None. + input_ids (torch.Tensor, optional): The input IDs tensor. Shape [seq_length, bsz]. + Defaults to None. Returns: probs (torch.Tensor): The probabilities of token to experts assignment. @@ -772,13 +801,14 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N logits = self.apply_z_loss(logits, padding_mask=padding_mask) # Calculate probs and routing_map for token dispatching - if self.routing_type == "sinkhorn": + if self.is_hash_layer: + assert input_ids is not None, ( + "input_ids is required for hash-based routing but was None. " + "Ensure --moe-n-hash-layers is set correctly and input_ids are passed." + ) + probs, routing_map = self._hash_routing(logits, input_ids) + elif self.routing_type == "sinkhorn": probs, routing_map = self.sinkhorn_load_balancing(logits) - elif self.routing_type == "quantile_balancing": - assert ( - padding_mask is None - ), "Quantile balancing routing does not support padding masks yet." - probs, routing_map = self.quantile_balancing(logits) else: probs, routing_map = topk_routing_with_score_function( logits, @@ -793,6 +823,20 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N router_replay=self.router_replay, ) + # Dropless HybridEP consumes the sparse routing map directly, so exclude padding + # rows before dispatch. Other dispatchers retain their existing fixed-route + # assumptions until they support sparse routing maps end to end. + use_dropless_hybridep = ( + self.config.moe_token_dispatcher_type == "flex" + and self.config.moe_flex_dispatcher_backend == "hybridep" + and self.config.moe_expert_capacity_factor is None + and self.config.moe_expert_rank_capacity_factor is None + ) + if padding_mask is not None and use_dropless_hybridep: + valid_tokens = (~padding_mask).unsqueeze(-1) + probs = probs * valid_tokens + routing_map = routing_map & valid_tokens + # Apply token dropping to probs and routing_map. if self.config.moe_expert_capacity_factor is not None: probs, routing_map = apply_router_token_dropping( @@ -819,6 +863,7 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N scores_for_aux_loss, routing_map_for_aux_loss, with_padding_mask=padding_mask is not None, + packed_seq_params=packed_seq_params, ) probs = self._apply_seq_aux_loss( probs, @@ -827,6 +872,7 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N seq_length, bsz, with_padding_mask=padding_mask is not None, + packed_seq_params=packed_seq_params, ) probs = self._apply_global_aux_loss( probs, @@ -846,15 +892,23 @@ def reset_global_aux_loss_tracker(self): self.global_tokens_per_expert.zero_() self.ga_steps.zero_() - def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def forward( + self, + input: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): """ Forward pass of the router. Args: input (torch.Tensor): Input tensor. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape [seq_length, bsz]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [seq_length, bsz]. True = padding, + False = valid. Defaults to None. + input_ids (torch.Tensor, optional): The input IDs tensor. Shape [seq_length, bsz]. + Defaults to None. """ self._maintain_float32_expert_bias() @@ -872,7 +926,12 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No logits, self.config.moe_router_force_biased, self.layer_number ) - probs, routing_map = self.routing(logits, padding_mask=padding_mask) + probs, routing_map = self.routing( + logits, + padding_mask=padding_mask, + input_ids=input_ids, + packed_seq_params=packed_seq_params, + ) return probs, routing_map @@ -904,6 +963,7 @@ def __init__( config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None, is_mtp_layer: bool = False, + layer_number: Optional[int] = None, ) -> None: """Initialize the specialized inference top-k router. @@ -921,7 +981,12 @@ def __init__( f"['sigmoid', 'softmax'], got '{config.moe_router_score_function}'" ) - super().__init__(config=config, pg_collection=pg_collection) + super().__init__( + config=config, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + layer_number=layer_number, + ) @staticmethod @torch.compile @@ -937,7 +1002,6 @@ def _compiled_topk_routing( fused, router_replay, dense_output, - precomputed_indices, ): return topk_routing_with_score_function( logits, @@ -951,23 +1015,12 @@ def _compiled_topk_routing( fused=fused, router_replay=router_replay, dense_output=dense_output, - precomputed_indices=precomputed_indices, ) def _forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): logits = self.gating(input).squeeze(1) # [num_tokens, num_experts] - # QB selects on (logits - qb_beta); at inference qb_beta is fixed, so it's per-token. - precomputed_indices = None - if self.qb_beta is not None: - precomputed_indices = (logits - self.qb_beta).topk(self.topk, dim=1).indices - - routing = ( - topk_routing_with_score_function - if is_batch_invariant_mode_enabled() - else self._compiled_topk_routing - ) - probs, top_indices = routing( + probs, top_indices = self._compiled_topk_routing( logits, self.topk, use_pre_softmax=self.config.moe_router_pre_softmax, @@ -979,16 +1032,22 @@ def _forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = N fused=self.config.moe_router_fusion, router_replay=self.router_replay, dense_output=True, - precomputed_indices=precomputed_indices, ) return probs.squeeze(1), top_indices.squeeze(1) - def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def forward( + self, + input: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): """Simplified forward pass for inference - returns dense tensors only. Args: input (torch.Tensor): Input tensor of shape [seq_length, bsz, hidden_size]. padding_mask (torch.Tensor, optional): Not used in inference. + input_ids (torch.Tensor, optional): Not used in inference. Returns: Tuple[torch.Tensor, torch.Tensor]: @@ -997,6 +1056,6 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No """ if not InferenceMode.is_active(): - return super().forward(input, padding_mask) + return super().forward(input, padding_mask, input_ids, packed_seq_params) return self._forward(input, padding_mask) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index d821c14effe..97eee00aac2 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -15,9 +15,11 @@ except ImportError: HAVE_EINOPS = False + from megatron.core import tensor_parallel from megatron.core.dist_checkpointing.mapping import ShardedObject from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.jit import jit_fuser from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, @@ -36,7 +38,6 @@ ) from megatron.core.transformer.attention import Attention, LinearProjBuilder from megatron.core.transformer.enums import AttnMaskType -from megatron.core.transformer.mla_qk_norm_config import QKNormConfigResolver from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.torch_norm import LayerNormBuilder from megatron.core.transformer.transformer_config import MLATransformerConfig @@ -125,6 +126,7 @@ class MLASelfAttentionSubmodules: linear_kv_up_proj: Union[ModuleSpec, type] = None linear_qkv_down_proj: Union[ModuleSpec, type] = None core_attention: Union[ModuleSpec, type] = None + linear_gate: Union[ModuleSpec, type] = None class MultiLatentAttention(Attention): @@ -144,6 +146,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ) -> None: # TODO(nschank): Restructure so that the Attention initializer knows which specific @@ -156,11 +159,16 @@ def __init__( attn_mask_type=attn_mask_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) self.config: MLATransformerConfig self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads + if self.config.gated_attention_proj_granularity == 'elementwise': + self.gate_projection_size = self.query_projection_size + else: + self.gate_projection_size = self.config.num_attention_heads self.q_head_dim = self.config.qk_head_dim + self.config.qk_pos_emb_head_dim @@ -204,6 +212,11 @@ def __init__( "'rope' and 'yarn'" ) + if self.config.experimental_attention_variant == "dsa": + core_attn_extra_kwargs = {"is_mtp_layer": is_mtp_layer} + else: + core_attn_extra_kwargs = {} + self.core_attention = build_module( submodules.core_attention, config=self.config, @@ -215,8 +228,29 @@ def __init__( v_channels=self.config.v_head_dim, cp_comm_type=cp_comm_type, pg_collection=self.pg_collection, + **core_attn_extra_kwargs, ) + if self.config.attention_output_gate: + if submodules.linear_gate is None: + raise ValueError("MLA output gating requires a linear_gate module spec.") + self.linear_gate = build_module( + submodules.linear_gate, + self.config.hidden_size, + self.gate_projection_size, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='mla_gate', + tp_group=self.pg_collection.tp, + name=(name + ".linear_gate") if name is not None else None, + ) + else: + self.linear_gate = None + # Output. self.linear_proj = submodules.linear_proj( self.query_projection_size, @@ -325,6 +359,8 @@ def forward( # hidden_states: [sq, b, h] + gate_input = hidden_states + inference_context = deprecate_inference_params(inference_context, inference_params) if inference_context and not inference_context.is_static_batching(): assert ( @@ -334,6 +370,26 @@ def forward( if self.config.cache_mla_latents: self.prepare_for_absorption() + # Set the right cp group for dynamic-cp. Mirrors Attention.forward: + # downstream RoPE uses self.pg_collection.cp, which must point at this + # microbatch's dynamic CP group. Restored before every return. + _orig_cp_group = self.pg_collection.cp + if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: + assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" + self.pg_collection.cp = packed_seq_params.cp_group + if ( + packed_seq_params is not None + and packed_seq_params.qkv_format == "thd" + and self.pg_collection.cp is not None + and get_pg_size(self.pg_collection.cp) > 1 + and packed_seq_params.cp_partition_mode != "zigzag" + ): + raise ValueError( + "MultiLatentAttention requires cp_partition_mode='zigzag', but " + f"packed_seq_params has {packed_seq_params.cp_partition_mode!r}. CP partition " + "conversion must be handled before entering MLA." + ) + # ===================== # Query, Key, and Value # ===================== @@ -377,10 +433,10 @@ def forward( # core attention computation # ================================== # Need corresponding TE change + needs_output_trim = False core_attn_manager = off_interface( self.offload_core_attention and self.training, query, "core_attn" ) - needs_output_trim = False if self.checkpoint_core_attention and self.training: core_attn_out = self._checkpointed_attention_forward( query, @@ -421,7 +477,6 @@ def forward( cu_kv_lengths, kv_lengths, block_table, - inference_context.is_decode_only(), ) # Only rearrange if not in absorption mode (Flash MLA handles format correctly) if not inference_context.is_decode_only(): @@ -457,6 +512,9 @@ def forward( self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out) self.qkv_up_checkpoint = None + if self.linear_gate is not None: + core_attn_out = self._project_and_apply_mla_output_gate(core_attn_out, gate_input) + # ================= # Output. [sq, b, h] # ================= @@ -465,8 +523,68 @@ def forward( output, bias = apply_module(self.linear_proj)(core_attn_out) output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out]) + self.pg_collection.cp = _orig_cp_group return output, bias + def _project_and_apply_mla_output_gate( + self, core_attn_out: torch.Tensor, gate_input: torch.Tensor + ) -> torch.Tensor: + """Project the gate after core attention, then apply it to the MLA output.""" + gate, _ = apply_module(self.linear_gate)(gate_input) + return self._apply_mla_output_gate(core_attn_out, gate) + + def _apply_mla_output_gate( + self, core_attn_out: torch.Tensor, gate: torch.Tensor + ) -> torch.Tensor: + """Apply an MLA output gate with the configured projection granularity.""" + if core_attn_out.shape[:-1] != gate.shape[:-1]: + raise ValueError( + "MLA output gate and core attention output must have matching token/batch " + f"dimensions, got {tuple(gate.shape)} and {tuple(core_attn_out.shape)}." + ) + + if self.config.gated_attention_proj_granularity == 'elementwise': + if core_attn_out.shape != gate.shape: + raise ValueError( + "Elementwise MLA output gating requires the gate and core attention output " + f"to have the same shape, got {tuple(gate.shape)} and " + f"{tuple(core_attn_out.shape)}." + ) + return self._apply_mla_elementwise_output_gate(core_attn_out, gate) + + if core_attn_out.size(-1) % gate.size(-1) != 0: + raise ValueError( + "Headwise MLA output gating requires the core attention output dimension " + f"({core_attn_out.size(-1)}) to be divisible by the number of local gates " + f"({gate.size(-1)})." + ) + return self._apply_mla_headwise_output_gate(core_attn_out, gate) + + @staticmethod + def _apply_mla_elementwise_output_gate( + core_attn_out: torch.Tensor, gate: torch.Tensor + ) -> torch.Tensor: + """Apply one gate per local MLA output element in the activation dtype.""" + # Keep the FP32 sigmoid followed by native-dtype multiplication used by the + # released model. Do not jit-fuse this helper: nvFuser can move the cast + # across the elementwise multiply and silently change the rounding/VJP. + gate = torch.sigmoid(gate.float()).to(core_attn_out.dtype) + return core_attn_out * gate + + @staticmethod + @jit_fuser + def _apply_mla_headwise_output_gate( + core_attn_out: torch.Tensor, gate: torch.Tensor + ) -> torch.Tensor: + """Apply one gate per local MLA attention head in the activation dtype.""" + output_shape = core_attn_out.shape + core_attn_out = core_attn_out.view(*output_shape[:2], gate.size(-1), -1) + # Compute the sigmoid in FP32 for a stable gate VJP, then cast back to the + # attention output dtype. + gate = torch.sigmoid(gate.float()).to(core_attn_out.dtype) + core_attn_out = core_attn_out * gate.unsqueeze(-1) + return core_attn_out.reshape(output_shape) + class MLASelfAttention(MultiLatentAttention): """MLA Self-attention layer class @@ -484,6 +602,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): if pg_collection is None: @@ -498,17 +617,14 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) - # Resolve which classes to use for Q and KV linear up projections and norms, based on - # QK-norm selection. - layer_classes = self._resolve_qk_norm_config(submodules) - if self.config.q_lora_rank is None: # Not projecting query self.linear_q_proj = build_module( - layer_classes["linear_q_proj"], + submodules.linear_q_proj, self.config.hidden_size, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -555,7 +671,7 @@ def __init__( ) self.linear_q_up_proj = build_module( - layer_classes["linear_q_up_proj"], + submodules.linear_q_up_proj, self.config.q_lora_rank, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -602,7 +718,7 @@ def __init__( ) self.linear_kv_up_proj = build_module( - layer_classes["linear_kv_up_proj"], + submodules.linear_kv_up_proj, self.config.kv_lora_rank, self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim), config=self.config, @@ -617,24 +733,18 @@ def __init__( ) if self.config.q_lora_rank is not None: - self.q_layernorm = layer_classes["q_layernorm"]( + self.q_layernorm = submodules.q_layernorm( hidden_size=self.config.q_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, ) - self.kv_layernorm = layer_classes["kv_layernorm"]( + self.kv_layernorm = submodules.kv_layernorm( hidden_size=self.config.kv_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, ) - def _resolve_qk_norm_config( - self, submodules - ) -> dict[str, ModuleSpec | type | LayerNormBuilder]: - """Resolve which Q/KV norm and up-projection implementations to build.""" - return QKNormConfigResolver(self.config, submodules).resolve() - def _qkv_down_projection(self, hidden_states): """Unfused q/kv down projection path.""" if self.config.q_lora_rank is not None: @@ -681,11 +791,6 @@ def get_query_key_value_tensors( assert ( hidden_states.ndim == 3 ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" - if packed_seq_params is not None: - assert ( - packed_seq_params.local_cp_size is None - ), "hybrid_context_parallel is not supported with MLA yet and is planned for future. \ - Please disable hybrid_context_parallel." inference_context = deprecate_inference_params(inference_context, inference_params) @@ -730,14 +835,9 @@ def get_query_key_value_tensors( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv rope_max_seqlen_q = packed_seq_params.max_seqlen_q rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv - rope_freqs_max_seqlen = ( - max(rope_max_seqlen_q, rope_max_seqlen_kv) - if rope_max_seqlen_q is not None and rope_max_seqlen_kv is not None - else None - ) else: cu_seqlens_q = cu_seqlens_kv = None - rope_freqs_max_seqlen = None + rope_max_seqlen_q = rope_max_seqlen_kv = None # ========================================= # QKV down projection and layernorm @@ -762,7 +862,7 @@ def get_query_key_value_tensors( # k_pos_emb: [s, b, qk_pos_emb_head_dim] k_pos_emb = gather_from_sequence_parallel_region(k_pos_emb, group=self.tp_group) - if packed_seq_params is not None: + if thd_packed_seq: # If sequence packing, TE expect [t, h, d] shaped qkv input. # In Megatron-Core, the qkv shape is [t, 1, h, d]. # So we need to reshape qkv from [t, 1, h, d] to [t, h, d]. @@ -917,7 +1017,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po sequence_start = inference_context.sequence_len_offset sequence_end = sequence_start + q_len rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end] - elif packed_seq_params is None or self.config.context_parallel_size == 1: + elif not thd_packed_seq or self.config.context_parallel_size == 1: # Shorten rotary_pos_emb to the sequence length when inference_params # is not provided. This makes sure we can run forward directly with # any sequence length. During training, the sequence length is always @@ -949,7 +1049,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, - max_seqlen=rope_freqs_max_seqlen, + max_seqlen=rope_max_seqlen_q, ) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( @@ -960,7 +1060,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, - max_seqlen=rope_freqs_max_seqlen, + max_seqlen=rope_max_seqlen_kv, ) # query: [num_tokens, n, (qk_head_dim + v_head_dim)] @@ -1092,6 +1192,15 @@ def backward_dw(self) -> NoReturn: """Execute weight gradient computation""" self._backward_kv_proj() self._backward_q_proj() + if self.linear_gate is not None: + self.linear_gate.backward_dw() + # For the 'dsa' experimental variant core_attention is DSAttention, whose + # indexer linears defer their wgrads under delay_wgrad_compute and need an + # explicit flush. For standard MLA the core is TEDotProductAttention, which + # owns no linears and defines no backward_dw — hence the guard. + core_attention_backward_dw = getattr(self.core_attention, "backward_dw", None) + if core_attention_backward_dw is not None: + core_attention_backward_dw() self._backward_output_proj() def _backward_kv_proj(self): @@ -1249,6 +1358,7 @@ def __init__( attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, pp_layer_offset: Optional[int] = None, name: str | None = None, ): @@ -1264,6 +1374,7 @@ def __init__( attention_type="self", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, pp_layer_offset=pp_layer_offset, name=name, ) @@ -1272,9 +1383,6 @@ def __init__( "FusedMLASelfAttention requires q_lora_rank to be set; " "fallback to MLASelfAttention for q_lora_rank=None." ) - # Resolve which linear class to use for Q and KV up projections, - # based on QK-norm selection. - layer_classes = self._resolve_qk_norm_config(submodules) qkv_down_proj_kwargs = {} if submodules.linear_qkv_down_proj in [TELinear]: @@ -1310,7 +1418,7 @@ def __init__( ) self.linear_q_up_proj = build_module( - layer_classes["linear_q_up_proj"], + submodules.linear_q_up_proj, self.config.q_lora_rank, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -1325,7 +1433,7 @@ def __init__( ) self.linear_kv_up_proj = build_module( - layer_classes["linear_kv_up_proj"], + submodules.linear_kv_up_proj, self.config.kv_lora_rank, self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim), config=self.config, @@ -1339,12 +1447,12 @@ def __init__( name=(name + ".linear_kv_up_proj") if name is not None else None, ) - self.q_layernorm = layer_classes["q_layernorm"]( + self.q_layernorm = submodules.q_layernorm( hidden_size=self.config.q_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, ) - self.kv_layernorm = layer_classes["kv_layernorm"]( + self.kv_layernorm = submodules.kv_layernorm( hidden_size=self.config.kv_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, @@ -1428,8 +1536,7 @@ def _clone_sharded_object_with_key(obj: ShardedObject, new_key: str) -> ShardedO sharded_state_dict[kv_extra_key] = fused_obj for key in list(sharded_state_dict.keys()): - suffix = key[len(fused_prefix) :] if key.startswith(fused_prefix) else "" - if key.startswith(fused_prefix) and not suffix.startswith("layer_norm_"): + if key.startswith(fused_prefix): del sharded_state_dict[key] fused_weight = self.linear_qkv_down_proj.weight @@ -1473,12 +1580,8 @@ def _clone_sharded_object_with_key(obj: ShardedObject, new_key: str) -> ShardedO return sharded_state_dict - def _synthetic_state_dict_key_suffixes(self): - """Return source checkpoint keys used to locate this module in a state dict.""" - return ("linear_q_down_proj.weight",) - - def _synthesize_fused_qkv_down_weight(self, state_dict, prefix): - """Materialize fused qkv-down weight from old separate q/kv checkpoint keys.""" + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Load state dict with automatic unfused->fused conversion.""" q_key = f"{prefix}linear_q_down_proj.weight" kv_key = f"{prefix}linear_kv_down_proj.weight" fused_key = f"{prefix}linear_qkv_down_proj.weight" @@ -1495,8 +1598,4 @@ def _as_tensor(x): state_dict.pop(f"{prefix}linear_q_down_proj.bias", None) state_dict.pop(f"{prefix}linear_kv_down_proj.bias", None) - def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): - """Load state dict with automatic unfused->fused conversion.""" - self._synthesize_fused_qkv_down_weight(state_dict, prefix) - return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 89bc09f64b3..c2063b58242 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1,5 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import inspect import logging import math import warnings @@ -9,7 +10,6 @@ import torch import torch.nn.functional as F -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 @@ -25,8 +25,10 @@ CudaGraphModule, CudaGraphScope, InferenceCudaGraphScope, + LayerType, ) from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout +from megatron.core.utils import experimental_api from .._rank_utils import log_single_rank from ..fusions.fused_bias_geglu import quick_gelu @@ -52,6 +54,7 @@ @dataclass +@experimental_api class TransformerConfig(ModelParallelConfig): """Configuration object for megatron-core transformers. @@ -88,13 +91,6 @@ class TransformerConfig(ModelParallelConfig): 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.""" - mtp_hybrid_override_pattern: Optional[str] = None """DEPRECATED: Use unified hybrid_layer_pattern instead. Legacy argument for loading old checkpoints. @@ -155,16 +151,6 @@ class TransformerConfig(ModelParallelConfig): 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.""" @@ -235,7 +221,7 @@ class TransformerConfig(ModelParallelConfig): 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).""" + is quick_gelu or weighted SwiGLU (MoE only).""" 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 @@ -277,7 +263,34 @@ class TransformerConfig(ModelParallelConfig): 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.""" + """Whether to apply output gating to attention layers. + + The gate projection granularity is controlled by + ``gated_attention_proj_granularity``. It is mutually exclusive with + ``head_wise_attn_gate``. + """ + + gated_attention_proj_granularity: Literal['elementwise', 'headwise'] = "elementwise" + """Projection granularity for ``attention_output_gate``. + + ``elementwise`` projects one gate per attention output element. ``headwise`` projects one + scalar gate per attention head and is currently supported only by Multi-Latent Attention. + """ + rotary_base_per_layer: Optional[List[float]] = None + """Per-layer RoPE theta values. Length must equal num_layers. When set, each + SelfAttention layer creates its own RotaryEmbedding with the corresponding base; + the shared model-level rotary_pos_emb is not created.""" + head_wise_attn_gate: bool = False + """Apply a per-head scalar output gate (Step-3.5-Flash g_proj): + num_attention_heads scalar gates fused as the trailing rows of + linear_qkv; sigmoid scales each head uniformly across head_dim. + Contrast with `attention_output_gate` (full head_dim gate, inline + per-group); the two are mutually exclusive. Self-attention only. + dist-ckpt TP resharding is handled by a ShardedTensorFactory in + SelfAttention.sharded_state_dict. Requires (validated in + __post_init__): num_attention_heads % tp == 0, + num_query_groups >= tp, and under fp8/fp4 a per-partition + linear_qkv_out_dim aligned to 16/32.""" test_mode: bool = False """Whether to run real-time tests.""" @@ -299,15 +312,22 @@ class TransformerConfig(ModelParallelConfig): #################### # attention variant #################### - experimental_attention_variant: Optional[Literal['gdn', 'gdn2', 'dsa', 'gated_delta_net']] = ( - None - ) - """Type of attention variant to use. Currently support gdn, gdn2 and dsa. - 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: Optional[ + Literal['gdn', 'kda', 'gated_delta_net', 'dsa', 'dsv4_hybrid'] + ] = None + """Type of experimental attention variant to use. + + ``gdn`` and ``kda`` select Gated DeltaNet and Kimi Delta Attention, respectively. + ``gated_delta_net`` is a deprecated compatibility alias for ``gdn`` and is normalized in + ``__post_init__`` with a ``DeprecationWarning``. + """ + + cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag" + """How THD sequence rows are partitioned across context-parallel ranks. + + Contiguous partitioning is defined only for THD inputs; BSHD context parallelism uses the + standard zigzag layout. + """ experimental_attention_variant_loss_scale_func: Optional[Callable[[torch.Tensor], None]] = None """Optional hook for experimental attention variants to receive the main loss scale.""" @@ -339,18 +359,16 @@ class TransformerConfig(ModelParallelConfig): top-k indices.""" dsa_kernel_backend: Literal["none", "tilelang", "cudnn"] = "none" - """Optional fused DSA kernel backend. - ``none`` disables fused DSA kernels. Explicit ``tilelang`` or ``cudnn`` enables only that - backend. Unsupported DSA layouts continue to use the PyTorch fallback.""" + """Optional fused ordinary-DSA kernel backend. Unsupported layouts use 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.""" + """Whether DSA indexer should apply Hadamard rotation before scoring.""" dsa_indexer_scoring_relu: bool = True - """Whether DSA indexer should apply ReLU to q@k^T scores before weighting.""" + """Whether DSA indexer should apply ReLU to q@k scores before weighting.""" dsa_indexer_k_norm_epsilon: Optional[float] = None """Optional epsilon override for the DSA indexer key LayerNorm.""" @@ -359,16 +377,14 @@ class TransformerConfig(ModelParallelConfig): """Whether DSA indexer key LayerNorm should run on fp32 inputs.""" #################### - # Compressed sparse attention + # DeepSeek-v4 hybrid attention #################### csa_window_size: int = 128 """Sliding window size for compressed sparse attention.""" - # TODO(#6402): consumed by DSv4 Hybrid attention orchestration, which selects the - # per-layer compression ratio and builds the compressed-KV rotary embedding. - # Neither field has a production reader in this primitive-only PR. csa_compress_ratios: Optional[List[int]] = None - """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...].""" + """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...]. A value of 0 is a + sliding-window-only layer (no compressor / no top-k indexer; the 'W' hybrid layer symbol).""" csa_compress_rotary_base: float = 40000.0 """RoPE base for compressed KV positions in compressed sparse attention.""" @@ -377,9 +393,16 @@ class TransformerConfig(ModelParallelConfig): """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be disabled.""" + apply_dsa_kernel_fusion: Optional[bool] = None + """Deprecated DSv4 fused-kernel switch. Use ``dsa_kernel_backend`` instead. For + ``dsv4_hybrid`` only, True maps to ``cudnn`` and False maps to ``none``.""" + #################### # linear attention #################### + linear_attention_type: Optional[str] = None + """Type of linear attention to use. + Deprecated. Use experimental_attention_variant instead.""" linear_attention_freq: Optional[Union[int, List[int]]] = None """Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) layers. @@ -402,6 +425,20 @@ class TransformerConfig(ModelParallelConfig): linear_num_value_heads: Optional[int] = 32 """Number of value and gate heads for the gated delta net.""" + kda_safe_gate: bool = False + """Whether the KDA kernel should use bounded gate values.""" + + kda_lower_bound: Optional[float] = None + """Optional lower bound for KDA's bounded gate values.""" + + gdn_pre_gated_delta_rule_fusion: bool = False + """Whether to use the streamed Triton fusion for GatedDeltaNet pre-GDR preprocessing.""" + + gdn_conv_pad_alignment: Optional[int] = None + """When set, pad packed GDN causal-conv inputs to this token alignment. + This is only valid without chunkwise CP: padding a chunk-local causal-conv input changes + the sequence seen by later chunks and therefore changes the GDN recurrence numerics.""" + #################### # initialization #################### @@ -539,13 +576,13 @@ class TransformerConfig(ModelParallelConfig): 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.""" + fused_residual_rmsnorm: bool = False + """If True, fuses residual connection and RMSNorm backward pass when TE is used.""" + #################### # activation recomputation #################### @@ -580,7 +617,7 @@ class TransformerConfig(ModelParallelConfig): 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". + "shared_experts", "mhc", "gdn", "gdn_norm_out". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -589,16 +626,17 @@ class TransformerConfig(ModelParallelConfig): "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". - "moe_act", "layernorm", "mla_up_proj", "gdn_norm_out", "gdp_in_proj", "gdp_qkv", and - "mhc" use output-discarding checkpointing, "core_attn", "mlp", "moe", and - "shared_experts" use normal checkpointing. + enable_hyper_connections=True. Cannot be used with "mlp". + "gdn": recompute the entire GDN-family layer, including GatedDeltaNet and KDA + (input projections, conv1d, gated delta rule, gated norm, CP all-to-all, and + out_proj). Requires a GDN-family experimental attention variant or a hybrid model. + "gdn_norm_out": recompute gated output normalization and layout restoration for + Gated DeltaNet-family layers, including GatedDeltaNet and KDA. + "moe_act", "layernorm", "mla_up_proj", "mhc", and "gdn_norm_out" use + output-discarding checkpointing, + "core_attn", "mlp", "moe", "shared_experts", and "gdn" use normal checkpointing. """ #################### @@ -731,7 +769,7 @@ class TransformerConfig(ModelParallelConfig): in the hidden_states gradient.""" moe_shared_expert_gate: bool = False - """Enable gate for shared expert. Only effective when + """Enable gate for shared expert. Only effective when moe-shared-expert-intermediate-size is set.""" moe_shared_expert_overlap: bool = False @@ -769,9 +807,6 @@ class TransformerConfig(ModelParallelConfig): 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". @@ -822,8 +857,8 @@ class TransformerConfig(ModelParallelConfig): 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.""" + """Scaling factor for routing score in top-k selection, only works when moe_router_pre_softmax + enabled. 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".""" @@ -833,9 +868,6 @@ class TransformerConfig(ModelParallelConfig): 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. @@ -847,12 +879,6 @@ class TransformerConfig(ModelParallelConfig): 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.""" @@ -864,36 +890,49 @@ class TransformerConfig(ModelParallelConfig): If negative, generates bias once per layer and reuses it (abs value is std). This is an experimental feature for benchmarking purposes.""" - use_grouped_gemm_for_dense_mlp: bool = False + moe_n_hash_layers: int = 0 + """Number of leading transformer layers that use hash-based MoE routing. + Layers with layer_number <= moe_n_hash_layers use a pre-computed tid2eid + lookup table for expert selection instead of learned top-k routing.""" + + actual_vocab_size: Optional[int] = None + """Padded actual vocabulary size. Required when moe_n_hash_layers > 0 for the + tid2eid lookup buffer in hash-based MoE routing.""" + + dense_grouped_gemm: 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. + use_grouped_gemm_for_dense_mlp: bool = False + """Alias of ``dense_grouped_gemm``. 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_use_grouped_tensor: bool = False - """Use Transformer Engine's native GroupedTensor path for grouped MoE GEMMs. + log_moe_overload_factor: bool = False + """When True, log MoE overload metrics (avg/max vs balanced token count per step; max cum + overload = peak cumulative actual tokens / peak cumulative balanced count over interleaved + fwd/bwd) to TensorBoard/W&B and console. Records tokens_per_expert.sum() after dispatch; + use for debugging.""" - 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_grouped_gemm: bool = False + """When there are multiple experts per rank, compress multiple local (potentially small) gemms + in a single kernel launch to improve the utilization and performance by leveraging the Grouped + GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm). """ 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``. + ``use_transformer_engine_op_fuser=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``.""" + parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=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. @@ -918,20 +957,22 @@ class TransformerConfig(ModelParallelConfig): 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" + moe_flex_dispatcher_backend: Literal['deepep', 'deepepv2', '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 + Options are "deepep", "deepepv2", "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_hybridep_pad_variable_tokens: bool = False + """Dynamically pad uneven local token counts to the HybridEP group maximum before dispatch. + + Enable this when local token counts can differ across ranks. When disabled, the caller must + guarantee equal token counts across the HybridEP communication group, for example by padding + THD inputs to a fixed maximum before dispatch. CUDA Graph inputs should be statically padded + upstream and leave this option disabled. """ moe_per_layer_logging: bool = False @@ -941,6 +982,11 @@ class TransformerConfig(ModelParallelConfig): """moe_expert_capacity_factor (float): The capacity factor for each expert, None means no token will be dropped. The default is None.""" + moe_expert_rank_capacity_factor: Optional[float] = None + """moe_expert_rank_capacity_factor (float): The capacity factor for each expert rank. Tokens + exceeding this budget will be dropped. 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 @@ -975,19 +1021,6 @@ class TransformerConfig(ModelParallelConfig): 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 @@ -1014,26 +1047,22 @@ class TransformerConfig(ModelParallelConfig): moe_hybridep_num_sms_preprocessing: int = 108 """Number of SMs to use for HybridEP preprocessing (metadata scan kernel).""" - moe_ncclep_zero_copy: bool = False + moe_ncclep_static_shape: bool = False + """For the 'ncclep' flex dispatcher: feed the experts the full fixed-size receive buffer + instead of narrowing to the (data-dependent) number of received tokens, removing the D2H sync + and dynamic shapes from the dispatch (required for CUDA-graph capture of the MoE A2A and for the + 1F1B EP comm overlap). The fused grouped GEMM consumes the ragged per-expert counts on device + and walks only the received tokens (no slack GEMM, no last-expert padding). This requires the + CuTe DSL / device-offset grouped GEMM, so it is only supported with the fused op + (use_transformer_engine_op_fuser, NVTE_CUTEDSL_FUSED_GROUPED_MLP=1) on sm100+ (Blackwell or + later); the dispatcher asserts this. On older GPUs leave it False (dynamic shape). Defaults to + False (narrow to the received tokens).""" + + moe_ncclep_use_symm_mem: 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).""" + staged-copy path. NOT SUPPORTED YET -- the dispatcher rejects this if set; the cross-stream + reuse ordering for the persistent symm-mem buffer is not implemented. Leave False.""" moe_mlp_glu_interleave_size: Optional[int] = None """When set, GLU activations in the MoE grouped MLP layer will use a @@ -1043,14 +1072,6 @@ class TransformerConfig(ModelParallelConfig): 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 ################## @@ -1059,6 +1080,8 @@ class TransformerConfig(ModelParallelConfig): 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". + This option controls standard attention layers. Linear-attention layers use + `linear_cp_mode` instead. "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 @@ -1070,11 +1093,18 @@ class TransformerConfig(ModelParallelConfig): and P2P communications in high-level CP groups (e.g., via IBLink). """ - linear_cp_layout: CPLayout = "zigzag" - """CP layout for linear-attention layers.""" - - attention_cp_layout: CPLayout = "zigzag" - """CP layout for softmax-attention layers.""" + linear_cp_mode: Optional[str] = "chunkwise" + """Context-parallel execution mode for linear-attention layers + (e.g. Gated Delta Net). Independent of `cp_comm_type`, which only controls standard attention. + Can be "chunkwise" or "headwise": + "chunkwise": Keep sequence chunks sharded across CP ranks and use CP-aware linear kernels + (e.g. chunk_gated_delta_rule + causal_conv1d with a CP context). This follows the chunkwise + DeltaNet idea of storing state at chunk boundaries and doing chunk-local matrix work, avoiding + a full per-token recurrent state materialization while keeping tensor-core-friendly matmuls. + See https://sustcsonglin.github.io/blog/2024/deltanet-2/#a-chunkwise-algorithm-for-deltanet. + "headwise": Scatter heads across the CP group with all-to-all (Ulysses-style); each rank runs + the linear-attention kernel on the full sequence for a shard of heads. Correct but memory-heavy. + """ ################## # Cuda Graphs @@ -1086,7 +1116,7 @@ class TransformerConfig(ModelParallelConfig): 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. + """For cuda_graph_impl "full_iteration" only. When True, full-iteration graph replay (training and evaluation) and optimizer graph capture/replay share the same CUDA graph memory pool.""" @@ -1097,9 +1127,7 @@ class TransformerConfig(ModelParallelConfig): 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.""" + """Number of warmup steps for CUDA graphs""" external_cuda_graph: bool = False """DEPRECATED and replaced by cuda_graph_impl. @@ -1138,6 +1166,14 @@ class TransformerConfig(ModelParallelConfig): 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__.""" + create_attention_mask_in_dataloader: bool = True + """Whether training data loaders create and pass an attention_mask tensor. + + This mirrors the training argument of the same name so Transformer Engine CUDA graph capture + can preserve the actual forward signature. When disabled for causal/no-mask attention, + the graph capture path should not allocate a synthetic global-sequence attention mask. + """ + inference_cuda_graph_scope: Optional[InferenceCudaGraphScope] = field( default=None, metadata={ @@ -1169,13 +1205,33 @@ class TransformerConfig(ModelParallelConfig): CudaGraphScope instances deserialized from pre-refactor checkpoints are converted to their string names before normalization so existing CUDA_GRAPH_MODULES_DEPRECATIONS handles them.""" + thd_max_packed_sequences: Optional[int] = field( + default=None, metadata={"argparse_meta": {"arg_names": ["--thd-max-packed-sequences"]}} + ) + """Optional maximum number of THD packed sequences per microbatch, including any + dummy sequence appended for a padding tail. The dp_balanced packing scheduler reserves + that dummy slot when THD padding appends one. When set, cu_seqlens tensors are padded + to this size + 1 in both eager and CUDA Graph modes. THD CUDA Graph requires this value. + + Sizing guidance: choose a value that comfortably covers the worst-case packing, + roughly ceil(max_seqlen_per_dp_cp_rank * cp_size / min_seq_len_after_filter). + Setting it too small results in more microbatches with smaller packs (wasted + token budget); setting it too large just allocates a slightly larger cu_seqlens + buffer.""" + + cuda_graph_dynamic_microbatches: bool = False + """Allow CUDA graph replay when runtime microbatch count varies across iterations. + This option is only meaningful for cuda_graph_impl=transformer_engine. For THD sequence + packing, capture uses a conservative upper bound on the packed microbatch count so graph + replay can cover iterations whose real packed microbatch count changes.""" + #################### # Hyper-Connection Configuration #################### - enable_mhc_connections: bool = False + enable_hyper_connections: bool = False """Enable mHC residual connections.""" - mhc_num_residual_streams: int = 4 + num_residual_streams: int = 4 """Number of residual streams (n in paper).""" mhc_sinkhorn_iterations: int = 20 @@ -1185,35 +1241,34 @@ class TransformerConfig(ModelParallelConfig): """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. + """Use unified fused kernels for mHC operations. + + When True, attempts to replace the reference mHC modules (SinkhornKnopp, + H_aggregate, H_post_bda, ProjRms) with fused/autograd implementations for + better performance on supported GPUs. Backend selection is internal and + op-specific: Triton for Sinkhorn and H_post_bda backward when available, + cuTile for the remaining fused kernels when available, then native torch + fallback. If every mHC operation uses the native torch fallback, + use_fused_mhc remains enabled and a rank-0 warning is emitted. The all-native + fallback is functionally equivalent, but may not provide fused backend + performance benefits. """ mhc_recompute_layer_num: Optional[int] = None - """Number of layers per MHC recompute block. + """Number of layers in each mHC recompute group. - 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 + Layers are grouped in their local transformer-block order. The last layer in each group leaves + its final MLP BDA output uncheckpointed and closes the current + ``CheckpointWithoutOutputManager``; a new manager is created for the next group. - If None, all layers in the transformer block share a single recompute block. + In the standard forward path, the group-ending layer registers a unified recompute hook on its + MLP BDA output. In the fine-grained expert-parallel overlap schedule, the group outputs are + discarded explicitly and a compute-stream schedule node replays the group before its backward + computation. - Must be a positive integer when set.""" + If ``None``, all local layers in the transformer block share one recompute group. The value must + be a positive integer when set. + """ #################### # miscellaneous @@ -1238,24 +1293,6 @@ class TransformerConfig(ModelParallelConfig): 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""" @@ -1313,6 +1350,12 @@ class TransformerConfig(ModelParallelConfig): """ Multimodal rope section is for channel dimension of temporal, height and width in rope calculation. """ + mrope_interleaved: bool = False + """When True, use the interleaved T/H/W MRoPE layout (Qwen3.5-VL style) where + H freqs occupy stride-3 positions {1,4,7,...} and W freqs occupy {2,5,8,...}. + When False (default), use the original section-based layout (Qwen2-VL style) + that cycles through T/H/W per mrope_section chunk.""" + is_hybrid_model: bool = False """ Indicates whether this is a hybrid model. """ @@ -1329,9 +1372,6 @@ class TransformerConfig(ModelParallelConfig): """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.""" @@ -1341,17 +1381,6 @@ class TransformerConfig(ModelParallelConfig): ) """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 the 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.""" @@ -1377,6 +1406,9 @@ class TransformerConfig(ModelParallelConfig): """Transformer implementation to use. Options are 'transformer_engine' for Transformer Engine and 'local' for MCore.""" + fallback_to_eager_attn: bool = False + """Whether to fallback to eager attention in TE implementation. + Suggested for when desired features are not available in TE implementation.""" ##################################### # Fine-grained Activation Offloading ##################################### @@ -1388,7 +1420,7 @@ class TransformerConfig(ModelParallelConfig): 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". + "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp". "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. @@ -1397,8 +1429,6 @@ class TransformerConfig(ModelParallelConfig): "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. """ min_offloaded_tensor_size: int = 1024 * 1024 """The minimum size of the tensor to be offloaded.""" @@ -1421,6 +1451,14 @@ class TransformerConfig(ModelParallelConfig): https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#activation-offload-fraction. """ + 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""" + moe_paged_stash: bool = False """If True, enable paged stash for all routed-expert activations needed for backward""" @@ -1431,7 +1469,7 @@ class TransformerConfig(ModelParallelConfig): """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%%).""" + (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. @@ -1439,84 +1477,19 @@ class TransformerConfig(ModelParallelConfig): 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""" - - def _validate_cp_layouts(self) -> None: - """Validate context-parallel layout settings.""" - if self.linear_cp_layout not in ("contiguous", "zigzag"): - raise ValueError( - "linear_cp_layout must be either 'contiguous' or 'zigzag', " - f"got {self.linear_cp_layout!r}" - ) - if self.attention_cp_layout not in ("contiguous", "zigzag"): - raise ValueError( - "attention_cp_layout must be either 'contiguous' or 'zigzag', " - f"got {self.attention_cp_layout!r}" - ) - 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.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}." - ) - if ( - self.linear_cp_layout == "contiguous" - and self.context_parallel_size > 1 - and (self.mtp_num_layers or 0) > 0 - ): - raise ValueError( - "linear_cp_layout='contiguous' with context parallelism does not yet support MTP." - ) - 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() - # 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. + # Imported lazily because the module-spec module imports TransformerConfig. 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.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.") - - if self.mtp_hsm and (self.mtp_num_layers is None or self.mtp_num_layers < 2): - raise ValueError("mtp_hsm=True requires mtp_num_layers >= 2.") - # 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: @@ -1535,11 +1508,6 @@ def __post_init__(self): 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 @@ -1559,6 +1527,17 @@ def __post_init__(self): if self.num_query_groups is None: self.num_query_groups = self.num_attention_heads + if self.gated_attention_proj_granularity not in ('elementwise', 'headwise'): + raise ValueError( + "gated_attention_proj_granularity must be either 'elementwise' or 'headwise', " + f"got {self.gated_attention_proj_granularity!r}." + ) + if self.gated_attention_proj_granularity == 'headwise' and not self.multi_latent_attention: + raise ValueError( + "Regular attention does not support headwise " + "gated_attention_proj_granularity; use 'elementwise'." + ) + if ( self.num_query_groups % self.tensor_model_parallel_size != 0 and self.tensor_model_parallel_size % self.num_query_groups != 0 @@ -1568,47 +1547,229 @@ def __post_init__(self): f"tensor_model_parallel_size ({self.tensor_model_parallel_size})." ) + if self.head_wise_attn_gate: + tp = self.tensor_model_parallel_size + if self.attention_output_gate: + raise ValueError( + "head_wise_attn_gate and attention_output_gate cannot both be enabled " + "(incompatible linear_qkv row layouts)." + ) + if self.num_attention_heads % tp != 0: + raise ValueError( + f"head_wise_attn_gate requires num_attention_heads " + f"({self.num_attention_heads}) divisible by tp ({tp})." + ) + # The gate peel-off in SelfAttention runs before the + # num_query_groups < world_size AllGather+reslice fallback, so the + # gate cannot ride that path without over-taking V/K rows. + if self.num_query_groups < tp: + raise ValueError( + f"head_wise_attn_gate requires num_query_groups " + f"({self.num_query_groups}) >= tp ({tp})." + ) + # TE FP8/FP4 GEMM requires per-partition output dim aligned to + # 16/32. attention_output_gate's wider gate gets this for free; + # head_wise's num_attention_heads-row increment can mis-align. + if self.fp8 is not None or self.fp4 is not None: + align = 32 if self.fp4 is not None else 16 + linear_qkv_out_dim = ( + self.kv_channels * self.num_attention_heads + + 2 * self.kv_channels * self.num_query_groups + + self.num_attention_heads + ) + per_partition = linear_qkv_out_dim // tp + if per_partition % align != 0: + fp_name = "fp4" if self.fp4 is not None else "fp8" + raise ValueError( + f"head_wise_attn_gate under {fp_name} requires per-partition " + f"linear_qkv output dim ({per_partition}) to be a multiple of " + f"{align}; got num_attention_heads={self.num_attention_heads}, " + f"num_query_groups={self.num_query_groups}, " + f"kv_channels={self.kv_channels}, tp={tp}." + ) + + if self.linear_attention_type is not None: + warnings.warn( + "linear_attention_type is deprecated, " + "use experimental_attention_variant instead." + ) + self.experimental_attention_variant = self.linear_attention_type + self.linear_attention_type = None + + if self.experimental_attention_variant is not None: + self.experimental_attention_variant = normalize_experimental_attention_variant( + self.experimental_attention_variant + ) + + if self.cp_partition_mode not in ("zigzag", "contiguous"): + raise ValueError(f"Unsupported cp_partition_mode: {self.cp_partition_mode}") + + if self.cp_partition_mode == "contiguous" and ( + self.context_parallel_size > 1 or self.dynamic_context_parallel + ): + if self.sequence_packing_scheduler is None: + raise ValueError( + "cp_partition_mode='contiguous' with context parallelism requires THD " + "inputs from a sequence_packing_scheduler; BSHD inputs are not supported. " + "The legacy non-scheduler CP batch slicing path only supports zigzag layout." + ) + if ( + (self.mtp_num_layers or 0) > 0 + and self.tensor_model_parallel_size > 1 + and self.sequence_parallel + ): + raise ValueError( + "MTP with tensor_model_parallel_size > 1, sequence_parallel=True, and " + "cp_partition_mode='contiguous' has a known token-side padding-mask layout " + "bug. This combination is temporarily unsupported and will be fixed in a " + "follow-up change." + ) + + if self.context_parallel_size > 1: + if self.cp_partition_mode == "contiguous": + if ( + self.multi_latent_attention + and self.experimental_attention_variant != "dsv4_hybrid" + ): + raise ValueError( + "cp_partition_mode='contiguous' is not supported with " + "multi_latent_attention outside dsv4_hybrid." + ) + if self.experimental_attention_variant != "dsv4_hybrid" and not ( + is_gated_delta_net_variant(self.experimental_attention_variant) + ): + raise ValueError( + "cp_partition_mode='contiguous' with context parallelism currently " + "requires experimental_attention_variant to be 'dsv4_hybrid', 'gdn', " + "or 'kda'." + ) + if ( + is_gated_delta_net_variant(self.experimental_attention_variant) + and self.linear_cp_mode == "headwise" + ): + raise ValueError( + "cp_partition_mode='contiguous' is incompatible with " + "GDN-family linear_cp_mode='headwise'." + ) + elif self.cp_partition_mode == "zigzag": + if self.experimental_attention_variant == "dsv4_hybrid": + raise ValueError( + "DSv4 Hybrid with context parallelism requires " + "cp_partition_mode='contiguous'." + ) + + # Normalize the deprecated DSv4 kernel switch only after all deprecated attention + # selectors have been folded into experimental_attention_variant, and immediately + # before the centralized attention-variant validation consumes dsa_kernel_backend. + if self.apply_dsa_kernel_fusion is not None: + if self.experimental_attention_variant != "dsv4_hybrid": + log_single_rank( + logger, + logging.WARNING, + "apply_dsa_kernel_fusion is deprecated and ignored outside " + "experimental_attention_variant='dsv4_hybrid'; use dsa_kernel_backend " + "instead.", + ) + else: + legacy_backend = "cudnn" if self.apply_dsa_kernel_fusion else "none" + if self.dsa_kernel_backend not in ("none", legacy_backend): + raise ValueError( + "Conflicting DSA kernel controls: " + f"apply_dsa_kernel_fusion={self.apply_dsa_kernel_fusion} maps to " + f"dsa_kernel_backend={legacy_backend!r}, but " + f"dsa_kernel_backend={self.dsa_kernel_backend!r} was also selected. " + "Remove apply_dsa_kernel_fusion and use dsa_kernel_backend only." + ) + log_single_rank( + logger, + logging.WARNING, + "apply_dsa_kernel_fusion is deprecated and will be removed in a future " + f"release; use dsa_kernel_backend={legacy_backend!r} instead.", + ) + self.dsa_kernel_backend = legacy_backend + 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": + if not self.is_hybrid_model: assert ( self.linear_attention_freq is not None - ), "linear_attention_freq must be set for linear gdn." + ), "linear_attention_freq must be set for linear attention." + if self.pad_packed_seq_alignment is not None: + tail_policy = self.thd_tail_padding_policy or 'append_dummy_seq' + assert tail_policy == 'append_dummy_seq', ( + "GDN-family attention with pad_packed_seq_alignment requires " + "thd_tail_padding_policy='append_dummy_seq'." + ) - # Check required parameters + # Required by both standalone GDN and hybrid KDA layers. assert ( self.linear_conv_kernel_dim is not None - ), "linear_conv_kernel_dim must be set for gated delta net." + ), "linear_conv_kernel_dim must be set for a GDN-family layer." assert ( self.linear_key_head_dim is not None - ), "linear_key_head_dim must be set for gated delta net." + ), "linear_key_head_dim must be set for a GDN-family layer." assert ( self.linear_value_head_dim is not None - ), "linear_value_head_dim must be set for gated delta net." + ), "linear_value_head_dim must be set for a GDN-family layer." assert ( self.linear_num_key_heads is not None - ), "linear_num_key_heads must be set for gated delta net." + ), "linear_num_key_heads must be set for a GDN-family layer." assert ( self.linear_num_value_heads is not None - ), "linear_num_value_heads must be set for gated delta net." + ), "linear_num_value_heads must be set for a GDN-family layer." + + if self.experimental_attention_variant == "kda": + if self.linear_num_key_heads != self.linear_num_value_heads: + raise ValueError("KDA requires equal key and value head counts.") + if self.linear_key_head_dim != self.linear_value_head_dim: + raise ValueError("KDA requires equal key and value head dimensions.") + if self.kda_safe_gate: + if self.kda_lower_bound is None: + raise ValueError("KDA requires kda_lower_bound when kda_safe_gate=True.") + if not (-5.0 <= self.kda_lower_bound < 0.0): + raise ValueError( + "KDA requires kda_lower_bound to be in [-5, 0) " + "when kda_safe_gate=True." + ) + 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})." ) + if ( + self.experimental_attention_variant == "kda" or self.context_parallel_size > 1 + ) and self.linear_cp_mode not in ("headwise", "chunkwise"): + raise ValueError( + f"linear_cp_mode must be either 'headwise' or 'chunkwise', " + f"got {self.linear_cp_mode!r}." + ) + if self.gdn_conv_pad_alignment is not None: + assert self.gdn_conv_pad_alignment > 0, ( + f"gdn_conv_pad_alignment must be positive when set, " + f"got {self.gdn_conv_pad_alignment}." + ) - # 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, ( + if self.context_parallel_size > 1: + if self.gdn_conv_pad_alignment is not None: + assert self.linear_cp_mode != "chunkwise", ( + "gdn_conv_pad_alignment is incompatible with " + "linear_cp_mode='chunkwise' when context_parallel_size > 1. " + "Padding chunk-local GDN causal-conv inputs can change later " + "chunk numerics." + ) + # Check tensor parallelism compatibility. Headwise CP splits linear-attention heads + # across CP ranks; chunkwise CP keeps all TP-local heads on each CP rank. + linear_head_parallel_size = self.tensor_model_parallel_size + if self.context_parallel_size > 1 and self.linear_cp_mode == "headwise": + linear_head_parallel_size *= self.context_parallel_size + assert self.linear_num_key_heads % linear_head_parallel_size == 0, ( f"{self.linear_num_key_heads=} must be a multiple of " - f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." + f"{linear_head_parallel_size=} for {self.linear_cp_mode=}." ) - assert self.linear_num_value_heads % tp_cp_size == 0, ( + assert self.linear_num_value_heads % linear_head_parallel_size == 0, ( f"{self.linear_num_value_heads=} must be a multiple of " - f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." + f"{linear_head_parallel_size=} for {self.linear_cp_mode=}." ) - elif self.experimental_attention_variant == "dsa": + if self.experimental_attention_variant == "dsa": _validate_dsa_kernel_backend_dependencies(self.dsa_kernel_backend) if self.add_bias_linear: raise ValueError( @@ -1624,6 +1785,108 @@ def __post_init__(self): "dsa_indexer_skip_topk_offset must be non-negative, got " f"{self.dsa_indexer_skip_topk_offset}." ) + 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." + ) + 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 length is num_layers + mtp_num_layers (the GPT path uses exactly this, + # where mtp_num_layers == #MTP transformer layers). On HybridModel an MTP "depth" + # can expand to MULTIPLE hybrid layers, so mtp_num_layers (= depth count) undercounts + # the real MTP layers and csa_compress_ratios must be at least long enough to index + # every MTP attention layer (num_layers + layer_number - 1). Hence ">=", not "==". + expected_len = self.num_layers + mtp_layers + assert len(self.csa_compress_ratios) >= expected_len, ( + f"csa_compress_ratios length ({len(self.csa_compress_ratios)}) must be at least " + f"num_layers + mtp_num_layers ({self.num_layers} + {mtp_layers} = {expected_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 not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention." + self.hetereogenous_dist_checkpoint = True + + 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, ( + f"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'." + ) + + from cudnn import DSA + + if ( + self.context_parallel_size > 1 or self.dynamic_context_parallel + ) and uses_ratio4_indexer: + required_wrappers = [DSA.indexer_forward_wrapper] + if indexer_loss_enabled and not self.dsa_indexer_use_sparse_loss: + required_wrappers.extend( + [ + DSA.dense_indexer_score_recompute_wrapper, + DSA.dense_attn_score_recompute_wrapper, + DSA.dense_indexer_backward_wrapper, + ] + ) + missing_offsets = [ + wrapper.__name__ + for wrapper in required_wrappers + if "q_causal_offsets" not in inspect.signature(wrapper).parameters + ] + if missing_offsets: + raise ValueError( + "DSv4 CP with ratio-4 fused DSA requires cuDNN Frontend wrappers " + "with q_causal_offsets support; missing from: " + f"{', '.join(missing_offsets)}. Install a compatible cuDNN Frontend " + "build or set dsa_kernel_backend='none'." + ) + + if self.gdn_pre_gated_delta_rule_fusion and self.experimental_attention_variant == "kda": + raise NotImplementedError( + "gdn_pre_gated_delta_rule_fusion is not implemented for KDA yet." + ) + + 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'." + ) if self.fp8: # cannot support first last layer bf16 with delayed scaling @@ -1714,16 +1977,10 @@ def __post_init__(self): "to avoid costly dtype conversions during decode." ) - # Gated linear units (SwiGLU/GeGLU) are supported by the torch and vllm - # grouped-GEMM backends only. - if self.gated_linear_unit and self.inference_grouped_gemm_backend not in ( - "torch", - "vllm", - ): + if self.gated_linear_unit: 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}'." + "--transformer-impl='inference_optimized' does not yet support " + "gated linear units (SwiGLU/GeGLU)." ) if self.fp8 == "mxfp8": @@ -1762,24 +2019,6 @@ def __post_init__(self): "Set inference_grouped_gemm_backend to 'torch' for MXFP8." ) - if self.batch_invariant_mode: - if self.inference_grouped_gemm_backend not in ( - InferenceGroupedGemmBackend.TORCH, - InferenceGroupedGemmBackend.VLLM, - ): - raise ValueError( - "batch_invariant_mode requires inference_grouped_gemm_backend " - "'torch' or 'vllm'." - ) - 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.") @@ -1828,23 +2067,20 @@ def __post_init__(self): "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: + if not self.use_transformer_engine_op_fuser: 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." + "moe_single_grouped_weight requires " + "use_transformer_engine_op_fuser=True. The non-op-fuser TE GroupedLinear " + "path splits the grouped parameter into per-expert tensors and does not " + "support single-grouped-weight training." ) - 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": + if self.moe_flex_dispatcher_backend in ("deepepv2", "hybridep"): raise ValueError("Only one backend is supported for flex token dispatcher.") self.moe_flex_dispatcher_backend = "deepep" warnings.warn( @@ -1854,10 +2090,10 @@ def __post_init__(self): 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" + self.moe_enable_deepep or self.moe_flex_dispatcher_backend in ("deepep", "deepepv2") ): raise ValueError( - "Flex token dispatcher with deepep backend does not support " + "Flex token dispatcher with deepep/deepepv2 backend does not support " "moe_pad_expert_input_to_capacity" ) @@ -1867,31 +2103,6 @@ def __post_init__(self): "moe_flex_dispatcher_backend='ncclep' requires " "moe_token_dispatcher_type='flex'." ) - if self.moe_use_grouped_tensor and not self.use_transformer_engine_op_fuser: - raise ValueError( - "moe_use_grouped_tensor=True without use_transformer_engine_op_fuser is " - "not yet supported with the NCCL-EP dispatcher. Use the TE op-fuser path " - "or select the alltoall, DeepEP, or HybridEP dispatcher." - ) - - 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 @@ -1978,11 +2189,10 @@ def __post_init__(self): 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." + "use_transformer_engine_op_fuser to be enabled." ) if self.cpu_offloading and ( @@ -2054,10 +2264,9 @@ def __post_init__(self): "mlp", "moe", "shared_experts", - "gdn_norm_out", - "gdp_in_proj", - "gdp_qkv", "mhc", + "gdn", + "gdn_norm_out", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -2076,14 +2285,31 @@ def __post_init__(self): "multi_latent_attention." ) - if "gdn_norm_out" in self.recompute_modules and ( - not is_gated_delta_net_variant(self.experimental_attention_variant) + if ( + "gdn_norm_out" in self.recompute_modules + and not self.is_hybrid_model + 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'." + f"GDN-family layers, but got {self.experimental_attention_variant=}." ) + if ( + "gdn" in self.recompute_modules + and not self.is_hybrid_model + and not is_gated_delta_net_variant(self.experimental_attention_variant) + ): + raise ValueError( + "gdn in recompute_modules is only supported with GDN-family layers, but got " + f"{self.experimental_attention_variant=} and {self.is_hybrid_model=}." + ) + + if "gdn" in self.recompute_modules and "gdn_norm_out" in self.recompute_modules: + raise ValueError( + "'gdn' and 'gdn_norm_out' in recompute_modules cannot be used together. " + "'gdn' recomputes the full GDN-family layer, including gated norm." + ) if "core_attn" in self.recompute_modules: warnings.warn( "If you are using transformer_engine as the transformer implementation, " @@ -2130,8 +2356,10 @@ def __post_init__(self): # 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 not self.enable_hyper_connections: + raise ValueError( + "'mhc' in recompute_modules requires enable_hyper_connections=True." + ) if "mlp" in self.recompute_modules: raise ValueError( "'mhc' and 'mlp' in recompute_modules cannot be used together. " @@ -2146,92 +2374,94 @@ def __post_init__(self): "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." - ) + if self.fine_grained_activation_offloading and self.offload_modules: + # mHC checkpoints wrap input_layernorm (inside attn_norm offload context) + # and pre_mlp_layernorm (inside mlp_norm offload context). The unified + # recompute trigger runs before GroupCommitFunction.backward() initializes + # the backward chunk, so tensor_pop hits a None chunk for these modules. + # Other offload modules (qkv_linear, core_attn, attn_proj, expert_fc1, + # moe_act) live inside self_attention/MLP which are NOT wrapped by mHC + # checkpoints, so they are safe to use with mHC recompute -- except + # under the attention-only CUDA Graph split, which narrows the three + # attention-scope modules further down in __post_init__ (the split + # predicate needs normalized cuda_graph_modules, which are not + # available yet here). + _MHC_CONFLICTING_OFFLOAD_MODULES = {"attn_norm", "mlp_norm"} + conflicting = _MHC_CONFLICTING_OFFLOAD_MODULES & set(self.offload_modules) + if conflicting: + raise ValueError( + f"'mhc' in recompute_modules is incompatible with " + f"offload_modules {conflicting}. The mHC recompute replay starts " + f"before the offloading backward chunk is initialized for these " + f"modules, causing tensor_pop on a None chunk. Remove " + f"{conflicting} from offload_modules or remove 'mhc' from " + f"recompute_modules." + ) - if self.enable_mhc_connections and not ( + use_mhc_recompute = ( self.recompute_granularity == "selective" and "mhc" in self.recompute_modules + ) + # NOTE: the mHC-recompute/CUDA-graph capability gate lives further down in + # __post_init__, after cuda_graph_modules normalization and the deprecated + # enable_cuda_graph/external_cuda_graph migration, so that string module + # forms and legacy flags reach the same gate. + + # The attention-only Transformer Engine split keeps the mHC producer and + # its checkpoint registration eager per microbatch (only input-layernorm + # + self-attention are captured), so it composes with EP a2a overlap: + # the captured attention backward reads the fixed arena slot that the + # group-end recompute node repopulates by direct kernel output. + # Full-iteration capture is also valid because it records the explicit + # schedule barrier and its fixed-address recompute kernels as part of + # the iteration graph. Other per-layer graph forms remain rejected. + # Normalize locally: cuda_graph_modules is only normalized in bulk further + # down in __post_init__, so a config built directly (rather than through + # megatron/training/arguments.py) still holds string forms like ["attn"] + # here. Comparing them raw would reject a combination the later gate + # accepts, and the string form is part of this field's declared type. + normalized_graph_modules, _, _ = normalize_cuda_graph_modules(self.cuda_graph_modules) + is_te_attn_split = ( + self.cuda_graph_impl == "transformer_engine" + and not self.enable_cuda_graph + and not self.external_cuda_graph + and list(normalized_graph_modules or []) == [CudaGraphModule.attn] + and list(self.recompute_modules) == ["mhc"] + ) + is_full_iteration = ( + self.cuda_graph_impl == "full_iteration" + and not self.enable_cuda_graph + and not self.external_cuda_graph + ) + if ( + self.overlap_moe_expert_parallel_comm + and use_mhc_recompute + and not is_te_attn_split + and not is_full_iteration + and ( + self.cuda_graph_impl != "none" or self.enable_cuda_graph or self.external_cuda_graph + ) ): + raise ValueError( + "mHC recompute with overlap_moe_expert_parallel_comm requires CUDA graphs " + "to be disabled, except for the attention-only Transformer Engine split " + "(cuda_graph_impl='transformer_engine', cuda_graph_modules=[attn], " + "recompute_modules=[mhc]) or full-iteration capture " + "(cuda_graph_impl='full_iteration', hidden_dropout=0, " + "attention_dropout=0), because other graph scopes cannot represent the " + "explicit schedule-owned recompute barrier." + ) + + if self.enable_hyper_connections and not use_mhc_recompute: 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}." - ) + # Validation for use_fused_mhc + if self.use_fused_mhc: + if not self.enable_hyper_connections: + raise ValueError("use_fused_mhc requires enable_hyper_connections=True.") if self.fine_grained_activation_offloading: assert ( @@ -2247,7 +2477,6 @@ def __post_init__(self): "attn_norm", "mlp_norm", "qkv_linear", - "gdp_qkv", } invalid_modules = set(self.offload_modules) - allowed_modules assert not invalid_modules, ( @@ -2260,17 +2489,6 @@ def __post_init__(self): "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 @@ -2291,7 +2509,10 @@ def __post_init__(self): assert ( self.delta_offload_bytes_across_pp_ranks >= 0 ), "delta_offload_bytes_across_pp_ranks must be non-negative." - + if self.fine_grained_offloading_max_inflight_offloads is not None: + assert ( + self.fine_grained_offloading_max_inflight_offloads >= 0 + ), "fine_grained_offloading_max_inflight_offloads must be non-negative when set." 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.") @@ -2301,39 +2522,21 @@ def __post_init__(self): "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." - ) + assert not self.cpu_offloading, "moe_paged_stash cannot be enabled with cpu_offloading." + assert self.moe_expert_rank_capacity_factor is not None, ( + "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}" - ) + assert not moe_offload_conflict, ( + "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 @@ -2577,32 +2780,18 @@ def __post_init__(self): 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.activation_func_clamp_value is not None: + # swiglu + if self.activation_func == F.silu and self.gated_linear_unit: + if self.num_moe_experts is None: + raise ValueError( + "activation_func_clamp_value for SwiGLU is only supported with MoE." + ) + 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: @@ -2611,7 +2800,18 @@ def __post_init__(self): "It is experimental and may change in future versions." ) else: - if self.rotary_interleaved: + fused_mrope_available = False + # Triton fused mRoPE supports split-half RoPE only. Keep rotary_interleaved + # configs on the TE validation path so the TE >= 2.3 check still applies. + if self.mrope_section is not None and not self.rotary_interleaved: + try: + from megatron.core.fusions.fused_mrope import is_fused_mrope_available + + fused_mrope_available = is_fused_mrope_available() + except ImportError: + fused_mrope_available = False + + if self.rotary_interleaved and not fused_mrope_available: if not is_te_min_version("2.3.0"): raise ValueError( "rotary_interleaved does not work with apply_rope_fusion for " @@ -2623,9 +2823,14 @@ def __post_init__(self): fused_apply_rotary_pos_emb_thd, ) - if fused_apply_rotary_pos_emb is None and fused_apply_rotary_pos_emb_thd is None: + if ( + fused_apply_rotary_pos_emb is None + and fused_apply_rotary_pos_emb_thd is None + and not fused_mrope_available + ): raise ValueError( - "apply_rope_fusion is not available. Please install TE >= 1.4." + "apply_rope_fusion is not available. Please install TE >= 1.4 " + "or Triton for fused mRoPE." ) if self.fused_single_qkv_rope: @@ -2733,6 +2938,38 @@ def __post_init__(self): "'sqrtsoftplus', or unset --moe-router-enable-expert-bias." ) + if self.moe_n_hash_layers > 0: + assert ( + self.actual_vocab_size is not None + ), "actual_vocab_size must be set when moe_n_hash_layers > 0." + 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)." + ) + # The embedding is always in layout[0][0] (PP rank 0, VPP rank 0). + # All hash MoE layers must be in the same virtual pipeline stage. + embedding_stage = self.pipeline_model_parallel_layout.layout[0][0] + n_decoders_with_embedding = embedding_stage.count(LayerType.decoder) + assert self.moe_n_hash_layers <= n_decoders_with_embedding, ( + f"Currently, All hash MoE layers must be in the same virtual pipeline stage " + f"as the embedding. The embedding stage has " + f"{n_decoders_with_embedding} decoder layers, but " + f"moe_n_hash_layers={self.moe_n_hash_layers}." + ) + assert ( + not self.overlap_moe_expert_parallel_comm + ), "overlap_moe_expert_parallel_comm does not support moe_n_hash_layers > 0 for now." + log_single_rank( + logger, + logging.WARNING, + f"Hash MoE layer initialized with placeholder round-robin tid2eid. " + f"For real training, you MUST either (a) load tid2eid from a " + f"pre-trained DSv4 checkpoint, or (b) provide a frequency-aware " + f"initialization (e.g., Sinkhorn-balanced over token frequency). " + f"Round-robin will cause severe expert imbalance.", + ) + 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"): @@ -2904,6 +3141,118 @@ def _scope_to_str(s): self.cuda_graph_impl == "full_iteration" and self.cuda_graph_modules ), 'cuda_graph_modules must be empty when cuda_graph_impl="full_iteration".' + # mHC selective recompute couples with CUDA graphs only through the guarded + # attention-only Transformer Engine split. This gate must stay below the + # cuda_graph_modules normalization and the deprecated flag migration above: + # earlier placement would compare unnormalized string module forms and let + # enable_cuda_graph/external_cuda_graph bypass the gate entirely. + if use_mhc_recompute and self.cuda_graph_impl != "none": + if self.cuda_graph_impl == "local": + # Intentionally fail-closed even for inference-only local-graph + # configs that carry leftover training recompute args: mHC + # recompute is inert outside training, but silently accepting + # the combination would mask misconfigured training runs. + raise ValueError( + "mHC recompute is not supported with cuda_graph_impl='local': " + "eager mHC recompute and its per-microbatch checkpoint " + "registration need host execution between captured segments, " + "which the local per-layer implementation does not provide. Use " + "cuda_graph_impl='transformer_engine' with " + "cuda_graph_modules=['attn'], cuda_graph_impl='full_iteration' " + "with dropout disabled, or disable CUDA graphs." + ) + if self.cuda_graph_impl == "full_iteration": + # Full-iteration capture records the whole eager iteration — + # including mHC checkpoint registration, recompute kernels, and + # storage rebinding — into one graph, so replays re-execute the + # recompute at fixed addresses by construction (no partial-graph + # bridge involved). The one mechanical hazard is RNG-consuming + # ops inside a checkpointed region: the recompute-time RNG rewind + # cannot run under stream capture, so a captured recompute would + # replay a different dropout mask than its captured forward. + if self.hidden_dropout != 0.0 or self.attention_dropout != 0.0: + raise ValueError( + "mHC recompute with cuda_graph_impl='full_iteration' requires " + "hidden_dropout=0 and attention_dropout=0: RNG state cannot be " + "rewound inside CUDA graph capture, so a captured recompute " + "would replay a different dropout mask than its forward pass." + ) + elif list(self.cuda_graph_modules or []) != [CudaGraphModule.attn] or list( + self.recompute_modules + ) != ["mhc"]: + raise ValueError( + "mHC recompute with Transformer Engine CUDA Graphs currently supports " + "only the initial attention-only split: cuda_graph_modules=[attn] and " + "recompute_modules=[mhc]. The eager mHC producer must remain outside " + "the captured consumer." + ) + if ( + self.cuda_graph_impl == "transformer_engine" + and self.fine_grained_activation_offloading + ): + # HyperConnectionTransformerLayer._te_cuda_graph_capture replaces + # TransformerLayer's implementation rather than extending it, so it + # never plants the offload synchronization edges the full-layer + # capture plants -- backward_record() on the graph input and + # forward_record() after capture. _set_offload_modules plants those + # exactly for the attention-scope modules under an attn-scope graph, + # so without them the offload copies race the captured attention. + attn_scope_offload = {"qkv_linear", "core_attn", "attn_proj"} & set( + self.offload_modules or [] + ) + if attn_scope_offload: + raise ValueError( + f"mHC recompute with attention-only TE CUDA Graphs is incompatible " + f"with offload_modules {sorted(attn_scope_offload)}. The split " + f"capture path omits the offload stream synchronization the " + f"full-layer capture path performs, so the offload copies can race " + f"the captured attention. Remove {sorted(attn_scope_offload)} from " + f"offload_modules, or drop 'attn' from cuda_graph_modules." + ) + # The replay half of the same gap is fixed, not rejected: see + # _replay_mhc_attention_consumer. + if self.virtual_pipeline_model_parallel_size is not None: + # VPP is admitted only together with EP overlap. Interleaving used + # to diverge here (grad norms ~1e8 from the first iteration) on a + # caching-allocator use-after-free: mHC post-processing ran inside + # the communication-stream combine node, so the recompute subgraph + # was allocated on the compute stream and read from another, a + # window the allocator cannot track. The fix -- the post node owning + # a compute-stream schedule node -- lives in the overlap schedule, + # so the non-overlap VPP path has never carried it and stays + # unvalidated. StaticBufferLoader itself is VPP-safe, since only the + # pre_process chunk carries a data iterator. + if not self.overlap_moe_expert_parallel_comm: + raise ValueError( + "mHC recompute supports interleaved pipeline (VPP) " + "schedules only together with " + "overlap_moe_expert_parallel_comm: the non-overlap VPP " + "path is unvalidated." + ) + + cuda_graph_captures_attention = self.cuda_graph_impl == "full_iteration" or ( + self.cuda_graph_impl in ("local", "transformer_engine") + and (not self.cuda_graph_modules or CudaGraphModule.attn in self.cuda_graph_modules) + ) + + cp_layout_conversion_required = is_gated_delta_net_variant( + self.experimental_attention_variant + ) + if ( + (self.context_parallel_size > 1 or self.dynamic_context_parallel) + and self.sequence_packing_scheduler is not None + and cuda_graph_captures_attention + and cp_layout_conversion_required + ): + raise ValueError( + "THD context parallel layout conversion is required for this model " + "configuration, but it is not supported by CUDA graph capture that includes " + "attention " + f"(experimental_attention_variant={self.experimental_attention_variant!r}, " + f"cuda_graph_impl={self.cuda_graph_impl!r}, " + f"cuda_graph_modules={self.cuda_graph_modules!r})." + ) + if self.cuda_graph_impl != "none": if self.cpu_offloading and self.cuda_graph_impl != "full_iteration": @@ -2923,6 +3272,12 @@ def _scope_to_str(s): if CudaGraphModule.moe_preprocess not in self.cuda_graph_modules: self.cuda_graph_modules.append(CudaGraphModule.moe_preprocess) + if self.cuda_graph_impl != "transformer_engine": + assert not self.cuda_graph_dynamic_microbatches, ( + "cuda_graph_dynamic_microbatches is only supported with " + "cuda_graph_impl=transformer_engine." + ) + assert ( CudaGraphModule.moe not in self.cuda_graph_modules or CudaGraphModule.moe_router not in self.cuda_graph_modules @@ -3014,27 +3369,6 @@ def _scope_to_str(s): "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": @@ -3074,17 +3408,11 @@ def _scope_to_str(s): "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" + assert self.fine_grained_offloading_max_inflight_offloads is not None, ( + "fine_grained_offloading_max_inflight_offloads must be set 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( @@ -3129,15 +3457,60 @@ def _scope_to_str(s): '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' + if self.recompute_granularity == 'full': + # Full recompute runs per layer segment; recompute_method / + # recompute_num_layers control the decoder grouping and are validated by + # the shared checks above. MTP under 'block' and the FP8/FP4 block window + # intentionally differ from the non-overlap checkpoint path. See + # megatron/core/models/common/model_chunk_schedule_plan.py. + # The replay is hand-rolled rather than a checkpoint primitive, so it + # cannot shard the retained segment input across TP ranks. + assert not self.distribute_saved_activations, ( + 'overlap_moe_expert_parallel_comm full recompute does not support ' + 'distribute_saved_activations: the segment replay does not go through ' + 'a checkpoint primitive, so the retained input is not sharded.' + ) + # MTP is segmented under both methods here, so unlike + # MultiTokenPredictionLayer._checkpointed_forward there is no 'block' + # exemption. The 'uniform' restriction below is config parity with that + # function's own assert, not a constraint of this schedule - multi-layer + # decoder segments hand the mHC bridge over correctly. + if self.mtp_num_layers and self.recompute_method == 'uniform': + assert ( + self.recompute_num_layers == 1 + ), 'recompute_num_layers must be 1 for MTP recompute' + # With dropout the forward RNG stream is interleaved across microbatches + # in the 1F1B schedule and cannot yet be faithfully replayed. + assert self.attention_dropout == 0.0 and self.hidden_dropout == 0.0, ( + 'overlap_moe_expert_parallel_comm full recompute currently requires ' + 'attention_dropout==0 and hidden_dropout==0 (RNG replay across the ' + 'interleaved 1F1B schedule is not yet supported for nonzero dropout)' + ) + # Delayed-scaling FP8 accumulates amax into a persistent + # amax_history across microbatches. Both the initial no_grad forward + # and the backward-time recompute run under the delayed FP8 autocast, + # and this path does not enter TE's activation-recompute + # (recompute_phase) machinery that te_checkpoint uses to snapshot and + # restore FP8 metadata, so the recompute would update amax_history a + # second time and silently diverge. Current-scaling recipes + # (tensorwise / mxfp8 / blockwise) derive scales from the current + # tensor each pass, so the replayed forward re-derives identical + # scales (validated bitwise-identical to recompute-off); only delayed + # scaling is rejected here. + assert not (self.fp8 and self.fp8_recipe == Fp8Recipe.delayed), ( + 'overlap_moe_expert_parallel_comm full recompute is not yet supported ' + 'with delayed-scaling FP8 (fp8_recipe="delayed"): the backward-time ' + 'recompute would double-update the persistent amax_history. Use a ' + 'current-scaling recipe (tensorwise / mxfp8 / blockwise) instead.' + ) + else: + 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' @@ -3150,11 +3523,9 @@ def _scope_to_str(s): 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.' + assert ( + self.mtp_num_layers is None or self.mtp_num_layers == 1 + ), 'MTP layernum only supports 1 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. @@ -3162,6 +3533,14 @@ def _scope_to_str(s): self.moe_token_dispatcher_type == 'flex' and self.moe_flex_dispatcher_backend == 'ncclep' ): + if not self.moe_ncclep_static_shape: + warnings.warn( + 'overlap_moe_expert_parallel_comm with ncclep and ' + 'moe_ncclep_static_shape=False: get_permuted_hidden_states_by_experts ' + 'does a device-to-host sync that serializes the 1F1B overlap (correct, ' + 'but loses the overlap benefit). Set moe_ncclep_static_shape=True for ' + 'the overlapped path (needs the fused op on sm100+).' + ) assert not ( self.fine_grained_activation_offloading and 'expert_fc1' in (self.offload_modules or []) @@ -3174,7 +3553,8 @@ def _scope_to_str(s): if self.cuda_graph_impl != "none": if self.cuda_graph_impl == "transformer_engine": assert ( - CudaGraphModule.moe not in self.cuda_graph_modules + self.cuda_graph_impl == "transformer_engine" + and 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 ' @@ -3246,6 +3626,16 @@ def _scope_to_str(s): "2.3.0.dev0+39c0e70" ), "Must have at least TE version 2.3 or higher to use symmetric memory all reduce" + if self.rotary_base_per_layer is not None: + if self.mtp_num_layers is not None: + total_num_layers = self.num_layers + self.mtp_num_layers + else: + total_num_layers = self.num_layers + assert len(self.rotary_base_per_layer) == total_num_layers, ( + f"rotary_base_per_layer length ({len(self.rotary_base_per_layer)}) " + f"must equal num_layers ({total_num_layers})" + ) + 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): @@ -3263,6 +3653,25 @@ def _scope_to_str(s): f"the number of layers ({self.num_layers})" ) + if self.fallback_to_eager_attn: + assert self.transformer_impl == "transformer_engine", ( + f"fallback_to_eager_attn is only available with transformer_engine implementation," + f" but got {self.transformer_impl=}." + ) + + if self.fallback_to_eager_attn or self.transformer_impl == "local": + if self.context_parallel_size > 1 and self.cp_comm_type is not None: + all_cp_comm_types_are_all_gather = ( + all(item == "all_gather" for item in self.cp_comm_type) + if isinstance(self.cp_comm_type, list) + else self.cp_comm_type == "all_gather" + ) + if not all_cp_comm_types_are_all_gather: + raise ValueError( + f"fallback_to_eager_attn only supports all_gather communication type " + f"for context parallelism, but got {self.cp_comm_type=} instead." + ) + if self.transformer_impl == "inference_optimized": assert self.normalization == "RMSNorm" assert not self.layernorm_zero_centered_gamma @@ -3270,23 +3679,6 @@ def _scope_to_str(s): 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 " @@ -3302,91 +3694,43 @@ def _scope_to_str(s): "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. + ), "Batch invariant mode only supports FlashAttention" + + if self.cuda_graph_impl != "none" and ( + self.sequence_packing_scheduler is not None or self.dynamic_context_parallel + ): + if self.thd_max_packed_sequences is None: + raise ValueError("THD CUDA Graph requires --thd-max-packed-sequences to be set.") assert ( - self.context_parallel_size == 1 - ), "Batch invariant mode does not support context parallelism" + self.pad_packed_seq_alignment is not None + ), "THD CUDA Graph requires --pad-packed-seq-alignment to be set." 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, - ) + self.pad_packed_seq_alignment == "max" + or self.pad_packed_seq_alignment == self.max_seqlen_per_dp_cp_rank + ), ( + "THD CUDA Graph requires --pad-packed-seq-alignment='max' " + "or --pad-packed-seq-alignment equal to max_seqlen_per_dp_cp_rank " + f"({self.max_seqlen_per_dp_cp_rank}), got {self.pad_packed_seq_alignment}." + ) - if self.transformer_impl != "inference_optimized": - assert self.moe_token_dispatcher_type == "alltoall", ( - "Batch-invariant MoE training requires " - "moe_token_dispatcher_type='alltoall'." - ) - # DeepGEMM is used by the "deepgemm"/"triton" backends, and by - # the torch inference grouped-GEMM path under any backend. The - # "te_native" backend with the vLLM inference backend (or the - # training path, where TE grouped GEMM stays native) does not - # need it. - needs_deepgemm = self.batch_invariant_backend in ("deepgemm", "triton") or ( - self.transformer_impl == "inference_optimized" - and self.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.TORCH - ) - 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]`." - ) - assert not ( - self.fp8 or self.fp4 - ), "Batch-invariant MoE is bf16-only. Disable fp8/fp4 to use it." - 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." - ) + # 'extend_last' THD tail padding with context parallelism requires the + # global metadata to be extended before CP slicing, which only the + # sequence-packing scheduler path performs. Restrict this combination. + if ( + self.sequence_packing_scheduler is None + and self.pad_packed_seq_alignment is not None + and (self.thd_tail_padding_policy or 'append_dummy_seq') == 'extend_last' + and (self.context_parallel_size > 1 or self.dynamic_context_parallel) + ): + raise ValueError( + "thd_tail_padding_policy='extend_last' with context parallelism is only " + "supported together with a sequence_packing_scheduler. Either enable a " + "sequence_packing_scheduler, or use thd_tail_padding_policy='append_dummy_seq'." + ) - # 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: @@ -3402,19 +3746,28 @@ def _scope_to_str(s): 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). + # Needed for passing variable sequences between pp stages. + self.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}'" + assert self.moe_token_dispatcher_type in ("alltoall", "flex"), ( + f"sequence_packing only supports moe_token_dispatcher_type in " + f"('alltoall', 'flex'), got '{self.moe_token_dispatcher_type}'" + ) + + supported_schedulers = ['dp_balanced', 'default_dynamic_cp'] + if ( + self.sequence_packing_scheduler is not None + and self.sequence_packing_scheduler not in supported_schedulers + ): + raise ValueError( + f"Unsupported scheduler: {self.sequence_packing_scheduler}. " + f"Available schedulers: {supported_schedulers}" ) @dataclass +@experimental_api class MLATransformerConfig(TransformerConfig): """Configuration object for megatron-core Multi-Latent Attention (MLA) transformers. @@ -3429,10 +3782,12 @@ class MLATransformerConfig(TransformerConfig): """Rank of Query tensor's low rank representation.""" kv_lora_rank: int = 512 - """Rank of Key and Value tensors' low rank representation.""" + """Rank of Key and Value tensors' low rank representation. + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_head_dim: int = 128 - """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim""" + """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.""" @@ -3470,6 +3825,12 @@ class MLATransformerConfig(TransformerConfig): mscale_all_dim: float = 0.0 """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" + o_groups: int = 8 + """Number of groups for grouped low-rank output projection (wo_a).""" + + o_lora_rank: int = 1024 + """Low-rank dimension per group for grouped output (wo_a). Used when o_groups > 0.""" + 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 @@ -3482,11 +3843,39 @@ class MLATransformerConfig(TransformerConfig): def __post_init__(self): super().__post_init__() - if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": + 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.attention_output_gate: - raise NotImplementedError("Output gate is not supported for MLA yet.") + if self.attention_output_gate and self.mla_down_proj_fusion: + # Fused MLA hides the post-input-LayerNorm activation inside the fused + # LayerNorm+linear module. Gated MLA must consume that activation as the + # gate input; using raw hidden_states would silently change the model. + # Keep this combination fail-fast until the fused API exposes the + # normalized activation. + raise ValueError( + "MLA output gating does not support fused down projections; " + "disable mla_down_proj_fusion to use the unfused path." + ) + + # 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." + log_single_rank( + logger, + logging.WARNING, + f"DSv4 hybrid mode is enabled, deriving qk_head_dim and kv_lora_rank from " + f"v_head_dim and qk_pos_emb_head_dim", + ) + derived = self.v_head_dim - self.qk_pos_emb_head_dim + self.qk_head_dim = derived + self.kv_lora_rank = derived if self.cache_mla_latents: assert ( diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py index 124083ead83..db927573166 100644 --- a/megatron/training/argument_utils.py +++ b/megatron/training/argument_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import ast import builtins @@ -17,7 +17,7 @@ import torch.nn.functional as F from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.spec_utils import import_module +from megatron.core.transformer.spec_utils import ModuleSpec, import_module from megatron.training.config import ( CheckpointConfig, DistributedInitConfig, @@ -38,17 +38,20 @@ # TODO: support arg renames + class TypeInferenceError(Exception): """Custom exception type to be conditionally handled by ArgumentGroupFactory.""" + pass + class ArgumentGroupFactory: """Utility that adds an argument group to an ArgumentParser based on the attributes of a dataclass. This utility uses dataclass metadata including type annotations and docstrings to automatically infer the type, default, and other argparse keyword arguments. - You can override or supplement the automatically inferred argparse kwargs for any + You can override or supplement the automatically inferred argparse kwargs for any dataclass field by providing an "argparse_meta" key in the field's metadata dict. The value should be a dict of kwargs that will be passed to ArgumentParser.add_argument(). These metadata kwargs take precedence over the automatically inferred values. @@ -77,13 +80,13 @@ class YourConfig: that require some customized or additional handling. Args: - src_cfg_class: The source dataclass type (not instance) whose fields will be - converted into command-line arguments. Each field's type annotation determines - the argument type, default values become argument defaults, and field-level + src_cfg_class: The source dataclass type (not instance) whose fields will be + converted into command-line arguments. Each field's type annotation determines + the argument type, default values become argument defaults, and field-level docstrings are extracted to populate argument help text. - exclude: Optional list of attribute names from `src_cfg_class` to exclude from + exclude: Optional list of attribute names from `src_cfg_class` to exclude from argument generation. Useful for omitting internal fields, computed properties, - or attributes that should be configured through other means. If None, all + or attributes that should be configured through other means. If None, all dataclass fields will be converted to command-line arguments. Default: None. """ @@ -97,7 +100,7 @@ def _format_arg_name(self, config_attr_name: str, prefix: Optional[str] = None) Args: config_attr_name: dataclass attribute name - prefix: prefix string to add to the dataclass attribute name. e.g. 'no' for bool + prefix: prefix string to add to the dataclass attribute name. e.g. 'no' for bool settings that are default True. A hyphen is added after the prefix. Default: None """ arg_name = config_attr_name @@ -112,6 +115,7 @@ def _get_enum_kwargs(self, config_type: enum.EnumMeta) -> dict[str, Any]: With these settings, the user must provide a valid enum value, e.g. 'flash', for `AttnBackend.flash`. """ + def enum_type_handler(cli_arg): return config_type[cli_arg] @@ -135,7 +139,9 @@ def _extract_type(self, config_type: type) -> dict[str, Any]: if origin in [types.UnionType, typing.Union]: # Handle Optional and Union - if type_tuple[1] == type(None): # Optional type. First element is value inside Optional[] + if type_tuple[1] == type( + None + ): # Optional type. First element is value inside Optional[] return self._extract_type(type_tuple[0]) else: raise TypeInferenceError(f"Unions not supported by argparse: {config_type}") @@ -146,17 +152,20 @@ def _extract_type(self, config_type: type) -> dict[str, Any]: kwargs["nargs"] = "+" return kwargs else: - raise TypeInferenceError(f"Multi-type lists not supported by argparse: {config_type}") + raise TypeInferenceError( + f"Multi-type lists not supported by argparse: {config_type}" + ) elif origin is typing.Literal: choices_types = [type(choice) for choice in type_tuple] - assert all([t == choices_types[0] for t in choices_types]), "Type of each choice in a Literal type should all be the same." + assert all( + [t == choices_types[0] for t in choices_types] + ), "Type of each choice in a Literal type should all be the same." kwargs = {"type": choices_types[0], "choices": type_tuple} return kwargs else: raise TypeInferenceError(f"Unsupported type: {config_type}") - def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: """Assemble kwargs for add_argument(). @@ -166,7 +175,9 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: argparse_kwargs = {} argparse_kwargs["arg_names"] = [self._format_arg_name(attribute.name)] argparse_kwargs["dest"] = attribute.name - argparse_kwargs["help"] = self.field_docstrings[attribute.name] if attribute.name in self.field_docstrings else "" + argparse_kwargs["help"] = ( + self.field_docstrings[attribute.name] if attribute.name in self.field_docstrings else "" + ) # dataclasses specifies that both should not be set if isinstance(attribute.default, type(dataclasses.MISSING)): @@ -180,7 +191,6 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: # save metadata here, but update at the end so the metadata has highest precedence attr_argparse_meta = attribute.metadata["argparse_meta"] - # if we cannot infer the argparse type, all of this logic may fail. we try to defer # to the developer-specified metadata if present try: @@ -188,12 +198,17 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: # use store_true or store_false action for enable/disable flags, which doesn't accept a 'type' if argparse_kwargs["type"] == bool: - argparse_kwargs["action"] = "store_true" if attribute.default == False else "store_false" + argparse_kwargs["action"] = ( + "store_true" if attribute.default == False else "store_false" + ) argparse_kwargs.pop("type") # add '--no-*' and '--disable-*' prefix if this is a store_false argument if argparse_kwargs["action"] == "store_false": - argparse_kwargs["arg_names"] = [self._format_arg_name(attribute.name, prefix="no"), self._format_arg_name(attribute.name, prefix="disable")] + argparse_kwargs["arg_names"] = [ + self._format_arg_name(attribute.name, prefix="no"), + self._format_arg_name(attribute.name, prefix="disable"), + ] except TypeInferenceError as e: if attr_argparse_meta is not None: print( @@ -205,7 +220,7 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: else: raise e - # metadata provided by field takes precedence + # metadata provided by field takes precedence if attr_argparse_meta is not None: argparse_kwargs.update(attr_argparse_meta) @@ -255,8 +270,12 @@ def _get_field_docstrings(self, src_cfg_class: type) -> dict[str, str]: if a_cond and b_cond: # These should be guaranteed by typechecks above, but assert just in case - assert isinstance(a.target.id, str), "Dataclass attribute not in the expected format. Name is not a string." - assert isinstance(b.value.value, str), "Dataclass attribute docstring is not a string." + assert isinstance( + a.target.id, str + ), "Dataclass attribute not in the expected format. Name is not a string." + assert isinstance( + b.value.value, str + ), "Dataclass attribute docstring is not a string." # Formatting docstring = inspect.cleandoc(b.value.value) @@ -274,6 +293,100 @@ def _get_field_docstrings(self, src_cfg_class: type) -> dict[str, str]: return field_docstrings +def _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, pattern): + """Normalize DSv4 hybrid compression ratios for args and TransformerConfig.""" + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + + variant = kw_args.get( + 'experimental_attention_variant', getattr(args, 'experimental_attention_variant', None) + ) + if variant != 'dsv4_hybrid': + return + + fixed_ratio_map = {Symbols.WINDOW: 0, Symbols.CSA: 4, Symbols.HCA: 128} + ratio_symbols = set(fixed_ratio_map) + sections = pattern.split(Symbols.MTP_SEPARATOR) + layers = ''.join(section.replace(Symbols.PIPE, '') for section in sections) + attn_symbols = [symbol for symbol in layers if symbol in ratio_symbols] + compact_len = len(attn_symbols) + full_len = len(layers) + + def padded_to_compact_and_full(provided): + compact = [] + full = [] + compact_iter = iter(provided) + for symbol in layers: + if symbol in ratio_symbols: + ratio = next(compact_iter) + expected = fixed_ratio_map[symbol] + assert ratio == expected, ( + f"csa_compress_ratios has ratio {ratio} for hybrid symbol " + f"'{symbol}', expected {expected}." + ) + compact.append(ratio) + full.append(ratio) + else: + full.append(0) + return compact, full + + # args keeps the compact CLI form; TransformerConfig gets the full form for per-layer indexing. + if getattr(args, 'csa_compress_ratios', None) is None: + compact_ratios = [fixed_ratio_map[symbol] for symbol in attn_symbols] + args.csa_compress_ratios, kw_args['csa_compress_ratios'] = padded_to_compact_and_full( + compact_ratios + ) + else: + provided = list(args.csa_compress_ratios) + if len(provided) == compact_len: + args.csa_compress_ratios, kw_args['csa_compress_ratios'] = padded_to_compact_and_full( + provided + ) + elif len(provided) == full_len: + compact = [] + for ratio, symbol in zip(provided, layers): + if symbol in ratio_symbols: + expected = fixed_ratio_map[symbol] + assert ratio == expected, ( + f"csa_compress_ratios has ratio {ratio} for hybrid symbol " + f"'{symbol}', expected {expected}." + ) + compact.append(ratio) + else: + assert ratio == 0, ( + f"csa_compress_ratios should not pad non-attention hybrid symbol " + f"'{symbol}' with non-zero ratio {ratio}." + ) + args.csa_compress_ratios = compact + kw_args['csa_compress_ratios'] = provided + else: + raise AssertionError( + f"csa_compress_ratios length ({len(provided)}) must equal either the " + f"number of W/C/H attention symbols ({compact_len}) or the legacy " + f"number of all layers in the hybrid pattern ({full_len}) for pattern " + f"'{pattern}'." + ) + + +def _resolve_dsa_kernel_backend_cli_default(args, kw_args): + """Resolve an omitted --dsa-kernel-backend flag to its effective backend. + + The CLI flag carries no static default so an omitted flag is distinguishable + from an explicit "none". DSv4 hybrid launches historically ran fused kernels + by default (the deprecated --no-dsa-kernel-fusion flag defaulted to True), so + an omitted flag resolves to "cudnn" for that variant unless the deprecated + switch was passed; every other configuration resolves to "none". Runs before + TransformerConfig.__post_init__ folds the deprecated switch into + dsa_kernel_backend, so explicit legacy values still win. + """ + if 'dsa_kernel_backend' not in kw_args or kw_args['dsa_kernel_backend'] is not None: + return + variant = kw_args.get('experimental_attention_variant') or kw_args.get('linear_attention_type') + if variant == 'dsv4_hybrid' and getattr(args, 'apply_dsa_kernel_fusion', None) is None: + kw_args['dsa_kernel_backend'] = 'cudnn' + else: + kw_args['dsa_kernel_backend'] = 'none' + + def core_transformer_config_from_args(args, config_class=None): from megatron.core.activations import squared_relu from megatron.core.fusions.fused_bias_geglu import quick_gelu @@ -293,7 +406,9 @@ def core_transformer_config_from_args(args, config_class=None): config_class = MLATransformerConfig if args.heterogeneous_layers_config_path is not None: - assert not args.multi_latent_attention, "Multi latent attention with heterogeneous layers is not supported." + assert ( + not args.multi_latent_attention + ), "Multi latent attention with heterogeneous layers is not supported." config_class = HeterogeneousTransformerConfig # Translate args to core transformer configuration @@ -306,9 +421,10 @@ def core_transformer_config_from_args(args, config_class=None): kw_args['pipeline_dtype'] = args.params_dtype kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm kw_args['num_moe_experts'] = args.num_experts + kw_args['actual_vocab_size'] = args.padded_vocab_size kw_args['rotary_interleaved'] = args.rotary_interleaved - kw_args['num_layers_in_first_pipeline_stage']= args.decoder_first_pipeline_num_layers - kw_args['num_layers_in_last_pipeline_stage']= args.decoder_last_pipeline_num_layers + kw_args['num_layers_in_first_pipeline_stage'] = args.decoder_first_pipeline_num_layers + kw_args['num_layers_in_last_pipeline_stage'] = args.decoder_last_pipeline_num_layers kw_args['fp8_param'] = args.fp8_param_gather kw_args['fp4_param'] = args.fp4_param_gather if args.swiglu: @@ -336,17 +452,31 @@ def core_transformer_config_from_args(args, config_class=None): # Pop 'rope_type' to let the config class use the default value. kw_args.pop('rope_type', None) else: - assert (args.multi_latent_attention or args.rope_type == 'rope'), ( - f'Common attention only support rope_type="rope", but got {args.rope_type}.' - ) + assert ( + args.multi_latent_attention or args.rope_type == 'rope' + ), f'Common attention only support rope_type="rope", but got {args.rope_type}.' if len(args.cp_comm_type) == 1: kw_args['cp_comm_type'] = args.cp_comm_type[0] if args.hybrid_layer_pattern is not None: kw_args['is_hybrid_model'] = True from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols - if Symbols.DS_ATTENTION in args.hybrid_layer_pattern: - kw_args['experimental_attention_variant'] = 'dsa' + + pattern = args.hybrid_layer_pattern + has_kda = Symbols.KDA in pattern + has_dsv4_csa = any(s in pattern for s in (Symbols.WINDOW, Symbols.CSA, Symbols.HCA)) + has_dsa = Symbols.DS_ATTENTION in pattern + if getattr(args, 'experimental_attention_variant', None) is None: + if has_dsv4_csa: + kw_args['experimental_attention_variant'] = 'dsv4_hybrid' + elif has_dsa: + kw_args['experimental_attention_variant'] = 'dsa' + elif has_kda: + kw_args['experimental_attention_variant'] = 'kda' + + _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, pattern) + + _resolve_dsa_kernel_backend_cli_default(args, kw_args) kw_args['inference_sampling_seed'] = args.seed @@ -430,16 +560,8 @@ def _default_config_from_args(cls: type, args: Namespace, return_instance: bool return kwargs -def gpt_config_from_args( - args: Namespace, config: TransformerConfig | None = None, model_config_cls: type = GPTModelConfig -) -> Any: - """Create a GPTModelConfig (or a compatible subclass) from the `args` Namespace. - - `model_config_cls` lets callers reuse this same arg-derivation logic for - subclasses that only override metadata (e.g. `builder`) and add no new fields, - such as `ModelOptModelConfig`. - """ - assert issubclass(model_config_cls, GPTModelConfig) +def gpt_config_from_args(args: Namespace, config: TransformerConfig | None = None) -> Any: + """Create a GPTModelConfig from the appropriate values in the `args` Namespace.""" kwargs = {} if config is None: @@ -457,7 +579,6 @@ def gpt_config_from_args( kwargs["transformer_layer_spec"] = import_module(args.spec) kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy - kwargs["logit_dtype"] = getattr(args, "logit_dtype", None) kwargs["position_embedding_type"] = args.position_embedding_type kwargs["rotary_percent"] = args.rotary_percent kwargs["rotary_base"] = args.rotary_base @@ -475,23 +596,17 @@ def gpt_config_from_args( kwargs["vocab_size"] = args.padded_vocab_size kwargs["should_pad_vocab"] = False else: - assert args.vocab_size is not None, "Either --padded-vocab-size or --vocab-size must be specified." + assert ( + args.vocab_size is not None + ), "Either --padded-vocab-size or --vocab-size must be specified." kwargs["vocab_size"] = args.vocab_size kwargs["should_pad_vocab"] = True - return model_config_cls(**kwargs) + return GPTModelConfig(**kwargs) -def hybrid_config_from_args( - args: Namespace, config: TransformerConfig | None = None, model_config_cls: type = HybridModelConfig -) -> Any: - """Create a HybridModelConfig (or a compatible subclass) from the `args` Namespace. - - `model_config_cls` lets callers reuse this same arg-derivation logic for - subclasses that only override metadata (e.g. `builder`) and add no new fields, - such as `ModelOptHybridModelConfig`. - """ - assert issubclass(model_config_cls, HybridModelConfig) +def hybrid_config_from_args(args: Namespace, config: TransformerConfig | None = None) -> Any: + """Create a HybridModelConfig from the appropriate values in the `args` Namespace.""" kwargs = {} if config is None: @@ -505,10 +620,15 @@ def hybrid_config_from_args( not transformer_cfg.inference_fuse_tp_communication ), "inference_fuse_tp_communication is not supported for HybridModel" elif args.spec is not None: - kwargs["hybrid_stack_spec"] = import_module(args.spec) + hybrid_stack_spec = import_module(args.spec) + # ModuleSpec implements __call__ to build the described module, so a + # generic callable check would instantiate static specs here before a + # ProcessGroupCollection exists. Only resolve callable spec factories. + if callable(hybrid_stack_spec) and not isinstance(hybrid_stack_spec, ModuleSpec): + hybrid_stack_spec = hybrid_stack_spec(transformer_cfg) + kwargs["hybrid_stack_spec"] = hybrid_stack_spec kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy - kwargs["logit_dtype"] = getattr(args, "logit_dtype", None) kwargs["hybrid_layer_pattern"] = args.hybrid_layer_pattern kwargs["position_embedding_type"] = args.position_embedding_type kwargs["rotary_percent"] = args.rotary_percent @@ -526,11 +646,13 @@ def hybrid_config_from_args( kwargs["vocab_size"] = args.padded_vocab_size kwargs["should_pad_vocab"] = False else: - assert args.vocab_size is not None, "Either --padded-vocab-size or --vocab-size must be specified." + assert ( + args.vocab_size is not None + ), "Either --padded-vocab-size or --vocab-size must be specified." kwargs["vocab_size"] = args.vocab_size kwargs["should_pad_vocab"] = True - return model_config_cls(**kwargs) + return HybridModelConfig(**kwargs) def pretrain_cfg_container_from_args(args: Namespace, model_cfg=None) -> PretrainConfigContainer: @@ -574,7 +696,6 @@ def pretrain_cfg_container_from_args(args: Namespace, model_cfg=None) -> Pretrai checkpoint=CheckpointConfig(**ckpt_kwargs), profiling=ProfilingConfig(**prof_kwargs), tokenizer=_default_config_from_args(TokenizerConfig, args), - rerun_state_machine=RerunStateMachineConfig(**rerunsm_kwargs), straggler=_default_config_from_args(StragglerDetectionConfig, args), ) @@ -595,9 +716,7 @@ def inference_cfg_from_args(args: Namespace) -> InferenceSetupConfig: return _default_config_from_args(InferenceSetupConfig, args) -def inference_cfg_container_from_args( - args: Namespace, model_cfg=None -) -> InferenceConfigContainer: +def inference_cfg_container_from_args(args: Namespace, model_cfg=None) -> InferenceConfigContainer: """Build an InferenceConfigContainer from the argparse arguments. This mirrors ``pretrain_cfg_container_from_args`` but assembles only the configs that diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index dc093508775..2d08bcf91dc 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Megatron arguments.""" @@ -11,10 +11,20 @@ from pathlib import Path import torch +import torch.nn.functional as F +from packaging.version import Version as PkgVersion +from megatron.core.activations import squared_relu +from megatron.core.dist_checkpointing.validation import StrictHandling +from megatron.core.fusions.fused_bias_geglu import quick_gelu +from megatron.core.model_parallel_config import _parse_pad_packed_seq_alignment from megatron.core.msc_utils import MultiStorageClientFeature +from megatron.core.quantization.utils import ( + kitchen_quantization_recipe_config, + load_quantization_recipe, +) from megatron.core.rerun_state_machine import RerunStateMachine -from megatron.core.transformer import TransformerConfig +from megatron.core.transformer import MLATransformerConfig, TransformerConfig from megatron.core.transformer.cuda_graph_config import ( ALLOWED_INFERENCE_SCOPES, get_deprecated_cuda_graph_modules_migration, @@ -23,6 +33,10 @@ validate_deprecated_cuda_graph_modules_migration_inputs, ) from megatron.core.transformer.enums import AttnBackend, CudaGraphModule, InferenceCudaGraphScope +from megatron.core.transformer.heterogeneous.heterogeneous_config import ( + HeterogeneousTransformerConfig, + MLPConfig, +) from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.utils import ( get_torch_version, @@ -30,9 +44,10 @@ is_te_min_version, is_torch_min_version, ) -from megatron.training.argument_utils import ( # noqa: F401 # pylint: disable=unused-import +from megatron.training.argument_utils import ( ArgumentGroupFactory, - core_transformer_config_from_args, + _normalize_dsv4_hybrid_csa_compress_ratios, + _resolve_dsa_kernel_backend_cli_default, ) from megatron.training.global_vars import set_global_variables from megatron.training.utils import ( @@ -44,7 +59,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): - """"Add Megatron-LM arguments to the given parser.""" + """ "Add Megatron-LM arguments to the given parser.""" # Standard arguments. parser = _add_network_size_args(parser) @@ -67,7 +82,6 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): parser = _add_experimental_attention_variant_args(parser) parser = _add_heterogeneous_args(parser) parser = _add_logging_args(parser) - parser = _add_logits_distillation_args(parser) parser = _add_straggler_detector_args(parser) parser = _add_workload_inspector_server_args(parser) parser = _add_inference_args(parser) @@ -80,18 +94,23 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): parser = _add_msc_args(parser) parser = _add_kitchen_quantization_arguments(parser) parser = _add_sft_args(parser) + parser = _add_varlen_dataset_args(parser) + parser = _add_logits_distillation_args(parser) parser = _add_fault_injector_args(parser) return parser + def parse_and_validate_args(extra_args_provider=None, ignore_unknown_args=False, args_defaults={}): args = parse_args(extra_args_provider, ignore_unknown_args) if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False): from megatron.training.checkpointing import load_args_from_checkpoint - assert args.load is not None or args.pretrained_checkpoint is not None, "--use-checkpoint-args requires --load or --pretrained-checkpoint argument" + assert ( + args.load is not None or args.pretrained_checkpoint is not None + ), "--use-checkpoint-args requires --load or --pretrained-checkpoint argument" assert args.non_persistent_ckpt_type != "local", ( "--use-checkpoint-args is not supported with --non_persistent_ckpt_type=local. " "Two-stage checkpoint loading is not implemented, and all arguments must be defined " @@ -116,8 +135,7 @@ def parse_and_validate_args(extra_args_provider=None, ignore_unknown_args=False, def parse_args(extra_args_provider=None, ignore_unknown_args=False): """Parse all arguments.""" - parser = argparse.ArgumentParser(description='Megatron-LM Arguments', - allow_abbrev=False) + parser = argparse.ArgumentParser(description='Megatron-LM Arguments', allow_abbrev=False) parser = add_megatron_arguments(parser) @@ -139,7 +157,6 @@ def parse_args(extra_args_provider=None, ignore_unknown_args=False): args = load_yaml(args.yaml_cfg) - # Args from environment args.rank = int(os.getenv('RANK', '0')) args.world_size = int(os.getenv("WORLD_SIZE", '1')) @@ -198,17 +215,22 @@ def validate_model_config_args_from_heterogeneous_config(args): args.heterogeneous_layers_config_path ).read_text() - hf_config_dict = types.SimpleNamespace(**json.loads(args.heterogeneous_layers_config_encoded_json)) - - assert hf_config_dict.hidden_act == "silu", ( - f"hidden_act in heterogeneous config is {hf_config_dict.hidden_act}, should be silu" + hf_config_dict = types.SimpleNamespace( + **json.loads(args.heterogeneous_layers_config_encoded_json) ) + assert ( + hf_config_dict.hidden_act == "silu" + ), f"hidden_act in heterogeneous config is {hf_config_dict.hidden_act}, should be silu" + n_kv_heads_in_group = [ - config["attention"]["n_heads_in_group"] for config in hf_config_dict.block_configs + config["attention"]["n_heads_in_group"] + for config in hf_config_dict.block_configs if config["attention"]["n_heads_in_group"] is not None ] - assert all(num == n_kv_heads_in_group[0] for num in n_kv_heads_in_group), "num query head must be consistent across all layers" + assert all( + num == n_kv_heads_in_group[0] for num in n_kv_heads_in_group + ), "num query head must be consistent across all layers" args_to_validate = { "swiglu": True, @@ -234,16 +256,17 @@ def validate_model_config_args_from_heterogeneous_config(args): incompatible_args[key] = (provided_value, value) if incompatible_args: - incompatible_args_str = ', '.join([ - f"{k}: {provided_value} (provided) != {value} (expected)" - for k, (provided_value, value) in incompatible_args.items() - ]) - raise ValueError( - f"Arguments differ from heterogeneous config: {incompatible_args_str}" + incompatible_args_str = ', '.join( + [ + f"{k}: {provided_value} (provided) != {value} (expected)" + for k, (provided_value, value) in incompatible_args.items() + ] ) + raise ValueError(f"Arguments differ from heterogeneous config: {incompatible_args_str}") + def _eval_pattern(pattern): - """ Validate and evaluate a string containing a Python list expression """ + """Validate and evaluate a string containing a Python list expression""" assert isinstance(pattern, str) # validate input, only allow comma, digits, [, ], (, ), +, and * @@ -266,12 +289,12 @@ def _normalize_cuda_graph_modules_args(args): args.cuda_graph_modules ) validate_deprecated_cuda_graph_modules_migration_inputs( - deprecated_scopes, - args.cuda_graph_impl, - args.inference_cuda_graph_scope, + deprecated_scopes, args.cuda_graph_impl, args.inference_cuda_graph_scope ) if used_full_scope: - warn_rank_0('full scope is deprecated. Use empty cuda_graph_modules to capture the whole layer.') + warn_rank_0( + '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( @@ -302,7 +325,7 @@ def _normalize_inference_cuda_graph_scope_arg(args): def no_rope_freq_type(x): - """ Controls which layers to skip performing Rotary Position Embedding. + """Controls which layers to skip performing Rotary Position Embedding. - An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. - A string "N": Same as above, but provided as a string - A string containing a Python list expression that defines a custom pattern, e.g.: @@ -324,6 +347,22 @@ def no_rope_freq_type(x): # it's a single int but in str return int(x) + +def compress_ratios_type(x): + """Per-layer compress ratios for compressed sparse attention. + + Accepts a string containing a Python list expression, e.g.: + "[0,0,4,128,4,128]" + "([0]+[4,128]*2)*3" + The result must be a list of integers. Each value represents the + compression ratio for the corresponding transformer layer. + """ + if isinstance(x, list): + return x + assert isinstance(x, str) + return _eval_pattern(x) + + def moe_freq_type(x): """Frequency between MoE layers and Dense layers. @@ -349,6 +388,7 @@ def moe_freq_type(x): # it's a single int but in str return int(x) + def la_freq_type(x): """Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) layers. @@ -374,6 +414,7 @@ def la_freq_type(x): # it's a single int but in str return int(x) + def tuple_type(x): """ Convert a string to a tuple of integers. @@ -396,8 +437,11 @@ def validate_args(args, defaults={}): args.exit_on_missing_checkpoint = True # Temporary - assert args.non_persistent_ckpt_type in ['global', 'local', None], \ - 'Currently only global and local checkpoints are supported' + assert args.non_persistent_ckpt_type in [ + 'global', + 'local', + None, + ], 'Currently only global and local checkpoints are supported' if args.non_persistent_ckpt_type == 'local': try: from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import ( @@ -411,31 +455,21 @@ def validate_args(args, defaults={}): update_use_dist_ckpt(args) - # GTP_remat counts toward total_model_size (an independent weight-shard axis), so the - # args.data_parallel_size below is the replicate degree (matches - # parallel_state). gtp_weight_remat_size is derived from --tensor-parallel-num-weight-shards. - from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards - (args.tensor_parallel_num_weight_shards, args.gtp_weight_remat_size) = ( - resolve_tensor_parallel_weight_shards( - args.tensor_model_parallel_size, - args.tensor_parallel_num_weight_shards, - getattr(args, "gtp_weight_remat_size", 1), - ) - ) total_model_size = ( args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size - * args.gtp_weight_remat_size ) # Total model size. - assert args.world_size % total_model_size == 0, ( - f"world size ({args.world_size}) is not divisible by total_model_size ({total_model_size=})" - ) + assert ( + args.world_size % total_model_size == 0 + ), f"world size ({args.world_size}) is not divisible by total_model_size ({total_model_size=})" if args.attention_backend == AttnBackend.local: - assert args.spec[0] == 'local' , '--attention-backend local is only supported with --spec local' + assert ( + args.spec[0] == 'local' + ), '--attention-backend local is only supported with --spec local' # Pipeline model parallel size. args.transformer_pipeline_model_parallel_size = args.pipeline_model_parallel_size @@ -444,13 +478,10 @@ def validate_args(args, defaults={}): args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size - * args.gtp_weight_remat_size ) args.data_parallel_size = args.world_size // total_model_size if args.perform_rl_step: - assert args.refit_method != 'nccl_m2n', 'nccl_m2n is unsupported by the built-in RL loop' - # ---------------------------------------------------------------- # CUDA graphs # @@ -481,13 +512,13 @@ def validate_args(args, defaults={}): # ---------------------------------------------------------------- # Persisting CGs only makes sense if we build any CGs. - assert not args.rl_persist_cuda_graphs or args.cuda_graph_impl != "none", ( - "--rl-persist-cuda-graphs is set but no CUDA graphs are being built." - ) + assert ( + not args.rl_persist_cuda_graphs or args.cuda_graph_impl != "none" + ), "--rl-persist-cuda-graphs is set but no CUDA graphs are being built." # Training CGs only makes sense if we build any CGs. - assert not args.rl_training_cuda_graphs or args.cuda_graph_impl != "none", ( - "--rl-training-cuda-graphs is set but no CUDA graphs are being built." - ) + assert ( + not args.rl_training_cuda_graphs or args.cuda_graph_impl != "none" + ), "--rl-training-cuda-graphs is set but no CUDA graphs are being built." # If CUDA graphs persist and KV cache memory address is not static, we need # either UVM or torch_memory_saver to maintain memory address stability for CGs. if args.rl_persist_cuda_graphs and args.rl_kv_cache_management_mode != "persist": @@ -501,8 +532,8 @@ def validate_args(args, defaults={}): # Offload mode requires CG persistence: CG recapture runs dummy forward # passes that corrupt the preserved KV data. - assert ( - (not args.rl_kv_cache_management_mode == "offload") or (args.rl_persist_cuda_graphs) + assert (not args.rl_kv_cache_management_mode == "offload") or ( + args.rl_persist_cuda_graphs ), "--rl-kv-cache-management-mode=offload requires --rl-persist-cuda-graphs" # There's no need to manually offload the KV cache with UVM. @@ -511,7 +542,7 @@ def validate_args(args, defaults={}): and args.rl_kv_cache_management_mode == "offload" ), "--rl-kv-cache-management-mode=offload is incompatible with UVM" # We currently cannot recapture CGs in offload mode. - assert not( + assert not ( not args.rl_persist_cuda_graphs and args.rl_kv_cache_management_mode == "offload" ), "Cannot recapture CUDA graphs while offloading KV cache." @@ -528,82 +559,75 @@ def validate_args(args, defaults={}): "installed. See https://github.com/fzyzcjy/torch_memory_saver." ) - if args.rl_max_inflight_requests is not None: - requests_per_batch = args.grpo_prompts_per_step * args.grpo_group_size - assert args.rl_generation_lag is None, \ - "--rl-generation-lag and --rl-max-inflight-requests are mutually exclusive." - assert args.rl_max_inflight_requests >= 1, \ - f"--rl-max-inflight-requests ({args.rl_max_inflight_requests}) must be >= 1." - if args.rl_max_inflight_requests > requests_per_batch: - assert args.rl_partial_rollouts, \ - f"--rl-max-inflight-requests above one training batch " \ - f"({requests_per_batch} requests) requires --rl-partial-rollouts." - # Total in-flight requests = (lag + 1) trainer batches of P * G requests each. - args.rl_generation_lag = args.rl_max_inflight_requests / requests_per_batch - 1 - if args.rl_generation_lag is None: - # With --rl-partial-rollouts the lag is autotuned from engine capacity - # at inference launch; otherwise generation is fully synchronous. - if not args.rl_partial_rollouts: - args.rl_generation_lag = 0 - else: - assert args.rl_generation_lag >= -1, \ - f"--rl-generation-lag ({args.rl_generation_lag}) must be >= -1." - if args.rl_generation_lag > 0: - assert args.rl_partial_rollouts, \ - "--rl-generation-lag requires --rl-partial-rollouts." - assert args.rl_submission_granularity == "B" or args.rl_partial_rollouts, \ - f"--rl-submission-granularity {args.rl_submission_granularity} requires " \ - "--rl-partial-rollouts." - assert args.rl_consumption_granularity != "R", \ - "--rl-consumption-granularity R is not currently supported." + submit_rollouts_at_rollout_granularity = args.rl_submission_granularity == "R" + if args.rl_generation_lag > 0: + assert args.rl_partial_rollouts, "--rl-generation-lag requires --rl-partial-rollouts." + if submit_rollouts_at_rollout_granularity: + assert ( + args.rl_partial_rollouts + ), "Rollout submission granularity requires streaming grouped rollouts." + assert ( + args.rl_consumption_granularity != "R" + ), "--rl-consumption-granularity R is not currently supported." assert not ( - args.rl_submission_granularity == "B" - and args.rl_consumption_granularity == "G" + args.rl_submission_granularity == "B" and args.rl_consumption_granularity == "G" ), "--rl-submission-granularity B with --rl-consumption-granularity G is not supported." args.grpo_samples_per_iteration = args.grpo_prompts_per_step * args.grpo_group_size if args.rl_use_sequence_packing: - assert args.micro_batch_size == 1, \ - "micro_batch_size must be 1 when using sequence packing. To increase compute per micro batch increase the sequence length." - - print_rank_0('using world size: {}, data-parallel size: {}, ' - 'context-parallel size: {}, ' - 'hierarchical context-parallel sizes: {}, ' - 'tensor-model-parallel size: {}, ' - 'pipeline-model-parallel size: {}'.format( - args.world_size, args.data_parallel_size, - args.context_parallel_size, - args.hierarchical_context_parallel_sizes, - args.tensor_model_parallel_size, - args.pipeline_model_parallel_size)) + assert ( + args.micro_batch_size == 1 + ), "micro_batch_size must be 1 when using sequence packing. To increase compute per micro batch increase the sequence length." + + print_rank_0( + 'using world size: {}, data-parallel size: {}, ' + 'context-parallel size: {}, ' + 'hierarchical context-parallel sizes: {}, ' + 'tensor-model-parallel size: {}, ' + 'pipeline-model-parallel size: {}'.format( + args.world_size, + args.data_parallel_size, + args.context_parallel_size, + args.hierarchical_context_parallel_sizes, + args.tensor_model_parallel_size, + args.pipeline_model_parallel_size, + ) + ) # Checks. if args.hierarchical_context_parallel_sizes: from numpy import prod + assert args.context_parallel_size == prod(args.hierarchical_context_parallel_sizes) if "a2a+p2p" in args.cp_comm_type: - assert args.hierarchical_context_parallel_sizes is not None, \ - "--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm" + assert ( + args.hierarchical_context_parallel_sizes is not None + ), "--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm" if args.expert_tensor_parallel_size is None: args.expert_tensor_parallel_size = args.tensor_model_parallel_size # Deprecated arguments. - assert args.batch_size is None, '--batch-size argument is no longer ' \ - 'valid, use --micro-batch-size instead' + assert args.batch_size is None, ( + '--batch-size argument is no longer ' 'valid, use --micro-batch-size instead' + ) del args.batch_size - assert args.warmup is None, '--warmup argument is no longer valid, use ' \ - '--lr-warmup-fraction instead' + assert args.warmup is None, ( + '--warmup argument is no longer valid, use ' '--lr-warmup-fraction instead' + ) del args.warmup - assert args.model_parallel_size is None, '--model-parallel-size is no ' \ - 'longer valid, use --tensor-model-parallel-size instead' + assert args.model_parallel_size is None, ( + '--model-parallel-size is no ' 'longer valid, use --tensor-model-parallel-size instead' + ) del args.model_parallel_size if args.checkpoint_activations: - print_rank_0('--checkpoint-activations is no longer valid, use --recompute-activations, ' - 'or, for more control, --recompute-granularity and --recompute-method.') + print_rank_0( + '--checkpoint-activations is no longer valid, use --recompute-activations, ' + 'or, for more control, --recompute-granularity and --recompute-method.' + ) exit() del args.checkpoint_activations @@ -634,12 +658,10 @@ def validate_args(args, defaults={}): del args.external_cuda_graph if getattr(args, 'cuda_graph_scope_deprecated', None) is not None: - assert not args.cuda_graph_modules, ( - "--cuda-graph-scope and --cuda-graph-modules cannot be used together." - ) - warn_rank_0( - '--cuda-graph-scope is deprecated, use --cuda-graph-modules instead.' - ) + assert ( + not args.cuda_graph_modules + ), "--cuda-graph-scope and --cuda-graph-modules cannot be used together." + warn_rank_0('--cuda-graph-scope is deprecated, use --cuda-graph-modules instead.') args.cuda_graph_modules = args.cuda_graph_scope_deprecated del args.cuda_graph_scope_deprecated @@ -647,10 +669,7 @@ def validate_args(args, defaults={}): # all subsequent validation sees fully-typed enum values. _normalize_cuda_graph_modules_args(args) _normalize_inference_cuda_graph_scope_arg(args) - assert ( - args.inference_cuda_graph_scope - in ALLOWED_INFERENCE_SCOPES[args.cuda_graph_impl] - ), ( + assert args.inference_cuda_graph_scope in ALLOWED_INFERENCE_SCOPES[args.cuda_graph_impl], ( "Invalid inference CUDA graph scope " f"{args.inference_cuda_graph_scope.name!r} for " f"--cuda-graph-impl={args.cuda_graph_impl!r}." @@ -662,31 +681,42 @@ def validate_args(args, defaults={}): # arguments that are passed to the program. We check this by # ensuring the arg is set to None. if getattr(args, key, None) is not None: - warn_rank_0('Overriding default arguments for {key}:{v} ' - 'with {key}:{v2}'.format(key=key, v=defaults[key], - v2=getattr(args, key))) + warn_rank_0( + 'Overriding default arguments for {key}:{v} ' + 'with {key}:{v2}'.format(key=key, v=defaults[key], v2=getattr(args, key)) + ) else: setattr(args, key, defaults[key]) if args.data_path is not None and args.split is None: legacy_default_split_value = '969, 30, 1' - warn_rank_0('Please specify --split when using --data-path. Using legacy default value ' - f'of "{legacy_default_split_value}"') + warn_rank_0( + 'Please specify --split when using --data-path. Using legacy default value ' + f'of "{legacy_default_split_value}"' + ) args.split = legacy_default_split_value use_data_path = (args.data_path is not None) or (args.data_args_path is not None) if use_data_path: # Exactly one of the two has to be None if we use it. assert (args.data_path is None) or (args.data_args_path is None) - use_per_split_data_path = any( - elt is not None - for elt in [args.train_data_path, args.valid_data_path, args.test_data_path]) or \ - args.per_split_data_args_path is not None + use_per_split_data_path = ( + any( + elt is not None + for elt in [args.train_data_path, args.valid_data_path, args.test_data_path] + ) + or args.per_split_data_args_path is not None + ) if use_per_split_data_path: - # Exactly one of the two has to be None if we use it. - assert any(elt is not None - for elt in [args.train_data_path, args.valid_data_path, args.test_data_path]) is False or \ - args.per_split_data_args_path is None + # Exactly one of the two has to be None if we use it. + assert ( + any( + elt is not None + for elt in [args.train_data_path, args.valid_data_path, args.test_data_path] + ) + is False + or args.per_split_data_args_path is None + ) if args.phase_transition_iterations: args.phase_transition_iterations = sorted( @@ -699,9 +729,7 @@ def validate_args(args, defaults={}): args, '_is_global_batch_size_explicitly_specified', args.global_batch_size is not None ) if args.step_batch_size_schedule is not None and is_global_batch_size_explicitly_specified: - raise ValueError( - 'Cannot specify both --step-batch-size-schedule and --global-batch-size' - ) + raise ValueError('Cannot specify both --step-batch-size-schedule and --global-batch-size') if args.global_batch_size is None: args.global_batch_size = args.micro_batch_size * args.data_parallel_size print_rank_0('setting global batch size to {}'.format(args.global_batch_size)) @@ -712,34 +740,36 @@ def validate_args(args, defaults={}): args.eval_global_batch_size = args.global_batch_size if args.eval_micro_batch_size is None: args.eval_micro_batch_size = args.micro_batch_size - # data_parallel_size is the replicate degree, so multiply the GTP-remat axis back in: evaluate() - # divides eval_global_batch_size by the same product to get its microbatch count, and without - # gtp_weight_remat_size here that division can silently floor (down to zero microbatches). - assert args.eval_global_batch_size % ( - args.eval_micro_batch_size * args.data_parallel_size * args.gtp_weight_remat_size - ) == 0, \ - f"eval_global_batch_size ({args.eval_global_batch_size}) must be divisible by " \ - f"eval_micro_batch_size ({args.eval_micro_batch_size}) * data_parallel_size ({args.data_parallel_size})" \ - f" * gtp_weight_remat_size ({args.gtp_weight_remat_size})" + assert ( + args.eval_global_batch_size % (args.eval_micro_batch_size * args.data_parallel_size) == 0 + ), ( + f"eval_global_batch_size ({args.eval_global_batch_size}) must be divisible by " + f"eval_micro_batch_size ({args.eval_micro_batch_size}) * data_parallel_size ({args.data_parallel_size})" + ) if args.perform_rl_step: num_generated_samples_per_inference_iteration = ( - args.grpo_samples_per_iteration * args.grpo_iterations) + args.grpo_samples_per_iteration * args.grpo_iterations + ) # Ensure that the number of prompts we collect is a multiple of the global batch size. - assert num_generated_samples_per_inference_iteration % args.global_batch_size == 0, \ - f"grpo_group_size * grpo_prompts_per_step * grpo_iterations should be divisible by global_batch_size" + assert ( + num_generated_samples_per_inference_iteration % args.global_batch_size == 0 + ), f"grpo_group_size * grpo_prompts_per_step * grpo_iterations should be divisible by global_batch_size" # For now only exit/checkpoint on iterations where we generate data. We don't currently # have a way to checkpoint the generated data. num_training_iterations_per_inference_iteration = ( - num_generated_samples_per_inference_iteration // args.global_batch_size) + num_generated_samples_per_inference_iteration // args.global_batch_size + ) if args.exit_interval is not None: - assert args.exit_interval % num_training_iterations_per_inference_iteration == 0, \ - f"exit_interval should be divisible by number of global batches per inference iteration." + assert ( + args.exit_interval % num_training_iterations_per_inference_iteration == 0 + ), f"exit_interval should be divisible by number of global batches per inference iteration." if args.save_interval is not None: - assert args.save_interval % num_training_iterations_per_inference_iteration == 0, \ - f"save_interval should be divisible by number of global batches per inference iteration." + assert ( + args.save_interval % num_training_iterations_per_inference_iteration == 0 + ), f"save_interval should be divisible by number of global batches per inference iteration." # === Hybrid layer pattern: deprecation handling and validation === @@ -772,6 +802,7 @@ def validate_args(args, defaults={}): get_hybrid_total_pipeline_segment_count, parse_hybrid_pattern, ) + sep = Symbols.MTP_SEPARATOR # Backward compat: convert legacy mtp_hybrid_override_pattern to unified format @@ -783,10 +814,14 @@ def validate_args(args, defaults={}): ): main_pattern = args.hybrid_layer_pattern or '' mtp_pattern = args.mtp_hybrid_override_pattern - args.hybrid_layer_pattern = main_pattern + sep + sep.join([mtp_pattern] * args.mtp_num_layers) + args.hybrid_layer_pattern = ( + main_pattern + sep + sep.join([mtp_pattern] * args.mtp_num_layers) + ) args.mtp_hybrid_override_pattern = None print_rank_0(f"Converted legacy MTP pattern to unified: {args.hybrid_layer_pattern}") + parsed_hybrid_pattern = parse_hybrid_pattern(args.hybrid_layer_pattern) + if args.hybrid_layer_pattern is not None: # Derive num_layers from pattern; hybrid_layer_pattern always overrides --num-layers when # both are present (e.g. when loading from checkpoint with --use-checkpoint-args). @@ -866,9 +901,8 @@ def validate_args(args, defaults={}): # Infer mtp_num_layers from unified pattern if args.hybrid_layer_pattern and sep in args.hybrid_layer_pattern: - parsed = parse_hybrid_pattern(args.hybrid_layer_pattern) - if parsed.mtp_pattern and parsed.mtp_num_depths > 0: - inferred_mtp_num_layers = parsed.mtp_num_depths + if parsed_hybrid_pattern.mtp_pattern and parsed_hybrid_pattern.mtp_num_depths > 0: + inferred_mtp_num_layers = parsed_hybrid_pattern.mtp_num_depths if args.mtp_num_layers is None: args.mtp_num_layers = inferred_mtp_num_layers elif args.mtp_num_layers != inferred_mtp_num_layers: @@ -877,7 +911,7 @@ def validate_args(args, defaults={}): f"MTP depth count ({inferred_mtp_num_layers}) in pattern " f"'{args.hybrid_layer_pattern}'. " f"Using the inferred value ({inferred_mtp_num_layers}).", - args.rank + args.rank, ) args.mtp_num_layers = inferred_mtp_num_layers @@ -888,25 +922,19 @@ def validate_args(args, defaults={}): + f"The supported position embedding types are rope and none." ) - if args.mtp_hsm and not (args.mtp_num_layers and args.mtp_num_layers >= 2): - warn_rank_0( - "--mtp-hsm needs at least two MTP layers to mix anything, but " - f"--mtp-num-layers is {args.mtp_num_layers}. Disabling Hidden State Mixing.", - args.rank, - ) - args.mtp_hsm = False - # Validate MTP args for hybrid vs non-hybrid models if args.hybrid_layer_pattern is not None: # Mamba/hybrid model MTP validation - if args.mtp_num_layers and not (args.hybrid_layer_pattern and sep in args.hybrid_layer_pattern): + if args.mtp_num_layers and not ( + args.hybrid_layer_pattern and sep in args.hybrid_layer_pattern + ): # Hybrid model wants MTP but no unified pattern - check for legacy args if args.mtp_hybrid_override_pattern is None: warn_rank_0( "Hybrid model with --mtp-num-layers but no MTP pattern. " "Use unified --hybrid-layer-pattern with '/' separator (e.g., 'M*M*/MM/MM') " "or legacy --mtp-hybrid-override-pattern for old checkpoints.", - args.rank + args.rank, ) else: # Non-hybrid (GPT) model MTP validation @@ -915,13 +943,15 @@ def validate_args(args, defaults={}): "--mtp-hybrid-override-pattern is for Mamba/hybrid models only. " "For GPT models, MTP replicates the main transformer layer structure. " "This argument will be ignored.", - args.rank + args.rank, ) - # Infer use of MLA from unified pattern - if args.hybrid_layer_pattern and ( - Symbols.MLA in args.hybrid_layer_pattern - or Symbols.DS_ATTENTION in args.hybrid_layer_pattern + # Infer use of MLA from the parsed main and MTP patterns before config-class selection. + if any( + symbol in pattern + for pattern in (parsed_hybrid_pattern.main_pattern, parsed_hybrid_pattern.mtp_pattern) + if pattern is not None + for symbol in Symbols.MLA_ATTENTION ): args.multi_latent_attention = True @@ -944,7 +974,9 @@ def validate_args(args, defaults={}): if args.pipeline_model_parallel_layout is not None: # Parse the input flattened layout to a list and get the vpp size. # We will validate the layout more carefully in the TransformerConfig constructor. - num_stages = PipelineParallelLayerLayout.get_num_stages_from_str(args.pipeline_model_parallel_layout) + num_stages = PipelineParallelLayerLayout.get_num_stages_from_str( + args.pipeline_model_parallel_layout + ) assert num_stages % args.pipeline_model_parallel_size == 0, ( f"The length of pipeline_model_parallel_layout must be divisible" f" by pipeline_model_parallel_size ({num_stages=}," @@ -953,10 +985,15 @@ def validate_args(args, defaults={}): args.virtual_pipeline_model_parallel_size = num_stages // args.pipeline_model_parallel_size if args.virtual_pipeline_model_parallel_size == 1: args.virtual_pipeline_model_parallel_size = None - elif args.num_layers_per_virtual_pipeline_stage is not None or args.num_virtual_stages_per_pipeline_rank is not None: + elif ( + args.num_layers_per_virtual_pipeline_stage is not None + or args.num_virtual_stages_per_pipeline_rank is not None + ): if args.num_virtual_stages_per_pipeline_rank is None: - assert args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None, \ - 'please use --num-virtual-stages-per-pipeline-rank to specify virtual pipeline parallel degree when enable uneven pipeline parallelism' + assert ( + args.decoder_first_pipeline_num_layers is None + and args.decoder_last_pipeline_num_layers is None + ), 'please use --num-virtual-stages-per-pipeline-rank to specify virtual pipeline parallel degree when enable uneven pipeline parallelism' if args.num_layers is not None: num_layers = args.num_layers else: @@ -968,14 +1005,19 @@ def validate_args(args, defaults={}): if args.account_for_loss_in_pipeline_split: num_layers += 1 - assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \ - 'number of layers of the model must be divisible pipeline model parallel size' - num_layers_per_pipeline_stage = num_layers // args.transformer_pipeline_model_parallel_size + assert ( + num_layers % args.transformer_pipeline_model_parallel_size == 0 + ), 'number of layers of the model must be divisible pipeline model parallel size' + num_layers_per_pipeline_stage = ( + num_layers // args.transformer_pipeline_model_parallel_size + ) - assert num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0, \ - 'number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage' - args.virtual_pipeline_model_parallel_size = num_layers_per_pipeline_stage // \ - args.num_layers_per_virtual_pipeline_stage + assert ( + num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0 + ), 'number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage' + args.virtual_pipeline_model_parallel_size = ( + num_layers_per_pipeline_stage // args.num_layers_per_virtual_pipeline_stage + ) else: args.virtual_pipeline_model_parallel_size = args.num_virtual_stages_per_pipeline_rank if args.virtual_pipeline_model_parallel_size == 1: @@ -985,7 +1027,10 @@ def validate_args(args, defaults={}): if args.hybrid_layer_pattern is None: args.virtual_pipeline_model_parallel_size = None - if args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None: + if ( + args.decoder_first_pipeline_num_layers is None + and args.decoder_last_pipeline_num_layers is None + ): # Divisibility check not applicable for T5 models which specify encoder_num_layers # and decoder_num_layers, or for hybrid models using --hybrid-layer-pattern. if args.num_layers is not None and args.hybrid_layer_pattern is None: @@ -997,59 +1042,73 @@ def validate_args(args, defaults={}): if args.account_for_loss_in_pipeline_split: num_layers += 1 - assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \ - 'Number of layers should be divisible by the pipeline-model-parallel size' + assert ( + num_layers % args.transformer_pipeline_model_parallel_size == 0 + ), 'Number of layers should be divisible by the pipeline-model-parallel size' if args.virtual_pipeline_model_parallel_size is not None: if args.overlap_p2p_comm: - assert args.pipeline_model_parallel_size > 1, \ - 'When interleaved schedule is used, pipeline-model-parallel size '\ + assert args.pipeline_model_parallel_size > 1, ( + 'When interleaved schedule is used, pipeline-model-parallel size ' 'should be greater than 1' + ) else: - assert args.pipeline_model_parallel_size > 2, \ - 'When interleaved schedule is used and p2p communication overlap is disabled, '\ - 'pipeline-model-parallel size should be greater than 2 to avoid having multiple '\ + assert args.pipeline_model_parallel_size > 2, ( + 'When interleaved schedule is used and p2p communication overlap is disabled, ' + 'pipeline-model-parallel size should be greater than 2 to avoid having multiple ' 'p2p sends and recvs between same 2 ranks per communication batch' + ) else: # Overlap P2P communication is disabled if not using the interleaved schedule. args.overlap_p2p_comm = False args.align_param_gather = False # Only print warning if PP size > 1. if args.rank == 0 and args.pipeline_model_parallel_size > 1: - print('WARNING: Setting args.overlap_p2p_comm and args.align_param_gather to False ' + print( + 'WARNING: Setting args.overlap_p2p_comm and args.align_param_gather to False ' 'since non-interleaved schedule does not support overlapping p2p communication ' - 'and aligned param AG') + 'and aligned param AG' + ) print_rank_0( f"Number of virtual stages per pipeline stage: {args.virtual_pipeline_model_parallel_size}" ) if args.overlap_param_gather: - assert args.use_distributed_optimizer or args.use_megatron_fsdp \ - or args.optimizer == 'dist_muon', \ - '--overlap-param-gather only supported with distributed optimizer, megatron fsdp, or dist_muon' - assert args.overlap_grad_reduce, \ - 'Must use --overlap-param-gather with --overlap-grad-reduce' + assert ( + args.use_distributed_optimizer + or args.use_megatron_fsdp + or args.optimizer == 'dist_muon' + ), '--overlap-param-gather only supported with distributed optimizer, megatron fsdp, or dist_muon' + assert ( + args.overlap_grad_reduce + ), 'Must use --overlap-param-gather with --overlap-grad-reduce' if args.use_torch_fsdp2: - assert is_torch_min_version("2.4.0"), \ - 'FSDP2 requires PyTorch >= 2.4.0 with FSDP 2 support.' - assert args.pipeline_model_parallel_size == 1, \ - '--use-torch-fsdp2 is not supported with pipeline parallelism' - assert args.expert_model_parallel_size == 1, \ - '--use-torch-fsdp2 is not supported with expert parallelism' - assert not args.use_distributed_optimizer, \ - "--use-torch-fsdp2 is not supported with MCore's distributed optimizer" - assert not args.gradient_accumulation_fusion, \ - '--use-torch-fsdp2 is not supported with gradient accumulation fusion' - assert args.ckpt_format in ('torch_dist', 'torch_dcp'), \ - '--use-torch-fsdp2 requires --ckpt-format torch_dist or torch_dcp' - assert args.untie_embeddings_and_output_weights, \ - '--use-torch-fsdp2 requires --untie-embeddings-and-output-weights' - assert not args.fp16, \ - '--use-torch-fsdp2 not supported with fp16 yet' - assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1", \ - 'FSDP always requires CUDA_DEVICE_MAX_CONNECTIONS value large than one' + assert is_torch_min_version("2.4.0"), 'FSDP2 requires PyTorch >= 2.4.0 with FSDP 2 support.' + assert ( + args.pipeline_model_parallel_size == 1 + ), '--use-torch-fsdp2 is not supported with pipeline parallelism' + assert ( + args.expert_model_parallel_size == 1 + ), '--use-torch-fsdp2 is not supported with expert parallelism' + assert ( + not args.use_distributed_optimizer + ), "--use-torch-fsdp2 is not supported with MCore's distributed optimizer" + assert ( + not args.gradient_accumulation_fusion + ), '--use-torch-fsdp2 is not supported with gradient accumulation fusion' + assert args.ckpt_format in ( + 'torch_dist', + 'torch_dcp', + ), '--use-torch-fsdp2 requires --ckpt-format torch_dist or torch_dcp' + assert ( + args.untie_embeddings_and_output_weights + ), '--use-torch-fsdp2 requires --untie-embeddings-and-output-weights' + assert not args.fp16, '--use-torch-fsdp2 not supported with fp16 yet' + assert ( + os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1" + ), 'FSDP always requires CUDA_DEVICE_MAX_CONNECTIONS value large than one' if args.fp8_param_gather and is_te_min_version("2.0.0"): args.fp8_param_gather = False @@ -1067,18 +1126,27 @@ def validate_args(args, defaults={}): ) if args.overlap_param_gather_with_optimizer_step: - assert args.use_distributed_optimizer, \ - '--overlap-param-gather-with-optimizer-step only supported with distributed optimizer' - assert args.overlap_param_gather, \ - 'Must use --overlap-param-gather-with-optimizer-step with --overlap-param-gather' - assert args.virtual_pipeline_model_parallel_size is not None, \ - '--overlap-param-gather-with-optimizer-step only supported with interleaved pipeline parallelism' - assert not args.use_dist_ckpt, \ - '--overlap-param-gather-with-optimizer-step not supported with distributed checkpointing yet' + assert ( + args.use_distributed_optimizer + ), '--overlap-param-gather-with-optimizer-step only supported with distributed optimizer' + assert ( + args.overlap_param_gather + ), 'Must use --overlap-param-gather-with-optimizer-step with --overlap-param-gather' + assert ( + args.virtual_pipeline_model_parallel_size is not None + ), '--overlap-param-gather-with-optimizer-step only supported with interleaved pipeline parallelism' + assert ( + not args.use_dist_ckpt + ), '--overlap-param-gather-with-optimizer-step not supported with distributed checkpointing yet' # Map string data-type to torch.dtype. dtype_map = { - 'fp32': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16, 'fp8': torch.uint8, 'auto': None, None: None, + 'fp32': torch.float32, + 'bf16': torch.bfloat16, + 'fp16': torch.float16, + 'fp8': torch.uint8, + 'auto': None, + None: None, } map_dtype = lambda d: d if isinstance(d, torch.dtype) else dtype_map[d] @@ -1089,19 +1157,24 @@ def validate_args(args, defaults={}): args.mamba_inference_conv_states_dtype = map_dtype(args.mamba_inference_conv_states_dtype) args.mamba_inference_ssm_states_dtype = map_dtype(args.mamba_inference_ssm_states_dtype) args.mamba_training_ssm_states_dtype = map_dtype(args.mamba_training_ssm_states_dtype) - args.logit_dtype = map_dtype(getattr(args, 'logit_dtype', None)) args.megatron_fsdp_main_params_dtype = map_dtype(args.megatron_fsdp_main_params_dtype) args.megatron_fsdp_main_grads_dtype = map_dtype(args.megatron_fsdp_main_grads_dtype) args.megatron_fsdp_grad_comm_dtype = map_dtype(args.megatron_fsdp_grad_comm_dtype) if args.fp8_param_gather: - assert args.use_distributed_optimizer or args.use_torch_fsdp2 or args.use_megatron_fsdp or not torch.is_grad_enabled(), \ - '--fp8-param-gather only supported with distributed optimizer, torch fsdp2, megatron fsdp, or inference mode' + assert ( + args.use_distributed_optimizer + or args.use_torch_fsdp2 + or args.use_megatron_fsdp + or not torch.is_grad_enabled() + ), '--fp8-param-gather only supported with distributed optimizer, torch fsdp2, megatron fsdp, or inference mode' # FP4 and FP8 are mutually exclusive if args.fp4 and args.fp8: - raise ValueError("--fp4-format and --fp8-format cannot be used simultaneously. Please choose one.") + raise ValueError( + "--fp4-format and --fp8-format cannot be used simultaneously. Please choose one." + ) # FP4 param requires FP4 mode if args.fp4_param_gather and not args.fp4: @@ -1109,7 +1182,9 @@ def validate_args(args, defaults={}): # FP4 requires TE >= 2.7.0.dev0 if args.fp4 and not is_te_min_version("2.7.0.dev0"): - raise ValueError("--fp4-format requires Transformer Engine >= 2.7.0.dev0 for NVFP4BlockScaling support.") + raise ValueError( + "--fp4-format requires Transformer Engine >= 2.7.0.dev0 for NVFP4BlockScaling support." + ) if ( args.fp8_recipe == 'mxfp8' @@ -1118,24 +1193,29 @@ def validate_args(args, defaults={}): ): raise ValueError("MXFP8 with inference optimized layers requires FlashInfer >= 0.6.4") + if args.inference_dynamic_batching_sampling_backend == 'flashinfer': + try: + import flashinfer # noqa: F401 + except ImportError as e: + raise ImportError( + "--inference-dynamic-batching-sampling-backend=flashinfer requires " + "the flashinfer package; install it or pass " + "--inference-dynamic-batching-sampling-backend=torch." + ) from e + if args.use_megatron_fsdp: # NOTE: The flag `use_custom_fsdp` is deprecated and will be removed in future versions. # Please use `use_megatron_fsdp` instead, as all functionality will be migrated there. # Future updates will drop support for `use_custom_fsdp` to avoid confusion. args.use_custom_fsdp = True - if args.megatron_fsdp_version == 2: - assert not args.use_distributed_optimizer, \ - '--megatron-fsdp-version 2 is not compatible with --use-distributed-optimizer' - else: - # Megatron-FSDP v1 requires the DistributedOptimizer. - if not args.use_distributed_optimizer: - warn_rank_0( - 'Megatron-FSDP v1 is only compatible with --use-distributed-optimizer. ' - 'Using DistributedOptimizer...', - args.rank, - ) - args.use_distributed_optimizer = True + # Megatron-FSDP requires the DistributedOptimizer. + if not args.use_distributed_optimizer: + warn_rank_0( + 'Megatron-FSDP is only compatible with --use-distributed-optimizer. Using DistributedOptimizer...', + args.rank, + ) + args.use_distributed_optimizer = True # Optimizer step MXFP8 buffer operation that is not relevant or supported for Megatron-FSDP. args.reuse_grad_buf_for_mxfp8_param_ag = False if args.moe_single_grouped_weight or args.moe_single_grouped_bias: @@ -1150,8 +1230,10 @@ def validate_args(args, defaults={}): "GroupedTensor param buffers." ) # Optimizer compatibility check. - assert args.optimizer in ('sgd', 'adam'), \ - f"Megatron-FSDP does not support the {args.optimizer} optimizer yet." + assert args.optimizer in ( + 'sgd', + 'adam', + ), f"Megatron-FSDP does not support the {args.optimizer} optimizer yet." if ( args.data_parallel_sharding_strategy in ["optim_grads_params", "optim_grads"] @@ -1163,21 +1245,40 @@ def validate_args(args, defaults={}): ) if args.data_parallel_sharding_strategy == "optim_grads_params": - assert args.check_weight_hash_across_dp_replicas_interval is None, \ - 'check_weight_hash_across_dp_replicas_interval is not supported with optim_grads_params' + assert ( + args.check_weight_hash_across_dp_replicas_interval is None + ), 'check_weight_hash_across_dp_replicas_interval is not supported with optim_grads_params' - assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1", \ - 'FSDP requires CUDA_DEVICE_MAX_CONNECTIONS > 1 or unset.' + assert ( + os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1" + ), 'FSDP requires CUDA_DEVICE_MAX_CONNECTIONS > 1 or unset.' + + assert ( + args.ckpt_format == "fsdp_dtensor" + ), "Megatron-FSDP requires the `fsdp_dtensor` checkpointing format." - assert args.ckpt_format == "fsdp_dtensor", \ - "Megatron-FSDP requires the `fsdp_dtensor` checkpointing format." + if args.megatron_fsdp_prefetch_recompute_forward_weights: + assert args.data_parallel_sharding_strategy == "optim_grads_params", ( + "--megatron-fsdp-prefetch-recompute-forward-weights is only supported " + 'with --data-parallel-sharding-strategy optim_grads_params.' + ) + assert args.recompute_granularity == "full", ( + "--megatron-fsdp-prefetch-recompute-forward-weights is only supported " + "with full activation recomputation." + ) + assert not args.overlap_moe_expert_parallel_comm, ( + "--megatron-fsdp-prefetch-recompute-forward-weights is not supported " + "with --overlap-moe-expert-parallel-comm." + ) if args.nccl_ub: # In Megatron-LM, required implementation for manual registration is already provided. # So we enable the manual registration by default when nccl-ub and use_megatron_fsdp is set. args.fsdp_manual_registration = True args.fsdp_double_buffer = True - warn_rank_0('FSDP double buffer and manual registration is enabled by default when --nccl-ub is enabled!') + warn_rank_0( + 'FSDP double buffer and manual registration is enabled by default when --nccl-ub is enabled!' + ) if args.megatron_fsdp_max_pool_double_buffer: # MaxPoolAllocator is a type of FSDP double buffer. @@ -1188,9 +1289,18 @@ def validate_args(args, defaults={}): "Meta device initialization (init_model_with_meta_device=True) is not " "supported or necessary for the 'no_shard' / 0 sharding strategy." ) + else: + assert not args.megatron_fsdp_prefetch_recompute_forward_weights, ( + "--megatron-fsdp-prefetch-recompute-forward-weights requires " "--use-megatron-fsdp." + ) + assert not args.megatron_fsdp_cache_param_bucket_views, ( + "--megatron-fsdp-cache-param-bucket-views requires " "--use-megatron-fsdp." + ) if args.fsdp_manual_registration: - assert args.use_megatron_fsdp, "FSDP manual registration is only supported with Megatron FSDP." + assert ( + args.use_megatron_fsdp + ), "FSDP manual registration is only supported with Megatron FSDP." assert args.nccl_ub, "FSDP manual registration is only supported with --nccl-ub argument." # Parameters dtype. @@ -1202,55 +1312,48 @@ def validate_args(args, defaults={}): # where NaNs in grads / loss are signal to the loss scaler. if not args.loss_scale: args.check_for_nan_in_loss_and_grad = False - warn_rank_0('Setting args.check_for_nan_in_loss_and_grad to False since ' - 'dynamic loss scaling is being used') + warn_rank_0( + 'Setting args.check_for_nan_in_loss_and_grad to False since ' + 'dynamic loss scaling is being used' + ) if args.bf16: assert not args.fp16 args.params_dtype = torch.bfloat16 # bfloat16 requires gradient accumulation and all-reduce to # be done in fp32. if args.accumulate_allreduce_grads_in_fp32: - assert args.main_grads_dtype == torch.float32, \ - "--main-grads-dtype can only be fp32 when --accumulate-allreduce-grads-in-fp32 is set" + assert ( + args.main_grads_dtype == torch.float32 + ), "--main-grads-dtype can only be fp32 when --accumulate-allreduce-grads-in-fp32 is set" if args.grad_reduce_in_bf16: args.accumulate_allreduce_grads_in_fp32 = False elif not args.accumulate_allreduce_grads_in_fp32 and args.main_grads_dtype == torch.float32: args.accumulate_allreduce_grads_in_fp32 = True print_rank_0('accumulate and all-reduce gradients in fp32 for bfloat16 data type.') - - if args.accumulate_allreduce_grads_in_fp32: - # The FP32-accumulation reduce-scatters only exist to recover FP32 accumulation from a - # lower-precision wire dtype. With FP32 main_grads the plain ring reduce-scatter already - # sums in FP32 on both axes, so they buy no precision and cost one extra - # unsharded-wgrad-sized scratch buffer per in-flight reduce-scatter. - for arg_name in ( - 'ddp_reduce_scatter_with_fp32_accumulation', - 'gtp_remat_reduce_scatter_with_fp32_accumulation', - ): - if getattr(args, arg_name, False): - setattr(args, arg_name, False) - warn_rank_0( - f"Setting args.{arg_name} to False since " - "--accumulate-allreduce-grads-in-fp32 already reduces in fp32" - ) - if args.cuda_graph_impl == "full_iteration": - assert not args.check_for_nan_in_loss_and_grad, \ - "--no-check-for-nan-in-loss-and-grad should be set with --cuda-graph-impl=full_iteration for training." + assert ( + not args.check_for_nan_in_loss_and_grad + ), "--no-check-for-nan-in-loss-and-grad should be set with --cuda-graph-impl=full_iteration for training." if args.inference_cuda_graph_scope == InferenceCudaGraphScope.block: if args.fp8 is not None: - assert args.transformer_impl == "inference_optimized", \ - "fp8 with --inference-cuda-graph-scope=block is only supported with " \ + assert args.transformer_impl == "inference_optimized", ( + "fp8 with --inference-cuda-graph-scope=block is only supported with " "--transformer-impl=inference_optimized" - assert args.fp8_recipe == "mxfp8", \ - "Only --fp8-recipe=mxfp8 is supported with --inference-cuda-graph-scope=block" + ) + assert ( + args.fp8_recipe == "mxfp8" + ), "Only --fp8-recipe=mxfp8 is supported with --inference-cuda-graph-scope=block" if args.cuda_graph_impl == 'local': - assert args.inference_dynamic_batching_num_cuda_graphs > 0 or args.inference_dynamic_batching_num_cuda_graphs == -1, \ - 'inference_dynamic_batching_num_cuda_graphs should be a positive integer or -1' \ + assert ( + args.inference_dynamic_batching_num_cuda_graphs > 0 + or args.inference_dynamic_batching_num_cuda_graphs == -1 + ), ( + 'inference_dynamic_batching_num_cuda_graphs should be a positive integer or -1' '-1 means that we will automatically determine the number of CUDA graphs to capture based on the `max_requests` value.' + ) print_rank_0('using {} for parameters ...'.format(args.params_dtype)) @@ -1272,43 +1375,39 @@ def validate_args(args, defaults={}): if args.train_iters and not args.skip_train: # If we use iteration-based training, make sure the # sample-based options are off. - assert args.train_samples is None, \ - 'expected iteration-based training' - assert args.lr_decay_samples is None, \ - 'expected iteration-based learning rate decay' - assert args.lr_warmup_samples == 0, \ - 'expected iteration-based learning rate warmup' + assert args.train_samples is None, 'expected iteration-based training' + assert args.lr_decay_samples is None, 'expected iteration-based learning rate decay' + assert args.lr_warmup_samples == 0, 'expected iteration-based learning rate warmup' if args.lr_warmup_fraction is not None: - assert args.lr_warmup_iters == 0, \ - 'can only specify one of lr-warmup-fraction and lr-warmup-iters' + assert ( + args.lr_warmup_iters == 0 + ), 'can only specify one of lr-warmup-fraction and lr-warmup-iters' # Sample-based training. if args.train_samples and not args.skip_train: # If we use sample-based training, make sure the # iteration-based options are off. - assert args.train_iters is None, \ - 'expected sample-based training' - assert args.lr_decay_iters is None, \ - 'expected sample-based learning rate decay' - assert args.lr_warmup_iters == 0, \ - 'expected sample-based learnig rate warmup' + assert args.train_iters is None, 'expected sample-based training' + assert args.lr_decay_iters is None, 'expected sample-based learning rate decay' + assert args.lr_warmup_iters == 0, 'expected sample-based learnig rate warmup' if args.lr_warmup_fraction is not None: - assert args.lr_warmup_samples == 0, \ - 'can only specify one of lr-warmup-fraction ' \ - 'and lr-warmup-samples' + assert args.lr_warmup_samples == 0, ( + 'can only specify one of lr-warmup-fraction ' 'and lr-warmup-samples' + ) if args.num_layers is not None: - assert args.encoder_num_layers is None, \ - 'cannot have both num-layers and encoder-num-layers specified' + assert ( + args.encoder_num_layers is None + ), 'cannot have both num-layers and encoder-num-layers specified' args.encoder_num_layers = args.num_layers else: - assert args.encoder_num_layers is not None, \ - 'either num-layers or encoder-num-layers should be specified' + assert ( + args.encoder_num_layers is not None + ), 'either num-layers or encoder-num-layers should be specified' args.num_layers = args.encoder_num_layers # Check required arguments. - required_args = ['num_layers', 'hidden_size', 'num_attention_heads', - 'max_position_embeddings'] + required_args = ['num_layers', 'hidden_size', 'num_attention_heads', 'max_position_embeddings'] for req_arg in required_args: _check_arg_is_not_none(args, req_arg) @@ -1329,9 +1428,10 @@ def validate_args(args, defaults={}): args.kv_channels = args.hidden_size // args.num_attention_heads if args.seq_length is not None and args.context_parallel_size > 1: - assert args.seq_length % (args.context_parallel_size * 2) == 0, \ - 'seq-length should be a multiple of 2 * context-parallel-size ' \ + assert args.seq_length % (args.context_parallel_size * 2) == 0, ( + 'seq-length should be a multiple of 2 * context-parallel-size ' 'if context-parallel-size > 1.' + ) if getattr(args, 'dataloader_inter_document_masking', False): # The dataset omits attention_mask when inter-document masking is @@ -1347,9 +1447,10 @@ def validate_args(args, defaults={}): args.seq_length = args.encoder_seq_length if args.seq_length is not None: - assert args.max_position_embeddings >= args.seq_length, \ - f"max_position_embeddings ({args.max_position_embeddings}) must be greater than " \ + assert args.max_position_embeddings >= args.seq_length, ( + f"max_position_embeddings ({args.max_position_embeddings}) must be greater than " f"or equal to seq_length ({args.seq_length})." + ) if args.decoder_seq_length is not None: assert args.max_position_embeddings >= args.decoder_seq_length if args.lr is not None: @@ -1368,8 +1469,9 @@ def validate_args(args, defaults={}): if args.fp16_lm_cross_entropy: assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.' if args.fp32_residual_connection: - assert args.fp16 or args.bf16, \ - 'residual connection in fp32 only supported when using fp16 or bf16.' + assert ( + args.fp16 or args.bf16 + ), 'residual connection in fp32 only supported when using fp16 or bf16.' if args.moe_grouped_gemm: dc = torch.cuda.get_device_capability() @@ -1398,30 +1500,33 @@ def validate_args(args, defaults={}): # Persistent fused layer norm. if not is_torch_min_version("1.11.0a0"): args.no_persist_layer_norm = True - print_rank_0('Persistent fused layer norm kernel is supported from ' - 'pytorch v1.11 (nvidia pytorch container paired with v1.11). ' - 'Defaulting to no_persist_layer_norm=True') + print_rank_0( + 'Persistent fused layer norm kernel is supported from ' + 'pytorch v1.11 (nvidia pytorch container paired with v1.11). ' + 'Defaulting to no_persist_layer_norm=True' + ) # Activation recomputing. if args.distribute_saved_activations: - assert args.tensor_model_parallel_size > 1, 'can distribute ' \ - 'recomputed activations only across tensor model ' \ - 'parallel groups' - assert args.recompute_granularity == 'full', \ - 'distributed recompute activations is only '\ - 'application to full recompute granularity' - assert args.recompute_method is not None, \ - 'for distributed recompute activations to work you '\ - 'need to use a recompute method ' - assert is_torch_min_version("1.10.0a0"), \ - 'distributed recompute activations are supported for pytorch ' \ - 'v1.10 and above (Nvidia Pytorch container >= 21.07). Current ' \ + assert args.tensor_model_parallel_size > 1, ( + 'can distribute ' 'recomputed activations only across tensor model ' 'parallel groups' + ) + assert args.recompute_granularity == 'full', ( + 'distributed recompute activations is only ' 'application to full recompute granularity' + ) + assert args.recompute_method is not None, ( + 'for distributed recompute activations to work you ' 'need to use a recompute method ' + ) + assert is_torch_min_version("1.10.0a0"), ( + 'distributed recompute activations are supported for pytorch ' + 'v1.10 and above (Nvidia Pytorch container >= 21.07). Current ' f'pytorch version is v{get_torch_version()}.' + ) if args.recompute_granularity == 'selective': - assert args.recompute_method is None, \ - 'recompute method is not yet supported for ' \ - 'selective recomputing granularity' + assert args.recompute_method is None, ( + 'recompute method is not yet supported for ' 'selective recomputing granularity' + ) # disable sequence parallelism when tp=1 # to avoid change in numerics when @@ -1435,36 +1540,90 @@ def validate_args(args, defaults={}): args.sequence_parallel = False if args.tp_comm_overlap: - assert args.sequence_parallel == True, 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' - - if args.hybrid_context_parallel: - assert not args.pipeline_model_parallel_size > 1, 'Hybrid context parallelism not supported with pipeline parallelism' - assert not args.enable_cuda_graph, 'Hybrid context parallelism not supported with CUDA Graph' - assert not args.use_megatron_fsdp, 'Hybrid context parallelism not supported with Megatron FSDP' - assert args.dataloader_type == 'single', 'Hybrid context parallelism only supported with single dataloader type' - assert args.calculate_per_token_loss, 'Hybrid context parallelism must be used with --calculate-per-token-loss' - - # Support for variable sequence lengths across batches/microbatches. - # set it if the dataloader supports generation of variable sequence lengths - # across batches/microbatches. Due to additional communication overhead - # during pipeline parallelism, it should not be set if sequence length - # is constant during training. - args.variable_seq_lengths = False - if args.mock_data and args.sft and args.sft_mock_dataset_config_json is None: - args.sft_mock_dataset_config_json = json.dumps( - { - "mode": "distribution", - "type": "lognormal", - "min_seq_len": args.seq_length // 2, - "max_seq_len": args.seq_length, - "mean_seq_len": args.seq_length // 4 * 3, - "lognormal_sigma": 1.1, - } + assert ( + args.sequence_parallel == True + ), 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' + + if args.dynamic_context_parallel: + assert ( + not args.enable_cuda_graph + ), 'Dynamic context parallelism not supported with CUDA Graph' + assert ( + not args.use_megatron_fsdp + ), 'Dynamic context parallelism not supported with Megatron FSDP' + assert ( + args.dataloader_type == 'single' + ), 'Dynamic context parallelism only supported with single dataloader type' + assert ( + args.calculate_per_token_loss + ), 'Dynamic context parallelism must be used with --calculate-per-token-loss' + if args.sequence_packing_scheduler is None: + args.sequence_packing_scheduler = 'default_dynamic_cp' + if args.sequence_packing_scheduler != 'default_dynamic_cp': + raise ValueError( + 'Dynamic context parallelism requires ' + 'sequence_packing_scheduler=default_dynamic_cp' + ) + + dp_cp_size = args.data_parallel_size * args.context_parallel_size + assert args.min_dynamic_context_parallel_size <= dp_cp_size, ( + f'min_dynamic_context_parallel_size ({args.min_dynamic_context_parallel_size}) ' + f'must be <= dp_size * cp_size ({dp_cp_size})' + ) + + import warnings + + warnings.warn( + f"Dynamic CP enabled: dp_size * context_parallel_size=" + f"{args.data_parallel_size * args.context_parallel_size} " + f"will be used as the maximum dynamic CP group size. " + f"Dynamic CP groups will range from " + f"min_dynamic_context_parallel_size={args.min_dynamic_context_parallel_size} " + f"to {args.data_parallel_size * args.context_parallel_size}." + ) + + if getattr(args, 'pad_packed_seq_alignment', None) is not None: + args.pad_packed_seq_alignment = _parse_pad_packed_seq_alignment( + args.pad_packed_seq_alignment ) + if args.max_seqlen_per_dp_cp_rank is None: + raise ValueError( + '--max-seqlen-per-dp-cp-rank must be set when ' + '--pad-packed-seq-alignment is enabled.' + ) + if args.pad_packed_seq_alignment != 'max': + if args.pad_packed_seq_alignment <= 0: + raise ValueError( + "--pad-packed-seq-alignment must be 'max' or a positive integer " "alignment." + ) + if args.pad_packed_seq_alignment > args.max_seqlen_per_dp_cp_rank: + raise ValueError( + '--pad-packed-seq-alignment must not exceed ' + f'--max-seqlen-per-dp-cp-rank ({args.max_seqlen_per_dp_cp_rank}), ' + f'got {args.pad_packed_seq_alignment}.' + ) + + if args.cuda_graph_impl != "none" and ( + args.sequence_packing_scheduler is not None or args.dynamic_context_parallel + ): + if getattr(args, 'pad_packed_seq_alignment', None) is None: + raise ValueError('THD CUDA Graph requires --pad-packed-seq-alignment to be set.') + if ( + args.pad_packed_seq_alignment != 'max' + and args.pad_packed_seq_alignment != args.max_seqlen_per_dp_cp_rank + ): + raise ValueError( + "THD CUDA Graph requires --pad-packed-seq-alignment='max' " + 'or --pad-packed-seq-alignment equal to ' + f'--max-seqlen-per-dp-cp-rank ({args.max_seqlen_per_dp_cp_rank}), ' + f'got {args.pad_packed_seq_alignment}.' + ) + # disable async_tensor_model_parallel_allreduce when # model parallel memory optimization is enabled - if (args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1) \ - and get_device_arch_version() < 10: + if ( + args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1 + ) and get_device_arch_version() < 10: # CUDA_DEVICE_MAX_CONNECTIONS requirement no longer exists since the Blackwell architecture if args.use_torch_fsdp2 or args.use_megatron_fsdp: fsdp_impl = "Torch-FSDP2" if args.use_torch_fsdp2 else "Megatron-FSDP" @@ -1487,9 +1646,10 @@ def validate_args(args, defaults={}): args.rank, ) else: - assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') == "1", \ - "Using tensor model parallelism or context parallelism require setting the environment variable " \ + assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') == "1", ( + "Using tensor model parallelism or context parallelism require setting the environment variable " "CUDA_DEVICE_MAX_CONNECTIONS to 1" + ) # Setting FSDP communication groups for high priority streams for Blackwell and later architectures # Assigning high priority to communication streams ensures that communication kernels are scheduled @@ -1497,123 +1657,9 @@ def validate_args(args, defaults={}): if args.use_torch_fsdp2 or args.use_megatron_fsdp and get_device_arch_version() >= 10: if 'dp_cp' not in args.high_priority_stream_groups: args.high_priority_stream_groups.append('dp_cp') - if args.expert_model_parallel_size > 1 and 'ep_dp' not in args.high_priority_stream_groups: + if args.expert_model_parallel_size > 1 and 'ep_dp' not in args.high_priority_stream_groups: args.high_priority_stream_groups.append('ep_dp') - - # Derive the internal gtp_weight_remat_size from the user-facing - # --tensor-parallel-num-weight-shards. gtp_weight_remat_size has no CLI flag (it is excluded - # from argument generation), so it is set here as a fresh attribute on args before it is - # consumed below (and in initialize/training, which read args.gtp_weight_remat_size directly). - # Mirrors ModelParallelConfig.__post_init__. - from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards - (args.tensor_parallel_num_weight_shards, args.gtp_weight_remat_size) = ( - resolve_tensor_parallel_weight_shards( - args.tensor_model_parallel_size, - args.tensor_parallel_num_weight_shards, - getattr(args, "gtp_weight_remat_size", 1), - ) - ) - # Same for the expert layers: derive the internal expert_gtp_weight_remat_size from the - # user-facing --expert-tensor-parallel-num-weight-shards (expert_tensor_parallel_size is - # defaulted earlier in validate_args). expert_gtp_weight_remat_size has no CLI flag. - (args.expert_tensor_parallel_num_weight_shards, args.expert_gtp_weight_remat_size) = ( - resolve_tensor_parallel_weight_shards( - args.expert_tensor_parallel_size, - args.expert_tensor_parallel_num_weight_shards, - getattr(args, "expert_gtp_weight_remat_size", 1), - ) - ) - - if args.gtp_weight_remat_size > 1 or args.expert_gtp_weight_remat_size > 1: - if args.fp4 and not args.fp4_param_gather: - raise ValueError( - "GTP (--tensor-parallel-num-weight-shards / " - "--expert-tensor-parallel-num-weight-shards > 1) with --fp4-format requires " - "--fp4-param-gather so NVFP4 weights are all-gathered as native NVFP4." - ) - gtp_weight_remat_size = args.gtp_weight_remat_size - egtp_weight_remat_size = args.expert_gtp_weight_remat_size - if get_device_arch_version() >= 10: - # Setting GTP communication groups for high priority streams for Blackwell and later - # architectures. Assigning high priority to communication streams ensures that - # communication kernels are scheduled with higher priority, minimizing the exposed - # communication when it is overlapped with other computation kernels. - if 'gtp_remat' not in args.high_priority_stream_groups: - args.high_priority_stream_groups.append('gtp_remat') - warn_rank_0("Setting 'gtp_remat' group for high priority streams.") - if ( - egtp_weight_remat_size > 1 - and 'expt_gtp_remat' not in args.high_priority_stream_groups - ): - args.high_priority_stream_groups.append('expt_gtp_remat') - warn_rank_0("Setting 'expt_gtp_remat' group for high priority streams.") - - # Sanity check for 'CUDA_GRAPHS_USE_NODE_PRIORITY'. - if args.cuda_graph_impl != "none": - assert os.environ.get('CUDA_GRAPHS_USE_NODE_PRIORITY') == "1", \ - 'GTP requires CUDA_GRAPHS_USE_NODE_PRIORITY=1 to make sure fine-grained GTP ' \ - 'comms can be well overlapped with GEMMs when CudaGraph is enabled for ' \ - 'Blackwell and later architecture.' - - # Sanity check for 'NCCL_PROTO'. - if os.environ.get('NCCL_PROTO', '').lower() == "simple": - warn_rank_0( - "Generally GTP prefers 'NCCL_PROTO=LL128 or LL' while get 'NCCL_PROTO=simple', " - "force setting NCCL_PROTO=Simple might introduce bad perf." - ) - - assert not args.ddp_average_in_collective, ( - "GTP requires --ddp-average-in-collective off (the default); averaged collectives " - "would need per-buffer 1/gtp_remat scaling." - ) - - assert args.ckpt_format in ('torch', 'torch_dist'), ( - f"GTP supports only --ckpt-format 'torch' (legacy) or 'torch_dist', got " - f"'{args.ckpt_format}'." - ) - assert not ( - getattr(args, 'dist_ckpt_optim_fully_reshardable', False) - and getattr(args, 'distrib_optim_fully_reshardable_mem_efficient', False) - ), ( - "GTP does not support the distributed-optimizer fully-reshardable + " - "mem-efficient checkpoint mode. Disable " - "--distrib-optim-fully-reshardable-mem-efficient (or " - "--dist-ckpt-optim-fully-reshardable)." - ) - - # GTP with the mxfp8 recipe requires --fp8-param-gather: GTP keeps no bf16 weight and - # relies on the optimizer maintaining the fp8 shard (the forward all-gathers fp8 and does - # not re-quantize). Without fp8-param-gather the fp8 forward weight would never be updated. - if getattr(args, 'fp8_recipe', None) == 'mxfp8': - assert getattr(args, 'fp8_param_gather', False), ( - "GTP + mxfp8 requires --fp8-param-gather (the optimizer maintains the fp8 shard; " - "GTP does not keep or re-quantize a bf16 weight)." - ) - # MXFP8 params cannot be mapped into the contiguous param buffer (TE's - # replace_raw_data does not support the MXFP8 tile-scaling layout), so the param - # all-gather must reuse the grad buffer instead. - assert getattr(args, 'reuse_grad_buf_for_mxfp8_param_ag', False), ( - "GTP + mxfp8 + --fp8-param-gather requires --reuse-grad-buf-for-mxfp8-param-ag " - "(MXFP8 params keep their own quantized storage; mapping them into the param " - "buffer via replace_raw_data is unsupported)." - ) - - # GTP symmetric memory registers pools with symmetric=True (NVLS needs symmetric - # windows), which contradicts --disable-symmetric-registration. - if getattr(args, 'gtp_remat_nccl_ub', False) or getattr(args, 'gtp_expert_remat_nccl_ub', False): - assert not getattr(args, 'disable_symmetric_registration', False), ( - "--gtp-remat-nccl-ub/--gtp-expert-remat-nccl-ub require symmetric window registration and " - "cannot be combined with --disable-symmetric-registration." - ) - if getattr(args, 'gtp_remat_reduce_scatter_with_fp32_accumulation', False): - print_rank_0( - "WARNING: --gtp-remat-nccl-ub/--gtp-expert-remat-nccl-ub take precedence over " - "--gtp-remat-reduce-scatter-with-fp32-accumulation on their groups: NVLS " - "symmetric reduce-scatters accumulate in fp32 in-switch " - "(NCCL multimem.ld_reduce .acc::f32)." - ) - # Disable bias gelu fusion if we are disabling bias altogether if not args.add_bias_linear: args.bias_gelu_fusion = False @@ -1623,22 +1669,24 @@ def validate_args(args, defaults={}): args.add_qkv_bias = True if args.qk_clip: - assert is_te_min_version("2.9.0"), \ - '--qk-clip is only supported with TE >= 2.9.0.' - assert 0.0 < args.qk_clip_alpha < 1.0, \ - '--qk-clip-alpha must be between 0.0 and 1.0 when using --qk-clip.' - assert args.qk_clip_threshold > 0, \ - '--qk-clip-threshold must be greater than 0 when using --qk-clip.' + assert is_te_min_version("2.9.0"), '--qk-clip is only supported with TE >= 2.9.0.' + assert ( + 0.0 < args.qk_clip_alpha < 1.0 + ), '--qk-clip-alpha must be between 0.0 and 1.0 when using --qk-clip.' + assert ( + args.qk_clip_threshold > 0 + ), '--qk-clip-threshold must be greater than 0 when using --qk-clip.' # decoupled log max attention logit check if args.log_max_attention_logit: - assert is_te_min_version("2.9.0"), \ - '--log-max-attention-logit is only supported with TE >= 2.9.0.' + assert is_te_min_version( + "2.9.0" + ), '--log-max-attention-logit is only supported with TE >= 2.9.0.' # Legacy RoPE arguments if args.use_rotary_position_embeddings: args.position_embedding_type = 'rope' - if args.position_embedding_type != 'rope': + if args.position_embedding_type not in ('rope', 'mrope'): args.apply_rope_fusion = False # Would just need to add 'NoPE' as a position_embedding_type to support this, but for now @@ -1654,8 +1702,9 @@ def validate_args(args, defaults={}): # MultiModal rotary embeddings arguments if args.position_embedding_type == "mrope": - assert args.mrope_section is not None, \ - '--mrope-section should be set when using --position-embedding-type mrope.' + assert ( + args.mrope_section is not None + ), '--mrope-section should be set when using --position-embedding-type mrope.' # MoE Spec check if args.num_experts == 0: @@ -1665,13 +1714,19 @@ def validate_args(args, defaults={}): warn_rank_0("moe_ffn_hidden_size is not set, using ffn_hidden_size for MoE instead.") # Expert parallelism check - if args.expert_model_parallel_size > 1: - assert args.num_experts is not None, "num_experts must be non None to use expert model parallelism" - assert args.num_experts % args.expert_model_parallel_size == 0, \ - "Number of experts should be a multiple of expert model parallel_size." + if args.expert_model_parallel_size > 1: + assert ( + args.num_experts is not None + ), "num_experts must be non None to use expert model parallelism" + assert ( + args.num_experts % args.expert_model_parallel_size == 0 + ), "Number of experts should be a multiple of expert model parallel_size." # MoE router check - if isinstance(args.moe_router_load_balancing_type, list) and len(args.moe_router_load_balancing_type) == 1: + if ( + isinstance(args.moe_router_load_balancing_type, list) + and len(args.moe_router_load_balancing_type) == 1 + ): args.moe_router_load_balancing_type = args.moe_router_load_balancing_type[0] if isinstance(args.moe_aux_loss_coeff, list) and len(args.moe_aux_loss_coeff) == 1: args.moe_aux_loss_coeff = args.moe_aux_loss_coeff[0] @@ -1679,40 +1734,81 @@ def validate_args(args, defaults={}): # torch_dcp (torch.distributed.checkpoint) checkpointing format checks. if args.ckpt_format == "torch_dcp": assert args.use_torch_fsdp2, "--ckpt-format torch_dcp is only tested with FSDP." - assert args.tensor_model_parallel_size <= 1, \ - "--ckpt-format torch_dcp is not tested with megatron tensor parallelism." - assert args.pipeline_model_parallel_size <= 1, \ - "--ckpt-format torch_dcp is not tested with megatron pipeline parallelism." + assert ( + args.tensor_model_parallel_size <= 1 + ), "--ckpt-format torch_dcp is not tested with megatron tensor parallelism." + assert ( + args.pipeline_model_parallel_size <= 1 + ), "--ckpt-format torch_dcp is not tested with megatron pipeline parallelism." # fsdp_dtensor checkpointing format checks. if args.ckpt_format == "fsdp_dtensor": - assert args.use_megatron_fsdp, "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP." - - # Scheduler-name and max-seqlen validation live in - # ModelParallelConfig.__post_init__; only the buffer-size check stays here - # because seq_length is not a core config field. The None case for - # max_seqlen_per_dp_cp_rank is rejected by the config check. - if args.sequence_packing_scheduler is not None: - args.variable_seq_lengths = True - # Packed microbatches carry different numbers of valid tokens, so the - # default per-microbatch loss averaging would weight tokens unevenly - # depending on how samples happened to be packed (same reasoning as - # the hybrid-context-parallel check above). - assert args.calculate_per_token_loss, ( - 'Sequence packing must be used with --calculate-per-token-loss' + assert ( + args.use_megatron_fsdp + ), "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP." + + # --use-varlen-dataset: independent of --sft. Cannot be combined with --sft + # because they are mutually-exclusive top-level dataset selectors that both + # drive the packed-sequence (THD) path. + if args.use_varlen_dataset: + assert not args.sft, ( + "--use-varlen-dataset and --sft are mutually exclusive; both " + "select the packed-sequence dataset family. Pick one." ) - if args.max_seqlen_per_dp_cp_rank is not None: - total_cp_ranks = args.context_parallel_size - assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, ( - f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) ' - f'must be >= single sequence max length ({args.seq_length})' + if args.varlen_sbhd_validation: + # ``--dynamic-context-parallel`` ⊥ ``--varlen-sbhd-validation`` is + # checked in ``GPTDatasetConfig.__post_init__``; only the + # scheduler check stays here, since ``sequence_packing_scheduler`` + # is a training-framework flag not stored on the dataset config. + assert args.sequence_packing_scheduler is None, ( + "--varlen-sbhd-validation does not use a sequence packing " + "scheduler; drop --sequence-packing-scheduler." ) + # SBHD validation is a real-data numerical-reference path only; + # MockVarlenDataset does not implement it. + assert not args.mock_data, ( + "--varlen-sbhd-validation is not supported with --mock-data; " + "SBHD validation requires a real dataset." + ) + else: + # VarlenDataset emits one unpacked sample per __getitem__; it + # relies on an upstream packing scheduler to group variable-length + # samples into THD batches. Auto-pick a default scheduler when + # the user did not request one explicitly: + # * ``--dynamic-context-parallel`` is already wired to + # ``default_dynamic_cp`` upstream (see the dynamic-cp block + # earlier in ``validate_args``). + # * Otherwise fall back to ``dp_balanced`` (static packing). + if args.sequence_packing_scheduler is None: + args.sequence_packing_scheduler = 'dp_balanced' + + # Packed-sequence buffer-size check. Placed after all scheduler auto-select + # logic (dynamic-cp and --use-varlen-dataset both set the scheduler above) + # so it validates the final resolved scheduler; the varlen path picks its + # default after the earlier generic validation has run. + if args.sequence_packing_scheduler is not None: + if args.sequence_packing_scheduler == 'dp_balanced': + total_cp_ranks = args.context_parallel_size + else: + total_cp_ranks = args.data_parallel_size * args.context_parallel_size + if args.max_seqlen_per_dp_cp_rank is None: + args.max_seqlen_per_dp_cp_rank = getattr(args, "max_seqlen_per_cp_rank", None) + if args.max_seqlen_per_dp_cp_rank is None: + args.max_seqlen_per_dp_cp_rank = ( + args.seq_length + total_cp_ranks - 1 + ) // total_cp_ranks + assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, ( + f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) ' + f'must be >= single sequence max length ({args.seq_length})' + ) # Data blend checks - assert args.mock_data + \ - bool(args.data_path) + \ - any([args.train_data_path, args.valid_data_path, args.test_data_path]) \ - <= 1, "A single data source must be provided in training mode, else None" + assert ( + args.mock_data + + bool(args.data_path) + + any([args.train_data_path, args.valid_data_path, args.test_data_path]) + <= 1 + ), "A single data source must be provided in training mode, else None" if args.fim_data: extra_tokens = [ @@ -1725,11 +1821,11 @@ def validate_args(args, defaults={}): assert not args.mock_data, "Mock dataset is not supported with FIM dataset." assert args.fim_rate, "--fim-rate should be specified." assert args.fim_spm_rate, "--fim-spm-rate should be specified." - assert all(token is not None for token in extra_tokens), "FIM extra tokens should be specified." + assert all( + token is not None for token in extra_tokens + ), "FIM extra tokens should be specified." - assert not ( - args.cross_entropy_loss_fusion and args.cross_entropy_fusion_impl == 'te' - ), ( + assert not (args.cross_entropy_loss_fusion and args.cross_entropy_fusion_impl == 'te'), ( "Transformer Engine cross entropy loss fusion is disabled due to stability issues. " "Use --cross-entropy-fusion-impl native, or omit --cross-entropy-loss-fusion." ) @@ -1769,17 +1865,62 @@ def validate_args(args, defaults={}): args.use_layer_wise_distributed_optimizer = True args.use_distributed_optimizer = False + if ( + args.optimizer == 'muon' + and (args.chunked_optimizer_state_offload or args.offload_optimizer_states) + and args.optimizer_state_offload_fraction > 0.0 + ): + assert args.use_layer_wise_distributed_optimizer, ( + "Muon optimizer state offload requires the LayerWise distributed optimizer " + "(--use-distributed-optimizer) so non-Muon parameter groups are routed through " + "a sibling DistributedOptimizer" + ) + assert not args.use_torch_fsdp2, "Emerging optimizer does not support Torch-FSDP2 for now." - assert not args.use_megatron_fsdp, "Emerging optimizer does not support Megatron-FSDP for now." - assert args.ckpt_format in ["torch", "torch_dist"], "Emerging optimizer supports torch and torch_dist checkpoint format." - - assert not ( - args.use_layer_wise_distributed_optimizer and args.moe_single_grouped_weight - ), ( - "The LayerWise distributed optimizer does not support --moe-single-grouped-weight: " - "Muon semantics for a single grouped [E, N, K] expert weight are not defined. " - "Disable --moe-single-grouped-weight or use Adam/DistributedOptimizer." - ) + assert ( + not args.use_megatron_fsdp + ), "Emerging optimizer does not support Megatron-FSDP for now." + assert args.ckpt_format in [ + "torch", + "torch_dist", + ], "Emerging optimizer supports torch and torch_dist checkpoint format." + + if args.use_layer_wise_distributed_optimizer: + if not args.use_layer_wise_param_layout: + # Decoupled compact LayerWise: fp8 parameter gather is supported via the FP8-aware + # whole-param all-gather. Only mxfp8/blockwise (fp4 out of scope); mxfp8 needs + # reuse_grad_buf. fp4 is rejected unconditionally -- the LayerWise gather routes + # buckets by is_float8tensor, so an NVFP4 param would silently take the raw + # flatten path. + assert not getattr(args, 'fp4_param_gather', False), ( + "Decoupled compact LayerWise DDP layout supports fp8 parameter gather only " + "(mxfp8 or blockwise); fp4_param_gather is out of scope." + ) + if args.fp8_param_gather: + assert args.fp8_recipe in ('mxfp8', 'blockwise'), ( + "fp8 parameter gather on the decoupled compact LayerWise DDP layout requires " + f"fp8_recipe in {{'mxfp8', 'blockwise'}}; got {args.fp8_recipe!r}." + ) + if args.fp8_recipe == 'mxfp8': + assert args.reuse_grad_buf_for_mxfp8_param_ag, ( + "mxfp8 + --fp8-param-gather on the decoupled compact LayerWise DDP layout " + "requires --reuse-grad-buf-for-mxfp8-param-ag (or use fp8_recipe='blockwise')." + ) + assert args.num_distributed_optimizer_instances == 1, ( + "the decoupled compact LayerWise DDP layout (the default; pass " + "--use-layer-wise-param-layout for the padded layout) requires " + "num_distributed_optimizer_instances == 1: the non-DistOpt LayerWise (Muon) buffers " + "only all-reduce within a single optimizer instance, so partial DistOpt (>1 " + "instance) would under-reduce Muon gradients across the full data-parallel domain." + ) + else: + # Padded LayerWise param layout: fp8/fp4 parameter gather is not supported here. + assert not args.fp8_param_gather and not getattr(args, 'fp4_param_gather', False), ( + "Layer-wise (Muon) distributed optimizer with the padded param layout does not " + "support FP8/FP4 parameter gather. Use the default compact decoupled layout (do " + "not pass --use-layer-wise-param-layout) for fp8 parameter gather, or " + "fp8_param_gather=False." + ) # Make sure all functionality that requires Gloo process groups is disabled. if not args.use_gloo_process_groups: @@ -1790,11 +1931,14 @@ def validate_args(args, defaults={}): assert args.use_dist_ckpt if args.dist_ckpt_optim_fully_reshardable: - assert not args.distrib_optim_fully_reshardable_mem_efficient, \ - '--distrib-optim-fully-reshardable-mem-efficient requires -enable-gloo-process-groups' + assert ( + not args.distrib_optim_fully_reshardable_mem_efficient + ), '--distrib-optim-fully-reshardable-mem-efficient requires -enable-gloo-process-groups' if args.fake_process_group: - assert args.moe_token_dispatcher_type != "flex", "Fake process group is not supported with flex token dispatcher." + assert ( + args.moe_token_dispatcher_type != "flex" + ), "Fake process group is not supported with flex token dispatcher." # Disable nan check for fake process group args.check_for_nan_in_loss_and_grad = False warn_rank_0('check_for_nan_in_loss_and_grad is set to False for fake process group.') @@ -1804,23 +1948,31 @@ def validate_args(args, defaults={}): # Checkpointing if args.ckpt_fully_parallel_save_deprecated and args.rank == 0: - print('--ckpt-fully-parallel-save flag is deprecated and has no effect.' - ' Use --no-ckpt-fully-parallel-save to disable parallel save.') + print( + '--ckpt-fully-parallel-save flag is deprecated and has no effect.' + ' Use --no-ckpt-fully-parallel-save to disable parallel save.' + ) if ( args.use_dist_ckpt and not args.ckpt_fully_parallel_save and args.use_distributed_optimizer and args.rank == 0 ): - print('Warning: With non-parallel ckpt save and DistributedOptimizer,' - ' it will be impossible to resume training with different parallelism.' - ' Consider removing flag --no-ckpt-fully-parallel-save.') + print( + 'Warning: With non-parallel ckpt save and DistributedOptimizer,' + ' it will be impossible to resume training with different parallelism.' + ' Consider removing flag --no-ckpt-fully-parallel-save.' + ) if args.use_dist_ckpt_deprecated and args.rank == 0: - print('--use-dist-ckpt is deprecated and has no effect.' - ' Use --ckpt-format to select the checkpoint format.') + print( + '--use-dist-ckpt is deprecated and has no effect.' + ' Use --ckpt-format to select the checkpoint format.' + ) if args.dist_ckpt_format_deprecated and args.rank == 0: - print('--dist-ckpt-format is deprecated and has no effect.' - ' Use --ckpt-format to select the checkpoint format.') + print( + '--dist-ckpt-format is deprecated and has no effect.' + ' Use --ckpt-format to select the checkpoint format.' + ) if args.use_dist_ckpt and args.ckpt_fully_parallel_load: if args.ckpt_fully_parallel_load_exchange_algo != "broadcast": @@ -1829,25 +1981,6 @@ def validate_args(args, defaults={}): "Other algorithms cannot guarantee numerical stability yet." ) - # The PG-distribution cache stores one file per parallelization group, keyed by - # that group's minimum global rank, and each file holds both the save and the - # load distribution. If save and load used different parallelization groups, - # two distinct groups could map to the same file (they can share a minimum - # global rank), so one side would silently read a distribution computed for the - # other group's layout. Require both sides to use the same group. - if args.ckpt_pg_tensors_cache_path is not None: - assert ( - args.ckpt_fully_parallel_save_process_group - == args.ckpt_fully_parallel_load_process_group - ), ( - "--ckpt-pg-tensors-cache-path requires --ckpt-fully-parallel-save-process-group and " - "--ckpt-fully-parallel-load-process-group to be identical, but got " - f"save='{args.ckpt_fully_parallel_save_process_group}' and " - f"load='{args.ckpt_fully_parallel_load_process_group}'. The cache is keyed by the " - "parallelization group, so mixing groups can read a distribution computed for a " - "different layout." - ) - if args.load_main_params_from_ckpt: assert args.no_load_optim, '--load-main-params-from-ckpt must be used with --no-load-optim.' @@ -1862,39 +1995,11 @@ def validate_args(args, defaults={}): if not args.async_save: args.async_strategy = "mcore" - if args.logits_save_dir is not None: - assert args.logits_save_top_k is not None, '--logits-save-top-k is required when --logits-save-dir is set.' - assert args.async_save, ( - '--logits-save-dir requires --async-save (and --use-persistent-ckpt-worker). ' - 'Logits are flushed as an async request in the checkpoint queue.' - ) - if not args.freeze_all_layers: - warn_rank_0( - '--logits-save-dir without --freeze-all-layers: the LM loss is still computed and ' - 'gradients will update the model while logits are dumped. This is intended only ' - 'when dumping logits during active training; for a frozen-teacher dump pass ' - '--freeze-all-layers.' - ) - - if args.freeze_all_layers: - if args.use_distributed_optimizer: - warn_rank_0( - '--freeze-all-layers incompatible with use_distributed_optimizer. Disabling use_distributed_optimizer.' - ) - args.use_distributed_optimizer = False - if args.overlap_param_gather: - warn_rank_0( - '--freeze-all-layers incompatible with overlap_param_gather. Disabling overlap_param_gather.' - ) - args.overlap_param_gather = False - - if args.override_ckpt_iteration is not None: - assert not args.finetune, "Cannot override checkpoint iteration together with finetune flag." - # Inference args if args.inference_batch_times_seqlen_threshold > -1: - assert args.pipeline_model_parallel_size > 1, \ - "--inference-batch-times-seqlen-threshold requires setting --pipeline-model-parallel-size > 1." + assert ( + args.pipeline_model_parallel_size > 1 + ), "--inference-batch-times-seqlen-threshold requires setting --pipeline-model-parallel-size > 1." assert ( args.cuda_graph_impl == "none" ), "Pipeline-parallel microbatched inference is incompatible with CUDA graphs" @@ -1915,10 +2020,16 @@ def validate_args(args, defaults={}): if args.skip_train and not args.perform_rl_step and not args.no_load_optim: args.no_load_optim = True warn_rank_0('enabling --no-load-optim when skipping training.') - if args.skip_train and args.perform_rl_step and args.no_load_optim and args.rl_offload_optimizer_during_inference: - assert False, \ - '--no-load-optim with --skip-train --perform-rl-step skips the optimizer; ' \ + if ( + args.skip_train + and args.perform_rl_step + and args.no_load_optim + and args.rl_offload_optimizer_during_inference + ): + assert False, ( + '--no-load-optim with --skip-train --perform-rl-step skips the optimizer; ' '--rl-offload-optimizer-during-inference is incompatible (no optimizer to offload).' + ) # Optimizer CPU offload check if args.optimizer_cpu_offload: @@ -1931,22 +2042,68 @@ def validate_args(args, defaults={}): "must be used in conjunction with `--fp8-recipe delayed`." ) + # Optimizer configuration normalizes the deprecated spelling. Include it here so + # cross-config validation still runs before that configuration is constructed. + if ( + args.chunked_optimizer_state_offload or args.offload_optimizer_states + ) and args.optimizer_state_offload_fraction > 0.0: + if args.optimizer_state_offload_chunk_size_mb == 0: + warn_rank_0( + "optimizer state offload is enabled with chunk size 0, so all " + "selected optimizer tensor state is temporarily restored together. Set " + "--optimizer-state-offload-chunk-size-mb to a positive value to bound the " + "tensor-state GPU window. Selected master weights always use one full " + "restore window regardless of chunk size." + ) + assert ( + not args.use_megatron_fsdp + ), "chunked optimizer state offload does not support Megatron-FSDP" + if not args.no_save_optim or not args.no_load_optim: + assert args.ckpt_format == 'torch_dist', ( + "chunked optimizer state offload requires --ckpt-format torch_dist when " + "optimizer state is saved or loaded because its sharded load hooks preserve " + "CPU canonical storage" + ) + if not args.no_save_optim: + assert not args.async_save, ( + "chunked optimizer state offload does not support --async-save when optimizer " + "state is saved because the background writer can retain tensors backed by " + "reusable pinned CPU buffers" + ) + assert ( + args.cuda_graph_impl != 'full_iteration' + ), "chunked optimizer state offload does not support full-iteration CUDA graphs" + if args.optimizer == 'muon' and args.fp8_param_gather and args.overlap_param_gather: + warn_rank_0( + "Muon compact FP8 parameter gather with chunked optimizer state offload " + "forces LayerWise-owned pre-forward parameter buckets to complete before " + "masters are offloaded; sibling DistributedOptimizer gathers and " + "optimizer-state transfers remain overlapped." + ) + if args.non_persistent_ckpt_type == "local": - assert args.non_persistent_local_ckpt_dir is not None, "Tried to use local checkpointing without specifying --local-ckpt-dir!" + assert ( + args.non_persistent_local_ckpt_dir is not None + ), "Tried to use local checkpointing without specifying --local-ckpt-dir!" if args.replication: - assert args.replication_jump is not None, "--replication requires the value of --replication-jump!" - assert args.non_persistent_ckpt_type == "local", f"--replication requires args.non_persistent_ckpt_type == 'local', but got: {args.non_persistent_ckpt_type}" + assert ( + args.replication_jump is not None + ), "--replication requires the value of --replication-jump!" + assert ( + args.non_persistent_ckpt_type == "local" + ), f"--replication requires args.non_persistent_ckpt_type == 'local', but got: {args.non_persistent_ckpt_type}" elif args.replication_jump: warn_rank_0("--replication-jump was specified despite not using replication. Ignoring.") args.replication_jump = None if args.delay_wgrad_compute: - assert args.transformer_impl == 'transformer_engine', \ - "Delaying wgrad compute is only supported with transformer_engine implementation" + assert ( + args.transformer_impl == 'transformer_engine' + ), "Delaying wgrad compute is only supported with transformer_engine implementation" if args.overlap_grad_reduce: - assert is_te_min_version("2.8.0"), ( - "overlap_grad_reduce is only supported with TE >= 2.8.0 when enabling delay_wgrad_compute" - ) + assert is_te_min_version( + "2.8.0" + ), "overlap_grad_reduce is only supported with TE >= 2.8.0 when enabling delay_wgrad_compute" wgrad_in_graph_scope = CudaGraphModule.attn in args.cuda_graph_modules or ( CudaGraphModule.moe_router in args.cuda_graph_modules and args.moe_shared_expert_intermediate_size is not None @@ -1973,11 +2130,13 @@ def validate_args(args, defaults={}): ) if args.fine_grained_activation_offloading: - assert args.transformer_impl == 'transformer_engine', \ - "Fine-grained activation offloading is only supported with transformer_engine implementation" + assert ( + args.transformer_impl == 'transformer_engine' + ), "Fine-grained activation offloading is only supported with transformer_engine implementation" if is_te_min_version("2.10.0"): - assert os.getenv("NVTE_CPU_OFFLOAD_V1", "0") == "1", \ - "For fine-grained activation offloading with TE >= 2.10.0, NVTE_CPU_OFFLOAD_V1 should be set to 1 to avoid offloading weights." + assert ( + os.getenv("NVTE_CPU_OFFLOAD_V1", "0") == "1" + ), "For fine-grained activation offloading with TE >= 2.10.0, NVTE_CPU_OFFLOAD_V1 should be set to 1 to avoid offloading weights." if args.mtp_num_layers: # MTP is compatible with position embedding types that use position_ids. @@ -2011,19 +2170,23 @@ def validate_args(args, defaults={}): ), '--cuda-graph-modules must be empty when --cuda-graph-impl=full_iteration.' if args.multi_latent_attention: - assert not args.group_query_attention, "Group query attention is mutually exclusive with multi latent attention." + assert ( + not args.group_query_attention + ), "Group query attention is mutually exclusive with multi latent attention." if args.mla_down_proj_fusion: - assert args.multi_latent_attention, "--mla-down-proj-fusion requires --multi-latent-attention" - - assert ( - not args.moe_use_norm_before_up_proj or args.moe_latent_size is not None - ), "--moe-use-norm-before-up-proj requires --moe-latent-size to be set." + assert ( + args.multi_latent_attention + ), "--mla-down-proj-fusion requires --multi-latent-attention" # MoE latent projections if args.moe_latent_size is not None: - assert args.moe_latent_size > 0, "MoE latent projection dimension has to be greater than zero." - assert args.num_experts is not None, "MoE latent projections are applicable only for MoE models." + assert ( + args.moe_latent_size > 0 + ), "MoE latent projection dimension has to be greater than zero." + assert ( + args.num_experts is not None + ), "MoE latent projections are applicable only for MoE models." # Print arguments. _print_args("arguments", args) @@ -2034,6 +2197,7 @@ def validate_args(args, defaults={}): def _print_args(title, args): """Print arguments.""" from megatron.training.utils import is_rank0 + if is_rank0(): print(f'------------------------ {title} ------------------------', flush=True) str_list = [] @@ -2049,296 +2213,524 @@ def _check_arg_is_not_none(args, arg): assert getattr(args, arg) is not None, '{} argument is None'.format(arg) +def core_transformer_config_from_args(args, config_class=None): + + # Config class. + config_class = config_class or TransformerConfig + + if args.multi_latent_attention: + config_class = MLATransformerConfig + + if args.heterogeneous_layers_config_path is not None: + assert ( + not args.multi_latent_attention + ), "Multi latent attention with heterogeneous layers is not supported." + config_class = HeterogeneousTransformerConfig + + # Translate args to core transformer configuration + kw_args = {} + for f in dataclasses.fields(config_class): + if hasattr(args, f.name): + kw_args[f.name] = getattr(args, f.name) + kw_args['persist_layer_norm'] = not args.no_persist_layer_norm + kw_args['deallocate_pipeline_outputs'] = True + kw_args['pipeline_dtype'] = args.params_dtype + kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm + kw_args['num_moe_experts'] = args.num_experts + kw_args['actual_vocab_size'] = args.padded_vocab_size + kw_args['rotary_interleaved'] = args.rotary_interleaved + kw_args['num_layers_in_first_pipeline_stage'] = args.decoder_first_pipeline_num_layers + kw_args['num_layers_in_last_pipeline_stage'] = args.decoder_last_pipeline_num_layers + kw_args['fp8_param'] = args.fp8_param_gather + kw_args['fp4_param'] = args.fp4_param_gather + if args.swiglu: + kw_args['activation_func'] = F.silu + kw_args['gated_linear_unit'] = True + kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion + else: + kw_args['bias_activation_fusion'] = args.bias_gelu_fusion + if args.squared_relu: + assert not args.swiglu + kw_args['activation_func'] = squared_relu + elif args.quick_geglu: + assert not args.swiglu + kw_args['gated_linear_unit'] = True + kw_args['activation_func'] = quick_gelu + if args.init_method_xavier_uniform: + kw_args['init_method'] = torch.nn.init.xavier_uniform_ + kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_ + if args.group_query_attention: + kw_args['num_query_groups'] = args.num_query_groups + else: + kw_args['num_query_groups'] = None + kw_args['config_logger_dir'] = args.config_logger_dir + if args.rope_type is None: + # Pop 'rope_type' to let the config class use the default value. + kw_args.pop('rope_type', None) + else: + assert ( + args.multi_latent_attention or args.rope_type == 'rope' + ), f'Common attention only support rope_type="rope", but got {args.rope_type}.' + + if len(args.cp_comm_type) == 1: + kw_args['cp_comm_type'] = args.cp_comm_type[0] + if args.hybrid_layer_pattern is not None: + kw_args['is_hybrid_model'] = True + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + + _pat = args.hybrid_layer_pattern + _has_kda = Symbols.KDA in _pat + _has_dsv4_csa = (Symbols.CSA in _pat) or (Symbols.HCA in _pat) or (Symbols.WINDOW in _pat) + _has_dsa = Symbols.DS_ATTENTION in _pat + if getattr(args, 'experimental_attention_variant', None) is None: + # 'C'/'H'/'W' run the DSv4 CompressedSparseAttention (CSA/HCA/window-only), which + # requires the full dsv4_hybrid contract (MLA, TP==1, no qk_clip, + # qk_head_dim/kv_lora_rank derivation). Set the variant so transformer_config runs + # that validation+derivation rather than silently skipping it. 'D' alone stays legacy + # DSv3 'dsa'. An explicit --experimental-attention-variant is always respected. + if _has_dsv4_csa: + kw_args['experimental_attention_variant'] = 'dsv4_hybrid' + elif _has_dsa: + kw_args['experimental_attention_variant'] = 'dsa' + elif _has_kda: + kw_args['experimental_attention_variant'] = 'kda' + # Normalize compact and legacy-padded ratios through the shared migration helper. + _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, _pat) + + _resolve_dsa_kernel_backend_cli_default(args, kw_args) + + kw_args['inference_sampling_seed'] = args.seed + + # handle quantization config + # NOTE: Kitchen arguments are only added to the namespace when + # Kitchen library is available. + if hasattr(args, "kitchen_config_file") and args.kitchen_config_file is not None: + kw_args['use_kitchen'] = True + kw_args['quant_recipe'] = load_quantization_recipe(args.kitchen_config_file) + elif hasattr(args, 'kitchen_recipe_number') and args.kitchen_recipe_number is not None: + kw_args['use_kitchen'] = True + kw_args['quant_recipe'] = kitchen_quantization_recipe_config(args.kitchen_recipe_number) + + kw_args['moe_latent_size'] = args.moe_latent_size + + if args.te_precision_config_file: + assert not 'quant_recipe' in kw_args, "Quantization recipe already configured." + # TODO(kwyss): Prohibit fp8_params or fp4_params with this flexibility + kw_args['quant_recipe'] = load_quantization_recipe(args.te_precision_config_file) + + if hasattr(args, "use_kitchen_attention"): + kw_args['use_kitchen_attention'] = args.use_kitchen_attention + if hasattr(args, "kitchen_attention_backend"): + kw_args['kitchen_attention_backend'] = args.kitchen_attention_backend + + # Return config. + return config_class(**kw_args) + + def _add_transformer_engine_args(parser): group = parser.add_argument_group(title='Transformer-Engine') # FP4 related arguments - group.add_argument('--fp4-param-gather', action='store_true', - help='Keep the compute param in fp4 (do not use any other intermediate ' - 'dtype) and perform the param all-gather in fp4.') + group.add_argument( + '--fp4-param-gather', + action='store_true', + help='Keep the compute param in fp4 (do not use any other intermediate ' + 'dtype) and perform the param all-gather in fp4.', + ) # FP8 related arguments - group.add_argument('--fp8-param-gather', action='store_true', - help='Keep the compute param in fp8 (do not use any other intermediate ' - 'dtype) and perform the param all-gather in fp8.') + group.add_argument( + '--fp8-param-gather', + action='store_true', + help='Keep the compute param in fp8 (do not use any other intermediate ' + 'dtype) and perform the param all-gather in fp8.', + ) # TE precision config file - group.add_argument('--te-precision-config-file', default=None, - help='Configuration file to select per-module precision overrides. ' - 'See TransformerEngineMixedPrecision.md') + group.add_argument( + '--te-precision-config-file', + default=None, + help='Configuration file to select per-module precision overrides. ' + 'See TransformerEngineMixedPrecision.md', + ) return parser + def _add_inference_args(parser): group = parser.add_argument_group(title='inference') - group.add_argument('--inference-batch-times-seqlen-threshold', - type=int, default=-1, - help='If (batch-size * sequence-length) is smaller than this threshold' - 'then batches will not be split up for pipelining.' - 'Requires setting --pipeline-model-parallel-size > 1.' - 'Setting this to -1 indicates that batch pipelining is not used.') - group.add_argument('--max-tokens-to-oom', - type=int, default=12000, - help='Maximum number of tokens during inference' - 'tokens here is # in prompt + # to generate' - 'Allows us to throw an error before OOM crashes server') - group.add_argument('--output-bert-embeddings', action='store_true', - help='Output Bert embeddings (via mean pooling) from ' - 'model, rather than its binary head output or entire ' - 'hidden batch.') - group.add_argument('--bert-embedder-type', default="megatron", - choices=["megatron", "huggingface"], - help='Select either Megatron or Huggingface as the ' - 'Bert embedder.') - group.add_argument('--cuda-graph-scope', nargs='+', type=_parse_cuda_graph_modules_arg, - default=None, dest='cuda_graph_scope_deprecated', - help=argparse.SUPPRESS) # hidden; use --cuda-graph-modules instead - group.add_argument('--cuda-graph-modules', nargs='+', type=_parse_cuda_graph_modules_arg, default=[], - help='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 validate_args. The deprecated values ' - '"full_iteration" and "full_iteration_inference" are also accepted and migrated ' - 'to the new API in validate_args.') - group.add_argument('--use-legacy-static-engine', action='store_true', default=False, - help='Use legacy static engine. (Current static engine uses dynamic engine under the hood)', - dest='use_legacy_static_engine') - group.add_argument('--inference-max-requests', type=int, default=8, - help='Maximum number of requests for inference.', - dest='inference_max_requests') - group.add_argument('--inference-max-seq-length', type=int, default=2560, - help='Maximum sequence length expected for inference (prefill + decode).', - dest='inference_max_seq_length') - group.add_argument('--inference-dynamic-batching', - action='store_true', default=False, - help='Enable dynamic batching mode.') - group.add_argument('--inference-dynamic-batching-buffer-size-gb', - type=float, default=40., - help='On-GPU portion of the shared KV cache block pool. ' - 'The total amount of memory allocated for the KV cache ' - '(CPU + GPU memory) depends on the value set for the ' - 'unified virtual memory (UVM) level (via ' - '`--inference-dynamic-batching-unified-memory-level`).' - 'If the UVM level is 0, then only GPU memory is used and ' - 'the total memory equals `buffer_size_gb`. If the UVM ' - 'level is 1, then additional memory is utilized on the ' - 'CPU and the total memory equals `buffer_size_gb + ' - 'paused_buffer_size_gb`.') - group.add_argument('--inference-dynamic-batching-paused-buffer-size-gb', - type=float, default=None, - help='Memory used to derive the paused-request block retention ' - 'budget. This does not reserve blocks from active requests: ' - 'active requests may use the entire shared pool of usable KV ' - 'cache blocks. Under allocation pressure, paused requests ' - 'retain blocks only within this budget and excess paused ' - 'requests may be evicted.') - group.add_argument('--inference-dynamic-batching-mamba-memory-ratio', type=float, default=None, - help='Percentage of memory buffer to allocate for Mamba states. ' - 'If not specified, allocates Mamba state tensors for each KV cache block. ' - 'Only used for hybrid models.') - group.add_argument('--inference-dynamic-batching-block-size', - type=int, default=256, - help='KV cache block size. ' - 'It should be a multiple of 256') - group.add_argument('--inference-dynamic-batching-max-requests', - type=int, default=None, - help='Override the inference context\'s `max_requests`. ' - 'By default, `max_requests` is set to the number of ' - 'blocks in the context\'s memory buffer.') - group.add_argument('--inference-dynamic-batching-max-tokens', - type=int, default=None, - help='Override the inference context\'s default `max_tokens`.') - group.add_argument('--inference-dynamic-batching-num-cuda-graphs', - type=int, default=16, - help='Maximum number of cuda graphs to capture, where the ' - 'cuda graph batch sizes range from 1 to `max_requests`. ' - '(See `dynamic_context.py` for details on how ' - '`max_requests` is computed). Due to rounding, the actual ' - 'number of cuda graphs may not equal this argument.' - 'The user can also pass -1, in which case we automatically determine the number of graphs ' \ - 'to capture based on the `max_requests`.') - group.add_argument('--inference-dynamic-batching-track-paused-request-events', - action='store_true', - help='Track paused request ids by adding \'paused\' events ' - 'to each request\'s event history. This has a very minor ' - 'impact on latency.') - group.add_argument('--inference-dynamic-batching-track-generated-token-events', - action='store_true', - help='Track per-token events with timestamps for each generated token. ' - 'When enabled, each generated token creates a GENERATED_TOKEN event ' - 'with a timestamp, useful for per-token latency analysis.') - group.add_argument('--decode-only-cuda-graphs', - action='store_true', default=False, - help='Only use cuda graphs for decode-only steps, not prefill and mixed steps.') - group.add_argument('--inference-cuda-graph-all-prefills', - action='store_true', default=False, - help='Extend prefill/mixed CUDA graph capture up to `max_tokens`. ' - 'By default, all graphs are limited by the decode limit of ' - '`max_requests * (num_speculative_tokens + 1)`.') - group.add_argument('--inference-cuda-graph-max-tokens', type=int, default=512, - dest='inference_cuda_graph_max_tokens', - help='Token ceiling for the largest captured prefill/mixed CUDA ' - 'graph (default: 512). Clamped to at least the decode limit ' - '`max_requests * (num_speculative_tokens + 1)` and at most ' - '`max_tokens`. Ignored when --inference-cuda-graph-all-prefills ' - 'is set (which extends capture to the full `max_tokens`).') - group.add_argument('--inference-dynamic-batching-unified-memory-level', - type=int, default=0, choices=[0, 1], - help='Set unified memory usage within the dynamic ' - 'inference context. The levels are: 0) no unified memory, ' - '1) allocate `memory_buffer` in unified memory. ' - 'Eventually, additional levels will be included to ' - 'control other tensors within the context.') - group.add_argument('--enable-chunked-prefill', dest='enable_chunked_prefill', - action='store_true', default=False, - help="Enable chunked prefill (disabled by default)") - group.add_argument('--num-speculative-tokens', type=int, default=0, - help='Number of speculative tokens generated during decode') - group.add_argument('--inference-dynamic-batching-prefix-caching', - dest='inference_dynamic_batching_enable_prefix_caching', - action=argparse.BooleanOptionalAction, - default=False, - help='Enable/disable prefix caching for dynamic batching inference. ' - 'When disabled, KV cache blocks cannot be shared between ' - 'requests with identical prompt prefixes.') - group.add_argument('--inference-dynamic-batching-prefix-caching-eviction-policy', - type=str, default='ref_zero', - choices=['ref_zero', 'lru'], - dest='inference_dynamic_batching_prefix_caching_eviction_policy', - help='Eviction policy for prefix caching blocks. ' - '"ref_zero" (default) immediately returns blocks to the ' - 'free pool when ref_count hits 0. "lru" keeps blocks ' - 'cached and evicts via LRU only when space is needed.') - group.add_argument('--inference-dynamic-batching-prefix-caching-coordinator-policy', - type=str, default='load_balanced', - choices=['longest_prefix', 'first_prefix_block', 'load_balanced'], - dest='inference_dynamic_batching_prefix_caching_coordinator_policy', - help='Coordinator routing policy for prefix caching. ' - '"load_balanced" (default) routes to the rank with the fewest ' - 'in-flight requests, ignoring prefix affinity. ' - '"first_prefix_block" routes based on the first block hash only. ' - '"longest_prefix" routes to the rank with the longest matching ' - 'prefix. "first_prefix_block" and "longest_prefix" both combine ' - 'prefix affinity with load balancing and fall back to ' - 'load-balanced routing when prefix caching is disabled or no ' - 'prefix match exists.') - group.add_argument('--inference-dynamic-batching-prefix-caching-routing-alpha', - type=float, default=0.5, - dest='inference_dynamic_batching_prefix_caching_routing_alpha', - help='Weight for prefix-aware routing score: ' - 'score = alpha * match + (1 - alpha) * normalized_load. ' - 'Higher alpha favors prefix cache hits; lower alpha ' - 'favors load balance. Default: 0.5.') - group.add_argument('--inference-dynamic-batching-media-cache-coordinator-policy', - type=str, default='affinity', - choices=['affinity', 'load_balanced'], - dest='inference_dynamic_batching_media_cache_coordinator_policy', - help='Coordinator routing policy for media caching. ' - '"affinity" prefers a rank assigned the same media; ' - '"load_balanced" ignores standalone media affinity.') - group.add_argument('--inference-dynamic-batching-media-cache-routing-weight', - type=float, default=1.0, - dest='inference_dynamic_batching_media_cache_routing_weight', - help='Media-cache hit weight in equivalent compact-prompt blocks. ' - 'Default: 1.0.') - group.add_argument('--inference-dynamic-batching-vision-embedding-cache-max-bytes', - type=int, default=0, - dest='inference_dynamic_batching_vision_embedding_cache_max_bytes', - help='Maximum GPU bytes retained per engine for reusable vision ' - 'embeddings. Zero disables the cache. Default: 0.') - group.add_argument('--inference-dynamic-batching-allow-stale-multimodal-embeddings', - action='store_true', - dest='inference_dynamic_batching_allow_stale_multimodal_embeddings', - help='Allow request-local and cached multimodal embeddings to survive ' - 'suspend/resume and generation-epoch changes. Use only when model ' - 'weights do not change across these boundaries.') - group.add_argument('--inference-dynamic-batching-prefix-caching-mamba-gb', - type=float, default=None, - dest='inference_dynamic_batching_prefix_caching_mamba_gb', - help='GPU memory budget (in GB) for the Mamba state cache ' - 'used by prefix caching on hybrid models. When set, Mamba ' - 'states at block boundaries are cached for reuse. This budget ' - 'covers both the durable cache (the ssm_states/conv_states ' - 'slots reused across requests) and the per-step extraction ' - 'scratch (the intermediate_ssm_out/intermediate_conv_out ' - 'buffers, sized to min(ceil(max_tokens / block_size), ' - '3 * max_requests) slots); the scratch is reserved first, so a ' - 'smaller max_tokens (or max_requests) shrinks the scratch and ' - 'leaves more durable slots.') - group.add_argument('--inference-dynamic-batching-cuda-graph-mixed-prefill-count', - type=int, default=16, - help='Number of mixed prefill requests to capture in a cuda graph.') - group.add_argument('--inference-dynamic-batching-cuda-graph-sizing-distribution', - type=str, default='exponential', - choices=['exponential', 'linear'], - dest='inference_dynamic_batching_cuda_graph_sizing_distribution', - help='Spacing of CUDA graph token counts. "exponential" (default) ' - 'halves from cuda_graph_max_tokens down to tp_size, giving a ' - 'log-spaced distribution with bounded relative padding. ' - '"linear" uses varying linear strides across the range.') - group.add_argument('--inference-dynamic-batching-sampling-backend', - type=str, default='torch', - choices=['torch', 'flashinfer'], - help='Which sampling kernels to use during inference. ' - 'Falls back to "torch" with a warning if "flashinfer" ' - 'is requested but the package is not installed.') - group.add_argument('--use-same-sampling-seed-across-dp-ranks', - action='store_false', dest='offset_sampling_seed_by_dp_rank', - default=True, - help='Use the same inference sampling seed on every data-parallel rank. ' - '--deterministic-mode also uses the same seed on every DP rank.') - group.add_argument('--inference-dynamic-batching-async-sched-mode', - type=str, default='legacy', - choices=['legacy', 'async'], - help='Async scheduling mode for dynamic batching. ' - '"legacy" (default) preserves the existing resolve-before-prepare ' - 'path. "async" overlaps asynchronous scheduling phases by reordering ' - 'them to prepare-before-resolve.') - group.add_argument('--inference-dynamic-batching-logprobs-mode', - type=str, default='raw_logprobs', - choices=['raw_logprobs', 'processed_logprobs'], - help='How returned inference log-probs are computed engine-wide. ' - '"raw_logprobs" (default) uses the unmodified model logits; ' - '"processed_logprobs" uses temperature and filters by top-k/top-p.') - group.add_argument('--inference-logging-step-interval', type=int, default=0, - help='Step interval for logging inference metrics. ' - 'Default to 0 to disable inference logging.') - group.add_argument('--inference-text-gen-server-logging', action=argparse.BooleanOptionalAction, - required=False, default=False, - help='Enable per-request logging in the inference text generation server.') - group.add_argument('--inference-wandb-logging', action=argparse.BooleanOptionalAction, - required=False, default=False, help='Enable inference wandb logging.') - group.add_argument("--inference-coordinator-port", type=int, - help="This port will be used to setup the inference coordinator on node-0") - group.add_argument('--mamba-inference-conv-states-dtype', type=str, - choices=['bf16', 'fp16', 'fp32'], default='bf16', - help='Dtype for the Mamba inference conv states tensor') - group.add_argument('--mamba-inference-ssm-states-dtype', type=str, - choices=['bf16', 'fp16', 'fp32'], default='bf16', - help='Dtype for the Mamba inference SSM states tensor') - group.add_argument('--inference-use-synchronous-zmq-collectives', action=argparse.BooleanOptionalAction, - required=False, default=False, help='Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for MoEs.') - group.add_argument('--inference-disable-ep-consensus', action=argparse.BooleanOptionalAction, - required=False, default=False, - help='Skip the EP-group consensus all-reduce in the inference engine control loop and step on local state only. ' - 'Pause/unpause take effect as soon as the signal is delivered to a rank. ' - 'Only safe when EP coordination is not required (e.g. ep_world_size == 1).') - group.add_argument('--inference-shards', type=str, default=None, metavar='SPEC', - help='Partition the world into independent inference models, each with ' - 'its own parallelism, e.g. "tp=2,role=prefill+tp=1,dp=2,role=decode". ' - 'Shards are separated by "+" or ";"; per-shard keys are ' - 'tp,pp,ep,expt_tp,dp (each defaults to 1) and must partition the full ' - 'world. Tagging shards role=prefill|decode enables disaggregated ' - 'inference (prefill hands KV to the decode pool); a dp>1 decode shard ' - 'is several independent decode instances.') + group.add_argument( + '--inference-batch-times-seqlen-threshold', + type=int, + default=-1, + help='If (batch-size * sequence-length) is smaller than this threshold' + 'then batches will not be split up for pipelining.' + 'Requires setting --pipeline-model-parallel-size > 1.' + 'Setting this to -1 indicates that batch pipelining is not used.', + ) + group.add_argument( + '--max-tokens-to-oom', + type=int, + default=12000, + help='Maximum number of tokens during inference' + 'tokens here is # in prompt + # to generate' + 'Allows us to throw an error before OOM crashes server', + ) + group.add_argument( + '--output-bert-embeddings', + action='store_true', + help='Output Bert embeddings (via mean pooling) from ' + 'model, rather than its binary head output or entire ' + 'hidden batch.', + ) + group.add_argument( + '--bert-embedder-type', + default="megatron", + choices=["megatron", "huggingface"], + help='Select either Megatron or Huggingface as the ' 'Bert embedder.', + ) + group.add_argument( + '--cuda-graph-scope', + nargs='+', + type=_parse_cuda_graph_modules_arg, + default=None, + dest='cuda_graph_scope_deprecated', + help=argparse.SUPPRESS, # hidden; use --cuda-graph-modules instead + ) + group.add_argument( + '--cuda-graph-modules', + nargs='+', + type=_parse_cuda_graph_modules_arg, + default=[], + help='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 validate_args. The deprecated values ' + '"full_iteration" and "full_iteration_inference" are also accepted and migrated ' + 'to the new API in validate_args.', + ) + group.add_argument( + '--use-legacy-static-engine', + action='store_true', + default=False, + help='Use legacy static engine. (Current static engine uses dynamic engine under the hood)', + dest='use_legacy_static_engine', + ) + group.add_argument( + '--inference-max-requests', + type=int, + default=8, + help='Maximum number of requests for inference.', + dest='inference_max_requests', + ) + group.add_argument( + '--inference-max-seq-length', + type=int, + default=2560, + help='Maximum sequence length expected for inference (prefill + decode).', + dest='inference_max_seq_length', + ) + group.add_argument( + '--inference-dynamic-batching', + action='store_true', + default=False, + help='Enable dynamic batching mode.', + ) + group.add_argument( + '--inference-dynamic-batching-buffer-size-gb', + type=float, + default=40.0, + help='Amount of on-GPU memory allocated for the KV cache. ' + 'The total amount of memory allocated for the KV cache ' + '(CPU + GPU memory) depends on the value set for the ' + 'unified virtual memory (UVM) level (via ' + '`--inference-dynamic-batching-unified-memory-level`).' + 'If the UVM level is 0, then only GPU memory is used and ' + 'the total memory equals `buffer_size_gb`. If the UVM ' + 'level is 1, then additional memory is utilized on the ' + 'CPU and the total memory equals `buffer_size_gb + ' + 'paused_buffer_size_gb`.', + ) + group.add_argument( + '--inference-dynamic-batching-paused-buffer-size-gb', + type=float, + default=None, + help='Amount of memory reserved for paused requests in ' + 'the dynamic inference context. Active requests are ' + 'paused when there are not enough active blocks available ' + 'to continue generating a request.', + ) + group.add_argument( + '--inference-dynamic-batching-mamba-memory-ratio', + type=float, + default=None, + help='Percentage of memory buffer to allocate for Mamba states. ' + 'If not specified, allocates Mamba state tensors for each KV cache block. ' + 'Only used for hybrid models.', + ) + group.add_argument( + '--inference-dynamic-batching-block-size', + type=int, + default=256, + help='KV cache block size. ' 'It should be a multiple of 256', + ) + group.add_argument( + '--inference-dynamic-batching-max-requests', + type=int, + default=None, + help='Override the inference context\'s `max_requests`. ' + 'By default, `max_requests` is set to the number of ' + 'blocks in the context\'s memory buffer.', + ) + group.add_argument( + '--inference-dynamic-batching-max-tokens', + type=int, + default=None, + help='Override the inference context\'s default `max_tokens`.', + ) + group.add_argument( + '--inference-dynamic-batching-num-cuda-graphs', + type=int, + default=16, + help='Maximum number of cuda graphs to capture, where the ' + 'cuda graph batch sizes range from 1 to `max_requests`. ' + '(See `dynamic_context.py` for details on how ' + '`max_requests` is computed). Due to rounding, the actual ' + 'number of cuda graphs may not equal this argument.' + 'The user can also pass -1, in which case we automatically determine the number of graphs ' + 'to capture based on the `max_requests`.', + ) + group.add_argument( + '--inference-dynamic-batching-track-paused-request-events', + action='store_true', + help='Track paused request ids by adding \'paused\' events ' + 'to each request\'s event history. This has a very minor ' + 'impact on latency.', + ) + group.add_argument( + '--inference-dynamic-batching-track-generated-token-events', + action='store_true', + help='Track per-token events with timestamps for each generated token. ' + 'When enabled, each generated token creates a GENERATED_TOKEN event ' + 'with a timestamp, useful for per-token latency analysis.', + ) + group.add_argument( + '--decode-only-cuda-graphs', + action='store_true', + default=False, + help='Only use cuda graphs for decode-only steps, not prefill and mixed steps.', + ) + group.add_argument( + '--inference-dynamic-batching-unified-memory-level', + type=int, + default=0, + choices=[0, 1], + help='Set unified memory usage within the dynamic ' + 'inference context. The levels are: 0) no unified memory, ' + '1) allocate `memory_buffer` in unified memory. ' + 'Eventually, additional levels will be included to ' + 'control other tensors within the context.', + ) + group.add_argument( + '--enable-chunked-prefill', + dest='enable_chunked_prefill', + action='store_true', + default=False, + help="Enable chunked prefill (disabled by default)", + ) + group.add_argument( + '--num-speculative-tokens', + type=int, + default=0, + help='Number of speculative tokens generated during decode', + ) + group.add_argument( + '--inference-dynamic-batching-prefix-caching', + dest='inference_dynamic_batching_enable_prefix_caching', + action=argparse.BooleanOptionalAction, + default=False, + help='Enable/disable prefix caching for dynamic batching inference. ' + 'When disabled, KV cache blocks cannot be shared between ' + 'requests with identical prompt prefixes.', + ) + group.add_argument( + '--inference-dynamic-batching-prefix-caching-eviction-policy', + type=str, + default='ref_zero', + choices=['ref_zero', 'lru'], + dest='inference_dynamic_batching_prefix_caching_eviction_policy', + help='Eviction policy for prefix caching blocks. ' + '"ref_zero" (default) immediately returns blocks to the ' + 'free pool when ref_count hits 0. "lru" keeps blocks ' + 'cached and evicts via LRU only when space is needed.', + ) + group.add_argument( + '--inference-dynamic-batching-prefix-caching-coordinator-policy', + type=str, + default='first_prefix_block', + choices=['longest_prefix', 'first_prefix_block', 'round_robin'], + dest='inference_dynamic_batching_prefix_caching_coordinator_policy', + help='Coordinator routing policy for prefix caching. ' + '"first_prefix_block" (default) routes based on the first ' + 'block hash only. "longest_prefix" routes to the rank with ' + 'the longest matching prefix. "round_robin" ignores prefix ' + 'affinity and cycles through ranks.', + ) + group.add_argument( + '--inference-dynamic-batching-prefix-caching-routing-alpha', + type=float, + default=0.5, + dest='inference_dynamic_batching_prefix_caching_routing_alpha', + help='Weight for prefix-aware routing score: ' + 'score = alpha * match + (1 - alpha) * normalized_load. ' + 'Higher alpha favors prefix cache hits; lower alpha ' + 'favors load balance. Default: 0.5.', + ) + group.add_argument( + '--inference-dynamic-batching-prefix-caching-mamba-gb', + type=float, + default=None, + dest='inference_dynamic_batching_prefix_caching_mamba_gb', + help='GPU memory budget (in GB) for the Mamba state cache ' + 'used by prefix caching on hybrid models. When set, Mamba ' + 'states at block boundaries are cached for reuse.', + ) + group.add_argument( + '--inference-dynamic-batching-cuda-graph-max-tokens', + type=int, + default=16384, + help='Maximum number of tokens to capture in a cuda graph.', + ) + group.add_argument( + '--inference-dynamic-batching-cuda-graph-mixed-prefill-count', + type=int, + default=16, + help='Number of mixed prefill requests to capture in a cuda graph.', + ) + group.add_argument( + '--inference-dynamic-batching-sampling-backend', + type=str, + default='torch', + choices=['torch', 'flashinfer'], + help='Which sampling kernels to use during inference. ' + 'Falls back to "torch" with a warning if "flashinfer" ' + 'is requested but the package is not installed.', + ) + group.add_argument( + '--inference-dynamic-batching-async-sched-mode', + type=str, + default='legacy', + choices=['legacy', 'serial'], + help='Async scheduling mode for dynamic batching. ' + '"legacy" (default) preserves the existing resolve-before-prepare ' + 'path. "serial" speculatively prepares and forwards decode-only ' + 'steps before resolving finished requests.', + ) + group.add_argument( + '--inference-dynamic-batching-logprobs-mode', + type=str, + default='raw_logprobs', + choices=['raw_logprobs', 'processed_logprobs'], + help='How returned inference log-probs are computed engine-wide. ' + '"raw_logprobs" (default) uses the unmodified model logits; ' + '"processed_logprobs" uses temperature and filters by top-k/top-p.', + ) + group.add_argument( + '--inference-logging-step-interval', + type=int, + default=0, + help='Step interval for logging inference metrics. ' + 'Default to 0 to disable inference logging.', + ) + group.add_argument( + '--inference-text-gen-server-logging', + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help='Enable per-request logging in the inference text generation server.', + ) + group.add_argument( + '--inference-wandb-logging', + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help='Enable inference wandb logging.', + ) + group.add_argument( + "--inference-coordinator-port", + type=int, + help="This port will be used to setup the inference coordinator on node-0", + ) + group.add_argument( + '--mamba-inference-conv-states-dtype', + type=str, + choices=['bf16', 'fp16', 'fp32'], + default='bf16', + help='Dtype for the Mamba inference conv states tensor', + ) + group.add_argument( + '--mamba-inference-ssm-states-dtype', + type=str, + choices=['bf16', 'fp16', 'fp32'], + default='bf16', + help='Dtype for the Mamba inference SSM states tensor', + ) + group.add_argument( + '--inference-use-synchronous-zmq-collectives', + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help='Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for MoEs.', + ) + group.add_argument( + '--inference-disable-ep-consensus', + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help='Skip the EP-group consensus all-reduce in the inference engine control loop and step on local state only. ' + 'Pause/unpause take effect as soon as the signal is delivered to a rank. ' + 'Only safe when EP coordination is not required (e.g. ep_world_size == 1).', + ) + group.add_argument( + '--inference-shards', + type=str, + default=None, + metavar='SPEC', + help='Partition the world into independent inference models, each with ' + 'its own parallelism, e.g. "tp=2,role=prefill+tp=1,dp=2,role=decode". ' + 'Shards are separated by "+" or ";"; per-shard keys are ' + 'tp,pp,ep,expt_tp,dp (each defaults to 1) and must partition the full ' + 'world. Tagging shards role=prefill|decode enables disaggregated ' + 'inference (prefill hands KV to the decode pool); a dp>1 decode shard ' + 'is several independent decode instances.', + ) + group.add_argument( + '--inference-cuda-graph-all-prefills', + action='store_true', + default=False, + help='Extend prefill/mixed CUDA graph capture up to `max_tokens`. ' + 'By default, all graphs are limited by the decode limit of ' + '`max_requests * (num_speculative_tokens + 1)`.', + ) + group.add_argument( + '--inference-dynamic-batching-cuda-graph-sizing-distribution', + type=str, + default='exponential', + choices=['exponential', 'linear'], + dest='inference_dynamic_batching_cuda_graph_sizing_distribution', + help='Spacing of CUDA graph token counts. "exponential" (default) ' + 'halves from cuda_graph_max_tokens down to tp_size, giving a ' + 'log-spaced distribution with bounded relative padding. ' + '"linear" uses varying linear strides across the range.', + ) return parser @@ -2366,6 +2758,7 @@ def _add_network_size_args(parser): "no_rope_freq", "moe_layer_freq", "linear_attention_freq", + "csa_compress_ratios", "moe_router_load_balancing_type", "moe_aux_loss_coeff", "cp_comm_type", @@ -2421,8 +2814,11 @@ def _add_network_size_args(parser): "barrier_with_L1_time", # args uses same var with a different name "num_moe_experts", + "actual_vocab_size", "fp8_param", "fp4_param", + # already generated by data args + "create_attention_mask_in_dataloader", # incompatible defaults in dataclass "gradient_accumulation_fusion", "overlap_p2p_comm", @@ -2431,249 +2827,418 @@ def _add_network_size_args(parser): "persist_layer_norm", "bias_dropout_fusion", "apply_rope_fusion", + "apply_dsa_kernel_fusion", + "dsa_kernel_backend", "mamba_training_ssm_states_dtype", - "max_seqlen_per_dp_cp_rank", - "hybrid_context_parallel", - "sequence_packing_scheduler", - # internal/derived: controlled only via --tensor-parallel-num-weight-shards - "gtp_weight_remat_size", - # internal/derived: controlled only via --expert-tensor-parallel-num-weight-shards - "expert_gtp_weight_remat_size", - "max_seqlen_per_dp_cp_rank", - "hybrid_context_parallel", - "sequence_packing_scheduler", ] transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude) transformer_group = transformer_factory.build_group(parser, "transformer configuration") group = parser.add_argument_group(title='network size') - group.add_argument('--encoder-num-layers', type=int, default=None, - help='Number of encoder transformer layers.') - group.add_argument('--decoder-num-layers', type=int, default=None, - help='Number of decoder transformer layers.') - group.add_argument('--group-query-attention', action='store_true', - help='Use group-query attention.') - group.add_argument('--window-size', type=tuple_type, default=None, - help='Window size for window attention. If not provided, ' - 'window attention will be disabled.') - group.add_argument('--window-attn-skip-freq', type=moe_freq_type, default=None, - help='Frequency of layers to skip window attention. Accepts either: ' - '- An integer N: Represents a (N-1):1 ratio, meaning one full attention layer ' - 'after (N-1) SWA layers. ' - '- A string containing a Python list expression that defines a custom pattern, ' - 'e.g.: "[1,1,1,0]*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' - 'where 1 indicates SWA and 0 indicates full attention. ') - group.add_argument('--max-position-embeddings', type=int, default=None, - help='Maximum number of position embeddings to use. ' - 'This is the size of position embedding.') - group.add_argument('--position-embedding-type', type=str, default='learned_absolute', - choices=['learned_absolute', 'rope', 'yarn', 'mrope', 'relative', 'none'], - help='Position embedding type.') - group.add_argument('--yarn-original-max-position-embeddings', type=int, default=None, - help='Original maximum position embeddings for YaRN RoPE frequency correction.') - group.add_argument('--yarn-beta-fast', type=float, default=None, - help='Beta fast for YaRN RoPE frequency correction.') - group.add_argument('--yarn-beta-slow', type=float, default=None, - help='Beta slow for YaRN RoPE frequency correction.') - group.add_argument('--yarn-correction-range-round-to-int', action='store_true', default=None, - help='Round YaRN correction range endpoints to integers.') - group.add_argument('--no-yarn-correction-range-round-to-int', action='store_false', - dest='yarn_correction_range_round_to_int', - help='Do not round YaRN correction range endpoints to integers.') - group.add_argument('--relative-attention-num-buckets', type=int, default=32, - help='Number of buckets for relative position embeddings.') - group.add_argument('--relative-attention-max-distance', type=int, default=128, - help='Maximum distance for relative position embeddings calculation.') - group.add_argument('--use-rotary-position-embeddings', action='store_true', - help='Use rotary positional embeddings or not. ' - 'Deprecated: use --position-embedding-type') - group.add_argument('--rotary-base', type=int, default=10000, - help='Base to use for rotary positional embeddings, default 10000') - group.add_argument('--rotary-percent', type=float, default=1.0, - help='Percent of rotary dimension to use, default 100%%') - group.add_argument('--rotary-seq-len-interpolation-factor', type=int, default=None, - help='Sequence length interpolation factor for rotary embeddings.') - group.add_argument('--use-rope-scaling', action='store_true', - help='Apply rope scaling as used in llama3.x') - group.add_argument('--rope-scaling-factor', type=float, default=8.0, - help='Rope scaling factor in llama3.x models') - group.add_argument('--no-rope-freq', type=no_rope_freq_type, default=None, - help='Controls which layers to skip performing Rotary Position Embedding. Accepts either: ' - '- An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. ' - '- A string containing a Python list expression that defines a custom pattern, e.g.: ' - '"([0]*3+[1]*1)*3" evaluates to [0,0,0,1,0,0,0,1,0,0,0,1] ' - 'where 1 indicates no-rope layer. This patten is equivalent to --no-rope-freq=4.' - 'By default this is disabled and set to None, indicating RoPE will be performed' - 'on every layer.' - ) - group.add_argument('--no-position-embedding', - action='store_false', - help='Disable position embedding. Deprecated: use --position-embedding-type', - dest='add_position_embedding') - group.add_argument('--make-vocab-size-divisible-by', type=int, default=128, - help='Pad the vocab size to be divisible by this value.' - 'This is added for computational efficieny reasons.') - group.add_argument('--openai-gelu', action='store_true', - help='Use OpenAIs GeLU implementation. This option' - 'should not be used unless for backward compatibility' - 'reasons.') - group.add_argument('--squared-relu', action='store_true', - help='Use squared relu activation instead of default gelu') - group.add_argument('--swiglu', action='store_true', - help='Use gated linear units and SiLU activation instead of default gelu') - group.add_argument('--quick-geglu', action='store_true', - help='Use quick geglu activation instead of default gelu') - group.add_argument('--onnx-safe', type=bool, required=False, - help='Use workarounds for known problems with ' - 'Torch ONNX exporter') - group.add_argument('--bert-no-binary-head', action='store_false', - help='Disable BERT binary head.', - dest='bert_binary_head') - group.add_argument('--untie-embeddings-and-output-weights', action='store_true', - help='Untie embeddings and output weights.') - return parser - -def _add_straggler_detector_args(parser): - from megatron.training.config import StragglerDetectionConfig - - straggler_factory = ArgumentGroupFactory(StragglerDetectionConfig) - group = straggler_factory.build_group(parser, "straggler") - - return parser - -def _add_workload_inspector_server_args(parser): - group = parser.add_argument_group(title='workload inspector') - group.add_argument('--run-workload-inspector-server', action='store_true', - help='If set, enables workload inspector server for on-demand profiling.') - return parser - -def _add_inprocess_restart_args(parser): - group = parser.add_argument_group(title='In-process restart') - - group.add_argument('--inprocess-restart', action='store_true', - help='Enables in-process restart.') - - group.add_argument('--inprocess-max-iterations', default=None, type=int, - help='Maximum number of in-process restart iterations.') - group.add_argument('--inprocess-monitor-thread-interval', default=1.0, type=float, - help='Monitoring interval (in seconds) for the monitoring thread.') - group.add_argument('--inprocess-monitor-process-interval', default=1.0, type=float, - help='Monitoring interval (in seconds) for the monitoring process.') - group.add_argument('--inprocess-progress-watchdog-interval', default=1.0, type=float, - help='Interval (in seconds) for automatic progress watchdog timestamp ' - 'updates.') - group.add_argument('--inprocess-heartbeat-interval', default=30, type=float, - help='Monitoring interval (in seconds) for detecting unresponsive ranks.') - - group.add_argument('--inprocess-soft-timeout', default=60, type=float, - help='Soft progress timeout (in seconds).') - group.add_argument('--inprocess-hard-timeout', default=90, type=float, - help='Hard progress timeout (in seconds).') - group.add_argument('--inprocess-heartbeat-timeout', default=60, type=float, - help='Timeout (in seconds) for a missing rank detection heartbeat.') - - group.add_argument('--inprocess-barrier-timeout', default=120, type=float, - help='Timeout (in seconds) for internal distributed barrier') - group.add_argument('--inprocess-completion-timeout', default=120, type=float, - help='Timeout (in seconds) for barrier on completion on all ranks') - - group.add_argument('--inprocess-last-call-wait', default=1, type=float, - help='Time interval (in seconds) for other ranks to report concurrent ' - 'terminal failures.') - group.add_argument('--inprocess-termination-grace-time', default=1, type=float, - help='Interval (in seconds) between SIGTERM and SIGKILL issued on hard ' - 'timeout') - - group.add_argument('--inprocess-granularity', default='node', type=str, - choices=['node', 'rank'], - help='Granularity for in-process restart.') - group.add_argument('--inprocess-active-world-size', - default=int(os.getenv('WORLD_SIZE', '1')), type=int, - help='The number of ranks initially executing the workload. ' - 'The remaining ranks from the allocation are set aside ' - 'as warm reserve.') - group.add_argument('--inprocess-empty-cuda-cache', action='store_true', - help='Release all unoccupied cached GPU memory on every in-process restart.') - return parser - -def _add_one_logger_args(parser): - group = parser.add_argument_group(title='one logger') - group.add_argument('--no-one-logger', action='store_false', - help='If set, disable using one_logger to track E2E metrics' - 'Note that one_logger is an internal tool and not ' - 'available externally. For installation, please go to ' - 'https://confluence.nvidia.com/display/MLWFO/Package+Repositories' - 'for more details', - dest='enable_one_logger') - group.add_argument('--one-logger-project', type=str, default='megatron-lm', - help='The one-logger project name. Will ignore if ' - '--no-one-logger is set') - group.add_argument('--one-logger-run-name', type=str, default=None, - help='The one-logger run name displayed. Will ignore if ' - '--no-one-logger is set') - group.add_argument('--one-logger-async', action='store_true', - help='If set, forces one_logger to use async mode.') - group.add_argument('--app-tag-run-name', type=str, default=None, - help='Jobs belonging to same training run, suppose to ' - 'have the same name. It will be used to track progress of ' - 'a training done over multiple different jobs') - group.add_argument('--app-tag-run-version', type=str, default='0.0.0', - help='The version of the training of which current job is ' - 'part of. It will be used to track the changes in the ' - 'application side which might change the performance ' - 'baseline') - return parser - - -def _add_ft_package_args(parser): - group = parser.add_argument_group(title='ft_package') - group.add_argument('--enable-ft-package', action='store_true', - help='If set, Fault Tolerance package is enabled. ' - 'Note: This feature is for Nvidia internal use only.') - group.add_argument('--calc-ft-timeouts', action='store_true', - help='If set, FT package will try to automatically compute the timeouts. ' - 'Note: This feature is for Nvidia internal use only.') - group.add_argument('--ft-num-warmup-iters', type=int, default=5, - help='Number of warmup iterations before monitoring step section and ' - 'out-of-section timeouts. The first N iterations are excluded from ' - 'timeout monitoring as they can be significantly slower than steady-state. ' - 'Default: 5. Note: This feature is for Nvidia internal use only.') - return parser - - -def _add_logging_args(parser): - from megatron.training.config import LoggerConfig - - log_factory = ArgumentGroupFactory(LoggerConfig, exclude = ["log_throughput_to_tensorboard", "throughput_window_size", "memory_keys", "log_l2_norm_grad_to_tensorboard", "log_runtime_to_tensorboard", "runtime_time_unit", "filter_warnings", "modules_to_filter", "set_level_for_all_loggers", "save_config_filepath"]) - group = log_factory.build_group(parser, title="logging") + group.add_argument( + '--encoder-num-layers', type=int, default=None, help='Number of encoder transformer layers.' + ) + group.add_argument( + '--decoder-num-layers', type=int, default=None, help='Number of decoder transformer layers.' + ) + group.add_argument( + '--group-query-attention', action='store_true', help='Use group-query attention.' + ) + group.add_argument( + '--window-size', + type=tuple_type, + default=None, + help='Window size for window attention. If not provided, ' + 'window attention will be disabled.', + ) + group.add_argument( + '--window-attn-skip-freq', + type=moe_freq_type, + default=None, + help='Frequency of layers to skip window attention. Accepts either: ' + '- An integer N: Represents a (N-1):1 ratio, meaning one full attention layer ' + 'after (N-1) SWA layers. ' + '- A string containing a Python list expression that defines a custom pattern, ' + 'e.g.: "[1,1,1,0]*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' + 'where 1 indicates SWA and 0 indicates full attention. ', + ) + group.add_argument( + '--max-position-embeddings', + type=int, + default=None, + help='Maximum number of position embeddings to use. ' + 'This is the size of position embedding.', + ) + group.add_argument( + '--position-embedding-type', + type=str, + default='learned_absolute', + choices=['learned_absolute', 'rope', 'yarn', 'mrope', 'relative', 'none'], + help='Position embedding type.', + ) + group.add_argument( + '--yarn-original-max-position-embeddings', + type=int, + default=None, + help='Original maximum position embeddings for YaRN RoPE frequency correction.', + ) + group.add_argument( + '--yarn-beta-fast', + type=float, + default=None, + help='Beta fast for YaRN RoPE frequency correction.', + ) + group.add_argument( + '--yarn-beta-slow', + type=float, + default=None, + help='Beta slow for YaRN RoPE frequency correction.', + ) + group.add_argument( + '--yarn-correction-range-round-to-int', + action='store_true', + default=None, + help='Round YaRN correction range endpoints to integers.', + ) + group.add_argument( + '--no-yarn-correction-range-round-to-int', + action='store_false', + dest='yarn_correction_range_round_to_int', + help='Do not round YaRN correction range endpoints to integers.', + ) + group.add_argument( + '--relative-attention-num-buckets', + type=int, + default=32, + help='Number of buckets for relative position embeddings.', + ) + group.add_argument( + '--relative-attention-max-distance', + type=int, + default=128, + help='Maximum distance for relative position embeddings calculation.', + ) + group.add_argument( + '--use-rotary-position-embeddings', + action='store_true', + help='Use rotary positional embeddings or not. ' + 'Deprecated: use --position-embedding-type', + ) + group.add_argument( + '--rotary-base', + type=int, + default=10000, + help='Base to use for rotary positional embeddings, default 10000', + ) + group.add_argument( + '--rotary-percent', + type=float, + default=1.0, + help='Percent of rotary dimension to use, default 100%%', + ) + group.add_argument( + '--rotary-seq-len-interpolation-factor', + type=int, + default=None, + help='Sequence length interpolation factor for rotary embeddings.', + ) + group.add_argument( + '--use-rope-scaling', action='store_true', help='Apply rope scaling as used in llama3.x' + ) + group.add_argument( + '--rope-scaling-factor', + type=float, + default=8.0, + help='Rope scaling factor in llama3.x models', + ) + group.add_argument( + '--no-rope-freq', + type=no_rope_freq_type, + default=None, + help='Controls which layers to skip performing Rotary Position Embedding. Accepts either: ' + '- An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. ' + '- A string containing a Python list expression that defines a custom pattern, e.g.: ' + '"([0]*3+[1]*1)*3" evaluates to [0,0,0,1,0,0,0,1,0,0,0,1] ' + 'where 1 indicates no-rope layer. This patten is equivalent to --no-rope-freq=4.' + 'By default this is disabled and set to None, indicating RoPE will be performed' + 'on every layer.', + ) + group.add_argument( + '--no-position-embedding', + action='store_false', + help='Disable position embedding. Deprecated: use --position-embedding-type', + dest='add_position_embedding', + ) + group.add_argument( + '--make-vocab-size-divisible-by', + type=int, + default=128, + help='Pad the vocab size to be divisible by this value.' + 'This is added for computational efficieny reasons.', + ) + group.add_argument( + '--openai-gelu', + action='store_true', + help='Use OpenAIs GeLU implementation. This option' + 'should not be used unless for backward compatibility' + 'reasons.', + ) + group.add_argument( + '--squared-relu', + action='store_true', + help='Use squared relu activation instead of default gelu', + ) + group.add_argument( + '--swiglu', + action='store_true', + help='Use gated linear units and SiLU activation instead of default gelu', + ) + group.add_argument( + '--quick-geglu', + action='store_true', + help='Use quick geglu activation instead of default gelu', + ) + group.add_argument( + '--onnx-safe', + type=bool, + required=False, + help='Use workarounds for known problems with ' 'Torch ONNX exporter', + ) + group.add_argument( + '--bert-no-binary-head', + action='store_false', + help='Disable BERT binary head.', + dest='bert_binary_head', + ) + group.add_argument( + '--untie-embeddings-and-output-weights', + action='store_true', + help='Untie embeddings and output weights.', + ) + return parser + + +def _add_straggler_detector_args(parser): + from megatron.training.config import StragglerDetectionConfig + + straggler_factory = ArgumentGroupFactory(StragglerDetectionConfig) + group = straggler_factory.build_group(parser, "straggler") + + return parser + - otel_group = parser.add_argument_group(title='opentelemetry') - otel_group.add_argument( - '--otel-enabled', +def _add_workload_inspector_server_args(parser): + group = parser.add_argument_group(title='workload inspector') + group.add_argument( + '--run-workload-inspector-server', action='store_true', - default=False, - help='Enable OpenTelemetry telemetry (traces and metrics). ' - 'See MEGATRON_OTEL_ENABLED env var for the env-var equivalent.', + help='If set, enables workload inspector server for on-demand profiling.', + ) + return parser + + +def _add_inprocess_restart_args(parser): + group = parser.add_argument_group(title='In-process restart') + + group.add_argument( + '--inprocess-restart', action='store_true', help='Enables in-process restart.' + ) + + group.add_argument( + '--inprocess-max-iterations', + default=None, + type=int, + help='Maximum number of in-process restart iterations.', + ) + group.add_argument( + '--inprocess-monitor-thread-interval', + default=1.0, + type=float, + help='Monitoring interval (in seconds) for the monitoring thread.', + ) + group.add_argument( + '--inprocess-monitor-process-interval', + default=1.0, + type=float, + help='Monitoring interval (in seconds) for the monitoring process.', + ) + group.add_argument( + '--inprocess-progress-watchdog-interval', + default=1.0, + type=float, + help='Interval (in seconds) for automatic progress watchdog timestamp ' 'updates.', + ) + group.add_argument( + '--inprocess-heartbeat-interval', + default=30, + type=float, + help='Monitoring interval (in seconds) for detecting unresponsive ranks.', + ) + + group.add_argument( + '--inprocess-soft-timeout', + default=60, + type=float, + help='Soft progress timeout (in seconds).', + ) + group.add_argument( + '--inprocess-hard-timeout', + default=90, + type=float, + help='Hard progress timeout (in seconds).', + ) + group.add_argument( + '--inprocess-heartbeat-timeout', + default=60, + type=float, + help='Timeout (in seconds) for a missing rank detection heartbeat.', ) - otel_group.add_argument( - '--otel-service-name', + + group.add_argument( + '--inprocess-barrier-timeout', + default=120, + type=float, + help='Timeout (in seconds) for internal distributed barrier', + ) + group.add_argument( + '--inprocess-completion-timeout', + default=120, + type=float, + help='Timeout (in seconds) for barrier on completion on all ranks', + ) + + group.add_argument( + '--inprocess-last-call-wait', + default=1, + type=float, + help='Time interval (in seconds) for other ranks to report concurrent ' + 'terminal failures.', + ) + group.add_argument( + '--inprocess-termination-grace-time', + default=1, + type=float, + help='Interval (in seconds) between SIGTERM and SIGKILL issued on hard ' 'timeout', + ) + + group.add_argument( + '--inprocess-granularity', + default='node', + type=str, + choices=['node', 'rank'], + help='Granularity for in-process restart.', + ) + group.add_argument( + '--inprocess-active-world-size', + default=int(os.getenv('WORLD_SIZE', '1')), + type=int, + help='The number of ranks initially executing the workload. ' + 'The remaining ranks from the allocation are set aside ' + 'as warm reserve.', + ) + group.add_argument( + '--inprocess-empty-cuda-cache', + action='store_true', + help='Release all unoccupied cached GPU memory on every in-process restart.', + ) + return parser + + +def _add_one_logger_args(parser): + group = parser.add_argument_group(title='one logger') + group.add_argument( + '--no-one-logger', + action='store_false', + help='If set, disable using one_logger to track E2E metrics' + 'Note that one_logger is an internal tool and not ' + 'available externally. For installation, please go to ' + 'https://confluence.nvidia.com/display/MLWFO/Package+Repositories' + 'for more details', + dest='enable_one_logger', + ) + group.add_argument( + '--one-logger-project', + type=str, + default='megatron-lm', + help='The one-logger project name. Will ignore if ' '--no-one-logger is set', + ) + group.add_argument( + '--one-logger-run-name', type=str, default=None, - help='Override OTEL_SERVICE_NAME for this training run.', + help='The one-logger run name displayed. Will ignore if ' '--no-one-logger is set', + ) + group.add_argument( + '--one-logger-async', + action='store_true', + help='If set, forces one_logger to use async mode.', ) - otel_group.add_argument( - '--otel-span-groups', + group.add_argument( + '--app-tag-run-name', type=str, default=None, - help='Comma-separated span-group spec controlling which OTel ' - 'instrumentation boundaries are active. Accepts preset keywords ' - '("default", "per_step", "full", "all") or individual group names ' - '("job", "checkpoint", "evaluate", "model_init", "load_checkpoint", ' - '"step", "forward_backward", "optimizer", "microbatch"), or a mix. ' - 'Defaults to "default" (coarse job/checkpoint/evaluate spans only). ' - 'Equivalent to MEGATRON_OTEL_SPAN_GROUPS env var.', + help='Jobs belonging to same training run, suppose to ' + 'have the same name. It will be used to track progress of ' + 'a training done over multiple different jobs', + ) + group.add_argument( + '--app-tag-run-version', + type=str, + default='0.0.0', + help='The version of the training of which current job is ' + 'part of. It will be used to track the changes in the ' + 'application side which might change the performance ' + 'baseline', + ) + return parser + + +def _add_ft_package_args(parser): + group = parser.add_argument_group(title='ft_package') + group.add_argument( + '--enable-ft-package', + action='store_true', + help='If set, Fault Tolerance package is enabled. ' + 'Note: This feature is for Nvidia internal use only.', + ) + group.add_argument( + '--calc-ft-timeouts', + action='store_true', + help='If set, FT package will try to automatically compute the timeouts. ' + 'Note: This feature is for Nvidia internal use only.', + ) + group.add_argument( + '--ft-num-warmup-iters', + type=int, + default=5, + help='Number of warmup iterations before monitoring step section and ' + 'out-of-section timeouts. The first N iterations are excluded from ' + 'timeout monitoring as they can be significantly slower than steady-state. ' + 'Default: 5. Note: This feature is for Nvidia internal use only.', + ) + return parser + + +def _add_logging_args(parser): + from megatron.training.config import LoggerConfig + + log_factory = ArgumentGroupFactory( + LoggerConfig, + exclude=[ + "log_throughput_to_tensorboard", + "throughput_window_size", + "memory_keys", + "log_l2_norm_grad_to_tensorboard", + "log_runtime_to_tensorboard", + "runtime_time_unit", + "filter_warnings", + "modules_to_filter", + "set_level_for_all_loggers", + "save_config_filepath", + ], ) + group = log_factory.build_group(parser, title="logging") return parser @@ -2681,639 +3246,1100 @@ def _add_logging_args(parser): def _add_regularization_args(parser): group = parser.add_argument_group(title='regularization') - group.add_argument('--weight-decay', type=float, default=0.01, - help='Weight decay coefficient for L2 regularization.') - group.add_argument('--apply-wd-to-qk-layernorm', action='store_true', - help='Apply weight decay to qk layernorm as a special case.') - group.add_argument('--clip-grad', type=float, default=1.0, - help='Gradient clipping based on global L2 norm.') - group.add_argument('--adam-beta1', type=float, default=0.9, - help='First coefficient for computing running averages ' - 'of gradient and its square') - group.add_argument('--adam-beta2', type=float, default=0.999, - help='Second coefficient for computing running averages ' - 'of gradient and its square') - group.add_argument('--adam-eps', type=float, default=1e-08, - help='Term added to the denominator to improve' - 'numerical stability') - group.add_argument('--sgd-momentum', type=float, default=0.9, - help='Momentum factor for sgd') - group.add_argument('--muon-momentum', type=float, default=0.9, - help='Momentum factor for Muon optimizer') - group.add_argument('--muon-no-split-qkv', action='store_false', default=True, - dest='muon_split_qkv', - help='Whether to split QKV parameters for Muon optimizer') - group.add_argument('--muon-nesterov', action='store_true', - help='Whether to use Nesterov-style momentum in the internal SGD') - group.add_argument('--muon-scale-mode', type=str, default='spectral', - choices=['spectral', 'unit_rms_norm', 'shape_scaling'], - help='Scale mode for Muon optimizer. With MuP, set ' - '--muon-scale-mode unit_rms_norm to use unit_rms_norm scaling, ' - 'or set --muon-scale-mode spectral to keep spectral scaling.') - group.add_argument('--muon-fp32-matmul-prec', type=str, default='medium', - choices=['low', 'medium', 'high'], - help='FP32 matmul precision for Newton-Schulz iteration') - group.add_argument('--muon-coefficient-type', type=str, default='quintic', - help='Newton-Schulz coefficient type for the Muon optimizer. ' - 'Valid types are discovered from the installed emerging_optimizers ' - 'package (e.g. simple, quintic, polar_express, aol). ' - 'Validated at optimizer creation time.') - group.add_argument('--muon-num-ns-steps', type=int, default=5, - help='Number of Newton-Schulz steps for Muon optimizer') - group.add_argument('--muon-tp-mode', type=str, default='duplicated', - choices=['blockwise', 'duplicated', 'distributed', 'auto'], - help='How to perform NS calculation for tensor model parallel weights. ' - 'blockwise orthogonalizes each shard on its own, so the update rule ' - 'depends on the parallelism config; duplicated and distributed both ' - 'orthogonalize the whole matrix and give TP-invariant results; auto ' - 'select between duplicated and distributed mode per-weight for ' - 'dense weights.') - group.add_argument('--muon-use-syrk', action='store_true', - help='Use the Triton SYRK kernel for the Gram matrix ' - 'in Newton-Schulz iteration.') - group.add_argument('--muon-extra-scale-factor', type=float, default=1.0, - help='Additional scale factor for the muon update') - group.add_argument('--muon-scalar-optimizer', type=str, default='adam', - choices=['adam', 'lion'], - help='Optimizer for scalar parameters (embeddings, biases, norms) ' - 'when using muon. Defaults to adam.') - group.add_argument('--lion-beta1', type=float, default=0.95, - help='First beta coefficient for Lion optimizer ' - '(used in sign update). Default: 0.95.') - group.add_argument('--lion-beta2', type=float, default=0.98, - help='Second beta coefficient for Lion optimizer ' - '(used in momentum EMA update). Default: 0.98.') - - group.add_argument('--no-weight-decay-cond-type', type=str, choices=['apply_wd_to_qk_layernorm'], - help='Type of no weight decay condition. Choices: ' - 'None (default): apply weight decay to 1D weights and biases.' - '"apply_wd_to_qk_layernorm": additionally apply weight decay to ' - 'qk layernorm as a special case.' - 'DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. ') + group.add_argument( + '--weight-decay', + type=float, + default=0.01, + help='Weight decay coefficient for L2 regularization.', + ) + group.add_argument( + '--apply-wd-to-qk-layernorm', + action='store_true', + help='Apply weight decay to qk layernorm as a special case.', + ) + group.add_argument( + '--clip-grad', type=float, default=1.0, help='Gradient clipping based on global L2 norm.' + ) + group.add_argument( + '--adam-beta1', + type=float, + default=0.9, + help='First coefficient for computing running averages ' 'of gradient and its square', + ) + group.add_argument( + '--adam-beta2', + type=float, + default=0.999, + help='Second coefficient for computing running averages ' 'of gradient and its square', + ) + group.add_argument( + '--adam-eps', + type=float, + default=1e-08, + help='Term added to the denominator to improve' 'numerical stability', + ) + group.add_argument('--sgd-momentum', type=float, default=0.9, help='Momentum factor for sgd') + group.add_argument( + '--muon-momentum', type=float, default=0.9, help='Momentum factor for Muon optimizer' + ) + group.add_argument( + '--muon-no-split-qkv', + action='store_false', + default=True, + dest='muon_split_qkv', + help='Whether to split QKV parameters for Muon optimizer', + ) + group.add_argument( + '--muon-nesterov', + action='store_true', + help='Whether to use Nesterov-style momentum in the internal SGD', + ) + group.add_argument( + '--muon-scale-mode', + type=str, + default='spectral', + choices=['spectral', 'unit_rms_norm', 'shape_scaling'], + help='Scale mode for Muon optimizer. With MuP, set ' + '--muon-scale-mode unit_rms_norm to use unit_rms_norm scaling, ' + 'or set --muon-scale-mode spectral to keep spectral scaling.', + ) + group.add_argument( + '--muon-fp32-matmul-prec', + type=str, + default='medium', + choices=['low', 'medium', 'high'], + help='FP32 matmul precision for Newton-Schulz iteration', + ) + group.add_argument( + '--muon-coefficient-type', + type=str, + default='quintic', + help='Newton-Schulz coefficient type for the Muon optimizer. ' + 'Valid types are discovered from the installed emerging_optimizers ' + 'package (e.g. simple, quintic, polar_express, aol). ' + 'Validated at optimizer creation time.', + ) + group.add_argument( + '--muon-num-ns-steps', + type=int, + default=5, + help='Number of Newton-Schulz steps for Muon optimizer', + ) + group.add_argument( + '--muon-tp-mode', + type=str, + default='blockwise', + choices=['blockwise', 'duplicated', 'distributed'], + help='How to perform NS calculation for tensor model parallel weights', + ) + group.add_argument( + '--muon-extra-scale-factor', + type=float, + default=1.0, + help='Additional scale factor for the muon update', + ) + group.add_argument( + '--muon-scalar-optimizer', + type=str, + default='adam', + choices=['adam', 'lion'], + help='Optimizer for scalar parameters (embeddings, biases, norms) ' + 'when using muon. Defaults to adam.', + ) + group.add_argument( + '--lion-beta1', + type=float, + default=0.95, + help='First beta coefficient for Lion optimizer ' '(used in sign update). Default: 0.95.', + ) + group.add_argument( + '--lion-beta2', + type=float, + default=0.98, + help='Second beta coefficient for Lion optimizer ' + '(used in momentum EMA update). Default: 0.98.', + ) + + group.add_argument( + '--no-weight-decay-cond-type', + type=str, + choices=['apply_wd_to_qk_layernorm'], + help='Type of no weight decay condition. Choices: ' + 'None (default): apply weight decay to 1D weights and biases.' + '"apply_wd_to_qk_layernorm": additionally apply weight decay to ' + 'qk layernorm as a special case.' + 'DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. ', + ) return parser def _add_rl_args(parser): group = parser.add_argument_group(title='rl') - group.add_argument('--perform-rl-step', action='store_true', - help="Use the RL training step.") - group.add_argument('--rl-prompts-per-eval', type=int, default=32, - help='Number of prompts to evaluate for for each RL task.' - 'This evaluation can be very expensive when using environments' - 'that evaluate pass@k so we default to a lower number.') + group.add_argument('--perform-rl-step', action='store_true', help="Use the RL training step.") + group.add_argument( + '--rl-prompts-per-eval', + type=int, + default=32, + help='Number of prompts to evaluate for for each RL task.' + 'This evaluation can be very expensive when using environments' + 'that evaluate pass@k so we default to a lower number.', + ) # TODO(rkirby): allow for "complete" evaluation when --rl-prompts-per-eval is set to -1 - group.add_argument('--grpo-prompts-per-step', type=int, default=32, - help="Number of GRPO groups (G in the paper).") - group.add_argument('--grpo-group-size', type=int, default=2, - help="Number of samples per a GRPO group.") - group.add_argument('--rl-generation-lag', type=float, default=None, - help='Number of trainer batches of rollout generation lag to allow ' - 'The number of in-flight trainer batches is this value plus one. ' - 'May be fractional or negative; the minimum of -1 keeps a single unit ' - 'of generation work in flight. If omitted, the lag is autotuned to the ' - 'inference engine\'s request capacity when --rl-partial-rollouts is ' - 'set, and is 0 otherwise. ' - 'Requires --rl-partial-rollouts when greater than 0. ' - 'Mutually exclusive with --rl-max-inflight-requests.') - group.add_argument('--rl-max-inflight-requests', type=int, default=None, - help='Maximum number of inference requests RL generation may keep inflight: ' - 'equivalent to (--rl-generation-lag + 1) training batches ' - 'of grpo_prompts_per_step * grpo_group_size requests each. ' - 'Requires --rl-partial-rollouts when above one training batch; ' - 'mutually exclusive with --rl-generation-lag.') + group.add_argument( + '--grpo-prompts-per-step', + type=int, + default=32, + help="Number of GRPO groups (G in the paper).", + ) + group.add_argument( + '--grpo-group-size', type=int, default=2, help="Number of samples per a GRPO group." + ) + group.add_argument( + '--rl-generation-lag', + type=int, + default=0, + help='Number of trainer batches of rollout generation lag to allow. ' + 'The number of in-flight trainer batches is this value plus one. ' + 'Requires --rl-partial-rollouts when greater than 0.', + ) # TODO: Refactor these string literals back to an enum after the megatron.training refactor. - group.add_argument('--rl-submission-granularity', type=str, - default="B", - choices=["R", "G", "B"], - help='Granularity for submitting rollout generation work. ' - 'R submits individual rollouts independently while still yielding ' - 'complete rollout groups to training. ' - 'G submits one rollout group at a time. ' - 'B submits grpo_prompts_per_step rollout groups together.') - group.add_argument('--rl-consumption-granularity', type=str, - default="B", - choices=["R", "G", "B"], - help='Granularity for consuming generated rollout groups. ' - 'G consumes groups as they complete. ' - 'B consumes complete trainer batches in submission order. ' - 'R is not currently supported.') - group.add_argument('--rl-durable-rollout-bank', action='store_true', - help='Persist completed rollout groups to a durable, write-through ' - 'ledger so they survive a SIGKILL (the SLURM time limit) and are ' - 'restored at restart instead of regenerated. No-op when unset.') - group.add_argument('--rl-rollout-bank-dir', type=str, default=None, - help='Directory for the durable rollout bank (on Lustre). Defaults to ' - '/rollout_bank so the bank stays coupled to the checkpoint.') - group.add_argument('--rl-rollout-bank-max-bytes', type=int, default=0, - help='Soft cap (bytes) on the rollout bank size; 0 = unbounded. On ' - 'exceed, a warning is logged. Compaction occurs at the next checkpoint ' - 'regardless of the cap and never blocks generation.') - group.add_argument('--grpo-iterations', type=int, default=2, - help="Number of iterations per a GRPO implementation.") - # As in DAPO, we keep upper/lower eps different. - # To have a vanilla GRPO, set them to be the same. - group.add_argument('--grpo-clamp-eps-lower', type=float, default=0.01, - help="Lower GRPO clipping bound.") - group.add_argument('--grpo-clamp-eps-upper', type=float, default=0.01, - help="Upper GRPO clipping bound. In vanilla implementation, equals to the lower one.") - group.add_argument('--grpo-kl-beta', type=float, default=0.001, - help="KL term weight in the GRPO loss.") - group.add_argument('--grpo-entropy-term-weight', type=float, default=0.0, - help="Entropy term weight in GRPO loss.") - group.add_argument('--grpo-filter-groups-with-same-reward', action='store_true', - help="Filter groups with same reward.") - group.add_argument('--langrl-env-config', type=str, default=None, - help="Path to YAML config file for RL environment configuration.") - group.add_argument('--rl-default-temperature', type=float, default=1.0, - help="Default temperature for model inference.") - group.add_argument('--rl-default-top-p', type=float, default=0, - help="Default top-p for model inference.") - group.add_argument('--rl-default-top-k', type=int, default=-1, - help="Default top-k for model inference.") - group.add_argument('--rl-offload-optimizer-during-inference', action='store_true', - help='Offload optimizer state to CPU during inference/rollout to save GPU memory') - group.add_argument('--rl-kv-cache-management-mode', type=str, default='persist', - choices=['persist', 'offload', 'recompute'], - help='KV cache management mode during RL training: ' - 'persist: leave KV cache in GPU memory (default), ' - 'offload: offload KV cache to CPU during training, ' - 'recompute: deallocate KV cache and recompute from scratch each cycle') - group.add_argument('--rl-persist-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='Persist CUDA graphs when the inference engine is suspended. ' - 'If False, CUDA graphs are deleted on suspend and re-captured on resume.') - group.add_argument('--rl-partial-rollouts', action=argparse.BooleanOptionalAction, default=False, - help='Allow inference to continue generating rollouts while training updates ' - 'the policy weights. This enables off-policy training where rollouts may ' - 'be generated with a stale version of the policy. Use ' - '--rl-generation-lag to control the degree of staleness.') - group.add_argument('--rl-inference-logprobs-is-correction', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='If set, use inference logprobs in importance sampling correction of the loss.') - group.add_argument('--rl-importance-sampling-truncation-coef', type=float, default=None, - help="If --inference-logprobs-is-correction is on and this coefficient is set, apply truncation for the IS correction at GRPO loss.") - group.add_argument('--rl-use-sequence-packing', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='Enable sequence packing') - group.add_argument('--rl-sequence-packing-max-sequences-per-bin', type=int, default=50, - help='Maximum number of sequences that can be packed into a single bin. ') - group.add_argument('--rl-sequence-packing-algo', type=str, default='fifo', - choices=['fifo', 'round-robin'], - help='Algorithm for distributing packed bins across ranks. ' - 'fifo: first-in-first-out sequential distribution, ' - 'round-robin: distribute bins cyclically across ranks for better load balancing') - group.add_argument('--rl-training-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, - default=False, - help='If set, do not toggle CUDA graphs on/off between inference and training phases.') - group.add_argument('--rl-inference-tensor-model-parallel-size', type=int, default=None, - help='Degree of tensor model parallelism for inference for RL.') group.add_argument( - '--rl-inference-pipeline-model-parallel-size', - type=int, - default=None, - help='Degree of pipeline model parallelism for inference for RL.', + '--rl-submission-granularity', + type=str, + default="B", + choices=["R", "G", "B"], + help='Granularity for submitting rollout generation work. ' + 'R submits individual rollouts independently while still yielding ' + 'complete rollout groups to training. ' + 'G submits one rollout group at a time. ' + 'B submits grpo_prompts_per_step rollout groups together.', + ) + group.add_argument( + '--rl-consumption-granularity', + type=str, + default="B", + choices=["R", "G", "B"], + help='Granularity for consuming generated rollout groups. ' + 'G consumes groups as they complete. ' + 'B consumes complete trainer batches in submission order. ' + 'R is not currently supported.', + ) + group.add_argument( + '--grpo-iterations', + type=int, + default=2, + help="Number of iterations per a GRPO implementation.", + ) + # As in DAPO, we keep upper/lower eps different. + # To have a vanilla GRPO, set them to be the same. + group.add_argument( + '--grpo-clamp-eps-lower', type=float, default=0.01, help="Lower GRPO clipping bound." + ) + group.add_argument( + '--grpo-clamp-eps-upper', + type=float, + default=0.01, + help="Upper GRPO clipping bound. In vanilla implementation, equals to the lower one.", + ) + group.add_argument( + '--grpo-kl-beta', type=float, default=0.001, help="KL term weight in the GRPO loss." + ) + group.add_argument( + '--grpo-entropy-term-weight', + type=float, + default=0.0, + help="Entropy term weight in GRPO loss.", + ) + group.add_argument( + '--grpo-filter-groups-with-same-reward', + action='store_true', + help="Filter groups with same reward.", + ) + group.add_argument( + '--langrl-env-config', + type=str, + default=None, + help="Path to YAML config file for RL environment configuration.", + ) + group.add_argument( + '--rl-default-temperature', + type=float, + default=1.0, + help="Default temperature for model inference.", + ) + group.add_argument( + '--rl-default-top-p', type=float, default=0, help="Default top-p for model inference." + ) + group.add_argument( + '--rl-default-top-k', type=int, default=-1, help="Default top-k for model inference." + ) + group.add_argument( + '--rl-offload-optimizer-during-inference', + action='store_true', + help='Offload optimizer state to CPU during inference/rollout to save GPU memory', + ) + group.add_argument( + '--rl-kv-cache-management-mode', + type=str, + default='persist', + choices=['persist', 'offload', 'recompute'], + help='KV cache management mode during RL training: ' + 'persist: leave KV cache in GPU memory (default), ' + 'offload: offload KV cache to CPU during training, ' + 'recompute: deallocate KV cache and recompute from scratch each cycle', + ) + group.add_argument( + '--rl-persist-cuda-graphs', + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help='Persist CUDA graphs when the inference engine is suspended. ' + 'If False, CUDA graphs are deleted on suspend and re-captured on resume.', + ) + group.add_argument( + '--rl-partial-rollouts', + action=argparse.BooleanOptionalAction, + default=False, + help='Allow inference to continue generating rollouts while training updates ' + 'the policy weights. This enables off-policy training where rollouts may ' + 'be generated with a stale version of the policy. Use --rl-generation-lag ' + 'to control the degree of staleness.', + ) + group.add_argument( + '--rl-inference-logprobs-is-correction', + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help='If set, use inference logprobs in importance sampling correction of the loss.', + ) + group.add_argument( + '--rl-importance-sampling-truncation-coef', + type=float, + default=None, + help="If --inference-logprobs-is-correction is on and this coefficient is set, apply truncation for the IS correction at GRPO loss.", + ) + group.add_argument( + '--rl-use-sequence-packing', + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help='Enable sequence packing', + ) + group.add_argument( + '--rl-sequence-packing-max-sequences-per-bin', + type=int, + default=50, + help='Maximum number of sequences that can be packed into a single bin. ', + ) + group.add_argument( + '--rl-sequence-packing-algo', + type=str, + default='fifo', + choices=['fifo', 'round-robin'], + help='Algorithm for distributing packed bins across ranks. ' + 'fifo: first-in-first-out sequential distribution, ' + 'round-robin: distribute bins cyclically across ranks for better load balancing', + ) + group.add_argument( + '--rl-training-cuda-graphs', + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help='If set, do not toggle CUDA graphs on/off between inference and training phases.', + ) + group.add_argument( + '--rl-inference-tensor-model-parallel-size', + type=int, + default=None, + help='Degree of tensor model parallelism for inference for RL.', + ) + group.add_argument( + '--rl-inference-pipeline-model-parallel-size', + type=int, + default=None, + help='Degree of pipeline model parallelism for inference for RL.', + ) + group.add_argument( + '--rl-inference-expert-model-parallel-size', + type=int, + default=None, + help='Degree of expert model parallelism for inference for RL.', + ) + group.add_argument( + '--rl-inference-expert-tensor-model-parallel-size', + type=int, + default=None, + help='Degree of expert tensor model parallelism for inference for RL. ' + 'For MoE models, this controls the TP size for expert layers specifically. ' + 'Defaults to training expert_tensor_parallel_size if not specified.', + ) + group.add_argument( + '--rl-inference-model-unified-memory-level', + type=int, + default=0, + choices=[0, 1], + help=( + 'Allocate the separate RL inference model parameters from a unified virtual memory (UVM) ' + 'CUDA mempool. Level 0 disables UVM (default). Level 1 enables UVM allocation so the ' + 'inference model weights can be prefetched to CPU when idle while keeping CUDA-graph-safe ' + 'device pointers.' + ), + ) + group.add_argument( + '--rl-offload-inference-model-weights-when-idle', + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help=( + 'When using a separate RL inference model, offload its weights to CPU when not doing rollout ' + 'inference, and restore to GPU right before inference. Works with two backends: ' + '1) UVM (when --rl-inference-model-unified-memory-level=1), or ' + '2) torch_memory_saver (when UVM is not enabled; requires torch_memory_saver to be installed).' + ), + ) + group.add_argument( + '--refit-method', + type=str, + default='gloo', + choices=['nccl', 'gloo', 'nvshmem'], + help=( + 'Method to refit the model weights between training and inference models during RL. ' + 'nccl: use NCCLCopyService to refit using NCCL; ' + 'gloo: use GlooCopyService over CPU; ' + 'nvshmem: use NVSHMEMCopyService to refit using the NVSHMEM.' + ), + ) + group.add_argument( + '--rl-verify-model-weights-swap', + action=argparse.BooleanOptionalAction, + default=False, + help='If set, verify that the model weights were correctly transferred by comparing forward pass outputs on' + 'the first swap of model weights.', + ) + + group.add_argument( + '--rl-skip-bos-token', + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help='Skip BOS token at the beginning of the sequences. Default is False.', + ) + group.add_argument( + '--rl-profile', + action='store_true', + default=False, + help='Enable RL profiling to collect detailed timer data (JSONL + CSV).', + ) + group.add_argument( + '--rl-profile-dir', + type=str, + default=None, + help='Directory to write RL profiling data. Defaults to {save}/profiles.', + ) + group.add_argument( + '--rl-inference-parsers', + nargs='*', + default=[], + help='List of response parsers to enable for RL inference ' + '(e.g. --rl-inference-parsers deepseek-r1-reasoning qwen3-coder-tool).', + ) + return parser + + +def _add_training_args(parser): + from megatron.training.config import ProfilingConfig, TrainingConfig + + prof_factory = ArgumentGroupFactory(ProfilingConfig) + prof_group = prof_factory.build_group(parser, "profiling") + + train_factory = ArgumentGroupFactory(TrainingConfig) + group = train_factory.build_group(parser, "training") + + group.add_argument( + '--batch-size', + type=int, + default=None, + help='Old batch size parameter, do not use. ' 'Use --micro-batch-size instead', + ) + group.add_argument( + '--recompute-activations', + action='store_true', + help='recompute activation to allow for training ' + 'with larger models, sequences, and batch sizes.', + ) + group.add_argument( + '--no-check-for-nan-in-loss-and-grad', + action='store_false', + help='Check for NaNs in loss and grad', + dest='check_for_nan_in_loss_and_grad', + ) + group.add_argument( + '--check-for-large-grads', + action='store_true', + help='Check for unexpectedly large grads', + dest='check_for_large_grads', + ) + group.add_argument( + '--result-rejected-tracker-filename', + type=str, + default=None, + help='Optional name of file tracking `result_rejected` events.', + ) + group.add_argument( + '--tp-comm-overlap-cfg', + type=str, + default=None, + help='Config file when tp_comm_overlap is enabled.', + ) + + # deprecated + group.add_argument( + '--checkpoint-activations', + action='store_true', + help='Checkpoint activation to allow for training ' + 'with larger models, sequences, and batch sizes.', + ) + group.add_argument( + '--no-masked-softmax-fusion', + action='store_false', + help='Disable fusion of query_key_value scaling, ' 'masking, and softmax.', + dest='masked_softmax_fusion', + ) + group.add_argument( + '--no-bias-gelu-fusion', + action='store_false', + help='Disable bias and gelu fusion.', + dest='bias_gelu_fusion', + ) + group.add_argument( + '--no-bias-swiglu-fusion', + action='store_false', + help='Disable bias and swiglu fusion, the fusion is ' + 'available only when using megatron-core.', + dest='bias_swiglu_fusion', + ) + group.add_argument( + '--no-bias-dropout-fusion', + action='store_false', + help='Disable bias and dropout fusion.', + dest='bias_dropout_fusion', + ) + group.add_argument( + '--no-rope-fusion', + action='store_false', + help='Disable rope fusion, the fusion is available ' 'only when using megatron-core.', + dest='apply_rope_fusion', + ) + group.add_argument( + '--rope-type', + type=str, + default=None, + choices=['rope', 'yarn'], + help='Type of rope to use. Note that MLA takes yarn by default, ' + 'and common attention takes rope by default.', + ) + group.add_argument( + '--use-flash-attn', + action='store_true', + help='use FlashAttention implementation of attention. ' 'https://arxiv.org/abs/2205.14135', + ) + group.add_argument( + '--optimizer', + type=str, + default='adam', + choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap', 'adaptive_muon'], + help='Optimizer function. ' + 'Note: dist_muon is deprecated; use --optimizer muon ' + 'with --use-distributed-optimizer instead.', + ) + group.add_argument( + '--optimizer-cpu-offload', action='store_true', help='Offload optimizer state to CPU' + ) + group.add_argument( + '--optimizer-cuda-graph', action='store_true', help='Enable CUDA graph for optimizer step' + ) + group.add_argument( + '--optimizer-offload-fraction', + type=float, + default=1.0, + help='Fraction used by --optimizer-cpu-offload. This is distinct from ' + '--optimizer-state-offload-fraction, and the two offload modes are mutually exclusive.', + ) + group.add_argument( + '--use-torch-optimizer-for-cpu-offload', + action='store_true', + help="Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu offload mode.", + ) + group.add_argument( + '--overlap-cpu-optimizer-d2h-h2d', + action='store_true', + default=False, + help='Overlap CPU optimizer step, gradients D2H and updated parameters H2D.', + ) + group.add_argument( + '--dump-param-to-param-group-map', + type=str, + default=None, + help="Path to a file containing parameter-to-parameter-group mapping. " + "Provide a JSON file that specifies which parameters belong to which " + "parameter group for global coordination.", + ) + group.add_argument( + '--no-pin-cpu-grads', + action='store_false', + dest='pin_cpu_grads', + help='Disable pinning of CPU memory for gradients.', + ) + group.add_argument( + '--no-pin-cpu-params', + action='store_false', + dest='pin_cpu_params', + help='Disable pinning of CPU memory for parameters.', + ) + group.add_argument( + '--chunked-optimizer-state-offload', + action='store_true', + help='Keep selected optimizer tensor states and master weights in CPU memory between ' + 'updates. GPU tensor-state updates run in bounded chunks, while selected master ' + 'weights use one full restore window.', + ) + group.add_argument( + '--optimizer-state-offload-chunk-size-mb', + type=int, + default=0, + help='Target size of each optimizer tensor-state staging buffer in MiB. Two buffers ' + 'can be live for overlap (roughly 2x this value per active offload manager). Zero ' + 'temporarily restores all selected tensor state for one full GPU update. A single ' + 'atomic parameter state may exceed this soft target and is reported at runtime. This ' + 'value does not bound the selected master-weight window, which is always restored in ' + 'full.', + ) + group.add_argument( + '--optimizer-state-offload-fraction', + type=float, + default=1.0, + help='Fraction of optimizer parameter bundles to offload. Selected bundles always ' + 'include both tensor state and any separate master weight. This is distinct from ' + '--optimizer-offload-fraction, which belongs to --optimizer-cpu-offload; the modes are ' + 'mutually exclusive. Zero disables optimizer state offload even when its flag is set.', + ) + group.add_argument( + '--offload-optimizer-states', + action='store_true', + dest='offload_optimizer_states', + help='Deprecated spelling for --chunked-optimizer-state-offload. It automatically ' + 'enables the replacement; the legacy-equivalent settings are chunk size 0 and offload ' + 'fraction 1.0 unless new tuning arguments are supplied. ' + 'With a nonzero fraction, optimizer-state checkpoint I/O requires --ckpt-format ' + 'torch_dist, optimizer-state saves must be synchronous, and optimizer or full-iteration ' + 'CUDA graphs are unsupported.', + ) + group.add_argument( + '--dataloader-type', + type=str, + default=None, + choices=['single', 'cyclic', 'external'], + help='Single pass vs multiple pass data loader', + ) + group.add_argument( + '--no-persist-layer-norm', + action='store_true', + help='Disable using persistent fused layer norm kernel. ' + 'This kernel supports only a set of hidden sizes. Please ' + 'check persist_ln_hidden_sizes if your hidden ' + 'size is supported.', + ) + group.add_argument( + '--no-gradient-accumulation-fusion', + action='store_false', + help='Disable fusing gradient accumulation to weight ' + 'gradient computation of linear layers', + dest='gradient_accumulation_fusion', + ) + group.add_argument( + '--use-mcore-models', + action='store_true', + dest='deprecated_use_mcore_models', + help='DEPRECATED. Use the implementation from megatron core.' + 'Now ignored and mcore models are the default.', + ) + + return parser + + +def _add_rerun_machine_args(parser): + from megatron.training.config import RerunStateMachineConfig + + rerun_factory = ArgumentGroupFactory(RerunStateMachineConfig, exclude=["check_for_nan_in_loss"]) + group = rerun_factory.build_group(parser, "rerun engine") + + return parser + + +def _add_initialization_args(parser): + from megatron.training.config import RNGConfig + + rng_factory = ArgumentGroupFactory(RNGConfig) + group = rng_factory.build_group(parser, "RNG and initialization") + + group.add_argument( + '--init-method-xavier-uniform', + action='store_true', + help='Enable Xavier uniform parameter initialization', + ) + + return parser + + +def _add_learning_rate_args(parser): + from megatron.training.config import SchedulerConfig + + sched_factory = ArgumentGroupFactory(SchedulerConfig, exclude=["no_weight_decay_cond_type"]) + group = sched_factory.build_group(parser, title="learning rate and weight decay") + + group.add_argument( + '--lr', + type=float, + default=None, + help='Initial learning rate. Depending on decay style ' + 'and initial warmup, the learning rate at each ' + 'iteration would be different.', + ) + group.add_argument( + '--warmup', + type=int, + default=None, + help='Old lr warmup argument, do not use. Use one of the' '--lr-warmup-* arguments above', + ) + group.add_argument( + '--min-lr', + type=float, + default=0.0, + help='Minimum value for learning rate. The scheduler' 'clip values below this threshold.', + ) + group.add_argument( + '--decoupled-lr', + type=float, + default=None, + help='Separate learning rate for the input and output layer', + ) + group.add_argument( + '--decoupled-min-lr', + type=float, + default=None, + help='Minimum value for learning rate for the input and output layer. The scheduler' + 'clip values below this threshold', + ) + group.add_argument( + '--freeze-all-layers', action='store_true', help='Freeze all layers of the model.' + ) + + return parser + + +def _add_checkpointing_args(parser): + from megatron.training.config import CheckpointConfig + + ckpt_factory = ArgumentGroupFactory( + CheckpointConfig, + exclude=[ + "most_recent_k", + "save_tokenizer_assets", + "save_optim", + "save_rng", + "load_optim", + "load_rng", + ], + ) + group = ckpt_factory.build_group(parser, "checkpointing") + + group.add_argument( + '--no-save-optim', action='store_true', default=None, help='Do not save current optimizer.' + ) + group.add_argument( + '--no-save-rng', action='store_true', default=None, help='Do not save current rng state.' + ) + group.add_argument( + '--no-load-optim', + action='store_true', + default=None, + help='Do not load optimizer when loading checkpoint.', + ) + group.add_argument( + '--no-load-rng', + action='store_true', + default=None, + help='Do not load rng state when loading checkpoint.', + ) + group.add_argument( + '--use-dist-ckpt', + action='store_true', + dest='use_dist_ckpt_deprecated', + help='Deprecated: see --ckpt-format.', + ) + group.add_argument( + '--dist-ckpt-format', + dest='dist_ckpt_format_deprecated', + help='Deprecated: see --ckpt-format.', + ) + group.add_argument( + '--dist-ckpt-workers', + type=int, + default=1, + help='Number of workers for distributed checkpointing. ' + 'Only used for async save. ' + 'If set to 1, the checkpointing is performed in a single process.', + ) + group.add_argument( + '--ckpt-fully-parallel-save', + action='store_true', + dest='ckpt_fully_parallel_save_deprecated', + help='Deprecated: see --no-ckpt-fully-parallel-save.', + ) + return parser + + +def _add_mixed_precision_args(parser): + group = parser.add_argument_group(title='mixed precision') + + group.add_argument( + '--grad-reduce-in-bf16', action='store_true', help='Reduce gradients in bfloat16.' + ) + group.add_argument( + '--loss-scale', + type=float, + default=None, + help='Static loss scaling, positive power of 2 ' + 'values can improve fp16 convergence. If None, dynamic' + 'loss scaling is used.', + ) + group.add_argument( + '--initial-loss-scale', + type=float, + default=2**32, + help='Initial loss-scale for dynamic loss scaling.', + ) + group.add_argument( + '--min-loss-scale', + type=float, + default=1.0, + help='Minimum loss scale for dynamic loss scaling.', + ) + group.add_argument( + '--loss-scale-window', + type=float, + default=1000, + help='Window over which to raise/lower dynamic scale.', + ) + group.add_argument( + '--hysteresis', type=int, default=2, help='hysteresis for dynamic loss scaling' + ) + group.add_argument( + '--attention-softmax-in-fp32', + action='store_true', + help='Run attention masking and softmax in fp32.', + ) + group.add_argument( + '--accumulate-allreduce-grads-in-fp32', + action='store_true', + help='Gradient accumulation and all-reduce in fp32.', + ) + group.add_argument( + '--fp16-lm-cross-entropy', + action='store_true', + help='Move the cross entropy unreduced loss calculation' 'for lm head to fp16.', + ) + group.add_argument( + '--reuse-grad-buf-for-mxfp8-param-ag', + action='store_true', + help='If True, reuse the grad buffer for MXFP8 parameter all-gather.', + ) + group.add_argument( + '--mamba-training-ssm-states-dtype', + type=str, + choices=['fp32', 'bf16'], + default=None, + help='Dtype of the materialized inter-chunk SSM states in Mamba training', + ) + + return parser + + +def _add_distributed_args(parser): + from megatron.training.config import DistributedInitConfig + + dist_init_factory = ArgumentGroupFactory(DistributedInitConfig) + group = dist_init_factory.build_group(parser, "distributed init") + + group.add_argument( + '--decoder-first-pipeline-num-layers', + type=int, + default=None, + help=( + 'The number of transformer layers on the first pipeline stage of the decoder. ' + 'Default None is even split of transformer layers across all pipeline stages' + ), + ) + group.add_argument( + '--decoder-last-pipeline-num-layers', + type=int, + default=None, + help=( + 'The number of transformer layers on the last pipeline stage of the decoder. ' + 'Default None is even split of transformer layers across all pipeline stages' + ), + ) + group.add_argument( + '--pipeline-model-parallel-layout', + type=str, + default=None, + help=( + 'A string that describes a custom pipeline model parallel layout. ' + 'e.g., "E|(t|)*3,m|m||L". E, L, t, m denotes embedding, loss, transformer ' + 'decoder layer, and mtp layer, respectively. Stages are split by "|". ' + 'Replicated stages or layers can be described with multiplication. ' + 'Commas can be used cosmetically. ' + 'Default None is not using this argument to set the layout.' + ), + ) + group.add_argument( + '--model-parallel-size', + type=int, + default=None, + help='Old model parallel argument, do not use. Use ' + '--tensor-model-parallel-size instead.', + ) + group.add_argument( + '--num-layers-per-virtual-pipeline-stage', + type=int, + default=None, + help='Number of layers per virtual pipeline stage', + ) + group.add_argument( + '--num-virtual-stages-per-pipeline-rank', + type=int, + default=None, + help='Number of virtual pipeline stages per pipeline parallelism rank', + ) + group.add_argument( + '--no-overlap-p2p-communication', + action='store_false', + help='overlap pipeline parallel communication with forward and backward chunks in 1F1B', + dest='overlap_p2p_comm', + ) + group.add_argument( + '--overlap-grad-reduce', + action='store_true', + default=False, + help='If set, overlap DDP grad reduce.', + ) + group.add_argument( + '--ddp-num-buckets', + type=int, + default=None, + help='Number of buckets for data-parallel communication', + ) + group.add_argument( + '--ddp-bucket-size', + type=int, + default=None, + help='Bucket size for data-parallel communication', + ) + group.add_argument( + '--ddp-pad-buckets-for-high-nccl-busbw', + action='store_true', + default=False, + help='If set, make sure the bucket size is divisible by a large power ' + 'of 2 (2^16) to ensure NCCL collectives have high bus bandwidth at large DP counts, ' + 'since NCCL message size (which for ring algorithms is bucket_size / dp_size) ' + 'apparently needs to be divisible by a power of 2 for high busbw.', + ) + group.add_argument( + '--ddp-reduce-scatter-with-fp32-accumulation', + action='store_true', + default=False, + help='If set, use a reduce-scatter implementation which sends lower-precision ' + 'values over the wire (using an all-to-all to keep total communication overhead in line ' + 'with the standard ring implementation) but performs accumulation locally in FP32.', + ) + group.add_argument( + '--ddp-param-name-patterns-for-fp32-local-accumulation', + nargs='+', + default=[], + help='List of param_name patterns (in Python\'s fnmatch format) to match against ' + 'to do local gradient accumulation in FP32. The special pattern \'all\' matches ' + 'every parameter.', + ) + group.add_argument( + '--ddp-average-in-collective', + action='store_true', + default=False, + help='If set, average directly in data-parallel communication collective.', + ) + group.add_argument( + '--disable-grad-buffers-cpu-backup', + action='store_true', + default=False, + help='If set, allocate DDP gradient buffers in a torch_memory_saver region ' + 'without CPU backup.', + ) + group.add_argument( + '--disable-param-buffers-cpu-backup', + action='store_true', + default=False, + help='If set, allocate DDP parameter buffers in a torch_memory_saver region ' + 'without CPU backup. Only applies when using the distributed optimizer.', + ) + group.add_argument( + '--overlap-param-gather', + action='store_true', + default=False, + help='If set, overlap param all-gather in distributed optimizer.', + ) + group.add_argument( + '--overlap-param-gather-with-optimizer-step', + action='store_true', + default=False, + help='If set, overlap param all-gather of first bucket with optimizer step.', + ) + group.add_argument( + '--no-align-param-gather', + action='store_false', + help='If not set, all PP stages will launch param all-gathers simultaneously. ' + 'Otherwise, each PP stage will independently launch as needed.', + dest='align_param_gather', + ) + group.add_argument( + '--use-distributed-optimizer', action='store_true', help='Use distributed optimizer.' + ) + group.add_argument( + '--use-nccl-ub', + action='store_true', + dest='nccl_ub', + help='Use the userbuffer registration for DP/FSDP communication buffers.' + 'This option will reduce GPU SM usage for the DP/FSDP communication,' + 'which is improving the performance of the overlapped computation.', + ) + group.add_argument( + '--disable-symmetric-registration', + action='store_true', + dest='disable_symmetric_registration', + default=False, + help='Disable symmetric (window) registration for NCCL userbuffer registration.' + 'This option will force to use conventional (local) userbuffer registration when use-nccl-ub is set.', + ) + group.add_argument( + '--fsdp-manual-registration', + action='store_true', + dest='fsdp_manual_registration', + default=False, + help='Manually register the FSDP communication buffers to NCCL user buffer.' + 'This option is only effective when use-megatron-fsdp and use-nccl-ub is set.', + ) + group.add_argument( + '--create-all-gather-group', + action='store_true', + help='Create a separate process group for all-gather operations ' + 'to overlap reduce-scatter and all-gather operations.', + ) + group.add_argument( + '--data-parallel-sharding-strategy', + type=str, + default='optim_grads_params', + choices=['no_shard', 'optim', 'optim_grads', 'optim_grads_params'], + help='Sharding strategy of data parallelism.', + ) + group.add_argument( + '--outer-dp-sharding-strategy', + type=str, + default='no_shard', + choices=['no_shard', 'optim'], + help='Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode. ' + 'Valid values are "no_shard" (DP Replication) and "optim" (Optimizer State Hybrid Sharding). ' + 'The "optim" option is only supported when --data-parallel-sharding-strategy is "optim_grads_params". ' + 'This option is only effective when Hybrid FSDP is enabled (i.e., when dp_outer_dim is not None). ' + 'Default: "no_shard".', + ) + group.add_argument( + '--no-gradient-reduce-div-fusion', + action='store_false', + dest='gradient_reduce_div_fusion', + help='If not set, fuse the division in gradient reduce.', ) group.add_argument( - '--rl-inference-expert-model-parallel-size', - type=int, - default=None, - help='Degree of expert model parallelism for inference for RL.', + '--fsdp-double-buffer', + action='store_true', + help="Enable double buffering for temporary memory needed for Megatron FSDP communications. " + "Double-buffering the communication memory improves memory management efficiency by " + "reusing previously allocated buffers, rather than creating new buffers for each FSDP communication. " + "This is required for user buffer registration and is enabled by default when using NCCL user buffers.", ) group.add_argument( - '--rl-inference-expert-tensor-model-parallel-size', + '--suggested-communication-unit-size', type=int, default=None, - help='Degree of expert tensor model parallelism for inference for RL. ' - 'For MoE models, this controls the TP size for expert layers specifically. ' - 'Defaults to training expert_tensor_parallel_size if not specified.', + help='Specifies the number of elements to communicate at once during FSDP (Fully Sharded Data Parallel) operations. ' + 'This flag also affects FSDP all-gather prefetch behavior. Setting a larger value increases the communication buffer size, ' + 'while a smaller value disables prefetching and may degrade performance. Adjust this value based on your system\'s memory ' + 'and performance requirements.', ) group.add_argument( - '--rl-inference-model-unified-memory-level', + '--keep-fp8-transpose-cache', + action='store_true', + help='If set, keep the fp8 transpose cache when using Megatron FSDP.', + ) + group.add_argument( + '--enable-full-sharding-in-hsdp', + action='store_true', + help='If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel (HSDP) mode.', + ) + group.add_argument( + '--num-distributed-optimizer-instances', type=int, - default=0, - choices=[0, 1], - help=( - 'Allocate the separate RL inference model parameters from a unified virtual memory (UVM) ' - 'CUDA mempool. Level 0 disables UVM (default). Level 1 enables UVM allocation so the ' - 'inference model weights can be prefetched to CPU when idle while keeping CUDA-graph-safe ' - 'device pointers.' - ), + default=1, + help='Number of Distributed Optimizer copies across Data Parallel domain.', ) group.add_argument( - '--rl-offload-inference-model-weights-when-idle', - action=argparse.BooleanOptionalAction, - required=False, - default=False, - help=( - 'When using a separate RL inference model, offload its weights to CPU when not doing rollout ' - 'inference, and restore to GPU right before inference. Works with two backends: ' - '1) UVM (when --rl-inference-model-unified-memory-level=1), or ' - '2) torch_memory_saver (when UVM is not enabled; requires torch_memory_saver to be installed).' - ), + '--torch-fsdp2-no-reshard-after-forward', + action='store_false', + dest='torch_fsdp2_reshard_after_forward', + help='Whether to reshard weights after forward pass when using PyTorch FSDP2. ' + 'Set to enable FSDP ZeRO-2.', ) - group.add_argument('--refit-method', type=str, default='gloo', - choices=['nccl', 'nccl_m2n', 'gloo', 'nvshmem', 'nixl'], - help=('Method to refit model weights. ' - 'nccl: use NCCLCopyService; ' - 'nccl_m2n: use the official NCCL M2N API from a non-RL ' - 'launcher such as the ReFIT benchmark; ' - 'gloo: use GlooCopyService over CPU; ' - 'nvshmem: use NVSHMEMCopyService; ' - 'nixl: use NixlCopyService.')) - group.add_argument('--rl-verify-model-weights-swap', action=argparse.BooleanOptionalAction, default=False, - help='If set, verify that the model weights were correctly transferred by comparing forward pass outputs on' - 'the first swap of model weights.') - - group.add_argument('--rl-skip-bos-token', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='Skip BOS token at the beginning of the sequences. Default is False.') - group.add_argument('--rl-profile', action='store_true', default=False, - help='Enable RL profiling to collect detailed timer data (JSONL + CSV).') - group.add_argument('--rl-profile-dir', type=str, default=None, - help='Directory to write RL profiling data. Defaults to {save}/profiles.') - group.add_argument('--rl-inference-parsers', nargs='*', default=[], - help='List of response parsers to enable for RL inference ' - '(e.g. --rl-inference-parsers deepseek-r1-reasoning qwen3-coder-tool).') - return parser - -def _add_training_args(parser): - from megatron.training.config import ProfilingConfig, TrainingConfig - - prof_factory = ArgumentGroupFactory(ProfilingConfig) - prof_group = prof_factory.build_group(parser, "profiling") - - train_factory = ArgumentGroupFactory(TrainingConfig) - group = train_factory.build_group(parser, "training") - - group.add_argument('--batch-size', type=int, default=None, - help='Old batch size parameter, do not use. ' - 'Use --micro-batch-size instead') - group.add_argument('--recompute-activations', action='store_true', - help='recompute activation to allow for training ' - 'with larger models, sequences, and batch sizes.') - group.add_argument('--no-check-for-nan-in-loss-and-grad', action='store_false', - help='Check for NaNs in loss and grad', - dest='check_for_nan_in_loss_and_grad') - group.add_argument('--check-for-large-grads', action='store_true', - help='Check for unexpectedly large grads', - dest='check_for_large_grads') - group.add_argument('--result-rejected-tracker-filename', type=str, default=None, - help='Optional name of file tracking `result_rejected` events.') - group.add_argument('--tp-comm-overlap-cfg', type=str, default=None, - help='Config file when tp_comm_overlap is enabled.') - - # deprecated - group.add_argument('--checkpoint-activations', action='store_true', - help='Checkpoint activation to allow for training ' - 'with larger models, sequences, and batch sizes.') - group.add_argument('--no-masked-softmax-fusion', - action='store_false', - help='Disable fusion of query_key_value scaling, ' - 'masking, and softmax.', - dest='masked_softmax_fusion') - group.add_argument('--no-bias-gelu-fusion', action='store_false', - help='Disable bias and gelu fusion.', - dest='bias_gelu_fusion') - group.add_argument('--no-bias-swiglu-fusion', action='store_false', - help='Disable bias and swiglu fusion, the fusion is ' - 'available only when using megatron-core.', - dest='bias_swiglu_fusion') - group.add_argument('--no-bias-dropout-fusion', action='store_false', - help='Disable bias and dropout fusion.', - dest='bias_dropout_fusion') - group.add_argument('--no-rope-fusion', action='store_false', - help='Disable rope fusion, the fusion is available ' - 'only when using megatron-core.', - dest='apply_rope_fusion') - group.add_argument('--rope-type', type=str, default=None, - choices=['rope', 'yarn'], - help='Type of rope to use. Note that MLA takes yarn by default, ' - 'and common attention takes rope by default.') - group.add_argument('--use-flash-attn', action='store_true', - help='use FlashAttention implementation of attention. ' - 'https://arxiv.org/abs/2205.14135') - group.add_argument('--optimizer', type=str, default='adam', - choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap', 'adaptive_muon'], - help='Optimizer function. ' - 'Note: dist_muon is deprecated; use --optimizer muon ' - 'with --use-distributed-optimizer instead.') - group.add_argument('--optimizer-cpu-offload', action='store_true', - help='Offload optimizer state to CPU') - group.add_argument('--optimizer-cuda-graph', action='store_true', - help='Enable CUDA graph for optimizer step') - group.add_argument('--optimizer-offload-fraction', type=float, default=1.0, - help='Ratio of optimizer state to offload to CPU') - group.add_argument('--use-torch-optimizer-for-cpu-offload', action='store_true', - help="Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu offload mode.") - group.add_argument('--overlap-cpu-optimizer-d2h-h2d', action='store_true', default=False, - help='Overlap CPU optimizer step, gradients D2H and updated parameters H2D.') - group.add_argument('--dump-param-to-param-group-map', type=str, default=None, - help="Path to a file containing parameter-to-parameter-group mapping. " - "Provide a JSON file that specifies which parameters belong to which " - "parameter group for global coordination.") - group.add_argument('--no-pin-cpu-grads', action='store_false', dest='pin_cpu_grads', - help='Disable pinning of CPU memory for gradients.') - group.add_argument('--no-pin-cpu-params', action='store_false', dest='pin_cpu_params', - help='Disable pinning of CPU memory for parameters.') - group.add_argument('--dataloader-type', type=str, default=None, - choices=['single', 'cyclic', 'external'], - help='Single pass vs multiple pass data loader') - group.add_argument('--no-persist-layer-norm', action='store_true', - help='Disable using persistent fused layer norm kernel. ' - 'This kernel supports only a set of hidden sizes. Please ' - 'check persist_ln_hidden_sizes if your hidden ' - 'size is supported.') - group.add_argument('--no-gradient-accumulation-fusion', - action='store_false', - help='Disable fusing gradient accumulation to weight ' - 'gradient computation of linear layers', - dest='gradient_accumulation_fusion') - group.add_argument('--use-mcore-models', action='store_true', - dest='deprecated_use_mcore_models', - help='DEPRECATED. Use the implementation from megatron core.' - 'Now ignored and mcore models are the default.') - - return parser - - -def _add_rerun_machine_args(parser): - from megatron.training.config import RerunStateMachineConfig - - rerun_factory = ArgumentGroupFactory(RerunStateMachineConfig, exclude=["check_for_nan_in_loss"]) - group = rerun_factory.build_group(parser, "rerun engine") - - return parser - - -def _add_initialization_args(parser): - from megatron.training.config import RNGConfig - - rng_factory = ArgumentGroupFactory(RNGConfig) - group = rng_factory.build_group(parser, "RNG and initialization") - - group.add_argument('--init-method-xavier-uniform', action='store_true', - help='Enable Xavier uniform parameter initialization') - - return parser - - -def _add_learning_rate_args(parser): - from megatron.training.config import SchedulerConfig - - sched_factory = ArgumentGroupFactory(SchedulerConfig, exclude=["no_weight_decay_cond_type"]) - group = sched_factory.build_group(parser, title="learning rate and weight decay") - - group.add_argument('--lr', type=float, default=None, - help='Initial learning rate. Depending on decay style ' - 'and initial warmup, the learning rate at each ' - 'iteration would be different.') - group.add_argument('--warmup', type=int, default=None, - help='Old lr warmup argument, do not use. Use one of the' - '--lr-warmup-* arguments above') - group.add_argument('--min-lr', type=float, default=0.0, - help='Minimum value for learning rate. The scheduler' - 'clip values below this threshold.') - group.add_argument('--decoupled-lr', type=float, default=None, - help='Separate learning rate for the input and output layer') - group.add_argument('--decoupled-min-lr', type=float, default=None, - help='Minimum value for learning rate for the input and output layer. The scheduler' - 'clip values below this threshold') - group.add_argument('--freeze-all-layers', action='store_true', - help='Freeze all layers of the model.') - - return parser - - -def _add_checkpointing_args(parser): - from megatron.training.config import CheckpointConfig - - ckpt_factory = ArgumentGroupFactory(CheckpointConfig, exclude=["most_recent_k", "save_optim", "save_rng", "load_optim", "load_rng"]) - group = ckpt_factory.build_group(parser, "checkpointing") - - group.add_argument('--no-save-optim', action='store_true', default=None, - help='Do not save current optimizer.') - group.add_argument('--no-save-rng', action='store_true', default=None, - help='Do not save current rng state.') - group.add_argument('--no-load-optim', action='store_true', default=None, - help='Do not load optimizer when loading checkpoint.') - group.add_argument('--no-load-rng', action='store_true', default=None, - help='Do not load rng state when loading checkpoint.') - group.add_argument('--override-ckpt-iteration', type=int, default=None, - help='Override the iteration stored in the loaded checkpoint. ' - 'Also resets consumed_train_samples accordingly so the ' - 'data loader replays samples from that iteration onward.') - group.add_argument('--use-dist-ckpt', action='store_true', - dest='use_dist_ckpt_deprecated', - help='Deprecated: see --ckpt-format.') - group.add_argument('--dist-ckpt-format', - dest='dist_ckpt_format_deprecated', - help='Deprecated: see --ckpt-format.') - group.add_argument('--dist-ckpt-workers', type=int, default=1, - help='Number of workers for distributed checkpointing. ' - 'Only used for async save. ' - 'If set to 1, the checkpointing is performed in a single process.') - group.add_argument('--ckpt-fully-parallel-save', action='store_true', - dest='ckpt_fully_parallel_save_deprecated', - help='Deprecated: see --no-ckpt-fully-parallel-save.') - group.add_argument('--ckpt-drop-redundant-extra-state', action='store_true', - default=False, - help='Drop TE `_extra_state` artifacts that carry no ' - 'irreplaceable state (no FP8, or block/current FP8 ' - 'scaling) from the distributed checkpoint, keeping them ' - 'local instead of writing them. Only delayed-scaling ' - '`_extra_state` (amax history + scale) is ever needed and ' - 'it is always persisted regardless of this flag. By ' - 'default (flag off) all `_extra_state` are persisted, ' - 'preserving the previous behavior. Older checkpoints load ' - 'unchanged either way.') - return parser - - -def _add_mixed_precision_args(parser): - group = parser.add_argument_group(title='mixed precision') - - group.add_argument('--grad-reduce-in-bf16', action='store_true', - help='Reduce gradients in bfloat16.') - group.add_argument('--loss-scale', type=float, default=None, - help='Static loss scaling, positive power of 2 ' - 'values can improve fp16 convergence. If None, dynamic' - 'loss scaling is used.') - group.add_argument('--initial-loss-scale', type=float, default=2**32, - help='Initial loss-scale for dynamic loss scaling.') - group.add_argument('--min-loss-scale', type=float, default=1.0, - help='Minimum loss scale for dynamic loss scaling.') - group.add_argument('--loss-scale-window', type=float, default=1000, - help='Window over which to raise/lower dynamic scale.') - group.add_argument('--hysteresis', type=int, default=2, - help='hysteresis for dynamic loss scaling') - group.add_argument('--attention-softmax-in-fp32', action='store_true', - help='Run attention masking and softmax in fp32.') - group.add_argument('--accumulate-allreduce-grads-in-fp32', - action='store_true', - help='Gradient accumulation and all-reduce in fp32.') - group.add_argument('--fp16-lm-cross-entropy', action='store_true', - help='Move the cross entropy unreduced loss calculation' - 'for lm head to fp16.') - group.add_argument('--output-logit-dtype', type=str, choices=['bf16', 'fp32'], default=None, - dest='logit_dtype', - help='Output dtype for the language-model output-layer GEMM. When the ' - 'requested dtype differs from the input dtype, Transformer Engine ' - 'general_gemm is used. By default, logits use the output-layer input dtype.') - group.add_argument('--reuse-grad-buf-for-mxfp8-param-ag', action='store_true', - help='If True, reuse the grad buffer for MXFP8 parameter all-gather.') - group.add_argument('--mamba-training-ssm-states-dtype', type=str, - choices=['fp32', 'bf16'], default=None, - help='Dtype of the materialized inter-chunk SSM states in Mamba training') - - return parser - - -def _add_distributed_args(parser): - from megatron.training.config import DistributedInitConfig - - dist_init_factory = ArgumentGroupFactory(DistributedInitConfig) - group = dist_init_factory.build_group(parser, "distributed init") - - group.add_argument('--decoder-first-pipeline-num-layers', - type=int, default=None, - help=('The number of transformer layers on the first pipeline stage of the decoder. ' - 'Default None is even split of transformer layers across all pipeline stages')) - group.add_argument('--decoder-last-pipeline-num-layers', - type=int, default=None, - help=('The number of transformer layers on the last pipeline stage of the decoder. ' - 'Default None is even split of transformer layers across all pipeline stages')) - group.add_argument('--pipeline-model-parallel-layout', - type=str, default=None, - help=('A string that describes a custom pipeline model parallel layout. ' - 'e.g., "E|(t|)*3,m|m||L". E, L, t, m denotes embedding, loss, transformer ' - 'decoder layer, and mtp layer, respectively. Stages are split by "|". ' - 'Replicated stages or layers can be described with multiplication. ' - 'Commas can be used cosmetically. ' - 'Default None is not using this argument to set the layout.')) - group.add_argument('--model-parallel-size', type=int, default=None, - help='Old model parallel argument, do not use. Use ' - '--tensor-model-parallel-size instead.') - group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None, - help='Number of layers per virtual pipeline stage') - group.add_argument('--num-virtual-stages-per-pipeline-rank', type=int, default=None, - help='Number of virtual pipeline stages per pipeline parallelism rank') - group.add_argument('--no-overlap-p2p-communication', action='store_false', - help='overlap pipeline parallel communication with forward and backward chunks in 1F1B', - dest='overlap_p2p_comm') - group.add_argument('--overlap-grad-reduce', action='store_true', - default=False, help='If set, overlap DDP grad reduce.') - group.add_argument('--ddp-num-buckets', type=int, default=None, - help='Number of buckets for data-parallel communication') - group.add_argument('--ddp-bucket-size', type=int, default=None, - help='Bucket size for data-parallel communication') - group.add_argument('--ddp-pad-buckets-for-high-nccl-busbw', action='store_true', - default=False, help='If set, make sure the bucket size is divisible by a large power ' - 'of 2 (2^16) to ensure NCCL collectives have high bus bandwidth at large DP counts, ' - 'since NCCL message size (which for ring algorithms is bucket_size / dp_size) ' - 'apparently needs to be divisible by a power of 2 for high busbw.') - group.add_argument('--ddp-reduce-scatter-with-fp32-accumulation', action='store_true', - default=False, help='If set, use a reduce-scatter implementation which sends lower-precision ' - 'values over the wire (using an all-to-all to keep total communication overhead in line ' - 'with the standard ring implementation) but performs accumulation locally in FP32.') - group.add_argument('--gtp-remat-reduce-scatter-with-fp32-accumulation', action='store_true', - default=False, help='Same trade as --ddp-reduce-scatter-with-fp32-accumulation, but for ' - 'the wgrad reduce-scatter GTP weight-remat performs over the gtp_remat axis: send ' - 'low-precision values over the wire via an all-to-all and accumulate locally in FP32. ' - 'Independent of the DDP flag (different collective, different process group). Costs one ' - 'extra unsharded-wgrad-sized scratch buffer per in-flight reduce-scatter.') - group.add_argument('--ddp-param-name-patterns-for-fp32-local-accumulation', nargs='+', default=[], - help='List of param_name patterns (in Python\'s fnmatch format) to match against ' - 'to do local gradient accumulation in FP32. The special pattern \'all\' matches ' - 'every parameter.') - group.add_argument('--ddp-average-in-collective', action='store_true', - default=False, help='If set, average directly in data-parallel communication collective.') - group.add_argument('--overlap-param-gather', action='store_true', - default=False, help='If set, overlap param all-gather in distributed optimizer.') - group.add_argument('--overlap-param-gather-with-optimizer-step', action='store_true', - default=False, help='If set, overlap param all-gather of first bucket with optimizer step.') - group.add_argument('--no-align-param-gather', action='store_false', - help='If not set, all PP stages will launch param all-gathers simultaneously. ' - 'Otherwise, each PP stage will independently launch as needed.', - dest='align_param_gather') - group.add_argument('--use-distributed-optimizer', action='store_true', - help='Use distributed optimizer.') - group.add_argument('--megatron-fsdp-version', type=int, default=1, choices=[1, 2], - help='Megatron-FSDP implementation version. Defaults to 1.') - group.add_argument('--no-use-layer-wise-param-layout', - action='store_false', - dest='use_layer_wise_param_layout', - help='Opt out of the precomputed LayerWise param layout. When set, ' - 'falls back to the legacy LayerWise ping-pong path: all params ' - '(including non-Muon embeddings, biases, layernorm) live in a single ' - 'LayerWise buffer and the optimizer uses the allgather_params() codepath. ' - 'The default (precomputed layout) routes non-Muon params through a ' - 'separate DistributedOptimizer with byte-level sharding, which is faster ' - 'and uses less padding but produces different bf16 reduction ordering ' - 'and so will not match legacy-path loss curves bit-for-bit.') - group.add_argument('--use-nccl-ub', action='store_true', dest='nccl_ub', - help='Use the userbuffer registration for DP/FSDP communication buffers.' - 'This option will reduce GPU SM usage for the DP/FSDP communication,' - 'which is improving the performance of the overlapped computation.') - group.add_argument('--disable-symmetric-registration', action='store_true', dest='disable_symmetric_registration', - default=False, help='Disable symmetric (window) registration for NCCL userbuffer registration.' - 'This option will force to use conventional (local) userbuffer registration when use-nccl-ub is set. ' - 'Cannot be combined with --gtp-remat-nccl-ub/--gtp-expert-remat-nccl-ub, which require symmetric windows.') - group.add_argument('--gtp-remat-nccl-ub', action='store_true', dest='gtp_remat_nccl_ub', - default=False, help='Register the wgrad reduce-scatter send buffers with NCCL symmetric ' - 'memory on the GTP group, independent of --use-nccl-ub (which covers the DP group).') - group.add_argument('--gtp-expert-remat-nccl-ub', action='store_true', dest='gtp_expert_remat_nccl_ub', - default=False, help='Like --gtp-remat-nccl-ub but for routed-expert (EGTP) groups.') - group.add_argument('--fsdp-manual-registration', action='store_true', dest='fsdp_manual_registration', - default=False, help='Manually register the FSDP communication buffers to NCCL user buffer.' - 'This option is only effective when use-megatron-fsdp and use-nccl-ub is set.') - group.add_argument('--create-all-gather-group', action='store_true', - help='Create a separate process group for all-gather operations ' - 'to overlap reduce-scatter and all-gather operations.') - group.add_argument('--data-parallel-sharding-strategy', type=str, default='optim_grads_params', - choices=['no_shard', 'optim', 'optim_grads', 'optim_grads_params'], - help='Sharding strategy of data parallelism.') - group.add_argument('--outer-dp-sharding-strategy', type=str, default='no_shard', - choices=['no_shard', 'optim'], - help='Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode. ' - 'Valid values are "no_shard" (DP Replication) and "optim" (Optimizer State Hybrid Sharding). ' - 'The "optim" option is only supported when --data-parallel-sharding-strategy is "optim_grads_params". ' - 'This option is only effective when Hybrid FSDP is enabled (i.e., when dp_outer_dim is not None). ' - 'Default: "no_shard".') - group.add_argument('--no-gradient-reduce-div-fusion', action='store_false', dest='gradient_reduce_div_fusion', - help='If not set, fuse the division in gradient reduce.') - group.add_argument('--fsdp-double-buffer', action='store_true', - help="Enable double buffering for temporary memory needed for Megatron FSDP communications. " - "Double-buffering the communication memory improves memory management efficiency by " - "reusing previously allocated buffers, rather than creating new buffers for each FSDP communication. " - "This is required for user buffer registration and is enabled by default when using NCCL user buffers.") - group.add_argument('--suggested-communication-unit-size', type=int, default=None, - help='Specifies the number of elements to communicate at once during FSDP (Fully Sharded Data Parallel) operations. ' - 'This flag also affects FSDP all-gather prefetch behavior. Setting a larger value increases the communication buffer size, ' - 'while a smaller value disables prefetching and may degrade performance. Adjust this value based on your system\'s memory ' - 'and performance requirements.') - group.add_argument('--keep-fp8-transpose-cache', action='store_true', - help='If set, keep the fp8 transpose cache when using Megatron FSDP.') - group.add_argument('--enable-full-sharding-in-hsdp', action='store_true', - help='If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel (HSDP) mode.') - group.add_argument('--num-distributed-optimizer-instances', type=int, default=1, - help='Number of Distributed Optimizer copies across Data Parallel domain.') - group.add_argument('--torch-fsdp2-no-reshard-after-forward', action='store_false', dest='torch_fsdp2_reshard_after_forward', - help='Whether to reshard weights after forward pass when using PyTorch FSDP2. ' - 'Set to enable FSDP ZeRO-2.') - group.add_argument('--cp-comm-type', nargs='+', type=str, default=["p2p"], - help='Inter-gpu communication type for context parallelism: ' - 'p2p, a2a, allgather or a2a+p2p. If a single string is provided, ' - 'all layers will share the same communication type. Users can also ' - 'specify separated types for each layer like ' - '--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p') - group.add_argument('--max-seqlen-per-dp-cp-rank', type=int, default=None, - help='Maximum sequence length per CP rank. This is used to calculate the ' - 'number of sub-samples assigned to each CP rank when using heterogeneous context parallel.') - group.add_argument('--hybrid-context-parallel', action='store_true', default=False, - help='Enables hybrid context parallel. This is used to balance the workload ' - 'of each CP rank when we use packed samples with variable sequence lengths. ' - 'Requires --max-seqlen-per-dp-cp-rank to be set.') - group.add_argument('--sequence-packing-scheduler', type=str, default=None, choices=['dp_balanced']) - group.add_argument('--fake-process-group', action='store_true', default=False, - help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ + group.add_argument( + '--cp-comm-type', + nargs='+', + type=str, + default=["p2p"], + help='Inter-gpu communication type for context parallelism: ' + 'p2p, a2a, allgather or a2a+p2p. If a single string is provided, ' + 'all layers will share the same communication type. Users can also ' + 'specify separated types for each layer like ' + '--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p', + ) + group.add_argument( + '--fake-process-group', + action='store_true', + default=False, + help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ This is quite useful for profiling memory usage of distributed training with just one GPU. \ - Setting WORLD_SIZE and RANK to the specific values for target distribtued scale.') + Setting WORLD_SIZE and RANK to the specific values for target distribtued scale.', + ) + group.add_argument( + '--use-layer-wise-param-layout', + action='store_true', + dest='use_layer_wise_param_layout', + default=False, + help='Opt INTO the padded shard-aligned LayerWise param layout. The default is the compact ' + 'decoupled layout, where LayerWise (Muon 2D) buffers use a no-padding DDP layout and locally ' + 'disable DistributedOptimizer (all-reduce grads + whole-param ping-pong + allgather_params), ' + 'while sibling buffers keep the byte-level DistributedOptimizer; this avoids the persistent ' + 'dp_size * max(shard_load) padding. The compact layout supports --fp8-param-gather (mxfp8 or ' + 'blockwise): the Muon buffers stay non-DistOpt and the param sync runs through the FP8-aware ' + 'allgather_params ping-pong path (stage fp32 master -> bf16, uneven all-gather, requantize ' + 'per rank); mxfp8 requires --reuse-grad-buf-for-mxfp8-param-ag. Pass this flag to restore the ' + 'padded layout (e.g. for bit-for-bit comparison; it uses a different bf16 reduction ordering).', + ) return parser @@ -3338,138 +4364,254 @@ def _add_tokenizer_args(parser): def _add_data_args(parser): group = parser.add_argument_group(title='data and dataloader') - group.add_argument('--data-path', nargs='*', default=None, - help='The weight and prefix list for a set of train, validation, and test' - 'datasets which split according to --split. The accepted formats are: ' - '(1) a single prefix, ' - '(2) a list of weight prefix pairs e.g. weight1 prefix1 weight2 prefix2, ' - '(3) a list of prefixes e.g. prefix1 prefix2. ' - 'For (3), weights are inferred from the lengths of the contributing datasets. ' - 'This argument is exclusive to the other independent --*-data-path arguments.') - group.add_argument('--phase-transition-iterations', type=str, default=None, - help='Comma-separated list of iterations where phase ' - 'transitions occur. Requires fixed global batch size across phases.') - group.add_argument('--split', type=str, default=None, - help='Comma-separated list of proportions for training,' - ' validation, and test split. For example the split ' - '`90,5,5` will use 90%% of data for training, 5%% for ' - 'validation and 5%% for test.') - group.add_argument('--train-data-path', nargs='*', default=None, - help='The weight and prefix list for an independent train dataset. ' - 'Follows the same pattern rules as --data-path.') - group.add_argument('--valid-data-path', nargs='*', default=None, - help='The weight and prefix list for an independent validation dataset. ' - 'Follows the same pattern rules as --data-path.') - group.add_argument('--test-data-path', nargs='*', default=None, - help='The weight and prefix list for an independent test dataset. ' - 'Follows the same pattern rules as --data-path.') - group.add_argument('--data-args-path', type=str, default=None, - help='Path to data-args. Instead of feeding `--data-path` ' - 'with weighted dataset, we pass in a file path from which ' - 'we read that argument. This is useful when the list of data is ' - 'too big.') - group.add_argument('--per-split-data-args-path', type=str, default=None, - help='Path to per-split-data-args. Instead of feeding ' - '`--(train|valid|test)-data-path` with weighted dataset, ' - 'we pass in a file path from which we read those arguments. ' - 'This is useful when the list of data is too big. Format is a ' - 'json file with `train`, `valid, `test` keys') - group.add_argument('--per-dataset-sequences-path', default=None, - help='Path to a json file with the sequences per dataset. Check the tools/build_sequences_per_dataset.py script to build this file.') - group.add_argument('--dataloader-fast-cache-load', action='store_true', - help='Option to use the fast cache loading path when building the datasets. Requires all the dataset caches to be built and stored in --data-cache-path.') - group.add_argument('--dataloader-defer-npy-index-mmap', action='store_true', - help='Defer the mmap of the dataset indexes (.npy files) until the first access. Requires all the dataset caches to be built and stored in --data-cache-path.') - group.add_argument('--data-cache-path', default=None, - help='Path to a directory to hold cached index files.') - group.add_argument('--no-mmap-bin-files', action='store_false', - help='Disable mmap-ing of .bin files.', - dest='mmap_bin_files') - group.add_argument('--mock-data', action='store_true', - help='Skip data loading and validation and opt for artificial ' - 'generation of mock data when an implementation is available.') - group.add_argument('--seq-length', type=int, default=None, - help='Maximum sequence length to process.') - group.add_argument('--encoder-seq-length', type=int, default=None, - help='Maximum encoder sequence length to process.' - 'This should be exclusive of --seq-length') - group.add_argument('--decoder-seq-length', type=int, default=None, - help="Maximum decoder sequence length to process.") - group.add_argument('--sample-rate', type=float, default=1.0, - help='sample rate for training data. Supposed to be 0 ' - ' < sample_rate < 1') - group.add_argument('--mask-prob', type=float, default=0.15, - help='Probability of replacing a token with mask.') - group.add_argument('--short-seq-prob', type=float, default=0.1, - help='Probability of producing a short sequence.') - group.add_argument('--num-workers', type=int, default=2, - help="Dataloader number of workers.") - group.add_argument('--reset-position-ids', action='store_true', - help='Reset posistion ids after end-of-document token.') - group.add_argument('--reset-attention-mask', action='store_true', - help='Reset self attention mask after ' - 'end-of-document token.') - group.add_argument('--dataloader-inter-document-masking', action='store_true', - help='Return cu_seqlens marking document boundaries ' - 'within each sample so that attention is restricted ' - 'to individual documents.') - group.add_argument('--eod-mask-loss', action='store_true', - help='Mask loss for the end of document tokens.') - group.add_argument('--no-create-attention-mask-in-dataloader', action='store_false', - help='If set, do not create attention_masks in dataloader.', - dest='create_attention_mask_in_dataloader') - group.add_argument('--num-dataset-builder-threads', type=int, default=1, - help='Number of parallel threads per rank for dataset builder') - group.add_argument('--object-storage-cache-path', type=str, default=None, - help='Path to cache index files when using s3 or msc dataloader') - group.add_argument('--mid-level-dataset-surplus', type=float, default=0.005, - help='The sample surplus to build for the mid-level datasets(s)') - group.add_argument('--allow-ambiguous-pad-tokens', action='store_true', - help='Whether to prevent pad tokens already present in the dataset ' - 'from being masked out when the pad token incorrectly shares the same id ' - 'with other special tokens in the tokenizer. Note that this argument has ' - 'no effect when the tokenizer correctly provides a unique id for the pad. ' - 'Masking out such ambiguous pad tokens results in training instability. ' - 'Such a scenario is best resolved by fixing the tokenizer; leaving this ' - 'option as False provides a workaround. ' - 'When left to the default of False, any token ids that collide with the ' - 'pad token id - as provided by the tokenizer - will not be masked out of ' - 'the loss calculation: it cannot be determined whether they are truly pad. ' - 'If instead this argument is set, the training flow will treat all tokens ' - 'that share the same id as the pad token as true pad tokens, potentially ' - 'causing severe training instability.') + group.add_argument( + '--data-path', + nargs='*', + default=None, + help='The weight and prefix list for a set of train, validation, and test' + 'datasets which split according to --split. The accepted formats are: ' + '(1) a single prefix, ' + '(2) a list of weight prefix pairs e.g. weight1 prefix1 weight2 prefix2, ' + '(3) a list of prefixes e.g. prefix1 prefix2. ' + 'For (3), weights are inferred from the lengths of the contributing datasets. ' + 'This argument is exclusive to the other independent --*-data-path arguments.', + ) + group.add_argument( + '--phase-transition-iterations', + type=str, + default=None, + help='Comma-separated list of iterations where phase ' + 'transitions occur. Requires fixed global batch size across phases.', + ) + group.add_argument( + '--split', + type=str, + default=None, + help='Comma-separated list of proportions for training,' + ' validation, and test split. For example the split ' + '`90,5,5` will use 90%% of data for training, 5%% for ' + 'validation and 5%% for test.', + ) + group.add_argument( + '--train-data-path', + nargs='*', + default=None, + help='The weight and prefix list for an independent train dataset. ' + 'Follows the same pattern rules as --data-path.', + ) + group.add_argument( + '--valid-data-path', + nargs='*', + default=None, + help='The weight and prefix list for an independent validation dataset. ' + 'Follows the same pattern rules as --data-path.', + ) + group.add_argument( + '--test-data-path', + nargs='*', + default=None, + help='The weight and prefix list for an independent test dataset. ' + 'Follows the same pattern rules as --data-path.', + ) + group.add_argument( + '--data-args-path', + type=str, + default=None, + help='Path to data-args. Instead of feeding `--data-path` ' + 'with weighted dataset, we pass in a file path from which ' + 'we read that argument. This is useful when the list of data is ' + 'too big.', + ) + group.add_argument( + '--per-split-data-args-path', + type=str, + default=None, + help='Path to per-split-data-args. Instead of feeding ' + '`--(train|valid|test)-data-path` with weighted dataset, ' + 'we pass in a file path from which we read those arguments. ' + 'This is useful when the list of data is too big. Format is a ' + 'json file with `train`, `valid, `test` keys', + ) + group.add_argument( + '--per-dataset-sequences-path', + default=None, + help='Path to a json file with the sequences per dataset. Check the tools/build_sequences_per_dataset.py script to build this file.', + ) + group.add_argument( + '--dataloader-fast-cache-load', + action='store_true', + help='Option to use the fast cache loading path when building the datasets. Requires all the dataset caches to be built and stored in --data-cache-path.', + ) + group.add_argument( + '--dataloader-defer-npy-index-mmap', + action='store_true', + help='Defer the mmap of the dataset indexes (.npy files) until the first access. Requires all the dataset caches to be built and stored in --data-cache-path.', + ) + group.add_argument( + '--data-cache-path', default=None, help='Path to a directory to hold cached index files.' + ) + group.add_argument( + '--no-mmap-bin-files', + action='store_false', + help='Disable mmap-ing of .bin files.', + dest='mmap_bin_files', + ) + group.add_argument( + '--mock-data', + action='store_true', + help='Skip data loading and validation and opt for artificial ' + 'generation of mock data when an implementation is available.', + ) + group.add_argument( + '--seq-length', type=int, default=None, help='Maximum sequence length to process.' + ) + group.add_argument( + '--encoder-seq-length', + type=int, + default=None, + help='Maximum encoder sequence length to process.' + 'This should be exclusive of --seq-length', + ) + group.add_argument( + '--decoder-seq-length', + type=int, + default=None, + help="Maximum decoder sequence length to process.", + ) + group.add_argument( + '--sample-rate', + type=float, + default=1.0, + help='sample rate for training data. Supposed to be 0 ' ' < sample_rate < 1', + ) + group.add_argument( + '--mask-prob', type=float, default=0.15, help='Probability of replacing a token with mask.' + ) + group.add_argument( + '--short-seq-prob', + type=float, + default=0.1, + help='Probability of producing a short sequence.', + ) + group.add_argument('--num-workers', type=int, default=2, help="Dataloader number of workers.") + group.add_argument( + '--reset-position-ids', + action='store_true', + help='Reset posistion ids after end-of-document token.', + ) + group.add_argument( + '--reset-attention-mask', + action='store_true', + help='Reset self attention mask after ' 'end-of-document token.', + ) + group.add_argument( + '--dataloader-inter-document-masking', + action='store_true', + help='Return cu_seqlens marking document boundaries ' + 'within each sample so that attention is restricted ' + 'to individual documents.', + ) + group.add_argument( + '--eod-mask-loss', action='store_true', help='Mask loss for the end of document tokens.' + ) + group.add_argument( + '--no-create-attention-mask-in-dataloader', + action='store_false', + help='If set, do not create attention_masks in dataloader.', + dest='create_attention_mask_in_dataloader', + ) + group.add_argument( + '--num-dataset-builder-threads', + type=int, + default=1, + help='Number of parallel threads per rank for dataset builder', + ) + group.add_argument( + '--object-storage-cache-path', + type=str, + default=None, + help='Path to cache index files when using s3 or msc dataloader', + ) + group.add_argument( + '--mid-level-dataset-surplus', + type=float, + default=0.005, + help='The sample surplus to build for the mid-level datasets(s)', + ) + group.add_argument( + '--allow-ambiguous-pad-tokens', + action='store_true', + help='Whether to prevent pad tokens already present in the dataset ' + 'from being masked out when the pad token incorrectly shares the same id ' + 'with other special tokens in the tokenizer. Note that this argument has ' + 'no effect when the tokenizer correctly provides a unique id for the pad. ' + 'Masking out such ambiguous pad tokens results in training instability. ' + 'Such a scenario is best resolved by fixing the tokenizer; leaving this ' + 'option as False provides a workaround. ' + 'When left to the default of False, any token ids that collide with the ' + 'pad token id - as provided by the tokenizer - will not be masked out of ' + 'the loss calculation: it cannot be determined whether they are truly pad. ' + 'If instead this argument is set, the training flow will treat all tokens ' + 'that share the same id as the pad token as true pad tokens, potentially ' + 'causing severe training instability.', + ) group.add_argument('--fim-data', action='store_true', help='Whether to use the FIM dataset.') - group.add_argument('--fim-rate', type=float, default=0.5, - help='Probability to convert a training sample into a FIM format.') - group.add_argument('--fim-spm-rate', type=float, default=0.5, - help='Probability that the a FIM sample uses the SPM format over the PSM format.') - group.add_argument('--fim-split-sample', type=str, default=None, - help='String around which to split the sample for FIM.') - group.add_argument('--fim-fragment-rate', type=float, default=None, - help='Rate of FIM on each fragment when --fim-split-sample is not None.') - group.add_argument('--fim-no-prefix', type=str, default=None, - help='Do not apply FIM to fragments that start with this prefix') - group.add_argument('--fim-prefix-token', type=str, default='', - help='FIM prefix token') - group.add_argument('--fim-middle-token', type=str, default='', - help='FIM middle token') - group.add_argument('--fim-suffix-token', type=str, default='', - help='FIM suffix token') - group.add_argument('--fim-pad-token', type=str, default='', - help='FIM PAD token') - group.add_argument('--fim-eod-token', type=str, default='<|endoftext|>', - help='FIM EOD token') + group.add_argument( + '--fim-rate', + type=float, + default=0.5, + help='Probability to convert a training sample into a FIM format.', + ) + group.add_argument( + '--fim-spm-rate', + type=float, + default=0.5, + help='Probability that the a FIM sample uses the SPM format over the PSM format.', + ) + group.add_argument( + '--fim-split-sample', + type=str, + default=None, + help='String around which to split the sample for FIM.', + ) + group.add_argument( + '--fim-fragment-rate', + type=float, + default=None, + help='Rate of FIM on each fragment when --fim-split-sample is not None.', + ) + group.add_argument( + '--fim-no-prefix', + type=str, + default=None, + help='Do not apply FIM to fragments that start with this prefix', + ) + group.add_argument( + '--fim-prefix-token', type=str, default='', help='FIM prefix token' + ) + group.add_argument( + '--fim-middle-token', type=str, default='', help='FIM middle token' + ) + group.add_argument( + '--fim-suffix-token', type=str, default='', help='FIM suffix token' + ) + group.add_argument('--fim-pad-token', type=str, default='', help='FIM PAD token') + group.add_argument('--fim-eod-token', type=str, default='<|endoftext|>', help='FIM EOD token') return parser def _add_autoresume_args(parser): group = parser.add_argument_group(title='autoresume') - group.add_argument('--adlr-autoresume', action='store_true', - help='Enable autoresume on adlr cluster.') - group.add_argument('--adlr-autoresume-interval', type=int, default=1000, - help='Intervals over which check for autoresume' - 'termination signal') + group.add_argument( + '--adlr-autoresume', action='store_true', help='Enable autoresume on adlr cluster.' + ) + group.add_argument( + '--adlr-autoresume-interval', + type=int, + default=1000, + help='Intervals over which check for autoresume' 'termination signal', + ) return parser @@ -3478,56 +4620,95 @@ def _add_biencoder_args(parser): group = parser.add_argument_group(title='biencoder') # network size - group.add_argument('--ict-head-size', type=int, default=None, - help='Size of block embeddings to be used in ICT and ' - 'REALM (paper default: 128)') - group.add_argument('--biencoder-projection-dim', type=int, default=0, - help='Size of projection head used in biencoder (paper' - ' default: 128)') - group.add_argument('--biencoder-shared-query-context-model', action='store_true', - help='Whether to share the parameters of the query ' - 'and context models or not') + group.add_argument( + '--ict-head-size', + type=int, + default=None, + help='Size of block embeddings to be used in ICT and ' 'REALM (paper default: 128)', + ) + group.add_argument( + '--biencoder-projection-dim', + type=int, + default=0, + help='Size of projection head used in biencoder (paper' ' default: 128)', + ) + group.add_argument( + '--biencoder-shared-query-context-model', + action='store_true', + help='Whether to share the parameters of the query ' 'and context models or not', + ) # checkpointing - group.add_argument('--ict-load', type=str, default=None, - help='Directory containing an ICTBertModel checkpoint') - group.add_argument('--bert-load', type=str, default=None, - help='Directory containing an BertModel checkpoint ' - '(needed to start ICT and REALM)') + group.add_argument( + '--ict-load', type=str, default=None, help='Directory containing an ICTBertModel checkpoint' + ) + group.add_argument( + '--bert-load', + type=str, + default=None, + help='Directory containing an BertModel checkpoint ' '(needed to start ICT and REALM)', + ) # data - group.add_argument('--titles-data-path', type=str, default=None, - help='Path to titles dataset used for ICT') - group.add_argument('--query-in-block-prob', type=float, default=0.1, - help='Probability of keeping query in block for ' - 'ICT dataset') - group.add_argument('--use-one-sent-docs', action='store_true', - help='Whether to use one sentence documents in ICT') - group.add_argument('--evidence-data-path', type=str, default=None, - help='Path to Wikipedia Evidence frm DPR paper') + group.add_argument( + '--titles-data-path', type=str, default=None, help='Path to titles dataset used for ICT' + ) + group.add_argument( + '--query-in-block-prob', + type=float, + default=0.1, + help='Probability of keeping query in block for ' 'ICT dataset', + ) + group.add_argument( + '--use-one-sent-docs', + action='store_true', + help='Whether to use one sentence documents in ICT', + ) + group.add_argument( + '--evidence-data-path', + type=str, + default=None, + help='Path to Wikipedia Evidence frm DPR paper', + ) # training - group.add_argument('--retriever-report-topk-accuracies', nargs='+', type=int, - default=[], help="Which top-k accuracies to report " - "(e.g. '1 5 20')") - group.add_argument('--retriever-score-scaling', action='store_true', - help='Whether to scale retriever scores by inverse ' - 'square root of hidden size') + group.add_argument( + '--retriever-report-topk-accuracies', + nargs='+', + type=int, + default=[], + help="Which top-k accuracies to report " "(e.g. '1 5 20')", + ) + group.add_argument( + '--retriever-score-scaling', + action='store_true', + help='Whether to scale retriever scores by inverse ' 'square root of hidden size', + ) # faiss index - group.add_argument('--block-data-path', type=str, default=None, - help='Where to save/load BlockData to/from') - group.add_argument('--embedding-path', type=str, default=None, - help='Where to save/load Open-Retrieval Embedding' - ' data to/from') + group.add_argument( + '--block-data-path', type=str, default=None, help='Where to save/load BlockData to/from' + ) + group.add_argument( + '--embedding-path', + type=str, + default=None, + help='Where to save/load Open-Retrieval Embedding' ' data to/from', + ) # indexer - group.add_argument('--indexer-batch-size', type=int, default=128, - help='How large of batches to use when doing indexing ' - 'jobs') - group.add_argument('--indexer-log-interval', type=int, default=1000, - help='After how many batches should the indexer ' - 'report progress') + group.add_argument( + '--indexer-batch-size', + type=int, + default=128, + help='How large of batches to use when doing indexing ' 'jobs', + ) + group.add_argument( + '--indexer-log-interval', + type=int, + default=1000, + help='After how many batches should the indexer ' 'report progress', + ) return parser @@ -3535,144 +4716,299 @@ def _add_vision_args(parser): group = parser.add_argument_group(title="vision") # general vision arguements - group.add_argument('--num-classes', type=int, default=1000, - help='num of classes in vision classificaiton task') - group.add_argument('--img-h', type=int, default=224, - help='Image height for vision classification task') - group.add_argument('--img-w', type=int, default=224, - help='Image height for vision classification task') - group.add_argument('--num-channels', type=int, default=3, - help='Number of channels in input image data') - group.add_argument('--patch-dim', type=int, default=16, - help='patch dimension') - group.add_argument('--classes-fraction', type=float, default=1.0, - help='training with fraction of classes.') - group.add_argument('--data-per-class-fraction', type=float, default=1.0, - help='training with fraction of data per class.') - group.add_argument('--no-data-sharding', action='store_false', - help='Disable data sharding.', - dest='data_sharding') - group.add_argument('--head-lr-mult', type=float, default=1.0, - help='learning rate multiplier for head during finetuning') + group.add_argument( + '--num-classes', type=int, default=1000, help='num of classes in vision classificaiton task' + ) + group.add_argument( + '--img-h', type=int, default=224, help='Image height for vision classification task' + ) + group.add_argument( + '--img-w', type=int, default=224, help='Image height for vision classification task' + ) + group.add_argument( + '--num-channels', type=int, default=3, help='Number of channels in input image data' + ) + group.add_argument('--patch-dim', type=int, default=16, help='patch dimension') + group.add_argument( + '--classes-fraction', type=float, default=1.0, help='training with fraction of classes.' + ) + group.add_argument( + '--data-per-class-fraction', + type=float, + default=1.0, + help='training with fraction of data per class.', + ) + group.add_argument( + '--no-data-sharding', + action='store_false', + help='Disable data sharding.', + dest='data_sharding', + ) + group.add_argument( + '--head-lr-mult', + type=float, + default=1.0, + help='learning rate multiplier for head during finetuning', + ) # pretraining type and backbone selection` - group.add_argument('--vision-pretraining', action='store_true', - help='flag to indicate vision pretraining') - group.add_argument('--vision-pretraining-type', type=str, default='classify', - choices=['classify', 'inpaint', 'dino'], - help='pretraining objectives') - group.add_argument('--vision-backbone-type', type=str, default='vit', - choices=['vit', 'mit', 'swin'], - help='backbone types types') - group.add_argument('--swin-backbone-type', type=str, default='tiny', - choices=['tiny', 'base', 'h3'], - help='pretraining objectives') + group.add_argument( + '--vision-pretraining', action='store_true', help='flag to indicate vision pretraining' + ) + group.add_argument( + '--vision-pretraining-type', + type=str, + default='classify', + choices=['classify', 'inpaint', 'dino'], + help='pretraining objectives', + ) + group.add_argument( + '--vision-backbone-type', + type=str, + default='vit', + choices=['vit', 'mit', 'swin'], + help='backbone types types', + ) + group.add_argument( + '--swin-backbone-type', + type=str, + default='tiny', + choices=['tiny', 'base', 'h3'], + help='pretraining objectives', + ) # inpainting arguments - group.add_argument('--mask-type', type=str, default='random', - choices=['random', 'row'], - help='mask types') - group.add_argument('--mask-factor', type=float, default=1.0, - help='mask size scaling parameter') + group.add_argument( + '--mask-type', type=str, default='random', choices=['random', 'row'], help='mask types' + ) + group.add_argument('--mask-factor', type=float, default=1.0, help='mask size scaling parameter') # dino arguments - group.add_argument('--iter-per-epoch', type=int, default=1250, - help='iterations per epoch') - group.add_argument('--dino-local-img-size', type=int, default=96, - help='Image size for vision classification task') - group.add_argument('--dino-local-crops-number', type=int, default=10, - help='Number of local crops') - group.add_argument('--dino-head-hidden-size', type=int, default=2048, - help='Hidden dimension size in dino head') - group.add_argument('--dino-bottleneck-size', type=int, default=256, - help='Bottle neck dimension in dino head ') - group.add_argument('--dino-freeze-last-layer', type=float, default=1, - help='Freezing last layer weights') - group.add_argument('--dino-norm-last-layer', action='store_true', - help='Disable Norm in last layer.') - group.add_argument('--dino-warmup-teacher-temp', type=float, default=0.04, - help='warump teacher temperature') - group.add_argument('--dino-teacher-temp', type=float, default=0.07, - help='teacher temperature') - group.add_argument('--dino-warmup-teacher-temp-epochs', type=int, default=30, - help='warmup teacher temperaure epochs') + group.add_argument('--iter-per-epoch', type=int, default=1250, help='iterations per epoch') + group.add_argument( + '--dino-local-img-size', + type=int, + default=96, + help='Image size for vision classification task', + ) + group.add_argument( + '--dino-local-crops-number', type=int, default=10, help='Number of local crops' + ) + group.add_argument( + '--dino-head-hidden-size', type=int, default=2048, help='Hidden dimension size in dino head' + ) + group.add_argument( + '--dino-bottleneck-size', type=int, default=256, help='Bottle neck dimension in dino head ' + ) + group.add_argument( + '--dino-freeze-last-layer', type=float, default=1, help='Freezing last layer weights' + ) + group.add_argument( + '--dino-norm-last-layer', action='store_true', help='Disable Norm in last layer.' + ) + group.add_argument( + '--dino-warmup-teacher-temp', type=float, default=0.04, help='warump teacher temperature' + ) + group.add_argument('--dino-teacher-temp', type=float, default=0.07, help='teacher temperature') + group.add_argument( + '--dino-warmup-teacher-temp-epochs', + type=int, + default=30, + help='warmup teacher temperaure epochs', + ) return parser + def _add_moe_args(parser): group = parser.add_argument_group(title="moe") # General arguments - group.add_argument('--num-experts', type=int, default=None, - help='Number of Experts in MoE (None means no MoE)') - group.add_argument('--moe-layer-freq', type=moe_freq_type, default=1, - help='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 string containing a Python list expression that defines a custom pattern, e.g.: ' - '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' - 'where 1 indicates an expert layer and 0 indicates a dense layer. ' - 'Examples: "([0]+[1]*23)": 1 dense layer followed by 23 expert layers, ' - '"([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.') - group.add_argument('--moe-use-upcycling', action='store_true', - help='Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. ' - 'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.') + group.add_argument( + '--num-experts', type=int, default=None, help='Number of Experts in MoE (None means no MoE)' + ) + group.add_argument( + '--moe-layer-freq', + type=moe_freq_type, + default=1, + help='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 string containing a Python list expression that defines a custom pattern, e.g.: ' + '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' + 'where 1 indicates an expert layer and 0 indicates a dense layer. ' + 'Examples: "([0]+[1]*23)": 1 dense layer followed by 23 expert layers, ' + '"([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.', + ) + group.add_argument( + '--moe-use-upcycling', + action='store_true', + help='Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. ' + 'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.', + ) # Router arguments - group.add_argument('--moe-router-load-balancing-type', nargs='+', type=str, - choices=['aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'quantile_balancing', 'none'], - default='aux_loss', - help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE; "quantile_balancing" (QB) uses dual coordinate descent on a per-expert bias to handle load balance internally; "none" implies no load balancing. The default is "aux_loss".') - group.add_argument('--moe-aux-loss-coeff', type=float, nargs='+', default=0.0, - help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.') + group.add_argument( + '--moe-router-load-balancing-type', + nargs='+', + type=str, + choices=['aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'none'], + default='aux_loss', + help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".', + ) + group.add_argument( + '--moe-aux-loss-coeff', + type=float, + nargs='+', + default=0.0, + help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.', + ) # Token dispatcher arguments # MoE communication overlap arguments - group.add_argument('--moe-upcycling-granularity', type=int, default=1, - help='This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. ' - 'For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy.') + group.add_argument( + '--moe-upcycling-granularity', + type=int, + default=1, + help='This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. ' + 'For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy.', + ) return parser + def _add_mla_args(parser): group = parser.add_argument_group(title="mla") - group.add_argument('--q-lora-rank', type=int, default=None, - help="Rank of Query tensor's low rank representation.") - group.add_argument('--kv-lora-rank', type=int, default=32, - help="Rank of Key and Value tensors' low rank representation.") - group.add_argument('--qk-head-dim', type=int, default=128, - help="Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim") - group.add_argument('--qk-pos-emb-head-dim', type=int, default=64, - help="Dimension of the position embedding in the QK projection.") - group.add_argument('--v-head-dim', type=int, default=128, - help="Dimension of the head in the V projection.") - group.add_argument('--rotary-scaling-factor', type=float, default=1.0, - help="Rotary scaling factor for the rotary embeddings.") - group.add_argument('--mscale', type=float, default=1.0, - help="Mscale for YaRN RoPE in multi-latent attention.") - group.add_argument('--mscale-all-dim', type=float, default=0.0, - help="Mscale all dimensions for YaRN RoPE in multi-latent attention.") - group.add_argument('--cache-mla-latents', action='store_true', default=False, - help="If set caches the mla down projected latents with mla flash decode.") + group.add_argument( + '--q-lora-rank', + type=int, + default=None, + help="Rank of Query tensor's low rank representation.", + ) + group.add_argument( + '--kv-lora-rank', + type=int, + default=32, + help="Rank of Key and Value tensors' low rank representation.", + ) + group.add_argument( + '--qk-head-dim', + type=int, + default=128, + help="Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim", + ) + group.add_argument( + '--qk-pos-emb-head-dim', + type=int, + default=64, + help="Dimension of the position embedding in the QK projection.", + ) + group.add_argument( + '--v-head-dim', type=int, default=128, help="Dimension of the head in the V projection." + ) + group.add_argument( + '--rotary-scaling-factor', + type=float, + default=1.0, + help="Rotary scaling factor for the rotary embeddings.", + ) + group.add_argument( + '--original-max-position-embeddings', + type=int, + default=4096, + help="Original maximum position embeddings for the original model, used by yarn.", + ) + group.add_argument( + '--mscale', type=float, default=1.0, help="Mscale for YaRN RoPE in multi-latent attention." + ) + group.add_argument( + '--mscale-all-dim', + type=float, + default=0.0, + help="Mscale all dimensions for YaRN RoPE in multi-latent attention.", + ) + group.add_argument( + '--o-groups', + type=int, + default=8, + help="Number of groups for grouped output (wo_a). 0 = single linear.", + ) + group.add_argument( + '--o-lora-rank', + type=int, + default=1024, + help="Low-rank dimension per group for grouped output (wo_a). Used when o-groups > 0.", + ) + group.add_argument( + '--cache-mla-latents', + action='store_true', + default=False, + help="If set caches the mla down projected latents with mla flash decode.", + ) group.add_argument( '--mla-down-proj-fusion', action='store_true', default=False, help="Enable fused q/kv down-projection and fused input layernorm when backend supports. " - "Otherwise fall back to the unfused MLA.", + "Otherwise fall back to the unfused MLA.", ) return parser + def _add_experimental_attention_variant_args(parser): group = parser.add_argument_group(title="experimental_attention_variant") # Linear attention - group.add_argument('--linear-attention-freq', type=la_freq_type, default=None, - help='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 string containing a Python list expression that defines a custom pattern, e.g.: ' - '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' - 'where 1 indicates an LA layer and 0 indicates a SDPA layer. ' - 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' - '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.') + group.add_argument( + '--linear-attention-freq', + type=la_freq_type, + default=None, + help='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 string containing a Python list expression that defines a custom pattern, e.g.: ' + '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' + 'where 1 indicates an LA layer and 0 indicates a SDPA layer. ' + 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' + '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.', + ) + group.add_argument( + '--csa-compress-ratios', + type=compress_ratios_type, + default=None, + help='Per-layer compress ratios for compressed sparse attention. ' + 'Accepts a string containing a Python list expression, e.g.: ' + '"[0,0,4,128,4,128]" or "([0]+[4,128]*2)*3". ' + 'Each value is the compression ratio for the corresponding ' + 'transformer layer (valid values: 0, 4, 128; 0 = sliding-window-only, the "W" ' + 'hybrid layer symbol). ' + 'For HybridModel with --hybrid-layer-pattern, the preferred compact form has one ' + 'entry per W/C/H attention symbol; legacy zero-padded one-entry-per-hybrid-layer ' + 'lists are also accepted.', + ) + group.add_argument( + '--no-dsa-kernel-fusion', + action='store_false', + default=None, + dest='apply_dsa_kernel_fusion', + help='Deprecated compatibility flag for DSv4 hybrid attention. ' + 'Use --dsa-kernel-backend none instead.', + ) + # Defined manually (and excluded from ArgumentGroupFactory) so that an omitted + # flag is distinguishable from an explicit "none": dsv4_hybrid launches resolve + # an omitted flag to "cudnn" in core_transformer_config_from_args, preserving + # the historical fused-by-default CLI behavior of --no-dsa-kernel-fusion. + group.add_argument( + '--dsa-kernel-backend', + type=str, + choices=['none', 'tilelang', 'cudnn'], + default=None, + help='Fused DSA kernel backend. When omitted, resolves to "cudnn" for ' + 'experimental_attention_variant="dsv4_hybrid" launches and to "none" ' + 'otherwise.', + ) + # Note: --dsa-indexer-{n-heads,head-dim,topk,loss-coeff,use-sparse-loss}, + # --csa-window-size, --csa-compress-rotary-base, --csa-dense-mode are + # auto-generated by ArgumentGroupFactory from TransformerConfig fields + # (none of them are in the exclude list at line 2500-2576). return parser + def _add_heterogeneous_args(parser): """ Heterogeneous models refer to transformer architectures where individual layers can differ @@ -3712,118 +5048,210 @@ def _add_heterogeneous_args(parser): ] """ group = parser.add_argument_group(title="heterogeneous architecture") - group.add_argument('--heterogeneous-layers-config-path', type=str, default=None, - help='Path to json file containing heterogeneous model configuration. ' - 'Use the format of the HuggingFace config files in llama nemotron ' - 'models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.') - group.add_argument('--heterogeneous-layers-config-encoded-json', type=str, default=None, - help='This is encoded json string of the heterogeneous model configuration. Used to keep the content ' - 'of the heterogeneous model specification in args when the model is loaded from a checkpoint. ' - 'Use the format of the HuggingFace config files in llama nemotron ' - 'models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.') + group.add_argument( + '--heterogeneous-layers-config-path', + type=str, + default=None, + help='Path to json file containing heterogeneous model configuration. ' + 'Use the format of the HuggingFace config files in llama nemotron ' + 'models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.', + ) + group.add_argument( + '--heterogeneous-layers-config-encoded-json', + type=str, + default=None, + help='This is encoded json string of the heterogeneous model configuration. Used to keep the content ' + 'of the heterogeneous model specification in args when the model is loaded from a checkpoint. ' + 'Use the format of the HuggingFace config files in llama nemotron ' + 'models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.', + ) return parser + def _add_experimental_args(parser): group = parser.add_argument_group(title='experimental') - group.add_argument('--enable-experimental', action='store_true', - help='Enable experimental features.') - group.add_argument('--spec', type=str, default=None, nargs='*', - help='Specify the pair ' - 'that returns a spec to customize a model, transformer ' - 'block, or transformer layer, depending on the use case.' - 'To use local spec specify local as the argument.' - 'For more details, see the model class, ' - '`transformer_block.py`, or `transformer_layer.py`') - group.add_argument('--hybrid-layer-pattern', type=str, default=None, - help='Specify a hybrid layer pattern using M (mamba), G (gdn), ' - '* (attention), D (dsa), - (mlp), E (moe). Use | to define pipeline ' - 'stage boundaries for flexible virtual pipeline parallel (fVPP). ' - 'Use / to separate MTP patterns. ' - 'Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". ' - 'When this flag is used, it is the sole indicator that a hybrid model ' - 'is being run.') - group.add_argument('--hybrid-override-pattern', type=str, default=None, - help='Deprecated. Use --hybrid-layer-pattern instead. ' - 'If specified, its value will be forwarded to --hybrid-layer-pattern.') - group.add_argument('--yaml-cfg', type=str, default=None, - help = 'Config file to add additional arguments') + group.add_argument( + '--enable-experimental', action='store_true', help='Enable experimental features.' + ) + group.add_argument( + '--spec', + type=str, + default=None, + nargs='*', + help='Specify the pair ' + 'that returns a spec to customize a model, transformer ' + 'block, or transformer layer, depending on the use case.' + 'To use local spec specify local as the argument.' + 'For more details, see the model class, ' + '`transformer_block.py`, or `transformer_layer.py`', + ) + group.add_argument( + '--hybrid-layer-pattern', + type=str, + default=None, + help='Specify a hybrid layer pattern using M (mamba), G (gdn), K (kda), ' + '* (attention), D (dsa), + (mla), - (mlp), E (moe). Use | to define pipeline ' + 'stage boundaries for flexible virtual pipeline parallel (fVPP). ' + 'Use / to separate MTP patterns. ' + 'Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". ' + 'When this flag is used, it is the sole indicator that a hybrid model ' + 'is being run.', + ) + group.add_argument( + '--hybrid-override-pattern', + type=str, + default=None, + help='Deprecated. Use --hybrid-layer-pattern instead. ' + 'If specified, its value will be forwarded to --hybrid-layer-pattern.', + ) + group.add_argument( + '--yaml-cfg', type=str, default=None, help='Config file to add additional arguments' + ) # Args of precision-aware optimizer. - group.add_argument('--use-precision-aware-optimizer', action='store_true', - help='Use the precision-aware optimizer in TransformerEngine, which allows ' - 'setting the main params and optimizer states to lower precision, such as ' - 'fp16, bf16 and fp8.') - group.add_argument('--main-grads-dtype', default='fp32', choices=['fp32', 'bf16'], - help='Dtype of main grads when enabling precision-aware-optimizer.') - group.add_argument('--main-params-dtype', default='fp32', choices=['fp32', 'fp16'], - help='Dtype of main params when enabling precision-aware-optimizer.') - group.add_argument('--exp-avg-dtype', default='fp32', choices=['fp32', 'fp16', 'bf16', 'fp8'], - help='Dtype of exp_avg (1st moment in adam optimizer) when enabling ' - 'precision-aware-optimizer. This dtype is used for storing the ' - 'optimizer state in memory during training but does not affect ' - 'the precision in the kernel computation.') - group.add_argument('--exp-avg-sq-dtype', default='fp32', choices=['fp32', 'fp16', 'bf16', 'fp8'], - help='Dtype of exp_avg_sq (2nd moment in adam optimizer) when enabling ' - 'precision-aware-optimizer. This dtype is used for storing the ' - 'optimizer state in memory during training but does not affect ' - 'the precision in the kernel computation.') + group.add_argument( + '--use-precision-aware-optimizer', + action='store_true', + help='Use the precision-aware optimizer in TransformerEngine, which allows ' + 'setting the main params and optimizer states to lower precision, such as ' + 'fp16, bf16 and fp8.', + ) + group.add_argument( + '--main-grads-dtype', + default='fp32', + choices=['fp32', 'bf16'], + help='Dtype of main grads when enabling precision-aware-optimizer.', + ) + group.add_argument( + '--main-params-dtype', + default='fp32', + choices=['fp32', 'fp16'], + help='Dtype of main params when enabling precision-aware-optimizer.', + ) + group.add_argument( + '--exp-avg-dtype', + default='fp32', + choices=['fp32', 'fp16', 'bf16', 'fp8'], + help='Dtype of exp_avg (1st moment in adam optimizer) when enabling ' + 'precision-aware-optimizer. This dtype is used for storing the ' + 'optimizer state in memory during training but does not affect ' + 'the precision in the kernel computation.', + ) + group.add_argument( + '--exp-avg-sq-dtype', + default='fp32', + choices=['fp32', 'fp16', 'bf16', 'fp8'], + help='Dtype of exp_avg_sq (2nd moment in adam optimizer) when enabling ' + 'precision-aware-optimizer. This dtype is used for storing the ' + 'optimizer state in memory during training but does not affect ' + 'the precision in the kernel computation.', + ) # Megatron-FSDP Arguments - group.add_argument('--megatron-fsdp-main-params-dtype', default='fp32', choices=['fp32', 'bf16', 'fp16', 'auto'], - help="Data type for the main weight buffer utilized for distributed optimization " - "and quantization with Megatron-FSDP. If 'auto', then the native model parameter " - "data-type will be used for the main weight data-type. Replaces --main-params-dtype.") - group.add_argument('--megatron-fsdp-main-grads-dtype', default='auto', choices=['fp32', 'bf16', 'fp16', 'auto'], - help="Data type for the main gradient buffer utilized for distributed optimization " - "with Megatron-FSDP. If 'auto', then the native model gradient data-type will " - "be used for the main gradient / accumulation data-type. Replaces --main-grads-dtype.") - group.add_argument("--megatron-fsdp-grad-comm-dtype", default='auto', choices=['fp32', 'fp16', 'bf16', 'auto'], - help="When using Megatron-FSDP, this controls the data-type used when communicating " - "model gradients during FSDP. If 'auto', then the main gradient data-type will " - "be used for the gradient communication / reduction data-type. When using NCCL " - "v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels.") - group.add_argument( - '--megatron-fsdp-enable-fine-grained-param-gather', - action='store_true', - default=False, - dest='megatron_fsdp_enable_fine_grained_param_gather', - help=( - 'If set, enables fine-grained parameter gathering for Megatron-FSDP. ' - 'This allows greater overlap between parameter all-gather operations and ' - 'forward computation, at the cost of additional communication calls. ' - 'For MXFP8, this helps save memory during fine-grained activation ' - 'recomputation, because MXFP8 forward and backward passes use different ' - 'parameter representations (rowwise data for forward, colwise data for ' - 'backward). Only the rowwise parameters of modules involved in ' - 'recomputation will be unsharded.' - ), - ) + group.add_argument( + '--megatron-fsdp-main-params-dtype', + default='fp32', + choices=['fp32', 'bf16', 'fp16', 'auto'], + help="Data type for the main weight buffer utilized for distributed optimization " + "and quantization with Megatron-FSDP. If 'auto', then the native model parameter " + "data-type will be used for the main weight data-type.", + ) + group.add_argument( + '--megatron-fsdp-main-grads-dtype', + default='auto', + choices=['fp32', 'bf16', 'fp16', 'auto'], + help="Data type for the main gradient buffer utilized for distributed optimization " + "with Megatron-FSDP. If 'auto', then the native model gradient data-type will " + "be used for the main gradient / accumulation data-type.", + ) + group.add_argument( + "--megatron-fsdp-grad-comm-dtype", + default='auto', + choices=['fp32', 'fp16', 'bf16', 'auto'], + help="When using Megatron-FSDP, this controls the data-type used when communicating " + "model gradients during FSDP. If 'auto', then the main gradient data-type will " + "be used for the gradient communication / reduction data-type. When using NCCL " + "v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels.", + ) + group.add_argument( + '--megatron-fsdp-prefetch-recompute-forward-weights', + action='store_true', + default=False, + dest='megatron_fsdp_prefetch_recompute_forward_weights', + help=( + 'If set, Megatron-FSDP prefetches rowwise weights needed by activation ' + 'recomputation during backward before prefetching backward transpose ' + 'weights.' + ), + ) + group.add_argument( + '--megatron-fsdp-cache-param-bucket-views', + action='store_true', + default=False, + dest='megatron_fsdp_cache_param_bucket_views', + help=( + 'If set, Megatron-FSDP caches parameter bucket views to reduce repeated ' + 'Python-side view setup when attaching module parameters to all-gather buckets.' + ), + ) + group.add_argument( + '--megatron-fsdp-enable-fine-grained-param-gather', + action='store_true', + default=False, + dest='megatron_fsdp_enable_fine_grained_param_gather', + help=( + 'If set, enables fine-grained parameter gathering for Megatron-FSDP. ' + 'This allows greater overlap between parameter all-gather operations and ' + 'forward computation, at the cost of additional communication calls. ' + 'For MXFP8, this helps save memory during fine-grained activation ' + 'recomputation, because MXFP8 forward and backward passes use different ' + 'parameter representations (rowwise data for forward, colwise data for ' + 'backward). Only the rowwise parameters of modules involved in ' + 'recomputation will be unsharded.' + ), + ) - group.add_argument("--megatron-fsdp-max-pool-double-buffer", action='store_true', - help="When using Megatron-FSDP double buffering, use the MaxPoolAllocator instead of " - "the FixedPoolAllocator to support asymmetrical FSDP unit configurations. Will " - "increase memory overhead to recycle buffers that fit all FSDP units. Enables " - "NCCL user buffer registration and CUDA graph replay for mixed-arch models.") - group.add_argument("--fsdp-db-use-persist-buf-on-alloc-fail", action='store_true', - help="When using Megatron-FSDP double buffering, persist non-unit modules that " - "are not included in the symmetric buffer pool. May be necessary for NCCL " - "UBR or CUDA Graphs on hybrid architectures.") + group.add_argument( + "--megatron-fsdp-max-pool-double-buffer", + action='store_true', + help="When using Megatron-FSDP double buffering, use the MaxPoolAllocator instead of " + "the FixedPoolAllocator to support asymmetrical FSDP unit configurations. Will " + "increase memory overhead to recycle buffers that fit all FSDP units. Enables " + "NCCL user buffer registration and CUDA graph replay for mixed-arch models.", + ) + group.add_argument( + "--fsdp-db-use-persist-buf-on-alloc-fail", + action='store_true', + help="When using Megatron-FSDP double buffering, persist non-unit modules that " + "are not included in the symmetric buffer pool. May be necessary for NCCL " + "UBR or CUDA Graphs on hybrid architectures.", + ) return parser def _add_msc_args(parser): group = parser.add_argument_group(title="msc") - group.add_argument('--enable-msc', default=False, action='store_true', dest='enable_msc', - help='Enable the usage of Multi-Storage Client (MSC) in Megatron Core. ' - 'Disabled by default; pass this flag to opt in.') - group.add_argument('--disable-msc', default=False, action='store_true', - dest='disable_msc_deprecated', - help='[DEPRECATED] MSC is disabled by default; this flag is a no-op ' - 'and will be removed in a future release.') + group.add_argument( + '--enable-msc', + default=False, + action='store_true', + dest='enable_msc', + help='Enable the usage of Multi-Storage Client (MSC) in Megatron Core. ' + 'Disabled by default; pass this flag to opt in.', + ) + group.add_argument( + '--disable-msc', + default=False, + action='store_true', + dest='disable_msc_deprecated', + help='[DEPRECATED] MSC is disabled by default; this flag is a no-op ' + 'and will be removed in a future release.', + ) return parser + def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): """Add quant-specific arguments to the main parser @@ -3854,70 +5282,155 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): ) return parser + def _add_sft_args(parser): group = parser.add_argument_group(title='sft') group.add_argument('--sft', action="store_true", help='Megatron SFT training') - group.add_argument('--sft-tokenizer-prompt-format', type=str, default="nemotron-h-aligned", - help='SFT prompt format.') - group.add_argument('--sft-mock-dataset-config-json', type=str, default=None, - help='This config provides the necessary information for the mock ' - 'dataset. Accepts either an inline JSON literal or a path to a JSON ' - 'file containing the same schema. You can either specify a CSV file ' - 'that contains sequence lengths, where each line stores the length of ' - 'a sequence, for example: {"mode":"file","path":"/path/to/file"}. ' - 'Alternatively, you can specify a distribution (currently only ' - 'supporting lognormal distribution) along with the required ' - 'parameters, for example, {"mode":"distribution","type":"lognormal",' - '"min_seq_len":1024,"max_seq_len":2048,"mean_seq_len":1536,' - '"lognormal_sigma":1.1}, where sigma controls the variability of the ' - 'lognormal distribution. If not specified and --mock-data is set, ' - 'defaults to a lognormal distribution with min_seq_len=seq_length//2, ' - 'max_seq_len=seq_length, mean_seq_len=seq_length*3//4, ' - 'lognormal_sigma=1.1.') + group.add_argument( + '--sft-tokenizer-prompt-format', + type=str, + default="nemotron-h-aligned", + help='SFT prompt format.', + ) + group.add_argument( + '--sft-mock-dataset-config-json', + type=str, + default=None, + help='This config provides the necessary information for the mock dataset. ' + 'Accepts either an inline JSON literal or a path to a JSON file containing ' + 'the same schema. You can either specify a CSV file that contains sequence lengths, ' + 'where each line stores the length of a sequence, for example: ' + '{"mode":"file","path":"/path/to/file"}. Alternatively, you can specify a distribution ' + '(currently only supporting lognormal distribution) along with the required parameters, ' + 'for example, {"mode":"distribution","type":"lognormal","min_seq_len":1024,' + '"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, where sigma controls ' + 'the variability of the lognormal distribution. ' + 'If not specified and --mock-data is set, defaults to a lognormal distribution with ' + 'min_seq_len=seq_length//2, max_seq_len=seq_length, mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.', + ) + return parser + + +def _add_varlen_dataset_args(parser): + group = parser.add_argument_group(title='varlen dataset') + group.add_argument( + '--use-varlen-dataset', + action="store_true", + help='Train with VarlenDataset, a variable-length packed (THD) dataset ' + 'that consumes instruction-tuning data from a HuggingFace Hub repo id, ' + 'a local parquet file, or a local jsonl file. Schema (alpaca / sharegpt ' + '/ openai-messages) is auto-detected from the dataset columns. ' + 'Mutually exclusive with --sft. Auto-picks a sequence packing ' + 'scheduler when none is given: ``dp_balanced`` by default, ' + '``default_dynamic_cp`` when ``--dynamic-context-parallel`` is set. ' + 'Combine with --mock-data for a synthetic lognormal sequence-length ' + 'distribution; see --varlen-mock-dataset-config-json.', + ) + group.add_argument( + '--varlen-sbhd-validation', + action="store_true", + help='Reference SBHD mode for THD numerical verification. When set, ' + 'VarlenDataset emits SBHD-style samples right-padded to ' + '--seq-length (no cu_seqlens, no packing scheduler), so the run can ' + 'be compared against the THD path to validate correctness. ' + 'Incompatible with --dynamic-context-parallel and ' + '--sequence-packing-scheduler.', + ) + group.add_argument( + '--varlen-mock-dataset-config-json', + type=str, + default=None, + help='Mock-dataset config for --use-varlen-dataset --mock-data. ' + 'Accepts either an inline JSON literal or a path to a JSON file containing ' + 'the same schema as --sft-mock-dataset-config-json: either ' + '{"mode":"file","path":"/path/to/lengths.csv"}, ' + '{"mode":"distribution","type":"lognormal","min_seq_len":1024,' + '"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, or ' + '{"mode":"verification","data_path":"/prefix/of/IndexedDataset"}. ' + 'If not specified, defaults to a lognormal distribution with ' + 'min_seq_len=seq_length//2, max_seq_len=seq_length, ' + 'mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.', + ) return parser + def _add_logits_distillation_args(parser): group = parser.add_argument_group(title='Logits Distillation') - group.add_argument('--logits-save-top-k', type=int, default=None, - help='Number of top logits to save.') - group.add_argument('--logits-save-top-p', type=float, default=None, - help='Top-P (nucleus) threshold applied after top-K ' - 'selection when saving logits. Only the smallest ' - 'set of entries whose cumulative probability mass ' - 'reaches this threshold is kept. Must be in (0, 1].') - group.add_argument('--logits-save-top-p-min-k', type=int, default=1, - help='Minimum number of entries kept per token when ' - 'top-P masking is active, regardless of ' - 'cumulative mass. Default: 1.') - group.add_argument('--logits-save-dir', type=str, default=None, - help='Directory to save logits.') - group.add_argument('--logits-save-dtype', type=str, default='fp16', - choices=['fp16', 'bf16', 'fp32'], - help='Dtype for on-disk top-K log-probabilities.') - group.add_argument('--logits-load-dir', type=str, default=None, - help='Directory to load logits.') - group.add_argument('--logits-load-decode-threads', type=int, default=4, - help='Number of decode threads for cached-logits zstd ' - 'decompression and torch.load processing.') - group.add_argument('--logits-load-prefetch-factor', type=int, default=3, - help='PyTorch DataLoader prefetch factor for decoded ' - 'cached-logits iterations. (Non-MSC only)') - group.add_argument('--logits-load-msc-prefetch-depth', type=int, default=2, - help='For MSC/object-storage logits tar shards, number ' - 'of whole tar shards to prefetch into the MSC ' - 'cache ahead of sequential tar consumption.') - group.add_argument('--logits-load-kd-loss-alpha', type=float, default=1.0, - help='KD loss alpha for loading logits. Total loss is calculated as ' - 'alpha * kd_loss + (1 - alpha) * lm_loss.') - group.add_argument('--logits-load-ignore-errors', action='store_true', - default=False, - help='When set, KD loss errors are logged as warnings and ' - 'training falls back to LM-only loss instead of crashing.') + group.add_argument( + '--logits-save-top-k', type=int, default=None, help='Number of top logits to save.' + ) + group.add_argument( + '--logits-save-top-p', + type=float, + default=None, + help='Top-P (nucleus) threshold applied after top-K ' + 'selection when saving logits. Only the smallest ' + 'set of entries whose cumulative probability mass ' + 'reaches this threshold is kept. Must be in (0, 1].', + ) + group.add_argument( + '--logits-save-top-p-min-k', + type=int, + default=1, + help='Minimum number of entries kept per token when ' + 'top-P masking is active, regardless of ' + 'cumulative mass. Default: 1.', + ) + group.add_argument( + '--logits-save-dir', type=str, default=None, help='Directory to save logits.' + ) + group.add_argument( + '--logits-save-dtype', + type=str, + default='fp16', + choices=['fp16', 'bf16', 'fp32'], + help='Dtype for on-disk top-K log-probabilities.', + ) + group.add_argument( + '--logits-load-dir', type=str, default=None, help='Directory to load logits.' + ) + group.add_argument( + '--logits-load-decode-threads', + type=int, + default=4, + help='Number of decode threads for cached-logits zstd ' + 'decompression and torch.load processing.', + ) + group.add_argument( + '--logits-load-prefetch-factor', + type=int, + default=3, + help='PyTorch DataLoader prefetch factor for decoded ' + 'cached-logits iterations. (Non-MSC only)', + ) + group.add_argument( + '--logits-load-msc-prefetch-depth', + type=int, + default=2, + help='For MSC/object-storage logits tar shards, number ' + 'of whole tar shards to prefetch into the MSC ' + 'cache ahead of sequential tar consumption.', + ) + group.add_argument( + '--logits-load-kd-loss-alpha', + type=float, + default=1.0, + help='KD loss alpha for loading logits. Total loss is calculated as ' + 'alpha * kd_loss + (1 - alpha) * lm_loss.', + ) + group.add_argument( + '--logits-load-ignore-errors', + action='store_true', + default=False, + help='When set, KD loss errors are logged as warnings and ' + 'training falls back to LM-only loss instead of crashing.', + ) return parser def _add_fault_injector_args(parser): from megatron.training.config import FaultInjectorConfig + ArgumentGroupFactory(FaultInjectorConfig).build_group(parser, "fault injector") return parser diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index e0fe107302c..a9f6fc97146 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -7,7 +7,6 @@ import multiprocessing import os import random -import re import shutil import sys import threading @@ -15,7 +14,7 @@ from argparse import Namespace from datetime import datetime from enum import Enum, auto -from logging import DEBUG, getLogger +from logging import getLogger from pathlib import Path from time import time from typing import Any, Dict, List, Optional, Union @@ -24,15 +23,8 @@ import torch from torch.distributed.checkpoint import FileSystemReader, default_planner -try: - from nemo.lens.helpers import managed_span as _otel_managed_span -except ImportError: - from megatron.core.telemetry.fallbacks import managed_span as _otel_managed_span - from megatron.core import dist_checkpointing, mpu, tensor_parallel -from megatron.core._rank_utils import safe_get_rank as get_rank_safe -from megatron.core.dist_checkpointing.dict_utils import dict_list_map_inplace -from megatron.core.dist_checkpointing.mapping import LocalNonpersistentObject, ShardedObject +from megatron.core.dist_checkpointing.mapping import ShardedObject from megatron.core.dist_checkpointing.strategies.async_utils import _disable_gc from megatron.core.dist_checkpointing.strategies.fully_parallel import ( FullyParallelLoadStrategyWrapper, @@ -43,23 +35,18 @@ TorchDistSaveShardedStrategy, get_async_strategy, ) -from megatron.core.msc_utils import MultiStorageClientFeature, maybe_msc +from megatron.core.msc_utils import MultiStorageClientFeature, open_file from megatron.core.num_microbatches_calculator import update_num_microbatches from megatron.core.optimizer import DistributedOptimizer -from megatron.core.post_training.modelopt.checkpointing import save_modelopt_state, save_sharded_modelopt_state from megatron.core.rerun_state_machine import get_rerun_state_machine -from megatron.core.tokenizers import MegatronTokenizer from megatron.core.utils import get_pg_rank, get_pg_size, unwrap_model -from megatron.training.argument_utils import _default_config_from_args -from megatron.training.config import TokenizerConfig -from megatron.training.global_vars import get_tokenizer from ..core.dist_checkpointing.utils import _clean_metadata_for_serialization from . import ft_integration, wandb_utils from .async_utils import get_save_and_finalize_callbacks, is_empty_async_queue, schedule_async_save from .global_vars import get_args from .one_logger_utils import on_save_checkpoint_start, on_save_checkpoint_success -from .utils import append_to_progress_log, is_last_rank, print_rank_0, print_rank_last, warn_rank_0 +from .utils import append_to_progress_log, is_last_rank, print_rank_0 try: from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( @@ -68,11 +55,9 @@ from megatron.core.transformer.fsdp_dtensor_checkpoint import ( handle_experts_in_state_dict, handle_fp8_extra_state_case, - handle_mla_down_proj_in_state_dict, - handle_mtp_in_state_dict, + handle_gdn_in_state_dict, handle_swiglu_in_state_dict, print_diff_in_state_dicts, - validate_fsdp_dtensor_model_load, ) HAVE_MEGATRON_FSDP = True @@ -82,6 +67,8 @@ # [ModelOpt]: Import try: + from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state + from megatron.post_training.utils import print_distributed_quant_summary has_nvidia_modelopt = True @@ -99,22 +86,6 @@ _deletion_processes = [] -def _maybe_compact_rollout_bank(iteration): - """Compact the rollout bank without making checkpointing depend eagerly on RL.""" - from megatron.rl.rl_utils import maybe_compact_rollout_bank - - maybe_compact_rollout_bank(iteration) - - -def _register_rollout_bank_compaction(async_save_request, iteration): - """Compact the rollout bank after an asynchronous checkpoint becomes durable.""" - - def rollout_bank_finalize_fn(iteration=iteration): - _maybe_compact_rollout_bank(iteration) - - async_save_request.add_finalize_fn(rollout_bank_finalize_fn) - - def finalize_deletion_processes(blocking=False): """Clean up deletion processes to prevent zombie processes. @@ -132,13 +103,9 @@ def finalize_deletion_processes(blocking=False): finished = [] for proc in _deletion_processes: if not proc.is_alive() or blocking: - if logger.isEnabledFor(DEBUG): - logger.debug( - "Joining deletion process %s (blocking=%s, is_alive=%s)", - proc.pid, - blocking, - proc.is_alive(), - ) + logger.debug( + f"Joining deletion process {proc.pid} (blocking={blocking}, is_alive={proc.is_alive()})" + ) proc.join() finished.append(proc) for proc in finished: @@ -148,7 +115,7 @@ def finalize_deletion_processes(blocking=False): def set_checkpoint_version(value): global _CHECKPOINT_VERSION if _CHECKPOINT_VERSION is not None: - assert _CHECKPOINT_VERSION == value, 'checkpoint versions do not match' + assert _CHECKPOINT_VERSION == value, "checkpoint versions do not match" _CHECKPOINT_VERSION = value @@ -173,15 +140,12 @@ def get_loaded_iteration(): return _LOADED_ITERATION -def check_checkpoint_args(checkpoint_args, skip_args: set[str] | None = None): +def check_checkpoint_args(checkpoint_args): """Ensure fixed arguments for a model are the same for the input arguments and the one retrieved from checkpoint.""" args = get_args() - skip_args = skip_args or set() def _compare(arg_name, old_arg_name=None, default=None): - if arg_name in skip_args: - return if old_arg_name is not None: ckpt_arg_name = old_arg_name else: @@ -192,17 +156,14 @@ def _compare(arg_name, old_arg_name=None, default=None): checkpoint_value = getattr(checkpoint_args, ckpt_arg_name) args_value = getattr(args, arg_name) error_message = ( - '{} value from checkpoint ({}) is not equal to the input argument value ({}).'.format( - arg_name, checkpoint_value, args_value - ) + '{} value from checkpoint ({}) is not equal to the ' + 'input argument value ({}).'.format(arg_name, checkpoint_value, args_value) ) assert checkpoint_value == args_value, error_message _compare('num_layers') _compare('hidden_size') _compare('num_attention_heads') - if hasattr(args, 'gdp_num_householder'): - _compare('gdp_num_householder', default=3) _compare('add_position_embedding', default=True) if args.vocab_file: _compare('max_position_embeddings') @@ -221,10 +182,22 @@ def _compare(arg_name, old_arg_name=None, default=None): _compare('pipeline_model_parallel_size') +def isfile(filename) -> bool: + if MultiStorageClientFeature.is_enabled(): + msc = MultiStorageClientFeature.import_package() + return msc.os.path.isfile(filename) + else: + return os.path.isfile(filename) + + def ensure_directory_exists(filename, check_parent=True): """Build filename's path if it does not already exists.""" dirname = os.path.dirname(filename) if check_parent else filename - maybe_msc.os.makedirs(dirname, exist_ok=True) + if MultiStorageClientFeature.is_enabled(): + msc = MultiStorageClientFeature.import_package() + msc.os.makedirs(dirname, exist_ok=True) + else: + os.makedirs(dirname, exist_ok=True) def get_checkpoint_name( @@ -237,7 +210,7 @@ def get_checkpoint_name( expert_parallel=None, expert_rank=None, return_base_dir=False, - basename='model_optim_rng.pt', + basename="model_optim_rng.pt", ): """Determine the directory name for this rank's checkpoint.""" if release: @@ -276,25 +249,25 @@ def get_checkpoint_name( return os.path.join(common_path, basename) -def get_load_checkpoint_path_by_args(args, load_arg='load'): +def get_load_checkpoint_path_by_args(args, load_arg="load"): """Get the checkpoint path based on the arguments.""" load_dir = getattr(args, load_arg) iteration, release = -1, False tracker_filename = 'because load directory is not defined' if load_dir is not None: tracker_filename = get_checkpoint_tracker_filename(load_dir) - if maybe_msc.os.path.isfile(tracker_filename): + if isfile(tracker_filename): iteration, release = read_metadata(tracker_filename) # Allow user to specify the loaded iteration. - if getattr(args, 'ckpt_step', None): + if getattr(args, "ckpt_step", None): iteration = args.ckpt_step return get_checkpoint_name(load_dir, iteration, release, return_base_dir=True) def get_distributed_optimizer_checkpoint_name(model_checkpoint_name): - return os.path.join(os.path.dirname(model_checkpoint_name), 'distrib_optim.pt') + return os.path.join(os.path.dirname(model_checkpoint_name), "distrib_optim.pt") def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): @@ -317,7 +290,7 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): expert_parallel=False, expert_rank=0, ) - if maybe_msc.os.path.isfile(filename): + if isfile(filename): return filename # Look for checkpoint with no pipelining and expert parallelism @@ -331,7 +304,7 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): expert_parallel=True, expert_rank=0, ) - if maybe_msc.os.path.isfile(filename): + if isfile(filename): return filename # Look for checkpoint with pipelining and no expert parallelism @@ -345,7 +318,7 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): expert_parallel=False, expert_rank=0, ) - if maybe_msc.os.path.isfile(filename): + if isfile(filename): return filename # Look for checkpoint with pipelining and expert parallelism @@ -359,7 +332,7 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): expert_parallel=True, expert_rank=0, ) - if maybe_msc.os.path.isfile(filename): + if isfile(filename): return filename # Look for a distributed checkpoint @@ -382,7 +355,7 @@ def checkpoint_exists(checkpoints_path): if checkpoints_path is None: return False path = get_checkpoint_tracker_filename(checkpoints_path) - return maybe_msc.os.path.isfile(path) + return isfile(path) def read_metadata(tracker_filename): @@ -391,7 +364,7 @@ def read_metadata(tracker_filename): iteration = -1 release = False - with maybe_msc.open(tracker_filename, 'r') as f: + with open_file(tracker_filename, 'r') as f: metastring = f.read().strip() try: iteration = int(metastring) @@ -430,24 +403,6 @@ def read_metadata(tracker_filename): return max_iter, release -def read_frozen_resume_iteration(save_dir): - """Resume iteration for a ``--freeze-all-layers`` run. - - Returns the integer recorded in ``save_dir``'s progress tracker - (``latest_checkpointed_iteration.txt``), or ``0`` when there is none -- i.e. a - fresh run, equivalent to ``--finetune`` on the first launch. Unlike - :func:`read_metadata` (which returns ``-1`` / errors on a missing or malformed - file), a missing tracker here simply means "start from the beginning". - """ - if save_dir is None: - return 0 - tracker_filename = get_checkpoint_tracker_filename(save_dir) - if not maybe_msc.os.path.isfile(tracker_filename): - return 0 - iteration, _release = read_metadata(tracker_filename) - return iteration - - def get_rng_state( ckpt_format: str, tp_group: torch.distributed.ProcessGroup, @@ -490,7 +445,7 @@ def get_rng_state( if dp_cp_group is not None else mpu.get_data_parallel_rank(with_context_parallel=True) ) - if ckpt_format == 'torch_dist': + if ckpt_format == "torch_dist": pp_rank = get_pg_rank(pp_group) pp_size = get_pg_size(pp_group) tp_rank = get_pg_rank(tp_group) @@ -502,10 +457,10 @@ def get_rng_state( (pp_rank, tp_rank), replica_id=dp_cp_rank, ) - elif ckpt_format == 'fsdp_dtensor': + elif ckpt_format == "fsdp_dtensor": pp_rank = get_pg_rank(pp_group) tp_rank = get_pg_rank(tp_group) - rng_state_list = {f'({pp_rank}, {tp_rank})': rng_state_list} + rng_state_list = {f"({pp_rank}, {tp_rank})": rng_state_list} return rng_state_list @@ -545,10 +500,10 @@ def _build_sharded_state_dict_metadata( args, 'use_layer_wise_distributed_optimizer', False ) - if has_distributed_optimizer and args.ckpt_format == 'fsdp_dtensor': + if has_distributed_optimizer and args.ckpt_format == "fsdp_dtensor": metadata['distrib_optim_sharding_type'] = 'fsdp_dtensor' - if has_distributed_optimizer and args.ckpt_format != 'fsdp_dtensor': + if has_distributed_optimizer and args.ckpt_format != "fsdp_dtensor": if args.dist_ckpt_optim_fully_reshardable: metadata['distrib_optim_sharding_type'] = 'fully_reshardable' metadata['distrib_optim_fully_reshardable_mem_efficient'] = ( @@ -574,7 +529,8 @@ def save_grads(save_dir, state_dict, iteration, grad_label): this can be handled in postprocessing.""" print_rank_0( - f' [{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] saving {grad_label} from iteration {iteration:7d}' + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving {grad_label} " + f"from iteration {iteration:7d}" ) if mpu.get_expert_data_parallel_rank() == 0: @@ -584,21 +540,22 @@ def save_grads(save_dir, state_dict, iteration, grad_label): tp_rank = mpu.get_tensor_model_parallel_rank() assert save_dir is not None assert iteration is not None - save_dir = os.path.join(save_dir, grad_label, f'iter_{iteration:07d}') + save_dir = os.path.join(save_dir, grad_label, f"iter_{iteration:07d}") os.makedirs(save_dir, exist_ok=True) # Save state_dict. - checkpoint_name = f'mp_rank_{tp_rank:02d}' + checkpoint_name = f"mp_rank_{tp_rank:02d}" if mpu.get_pipeline_model_parallel_world_size() > 1: - checkpoint_name += f'_{pp_rank:03d}' + checkpoint_name += f"_{pp_rank:03d}" if mpu.get_expert_model_parallel_world_size() > 1: - checkpoint_name += f'_{ep_rank:03d}' - full_save_path = os.path.join(save_dir, f'{checkpoint_name}.pth') + checkpoint_name += f"_{ep_rank:03d}" + full_save_path = os.path.join(save_dir, f"{checkpoint_name}.pth") # Convert back to dict (e.g., from collections.defaultdict) for easy loading later. torch.save(dict(state_dict), full_save_path) print_rank_0( - f' [{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] saved {grad_label} from iteration {iteration:7d}' + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saved {grad_label} " + f"from iteration {iteration:7d}" ) @@ -647,6 +604,27 @@ def save_checkpoint( start_ckpt = time() args = get_args() + # Keep the safety invariant even when arguments came through a validation path that does not + # perform the flat CLI cross-check, or were mutated after startup. + optimizer_config = getattr(optimizer, "config", None) if optimizer is not None else None + chunked_optimizer_state_offload = bool( + optimizer_config is not None + and getattr(optimizer_config, "chunked_optimizer_state_offload", False) + and getattr(optimizer_config, "optimizer_state_offload_fraction", 1.0) > 0.0 + ) + if chunked_optimizer_state_offload and not args.no_save_optim: + if args.async_save: + raise RuntimeError( + "chunked optimizer state offload does not support --async-save when optimizer " + "state is saved because reusable pinned CPU buffers cannot be handed to a " + "background writer" + ) + if getattr(args, "ckpt_format", None) != "torch_dist": + raise RuntimeError( + "chunked optimizer state offload requires --ckpt-format torch_dist; checkpoint " + "arguments may have overwritten the startup setting during resume" + ) + if args.async_save and not is_empty_async_queue(): print_rank_0( 'WARNING: Starting a checkpoint save before previous has finished. Consider increasing the checkpoint interval.' @@ -661,10 +639,6 @@ def save_checkpoint( # Only rank zero of the data parallel writes to the disk. model = unwrap_model(model) - # --freeze-all-layers: weights are frozen, so skip the weight write and its finalizers; dump - # progress is recorded separately via the post-logits finalize (see the async block below). - skip_weight_ckpt = getattr(args, 'freeze_all_layers', False) - # Handle non_persistent_ckpt flag. Besides overwriting `args.save` and # `args.use_dist_ckpt`, non-persistent global ckpt requires no additional logic ckpt_type = CheckpointType.GLOBAL if args.use_dist_ckpt else CheckpointType.LEGACY @@ -686,13 +660,13 @@ def save_checkpoint( save_dir = checkpointing_context['local_checkpoint_manager'].local_ckpt_dir else: raise NotImplementedError( - f'Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})' + f"Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})" ) ckpt_format = args.ckpt_format if ckpt_type == CheckpointType.GLOBAL else 'torch' print_rank_0( - f'[{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] saving checkpoint ' - f'at iteration {iteration:7d} to {save_dir} in {ckpt_format} format' + f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving checkpoint " + f"at iteration {iteration:7d} to {save_dir} in {ckpt_format} format" ) # Collect rng state across data parallel ranks. @@ -728,14 +702,9 @@ def save_checkpoint( return_base_dir=return_base_dir, ) - # Save dataloader state if the external dataloader supports it. + # Save dataloader state if the dataloader supports it (currently only Megatron Energon). maybe_save_dataloader_state( - train_data_iterator, - iteration, - getattr(args, 'dataloader_save', None), - tp_group=tp_group, - pp_group=pp_group, - dp_group=dp_group, + train_data_iterator, iteration, getattr(args, "dataloader_save", None) ) # Save distributed optimizer's custom parameter state. @@ -751,10 +720,14 @@ def save_checkpoint( optimizer.save_parameter_state(optim_checkpoint_name) # LayerWiseDistributedOptimizer save optimizer state to file on different ranks - if getattr(args, 'use_layer_wise_distributed_optimizer', False) and args.ckpt_format == 'torch': + if ( + getattr(args, "use_layer_wise_distributed_optimizer", False) + and not args.no_save_optim + and args.ckpt_format == 'torch' + ): dp_rank = mpu.get_data_parallel_rank() optim_checkpoint_name = os.path.join( - os.path.dirname(checkpoint_name), f'layer_wise_optimizer_{dp_rank}.pt' + os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt" ) ensure_directory_exists(optim_checkpoint_name) if not optimizer.is_stub_optimizer: @@ -791,7 +764,7 @@ def save_checkpoint( # exactly one rank. Neither dp_rank==0 nor edp_rank==0 alone covers all shards when # the dense and expert parallelism layouts disagree (e.g. TP > EP*ETP); the union # does, with at most one rank per (tp_rank, ep_rank) inside any DP group. - if not skip_weight_ckpt and ( + if ( not torch.distributed.is_initialized() or ckpt_type != CheckpointType.LEGACY or dp_rank == 0 @@ -806,21 +779,20 @@ def save_checkpoint( ) else: sharded_sd_metadata = None - with _otel_managed_span('checkpoint', 'megatron.checkpoint.save.state_dict', is_goodput_span=True): - state_dict = generate_state_dict( - args, - model, - optimizer, - opt_param_scheduler, - rng_state, - iteration=iteration, - optim_sd_kwargs=dict(metadata=sharded_sd_metadata), - model_sd_kwargs=dict(metadata=sharded_sd_metadata), - rerun_state=rerun_state, - ) + state_dict = generate_state_dict( + args, + model, + optimizer, + opt_param_scheduler, + rng_state, + iteration=iteration, + optim_sd_kwargs=dict(metadata=sharded_sd_metadata), + model_sd_kwargs=dict(metadata=sharded_sd_metadata), + rerun_state=rerun_state, + ) state_dict['num_floating_point_operations_so_far'] = num_floating_point_operations_so_far - if ckpt_type == CheckpointType.GLOBAL and ckpt_format == 'torch_dist': + if ckpt_type == CheckpointType.GLOBAL and ckpt_format == "torch_dist": if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: # TODO Handle non-empty directories (e.g., after a crash during saving). ensure_directory_exists(checkpoint_name, check_parent=False) @@ -852,11 +824,11 @@ def save_checkpoint( checkpointing_context['load_strategy'], 'cached_global_metadata', None ) if cached_global_metadata is not None: - logger.debug('Plugging in the read metadata from the load strategy...') + logger.debug("Plugging in the read metadata from the load strategy...") save_strategy.cached_global_metadata = cached_global_metadata else: logger.debug( - 'Failed to plug in the read metadata from the load strategy...' + "Failed to plug in the read metadata from the load strategy..." ) if args.ckpt_fully_parallel_save: @@ -873,67 +845,56 @@ def save_checkpoint( else mpu.get_expert_data_parallel_group() ) save_strategy = FullyParallelSaveStrategyWrapper( - save_strategy, - process_group, - args.ckpt_assume_constant_structure, - validate_access_integrity=args.ckpt_load_validate_sharding_integrity, - pg_cache_path=getattr(args, 'ckpt_pg_tensors_cache_path', None), - pg_cache_create=getattr(args, 'ckpt_pg_tensors_cache_create', False), + save_strategy, process_group, args.ckpt_assume_constant_structure ) - # Allow opting out of save-side sharding validation entirely (even on - # the first save) via --no-ckpt-load-validate-sharding-integrity. This - # skips the world-wide determine_global_metadata all_gather_object. - if not args.ckpt_load_validate_sharding_integrity: - validate_sharding_integrity = False # Store save strategy for future checkpoint saves if checkpointing_context is not None: checkpointing_context['save_strategy'] = save_strategy end_ckpt = time() logger.debug( - f'rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt ' + f"rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt " + ) + async_save_request = dist_checkpointing.save( + state_dict, + checkpoint_name, + save_strategy, + async_sharded_save=args.async_save, + validate_access_integrity=validate_sharding_integrity, + preprocess_common_before_consistancy_check=preprocess_common_state_dict_fn, + content_metadata=_clean_metadata_for_serialization(sharded_sd_metadata), + async_strategy=args.async_strategy, + verify_integrity=args.verify_integrity, ) - with _otel_managed_span('checkpoint', 'megatron.checkpoint.save.io_write', is_goodput_span=True): - async_save_request = dist_checkpointing.save( - state_dict, - checkpoint_name, - save_strategy, - async_sharded_save=args.async_save, - validate_access_integrity=validate_sharding_integrity, - preprocess_common_before_consistancy_check=preprocess_common_state_dict_fn, - content_metadata=_clean_metadata_for_serialization(sharded_sd_metadata), - async_strategy=args.async_strategy, - verify_integrity=args.verify_integrity, - ) # [ModelOpt]: save sharded modelopt_state if has_nvidia_modelopt: save_sharded_modelopt_state(model, checkpoint_name, (args.ckpt_format, 1)) - elif ckpt_type == CheckpointType.GLOBAL and ckpt_format in ['torch_dcp', 'fsdp_dtensor']: + elif ckpt_type == CheckpointType.GLOBAL and ckpt_format in ["torch_dcp", "fsdp_dtensor"]: if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: # TODO Handle non-empty directories (e.g., after a crash during saving). ensure_directory_exists(checkpoint_name, check_parent=False) - if ckpt_format == 'fsdp_dtensor': + if ckpt_format == "fsdp_dtensor": state_dict = preprocess_fsdp_dtensor_state_dict(args, state_dict, model[0]) if args.async_save: planner = torch.distributed.checkpoint.DefaultSavePlanner() coordinator_rank = 0 _, async_modules = get_async_strategy(args.async_strategy) - FileSystemWriterAsync = async_modules['FileSystemWriterAsync'] - save_state_dict_async_plan = async_modules['save_state_dict_async_plan'] + FileSystemWriterAsync = async_modules["FileSystemWriterAsync"] + save_state_dict_async_plan = async_modules["save_state_dict_async_plan"] _cpu_shm = getattr(args, 'async_ckpt_use_cpu_shm', False) _writer_kwargs = {} if _cpu_shm: if ( - 'use_cpu_shm_for_gpu_tensors' + "use_cpu_shm_for_gpu_tensors" in inspect.signature(FileSystemWriterAsync.__init__).parameters ): - _writer_kwargs['use_cpu_shm_for_gpu_tensors'] = True + _writer_kwargs["use_cpu_shm_for_gpu_tensors"] = True else: raise AssertionError( - 'Installed nvidia-resiliency-ext does not support ' - 'use_cpu_shm_for_gpu_tensors. Update nvidia-resiliency-ext ' - 'to use --async-ckpt-use-cpu-shm.' + "Installed nvidia-resiliency-ext does not support " + "use_cpu_shm_for_gpu_tensors. Update nvidia-resiliency-ext " + "to use --async-ckpt-use-cpu-shm." ) fs_storage_writer = FileSystemWriterAsync( checkpoint_name, @@ -950,6 +911,9 @@ def save_checkpoint( planner=planner, enable_cache=args.ckpt_assume_constant_structure, ) + async_save_request = get_save_and_finalize_callbacks( + fs_storage_writer, save_state_dict_ret + ) async_save_request = get_save_and_finalize_callbacks( fs_storage_writer, save_state_dict_ret, args.async_strategy ) @@ -968,7 +932,7 @@ def save_checkpoint( end_ckpt = time() logger.debug( - f'rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt ' + f"rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt " ) if ckpt_type == CheckpointType.LOCAL: try: @@ -978,7 +942,7 @@ def save_checkpoint( except ModuleNotFoundError: raise RuntimeError( "The 'nvidia_resiliency_ext' module is required for local " - 'checkpointing but was not found. Please ensure it is installed.' + "checkpointing but was not found. Please ensure it is installed." ) if (sharded_sd_metadata or {}).get('distrib_optim_sharding_type') in [ 'fully_reshardable', @@ -1015,7 +979,6 @@ def save_checkpoint( # Save. ensure_directory_exists(checkpoint_name) torch.save(state_dict, checkpoint_name) - start_misc = time() if ckpt_type != CheckpointType.LOCAL: if not args.async_save: @@ -1025,18 +988,15 @@ def save_checkpoint( torch.distributed.barrier() # And update the latest iteration - if not skip_weight_ckpt and ( - not torch.distributed.is_initialized() - or torch.distributed.get_rank() == 0 - ): + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: tracker_filename = get_checkpoint_tracker_filename(save_dir) if ckpt_type == CheckpointType.LOCAL: def iter_finalize_fn(): print_rank_0( - f' [{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] successfully ' - f'saved local checkpoint from iteration {iteration:7d}' + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully " + f"saved local checkpoint from iteration {iteration:7d}" ) if args.log_progress and args.async_save: append_to_progress_log( @@ -1068,8 +1028,6 @@ def _rank_and_size(explicit_rank, group, mpu_rank_fn, mpu_size_fn): mpu.get_pipeline_model_parallel_rank, mpu.get_pipeline_model_parallel_world_size, ) - gtp_remat_rank = mpu.get_gtp_weight_remat_rank() + 1 - gtp_remat_size_to_print = mpu.get_gtp_weight_remat_world_size() def iter_finalize_fn(): prev_iteration = 0 @@ -1077,29 +1035,19 @@ def iter_finalize_fn(): args, 'save_retain_interval', None ) # For backwards compatibility of tests. if save_retain_interval is not None: - if maybe_msc.os.path.exists(tracker_filename): - with maybe_msc.open(tracker_filename, 'r') as f: + if os.path.exists( + tracker_filename + ): # TODO: Make this work with MSC remote paths? + with open_file(tracker_filename, 'r') as f: prev_iteration = int(f.read().strip()) - with maybe_msc.open(tracker_filename, 'w') as f: - f.write('release' if release else str(iteration)) + with open_file(tracker_filename, 'w') as f: + f.write("release" if release else str(iteration)) print_rank_0( - f' [{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] successfully saved ' - f'checkpoint from iteration {int(iteration):7d} to {args.save} ' - f'[ t {tensor_mp_rank}/{tp_size_to_print}, ' - f'gtp_remat {gtp_remat_rank}/{gtp_remat_size_to_print}, ' - f'p {pipeline_mp_rank}/{pp_size_to_print} ]' + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully saved " + f"checkpoint from iteration {int(iteration):7d} to {args.save} " + f"[ t {tensor_mp_rank}/{tp_size_to_print}, " + f"p {pipeline_mp_rank}/{pp_size_to_print} ]" ) - # Save tokenizer files for torch_dist checkpoints (if enabled) - if ( - args.save_tokenizer_assets - and args.ckpt_format == 'torch_dist' - and iteration > 0 - ): - checkpoint_name = get_checkpoint_name( - save_dir, iteration=iteration, return_base_dir=True - ) - config = _default_config_from_args(TokenizerConfig, args) - save_tokenizer_assets(get_tokenizer(), config, checkpoint_name) if args.log_progress and args.async_save: append_to_progress_log( args.save, f'Saved async checkpoint\tIteration: {iteration}', barrier=False @@ -1158,9 +1106,7 @@ def iter_finalize_fn(): iter_finalize_fn() # Additional callback for one_logger (last rank) - if not skip_weight_ckpt and ( - not torch.distributed.is_initialized() or is_last_rank() - ): + if not torch.distributed.is_initialized() or is_last_rank(): def onelogger_finalize_fn(): on_save_checkpoint_success(productive_metrics, args.async_save) @@ -1172,9 +1118,7 @@ def onelogger_finalize_fn(): onelogger_finalize_fn() # Additional callback for wandb (last rank) - if not skip_weight_ckpt and ( - not torch.distributed.is_initialized() or is_last_rank() - ): + if not torch.distributed.is_initialized() or is_last_rank(): def wandb_finalize_fn(): wandb_utils.on_save_checkpoint_success( @@ -1188,64 +1132,14 @@ def wandb_finalize_fn(): wandb_finalize_fn() if args.async_save: - # Schedule logits flush AFTER the checkpoint request so the persistent - # worker processes checkpoint preload first (unblocking the main - # thread), then writes logits in the background. Finalize_fns are - # moved from the checkpoint request to the logits request so that - # "success" callbacks only fire after both writes are confirmed. - from megatron.training.distillation import get_logits_saver - - logits_saver = get_logits_saver() - if logits_saver is not None: - # In frozen-dump mode there is no checkpoint request (async_save_request is None); the - # logits request then owns its own finalize_fns. - if async_save_request is not None: - logits_finalize_fns = async_save_request.finalize_fns.copy() - async_save_request.finalize_fns.clear() - else: - logits_finalize_fns = [] - # Record run progress AFTER the logits tar is confirmed written, so a resumed job - # never skips or replays a window. Written by a single rank -- the last rank, which - # lives on the last pipeline stage where the logits saver is attached (get_logits_saver - # is None on earlier stages, including global rank 0 when PP > 1). - if skip_weight_ckpt and ( - not torch.distributed.is_initialized() or is_last_rank() - ): - - def progress_finalize_fn(): - tracker_filename = get_checkpoint_tracker_filename(args.save) - with maybe_msc.open(tracker_filename, 'w') as f: - f.write(str(iteration)) - print_rank_last(f" recorded logits-dump progress: iteration " - f"{iteration} to {tracker_filename}") - - logits_finalize_fns.append(progress_finalize_fn) - async_request_cls = get_async_strategy(args.async_strategy)[1]['AsyncRequest'] - async_logits_request = async_request_cls( - async_fn=logits_saver._write_batched_tar, - async_fn_args=logits_saver.take_pending_data(), - finalize_fns=logits_finalize_fns, - ) - - if ( - getattr(args, "rl_durable_rollout_bank", False) - and async_save_request is not None - and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0) - ): - # Compact the bank once this async checkpoint becomes durable, so the - # bank's compacted-through T tracks the model checkpoint. Rank-0 only. - _register_rollout_bank_compaction(async_save_request, iteration) - if async_save_request is not None: - schedule_async_save(async_save_request) - if logits_saver is not None: - schedule_async_save(async_logits_request) + schedule_async_save(async_save_request) print_rank_0( - f' [{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] scheduled ' - f'an async checkpoint save at iteration {iteration:7d} to {save_dir}' + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] scheduled " + f"an async checkpoint save at iteration {iteration:7d} to {save_dir}" ) end_misc = time() - logger.debug(f'rank: {rank}, takes {end_misc - start_misc} to finalize ckpt save ') + logger.debug(f"rank: {rank}, takes {end_misc - start_misc} to finalize ckpt save ") if not args.async_save: # Add a barrier so that all ranks wait for finalization to complete @@ -1253,151 +1147,9 @@ def progress_finalize_fn(): if torch.distributed.is_initialized(): torch.distributed.barrier() - # Durable rollout bank: compact at the checkpoint boundary so the bank's - # compacted-through T tracks this (now durable) checkpoint. Only on sync - # saves; async saves compact in their durability finalize callback above. - # Rank-0 no-op otherwise. - if getattr(args, "rl_durable_rollout_bank", False): - _maybe_compact_rollout_bank(iteration) - ft_integration.on_checkpointing_end(is_async_finalization=False) -def save_tokenizer_assets( - tokenizer: MegatronTokenizer, - config: TokenizerConfig, - checkpoint_path: str, -) -> None: - """Save tokenizer files to the checkpoint directory. - - Always saves tokenizer files to ensure checkpoints are self-contained - and portable. Handles both HuggingFace tokenizers and file-based tokenizers. - Compatible with MultiStorageClient for cloud storage support. - - Args: - tokenizer: The tokenizer instance to save. - config: Tokenizers config. - checkpoint_path: The checkpoint directory path. - """ - if tokenizer is None: - return - - # Only rank 0 saves tokenizer files - rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 - if rank != 0: - return - - def resolve_path(path_str: str) -> str: - """Resolve relative paths to absolute paths.""" - if not path_str: - return path_str - path_obj = Path(path_str) - if path_obj.is_absolute(): - return path_str - # Resolve relative to current working directory - return str(path_obj.resolve()) - - try: - # Check if MultiStorageClient is enabled - if MultiStorageClientFeature .is_enabled(): - msc = MultiStorageClientFeature .import_package() - checkpoint_path_obj = msc.Path(checkpoint_path) - tokenizer_dir = checkpoint_path_obj / "tokenizer" - tokenizer_dir.mkdir(parents=True, exist_ok=True) - use_msc = True - else: - tokenizer_dir = os.path.join(checkpoint_path, "tokenizer") - os.makedirs(tokenizer_dir, exist_ok=True) - use_msc = False - - tokenizer_type = config.tokenizer_type - - # Handle HuggingFace and Multimodal tokenizers - if tokenizer_type in ("HuggingFaceTokenizer", "MultimodalTokenizer"): - if use_msc: - import tempfile - - with tempfile.TemporaryDirectory() as tmp_dir: - if hasattr(tokenizer, "save_pretrained"): - tokenizer.save_pretrained(tmp_dir) - elif hasattr(tokenizer, "_tokenizer") and hasattr(tokenizer._tokenizer, "save_pretrained"): - tokenizer._tokenizer.save_pretrained(tmp_dir) - else: - logger.debug(f"{tokenizer_type} does not support save_pretrained(), skipping tokenizer save") - return - - logger.debug(f"Saving {tokenizer_type} files to {tokenizer_dir}") - for filename in os.listdir(tmp_dir): - src_path = os.path.join(tmp_dir, filename) - if os.path.isfile(src_path): - dest_path = tokenizer_dir / filename - with open(src_path, "rb") as src_f: - with msc.open(str(dest_path), "wb") as dest_f: - dest_f.write(src_f.read()) - else: - logger.debug(f"Saving {tokenizer_type} files to {tokenizer_dir}") - if hasattr(tokenizer, "save_pretrained"): - tokenizer.save_pretrained(tokenizer_dir) - elif hasattr(tokenizer, "_tokenizer") and hasattr(tokenizer._tokenizer, "save_pretrained"): - tokenizer._tokenizer.save_pretrained(tokenizer_dir) - return - - # Handle file-based tokenizers - resolve all paths - files_to_copy = [] - - if tokenizer_type in ("BertWordPieceLowerCase", "BertWordPieceCase"): - if config.vocab_file: - resolved_path = resolve_path(config.vocab_file) - files_to_copy.append(("vocab_file", resolved_path, "vocab.txt")) - - elif tokenizer_type == "GPT2BPETokenizer": - if config.vocab_file: - resolved_path = resolve_path(config.vocab_file) - files_to_copy.append(("vocab_file", resolved_path, "vocab.json")) - if config.merge_file: - resolved_path = resolve_path(config.merge_file) - files_to_copy.append(("merge_file", resolved_path, "merges.txt")) - - elif tokenizer_type in ("SentencePieceTokenizer", "GPTSentencePieceTokenizer", "Llama2Tokenizer"): - if config.tokenizer_model: - resolved_path = resolve_path(config.tokenizer_model) - files_to_copy.append(("tokenizer_model", resolved_path, "tokenizer.model")) - - elif tokenizer_type == "TikTokenizer": - if config.tokenizer_model: - resolved_path = resolve_path(config.tokenizer_model) - files_to_copy.append(("tokenizer_model", resolved_path, "tokenizer.json")) - - elif tokenizer_type == "NullTokenizer": - logger.debug(f"{tokenizer_type} requires no file artifacts") - return - - # Copy the files - if files_to_copy: - logger.debug(f"Saving {tokenizer_type} files to {tokenizer_dir}") - for config_attr, source_path, dest_filename in files_to_copy: - if source_path and os.path.exists(source_path): - if use_msc: - dest_path = tokenizer_dir / dest_filename - with open(source_path, "rb") as src_f: - with msc.open(str(dest_path), "wb") as dest_f: - dest_f.write(src_f.read()) - logger.debug(f"Copied {config_attr}: {source_path} -> {dest_path}") - else: - dest_path = os.path.join(tokenizer_dir, dest_filename) - shutil.copy2(source_path, dest_path) - logger.debug(f"Copied {config_attr}: {source_path} -> {dest_path}") - else: - logger.debug(f"{config_attr} not found at resolved path: {source_path}") - - except Exception as e: - if get_rank_safe() == 0: - logger.error(f"Failed to save tokenizer files: {e}") - import traceback - - logger.error(traceback.format_exc()) - - @_disable_gc() def _async_delete_checkpoint_impl( save_path, @@ -1432,7 +1184,8 @@ def _async_delete_checkpoint_impl( try: shutil.rmtree(checkpoint_name) # TODO: Make this work with MSC remote paths? print( - f' successfully deleted checkpoint from iteration {iteration_to_delete:7d} at {save_path}', + f' successfully deleted checkpoint from iteration {iteration_to_delete:7d} ' + f'at {save_path}', flush=True, ) if log_progress: @@ -1454,7 +1207,7 @@ def cleanup_old_non_persistent_checkpoint(save_dir, leave_ckpt_num=1, do_async=F return save_dir = Path(save_dir) - iter_prefix = 'iter_' + iter_prefix = "iter_" iter_ckpts = save_dir.rglob(f'{iter_prefix}*') sorted_iter_ckpts = sorted( iter_ckpts, key=lambda ckpt_name: int(ckpt_name.name[len(iter_prefix) :]) @@ -1475,13 +1228,12 @@ def remove_iter_ckpts(_iter_ckpts): remove_iter_ckpts(rm_iter_ckpts) -def maybe_save_dataloader_state( - train_iterator, iteration, dataloader_save_path, *, tp_group=None, pp_group=None, dp_group=None -): +def maybe_save_dataloader_state(train_iterator, iteration, dataloader_save_path): """Saves dataloader state if the dataloader supports it. - External dataloaders use this to store state at a specific iteration. The Megatron built-in - dataloader creates index files upfront to track its state. + Currently, this is only used by Megatron Energon dataloader (multimodal) to store its state at a + specific iteration. The Megatron built-in dataloader (text-only) creates index files upfront + to track its state. If the provided dataloader has `save_state` method, then it is called to save the state. Otherwise, no state is saved. @@ -1490,63 +1242,39 @@ def maybe_save_dataloader_state( train_iterator (iterable): Train dataloader. iteration (int): Current iteration. dataloader_save_path (str): Path where the dataloader state is saved. - tp_group (ProcessGroup): Tensor-parallel group, or MPU fallback when unset. - pp_group (ProcessGroup): Pipeline-parallel group, or MPU fallback when unset. - dp_group (ProcessGroup): Data-parallel group, or MPU fallback when unset. """ # If no dataloader or saving path is provided, exit early, otherwise, raise an error. - if train_iterator is None or dataloader_save_path is None or dataloader_save_path == '': + if train_iterator is None or dataloader_save_path is None or dataloader_save_path == "": return # If dataloader doesn't support saving state, raise an error. - if not hasattr(train_iterator.iterable, 'save_state'): + if not hasattr(train_iterator.iterable, "save_state"): raise RuntimeError( - f'Could not find a save_state for the train_iterator of type {type(train_iterator)}' + f"Could not find a save_state for the train_iterator of type {type(train_iterator)}" ) # Save dataloader state for each data parallel rank only once. first_rank = ( - get_pg_rank(pp_group) == 0 - if pp_group is not None - else mpu.is_pipeline_first_stage(ignore_virtual=True) - ) and ( - get_pg_rank(tp_group) == 0 - if tp_group is not None - else mpu.get_tensor_model_parallel_rank() == 0 + mpu.is_pipeline_first_stage(ignore_virtual=True) + and mpu.get_tensor_model_parallel_rank() == 0 ) if not first_rank: return - dp_rank = get_pg_rank(dp_group) if dp_group is not None else mpu.get_data_parallel_rank() - train_dataloader_state_dict = train_iterator.iterable.save_state() + dp_rank = mpu.get_data_parallel_rank() if dp_rank == 0: - print(f'saving dataloader checkpoint at iteration {iteration} to {dataloader_save_path}') + print(f"saving dataloader checkpoint at iteration {iteration} to {dataloader_save_path}") + train_dataloader_state_dict = train_iterator.iterable.save_state() data_state_save_path = get_checkpoint_name( - dataloader_save_path, - iteration, - pipeline_parallel=( - get_pg_size(pp_group) > 1 - if pp_group is not None - else mpu.get_pipeline_model_parallel_world_size() > 1 - ), - # Dataloader state is sharded only by DP rank. Keep it in the canonical TP0/PP0 directory. - tensor_rank=0, - pipeline_rank=0, - expert_parallel=False, - expert_rank=0, - basename=f'train_dataloader_dprank{dp_rank:03d}.pt', + dataloader_save_path, iteration, basename=f'train_dataloader_dprank{dp_rank:03d}.pt' ) - data_parallel_group = dp_group if dp_group is not None else mpu.get_data_parallel_group() - torch.distributed.barrier(group=data_parallel_group) + torch.distributed.barrier(group=mpu.get_data_parallel_group()) - if dp_rank == 0: + if mpu.get_data_parallel_rank() == 0: ensure_directory_exists(data_state_save_path) - torch.distributed.barrier(group=data_parallel_group) - - if train_dataloader_state_dict is None: - return + torch.distributed.barrier(group=mpu.get_data_parallel_group()) dataloader_save_dict = {} dataloader_save_dict['dataloader_state_dict'] = train_dataloader_state_dict @@ -1574,17 +1302,17 @@ def generate_state_dict( state_dict['iteration'] = iteration for i in range(len(model)): - key = 'model' + key = "model" if len(model) > 1: - key = f'model{i}' + key = f"model{i}" - if args.ckpt_format == 'torch_dist': + if args.ckpt_format == "torch_dist": model_sd = model[i].sharded_state_dict( **( model_sd_kwargs or { - 'metadata': { - 'dp_cp_group': mpu.get_data_parallel_group(with_context_parallel=True) + "metadata": { + "dp_cp_group": mpu.get_data_parallel_group(with_context_parallel=True) } } ) @@ -1597,25 +1325,26 @@ def generate_state_dict( # Optimizer stuff. if not args.no_save_optim: if optimizer is not None and not optimizer.is_stub_optimizer: - if args.ckpt_format == 'torch_dist': + + if args.ckpt_format == "torch_dist": optimizer_sd = optimizer.sharded_state_dict( state_dict, **( optim_sd_kwargs or { - 'metadata': { - 'dp_cp_group': mpu.get_data_parallel_group( + "metadata": { + "dp_cp_group": mpu.get_data_parallel_group( with_context_parallel=True ) } } ), ) - elif args.ckpt_format == 'fsdp_dtensor': + elif args.ckpt_format == "fsdp_dtensor": if optim_sd_kwargs is None: optim_sd_kwargs = {} - if 'metadata' not in optim_sd_kwargs: - optim_sd_kwargs['metadata'] = {} + if "metadata" not in optim_sd_kwargs: + optim_sd_kwargs["metadata"] = {} optim_sd_kwargs['metadata'].update(_build_sharded_state_dict_metadata(args)) optimizer_sd = optimizer.sharded_state_dict(state_dict, **optim_sd_kwargs) else: @@ -1632,124 +1361,35 @@ def generate_state_dict( # RNG states. if not args.no_save_rng and rng_state: - state_dict['rng_state'] = rng_state - - # Optionally avoid persisting TE `_extra_state` artifacts that carry no - # irreplaceable state (see `_localize_redundant_extra_states`). Opt-in via - # --ckpt-drop-redundant-extra-state (off by default, preserving prior behavior). - # Applied here so the SAME decision is used for both the save state dict and - # the load *request* (`generate_state_dict` is the shared chokepoint), which - # keeps the two symmetric and backward compatible. - if args.ckpt_format == "torch_dist" and getattr( - args, "ckpt_drop_redundant_extra_state", False - ): - _localize_redundant_extra_states(state_dict) + state_dict["rng_state"] = rng_state return state_dict -# Byte markers of the dict keys that TE `get_extra_state` writes ONLY under -# `if recipe.delayed():` (see TransformerEngine `module/base.py`). Their presence -# in the serialized payload is a robust, unpickle-free signal that the -# `_extra_state` carries persistent delayed-scaling state (amax history + scale). -_FP8_PERSISTENT_EXTRA_STATE_MARKERS = (b"amax_history_fwd", b"scale_fwd", b"scale_bwd") - - -def _fp8_extra_state_is_persistent(data) -> bool: - """Whether a TE `_extra_state` payload must be checkpointed. - - A `_extra_state` is worth persisting only when it carries state that cannot - be reconstructed from the run config at model-build time, i.e. the - delayed-scaling FP8 amax history + scale. Concretely: - - * ``None`` / empty ``uint8`` tensor -> FP8 disabled -> NOT persistent - * non-empty, no delayed-scaling key -> block/current scaling (NVFP4, MXFP8, - Float8CurrentScaling): recipe + scalars only, all config-derived - -> NOT persistent - * non-empty, has ``scale_fwd`` / ``amax_history_fwd`` / ``scale_bwd`` - -> delayed scaling -> persistent - - The check operates on the raw serialized bytes (the uint8 tensor that - ``get_extra_state`` returns) so it needs no unpickling and no TE import. - Anything unrecognized defaults to persistent, so real state is never - silently dropped. - """ - if data is None: - return False - if isinstance(data, torch.Tensor): - if data.numel() == 0: - return False - try: - raw = data.detach().cpu().contiguous().numpy().tobytes() - except Exception: - return True - return any(marker in raw for marker in _FP8_PERSISTENT_EXTRA_STATE_MARKERS) - # Unknown payload type: keep it persistent to be safe. - return True - - -def _localize_redundant_extra_states(state_dict): - """Rewrite non-persistent TE `_extra_state` ShardedObjects as local objects. - - For checkpoints with no FP8, or with block/current FP8 scaling (NVFP4, - MXFP8, Float8CurrentScaling), every `_extra_state` is either empty or a - recipe-only blob that the freshly-built model reproduces from config — see - `_fp8_extra_state_is_persistent`. Such artifacts are pure overhead: in a - large MoE model they account for the overwhelming majority of the - checkpoint's ShardedObjects (e.g. ~50k on the 55B hybrid-MoE run). - - Wrapping them in ``LocalNonpersistentObject`` (instead of ``ShardedObject``) - means, via the dist-checkpointing pipeline: - - * SAVE: ``save_preprocess`` drops them, so they are never written to disk. - * LOAD: ``load_preprocess`` unwraps them back into the loaded state dict with - the local (freshly-built) value, so the module's ``_extra_state`` key is - still present (strict ``load_state_dict`` stays happy) and - ``set_extra_state`` no-ops on the empty/recipe payload. - - Because this runs in ``generate_state_dict`` (shared by save and the load - request), the decision is symmetric: dropped keys are never *requested*, so - loading an older checkpoint that still stores them simply does not read them - (they are neither "missing" nor "unexpected" in the request) — backward - compatible. Delayed-scaling `_extra_state` stays a ``ShardedObject`` and is - saved/loaded exactly as before. - """ - - def _maybe_localize(value): - if ( - isinstance(value, ShardedObject) - and value.key.endswith("_extra_state") - and not _fp8_extra_state_is_persistent(value.data) - ): - return LocalNonpersistentObject(value.data) - return value - - dict_list_map_inplace(_maybe_localize, state_dict) - - def preprocess_fsdp_dtensor_state_dict(args, raw_state_dict, model): state_dict = raw_state_dict.copy() - handle_fp8_extra_state_case(state_dict['model']) - - def apply(handler): - """Run a state dict handler over the model and, when present, the optimizer state.""" - model_state_dict, optimizer_state_dict = handler( - model, state_dict['model'], state_dict.get('optimizer') - ) - state_dict['model'] = model_state_dict - if 'optimizer' in state_dict: - state_dict['optimizer'] = optimizer_state_dict - + handle_fp8_extra_state_case(state_dict["model"]) if args.swiglu: - apply(handle_swiglu_in_state_dict) - # Split a fused MLA q/kv down-projection (mla_down_proj_fusion) back into the unfused - # layout used on disk. No-op for unfused models. - apply(handle_mla_down_proj_in_state_dict) + if "optimizer" in state_dict: + model_state_dict, optimizer_state_dict = handle_swiglu_in_state_dict( + model, state_dict["model"], state_dict["optimizer"] + ) + state_dict["model"] = model_state_dict + state_dict["optimizer"] = optimizer_state_dict + else: + model_state_dict, _ = handle_swiglu_in_state_dict(model, state_dict["model"], None) + state_dict["model"] = model_state_dict + if "optimizer" in state_dict: + model_state_dict, optimizer_state_dict = handle_gdn_in_state_dict( + model, state_dict["model"], state_dict["optimizer"] + ) + state_dict["model"] = model_state_dict + state_dict["optimizer"] = optimizer_state_dict + else: + model_state_dict, _ = handle_gdn_in_state_dict(model, state_dict["model"], None) + state_dict["model"] = model_state_dict if args.num_experts: - state_dict['model'] = handle_experts_in_state_dict(state_dict['model'], args.num_experts) - # Rename the MTP inner layer to the name used on disk. Runs last because the handlers - # above resolve keys against live module paths, which still use the new name. - apply(handle_mtp_in_state_dict) + state_dict["model"] = handle_experts_in_state_dict(state_dict["model"], args.num_experts) preprocess_state_dict_for_uneven_dtensor(state_dict) return state_dict @@ -1812,7 +1452,7 @@ def fix_query_key_value_ordering(model, checkpoint_version): elif checkpoint_version == 1.0: fixed_param = _transpose_first_dim(param.data, 3, False, model) else: - print_rank_0(f'Invalid checkpoint version {checkpoint_version}.') + print_rank_0(f"Invalid checkpoint version {checkpoint_version}.") sys.exit() param.data.copy_(fixed_param) if name.endswith(('.key_value.weight', '.key_value.bias')): @@ -1821,36 +1461,36 @@ def fix_query_key_value_ordering(model, checkpoint_version): elif checkpoint_version == 1.0: fixed_param = _transpose_first_dim(param.data, 2, False, model) else: - print_rank_0(f'Invalid checkpoint version {checkpoint_version}.') + print_rank_0(f"Invalid checkpoint version {checkpoint_version}.") sys.exit() param.data.copy_(fixed_param) print_rank_0( - ' successfully fixed query-key-values ordering for checkpoint version {}'.format( - checkpoint_version - ) + " successfully fixed query-key-values ordering for" + " checkpoint version {}".format(checkpoint_version) ) def _get_non_persistent_iteration(non_persistent_global_dir, args, checkpointing_context=None): if args.non_persistent_ckpt_type is None: return -1 - elif args.non_persistent_ckpt_type == 'global': + elif args.non_persistent_ckpt_type == "global": tracker_filename = get_checkpoint_tracker_filename(non_persistent_global_dir) - if maybe_msc.os.path.isfile(tracker_filename): + if isfile(tracker_filename): iteration, release = read_metadata(tracker_filename) if release: - raise RuntimeError("Non-persistent checkpoint can't be a release checkpoint") + raise RuntimeError('Non-persistent checkpoint can\'t be a release checkpoint') else: iteration = -1 print_rank_0('WARNING: could not find the metadata file {}'.format(tracker_filename)) print_rank_0(' will not load any non-persistent checkpoint') return iteration - elif args.non_persistent_ckpt_type == 'local': + elif args.non_persistent_ckpt_type == "local": return checkpointing_context['local_checkpoint_manager'].find_latest() else: - assert ( - False - ), f'Please use local or global non-persistent checkpoints(got: {args.non_persistent_ckpt_type})' + assert False, ( + 'Please use local or global non-persistent checkpoints' + f'(got: {args.non_persistent_ckpt_type})' + ) def _load_non_persistent_base_checkpoint( @@ -1867,7 +1507,7 @@ def _load_non_persistent_base_checkpoint( Depending on the non_persistent_ckpt_type, different logic may be required. """ assert args.non_persistent_ckpt_type is not None - if args.non_persistent_ckpt_type == 'global': + if args.non_persistent_ckpt_type == "global": if not rank0: print_rank_0( f'Loading from a non-persistent checkpoint (non-persistent iter {non_persistent_iteration})' @@ -1883,7 +1523,7 @@ def _load_non_persistent_base_checkpoint( dp_cp_group=dp_cp_group, expt_dp_group=expt_dp_group, ) - elif args.non_persistent_ckpt_type == 'local': + elif args.non_persistent_ckpt_type == "local": intermediate_state_dict, checkpoint_name = checkpointing_context[ 'local_checkpoint_manager' ].load() @@ -1899,7 +1539,7 @@ def _load_non_persistent_base_checkpoint( return state_dict, checkpoint_name, False, CheckpointType.LOCAL else: raise NotImplementedError( - f'Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})' + f"Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})" ) @@ -1945,21 +1585,14 @@ def _load_global_dist_base_checkpoint( ) else: raise ValueError( - f'Invalid load process group: {args.ckpt_fully_parallel_load_process_group}' + f"Invalid load process group: {args.ckpt_fully_parallel_load_process_group}" ) load_strategy = FullyParallelLoadStrategyWrapper( - load_strategy, - process_group, - exchange_algo=args.ckpt_fully_parallel_load_exchange_algo, - per_rank_object_load=getattr( - args, 'ckpt_fully_parallel_load_per_rank_objects', False - ), - pg_cache_path=getattr(args, 'ckpt_pg_tensors_cache_path', None), - pg_cache_create=getattr(args, 'ckpt_pg_tensors_cache_create', False), + load_strategy, process_group, exchange_algo=args.ckpt_fully_parallel_load_exchange_algo ) if checkpointing_context is not None: - checkpointing_context['load_strategy'] = load_strategy + checkpointing_context["load_strategy"] = load_strategy state_dict = dist_checkpointing.load( sharded_state_dict, checkpoint_name, @@ -1973,21 +1606,26 @@ def _load_global_dist_base_checkpoint( def _get_checkpoint_format(checkpoint_name, args): """Get the format of an existing checkpoint.""" - checkpoint_dir = maybe_msc.Path(checkpoint_name) - is_torch_ckpt = any([f.name.startswith('mp_rank_0') for f in checkpoint_dir.iterdir()]) - is_torch_dcp = checkpoint_dir.joinpath('.metadata').exists() + if MultiStorageClientFeature.is_enabled(): + msc = MultiStorageClientFeature.import_package() + checkpoint_dir = msc.Path(checkpoint_name) + is_torch_ckpt = any([f.name.startswith("mp_rank_0") for f in checkpoint_dir.iterdir()]) + is_torch_dcp = checkpoint_dir.joinpath(".metadata").exists() + else: + is_torch_ckpt = any([f.startswith("mp_rank_0") for f in os.listdir(checkpoint_name)]) + is_torch_dcp = os.path.exists(os.path.join(checkpoint_name, ".metadata")) ckpt_format = None if dist_checkpointing.check_is_distributed_checkpoint(checkpoint_name): - ckpt_format = 'torch_dist' + ckpt_format = "torch_dist" elif is_torch_ckpt: - ckpt_format = 'torch' + ckpt_format = "torch" elif is_torch_dcp: - ckpt_format = 'torch_dcp' - if getattr(args, 'use_megatron_fsdp', False): - ckpt_format = 'fsdp_dtensor' + ckpt_format = "torch_dcp" + if getattr(args, "use_megatron_fsdp", False): + ckpt_format = "fsdp_dtensor" else: - raise NotImplementedError(f'unknown checkpoint format in {checkpoint_name}') + raise NotImplementedError(f"unknown checkpoint format in {checkpoint_name}") return ckpt_format @@ -2000,7 +1638,6 @@ def _load_base_checkpoint( checkpointing_context=None, dp_cp_group=None, expt_dp_group=None, - gpt_compat_layer_maps=None, ): """Load the base state_dict from the given directory @@ -2019,11 +1656,11 @@ def _load_base_checkpoint( tracker_filename = 'because load directory is not defined' if load_dir is not None: tracker_filename = get_checkpoint_tracker_filename(load_dir) - if maybe_msc.os.path.isfile(tracker_filename): + if isfile(tracker_filename): iteration, release = read_metadata(tracker_filename) # Allow user to specify the loaded iteration. - if getattr(args, 'ckpt_step', None): + if getattr(args, "ckpt_step", None): iteration = args.ckpt_step # Record the iteration loaded (stored separately from args to avoid @@ -2053,25 +1690,19 @@ def _load_base_checkpoint( print_rank_0(' will not load any checkpoints and will start from random') # Conditionally exit if checkpoint not found. if args.exit_on_missing_checkpoint: - print_rank_0( - ">> '--exit-on-missing-checkpoint' is set but no checkpoint was found under " - f"load directory '{load_dir}' (missing metadata/tracker file " - f"'{tracker_filename}'). Exiting with a non-zero status. <<" - ) + print_rank_0(">> '--exit-on-missing-checkpoint' set ... exiting. <<") if torch.distributed.is_initialized(): torch.distributed.barrier() - # Exit non-zero so that callers (e.g. CI harnesses) detect the missing - # checkpoint as a failure instead of silently treating exit code 0 as success. - sys.exit(1) + sys.exit() - return None, '', False, None + return None, "", False, None # Determine the type of the checkpoint on disk. checkpoint_name = get_checkpoint_name(load_dir, iteration, release, return_base_dir=True) ckpt_format = _get_checkpoint_format(checkpoint_name, args) if not rank0: - dist_infix = 'distributed ' if ckpt_format == 'torch_dist' else '' + dist_infix = "distributed " if ckpt_format == "torch_dist" else "" if release: print_rank_0(f' loading release {dist_infix}checkpoint from {load_dir}') else: @@ -2082,7 +1713,7 @@ def _load_base_checkpoint( ckpt_type = None # Handle global distributed checkpoint - if ckpt_format == 'torch_dist': + if ckpt_format == "torch_dist": return _load_global_dist_base_checkpoint( load_dir, args, @@ -2094,7 +1725,7 @@ def _load_base_checkpoint( dp_cp_group=dp_cp_group, expt_dp_group=expt_dp_group, ) - elif ckpt_format == 'torch': + elif ckpt_format == "torch": ckpt_type = CheckpointType.LEGACY # Handle global legacy checkpoint if rank0: @@ -2104,18 +1735,18 @@ def _load_base_checkpoint( load_dir, iteration, release, return_base_dir=False ) try: - state_dict = torch.load(checkpoint_name, map_location='cpu') + state_dict = torch.load(checkpoint_name, map_location='cpu', weights_only=True) except Exception as e: print('could not load the checkpoint') print(e) sys.exit() - elif ckpt_format == 'torch_dcp': + elif ckpt_format == "torch_dcp": ckpt_type = CheckpointType.TORCH_DCP if rank0: # _load_base_checkpoint is called from load_args_from_checkpoint. torch.distributed is not initialized. # Load only metadata. - state_dict = {'args': None, 'iteration': None} + state_dict = {"args": None, "iteration": None} torch.distributed.checkpoint.load(state_dict=state_dict, checkpoint_id=checkpoint_name) else: # _load_base_checkpoint is called from load_checkpoint with a proper state dict. @@ -2126,72 +1757,63 @@ def _load_base_checkpoint( torch.distributed.checkpoint.load_state_dict( state_dict=state_dict, storage_reader=fs_storage_reader ) - elif ckpt_format == 'fsdp_dtensor': - assert HAVE_MEGATRON_FSDP, 'Should not be called if Megatron-FSDP is not available.' + elif ckpt_format == "fsdp_dtensor": + assert HAVE_MEGATRON_FSDP, "Should not be called if Megatron-FSDP is not available." if rank0: - state_dict = {'args': None, 'iteration': None, 'checkpoint_version': None} - torch.distributed.checkpoint.load(state_dict=state_dict, checkpoint_id=checkpoint_name) - return state_dict, checkpoint_name, release, CheckpointType.FSDP_DTENSOR + return {}, checkpoint_name, release, CheckpointType.FSDP_DTENSOR state_dict = sharded_state_dict raw_optimizer_state_dict = ( - state_dict['optimizer'].copy() if 'optimizer' in state_dict else None + state_dict["optimizer"].copy() if "optimizer" in state_dict else None ) - raw_model_state_dict = state_dict['model'].copy() if 'model' in state_dict else None - model = state_dict.pop('_model') + raw_model_state_dict = state_dict["model"].copy() if "model" in state_dict else None + model = state_dict.pop("_model") state_dict = preprocess_fsdp_dtensor_state_dict(args, state_dict, model[0]) - fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(checkpoint_name) - state_dict_metadata = fs_storage_reader.read_metadata().state_dict_metadata - if gpt_compat_layer_maps is not None: - from megatron.core.dist_checkpointing.gpt_checkpoint_interop import ( - retarget_fsdp_state_dict_to_gpt_checkpoint, - ) - - state_dict['model'] = retarget_fsdp_state_dict_to_gpt_checkpoint( - state_dict['model'], - gpt_compat_layer_maps, - tuple(key for key in state_dict_metadata if key.startswith('model.')), - checkpoint_prefix='model', - ) - if 'optimizer' in state_dict: - state_dict['optimizer'] = retarget_fsdp_state_dict_to_gpt_checkpoint( - state_dict['optimizer'], - gpt_compat_layer_maps, - tuple(key for key in state_dict_metadata if key.startswith('optimizer.')), - checkpoint_prefix='optimizer', - ) ckpt_type = CheckpointType.FSDP_DTENSOR + fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(checkpoint_name) allow_partial_load = not getattr(args, 'strict_fsdp_dtensor_load', False) + + # Stash state_dict keys that have no presence in the checkpoint. + # For Megatron-saved checkpoints, DCP stores all entries (tensors + # as TensorStorageMetadata, non-tensors as BytesStorageMetadata), + # so nothing needs to be stashed. + # For Bridge-converted checkpoints, metadata keys like 'args', + # 'iteration', etc. don't exist — stash them so DCP doesn't + # error, then restore placeholder values after load. + state_dict_metadata = fs_storage_reader.read_metadata().state_dict_metadata + _stashed_keys = {} + for key in list(state_dict.keys()): + # Keep the key if it exists as a top-level metadata entry + # (e.g. 'args', 'iteration') OR if any flattened sub-key + # exists (e.g. 'model.layer.weight' for top-level 'model'). + if key not in state_dict_metadata and not any( + k.startswith(key + '.') for k in state_dict_metadata + ): + _stashed_keys[key] = state_dict.pop(key) + if allow_partial_load: rank = torch.distributed.get_rank() import time as _time _time.sleep(rank * 0.001) # Make that logs of different ranks do not overlap print_diff_in_state_dicts(state_dict_metadata, state_dict) - # A partial load silently skips model weights the checkpoint does not have, - # leaving them at their initialized values. Report that, unless this is a - # fine-tune where loading only part of the model is the intent. - if not getattr(args, 'finetune', False): - validate_fsdp_dtensor_model_load( - state_dict_metadata, - state_dict, - checkpoint_name, - strict=args.dist_ckpt_strictness, - ) planner = default_planner.DefaultLoadPlanner(allow_partial_load=allow_partial_load) torch.distributed.checkpoint.load_state_dict( state_dict=state_dict, storage_reader=fs_storage_reader, planner=planner ) + # Restore stashed keys (Bridge fallback placeholders). + state_dict.update(_stashed_keys) + if raw_optimizer_state_dict is not None: - state_dict['optimizer'] = raw_optimizer_state_dict + state_dict["optimizer"] = raw_optimizer_state_dict if raw_model_state_dict is not None: - state_dict['model'] = raw_model_state_dict + state_dict["model"] = raw_model_state_dict else: - raise NotImplementedError(f'checkpoint format {ckpt_format} not supported') + raise NotImplementedError(f"checkpoint format {ckpt_format} not supported") return state_dict, checkpoint_name, release, ckpt_type @@ -2263,10 +1885,10 @@ def _set_arg(arg_name, old_arg_name=None, force=False): checkpoint_value = getattr(checkpoint_args, arg_name, None) if checkpoint_value is not None: - print_rank_0(f'Setting {arg_name} to {checkpoint_value} from checkpoint') + print_rank_0(f"Setting {arg_name} to {checkpoint_value} from checkpoint") setattr(args, arg_name, checkpoint_value) else: - print_rank_0(f'Checkpoint did not provide arguments {arg_name}') + print_rank_0(f"Checkpoint did not provide arguments {arg_name}") # Model args. _set_arg('num_layers') @@ -2325,10 +1947,6 @@ def _set_arg(arg_name, old_arg_name=None, force=False): _set_arg('mamba_head_dim', force=True) _set_arg('mamba_num_groups', force=True) _set_arg('mamba_num_heads', force=True) - # GDP checkpoints created before this argument existed always used three reflections. - if not hasattr(checkpoint_args, 'gdp_num_householder'): - setattr(checkpoint_args, 'gdp_num_householder', 3) - _set_arg('gdp_num_householder', force=True) # We need to be able to override hybrid_layer_pattern from the command-line so that different # pipelining can be specified when re-loading a model (e.g. for inference or post-training). _set_arg('hybrid_layer_pattern') @@ -2339,7 +1957,6 @@ def _set_arg(arg_name, old_arg_name=None, force=False): # MoE latent projection. _set_arg('moe_latent_size', force=True) - _set_arg('moe_use_norm_before_up_proj', force=True) # Tokenizer args. if args.use_tokenizer_model_from_checkpoint_args: @@ -2366,95 +1983,6 @@ def _set_arg(arg_name, old_arg_name=None, force=False): return args, checkpoint_args -def _maybe_setup_gpt_to_hybrid_load(args, ckpt_args, model): - """Detect a GPT (pure transformer) checkpoint being loaded into a HybridModel run. - - Returns ``(layer_maps, load_optim)`` where ``layer_maps`` is used to retarget - the run's sharded state dict at the GPT checkpoint's keys (see - ``megatron.core.dist_checkpointing.gpt_checkpoint_interop``) and ``load_optim`` is - True when the GPT optimizer state should be translated and loaded as well. - Returns ``(None, False)`` when checkpoint and runtime already agree. Raises - RuntimeError for combinations that cannot be loaded. - """ - from megatron.core.dist_checkpointing.gpt_checkpoint_interop import gpt_compatible_layer_maps - from megatron.core.models.hybrid.hybrid_model import HybridModel - from megatron.core.models.mimo.model.base import MimoModel - - def _contains_hybrid_model(module): - # Megatron-FSDP and Float16Module both retain the wrapped module under - # ``module`` but are intentionally not handled by the regular - # ``unwrap_model`` helper. Multimodal wrappers (e.g. LLaVAModel) attach - # the language model under ``language_model`` instead. - while module is not None: - if isinstance(module, HybridModel): - return True - if isinstance(module, MimoModel): - language_module = module.mimo_config.language_model_spec.module - if isinstance(language_module, type) and issubclass(language_module, HybridModel): - return True - inner = getattr(module, 'language_model', None) - if inner is not None and isinstance(inner, HybridModel): - return True - module = getattr(module, 'module', None) - return False - - runtime_is_hybrid = any(_contains_hybrid_model(m) for m in model) - ckpt_pattern = getattr(ckpt_args, 'hybrid_layer_pattern', None) or getattr( - ckpt_args, 'hybrid_override_pattern', None - ) - if runtime_is_hybrid == bool(ckpt_pattern): - return None, False - if not runtime_is_hybrid: - raise RuntimeError( - f'The checkpoint was saved by a hybrid model run (hybrid layer pattern ' - f'{ckpt_pattern!r}) but the current run builds a non-hybrid model. Load it ' - f'with the hybrid training entrypoint instead.' - ) - - # GPT checkpoint feeding a HybridModel run: translate the sharded state - # dict at load time instead of converting the checkpoint on disk. - if not args.hybrid_layer_pattern: - raise RuntimeError( - 'Loading a GPT checkpoint into a hybrid model requires ' - '--hybrid-layer-pattern so checkpoint layers can be paired with ' - 'hybrid layer positions.' - ) - try: - layer_maps = gpt_compatible_layer_maps(args.hybrid_layer_pattern) - except ValueError as exc: - raise RuntimeError(f'Cannot load a GPT checkpoint into this hybrid model: {exc}') from exc - - ckpt_num_layers = getattr(ckpt_args, 'num_layers', None) - if ckpt_num_layers is not None and ckpt_num_layers != layer_maps.num_gpt_layers: - raise RuntimeError( - f'Hybrid layer pattern {args.hybrid_layer_pattern!r} pairs with a GPT ' - f'checkpoint of {layer_maps.num_gpt_layers} layers, but the checkpoint ' - f'has num_layers={ckpt_num_layers}.' - ) - - # The optimizer state is loaded unless the user opts out or the GPT run saved - # no optimizer state. Fresh layers (e.g. Mamba) have no counterpart in the GPT - # checkpoint, so their optimizer state stays freshly initialized; warn about it. - load_optim = not args.no_load_optim and not getattr(ckpt_args, 'no_save_optim', False) - if load_optim and layer_maps.fresh_init: - print_rank_0( - f'> WARNING: {len(layer_maps.fresh_init)} hybrid layer(s) have no GPT ' - f'counterpart (e.g. Mamba positions); their weights and optimizer state ' - f'start from a fresh initialization while the GPT-sourced attention and ' - f'MLP layers load their optimizer state. Pass --no-load-optim to start ' - f"every layer's optimizer state fresh." - ) - - print_rank_0( - f'> loading a GPT checkpoint into the hybrid model: {layer_maps.num_gpt_layers} ' - f'GPT layers feed {len(layer_maps.attention_to_gpt)} attention and ' - f'{len(layer_maps.mlp_to_gpt)} MLP positions; {len(layer_maps.fresh_init)} ' - f'hybrid layers keep their fresh initialization' - + ('; optimizer state will be loaded.' if load_optim else '; optimizer state starts fresh.') - ) - return layer_maps, load_optim - - def load_checkpoint( ddp_model, optimizer, @@ -2484,22 +2012,6 @@ def load_checkpoint( args = get_args() load_dir = getattr(args, load_arg) - # --freeze-all-layers: nothing trains, so load the model in --load weights-only (finetune-style) - # and auto-resume the data position by feeding this run's own progress tracker -- written to - # --save on the previous launch -- into the standard --override-ckpt-iteration path. Reading - # progress from --save (this job's output) rather than --load lets --load stay pinned to a - # fixed checkpoint across resubmits. An explicit --override-ckpt-iteration wins. (The main use - # today is offline-KD teacher-logit dumps.) - if getattr(args, 'freeze_all_layers', False): - # Weights only: don't adopt the loaded checkpoint's optimizer / LR-scheduler / rng, or run - # check_checkpoint_args against a checkpoint from a different run (finetune gates all of - # those; --freeze-all-layers alone would still load the scheduler and assert on arg drift). - args.finetune = True - if args.override_ckpt_iteration is None: - progress_iteration = read_frozen_resume_iteration(args.save) - if progress_iteration > 0: - args.override_ckpt_iteration = progress_iteration - # Finetuning directories pretrained_dir = getattr(args, 'pretrained_checkpoint', None) if pretrained_dir is not None and not checkpoint_exists(load_dir): @@ -2508,63 +2020,47 @@ def load_checkpoint( ) load_dir = pretrained_dir if not checkpoint_exists(load_dir): - raise FileNotFoundError('No checkpoint found in load directory or pretrained directory') + raise FileNotFoundError("No checkpoint found in load directory or pretrained directory") args.finetune = True model = unwrap_model(ddp_model) ckpt_format = args.ckpt_format - state_dict = None - release = False - if args.auto_detect_ckpt_format or ckpt_format in ('torch_dist', 'fsdp_dtensor'): - with _otel_managed_span('load_checkpoint', 'megatron.checkpoint.load.io_read', is_goodput_span=True): - state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( - load_dir, args, rank0=True, checkpointing_context=checkpointing_context - ) + if args.auto_detect_ckpt_format or ckpt_format == "torch_dist": + state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( + load_dir, args, rank0=True, checkpointing_context=checkpointing_context + ) ckpt_format = None if ckpt_type == CheckpointType.TORCH_DCP: - ckpt_format = 'torch_dcp' + ckpt_format = "torch_dcp" elif ckpt_type == CheckpointType.FSDP_DTENSOR: - ckpt_format = 'fsdp_dtensor' + ckpt_format = "fsdp_dtensor" elif ckpt_type == CheckpointType.LEGACY: - ckpt_format = 'torch' + ckpt_format = "torch" elif ckpt_type in [CheckpointType.LOCAL, CheckpointType.GLOBAL]: - ckpt_format = 'torch_dist' + ckpt_format = "torch_dist" elif ckpt_type == None: pass # Not loaded. else: - raise NotImplementedError(f'checkpoint format {ckpt_format} not supported') + raise NotImplementedError(f"checkpoint format {ckpt_format} not supported") load_kwargs = {} ignore_rng_state = False ignore_rerun_state = True - ckpt_args = types.SimpleNamespace() - if ( - ckpt_format in ('torch_dist', 'fsdp_dtensor') - and state_dict is not None - and 'args' in state_dict - ): - ckpt_args = state_dict.get('args') or types.SimpleNamespace() - - # Both model-space torch_dist and fsdp_dtensor checkpoints carry model-keyed - # optimizer state that can be retargeted from GPTModel to HybridModel. - gpt_compat_layer_maps, gpt_compat_load_optim = ( - _maybe_setup_gpt_to_hybrid_load(args, ckpt_args, model) - if ckpt_format in ('torch_dist', 'fsdp_dtensor') and state_dict is not None - else (None, False) - ) - gpt_compat_load_optim = gpt_compat_load_optim and not release + if ckpt_format == "torch_dist": + ckpt_args = types.SimpleNamespace() + if state_dict is not None and "args" in state_dict: + ckpt_args = state_dict.get("args") - if ckpt_format == 'torch_dist': - if not hasattr(ckpt_args, 'tensor_model_parallel_size'): - print_rank_0('WARNING: TP size not found in checkpoint args, using 1 as default.') - if not hasattr(ckpt_args, 'pipeline_model_parallel_size'): - print_rank_0('WARNING: PP size not found in checkpoint args, using 1 as default.') + if not hasattr(ckpt_args, "tensor_model_parallel_size"): + print_rank_0("WARNING: TP size not found in checkpoint args, using 1 as default.") + if not hasattr(ckpt_args, "pipeline_model_parallel_size"): + print_rank_0("WARNING: PP size not found in checkpoint args, using 1 as default.") ckpt_tp_pp = ( - getattr(ckpt_args, 'tensor_model_parallel_size', 1), - getattr(ckpt_args, 'pipeline_model_parallel_size', 1), + getattr(ckpt_args, "tensor_model_parallel_size", 1), + getattr(ckpt_args, "pipeline_model_parallel_size", 1), ) run_tp_pp = (args.tensor_model_parallel_size, args.pipeline_model_parallel_size) @@ -2572,7 +2068,7 @@ def load_checkpoint( run_world_size = getattr(args, 'world_size', 0) ckpt_dp = getattr(ckpt_args, 'data_parallel_size', 0) run_dp = getattr(args, 'data_parallel_size', 0) - mismatch_msg = '(TP, PP) mismatch after resume ({} vs {} from checkpoint)'.format( + mismatch_msg = "(TP, PP) mismatch after resume ({} vs {} from checkpoint)".format( run_tp_pp, ckpt_tp_pp ) @@ -2599,7 +2095,7 @@ def load_checkpoint( ignore_rng_state = True gen_sd_rng_state = None if ckpt_tp_pp != run_tp_pp: - print_rank_0('{}: RNG state will be ignored'.format(mismatch_msg)) + print_rank_0("{}: RNG state will be ignored".format(mismatch_msg)) if ckpt_type == CheckpointType.LOCAL: sharded_sd_metadata = _build_sharded_state_dict_metadata(args, dp_cp_group=dp_cp_group) @@ -2611,12 +2107,10 @@ def load_checkpoint( f'sharded_state_dict metadata loaded from the checkpoint: {sharded_sd_metadata}' ) - # Determine if optimizer state will be loaded. For a GPT->hybrid load the - # optimizer state is retargeted at the GPT checkpoint even under --finetune, - # which independently controls iteration and LR-schedule reset semantics. + # Determine if optimizer state will be loaded if ( not release - and (not args.finetune or gpt_compat_load_optim) + and not args.finetune and not args.no_load_optim and not getattr(ckpt_args, 'no_save_optim', False) ): @@ -2635,32 +2129,15 @@ def load_checkpoint( else 'dp_zero_gather_scatter' ) } - # Retargeting optimizer state onto a GPT checkpoint only works for the - # model-space formats, where each optimizer ShardedTensor carries the - # model param's key and sharding. Bucket-space formats key state by a - # flat buffer layout that the extra hybrid layers reshuffle. - if gpt_compat_load_optim and sharded_sd_metadata[ - 'distrib_optim_sharding_type' - ] not in ('fully_reshardable', 'fully_sharded_model_space'): - raise RuntimeError( - 'Loading optimizer state from a GPT checkpoint into a hybrid model ' - 'is only supported for model-space distributed-optimizer checkpoints ' - "(sharding type 'fully_reshardable' or 'fully_sharded_model_space'), " - f'but the checkpoint uses ' - f'{sharded_sd_metadata["distrib_optim_sharding_type"]!r}. Re-save the ' - 'GPT checkpoint with --dist-ckpt-optim-fully-reshardable, or pass ' - '--no-load-optim ' - 'to start from a fresh optimizer.' - ) if ( ckpt_tp_pp != run_tp_pp and sharded_sd_metadata['distrib_optim_sharding_type'] not in DistributedOptimizer.checkpoint_fully_reshardable_formats ): raise RuntimeError( - f'{mismatch_msg}: not supported for DistributedOptimizer with sharding type' - f' {sharded_sd_metadata["distrib_optim_sharding_type"]}.' - f' Please use `--ckpt-fully-parallel-save` flag during checkpoint saving.' + f"{mismatch_msg}: not supported for DistributedOptimizer with sharding type" + f" {sharded_sd_metadata['distrib_optim_sharding_type']}." + f" Please use `--ckpt-fully-parallel-save` flag during checkpoint saving." ) # Check if fully parallel load is compatible with sharding type @@ -2670,8 +2147,8 @@ def load_checkpoint( == 'dp_zero_gather_scatter' ): raise RuntimeError( - 'Fully parallel load is not supported for dp_zero_gather_scatter checkpoints. ' - 'Please remove --ckpt-fully-parallel-load flag' + "Fully parallel load is not supported for dp_zero_gather_scatter checkpoints. " + "Please remove --ckpt-fully-parallel-load flag" ) else: gen_sd_optim = None @@ -2684,7 +2161,7 @@ def load_checkpoint( # Ensure we have a dict before updating to avoid NoneType AttributeError. if sharded_sd_metadata is None: sharded_sd_metadata = {} - sharded_sd_metadata['dp_cp_group'] = dp_cp_group + sharded_sd_metadata["dp_cp_group"] = dp_cp_group optim_sd_kwargs = dict(metadata=sharded_sd_metadata, is_loading=True) model_sd_kwargs = dict(metadata=sharded_sd_metadata) @@ -2706,7 +2183,7 @@ def load_checkpoint( ) ignore_rerun_state = False if ckpt_world_size != run_world_size or ckpt_tp_pp != run_tp_pp or ckpt_dp != run_dp: - print_rank_0('Job sharding has changed: Rerun state will be ignored') + print_rank_0("Job sharding has changed: Rerun state will be ignored") # [ModelOpt]: Initial loading from non-resume sharded checkpoint to a Distillation Model # will result in key mismatch with loss modules potentially containing parameters, since @@ -2714,7 +2191,7 @@ def load_checkpoint( with ( contextlib.ExitStack() as stack ): # Allows multiple context managers for each model shard - if args.finetune and hasattr(model[0], 'hide_loss_modules'): + if args.finetune and hasattr(model[0], "hide_loss_modules"): for m in model: stack.enter_context(m.hide_loss_modules()) load_kwargs['sharded_state_dict'] = generate_state_dict( @@ -2727,42 +2204,26 @@ def load_checkpoint( model_sd_kwargs=model_sd_kwargs, rerun_state=gen_sd_rerun_state, ) - - if gpt_compat_layer_maps is not None: - from megatron.core.dist_checkpointing.gpt_checkpoint_interop import ( - retarget_sharded_state_dict_to_gpt_checkpoint, - ) - - # The optimizer sharded state dict is built from the (hybrid) model sharded - # state dict, so its entries carry the same ``decoder.layers..`` keys and - # sharding; the same retargeting points them at the GPT checkpoint too. - for sd_key, sub_sd in load_kwargs['sharded_state_dict'].items(): - is_model = sd_key == 'model' or re.fullmatch(r'model\d+', sd_key) - is_optim = gpt_compat_load_optim and ( - sd_key == 'optimizer' or re.fullmatch(r'optimizer\d+', sd_key) - ) - if is_model or is_optim: - retarget_sharded_state_dict_to_gpt_checkpoint(sub_sd, gpt_compat_layer_maps) - elif args.ckpt_format == 'torch_dcp': + elif args.ckpt_format == "torch_dcp": model_sd = model[0].state_dict() optimizer_sd = optimizer.state_dict(is_loading=True) if tp_group is None and pp_group is None: tp_group = mpu.get_tensor_model_parallel_group() pp_group = mpu.get_pipeline_model_parallel_group() sharded_state_dict = { - 'model': model_sd, - 'optimizer': optimizer_sd, - 'args': None, - 'iteration': 1, - 'rng_state': get_rng_state( + "model": model_sd, + "optimizer": optimizer_sd, + "args": None, + "iteration": 1, + "rng_state": get_rng_state( args.ckpt_format, tp_group, pp_group, dp_cp_group=dp_cp_group, dp_group=dp_group ), - 'checkpoint_version': None, - 'opt_param_scheduler': opt_param_scheduler.state_dict(), - 'num_floating_point_operations_so_far': 0, + "checkpoint_version": None, + "opt_param_scheduler": opt_param_scheduler.state_dict(), + "num_floating_point_operations_so_far": 0, } - load_kwargs['sharded_state_dict'] = sharded_state_dict - elif args.ckpt_format == 'fsdp_dtensor': + load_kwargs["sharded_state_dict"] = sharded_state_dict + elif args.ckpt_format == "fsdp_dtensor": reader = FileSystemReader(get_load_checkpoint_path_by_args(args)) try: state_dict_metadata = reader.read_metadata().state_dict_metadata @@ -2774,7 +2235,7 @@ def load_checkpoint( gen_sd_rng_state = None gen_sd_optim = None if not args.finetune: - if 'rerun_state_machine' in state_dict_metadata: + if "rerun_state_machine" in state_dict_metadata: gen_sd_rerun_state = get_rerun_state_machine().state_dict( data_iterator=None, ckpt_format=ckpt_format, force=True ) @@ -2782,9 +2243,8 @@ def load_checkpoint( gen_sd_rng_state = get_rng_state( args.ckpt_format, tp_group, pp_group, dp_cp_group=dp_cp_group, dp_group=dp_group ) - if (not args.finetune or gpt_compat_load_optim) and not args.no_load_optim: - gen_sd_optim = optimizer - if not args.finetune: + if not args.no_load_optim: + gen_sd_optim = optimizer gen_sd_opt_param_scheduler = opt_param_scheduler optim_sd_kwargs = dict( @@ -2792,43 +2252,18 @@ def load_checkpoint( is_loading=True, ) - # Megatron-FSDP materializes optimizer slots with a dummy zero-gradient - # step while building a loading state dict. A normal full resume - # overwrites every model parameter afterward, but GPT->Hybrid leaves - # fresh-only layers untouched. Temporarily zero the optimizer LR so the - # shape-materialization step cannot weight-decay or momentum-update - # those fresh model parameters. - optimizer_lrs = [] - if gpt_compat_layer_maps is not None and gen_sd_optim is not None: - pending_optimizers = [gen_sd_optim] - while pending_optimizers: - pending_optimizer = pending_optimizers.pop() - if hasattr(pending_optimizer, 'chained_optimizers'): - pending_optimizers.extend(pending_optimizer.chained_optimizers) - continue - inner_optimizer = getattr(pending_optimizer, 'optimizer', None) - if inner_optimizer is None: - continue - for param_group in inner_optimizer.param_groups: - optimizer_lrs.append((param_group, param_group.get('lr'))) - param_group['lr'] = 0.0 - - try: - state_dict = generate_state_dict( - args, - model=model, - optimizer=gen_sd_optim, - opt_param_scheduler=gen_sd_opt_param_scheduler, - rng_state=gen_sd_rng_state, - optim_sd_kwargs=optim_sd_kwargs, - rerun_state=gen_sd_rerun_state, - iteration=1, - ) - finally: - for param_group, lr in optimizer_lrs: - param_group['lr'] = lr - state_dict['_model'] = model - load_kwargs['sharded_state_dict'] = state_dict + state_dict = generate_state_dict( + args, + model=model, + optimizer=gen_sd_optim, + opt_param_scheduler=gen_sd_opt_param_scheduler, + rng_state=gen_sd_rng_state, + optim_sd_kwargs=optim_sd_kwargs, + rerun_state=gen_sd_rerun_state, + iteration=1, + ) + state_dict["_model"] = model + load_kwargs["sharded_state_dict"] = state_dict state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( load_dir, @@ -2837,7 +2272,6 @@ def load_checkpoint( checkpointing_context=checkpointing_context, dp_cp_group=dp_cp_group, expt_dp_group=expt_dp_group, - gpt_compat_layer_maps=gpt_compat_layer_maps, **load_kwargs, ) @@ -2850,9 +2284,9 @@ def load_checkpoint( set_checkpoint_version(state_dict.get('checkpoint_version', 0)) # Convert to regular torch tensor to DTensor. - if ckpt_type == CheckpointType.LEGACY and args.ckpt_format == 'torch_dcp': - dtensor_state_dict = _to_dtensor(ddp_model, state_dict['model']) - state_dict['model'] = dtensor_state_dict + if ckpt_type == CheckpointType.LEGACY and args.ckpt_format == "torch_dcp": + dtensor_state_dict = _to_dtensor(ddp_model, state_dict["model"]) + state_dict["model"] = dtensor_state_dict # Set iteration. if args.finetune or release: @@ -2865,9 +2299,8 @@ def load_checkpoint( iteration = state_dict['total_iters'] except KeyError: print_rank_0( - 'A metadata file exists but unable to load iteration from checkpoint {}, exiting'.format( - checkpoint_name - ) + 'A metadata file exists but unable to load ' + 'iteration from checkpoint {}, exiting'.format(checkpoint_name) ) sys.exit() num_floating_point_operations_so_far = state_dict.get('num_floating_point_operations_so_far', 0) @@ -2875,12 +2308,7 @@ def load_checkpoint( # Check arguments. if 'args' in state_dict and not args.finetune: checkpoint_args = state_dict['args'] - # A GPT block is split into separate attention and MLP positions in - # HybridModel, so num_layers intentionally differs even for an - # architecture-preserving load. Keep every other resume-time argument - # compatibility check. - skip_args = {'num_layers'} if gpt_compat_layer_maps is not None else None - check_checkpoint_args(checkpoint_args, skip_args=skip_args) + check_checkpoint_args(checkpoint_args) args.consumed_train_samples = getattr(checkpoint_args, 'consumed_train_samples', 0) args.skipped_train_samples = getattr(checkpoint_args, 'skipped_train_samples', 0) update_num_microbatches(consumed_samples=args.consumed_train_samples, verbose=True) @@ -2888,64 +2316,18 @@ def load_checkpoint( else: print_rank_0('could not find arguments in the checkpoint ...') - # --override-ckpt-iteration: rewind the data loader to this iteration, operating on `args` - # (not state_dict) so it also works on checkpoints with no saved `args` (release / HF). The - # GBS-match check applies only when adopting this checkpoint's own args (not a finetune load). - if getattr(args, 'override_ckpt_iteration', None) is not None: - if 'args' in state_dict and not args.finetune: - ckpt_global_batch_size = getattr(state_dict['args'], 'global_batch_size', None) - if ( - ckpt_global_batch_size is not None - and ckpt_global_batch_size != args.global_batch_size - ): - warn_rank_0( - '--override-ckpt-iteration recomputes consumed_train_samples = target_iter * ' - f'current global_batch_size ({args.global_batch_size}), but this checkpoint ' - f'was saved at global_batch_size {ckpt_global_batch_size}. If the target ' - "iteration is in the checkpoint run's units, the data loader resumes from a " - 'shifted sample offset.' - ) - iteration = args.override_ckpt_iteration - args.consumed_train_samples = iteration * args.global_batch_size - args.skipped_train_samples = 0 - update_num_microbatches(consumed_samples=args.consumed_train_samples, verbose=True) - print_rank_0(f'--override-ckpt-iteration: start at iteration {iteration} ' - f'(consumed_train_samples {args.consumed_train_samples})') - def load_model_state_dict(module, state_dict, strict: bool): """Helper function to load state dict with fallback for missing extra states.""" - # GTP native-FP8 weights: load_state_dict's copy_ re-quantizes into the FP8 param, which - # TE's IsMXFP8Tensor check rejects for our subclass. Present the base FP8 class for it. - from megatron.core.tensor_parallel.gtp_api import HAVE_GTP - - if HAVE_GTP: - from megatron.core.tensor_parallel.gtp_api import gtp_native_fp8_load_context - - load_ctx = lambda: gtp_native_fp8_load_context(module) - else: - from contextlib import nullcontext - - load_ctx = nullcontext try: - with load_ctx(): - module.load_state_dict(state_dict, strict=strict) + module.load_state_dict(state_dict, strict=strict) except Exception as e: if strict: # Fallback support for backward compatibility breaking changes in TransformerEngine - with load_ctx(): - load_return = module.load_state_dict(state_dict, strict=False) - print(f'load_return: {load_return}') - - # Megatron-FSDP DTensors are loaded into the model buffers in-place above. - # Replaying the translated raw state dict through ``load_state_dict`` would - # reapply HybridModel-local entries that were intentionally omitted (for - # example fresh Mamba layers). - gpt_fsdp_model_loaded_in_place = ( - ckpt_format == 'fsdp_dtensor' and gpt_compat_layer_maps is not None - ) + load_return = module.load_state_dict(state_dict, strict=False) + print(f"load_return: {load_return}") # Model. - if not skip_load_to_model_and_opt and not gpt_fsdp_model_loaded_in_place: + if not skip_load_to_model_and_opt: if len(ddp_model) == 1: load_model_state_dict(ddp_model[0], state_dict['model'], strict) else: @@ -2961,17 +2343,17 @@ def load_model_state_dict(module, state_dict, strict: bool): fix_query_key_value_ordering(model, checkpoint_version) # Optimizer. - if not release and (not args.finetune or gpt_compat_load_optim) and not args.no_load_optim: + if not release and not args.finetune and not args.no_load_optim: try: # Load state dict. if ( - getattr(args, 'use_layer_wise_distributed_optimizer', False) + getattr(args, "use_layer_wise_distributed_optimizer", False) and args.ckpt_format == 'torch' ): # LayerWiseDistributedOptimizer load optimizer state from file on different ranks dp_rank = mpu.get_data_parallel_rank() optim_checkpoint_name = os.path.join( - os.path.dirname(checkpoint_name), f'layer_wise_optimizer_{dp_rank}.pt' + os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt" ) optimizer.load_state_dict_from_file(optim_checkpoint_name) elif ( @@ -2983,11 +2365,11 @@ def load_model_state_dict(module, state_dict, strict: bool): # Load distributed optimizer's custom parameter state. # For distributed checkpoint it's already loaded in load_state_dict above - is_torch_dist = ckpt_format == 'torch_dist' + is_torch_dist = ckpt_format == "torch_dist" if ( args.use_distributed_optimizer and not is_torch_dist - and ckpt_format not in ['torch_dcp', 'fsdp_dtensor'] + and ckpt_format not in ["torch_dcp", "fsdp_dtensor"] ): # NOTE: this is a manual read of the tracker file. # This code should not be reached when reading from a non_persistent checkpoint @@ -3003,53 +2385,12 @@ def load_model_state_dict(module, state_dict, strict: bool): update_legacy_format=args.ckpt_convert_update_legacy_dist_opt_format, ) - # Load scheduler unless --finetune requests a fresh iteration and LR schedule. - if opt_param_scheduler is not None and not args.finetune: + # Load scheduler. + if opt_param_scheduler is not None: if 'lr_scheduler' in state_dict: # backward compatbility opt_param_scheduler.load_state_dict(state_dict['lr_scheduler']) else: opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler']) - - # Optimizer state dict can overwrite per-group max_lr/min_lr values. - # If scheduler override is requested, re-apply runtime lr bounds from - # args so scheduler math does not stay pinned to checkpoint lr settings. - if getattr(args, "override_opt_param_scheduler", False): - if ( - optimizer is None - or getattr(optimizer, "is_stub_optimizer", False) - or opt_param_scheduler is None - ): - print_rank_0( - " > WARNING: --override-opt-param-scheduler is set, but optimizer or " - "opt_param_scheduler is not available; skipping override of scheduler " - "and optimizer param_group max_lr/min_lr." - ) - else: - for param_group in optimizer.param_groups: - if param_group.get("is_decoupled_lr", False): - max_lr = getattr(args, "decoupled_lr", None) - min_lr = getattr(args, "decoupled_min_lr", None) - else: - max_lr = args.lr - min_lr = args.min_lr - if max_lr is not None: - param_group["max_lr"] = max_lr - if min_lr is not None: - param_group["min_lr"] = min_lr - # Synchronize scheduler num_steps with consumed_train_samples - # to ensure lr calculation is based on current training progress - if opt_param_scheduler.num_steps != args.consumed_train_samples: - print_rank_0( - f" > WARNING: scheduler num_steps ({opt_param_scheduler.num_steps}) " - f"differs from consumed_train_samples ({args.consumed_train_samples}). " - f"Resetting scheduler num_steps to match consumed_train_samples." - ) - opt_param_scheduler.num_steps = args.consumed_train_samples - opt_param_scheduler.step(increment=0) - print_rank_0( - " > restored optimizer param_group max_lr/min_lr from runtime config " - "because --override-opt-param-scheduler is set" - ) except KeyError as e: print_rank_0( 'Unable to load optimizer from checkpoint {}. ' @@ -3071,7 +2412,7 @@ def load_model_state_dict(module, state_dict, strict: bool): if 'rerun_state_machine' in state_dict: get_rerun_state_machine().load_state_dict(state_dict['rerun_state_machine']) except Exception as e: - print_rank_0(f'Unable to restore RerunMachine from checkpoint: {e}. Skipping.') + print_rank_0(f"Unable to restore RerunMachine from checkpoint: {e}. Skipping.") # rng states. if not release and not args.finetune and not args.no_load_rng and not ignore_rng_state: @@ -3079,7 +2420,7 @@ def load_model_state_dict(module, state_dict, strict: bool): cuda_rng_tracker = tensor_parallel.get_cuda_rng_tracker() graph_safe_rng = tensor_parallel.is_graph_safe_cuda_rng_tracker(cuda_rng_tracker) if 'rng_state' in state_dict: - if args.ckpt_format == 'fsdp_dtensor': + if args.ckpt_format == "fsdp_dtensor": # FSDP DTensor checkpoints store rng_state in a different format. tp_rank = ( get_pg_rank(tp_group) @@ -3091,10 +2432,10 @@ def load_model_state_dict(module, state_dict, strict: bool): if pp_group is not None else mpu.get_pipeline_model_parallel_rank() ) - if f'({pp_rank}, {tp_rank})' in state_dict['rng_state']: - rng_state = state_dict['rng_state'][f'({pp_rank}, {tp_rank})'] + if f"({pp_rank}, {tp_rank})" in state_dict['rng_state']: + rng_state = state_dict['rng_state'][f"({pp_rank}, {tp_rank})"] else: - print_rank_0('WARNING: RNG state not found for current TP/PP rank') + print_rank_0("WARNING: RNG state not found for current TP/PP rank") rng_state = next(iter(state_dict['rng_state'].values())) else: rng_state = state_dict['rng_state'] @@ -3160,12 +2501,9 @@ def load_model_state_dict(module, state_dict, strict: bool): if pp_group is not None else mpu.get_pipeline_model_parallel_world_size() ) - _gtp_remat_r = mpu.get_gtp_weight_remat_rank() - _gtp_remat_w = mpu.get_gtp_weight_remat_world_size() print_rank_0( f' successfully loaded checkpoint from {load_dir} ' f'[ t {_tp_r + 1}/{_tp_w}, ' - f'gtp_remat {_gtp_remat_r + 1}/{_gtp_remat_w}, ' f'p {_pp_r + 1}/{_pp_w} ] ' f'at iteration {iteration}' ) @@ -3194,7 +2532,28 @@ def load_model_state_dict(module, state_dict, strict: bool): log_printed = True if has_nvidia_modelopt: - print_distributed_quant_summary(model, msg='After loading checkpoint') + print_distributed_quant_summary(model, msg="After loading checkpoint") + + # Load teacher model in Distillation mode. + if getattr(args, "export_kd_teacher_load", None): + from megatron.post_training.checkpointing import load_modelopt_checkpoint + + unwrapped_model = unwrap_model(model)[0] + # Note: load_modelopt_checkpoint may call this function so we prevent infinite recursion. + if hasattr(unwrapped_model, 'teacher_model'): + teacher = unwrapped_model.teacher_model + print_rank_0( + f"Loading teacher as {type(teacher).__name__} from {args.export_kd_teacher_load} ..." + ) + # [WAR]: To avoid error out on loading teacher's checkpoint, we temporarily + # set args.finetune to True while loading the teacher checkpoint. + original_args_finetune, original_ckpt_format = args.finetune, args.ckpt_format + args.finetune = True + if args.export_kd_teacher_ckpt_format is not None: + args.ckpt_format = args.export_kd_teacher_ckpt_format + load_modelopt_checkpoint([teacher], load_arg='export_kd_teacher_load') + args.finetune, args.ckpt_format = original_args_finetune, original_ckpt_format + print_rank_0("... teacher loaded successfully.") return iteration, num_floating_point_operations_so_far @@ -3205,7 +2564,7 @@ def _to_dtensor(wrapped_model, model_state_dict): new_model_sd = dict() for k, v in model_state_dict.items(): # FP8 extra state cannot be converted to dtensor yet. - if '_extra_state' in k: + if "_extra_state" in k: new_model_sd[k] = v else: new_model_sd[k] = torch.distributed.tensor.distribute_tensor(v, device_mesh) @@ -3229,7 +2588,7 @@ def load_biencoder_checkpoint( tracker_filename = get_checkpoint_tracker_filename(load_path) - with maybe_msc.open(tracker_filename, 'r') as f: + with open_file(tracker_filename, 'r') as f: iteration = int(f.read().strip()) checkpoint_name = get_checkpoint_name( @@ -3243,7 +2602,7 @@ def load_biencoder_checkpoint( ) ) - state_dict = torch.load(checkpoint_name, map_location='cpu') + state_dict = torch.load(checkpoint_name, map_location='cpu', weights_only=True) ret_state_dict = state_dict['model'] if only_query_model: diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 7f70205a51c..e9169d77144 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -334,6 +334,11 @@ def forward_step(data_iterator, model: HybridModel): cu_seqlens_for_params = cu_seqlens packed_seq_params = PackedSeqParams( qkv_format="thd", + cp_partition_mode=( + "contiguous" + if args.linear_cp_mode == "chunkwise" and args.context_parallel_size > 1 + else "zigzag" + ), cu_seqlens_q=cu_seqlens_for_params, cu_seqlens_kv=cu_seqlens_for_params, cu_seqlens_q_padded=cu_seqlens_padded, diff --git a/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py b/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py index 2dcf06c6b57..9eb573e63de 100644 --- a/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py +++ b/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py @@ -45,27 +45,26 @@ def collect_train_test_metrics( if golden_value_key in [ "iteration-time", + "grad-norm", "mem-allocated-bytes", "mem-max-allocated-bytes", "lm loss", "num-zeros", "mtp_1 loss", - "mtp_2 loss", - "total loss", + "z_loss", ] } if output_path is not None: - serialized_summaries = json.dumps( - { - golden_value_key: golden_values.model_dump() - for golden_value_key, golden_values in summaries.items() - }, - indent=4, - allow_nan=False, - ) with open(output_path, "w") as fh: - fh.write(serialized_summaries) + json.dump( + { + golden_value_key: golden_values.model_dump() + for golden_value_key, golden_values in summaries.items() + }, + fh, + indent=4, + ) if __name__ == "__main__": diff --git a/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py b/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py index 832558adaad..a39f7ebb92e 100644 --- a/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py +++ b/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py @@ -13,15 +13,15 @@ CHECK_THRESHOLDS = { "iteration-time": [common.ApproximateTest(atol=0, rtol=0.05)], + "grad-norm": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], "mem-allocated-bytes": [common.ApproximateTest(atol=0, rtol=0.05)], "mem-max-allocated-bytes": [common.ApproximateTest(atol=0, rtol=0.05)], "lm loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], "mtp_1 loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], - "mtp_2 loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], + "z_loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], "num-zeros": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.20)], "generated_tokens": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], "logprobs": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], - "total loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], } diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp1_ep8_cp2_kda_gated_mla_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp1_ep8_cp2_kda_gated_mla_1N8G/model_config.yaml new file mode 100644 index 00000000000..5b6bfd80e9d --- /dev/null +++ b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp1_ep8_cp2_kda_gated_mla_1N8G/model_config.yaml @@ -0,0 +1,128 @@ +ENV_VARS: + CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_NONDETERMINISTIC_ALGO: 1 + NON_DETERMINSTIC_RESULTS: 1 + NCCL_ALGO: Ring + CUBLAS_WORKSPACE_CONFIG: :4096:8 + PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True + ENABLE_LIGHTWEIGHT_MODE: true + +MODEL_ARGS: + # Reduced Hybrid stack: 3 KDA, 1 MLA, 1 dense MLP, 3 MoE, and 1 MTP depth. + --spec: "[megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec]" + --hybrid-layer-pattern: K-KEKE+E/+E + + # Self-contained mock-data setup. + --mock-data: true + --tokenizer-type: NullTokenizer + --vocab-size: 4096 + --make-vocab-size-divisible-by: 128 + --seq-length: 256 + --max-position-embeddings: 256 + --micro-batch-size: 1 + --global-batch-size: 8 + --train-iters: 50 + --eval-interval: 1000000000 + --eval-iters: 0 + + # Reduced model dimensions exercise the attention and FFN paths. + --hidden-size: 512 + --num-attention-heads: 8 + --kv-channels: 64 + --ffn-hidden-size: 1536 + --normalization: RMSNorm + --norm-epsilon: 1.0e-6 + --swiglu: true + --disable-bias-linear: true + --untie-embeddings-and-output-weights: true + --position-embedding-type: rope + --rotary-percent: 0.5 + --rotary-base: 6000000 + --transformer-impl: transformer_engine + --attention-backend: fused + + # One-node topology exercises the EP=8 and CP=2 paths. + --tensor-model-parallel-size: 1 + --pipeline-model-parallel-size: 1 + --expert-model-parallel-size: 8 + --context-parallel-size: 2 + --linear-cp-mode: headwise + --cp-comm-type: p2p + --distributed-backend: nccl + + # Grouped sigmoid MoE routing with 32 experts. + --num-experts: 32 + --moe-ffn-hidden-size: 256 + --moe-shared-expert-intermediate-size: 256 + --moe-router-topk: 8 + --moe-router-score-function: sigmoid + --moe-router-dtype: fp32 + --moe-router-topk-scaling-factor: 2.5 + --moe-router-num-groups: 8 + --moe-router-group-topk: 4 + --moe-router-enable-expert-bias: true + --moe-router-bias-update-rate: 0.0 + --moe-router-load-balancing-type: none + --moe-z-loss-coeff: 2.9e-6 + --moe-token-dispatcher-type: alltoall + --moe-grouped-gemm: true + + # Kimi Delta Attention. + --linear-conv-kernel-dim: 4 + --linear-key-head-dim: 64 + --linear-value-head-dim: 64 + --linear-num-key-heads: 8 + --linear-num-value-heads: 8 + --kda-safe-gate: true + --kda-lower-bound: -5.0 + + # Gated MLA with native unequal QK/V head dimensions. + --multi-latent-attention: true + --q-lora-rank: 128 + --kv-lora-rank: 256 + --qk-head-dim: 64 + --qk-pos-emb-head-dim: 32 + --v-head-dim: 64 + --qk-layernorm: true + --attention-output-gate: true + --gated-attention-proj-granularity: headwise + + # One MTP depth represented by the pattern suffix after '/'. + --mtp-num-layers: 1 + --mtp-loss-scaling-factor: 0.1 + + --recompute-granularity: full + --recompute-method: uniform + --recompute-num-layers: 1 + --bf16: true + --optimizer: adam + --lr: 1.0e-4 + --min-lr: 1.0e-4 + --lr-decay-style: constant + --weight-decay: 0.1 + --clip-grad: 1.0 + --attention-dropout: 0.0 + --hidden-dropout: 0.0 + --seed: 1234 + + --log-interval: 1 + --log-num-zeros-in-grad: true + --log-params-norm: true + --log-timers-to-tensorboard: true + --log-memory-to-tensorboard: true + --tensorboard-dir: ${TENSORBOARD_PATH} + --save: ${CHECKPOINT_SAVE_PATH} + --load: ${CHECKPOINT_LOAD_PATH} + --ckpt-format: torch_dist + --timing-log-level: 0 + +TEST_TYPE: ckpt-resume +METRICS: + - lm loss + - mtp_1 loss + - grad-norm + - z_loss + - num-zeros + - mem-allocated-bytes + - mem-max-allocated-bytes +LAUNCHER: ft_launcher diff --git a/tests/test_utils/recipes/h100/mamba.yaml b/tests/test_utils/recipes/h100/mamba.yaml index b4596e2d71a..8ca2fed4168 100644 --- a/tests/test_utils/recipes/h100/mamba.yaml +++ b/tests/test_utils/recipes/h100/mamba.yaml @@ -87,6 +87,12 @@ products: # - environment: [lts] # disabled until triton is bumped # scope: [nightly] + - test_case: [hybrid_mr_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp] + products: + - environment: [dev] + scope: [mr, mr-github, mr-github-slim] + platforms: [dgx_h100] + - test_case: [hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G] products: - environment: [dev] @@ -103,9 +109,9 @@ products: # - environment: [lts] # disabled until triton is bumped # scope: [nightly] - - test_case: [hybrid_nemotron_v3_pico_7b_a1b_tp1_ep8_QAD_dgx_h100_1N8G] + - test_case: [hybrid_mr_mcore_te_tp1_pp1_ep8_cp2_kda_gated_mla_1N8G] products: - # Disabled while the deterministic total-loss mismatch is investigated. + # Enable after the official H100 functional run produces golden values. - environment: [dev] - scope: [mr-broken, mr-github-broken] + scope: [mr, mr-github, mr-github-slim] platforms: [dgx_h100] diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index b7789575014..e1da4151cbe 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -179,6 +179,22 @@ def test_initialize_dynamic_context(self, is_hybrid_model: bool): # Check initializations to -1 assert torch.all(dynamic_context.request_ids == -1) + @pytest.mark.internal + def test_kda_inference_fails_early(self): + with pytest.raises(NotImplementedError, match="KDA layers"): + self._get_dynamic_context( + params_dtype=torch.float32, + num_layers=1, + kv_channels=8, + num_attention_heads=2, + max_sequence_length=128, + buffer_size_gb=0.03, + block_size_tokens=64, + max_tokens=None, + is_hybrid_model=True, + layer_type_list=[Symbols.KDA], + ) + @pytest.mark.internal def test_is_static_batching(self): diff --git a/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py b/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py index d9c4c14b1f7..3b277d4f5b8 100644 --- a/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py +++ b/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from unittest.mock import MagicMock, patch @@ -8,7 +8,10 @@ from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_block import TransformerBlockSubmodules -from megatron.core.transformer.transformer_layer import TransformerLayer +from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, + TransformerLayer, +) # --------------------------------------------------------------------------- # Helpers: fake backend and config builders @@ -82,6 +85,7 @@ def _make_config(**overrides): use_kitchen_attention=False, kitchen_attention_backend="sdpa", fallback_to_eager_attn=False, + enable_hyper_connections=False, ) defaults.update(overrides) cfg = MagicMock() @@ -108,9 +112,10 @@ def _fn(variant): "variant, expected", [ ("gdn", True), - ("gdn2", True), + ("kda", True), ("gated_delta_net", True), ("dsa", False), + ("dsv4_hybrid", False), (None, False), ("some_unknown_variant", False), ], @@ -120,6 +125,30 @@ def test_variants(self, variant, expected): assert self._fn(variant) is expected +class TestNormalizeExperimentalAttentionVariant: + """Validate canonical variant names and deprecated aliases.""" + + @staticmethod + def _fn(variant): + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + normalize_experimental_attention_variant, + ) + + return normalize_experimental_attention_variant(variant) + + @pytest.mark.parametrize("variant", ["gdn", "kda", "dsa", "dsv4_hybrid", None]) + def test_canonical_names_are_unchanged(self, variant): + """Canonical and unrelated variant names pass through unchanged.""" + assert self._fn(variant) == variant + + def test_legacy_gated_delta_net_alias_warns_and_normalizes(self): + """The legacy spelling warns and resolves to canonical GDN.""" + with pytest.warns(DeprecationWarning, match="Use 'gdn' instead"): + normalized = self._fn("gated_delta_net") + + assert normalized == "gdn" + + # =================================================================== # Tests for get_moe_layer_pattern # =================================================================== @@ -247,6 +276,26 @@ def test_returns_correct_module_spec(self): assert spec.module is GatedDeltaNet assert spec.metainfo == {"fuse_input_layernorm": True} + def test_kda_uses_direct_projection_submodules(self): + """KDA uses separate input and beta projections without fused input norm.""" + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_gated_delta_net_module_spec, + ) + from megatron.core.ssm.gated_delta_net import ( + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, + ) + + backend = _make_backend() + cfg = _make_config(experimental_attention_variant="kda", normalization="RMSNorm") + spec = get_gated_delta_net_module_spec(cfg, backend=backend) + + assert spec.module is KimiDeltaAttention + assert isinstance(spec.submodules, KimiDeltaAttentionSubmodules) + assert spec.submodules.in_proj == _FakeColumnParallelLinear + assert spec.submodules.beta_proj == _FakeColumnParallelLinear + assert spec.metainfo == {"fuse_input_layernorm": False} + def test_submodules_use_backend_modules(self): """Verify backend-provided projection/norm modules are wired into submodules.""" from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( @@ -420,9 +469,10 @@ class TestGetExperimentalAttentionVariantModuleSpec: "variant, target_fn", [ ("gdn", "get_gated_delta_net_module_spec"), - ("gdn2", "get_gated_delta_net_module_spec"), + ("kda", "get_gated_delta_net_module_spec"), ("gated_delta_net", "get_gated_delta_net_module_spec"), ("dsa", "get_dsa_module_spec_for_backend"), + ("dsv4_hybrid", "get_dsv4_hybrid_module_spec_for_backend"), ], ) def test_dispatches_to_variant_handler(self, variant, target_fn): @@ -539,7 +589,7 @@ def test_hybrid_attention_pattern(self): assert specs[2].submodules.self_attention is exp_attn_spec assert specs[3].submodules.self_attention is std_attn_spec - def test_hybrid_moe_pattern(self): + def test_hybrid_moe_pattern_with_mhc(self): """Verify MLP alternates between MoE and dense specs per moe_layer_freq pattern.""" from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( get_transformer_layer_with_experimental_attention_variant_spec, @@ -551,6 +601,7 @@ def test_hybrid_moe_pattern(self): num_moe_experts=8, moe_layer_freq=2, normalization="RMSNorm", + enable_hyper_connections=True, ) backend = _make_backend() attn_spec = self._make_attention_spec(fuse_input_layernorm=False) @@ -574,6 +625,8 @@ def test_hybrid_moe_pattern(self): assert specs[1].submodules.mlp is dense_spec assert specs[2].submodules.mlp is moe_spec assert specs[3].submodules.mlp is dense_spec + for s in specs: + assert s.module is HyperConnectionTransformerLayer # =================================================================== diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 2eac69295e1..1113c296e34 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -141,6 +141,7 @@ "fp8_wgrad": True, "fused_residual_rmsnorm": False, "fused_single_qkv_rope": False, + "gated_attention_proj_granularity": "elementwise", "gated_linear_unit": False, "gdp_cutedsl_kernel": False, "gdp_num_chunk_states_to_recompute": 2, @@ -165,6 +166,8 @@ "init_method_std": 0.014, "init_model_with_meta_device": False, "is_hybrid_model": True, + "kda_lower_bound": None, + "kda_safe_gate": False, "kitchen_attention_backend": "sdpa", "kv_channels": 128, "layernorm_epsilon": 1e-05, diff --git a/tests/unit_tests/optimizer/test_qk_clip.py b/tests/unit_tests/optimizer/test_qk_clip.py new file mode 100644 index 00000000000..a82bef0e502 --- /dev/null +++ b/tests/unit_tests/optimizer/test_qk_clip.py @@ -0,0 +1,98 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import torch + +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.optimizer.qk_clip import clip_qk +from megatron.core.transformer.attention import Attention + + +class _FakeCoreAttention(torch.nn.Module): + def __init__(self, max_logits): + super().__init__() + self.current_max_attn_logits = torch.tensor(max_logits, dtype=torch.float32) + + +class _FakeAttention(Attention): + def __init__(self, max_logits): + torch.nn.Module.__init__(self) + self.core_attention = _FakeCoreAttention(max_logits) + self.clip_calls = 0 + + def clip_qk(self): + self.clip_calls += 1 + self.core_attention.current_max_attn_logits = None + + def get_query_key_value_tensors(self, *args, **kwargs): + raise NotImplementedError + + +class _FakeQKLikeModule(torch.nn.Module): + """A non-Attention module that happens to expose the same attribute names.""" + + def __init__(self): + super().__init__() + self.core_attention = _FakeCoreAttention([1000.0]) + self.clip_calls = 0 + + def clip_qk(self): + self.clip_calls += 1 + + +class _FakeHybridModel(HybridModel): + def __init__(self, decoder_attention, mtp_attention): + torch.nn.Module.__init__(self) + self.decoder = torch.nn.ModuleList([decoder_attention]) + self.mtp = torch.nn.Module() + self.mtp.mtp_model_layer = torch.nn.ModuleList([mtp_attention]) + self.unrelated = _FakeQKLikeModule() + + +class _ModelWrapper(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.module = torch.nn.Module() + self.module.module = model + + +def test_clip_qk_reaches_nested_mtp(monkeypatch): + decoder_attention = _FakeAttention([80.0, 120.0]) + mtp_attention = _FakeAttention([150.0, 90.0]) + model = _FakeHybridModel(decoder_attention, mtp_attention) + model_chunk = _ModelWrapper(model) + all_reduce_calls = [] + + monkeypatch.setattr( + torch.distributed, + "all_reduce", + lambda tensor, **kwargs: all_reduce_calls.append((tensor, kwargs)), + ) + monkeypatch.setattr( + "megatron.core.optimizer.qk_clip.parallel_state.get_data_parallel_group", + lambda **kwargs: None, + ) + + assert clip_qk([model_chunk]) == 150.0 + assert decoder_attention.clip_calls == 1 + assert mtp_attention.clip_calls == 1 + assert model.unrelated.clip_calls == 0 + assert len(all_reduce_calls) == 2 + + +def test_clip_qk_log_only_resets_nested_mtp(monkeypatch): + decoder_attention = _FakeAttention([80.0]) + mtp_attention = _FakeAttention([150.0]) + model = _FakeHybridModel(decoder_attention, mtp_attention) + model_chunk = _ModelWrapper(model) + + monkeypatch.setattr(torch.distributed, "all_reduce", lambda *args, **kwargs: None) + monkeypatch.setattr( + "megatron.core.optimizer.qk_clip.parallel_state.get_data_parallel_group", + lambda **kwargs: None, + ) + + assert clip_qk([model_chunk], log_max_only=True) == 150.0 + assert decoder_attention.clip_calls == 0 + assert mtp_attention.clip_calls == 0 + assert decoder_attention.core_attention.current_max_attn_logits is None + assert mtp_attention.core_attention.current_max_attn_logits is None diff --git a/tests/unit_tests/ssm/gated_delta_net/test_gdn.py b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py new file mode 100644 index 00000000000..325b1f5c2d1 --- /dev/null +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py @@ -0,0 +1,841 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import copy +import inspect +import os +from unittest import mock + +import pytest +import torch +import torch.nn.functional as F + +from megatron.core import parallel_state +from megatron.core.dist_checkpointing.mapping import ShardedTensorFactory +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_experimental_attention_variant_module_spec, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gated_delta_net import GatedDeltaNet, torch_chunk_gated_delta_rule +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from tests.unit_tests.test_utilities import Utils +from tests.unit_tests.transformer.test_multi_latent_attention import ( + make_test_packed_seq_params, + make_test_packed_seq_params_with_padding, +) + +try: + import fla + + HAVE_FLA = True +except ImportError: + HAVE_FLA = False + +# https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-multi-rank-gpu-enable +# NVLS doesn't support one single GPU to be shared by multiple ranks, so disable this in test +os.environ.update({"NCCL_NVLS_ENABLE": "0"}) + +try: + from causal_conv1d.cpp_functions import causal_conv1d_bwd_function +except ImportError: + HAVE_FUSED_PRE_GDR = False +else: + HAVE_FUSED_PRE_GDR = callable(causal_conv1d_bwd_function) + + +def _make_gdn_config(**overrides): + config_kwargs = { + "hidden_size": 128, + "linear_conv_kernel_dim": 2, + "linear_key_head_dim": 32, + "linear_value_head_dim": 32, + "linear_num_key_heads": 4, + "linear_num_value_heads": 8, + "num_layers": 1, + "normalization": "RMSNorm", + "use_cpu_initialization": True, + "layernorm_zero_centered_gamma": True, + "num_attention_heads": 8, + "activation_func": F.silu, + "bf16": True, + "experimental_attention_variant": "gdn", + "linear_attention_freq": [1], + "transformer_impl": "transformer_engine", + } + config_kwargs.update(overrides) + return TransformerConfig(**config_kwargs) + + +def _set_gdn_test_cp_partition_mode(packed_seq_params, cp_size, linear_cp_mode): + if cp_size <= 1: + return packed_seq_params + if linear_cp_mode == "headwise": + packed_seq_params.cp_partition_mode = "zigzag" + elif linear_cp_mode == "chunkwise": + packed_seq_params.cp_partition_mode = "contiguous" + else: + raise ValueError(f"Invalid linear CP mode: {linear_cp_mode}") + return packed_seq_params + + +def test_gdn_pre_gated_delta_rule_fusion_defaults_to_disabled(): + config = _make_gdn_config() + assert not config.gdn_pre_gated_delta_rule_fusion + + +def test_gdn_pre_gated_delta_rule_fusion_accepts_gdn_variant(): + config = _make_gdn_config(gdn_pre_gated_delta_rule_fusion=True) + assert config.gdn_pre_gated_delta_rule_fusion + + +def test_gdn_pre_gated_delta_rule_fusion_requires_gdn_variant(): + with pytest.raises(ValueError, match="experimental_attention_variant='gdn'"): + _make_gdn_config( + experimental_attention_variant=None, + linear_attention_freq=None, + gdn_pre_gated_delta_rule_fusion=True, + ) + + +def test_gdn_norm_out_recompute_accepts_gdn_variant(): + config = _make_gdn_config(recompute_granularity="selective", recompute_modules=["gdn_norm_out"]) + assert "gdn_norm_out" in config.recompute_modules + + +def test_gdn_norm_out_recompute_rejects_non_hybrid_non_gdn_config(): + with pytest.raises(ValueError, match="gdn_norm_out in recompute_modules"): + _make_gdn_config( + experimental_attention_variant=None, + linear_attention_freq=None, + recompute_granularity="selective", + recompute_modules=["gdn_norm_out"], + ) + + +def test_gdn_and_norm_out_recompute_are_mutually_exclusive(): + with pytest.raises(ValueError, match="'gdn' and 'gdn_norm_out'"): + _make_gdn_config( + recompute_granularity="selective", recompute_modules=["gdn", "gdn_norm_out"] + ) + + +def test_gdn_norm_out_recompute_accepts_non_experimental_hybrid_config(): + # Hybrid specs select GDN/KDA per layer without setting a global + # experimental_attention_variant, so the selector must remain valid here. + config = _make_gdn_config( + experimental_attention_variant=None, + is_hybrid_model=True, + linear_attention_freq=None, + recompute_granularity="selective", + recompute_modules=["gdn_norm_out"], + ) + + assert config.recompute_modules == ["gdn_norm_out"] + + +def test_gdn_conv_pad_alignment_rejects_chunkwise_cp(): + with pytest.raises(AssertionError, match="gdn_conv_pad_alignment is incompatible"): + _make_gdn_config( + context_parallel_size=2, linear_cp_mode="chunkwise", gdn_conv_pad_alignment=4096 + ) + + +def test_gdn_chunkwise_cp_head_divisibility_ignores_cp_size(): + config = _make_gdn_config( + tensor_model_parallel_size=2, + context_parallel_size=4, + linear_cp_mode="chunkwise", + linear_num_key_heads=4, + linear_num_value_heads=8, + ) + assert config.linear_cp_mode == "chunkwise" + + +def test_torch_chunk_gated_delta_rule_preserves_public_signature(): + signature = inspect.signature(torch_chunk_gated_delta_rule) + assert tuple(signature.parameters) == ( + "q", + "k", + "v", + "g", + "beta", + "chunk_size", + "initial_state", + "output_final_state", + "use_qk_l2norm_in_kernel", + "cu_seqlens", + "cp_context", + "scale", + ) + + +def test_gdn_headwise_cp_head_divisibility_includes_cp_size(): + with pytest.raises(AssertionError, match="linear_head_parallel_size"): + _make_gdn_config( + tensor_model_parallel_size=2, + context_parallel_size=4, + linear_cp_mode="headwise", + linear_num_key_heads=4, + linear_num_value_heads=8, + ) + + +@pytest.mark.parametrize( + ("tp_size", "sp", "cp_size", "linear_cp_mode"), + [ + # cp_size=1: the CP path is inactive, so linear_cp_mode choice is irrelevant. + # Cover the "chunkwise" default and skip the "headwise" variants for brevity. + (1, False, 1, None), + (2, False, 1, None), + (2, True, 1, None), + # cp_size=2: exercise both CP paths. + (1, False, 2, "headwise"), + (2, False, 2, "headwise"), + (2, True, 2, "headwise"), + (1, False, 2, "chunkwise"), + (2, False, 2, "chunkwise"), + (2, True, 2, "chunkwise"), + ], +) +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +class TestGatedDeltaNet: + + @pytest.fixture(scope='function', autouse=True) + def setup_method(self, tp_size, sp, cp_size, linear_cp_mode): + # Initialize parallel and random seed + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp_size, + pipeline_model_parallel_size=1, + context_parallel_size=cp_size, + ) + model_parallel_cuda_manual_seed(123) + self.tp_size = tp_size + self.cp_size = cp_size + self.sp_size = tp_size if sp else 1 + self.linear_cp_mode = linear_cp_mode + if self.linear_cp_mode == "headwise": + self.cp_size_chunkwise = 1 + self.cp_size_headwise = self.cp_size + elif self.linear_cp_mode == "chunkwise": + self.cp_size_chunkwise = self.cp_size + self.cp_size_headwise = 1 + elif self.cp_size == 1: + self.cp_size_chunkwise = 1 + self.cp_size_headwise = 1 + else: + raise ValueError(f"Invalid linear CP mode: {self.linear_cp_mode}") + + # Get TP and CP process groups from device mesh + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + tp_cp_group = parallel_state.get_tensor_and_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group, tp_cp=tp_cp_group) + + # Initialize model, with the same config as Qwen Next except `num_layers` + self.transformer_config = TransformerConfig( + hidden_size=2048, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + num_layers=1, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + num_attention_heads=16, + num_query_groups=2, + activation_func=F.silu, + bf16=True, + tensor_model_parallel_size=tp_size, + sequence_parallel=sp, + context_parallel_size=cp_size, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + linear_cp_mode=self.linear_cp_mode, + transformer_impl="transformer_engine", + ) + gdn_submodules = get_experimental_attention_variant_module_spec( + config=self.transformer_config + ).submodules + + self.gdn = GatedDeltaNet( + self.transformer_config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + self.gdn = self.gdn.cuda().bfloat16() + + def teardown_method(self): + Utils.destroy_model_parallel() + + def test_gpu_forward(self): + gdn = self.gdn + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + attention_mask = None + + output, bias = gdn(hidden_states, attention_mask) + + assert output.dim() == 3, f"Output too many dimensions ({output.shape=})" + assert output.shape[0] == seq_length // self.sp_size // self.cp_size, ( + f"Output shape {output.shape[0]=} mismatch with " + f" {seq_length=} // {self.sp_size=} // {self.cp_size=}." + ) + assert ( + output.shape[1] == micro_batch_size + ), f"Output shape {output.shape[1]=} mismatch with {micro_batch_size=}" + assert ( + output.shape[2] == gdn.config.hidden_size + ), f"Output shape {output.shape[2]=} mismatch with {gdn.config.hidden_size=}" + assert ( + output.dtype == hidden_states.dtype + ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + + @pytest.mark.flaky_in_dev # Issue #5473 + def test_selective_recompute_gdn(self): + """Whole-module 'gdn' recompute must match the non-recompute forward and gradients. + + The same module/input is run twice (recompute off, then on); the forward output and + all parameter / input gradients must agree within a tight tolerance (rtol/atol=1e-4). + The recompute path is run-to-run deterministic on these kernels (empirically bitwise), + so a tolerance well below the bf16 floor is expected to hold. + """ + gdn = self.gdn + gdn.train() + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + torch.manual_seed(1234) + base_input = torch.randn( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + def run(recompute): + gdn.recompute_gdn = recompute + gdn.zero_grad(set_to_none=True) + hidden_states = base_input.clone().detach().requires_grad_(True) + output, _ = gdn(hidden_states, None) + output.float().square().mean().backward() + param_grads = { + name: param.grad.detach().clone() + for name, param in gdn.named_parameters() + if param.grad is not None + } + return output.detach().clone(), hidden_states.grad.detach().clone(), param_grads + + try: + out_ref, dinput_ref, pgrad_ref = run(recompute=False) + out_rc, dinput_rc, pgrad_rc = run(recompute=True) + finally: + gdn.recompute_gdn = False + + torch.testing.assert_close(out_rc, out_ref, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(dinput_rc, dinput_ref, rtol=1e-4, atol=1e-4) + assert pgrad_ref.keys() == pgrad_rc.keys(), "recompute changed the set of grad params" + assert len(pgrad_ref) > 0, "expected at least one parameter gradient" + for name in pgrad_ref: + torch.testing.assert_close( + pgrad_rc[name], + pgrad_ref[name], + rtol=1e-4, + atol=1e-4, + msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}", + ) + + @pytest.mark.flaky_in_dev # Issue #5473 + def test_selective_recompute_gdn_norm_out(self): + """Output-discarding 'gdn_norm_out' recompute must preserve outputs and gradients.""" + gdn = self.gdn + gdn.train() + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + torch.manual_seed(1234) + base_input = torch.randn( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + def run(recompute_norm_out): + gdn.recompute_gdn = False + gdn.recompute_norm_out = recompute_norm_out + gdn.norm_out_checkpoint = None + gdn.zero_grad(set_to_none=True) + hidden_states = base_input.clone().detach().requires_grad_(True) + output, _ = gdn(hidden_states, None) + output.float().square().mean().backward() + param_grads = { + name: param.grad.detach().clone() + for name, param in gdn.named_parameters() + if param.grad is not None + } + return output.detach().clone(), hidden_states.grad.detach().clone(), param_grads + + try: + out_ref, dinput_ref, pgrad_ref = run(recompute_norm_out=False) + out_rc, dinput_rc, pgrad_rc = run(recompute_norm_out=True) + finally: + gdn.recompute_norm_out = False + gdn.norm_out_checkpoint = None + + torch.testing.assert_close(out_rc, out_ref, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(dinput_rc, dinput_ref, rtol=1e-4, atol=1e-4) + assert pgrad_ref.keys() == pgrad_rc.keys(), "recompute changed the set of grad params" + assert len(pgrad_ref) > 0, "expected at least one parameter gradient" + for name in pgrad_ref: + torch.testing.assert_close( + pgrad_rc[name], + pgrad_ref[name], + rtol=1e-4, + atol=1e-4, + msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}", + ) + + def test_gpu_forward_rejects_sbhd_chunkwise_cp_batch_gt_one(self): + if not (self.linear_cp_mode == "chunkwise" and self.cp_size > 1): + pytest.skip("Only chunkwise CP with CP>1 uses the FLA CP batch guard.") + + gdn = self.gdn + + micro_batch_size = 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + with pytest.raises(ValueError, match="requires micro_batch_size == 1"): + gdn(hidden_states, None) + + def test_gpu_forward_rejects_sbhd_conv_padding(self): + gdn = self.gdn + gdn.config.gdn_conv_pad_alignment = 4096 + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + expected_error = ( + "incompatible with GDN chunkwise CP" + if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 + else "only supported with packed sequence" + ) + with pytest.raises(ValueError, match=expected_error): + gdn(hidden_states, None) + + def test_deterministic_mode(self): + if self.cp_size > 1: + pytest.skip( + "deterministic_mode uses torch_chunk_gated_delta_rule, which does not support CP." + ) + + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + det_config = copy.deepcopy(self.transformer_config) + det_config.deterministic_mode = True + + gdn_submodules = get_experimental_attention_variant_module_spec( + config=det_config + ).submodules + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + gdn = ( + GatedDeltaNet( + det_config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + .cuda() + .bfloat16() + ) + + # deterministic_mode must select the torch-native kernel, not FLA. + assert gdn.gated_delta_rule is torch_chunk_gated_delta_rule + + micro_batch_size = 2 + seq_length = 64 + torch.manual_seed(0) + base_input = torch.randn( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + def run(): + hidden_states = base_input.clone().requires_grad_(True) + output, _ = gdn(hidden_states, None) + output.float().sum().backward() + grads = { + name: param.grad.detach().clone() + for name, param in gdn.named_parameters() + if param.grad is not None + } + gdn.zero_grad(set_to_none=True) + return output.detach().clone(), grads, hidden_states.grad.detach().clone() + + out1, grads1, input_grad1 = run() + out2, grads2, input_grad2 = run() + + rank = torch.distributed.get_rank() + assert torch.equal(out1, out2), f"Output not reproducible ({rank=})" + assert torch.equal(input_grad1, input_grad2), f"Input grad not reproducible ({rank=})" + assert set(grads1.keys()) == set(grads2.keys()) + for name in grads1: + assert torch.equal( + grads1[name], grads2[name] + ), f"Grad not reproducible for {name} ({rank=})" + + def test_module_construction(self): + gdn = self.gdn + assert gdn.in_proj_dim == 2 * gdn.qk_dim + 2 * gdn.v_dim + 2 * gdn.num_value_heads + assert gdn.A_log.shape == (gdn.num_value_heads // self.tp_size,) + assert gdn.dt_bias.shape == (gdn.num_value_heads // self.tp_size,) + + def test_sharded_state_dict_splits_gdn_parameters(self): + sharded_sd = self.gdn.sharded_state_dict(prefix="gdn.") + + in_proj_weight = sharded_sd["gdn.in_proj.weight"] + conv1d_weight = sharded_sd["gdn.conv1d.weight"] + assert isinstance(in_proj_weight, ShardedTensorFactory) + assert isinstance(conv1d_weight, ShardedTensorFactory) + + in_proj_chunks = in_proj_weight.build() + assert tuple(chunk.key for chunk in in_proj_chunks) == tuple( + f"gdn.in_proj.weight.{name}" for name in self.gdn.in_proj_split_names + ) + assert sum(chunk.data.numel() for chunk in in_proj_chunks) == in_proj_weight.data.numel() + + conv1d_chunks = conv1d_weight.build() + assert tuple(chunk.key for chunk in conv1d_chunks) == ( + "gdn.conv1d.weight.query", + "gdn.conv1d.weight.key", + "gdn.conv1d.weight.value", + ) + assert sum(chunk.data.numel() for chunk in conv1d_chunks) == conv1d_weight.data.numel() + + def test_jit_compiled_helpers(self): + import torch._dynamo + + gdn = self.gdn + batch = 2 + seq_len = 16 + + device = torch.cuda.current_device() + num_v_heads_local = gdn.num_value_heads // gdn.tp_size // self.cp_size_headwise + qk_dim_local = gdn.qk_dim_local_tp // self.cp_size_headwise + v_dim_local = gdn.v_dim_local_tp // self.cp_size_headwise + + qkv = torch.randn( + batch, seq_len, 2 * qk_dim_local + v_dim_local, device=device, dtype=torch.bfloat16 + ) + gate = torch.randn( + batch, + seq_len, + num_v_heads_local, + gdn.value_head_dim, + device=device, + dtype=torch.bfloat16, + ) + gate_feats = ( + torch.randn(batch, seq_len, num_v_heads_local, device=device, dtype=torch.bfloat16), + torch.randn(batch, seq_len, num_v_heads_local, device=device, dtype=torch.bfloat16), + ) # beta, alpha + + # Disable dynamo so coverage.py can trace through the method bodies, + # which are normally wrapped by @jit_fuser (torch.compile). + A_log_mock = torch.randn(num_v_heads_local, device=device, dtype=torch.bfloat16) + dt_bias_mock = torch.randn(num_v_heads_local, device=device, dtype=torch.bfloat16) + + with torch._dynamo.config.patch(disable=True): + kernel_inputs = gdn._prepare_input_for_gated_delta_rule( + qkv, + gate, + A_log_mock, + dt_bias_mock, + batch, + seq_len, + *gate_feats, + cp_size_headwise=self.cp_size_headwise, + ) + + query = kernel_inputs["q"] + key = kernel_inputs["k"] + value = kernel_inputs["v"] + g = kernel_inputs["g"] + gate_out = kernel_inputs["gate"] + beta_out = kernel_inputs["beta"] + + assert query.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) + assert key.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) + assert value.shape == (batch, seq_len, num_v_heads_local, gdn.value_head_dim) + for t in (query, key, value, gate_out, beta_out): + assert t.is_contiguous() + + assert g.dtype == torch.float32 + assert g.shape == (batch, seq_len, num_v_heads_local) + assert beta_out.shape == (batch, seq_len, num_v_heads_local) + + def test_fused_pre_gated_delta_rule_headwise_cp_uses_cp_local_parameters(self): + if not HAVE_FUSED_PRE_GDR: + pytest.skip("causal-conv1d fused backward is not installed.") + if not (self.linear_cp_mode == "headwise" and self.cp_size > 1): + pytest.skip("Only headwise CP with CP>1 needs CP-local fused pre-GDR params.") + + gdn = self.gdn + batch = 2 + seq_len = 16 + qk_channels = gdn.qk_dim_local_tp // self.cp_size_headwise + v_channels = gdn.v_dim_local_tp // self.cp_size_headwise + num_key_heads = qk_channels // gdn.key_head_dim + num_value_heads = v_channels // gdn.value_head_dim + qkvzba_dim = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + qkvzba = torch.randn( + seq_len, batch, qkvzba_dim, device=torch.cuda.current_device(), dtype=torch.bfloat16 + ) + captured = {} + + def fake_fused_streamed_pre_gated_delta_rule( + qkvzba_arg, + conv1d_weight, + conv1d_bias, + A_log, + dt_bias, + *, + num_key_heads, + num_value_heads, + **kwargs, + ): + captured.update( + { + "qkvzba": qkvzba_arg, + "conv1d_weight": conv1d_weight, + "conv1d_bias": conv1d_bias, + "A_log": A_log, + "dt_bias": dt_bias, + "num_key_heads": num_key_heads, + "num_value_heads": num_value_heads, + "cp_group": kwargs["cp_group"], + } + ) + return tuple(torch.empty(0, device=qkvzba_arg.device) for _ in range(6)) + + with mock.patch( + "megatron.core.fusions.fused_pre_gated_delta_rule." + "fused_streamed_pre_gated_delta_rule", + side_effect=fake_fused_streamed_pre_gated_delta_rule, + ): + gdn._fused_streamed_pre_gated_delta_rule(qkvzba, cp_group_headwise=gdn.pg_collection.cp) + + assert captured["qkvzba"] is qkvzba + assert captured["conv1d_weight"].shape == ( + 2 * qk_channels + v_channels, + 1, + gdn.conv_kernel_dim, + ) + assert captured["conv1d_bias"] is None + assert captured["A_log"].shape == (num_value_heads,) + assert captured["dt_bias"].shape == (num_value_heads,) + assert captured["num_key_heads"] == num_key_heads + assert captured["num_value_heads"] == num_value_heads + assert captured["cp_group"] is None + + def test_gpu_forward_thd_correctness(self): + if self.sp_size > 1: + pytest.skip("Sequence parallel is not supported for this test case.") + if self.cp_size > 1 and self.linear_cp_mode == "chunkwise": + pytest.skip("Chunkwise CP is not supported for this test case.") + + atol, rtol = 3e-4, 3e-4 + + # Input shape + sequence_length = 32 + micro_batch_size = 4 + cu_seqlens = [0, 32, 64, 96, 128] + # sbhd input shape: [sequence length, batch size, hidden size] + sub_sequence_length = sequence_length // self.cp_size + hidden_states_sbhd = torch.rand( + (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size) + ) + attention_mask_sbhd = None + hidden_states_sbhd = hidden_states_sbhd.cuda().bfloat16() + # thd input shape: [sequence length * batch size, 1, hidden size] + hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() + hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) + attention_mask_thd = None + packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) + _set_gdn_test_cp_partition_mode(packed_seq_params, self.cp_size, self.linear_cp_mode) + + # THD format + output_thd, _ = self.gdn( + hidden_states_thd, attention_mask_thd, packed_seq_params=packed_seq_params + ) + # SBHD format + output_sbhd, _ = self.gdn(hidden_states_sbhd, attention_mask_sbhd) + output_sbhd_T = output_sbhd.transpose(0, 1).contiguous().view(*output_thd.shape) + + rank = torch.distributed.get_rank() + assert output_thd.shape[0] == sub_sequence_length * micro_batch_size + assert output_thd.shape[1] == 1 + assert output_thd.shape[2] == self.gdn.config.hidden_size + torch.testing.assert_close( + output_sbhd_T, + output_thd, + atol=atol, + rtol=rtol, + msg=lambda msg: f"Output mismatch ({rank=}): {msg}", + ) + + def test_gpu_forward_thd_padding_correctness(self): + if self.sp_size > 1: + pytest.skip("Sequence parallel is not supported for this test case.") + if self.cp_size > 1 and self.linear_cp_mode == "chunkwise": + pytest.skip("Chunkwise CP is not supported for this test case.") + + atol, rtol = 3e-4, 3e-4 + sequence_length = 32 + micro_batch_size = 4 + + # sbhd input shape: [sequence length, batch size, hidden size] + sub_sequence_length = sequence_length // self.cp_size + hidden_states_sbhd = torch.rand( + (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + output_sbhd, _ = self.gdn(hidden_states_sbhd, None) + + # thd input shape: [sequence length * batch size, 1, hidden size] + hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() + hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) + output_bshd = output_sbhd.transpose(0, 1).contiguous() + + rank = torch.distributed.get_rank() + + # A) padded branch: prefer *_padded when available. + padded_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 128] + ) + _set_gdn_test_cp_partition_mode(padded_params, self.cp_size, self.linear_cp_mode) + output_thd_padded, _ = self.gdn(hidden_states_thd, None, packed_seq_params=padded_params) + output_thd2bshd = output_thd_padded.view(*output_bshd.shape) + torch.testing.assert_close( + output_bshd[:, :30, :], + output_thd2bshd[:, :30, :], + atol=atol, + rtol=rtol, + msg=lambda msg: f"THD padded output mismatch ({rank=}): {msg}", + ) + + # B) no-padded branch: use actual cu_seqlens when it matches total_sequence_length. + no_padding_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 128]) + _set_gdn_test_cp_partition_mode(no_padding_params, self.cp_size, self.linear_cp_mode) + output_thd_no_padding, _ = self.gdn( + hidden_states_thd, None, packed_seq_params=no_padding_params + ) + assert output_thd_no_padding.shape == output_thd_padded.shape + + # C) explicit causal-conv padding is only applied to packed inputs and + # should not affect the original unpadded token outputs. + self.gdn.config.gdn_conv_pad_alignment = 48 + output_thd_conv_pad, _ = self.gdn( + hidden_states_thd, None, packed_seq_params=no_padding_params + ) + self.gdn.config.gdn_conv_pad_alignment = None + assert output_thd_conv_pad.shape == output_thd_no_padding.shape + torch.testing.assert_close( + output_thd_conv_pad, + output_thd_no_padding, + atol=atol, + rtol=rtol, + msg=lambda msg: f"THD conv-padded output mismatch ({rank=}): {msg}", + ) + + # D) padded mismatch branch: if *_padded[-1] mismatches total_sequence_length, should raise. + padded_mismatch_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 126] + ) + _set_gdn_test_cp_partition_mode(padded_mismatch_params, self.cp_size, self.linear_cp_mode) + with pytest.raises(ValueError, match="does not match"): + self.gdn(hidden_states_thd, None, packed_seq_params=padded_mismatch_params) + + # E) actual mismatch branch without *_padded: should raise. + actual_mismatch_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 129]) + _set_gdn_test_cp_partition_mode(actual_mismatch_params, self.cp_size, self.linear_cp_mode) + with pytest.raises(ValueError, match="does not match"): + self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) + + +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +class TestGDNCuSeqlensResolve: + + @pytest.fixture + def mock_gdn(self): + class MockGDN: + _resolve_cu_seqlens = GatedDeltaNet._resolve_cu_seqlens + + return MockGDN() + + def test_padded_preferred_when_available(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + padded = torch.tensor([0, 504, 1008], dtype=torch.int32) + result = mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q", cp_size=2) + assert torch.equal(result, padded) + + def test_actual_used_when_no_padding(self, mock_gdn): + actual = torch.tensor([0, 504, 1008], dtype=torch.int32) + result = mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + assert torch.equal(result, actual) + + def test_raises_when_padding_mismatch(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_raises_when_padded_mismatches_total(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + padded = torch.tensor([0, 504, 1004], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_raises_when_not_divisible_by_cp_size(self, mock_gdn): + actual = torch.tensor([0, 505, 1008], dtype=torch.int32) + with pytest.raises(ValueError, match="must be divisible by cp_size"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_cp1_still_validates_total(self, mock_gdn): + mock_gdn.cp_size = 1 + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=1) diff --git a/tests/unit_tests/ssm/gated_delta_net/test_kda.py b/tests/unit_tests/ssm/gated_delta_net/test_kda.py new file mode 100644 index 00000000000..5427b19092b --- /dev/null +++ b/tests/unit_tests/ssm/gated_delta_net/test_kda.py @@ -0,0 +1,424 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import copy +import os +from dataclasses import replace + +import pytest +import torch +import torch.nn.functional as F + +from megatron.core import parallel_state +from megatron.core.models.common.embeddings.rope_utils import ( + get_pos_emb_on_this_cp_rank as get_tensor_on_this_cp_rank, +) +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gated_delta_net import HAVE_FLA_KDA, KimiDeltaAttention +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.module import Float16Module +from tests.unit_tests.test_utilities import Utils +from tests.unit_tests.transformer.test_attention import _test_parallel_attention_correctness +from tests.unit_tests.transformer.test_multi_latent_attention import ( + make_test_packed_seq_params, + make_test_packed_seq_params_with_padding, +) + +os.environ.update({"NCCL_NVLS_ENABLE": "0"}) + +_KDA_INPUT_GRAD_ATOL = 3.2e-2 +_KDA_INPUT_GRAD_RELATIVE_L2 = 5e-3 + + +def _make_config( + tp_size: int = 1, + cp_size: int = 1, + sequence_parallel: bool = False, + params_dtype: torch.dtype = torch.bfloat16, + linear_cp_mode: str = "headwise", +) -> TransformerConfig: + return TransformerConfig( + hidden_size=128, + num_layers=1, + num_attention_heads=4, + num_query_groups=4, + kv_channels=32, + linear_conv_kernel_dim=4, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=4, + linear_num_value_heads=4, + kda_safe_gate=True, + kda_lower_bound=-5.0, + normalization="RMSNorm", + activation_func=F.silu, + add_bias_linear=False, + bf16=params_dtype == torch.bfloat16, + params_dtype=params_dtype, + use_cpu_initialization=True, + tensor_model_parallel_size=tp_size, + context_parallel_size=cp_size, + sequence_parallel=sequence_parallel, + linear_cp_mode=linear_cp_mode, + experimental_attention_variant="kda", + is_hybrid_model=True, + transformer_impl="transformer_engine", + ) + + +def test_kda_norm_out_recompute_config_accepts_hybrid_model(): + config = replace( + _make_config(), + is_hybrid_model=True, + recompute_granularity="selective", + recompute_modules=["gdn_norm_out"], + ) + + assert config.recompute_modules == ["gdn_norm_out"] + + +def test_kda_full_recompute_config_accepts_hybrid_model(): + config = replace( + _make_config(), + is_hybrid_model=True, + recompute_granularity="selective", + recompute_modules=["gdn"], + ) + + assert config.recompute_modules == ["gdn"] + + +def _build_kda(config: TransformerConfig) -> KimiDeltaAttention: + spec = hybrid_stack_spec.submodules.kda_layer.submodules.self_attention + return KimiDeltaAttention( + config, + submodules=spec.submodules, + layer_number=1, + conv_init=1.0, + pg_collection=ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ), + ).cuda() + + +def test_kda_config_rejects_unsupported_head_layouts(): + with pytest.raises(ValueError, match="equal key and value head counts"): + replace(_make_config(), linear_num_value_heads=8) + + with pytest.raises(ValueError, match="equal key and value head dimensions"): + replace(_make_config(), linear_value_head_dim=64) + + with pytest.raises(ValueError, match="must be either 'headwise' or 'chunkwise'"): + _make_config(linear_cp_mode="tokenwise") + + with pytest.raises(AssertionError, match="linear_head_parallel_size"): + _make_config(cp_size=3, linear_cp_mode="headwise") + + chunkwise_config = _make_config(cp_size=3, linear_cp_mode="chunkwise") + assert chunkwise_config.linear_cp_mode == "chunkwise" + + +@pytest.mark.parametrize("lower_bound", [None, -5.1, 0.0, 1.0, float("nan")]) +def test_kda_rejects_invalid_safe_gate_lower_bound(lower_bound): + with pytest.raises(ValueError, match="kda_lower_bound"): + replace(_make_config(), kda_safe_gate=True, kda_lower_bound=lower_bound) + + +@pytest.mark.parametrize("lower_bound", [-5.0, -1.0, -1e-6]) +def test_kda_accepts_valid_safe_gate_lower_bound(lower_bound): + config = replace(_make_config(), kda_safe_gate=True, kda_lower_bound=lower_bound) + + assert config.kda_lower_bound == lower_bound + + +def test_kda_ignores_lower_bound_when_safe_gate_is_disabled(): + config = replace(_make_config(), kda_safe_gate=False, kda_lower_bound=None) + + assert config.kda_lower_bound is None + + +def test_kda_rejects_invalid_packed_boundaries(): + q = torch.tensor([0, 8, 16], dtype=torch.int32) + KimiDeltaAttention._validate_packed_cu_seqlens(q, q.clone()) + + with pytest.raises(ValueError, match="cu_seqlens_q to equal cu_seqlens_kv"): + KimiDeltaAttention._validate_packed_cu_seqlens( + q, torch.tensor([0, 7, 16], dtype=torch.int32) + ) + with pytest.raises(ValueError, match="at least one sequence in both Q and KV"): + empty = torch.tensor([0], dtype=torch.int32) + KimiDeltaAttention._validate_packed_cu_seqlens(empty, empty.clone()) + + +@pytest.mark.internal +@pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") +def test_kda_forward_backward(): + Utils.initialize_distributed() + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=1 + ) + try: + model_parallel_cuda_manual_seed(123) + config = _make_config() + kda = _build_kda(config) + assert kda.in_proj_split_names == ["query", "key", "value", "g", "gate"] + assert kda.A_log.dtype == torch.float32 + assert kda.dt_bias.dtype == torch.float32 + assert getattr(kda.A_log, "keep_in_fp32", False) + assert getattr(kda.dt_bias, "keep_in_fp32", False) + + hidden_states = torch.randn( + (32, 2, config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + requires_grad=True, + ) + output, bias = kda(hidden_states, None) + assert bias is None + assert output.shape == hidden_states.shape + assert output.dtype == hidden_states.dtype + output.float().square().mean().backward() + assert hidden_states.grad is not None + assert torch.isfinite(hidden_states.grad).all() + for name, parameter in kda.named_parameters(): + if parameter.grad is not None: + assert torch.isfinite(parameter.grad).all(), f"non-finite gradient in {name}" + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.internal +@pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") +def test_kda_gate_parameters_remain_fp32_after_model_cast(): + Utils.initialize_model_parallel(1, 1) + try: + model_parallel_cuda_manual_seed(123) + config = _make_config(params_dtype=torch.float32) + kda = _build_kda(config) + config.bf16 = True + wrapped = Float16Module(config=config, module=kda) + assert wrapped.module.in_proj.weight.dtype == torch.bfloat16 + assert wrapped.module.beta_proj.weight.dtype == torch.bfloat16 + assert wrapped.module.A_log.dtype == torch.float32 + assert wrapped.module.dt_bias.dtype == torch.float32 + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.internal +@pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") +def test_kda_packed_matches_unpacked_cp1(): + Utils.initialize_model_parallel(1, 1) + try: + model_parallel_cuda_manual_seed(321) + kda = _build_kda(_make_config()) + unpacked = torch.randn( + (16, 2, kda.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + packed = unpacked.transpose(0, 1).contiguous().reshape(-1, 1, kda.config.hidden_size) + packed_seq_params = make_test_packed_seq_params(cu_seqlens=[0, 16, 32]) + + unpacked_output, _ = kda(unpacked, None) + packed_output, _ = kda(packed, None, packed_seq_params=packed_seq_params) + torch.testing.assert_close( + packed_output[:, 0], + unpacked_output.transpose(0, 1).reshape(-1, kda.config.hidden_size), + atol=5e-3, + rtol=5e-3, + ) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.internal +@pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") +@pytest.mark.parametrize("recompute_module", ["gdn_norm_out", "gdn"]) +def test_kda_selective_recompute(recompute_module): + Utils.initialize_model_parallel(1, 1) + try: + + def run(module, hidden_states): + output, _ = module(hidden_states, None) + output.float().sum().backward() + grads = { + name: parameter.grad.detach().clone() + for name, parameter in module.named_parameters() + if parameter.grad is not None + } + return output.detach(), hidden_states.grad.detach().clone(), grads + + base_config = _make_config() + recompute_config = copy.deepcopy(base_config) + recompute_config.recompute_granularity = "selective" + recompute_config.recompute_modules = [recompute_module] + + model_parallel_cuda_manual_seed(42) + base_kda = _build_kda(base_config) + model_parallel_cuda_manual_seed(42) + recompute_kda = _build_kda(recompute_config) + recompute_kda.load_state_dict(base_kda.state_dict()) + + torch.manual_seed(42) + hidden_states = torch.randn( + (64, 2, base_config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + base_input = hidden_states.detach().clone().requires_grad_(True) + recompute_input = hidden_states.detach().clone().requires_grad_(True) + + base_output, base_input_grad, base_grads = run(base_kda, base_input) + recompute_output, recompute_input_grad, recompute_grads = run( + recompute_kda, recompute_input + ) + + assert base_kda.recompute_norm_out is False + assert base_kda.recompute_gdn is False + assert base_kda.norm_out_checkpoint is None + assert recompute_kda.recompute_norm_out is (recompute_module == "gdn_norm_out") + assert recompute_kda.recompute_gdn is (recompute_module == "gdn") + assert (recompute_kda.norm_out_checkpoint is not None) is ( + recompute_module == "gdn_norm_out" + ) + torch.testing.assert_close(recompute_output, base_output, atol=0, rtol=0) + torch.testing.assert_close(recompute_input_grad, base_input_grad, atol=0, rtol=0) + assert recompute_grads.keys() == base_grads.keys() + for name in base_grads: + # Recomputed BF16 kernels can choose a different reduction order. Keep output and + # input-gradient checks bitwise while allowing only BF16-scale parameter-gradient + # rounding. + torch.testing.assert_close( + recompute_grads[name], + base_grads[name], + atol=2e-4, + rtol=5e-3, + msg=f"Selective recompute gradient mismatch for {name}", + ) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.internal +@pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") +def test_kda_cp2_packed_with_physical_padding_matches_cp1(): + """Match CP2 packed KDA against CP1 when physical and actual boundaries differ.""" + + sequence_length = 32 + micro_batch_size = 4 + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=1 + ) + try: + packed_seq_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 128] + ) + torch.manual_seed(2026) + hidden_states_sbhd = torch.rand( + (sequence_length, micro_batch_size, 128), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + model_parallel_cuda_manual_seed(2026) + reference_kda = _build_kda(_make_config(cp_size=1)) + reference_state = { + name: tensor.detach().clone() for name, tensor in reference_kda.state_dict().items() + } + reference_input = ( + hidden_states_sbhd.transpose(0, 1) + .contiguous() + .view(-1, 1, hidden_states_sbhd.shape[-1]) + .requires_grad_(True) + ) + reference_output, reference_bias = reference_kda( + reference_input, None, packed_seq_params=packed_seq_params + ) + assert reference_bias is None + reference_output.float().sum().backward() + reference_output_sbhd = ( + reference_output.detach() + .view(micro_batch_size, sequence_length, -1) + .transpose(0, 1) + .contiguous() + ) + reference_input_grad_sbhd = ( + reference_input.grad.detach() + .view(micro_batch_size, sequence_length, -1) + .transpose(0, 1) + .contiguous() + ) + finally: + Utils.destroy_model_parallel() + + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=2 + ) + try: + model_parallel_cuda_manual_seed(2026) + parallel_kda = _build_kda(_make_config(cp_size=2)) + parallel_kda.load_state_dict(reference_state) + cp_group = parallel_state.get_context_parallel_group() + + def get_packed_tensor_on_this_cp_rank(tensor): + local_tensor = get_tensor_on_this_cp_rank(tensor, 0, cp_group) + return local_tensor.transpose(0, 1).contiguous().view(-1, 1, tensor.shape[-1]) + + parallel_input = get_packed_tensor_on_this_cp_rank(hidden_states_sbhd) + parallel_input = parallel_input.detach().requires_grad_(True) + parallel_output, parallel_bias = parallel_kda( + parallel_input, None, packed_seq_params=packed_seq_params + ) + assert parallel_bias is None + parallel_output.float().sum().backward() + + expected_output = get_packed_tensor_on_this_cp_rank(reference_output_sbhd) + expected_input_grad = get_packed_tensor_on_this_cp_rank(reference_input_grad_sbhd) + torch.testing.assert_close(parallel_output, expected_output, atol=1e-2, rtol=1e-2) + grad_diff = (parallel_input.grad - expected_input_grad).float() + relative_l2 = torch.linalg.vector_norm(grad_diff) / torch.linalg.vector_norm( + expected_input_grad.float() + ).clamp_min(1e-12) + assert relative_l2 < _KDA_INPUT_GRAD_RELATIVE_L2 + # The fixed-seed CP comparison differs by at most one observed BF16 bin + # (0.03125) while aggregate relative L2 remains below 0.5%. + torch.testing.assert_close( + parallel_input.grad, expected_input_grad, atol=_KDA_INPUT_GRAD_ATOL, rtol=1e-2 + ) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.parametrize( + ("sequence_packing", "tp", "sp", "cp", "linear_cp_mode"), + [ + (False, 2, True, 1, "headwise"), + (True, 1, False, 2, "headwise"), + (True, 2, True, 2, "headwise"), + (True, 1, False, 2, "chunkwise"), + (True, 2, True, 2, "chunkwise"), + ], +) +@pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") +def test_parallel_kda_correctness(tmp_path_dist_ckpt, sequence_packing, tp, sp, cp, linear_cp_mode): + config = _make_config(linear_cp_mode=linear_cp_mode) + _test_parallel_attention_correctness( + transformer_config=config, + transformer_layer_spec=hybrid_stack_spec.submodules.kda_layer, + tmp_path_dist_ckpt=tmp_path_dist_ckpt, + atol=1e-2, + rtol=1e-2, + # Parallel projection/reduction order can move individual BF16 values by + # one observed bin; the dedicated padded-CP test above also bounds relative L2. + input_grad_atol=_KDA_INPUT_GRAD_ATOL, + input_grad_rtol=1e-2, + tp=tp, + sp=sp, + cp=cp, + seed=42, + sequence_length=128, + micro_batch_size=2, + sequence_packing=sequence_packing, + ) diff --git a/tests/unit_tests/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index 322f1d78911..7db5eb2170e 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -1,17 +1,13 @@ -# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import pytest import torch -from megatron.core.extensions.transformer_engine import TEDotProductAttention -from megatron.core.models.hybrid.hybrid_block import HybridStack +from megatron.core.models.hybrid.hybrid_block import HybridStack, HyperConnectionHybridLayer from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols, validate_segment_layers -from megatron.core.models.hybrid.hybrid_layer_specs import ( - hybrid_inference_stack_spec, - hybrid_stack_spec, -) +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.gated_delta_net import GatedDeltaNet +from megatron.core.ssm.gated_delta_net import HAVE_FLA_KDA, GatedDeltaNet, KimiDeltaAttention from megatron.core.ssm.mamba_layer import MambaLayer from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig @@ -21,7 +17,6 @@ ) from megatron.core.transformer.experimental_attention_variant.dsa import DSAttention from megatron.core.transformer.mlp import MLP -from megatron.core.transformer.multi_latent_attention import MLASelfAttention from megatron.core.transformer.transformer_config import MLATransformerConfig from megatron.core.transformer.transformer_layer import TransformerLayer from tests.unit_tests.test_utilities import Utils @@ -37,8 +32,13 @@ def setup_method(self, method): def get_pg_collection(self): return ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'pp', 'cp']) - def get_hybrid_block(self, layer_pattern, **config_kwargs): + def get_mamba_block(self, layer_pattern, enable_hyper_connections=False): layer_type_list = validate_segment_layers(layer_pattern) + mhc_kwargs = ( + {"enable_hyper_connections": True, "hidden_dropout": 0.0, "mhc_sinkhorn_iterations": 5} + if enable_hyper_connections + else {} + ) transformer_config = TransformerConfig( hidden_size=256, # The Mamba layer places several constraints on this # Need to specify num_attention_heads and num_layers or TransformerConfig @@ -46,7 +46,7 @@ def get_hybrid_block(self, layer_pattern, **config_kwargs): num_layers=len(layer_type_list), num_attention_heads=4, use_cpu_initialization=True, - **config_kwargs, + **mhc_kwargs, ) modules = hybrid_stack_spec.submodules return HybridStack( @@ -57,29 +57,16 @@ def get_hybrid_block(self, layer_pattern, **config_kwargs): pg_collection=self.get_pg_collection(), ) - def get_dsa_hybrid_block(self, layer_pattern): + def get_hybrid_block(self, layer_pattern, **config_kwargs): layer_type_list = validate_segment_layers(layer_pattern) - transformer_config = MLATransformerConfig( + transformer_config = TransformerConfig( hidden_size=256, # The Mamba layer places several constraints on this # Need to specify num_attention_heads and num_layers or TransformerConfig # will generate errors. num_layers=len(layer_type_list), - num_attention_heads=16, + num_attention_heads=4, use_cpu_initialization=True, - bf16=True, - params_dtype=torch.bfloat16, - q_lora_rank=64, - kv_lora_rank=64, - qk_head_dim=64, - qk_pos_emb_head_dim=32, - v_head_dim=64, - rope_type='rope', - rotary_base=10000, - rotary_percent=1.0, - dsa_indexer_n_heads=8, - dsa_indexer_head_dim=64, - dsa_indexer_topk=32, - add_bias_linear=False, + **config_kwargs, ) modules = hybrid_stack_spec.submodules return HybridStack( @@ -90,8 +77,13 @@ def get_dsa_hybrid_block(self, layer_pattern): pg_collection=self.get_pg_collection(), ) - def get_mla_hybrid_block(self, layer_pattern): + def get_dsa_mamba_block(self, layer_pattern, enable_hyper_connections=False): layer_type_list = validate_segment_layers(layer_pattern) + mhc_kwargs = ( + {"enable_hyper_connections": True, "hidden_dropout": 0.0, "mhc_sinkhorn_iterations": 5} + if enable_hyper_connections + else {} + ) transformer_config = MLATransformerConfig( hidden_size=256, # The Mamba layer places several constraints on this # Need to specify num_attention_heads and num_layers or TransformerConfig @@ -109,6 +101,11 @@ def get_mla_hybrid_block(self, layer_pattern): rope_type='rope', rotary_base=10000, rotary_percent=1.0, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=32, + add_bias_linear=False, + **mhc_kwargs, ) modules = hybrid_stack_spec.submodules return HybridStack( @@ -246,6 +243,161 @@ def test_layer_types(self): assert isinstance(layers[2], TransformerLayer) assert isinstance(layers[2].mlp, MLP) + def test_hyper_connection_layer_wrappers(self): + """mHC wraps each hybrid layer while preserving the layer type underneath.""" + layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP + block = self.get_mamba_block(layer_pattern, enable_hyper_connections=True) + layers = block.layers + assert all(isinstance(layer, HyperConnectionHybridLayer) for layer in layers) + assert isinstance(layers[0].inner_layer, MambaLayer) + assert isinstance(layers[1].inner_layer, TransformerLayer) + assert isinstance(layers[1].inner_layer.self_attention, SelfAttention) + assert isinstance(layers[2].inner_layer, TransformerLayer) + assert isinstance(layers[2].inner_layer.mlp, MLP) + + def test_hyper_connection_recompute_plan_for_hybrid_layers(self): + """HybridStack creates per-layer mHC recompute managers when requested.""" + layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP + layer_type_list = validate_segment_layers(layer_pattern) + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=len(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + hidden_dropout=0.0, + mhc_sinkhorn_iterations=5, + recompute_granularity="selective", + recompute_modules=["core_attn", "mhc"], + ) + block = HybridStack( + transformer_config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + pg_collection=self.get_pg_collection(), + ) + + managers, block_ends = block._build_mhc_recompute_layer_plan(use_mhc_recompute=True) + assert len(managers) == len(block.layers) + assert all(manager is not None for manager in managers) + assert block_ends[-1] is True + + def test_hyper_connection_gpu_forward(self): + """mHC-enabled HybridStack expands internally and contracts back at the output.""" + layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP + block = self.get_mamba_block(layer_pattern, enable_hyper_connections=True) + block.cuda() + micro_batch_size = 2 + sequence_length = 32 + hidden_states = torch.ones((sequence_length, micro_batch_size, block.config.hidden_size)) + hidden_states = hidden_states.cuda() + attention_mask = torch.ones( + (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool + ) + attention_mask = attention_mask.cuda() + output = block(hidden_states, attention_mask=attention_mask) + assert output.shape[0] == sequence_length + assert output.shape[1] == micro_batch_size + assert output.shape[2] == block.config.hidden_size + assert output.dtype == torch.float32 + + def test_hyper_connection_gdn_gpu_forward(self): + """mHC runs through GDN, attention, and Mamba hybrid layers.""" + layer_pattern = Symbols.GDN + Symbols.ATTENTION + Symbols.MAMBA + layer_type_list = validate_segment_layers(layer_pattern) + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=len(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + activation_func=torch.nn.functional.silu, + enable_hyper_connections=True, + hidden_dropout=0.0, + mhc_sinkhorn_iterations=5, + ) + block = HybridStack( + transformer_config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + pg_collection=self.get_pg_collection(), + ) + block.cuda() + micro_batch_size = 2 + sequence_length = 32 + hidden_states = torch.ones((sequence_length, micro_batch_size, block.config.hidden_size)) + hidden_states = hidden_states.cuda() + attention_mask = torch.ones( + (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool + ).cuda() + output = block(hidden_states, attention_mask=attention_mask) + assert output.shape == (sequence_length, micro_batch_size, block.config.hidden_size) + + def test_hyper_connection_dsa_layer_wrappers(self): + """mHC wraps DeepSeek-style DSA and MLP split layers.""" + layer_pattern = Symbols.MAMBA + Symbols.DS_ATTENTION + Symbols.MLP + block = self.get_dsa_mamba_block(layer_pattern, enable_hyper_connections=True) + layers = block.layers + assert all(isinstance(layer, HyperConnectionHybridLayer) for layer in layers) + assert isinstance(layers[0].inner_layer, MambaLayer) + assert isinstance(layers[1].inner_layer, TransformerLayer) + assert isinstance(layers[1].inner_layer.self_attention, AbsorbedMLASelfAttention) + assert isinstance(layers[1].inner_layer.self_attention.core_attention, DSAttention) + assert isinstance(layers[2].inner_layer, TransformerLayer) + assert isinstance(layers[2].inner_layer.mlp, MLP) + + def test_hyper_connection_pipeline_boundary_shapes(self): + """HybridStack keeps n-stream tensors between PP stages and contracts at the end.""" + layer_type_list = validate_segment_layers(Symbols.MAMBA) + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=len(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + hidden_dropout=0.0, + mhc_sinkhorn_iterations=5, + ) + modules = hybrid_stack_spec.submodules + first_stage = HybridStack( + transformer_config, + modules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + post_process=False, + pg_collection=self.get_pg_collection(), + ).cuda() + last_stage = HybridStack( + transformer_config, + modules, + pre_process=False, + layer_type_list=layer_type_list, + pp_layer_offset=1, + post_process=True, + pg_collection=self.get_pg_collection(), + ).cuda() + + micro_batch_size = 2 + sequence_length = 32 + hidden_states = torch.ones( + (sequence_length, micro_batch_size, transformer_config.hidden_size), device='cuda' + ) + attention_mask = torch.ones( + (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool, device='cuda' + ) + + pp_hidden = first_stage(hidden_states, attention_mask=attention_mask) + assert pp_hidden.shape == ( + sequence_length, + micro_batch_size, + transformer_config.hidden_size * transformer_config.num_residual_streams, + ) + + last_stage.set_input_tensor(pp_hidden.detach()) + output = last_stage(hidden_states, attention_mask=attention_mask) + assert output.shape == (sequence_length, micro_batch_size, transformer_config.hidden_size) + def test_invalid_layer_types_cause_failure(self): invalid_symbol = 'X' assert invalid_symbol not in Symbols.VALID_LAYERS # sanity check. @@ -268,11 +420,20 @@ def test_gdn_layer_types(self): assert isinstance(layers[1].self_attention, SelfAttention) assert isinstance(layers[2], MambaLayer) - def test_gdn_inference_spec(self): - """The inference stack must materialize GDN rather than its IdentityOp default.""" - gdn_spec = hybrid_inference_stack_spec.submodules.gdn_layer - assert gdn_spec.module is TransformerLayer - assert gdn_spec.submodules.self_attention.module is GatedDeltaNet + @pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") + def test_kda_layer_type(self): + """K builds a TransformerLayer wrapping KimiDeltaAttention.""" + block = self.get_hybrid_block( + Symbols.KDA, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=4, + activation_func=torch.nn.functional.silu, + add_bias_linear=False, + ) + assert isinstance(block.layers[0], TransformerLayer) + assert isinstance(block.layers[0].self_attention, KimiDeltaAttention) def test_gdn_gpu_forward(self): """Test GPU forward pass with GDN, attention, and Mamba layers.""" @@ -311,7 +472,7 @@ def test_gdn_gpu_forward(self): def test_dsa_layer_types(self): """D symbol creates a TransformerLayer with absorbed MLA and DSA core attention.""" layer_pattern = Symbols.MAMBA + Symbols.DS_ATTENTION + Symbols.MAMBA - block = self.get_dsa_hybrid_block(layer_pattern) + block = self.get_dsa_mamba_block(layer_pattern) layers = block.layers assert isinstance(layers[0], MambaLayer) assert isinstance(layers[1], TransformerLayer) @@ -323,22 +484,4 @@ def test_mixed_attention_and_dsa_layer_types(self): """* and D in the same block fail.""" layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.DS_ATTENTION + Symbols.MAMBA with pytest.raises(ValueError): - block = self.get_dsa_hybrid_block(layer_pattern) - - def test_mla_layer_types(self): - """+ symbol creates a TransformerLayer with MLASelfAttention but - standard (non-DSA) core attention.""" - layer_pattern = Symbols.MAMBA + Symbols.MLA + Symbols.MAMBA - block = self.get_mla_hybrid_block(layer_pattern) - layers = block.layers - assert isinstance(layers[0], MambaLayer) - assert isinstance(layers[1], TransformerLayer) - assert isinstance(layers[1].self_attention, MLASelfAttention) - assert isinstance(layers[1].self_attention.core_attention, TEDotProductAttention) - assert isinstance(layers[2], MambaLayer) - - def test_mixed_attention_and_mla_layer_types(self): - """* and + in the same block fail (same reason as * and D).""" - layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLA + Symbols.MAMBA - with pytest.raises(ValueError): - block = self.get_mla_hybrid_block(layer_pattern) + block = self.get_dsa_mamba_block(layer_pattern) diff --git a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py index 8b4c181ee30..e8383752e12 100644 --- a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py +++ b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py @@ -78,6 +78,7 @@ def test_valid_patterns(self): ("GGG*GGG*", ['G', 'G', 'G', '*', 'G', 'G', 'G', '*']), ("GEGEGE*E", ['G', 'E', 'G', 'E', 'G', 'E', '*', 'E']), ("MDMD", ['M', 'D', 'M', 'D']), + ("K+K", ['K', '+', 'K']), ("M+M+", ['M', '+', 'M', '+']), ] for pattern, expected in test_cases: @@ -102,11 +103,20 @@ def test_invalid_symbols_cause_failure(self): with pytest.raises(ValueError): # Not allowed to have both standard Attention and MLA/DSA validate_segment_layers("MDM*-") + + def test_window_symbol(self): + """'W' (sliding-window-only DSv4 attention) is a first-class MLA layer symbol.""" + assert Symbols.WINDOW == "W" + assert Symbols.WINDOW in Symbols.VALID_LAYERS + assert Symbols.WINDOW in Symbols.MLA_ATTENTION + # Parses and is counted like any other layer. + assert validate_segment_layers("WEWE") == ["W", "E", "W", "E"] + assert get_hybrid_layer_counts("WEWE")["W"] == 2 + # 'W' coexists with the other MLA-family attentions (C/H/D)... + validate_segment_layers("WEDECEHE") + # ...but not with standard '*' attention. with pytest.raises(ValueError): - # Not allowed to have both standard Attention and MLA (same reason - # as DSA: * uses the model-level rotary_pos_emb while + uses MLA's - # own decoupled RoPE). - validate_segment_layers("M+M*-") + validate_segment_layers("MWM*-") @pytest.mark.internal @@ -169,8 +179,6 @@ def test_main_pattern_only(self): ("GEGEGE*E", "GEGEGE*E"), ("MDMD", "MDMD"), ("DM", "DM"), - ("M+M+", "M+M+"), - ("+M", "+M"), ] for pattern, expected_main in test_cases: result = parse_hybrid_pattern(pattern) @@ -295,8 +303,6 @@ def test_complex_patterns(self): ("GEGEGE*E/GG/GG", "GEGEGE*E", "GG", 2), # DSA in main pattern with MTP ("MDMD/MD/MD", "MDMD", "MD", 2), - # MLA in main pattern with MTP - ("M+M+/M+/M+", "M+M+", "M+", 2), ] for pattern, expected_main, expected_mtp, expected_depths in test_cases: result = parse_hybrid_pattern(pattern) @@ -316,11 +322,15 @@ class TestGetHybridLayerCounts: def test_simple_pattern(self): assert get_hybrid_layer_counts("M*M*") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 2, 'D': 0, 'G': 0, - 'M': 2, + 'K': 0, '+': 0, + 'M': 2, '-': 0, 'E': 0, } @@ -328,50 +338,57 @@ def test_simple_pattern(self): def test_all_layer_types(self): # Not allowed to have both standard Attention and MLA/DSA, so we do separate asserts. assert get_hybrid_layer_counts("MG*-E") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 1, 'D': 0, 'G': 1, - 'M': 1, + 'K': 0, '+': 0, + 'M': 1, '-': 1, 'E': 1, } assert get_hybrid_layer_counts("MGD-E") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 1, 'G': 1, + 'K': 0, 'M': 1, '+': 0, '-': 1, 'E': 1, } - assert get_hybrid_layer_counts("MG+-E") == { - '*': 0, - 'D': 0, - 'G': 1, - 'M': 1, - '+': 1, - '-': 1, - 'E': 1, - } def test_with_pipes(self): # Pipes should be skipped in counting assert get_hybrid_layer_counts("M*|M*") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 2, 'D': 0, 'G': 0, - 'M': 2, + 'K': 0, '+': 0, + 'M': 2, '-': 0, 'E': 0, } assert get_hybrid_layer_counts("M-M-|M-M*-") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 1, 'D': 0, 'G': 0, - 'M': 4, + 'K': 0, '+': 0, + 'M': 4, '-': 4, 'E': 0, } @@ -379,11 +396,15 @@ def test_with_pipes(self): def test_with_mtp(self): # MTP pattern "MM" repeated 2 depths -> 4 extra mamba layers assert get_hybrid_layer_counts("M*M*/MM/MM") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 2, 'D': 0, 'G': 0, - 'M': 6, + 'K': 0, '+': 0, + 'M': 6, '-': 0, 'E': 0, } @@ -392,22 +413,30 @@ def test_with_pipes_and_mtp(self): # Main: M-M-|M-M*- -> 1 attn, 4 mamba, 4 mlp # MTP: MM x 2 depths -> +4 mamba assert get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 1, 'D': 0, 'G': 0, - 'M': 8, + 'K': 0, '+': 0, + 'M': 8, '-': 4, 'E': 0, } def test_moe_pattern(self): assert get_hybrid_layer_counts("MEME") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 0, 'G': 0, - 'M': 2, + 'K': 0, '+': 0, + 'M': 2, '-': 0, 'E': 2, } @@ -415,22 +444,30 @@ def test_moe_pattern(self): def test_mtp_with_attention(self): # MTP pattern "*M" repeated 3 depths -> 3 attn + 3 mamba from MTP assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 3, 'D': 0, 'G': 0, - 'M': 7, + 'K': 0, '+': 0, + 'M': 7, '-': 0, 'E': 0, } def test_gdn_pattern(self): assert get_hybrid_layer_counts("GMGM") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 0, 'G': 2, - 'M': 2, + 'K': 0, '+': 0, + 'M': 2, '-': 0, 'E': 0, } @@ -438,20 +475,28 @@ def test_gdn_pattern(self): def test_gdn_hybrid_pattern(self): # GDN + Mamba + Attention assert get_hybrid_layer_counts("G*GM*") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 2, 'D': 0, 'G': 2, - 'M': 1, + 'K': 0, '+': 0, + 'M': 1, '-': 0, 'E': 0, } def test_dsa_pattern(self): assert get_hybrid_layer_counts("DMDM") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 2, 'G': 0, + 'K': 0, 'M': 2, '+': 0, '-': 0, @@ -460,22 +505,45 @@ def test_dsa_pattern(self): def test_mla_pattern(self): assert get_hybrid_layer_counts("+M+M") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 0, 'G': 0, + 'K': 0, 'M': 2, '+': 2, '-': 0, 'E': 0, } + def test_kda_pattern(self): + assert get_hybrid_layer_counts("K+K/K") == { + 'C': 0, + 'H': 0, + 'W': 0, + '*': 0, + 'D': 0, + 'G': 0, + 'K': 3, + 'M': 0, + '+': 1, + '-': 0, + 'E': 0, + } + def test_empty_pattern(self): assert get_hybrid_layer_counts("") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 0, 'G': 0, - 'M': 0, + 'K': 0, '+': 0, + 'M': 0, '-': 0, 'E': 0, } @@ -761,7 +829,7 @@ def test_standard_layer_types(self): """Standard symbols each produce a single-entry map at local index 0.""" maps = get_layer_maps_from_layer_type_list(["*", "M", "-", "E"]) # We always get all symbols returned, not only those contained in the pattern. - assert len(maps) == 7 + assert len(maps) == len(Symbols.VALID_LAYERS) attention_map, mamba_map, mlp_map, moe_map = operator.itemgetter( Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE )(maps) @@ -769,6 +837,13 @@ def test_standard_layer_types(self): assert mamba_map == {1: 0} assert mlp_map == {2: 0} assert moe_map == {3: 0} + assert maps[Symbols.KDA] == {} + + def test_kda(self): + """KDA layers are tracked independently from MLA and standard attention.""" + maps = get_layer_maps_from_layer_type_list(["K", "+", "K"]) + assert maps[Symbols.KDA] == {0: 0, 2: 1} + assert maps[Symbols.MLA] == {1: 0} def test_dsa(self): """D (DSA) layers are treated as separate layers for KV cache mapping.""" @@ -804,39 +879,3 @@ def test_all_mamba(self): assert mamba_map == {0: 0, 1: 1, 2: 2} assert mlp_map == {} assert moe_map == {} - - def test_mla(self): - """+ (MLA) layers are mapped independently of other attention types.""" - maps = get_layer_maps_from_layer_type_list(["+", "M", "+", "M"]) - attention_map, dsa_map, mamba_map, mla_map, mlp_map, moe_map = operator.itemgetter( - Symbols.ATTENTION, - Symbols.DS_ATTENTION, - Symbols.MAMBA, - Symbols.MLA, - Symbols.MLP, - Symbols.MOE, - )(maps) - assert attention_map == {} - assert dsa_map == {} - assert mla_map == {0: 0, 2: 1} - assert mamba_map == {1: 0, 3: 1} - assert mlp_map == {} - assert moe_map == {} - - def test_mixed_dsa_and_mla(self): - """D and + can coexist (both are MLA-based and use decoupled RoPE).""" - maps = get_layer_maps_from_layer_type_list(["D", "+", "M", "-"]) - attention_map, dsa_map, mamba_map, mla_map, mlp_map, moe_map = operator.itemgetter( - Symbols.ATTENTION, - Symbols.DS_ATTENTION, - Symbols.MAMBA, - Symbols.MLA, - Symbols.MLP, - Symbols.MOE, - )(maps) - assert attention_map == {} - assert dsa_map == {0: 0} - assert mla_map == {1: 0} - assert mamba_map == {2: 0} - assert mlp_map == {3: 0} - assert moe_map == {} diff --git a/tests/unit_tests/test_argument_utils.py b/tests/unit_tests/test_argument_utils.py index 8ff6719af72..b6492874abb 100644 --- a/tests/unit_tests/test_argument_utils.py +++ b/tests/unit_tests/test_argument_utils.py @@ -1,24 +1,23 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import signal -import sys from argparse import ArgumentError, ArgumentParser, Namespace from dataclasses import dataclass, field from typing import Callable, Literal, Optional, Union from unittest.mock import MagicMock, patch import pytest -import torch from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig from megatron.core.optimizer import OptimizerConfig +from megatron.core.transformer.spec_utils import ModuleSpec from megatron.training.argument_utils import ( ArgumentGroupFactory, TypeInferenceError, - core_transformer_config_from_args, + _normalize_dsv4_hybrid_csa_compress_ratios, + hybrid_config_from_args, pretrain_cfg_container_from_args, ) -from megatron.training.arguments import add_megatron_arguments, parse_args, validate_args from megatron.training.config import PretrainConfigContainer @@ -45,57 +44,6 @@ class DummyConfig: """Setting with enum type to test enum handling""" -@dataclass(init=False) -class CapturingTransformerConfig: - """Minimal config that records kwargs produced by core_transformer_config_from_args.""" - - moe_use_norm_before_up_proj: bool = False - - def __init__(self, **kwargs): - self.__dict__.update(kwargs) - - -def test_moe_norm_flag_reaches_transformer_config(): - """The generated LatentMoE norm flag should populate the model config.""" - parser = ArgumentParser() - add_megatron_arguments(parser) - - default_args = parser.parse_args([]) - assert default_args.moe_use_norm_before_up_proj is False - - enabled_args = parser.parse_args(["--moe-use-norm-before-up-proj"]) - # params_dtype has no CLI flag of its own; validate_args normally derives it. - enabled_args.params_dtype = torch.float32 - config = core_transformer_config_from_args( - enabled_args, config_class=CapturingTransformerConfig - ) - - assert config.moe_use_norm_before_up_proj is True - - -def test_moe_norm_flag_requires_latent_size(monkeypatch): - """validate_args should reject the LatentMoE norm flag without a latent size.""" - monkeypatch.setattr(sys, 'argv', ['test_argument_utils.py']) - args = parse_args() - args.num_layers = 2 - args.hidden_size = 128 - args.num_attention_heads = 4 - args.max_position_embeddings = 1024 - args.seq_length = 1024 - args.micro_batch_size = 1 - # Let validate_args derive a global batch size that is valid for the - # active data-parallel size in distributed unit-test jobs. - args.train_iters = 1 - args.lr = 1e-4 - args.tokenizer_type = 'NullTokenizer' - args.vocab_size = 1024 - args.moe_use_norm_before_up_proj = True - args.moe_latent_size = None - - with pytest.raises(AssertionError, match="--moe-use-norm-before-up-proj requires"): - validate_args(args) - - @dataclass class ConfigWithOptional: """Config with optional fields.""" @@ -709,12 +657,6 @@ def test_handled_unsupported_unions(self): class TestMegatronNetworkArgumentGeneration: """Test Megatron's TransformerConfig-derived argument group.""" - @staticmethod - def _parser() -> ArgumentParser: - from megatron.training.arguments import _add_network_size_args - - return _add_network_size_args(ArgumentParser(exit_on_error=False)) - def test_transformer_callback_fields_are_not_registered_as_cli_args(self): """Callback fields are runtime hooks, not CLI-provided values.""" from megatron.training.arguments import _add_network_size_args @@ -738,39 +680,112 @@ def test_transformer_callback_fields_are_not_registered_as_cli_args(self): for field_name in callback_fields: assert not hasattr(args, field_name) - def test_mhc_fused_backend_is_exposed_as_config_choice(self): - assert self._parser().parse_args([]).mhc_fused_backend == "auto" - - args = self._parser().parse_args(["--mhc-fused-backend", "native"]) + def test_hybrid_layer_pattern_help_lists_every_public_symbol(self): + """The public help must describe every symbol accepted by the pattern parser.""" + from megatron.training.arguments import _add_experimental_args + + parser = _add_experimental_args(ArgumentParser()) + action = next(action for action in parser._actions if action.dest == "hybrid_layer_pattern") + for symbol_and_name in ( + "M (mamba)", + "G (gdn)", + "K (kda)", + "* (attention)", + "D (dsa)", + "+ (mla)", + "- (mlp)", + "E (moe)", + ): + assert symbol_and_name in action.help + + +class TestDsv4HybridCsaCompressRatioNormalization: + """Test the shared DSv4 hybrid compression-ratio normalization contract.""" + + @pytest.mark.parametrize( + ("provided", "expected_config_ratios"), + [ + (None, [0, 0, 0, 4, 128, 0]), + ([0, 4, 128], [0, 0, 0, 4, 128, 0]), + ([0, 0, 0, 4, 128, 0], [0, 0, 0, 4, 128, 0]), + ], + ) + def test_normalizes_default_compact_and_padded_ratios(self, provided, expected_config_ratios): + args = Namespace(experimental_attention_variant='dsv4_hybrid', csa_compress_ratios=provided) + kw_args = {} + + _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, "-W|EC/H-") + + assert args.csa_compress_ratios == [0, 4, 128] + assert kw_args['csa_compress_ratios'] == expected_config_ratios + + @pytest.mark.parametrize( + ("provided", "message"), + [ + ([0, 8, 128], "ratio 8.*symbol 'C'.*expected 4"), + ([1, 0, 0, 4, 128, 0], "non-attention hybrid symbol '-'.*non-zero ratio 1"), + ([0, 4], r"length \(2\).*W/C/H attention symbols \(3\)"), + ], + ) + def test_rejects_invalid_ratios(self, provided, message): + args = Namespace(experimental_attention_variant='dsv4_hybrid', csa_compress_ratios=provided) - assert args.mhc_fused_backend == "native" + with pytest.raises(AssertionError, match=message): + _normalize_dsv4_hybrid_csa_compress_ratios(args, {}, "-W|EC/H-") - def test_mhc_fused_backend_rejects_unknown_choice(self): - with pytest.raises(ArgumentError, match="invalid choice"): - self._parser().parse_args(["--mhc-fused-backend", "cuda"]) +class TestHybridConfigFromArgs: + """Test static and config-aware hybrid stack spec resolution.""" -class TestMegatronMixedPrecisionArguments: - """Test language-model logit dtype CLI choices.""" + @staticmethod + def _args(): + return Namespace( + spec=["test_module", "test_spec"], + fp16_lm_cross_entropy=False, + hybrid_layer_pattern="M", + position_embedding_type="none", + rotary_percent=1.0, + rotary_base=10000, + make_vocab_size_divisible_by=128, + rotary_seq_len_interpolation_factor=None, + max_position_embeddings=1024, + untie_embeddings_and_output_weights=False, + padded_vocab_size=128, + ) @staticmethod - def _parser() -> ArgumentParser: - from megatron.training.arguments import _add_mixed_precision_args + def _transformer_config(): + config = MagicMock() + config.transformer_impl = "transformer_engine" + config.inference_fuse_tp_communication = False + return config + + @patch( + "megatron.training.argument_utils.HybridModelConfig", side_effect=lambda **kwargs: kwargs + ) + @patch("megatron.training.argument_utils.import_module") + def test_preserves_static_module_spec(self, mock_import_module, _mock_model_config): + static_spec = ModuleSpec(module=object) + mock_import_module.return_value = static_spec - return _add_mixed_precision_args(ArgumentParser(exit_on_error=False)) + config = hybrid_config_from_args(self._args(), config=self._transformer_config()) - def test_logit_dtype_defaults_to_input_dtype(self): - args = self._parser().parse_args([]) - assert args.logit_dtype is None + assert config["hybrid_stack_spec"] is static_spec - @pytest.mark.parametrize("dtype", ["bf16", "fp32"]) - def test_logit_dtype_accepts_supported_choices(self, dtype): - args = self._parser().parse_args(["--output-logit-dtype", dtype]) - assert args.logit_dtype == dtype + @patch( + "megatron.training.argument_utils.HybridModelConfig", side_effect=lambda **kwargs: kwargs + ) + @patch("megatron.training.argument_utils.import_module") + def test_resolves_config_aware_spec_factory(self, mock_import_module, _mock_model_config): + static_spec = ModuleSpec(module=object) + spec_factory = MagicMock(return_value=static_spec) + mock_import_module.return_value = spec_factory + transformer_config = self._transformer_config() - def test_logit_dtype_rejects_fp16(self): - with pytest.raises(ArgumentError, match="invalid choice"): - self._parser().parse_args(["--output-logit-dtype", "fp16"]) + config = hybrid_config_from_args(self._args(), config=transformer_config) + + spec_factory.assert_called_once_with(transformer_config) + assert config["hybrid_stack_spec"] is static_spec # --------------------------------------------------------------------------- @@ -802,6 +817,29 @@ def _make_args(**overrides): return Namespace(**defaults) +def _make_ddp_args(**overrides): + """Build the minimum Namespace required by get_megatron_ddp_config.""" + defaults = { + "accumulate_allreduce_grads_in_fp32": False, + "check_for_nan_in_loss_and_grad": False, + "check_for_large_grads": False, + "ddp_num_buckets": None, + "ddp_bucket_size": None, + "ddp_pad_buckets_for_high_nccl_busbw": False, + "ddp_reduce_scatter_with_fp32_accumulation": False, + "ddp_param_name_patterns_for_fp32_local_accumulation": [], + "ddp_average_in_collective": False, + "megatron_fsdp_main_params_dtype": None, + "megatron_fsdp_main_grads_dtype": None, + "megatron_fsdp_grad_comm_dtype": None, + "use_precision_aware_optimizer": False, + "use_megatron_fsdp": False, + "cuda_graph_impl": "none", + } + defaults.update(overrides) + return Namespace(**defaults) + + @pytest.fixture def mock_optimizer_config(): return MagicMock(spec=OptimizerConfig) @@ -857,6 +895,37 @@ def test_ddp_config_comes_from_helper(self, mock_opt, mock_ddp): assert result.ddp is mock_ddp_instance +class TestDDPConfigMapping: + """Test argparse Namespace → DistributedDataParallelConfig mapping.""" + + def test_cpu_backup_disable_fields_map_from_args(self): + from megatron.training.training import get_megatron_ddp_config + + args = _make_ddp_args( + disable_grad_buffers_cpu_backup=True, disable_param_buffers_cpu_backup=True + ) + + config = get_megatron_ddp_config(args) + + assert config.disable_grad_buffers_cpu_backup is True + assert config.disable_param_buffers_cpu_backup is True + + def test_cpu_backup_disable_cli_flags_map_to_ddp_config(self): + from megatron.training.arguments import _add_distributed_args + from megatron.training.training import get_megatron_ddp_config + + parser = _add_distributed_args(ArgumentParser()) + parsed_args = parser.parse_args( + ["--disable-grad-buffers-cpu-backup", "--disable-param-buffers-cpu-backup"] + ) + args = _make_ddp_args(**vars(parsed_args)) + + config = get_megatron_ddp_config(args) + + assert config.disable_grad_buffers_cpu_backup is True + assert config.disable_param_buffers_cpu_backup is True + + class TestCheckpointConfigMapping: """Test the boolean inversions and custom dest mappings for CheckpointConfig.""" diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index d5c625ca8ca..e4cec83056b 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -1,6 +1,7 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Note: --ckpt-format torch_dist has tests in tests/unit_tests/dist_checkpointing. import os +from pathlib import Path from types import SimpleNamespace from typing import Optional from unittest import mock @@ -24,9 +25,7 @@ _build_sharded_state_dict_metadata, _load_base_checkpoint, get_checkpoint_tracker_filename, - load_args_from_checkpoint, load_checkpoint, - maybe_save_dataloader_state, read_metadata, save_checkpoint, ) @@ -76,130 +75,6 @@ def sharded_state_dict(self, *args, metadata: Optional[dict] = None, **kwargs): return self.state_dict() -def test_maybe_save_dataloader_state_uses_explicit_process_groups(tmp_path): - """Dataloader checkpoints use the supplied module groups and canonical model-parallel path.""" - groups = { - "tp": SimpleNamespace(rank=0, size=2), - "pp": SimpleNamespace(rank=0, size=2), - "dp": SimpleNamespace(rank=3, size=4), - } - barriers = [] - saved = [] - iterator = SimpleNamespace( - iterable=SimpleNamespace(save_state=lambda: {"global_sequence_id": 16}) - ) - - with ( - mock.patch( - "megatron.training.checkpointing.get_pg_rank", side_effect=lambda group: group.rank - ), - mock.patch( - "megatron.training.checkpointing.get_pg_size", side_effect=lambda group: group.size - ), - mock.patch( - "megatron.training.checkpointing.torch.distributed.barrier", - side_effect=lambda group: barriers.append(group), - ), - mock.patch( - "megatron.training.checkpointing.torch.save", - side_effect=lambda state, path: saved.append((state, path)), - ), - ): - maybe_save_dataloader_state( - iterator, - 2, - tmp_path, - tp_group=groups["tp"], - pp_group=groups["pp"], - dp_group=groups["dp"], - ) - - assert barriers == [groups["dp"], groups["dp"]] - assert saved[0][0] == {"dataloader_state_dict": {"global_sequence_id": 16}} - assert saved[0][1] == str( - tmp_path / "iter_0000002" / "mp_rank_00_000" / "train_dataloader_dprank003.pt" - ) - - -def test_maybe_save_dataloader_state_skips_empty_state_after_barriers(tmp_path): - """Ranks without dataloader state participate in barriers but do not write a file.""" - group = SimpleNamespace(rank=0, size=1) - iterator = SimpleNamespace(iterable=SimpleNamespace(save_state=lambda: None)) - barriers = [] - - with ( - mock.patch( - "megatron.training.checkpointing.get_pg_rank", - side_effect=lambda process_group: process_group.rank, - ), - mock.patch( - "megatron.training.checkpointing.get_pg_size", - side_effect=lambda process_group: process_group.size, - ), - mock.patch( - "megatron.training.checkpointing.torch.distributed.barrier", - side_effect=lambda group: barriers.append(group), - ), - mock.patch("megatron.training.checkpointing.torch.save") as save, - ): - maybe_save_dataloader_state( - iterator, 2, tmp_path, tp_group=group, pp_group=group, dp_group=group - ) - - assert barriers == [group, group] - save.assert_not_called() - - -class MockOptParamScheduler(MockState): - def __init__(self, state_dict): - super().__init__(state_dict) - self.num_steps = state_dict.get("num_steps", 0) - self.step_calls = [] - - def load_state_dict(self, state_dict): - super().load_state_dict(state_dict) - self.num_steps = state_dict.get("num_steps", self.num_steps) - - def step(self, increment=1): - self.step_calls.append(increment) - - -class MockOptimizer(MockState): - def state_dict(self, is_loading=False): - state_dict = super().state_dict(is_loading=is_loading).copy() - state_dict["param_groups"] = [param_group.copy() for param_group in self.param_groups] - return state_dict - - def load_state_dict(self, state_dict): - super().load_state_dict(state_dict) - self.param_groups = [param_group.copy() for param_group in state_dict["param_groups"]] - - -@pytest.mark.parametrize( - ("checkpoint_args", "configured_num_householder", "expected_num_householder"), - [(SimpleNamespace(gdp_num_householder=5), 3, 5), (SimpleNamespace(), 5, 3)], -) -def test_load_args_restores_gdp_num_householder_from_checkpoint( - checkpoint_args, configured_num_householder, expected_num_householder -): - args = SimpleNamespace( - load="checkpoint", - iteration=0, - gdp_num_householder=configured_num_householder, - use_tokenizer_model_from_checkpoint_args=False, - use_mp_args_from_checkpoint_args=False, - ) - state_dict = {"args": checkpoint_args, "iteration": 12} - - with mock.patch( - "megatron.training.checkpointing._load_base_checkpoint", - return_value=(state_dict, "checkpoint", False, CheckpointType.LEGACY), - ): - restored_args, _ = load_args_from_checkpoint(args) - - assert restored_args.gdp_num_householder == expected_num_householder - - def create_checkpoint(load_path, ckpt_format): """Setup a dummy checkpoint directory.""" iteration = 123 @@ -242,7 +117,6 @@ def create_args(): args.auto_detect_ckpt_format = False args.ckpt_convert_update_legacy_dist_opt_format = False args.ckpt_step = None - args.override_opt_param_scheduler = False args.swiglu = True args.num_experts = 1 args.verify_integrity = False @@ -250,6 +124,87 @@ def create_args(): yield args +def test_async_save_rejects_reusable_chunked_optimizer_buffers(create_args): + """Checkpoint entry must preserve the CLI guard when startup validation was bypassed.""" + + create_args.async_save = True + create_args.ckpt_format = 'torch_dist' + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + pytest.raises(RuntimeError, match="does not support --async-save"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_rejects_mutated_format_with_chunked_optimizer_offload(create_args): + """Checkpoint entry revalidates the format after resume-time argument restoration.""" + + create_args.ckpt_format = 'torch' + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + pytest.raises(RuntimeError, match="requires --ckpt-format torch_dist"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_treats_zero_offload_fraction_as_disabled(create_args): + """A zero fraction must not impose async-save or distributed-format restrictions.""" + + create_args.async_save = True + create_args.ckpt_format = 'torch' + optimizer = SimpleNamespace( + config=SimpleNamespace( + chunked_optimizer_state_offload=True, optimizer_state_offload_fraction=0.0 + ) + ) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + mock.patch( + "megatron.training.checkpointing.is_empty_async_queue", + side_effect=RuntimeError("continued past optimizer offload guard"), + ), + pytest.raises(RuntimeError, match="continued past optimizer offload guard"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_skips_offload_guards_without_optimizer_state(create_args): + """Model-only saves do not expose reusable optimizer buffers to the async writer.""" + + create_args.async_save = True + create_args.ckpt_format = 'torch_dist' + create_args.no_save_optim = True + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + mock.patch( + "megatron.training.checkpointing.is_empty_async_queue", + side_effect=RuntimeError("continued past optimizer offload guard"), + ), + pytest.raises(RuntimeError, match="continued past optimizer offload guard"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_allows_legacy_format_without_optimizer_state(create_args): + """A model-only save is not forced through optimizer-specific distributed hooks.""" + + create_args.ckpt_format = 'torch' + create_args.no_save_optim = True + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + mock.patch( + "megatron.training.checkpointing.on_save_checkpoint_start", + side_effect=RuntimeError("continued past optimizer offload guard"), + ), + pytest.raises(RuntimeError, match="continued past optimizer offload guard"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + @pytest.fixture def create_ckpt_load_args(create_args): """Setup dummy args allowing checkpoint load.""" @@ -290,6 +245,44 @@ def init_model_parallel(): unset_num_microbatches_calculator() +def test_legacy_checkpoint_load_is_restricted(create_ckpt_load_args, tmp_path): + """Legacy checkpoint loading must use the restricted torch deserializer.""" + + load_dir = tmp_path / "test_legacy_checkpoint_load" + load_dir.mkdir() + create_checkpoint(load_dir, "torch") + create_ckpt_load_args.ckpt_format = "torch" + + with mock.patch( + "megatron.training.checkpointing.torch.load", + return_value={"args": "dummy", "iteration": 123}, + ) as load: + _load_base_checkpoint(load_dir, create_ckpt_load_args, rank0=True) + + load.assert_called_once() + assert load.call_args.kwargs["weights_only"] is True + + +def test_legacy_checkpoint_load_rejects_unsafe_pickle(create_ckpt_load_args, tmp_path): + """Loading a crafted legacy checkpoint must not execute its pickle payload.""" + + class UnsafeCheckpointPayload: + def __reduce__(self): + return Path.write_text, (tmp_path / "executed", "unsafe") + + load_dir = tmp_path / "test_unsafe_legacy_checkpoint" + load_dir.mkdir() + create_checkpoint(load_dir, "torch") + create_ckpt_load_args.ckpt_format = "torch" + checkpoint_name = load_dir / "iter_0000123" / "mp_rank_00" / "model_optim_rng.pt" + torch.save({"payload": UnsafeCheckpointPayload()}, checkpoint_name) + + with pytest.raises(SystemExit): + _load_base_checkpoint(load_dir, create_ckpt_load_args, rank0=True) + + assert not (tmp_path / "executed").exists() + + @pytest.mark.parametrize("ckpt_format", ["torch_dcp"]) def test_load_base_checkpoint( init_model_parallel, create_ckpt_load_args, ckpt_format, tmp_path_dist_ckpt @@ -360,7 +353,6 @@ def test_save_checkpoint(init_model_parallel, create_args, tmp_path_dist_ckpt, c with TempNamedDir(tmp_path_dist_ckpt / "test_save_checkpoint", sync=True) as save_dir: args.save = save_dir - args.save_tokenizer_assets = False set_args(args) save_checkpoint( @@ -397,7 +389,6 @@ def test_load_checkpoint( with TempNamedDir(tmp_path_dist_ckpt / "test_load_checkpoint", sync=True) as ckpt_dir: args.load = ckpt_dir args.save = ckpt_dir - args.save_tokenizer_assets = False set_args(args) # Create and save a checkpoint first. @@ -433,79 +424,6 @@ def test_load_checkpoint( assert new_opt_param_scheduler.state_dict() == opt_param_scheduler.state_dict() -@pytest.mark.parametrize("ckpt_format", ["torch"]) -def test_load_checkpoint_override_opt_param_scheduler( - init_model_parallel, create_ckpt_load_args, tmp_path_dist_ckpt, ckpt_format -): - """Test override_opt_param_scheduler behavior during checkpoint load.""" - args = create_ckpt_load_args - args.ckpt_format = ckpt_format - args.use_distributed_optimizer = False - args.use_dist_ckpt = ckpt_format != "torch" - args.override_opt_param_scheduler = True - args.lr = 1.0 - args.min_lr = 0.1 - args.decoupled_lr = 0.5 - args.decoupled_min_lr = 0.05 - args.consumed_train_samples = 42 - - with TempNamedDir( - tmp_path_dist_ckpt / "test_load_checkpoint_override_opt_param_scheduler", sync=True - ) as ckpt_dir: - args.load = ckpt_dir - args.save = ckpt_dir - args.save_tokenizer_assets = False - set_args(args) - - # Create and save a checkpoint first. - iteration = 123 - config = TransformerConfig(num_layers=1, kv_channels=1) - model = MockModel(config) - - optimizer = MockOptimizer({"optimizer": "optimizer_state"}) - optimizer.param_groups = [ - {"is_decoupled_lr": False, "max_lr": -1.0, "min_lr": -1.0}, - {"is_decoupled_lr": True, "max_lr": -1.0, "min_lr": -1.0}, - ] - opt_param_scheduler = MockOptParamScheduler( - {"opt_param_scheduler": "scheduler_state", "num_steps": 3} - ) - num_floating_point_operations_so_far = 456 - - save_checkpoint( - iteration, [model], optimizer, opt_param_scheduler, num_floating_point_operations_so_far - ) - - # Create new model, optimizer, and scheduler instances to load into. - new_model = MockModel(config) - new_optimizer = MockOptimizer({"optimizer": "dummy1"}) - new_optimizer.param_groups = [ - {"is_decoupled_lr": False, "max_lr": -2.0, "min_lr": -2.0}, - {"is_decoupled_lr": True, "max_lr": -2.0, "min_lr": -2.0}, - ] - new_opt_param_scheduler = MockOptParamScheduler( - {"opt_param_scheduler": "dummy2", "num_steps": 0} - ) - - # Load checkpoint and verify runtime overrides are restored. - loaded_iter, loaded_flops = load_checkpoint( - [new_model], new_optimizer, new_opt_param_scheduler, strict=True - ) - assert loaded_iter == iteration - assert loaded_flops == num_floating_point_operations_so_far - assert new_optimizer.param_groups[0]["max_lr"] == args.lr - assert new_optimizer.param_groups[0]["min_lr"] == args.min_lr - assert new_optimizer.param_groups[1]["max_lr"] == args.decoupled_lr - assert new_optimizer.param_groups[1]["min_lr"] == args.decoupled_min_lr - assert new_opt_param_scheduler.num_steps == args.consumed_train_samples - assert new_opt_param_scheduler.step_calls[-1] == 0 - - # Ensure loading without optimizer/scheduler remains safe. - loaded_iter_none, loaded_flops_none = load_checkpoint([new_model], None, None, strict=True) - assert loaded_iter_none == iteration - assert loaded_flops_none == num_floating_point_operations_so_far - - def test_dist_checkpoint_versioning(init_model_parallel, tmp_path_dist_ckpt, create_ckpt_load_args): """Test distributed checkpoint versioning.""" args = create_ckpt_load_args @@ -518,7 +436,6 @@ def test_dist_checkpoint_versioning(init_model_parallel, tmp_path_dist_ckpt, cre ) as ckpt_dir: args.load = ckpt_dir args.save = ckpt_dir - args.save_tokenizer_assets = False set_args(args) # Create and save a checkpoint first. diff --git a/tests/unit_tests/test_num_floating_point_operations.py b/tests/unit_tests/test_num_floating_point_operations.py index df5e4191843..f1872e3f744 100644 --- a/tests/unit_tests/test_num_floating_point_operations.py +++ b/tests/unit_tests/test_num_floating_point_operations.py @@ -54,6 +54,7 @@ def _make_gpt_args( args.group_query_attention = False args.num_query_groups = num_attention_heads args.attention_output_gate = False + args.gated_attention_proj_granularity = "elementwise" args.multi_latent_attention = False # MoE / MTP disabled. args.num_experts = None @@ -96,7 +97,32 @@ def _make_hybrid_args(*, num_layers=4, hidden_size=512, num_attention_heads=8, s args.mamba_head_dim = 64 args.mamba_num_groups = 8 args.mamba_num_heads = 128 - args.gdp_num_householder = 3 + return args + + +def _make_kda_hybrid_args(): + """Minimal KDA dimensions for hybrid FLOPs tests.""" + args = _make_hybrid_args() + args.hybrid_layer_pattern = "K" + args.linear_key_head_dim = 32 + args.linear_value_head_dim = 32 + args.linear_num_key_heads = 8 + args.linear_num_value_heads = 8 + args.linear_conv_kernel_dim = 4 + return args + + +def _make_mla_hybrid_args(): + """Minimal head-wise gated MLA dimensions for hybrid FLOPs tests.""" + args = _make_hybrid_args() + args.hybrid_layer_pattern = "+" + args.q_lora_rank = 128 + args.qk_head_dim = 48 + args.qk_pos_emb_head_dim = 16 + args.kv_lora_rank = 256 + args.v_head_dim = 64 + args.attention_output_gate = True + args.gated_attention_proj_granularity = "headwise" return args @@ -255,35 +281,237 @@ def test_hybrid_attention_layers_count(self): expected_delta = num_attn_layers * expected_delta_per_layer_per_unit_sum * bshd_sum assert flops_doubled - flops_bshd == expected_delta + def test_mla_attention_variants_are_counted(self): + """Regression: DSv4 MLA attention variants (CSA 'C', HCA 'H', Window 'W', + DS_ATTENTION 'D') must contribute attention FLOPs. Previously these symbols + were omitted, so a pattern made entirely of C/H/W layers reported zero + attention FLOPs -- roughly halving the throughput estimate.""" + args = _make_hybrid_args() + args.multi_latent_attention = True + # MLA dims (DSv4-like; values only need to be self-consistent and positive). + args.q_lora_rank = 256 + args.kv_lora_rank = 32 + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.v_head_dim = 64 + batch_size = 4 -class TestGatedDeltaProductFlops: - """GDP FLOPs must use the Householder count from the model configuration.""" + # MLA attention variants (C, H) must be counted exactly like dense MLA '+'. + # Both patterns are 2 MLA + 2 Mamba layers, so they yield identical FLOPs. + # Regular attention '*' is intentionally not the reference: hybrid grammar + # assigns different projection work to '*' and '+'. + args.hybrid_layer_pattern = "CMHM" + flops_mla_variants = num_floating_point_operations(args, batch_size) + args.hybrid_layer_pattern = "+M+M" + flops_dense_mla = num_floating_point_operations(args, batch_size) + + assert flops_mla_variants == flops_dense_mla + # And attention must be a non-trivial contributor: dropping the two + # attention layers entirely (all Mamba) changes the estimate. + args.hybrid_layer_pattern = "MMMM" + flops_no_attn = num_floating_point_operations(args, batch_size) + assert flops_mla_variants != flops_no_attn + + +class TestHybridMatchesStandard: + """The hybrid-model path (``hybrid_flops`` / ``mla_attn_layer_flops``) must + agree with the standard-model path (``transformer_flops``) for an equivalent + architecture. + + A dense Transformer with ``N`` layers has ``N`` attention sub-layers and + ``N`` MLP sub-layers. The hybrid path models those as separate layers, so the + matching hybrid pattern is ``N`` attention layers interleaved with ``N`` MLP + layers (``"*-" * N``). With identical hidden / head / ffn / vocab dims the + two estimates must be bit-for-bit equal -- this pins ``mla_attn_layer_flops`` + (and the dense ``attn_layer_flops``) against the reference MLA / MHA / GQA + terms in ``transformer_flops`` (see GitHub PR #5358), both for BSHD and for + the THD ``sum(L_i^2)`` / real-token scaling. + """ - def test_householder_count_changes_flops(self): - args = _make_hybrid_args() - args.spec = ["megatron.core.models.hybrid.hybrid_layer_specs", "gdp_stack_spec"] + @staticmethod + def _equivalent_pair(configure, num_layers=3): + """Build (standard_args, hybrid_args) describing the same architecture.""" + standard = _make_gpt_args(num_layers=num_layers) + configure(standard) + + hybrid = _make_gpt_args(num_layers=num_layers) + configure(hybrid) + # The hybrid path reads Mamba dims even when no Mamba layers are present. + hybrid.mamba_state_dim = 128 + hybrid.mamba_head_dim = 64 + hybrid.mamba_num_groups = 8 + hybrid.mamba_num_heads = 128 + # Match the standard model's attention implementation explicitly: + # '*' is regular attention and '+' is dense MLA in the hybrid grammar. + attention_symbol = "+" if hybrid.multi_latent_attention else "*" + hybrid.hybrid_layer_pattern = (attention_symbol + "-") * num_layers + return standard, hybrid + + def _assert_match(self, configure): + standard, hybrid = self._equivalent_pair(configure) batch_size = 4 - flops_m3 = num_floating_point_operations(args, batch_size) - args.gdp_num_householder = 4 - flops_m4 = num_floating_point_operations(args, batch_size) + # BSHD: closed-form defaults. + assert num_floating_point_operations(standard, batch_size) == num_floating_point_operations( + hybrid, batch_size + ) + + # THD: both paths must consume the ragged stats identically. + s = standard.seq_length + thd_sum_sq = batch_size * 4 * (s // 4) ** 2 + thd_tokens = batch_size * s * 3 // 4 + assert num_floating_point_operations( + standard, + batch_size, + seqlen_squared_sum_in_batch=thd_sum_sq, + total_real_tokens_in_batch=thd_tokens, + ) == num_floating_point_operations( + hybrid, + batch_size, + seqlen_squared_sum_in_batch=thd_sum_sq, + total_real_tokens_in_batch=thd_tokens, + ) + def test_dense_mha(self): + self._assert_match(lambda args: None) + + def test_gqa(self): + def configure(args): + args.group_query_attention = True + args.num_query_groups = 2 + + self._assert_match(configure) + + def test_mla_with_q_lora(self): + def configure(args): + args.multi_latent_attention = True + args.group_query_attention = False + args.q_lora_rank = 256 + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.kv_lora_rank = 256 + args.v_head_dim = 64 + + self._assert_match(configure) + + def test_mla_without_q_lora(self): + def configure(args): + args.multi_latent_attention = True + args.group_query_attention = False + args.q_lora_rank = None + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.kv_lora_rank = 256 + args.v_head_dim = 64 + + self._assert_match(configure) + + @pytest.mark.parametrize("gate_granularity", ("elementwise", "headwise")) + def test_mla_output_gate(self, gate_granularity): + def configure(args): + args.multi_latent_attention = True + args.group_query_attention = False + args.q_lora_rank = 256 + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.kv_lora_rank = 256 + args.v_head_dim = 64 + args.attention_output_gate = True + args.gated_attention_proj_granularity = gate_granularity + + self._assert_match(configure) + + +class TestKimiDeltaAttentionFlops: + """KDA layers must contribute their projection and kernel work.""" + + def test_direct_projection_formula(self): + args = _make_kda_hybrid_args() + batch_size = 2 total_tokens = batch_size * args.seq_length - d_inner = args.mamba_num_heads * args.mamba_head_dim - group_state_dim = args.mamba_num_groups * args.mamba_state_dim - forward_delta_per_layer = ( + qk_dim = args.linear_key_head_dim * args.linear_num_key_heads + v_dim = args.linear_value_head_dim * args.linear_num_value_heads + in_proj_dim = 3 * qk_dim + 2 * v_dim + kda_forward = ( 2 * total_tokens * ( - args.hidden_size * (d_inner + group_state_dim + args.mamba_num_heads) - + 4 * (d_inner + group_state_dim) + args.hidden_size * (in_proj_dim + args.linear_num_key_heads) + + args.linear_conv_kernel_dim * (2 * qk_dim + v_dim) + + args.hidden_size * v_dim ) - + 4 * total_tokens * d_inner * args.mamba_state_dim ) - num_gdp_layers = 2 - expected_delta = 3 * num_gdp_layers * forward_delta_per_layer + kda_forward += ( + 8 + * total_tokens + * args.linear_num_key_heads + * args.linear_key_head_dim + * args.linear_value_head_dim + ) + logits_forward = 2 * total_tokens * args.hidden_size * args.padded_vocab_size + + assert num_floating_point_operations(args, batch_size) == 3 * (kda_forward + logits_forward) + + +class TestMLAHeadwiseOutputGateFlops: + """Head-wise gated MLA must account for gate projection and ragged attention work.""" + + def test_uses_ragged_attention_work(self): + args = _make_mla_hybrid_args() + batch_size = 2 + bshd_sum = batch_size * args.seq_length**2 + flops_full = num_floating_point_operations( + args, batch_size, seqlen_squared_sum_in_batch=bshd_sum + ) + flops_half = num_floating_point_operations( + args, batch_size, seqlen_squared_sum_in_batch=bshd_sum // 2 + ) + expected_delta = ( + 3 + * (bshd_sum // 2) + * args.num_attention_heads + * (args.qk_head_dim + args.qk_pos_emb_head_dim + args.v_head_dim) + ) + assert flops_full - flops_half == expected_delta + + def test_headwise_gate_projection_work(self): + args = _make_mla_hybrid_args() + batch_size = 2 + total_tokens = batch_size * args.seq_length + + args.attention_output_gate = False + ungated_flops = num_floating_point_operations(args, batch_size) + args.attention_output_gate = True + gated_flops = num_floating_point_operations(args, batch_size) + + expected_delta = 3 * 2 * total_tokens * args.hidden_size * args.num_attention_heads + assert gated_flops - ungated_flops == expected_delta + - assert flops_m4 - flops_m3 == expected_delta +class TestHybridAttentionOutputGateFlops: + """Regular attention output gating must use the regular attention projection width.""" + + @pytest.mark.parametrize( + ("gate_granularity", "gate_projection_size"), (("elementwise", 512), ("headwise", 8)) + ) + def test_projection_work(self, gate_granularity, gate_projection_size): + args = _make_hybrid_args() + args.hybrid_layer_pattern = "*" + # The model-wide MLA flag is also enabled in mixed '*' / '+' models, but '*' remains + # regular attention. Use a different MLA value-head width to make that distinction visible. + args.multi_latent_attention = True + args.v_head_dim = 32 + batch_size = 2 + total_tokens = batch_size * args.seq_length + + args.attention_output_gate = False + ungated_flops = num_floating_point_operations(args, batch_size) + args.attention_output_gate = True + args.gated_attention_proj_granularity = gate_granularity + gated_flops = num_floating_point_operations(args, batch_size) + + expected_delta = 3 * 2 * total_tokens * args.hidden_size * gate_projection_size + assert gated_flops - ungated_flops == expected_delta class TestPaddingRemoval: @@ -675,3 +903,272 @@ def test_dedup_across_topology(self, tp, cp, pp): f"topology tp={tp} cp={cp} pp={pp} dp={dp_size}: " f"got seqlen_squared_sum={seqlen_squared_sum}, expected {expected_sum_sq}" ) + + +def _make_dsv4_args(): + """Minimal args for a DSv4-hybrid MLA model with sparse attention. + + 4 layers with compress_ratios [0, 4, 128, 128] (1 r0, 1 r4, 2 r128). + No MoE / MTP to keep the golden reference simple. + """ + args = _make_gpt_args( + num_layers=4, + hidden_size=512, + num_attention_heads=8, + seq_length=256, + ffn_hidden_size=2048, + padded_vocab_size=1024, + ) + args.multi_latent_attention = True + args.group_query_attention = False + args.q_lora_rank = 128 + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.kv_lora_rank = 64 + args.v_head_dim = 64 + args.o_lora_rank = 64 + args.o_groups = 2 + args.experimental_attention_variant = "dsv4_hybrid" + args.csa_window_size = 64 + args.csa_compress_ratios = [0, 4, 128, 128] + args.dsa_indexer_n_heads = 4 + args.dsa_indexer_head_dim = 32 + args.dsa_indexer_topk = 16 + return args + + +def _dsv4_golden_flops(args, total_tokens, seqlen_squared_sum): + """Independent golden calculator for DSv4-hybrid FLOPs. + + Reimplements the formula from ``num_floating_point_operations`` so that the + test does not just call the same code twice. Assumes no MoE / MTP. + """ + fwd_bwd = 3 + fma = 2 + ffn_exp = 3 if args.swiglu else 2 + + # ---- MLA projections (token-linear, per layer) ---- + q_term = args.q_lora_rank * (args.hidden_size + args.num_attention_heads * args.v_head_dim + 1) + kv_term = args.hidden_size * args.v_head_dim + args.v_head_dim + o_term = ( + args.num_attention_heads * args.v_head_dim * args.o_lora_rank + + args.o_groups * args.o_lora_rank * args.hidden_size + ) + mla_proj_per_layer = fwd_bwd * fma * (q_term + kv_term + o_term) + + # ---- DSv4 sparse attention extra (token-linear + L^2) ---- + ratios = args.csa_compress_ratios + n_r0 = sum(1 for r in ratios if r == 0) + n_r4 = sum(1 for r in ratios if r == 4) + n_r128 = sum(1 for r in ratios if r == 128) + nh = args.num_attention_heads + vhd = args.v_head_dim + w = args.csa_window_size + + # Token-linear sparse attention + sparse_r0 = n_r0 * nh * w * vhd * 2 + sparse_r128_win = n_r128 * nh * w * vhd * 2 + if n_r4 > 0: + eff_topk = min(args.dsa_indexer_topk, args.seq_length // 4) + avg_comp_4 = eff_topk * (1 - eff_topk * 4 / (2 * args.seq_length)) + sparse_r4 = n_r4 * nh * (w + avg_comp_4) * vhd * 2 + idx_tok = ( + n_r4 * args.hidden_size * (2 * args.dsa_indexer_head_dim) * 2 + + n_r4 * args.q_lora_rank * args.dsa_indexer_n_heads * args.dsa_indexer_head_dim + + n_r4 * args.hidden_size * args.dsa_indexer_n_heads + ) + idx_core = n_r4 * args.dsa_indexer_n_heads * args.dsa_indexer_head_dim / 4 + else: + sparse_r4, idx_tok, idx_core = 0, 0, 0 + + compressor = n_r4 * args.hidden_size * (2 * vhd) * 2 + n_r128 * args.hidden_size * (1 * vhd) * 2 + dsv4_token = fwd_bwd * fma * (sparse_r0 + sparse_r4 + sparse_r128_win + compressor + idx_tok) + # L^2 core: r=128 compressed-KV + r=4 indexer scoring + r128_core = n_r128 * nh * vhd / 128 + dsv4_core = fwd_bwd * fma * (r128_core + idx_core) + + # ---- Aggregation ---- + num_layers = args.num_layers + self_attn_term = mla_proj_per_layer * num_layers + dsv4_token + self_attn_core_term = dsv4_core # standard core is 0 for DSv4 + + mlp = fwd_bwd * fma * args.hidden_size * (args.ffn_hidden_size * ffn_exp * num_layers) + logit = fwd_bwd * fma * args.hidden_size * args.padded_vocab_size + + return total_tokens * (mlp + self_attn_term + logit) + seqlen_squared_sum * self_attn_core_term + + +class TestDSv4Hybrid: + """DSv4 hybrid sparse-attention FLOPs against an independent golden calculator.""" + + def test_bshd(self): + """BSHD (uniform sequences) must match the golden calculator.""" + args = _make_dsv4_args() + batch_size = 2 + total_tokens = batch_size * args.seq_length + sum_sq = batch_size * args.seq_length**2 + + flops = num_floating_point_operations(args, batch_size) + expected = _dsv4_golden_flops(args, total_tokens, sum_sq) + assert flops == expected + + def test_thd(self): + """THD (packed variable-length subsequences) must match the golden + calculator and be strictly less than BSHD due to the L^2 sparse-attention + components (r=128 compressed-KV, r=4 indexer scoring).""" + args = _make_dsv4_args() + batch_size = 2 + packed_lengths = [64, 64, 128, 256] + total_tokens = sum(packed_lengths) + thd_sum_sq = sum(L**2 for L in packed_lengths) + + flops = num_floating_point_operations( + args, + batch_size, + seqlen_squared_sum_in_batch=thd_sum_sq, + total_real_tokens_in_batch=total_tokens, + ) + expected = _dsv4_golden_flops(args, total_tokens, thd_sum_sq) + assert flops == expected + # THD must be strictly less than BSHD. + bshd_flops = num_floating_point_operations(args, batch_size) + assert flops < bshd_flops + + +# ``compress_ratio -> hybrid attention symbol``: Window (r0) / CSA (r4) / HCA (r128). +_RATIO_TO_SYMBOL = {0: "W", 4: "C", 128: "H"} + + +def _make_dsv4_pair(*, ratios, mtp, moe): + """Build (standard_args, hybrid_args) for the SAME DSv4-hybrid model. + + The standard-model path expresses it as ``len(ratios)`` fused layers + (MLA attention + FFN per layer, ``experimental_attention_variant`` = + ``dsv4_hybrid``); the hybrid-model path expresses the identical model as a + pattern with one attention layer (W/C/H per ratio) and one FFN layer (MoE + ``E`` or dense ``-``) each. With one MTP depth, the standard path appends a + fused MTP layer (ratio 0 / Window) and the hybrid path appends ``/W``. + """ + ffn_symbol = "E" if moe else "-" + + def _configure(args): + args.multi_latent_attention = True + args.group_query_attention = False + args.q_lora_rank = 256 + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.kv_lora_rank = 64 + args.v_head_dim = 64 + args.o_lora_rank = 64 + args.o_groups = 2 + args.experimental_attention_variant = "dsv4_hybrid" + args.csa_window_size = 64 + args.dsa_indexer_n_heads = 4 + args.dsa_indexer_head_dim = 32 + args.dsa_indexer_topk = 16 + if moe: + args.num_experts = 256 + args.moe_layer_freq = 1 + args.moe_router_topk = 6 + args.moe_ffn_hidden_size = 2048 + args.moe_shared_expert_intermediate_size = 2048 + + # ---- standard-model args (fused layers) ---- + standard = _make_gpt_args( + num_layers=len(ratios), + hidden_size=512, + num_attention_heads=8, + seq_length=256, + ffn_hidden_size=2048, + padded_vocab_size=1024, + ) + _configure(standard) + # ``csa_compress_ratios`` must cover main + MTP layers; the MTP layer is a + # Window (ratio 0) attention layer. + standard.csa_compress_ratios = list(ratios) + [0] * mtp + standard.mtp_num_layers = mtp if mtp else None + + # ---- hybrid-model args (split layers, same dims) ---- + hybrid = _make_gpt_args( + num_layers=len(ratios), + hidden_size=512, + num_attention_heads=8, + seq_length=256, + ffn_hidden_size=2048, + padded_vocab_size=1024, + ) + _configure(hybrid) + hybrid.mamba_state_dim = 128 + hybrid.mamba_head_dim = 64 + hybrid.mamba_num_groups = 8 + hybrid.mamba_num_heads = 128 + main_pattern = "".join(_RATIO_TO_SYMBOL[r] + ffn_symbol for r in ratios) + mtp_pattern = ("/" + "W" + ffn_symbol) * mtp # MTP layer = Window attention + hybrid.hybrid_layer_pattern = main_pattern + mtp_pattern + hybrid.mtp_num_layers = mtp if mtp else None + return standard, hybrid + + +class TestDSv4HybridMatchesStandard: + """The hybrid-model path and the standard-model path must agree for a DSv4 + model with sparse attention. + + Regression for GitHub PR #5485: ``hybrid_flops`` previously routed DSv4 + attention through the plain full-MLA helper, which (a) computes the full + ``O(L^2)`` core attention that DSv4 actually replaces with sparse attention, + and (b) uses the dense output projection instead of DSv4's grouped low-rank + one. That inflated the hybrid-model FLOPs estimate by ~30% vs. the + (correct) standard-model ``dsv4_hybrid`` estimate, even though both describe + the same model. Both paths now share ``_dsv4_hybrid_self_attention_flops``, + so they must match bit-for-bit -- for BSHD and THD, with and without MTP, + and for dense and MoE FFN layers. + """ + + @staticmethod + def _assert_match(*, ratios, mtp, moe): + standard, hybrid = _make_dsv4_pair(ratios=ratios, mtp=mtp, moe=moe) + batch_size = 2 + + # BSHD (closed-form defaults). + assert num_floating_point_operations(standard, batch_size) == num_floating_point_operations( + hybrid, batch_size + ) + + # THD: both paths must consume the ragged stats identically. + s = standard.seq_length + thd_sum_sq = batch_size * 4 * (s // 4) ** 2 + thd_tokens = batch_size * s * 3 // 4 + assert num_floating_point_operations( + standard, + batch_size, + seqlen_squared_sum_in_batch=thd_sum_sq, + total_real_tokens_in_batch=thd_tokens, + ) == num_floating_point_operations( + hybrid, + batch_size, + seqlen_squared_sum_in_batch=thd_sum_sq, + total_real_tokens_in_batch=thd_tokens, + ) + + def test_dense_no_mtp(self): + self._assert_match(ratios=[0, 4, 128], mtp=0, moe=False) + + def test_dense_with_mtp(self): + self._assert_match(ratios=[0, 4, 128], mtp=1, moe=False) + + def test_moe_no_mtp(self): + self._assert_match(ratios=[0, 4, 128, 128], mtp=0, moe=True) + + def test_moe_with_mtp(self): + # Production-like: MoE FFN, window/CSA/HCA mix, one MTP depth. + self._assert_match(ratios=[0, 0, 4, 128, 4, 128], mtp=1, moe=True) + + def test_hybrid_overcount_is_gone(self): + """The OLD bug made the hybrid estimate strictly larger than the standard + one (phantom full-MLA O(L^2) core attention). Pin that they are equal, + not merely close.""" + standard, hybrid = _make_dsv4_pair(ratios=[0, 4, 128, 128], mtp=1, moe=True) + batch_size = 2 + std_flops = num_floating_point_operations(standard, batch_size) + hyb_flops = num_floating_point_operations(hybrid, batch_size) + assert hyb_flops == std_flops diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py index d6c8fcbd919..24685b5c0f9 100644 --- a/tests/unit_tests/transformer/moe/test_routers.py +++ b/tests/unit_tests/transformer/moe/test_routers.py @@ -9,6 +9,7 @@ from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules +from megatron.core.transformer.moe.moe_logging import MoEMetricsTracker from megatron.core.transformer.moe.moe_utils import ( get_updated_expert_bias, router_gating_linear, @@ -148,6 +149,27 @@ def test_aux_loss(self): out.sum().mul_(0).backward() assert self.sequential_mlp.router.weight.grad.abs().sum() > 0 + @pytest.mark.internal + def test_nested_hybrid_mtp_z_loss_uses_mtp_depth_for_metric_index(self, monkeypatch): + self.router.config.moe_z_loss_coeff = 1.0 + self.router.config.mtp_num_layers = 1 + self.router.is_mtp_layer = True + self.router.layer_number = 2 + self.router.mtp_layer_number = 1 + + tracker = MoEMetricsTracker() + monkeypatch.setattr( + "megatron.core.transformer.moe.router.get_moe_metrics_tracker", lambda: tracker + ) + logits = torch.randn(8, self.router.config.num_moe_experts, requires_grad=True) + + self.router.apply_z_loss(logits) + + values = tracker.metrics["z_loss"].values + assert values.shape == (self.router.config.num_layers + 1,) + assert torch.count_nonzero(values[:-1]) == 0 + assert values[-1] > 0 + @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_router_with_padding_mask(self): diff --git a/tests/unit_tests/transformer/test_attention.py b/tests/unit_tests/transformer/test_attention.py index e82ae7e458a..9b2d3954a18 100644 --- a/tests/unit_tests/transformer/test_attention.py +++ b/tests/unit_tests/transformer/test_attention.py @@ -1,18 +1,19 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import copy from unittest import mock +import einops import pytest import torch from packaging import version +from torch.nn import functional as F import megatron.core.parallel_state as parallel_state +from megatron.core.datasets.data_schedule_utils import get_cp_slice_for_thd from megatron.core.hyper_comm_grid import HyperCommGrid -from megatron.core.models.common.embeddings.rope_utils import ( - get_pos_emb_on_this_cp_rank as get_tensor_on_this_cp_rank, -) from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_local_spec, get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_with_transformer_engine_submodules, ) @@ -21,6 +22,10 @@ from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig from megatron.core.transformer.attention import SelfAttention +from megatron.core.transformer.dot_product_attention_context_parallel import ( + AttentionFuncionWithContextParallel, + to_zz_mask_attn_bias, +) from megatron.core.transformer.enums import AttnMaskType from megatron.core.utils import is_te_min_version, unwrap_model from megatron.training.arguments import parse_args @@ -44,10 +49,19 @@ @pytest.mark.parametrize("output_gate", [False, True]) +@pytest.mark.parametrize( + ("transformer_impl", "fallback_to_eager_attn"), + [("transformer_engine", False), ("transformer_engine", True), ("native", False)], +) class TestParallelAttention: @pytest.fixture(scope='function', autouse=True) - def setup_method(self, output_gate): + def setup_method(self, output_gate, transformer_impl, fallback_to_eager_attn): + if output_gate: + if transformer_impl == "native": + pytest.skip("Native implementation does not support output gate.") + if fallback_to_eager_attn: + pytest.skip("No need to test output gate for fallback_to_eager_attn = True.") Utils.initialize_model_parallel(1, 1) model_parallel_cuda_manual_seed(123) self.transformer_config = TransformerConfig( @@ -58,11 +72,17 @@ def setup_method(self, output_gate): bf16=True, params_dtype=torch.bfloat16, attention_output_gate=output_gate, + transformer_impl=transformer_impl, + fallback_to_eager_attn=fallback_to_eager_attn, ) + if transformer_impl == "transformer_engine": + attn_layer_spec = ( + get_gpt_layer_with_transformer_engine_submodules().self_attention.submodules + ) + else: + attn_layer_spec = get_gpt_layer_local_spec().submodules.self_attention.submodules self.parallel_attention = SelfAttention( - self.transformer_config, - get_gpt_layer_with_transformer_engine_submodules().self_attention.submodules, - layer_number=1, + self.transformer_config, attn_layer_spec, layer_number=1 ) def teardown_method(self): @@ -73,10 +93,19 @@ def test_constructor(self): assert self.parallel_attention.layer_number == 1 num_weights = sum([p.numel() for p in self.parallel_attention.parameters()]) + + hidden_size = self.transformer_config.hidden_size + standard_num_weights = ( + hidden_size * hidden_size * 4 + hidden_size * 4 # QKVO weight # QKVO bias + ) if self.transformer_config.attention_output_gate: - assert num_weights == 82816 - else: - assert num_weights == 66304 + standard_num_weights += hidden_size * hidden_size + hidden_size # Gate weight and bias + if self.transformer_config.transformer_impl == "transformer_engine": + standard_num_weights += hidden_size * 2 # fused pre layernorm weight and bias + + assert ( + num_weights == standard_num_weights + ), f"{num_weights=} does not match {standard_num_weights=}." def test_cpu_forward(self): # we can't currently do this because the global memory buffer is on GPU @@ -111,6 +140,8 @@ def test_gpu_forward(self): @pytest.mark.parametrize("rotary_interleaved", [True, False]) @pytest.mark.parametrize("fused_qkv_rope", [True, False]) def test_fused_rope_gpu_forward(self, rotary_interleaved, fused_qkv_rope): + if self.transformer_config.fallback_to_eager_attn: + pytest.skip("No need to test fused RoPE for fallback_to_eager_attn = True.") self.parallel_attention.config.apply_rope_fusion = True if rotary_interleaved and not is_te_min_version("2.3.0"): pytest.skip("Only TE >= 2.3.0 supports interleaved fused RoPE.") @@ -374,6 +405,199 @@ class TestSelfAttention: @pytest.fixture(scope='function', autouse=True) def setup_method(self, output_gate): + self.output_gate = output_gate + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self): + Utils.destroy_model_parallel() + + def test_clip_qk_disabled_raises_error(self): + """Test that clip_qk raises ValueError when qk_clip is not enabled.""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + qk_clip=False, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + + with pytest.raises(ValueError, match="qk_clip option needs to be enabled"): + attention.clip_qk() + + def test_clip_qk_none_logits_raises_error(self): + """Test that clip_qk raises ValueError when current_max_attn_logits is None.""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + qk_clip=True, + qk_clip_threshold=100.0, + qk_clip_alpha=0.5, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + + with pytest.raises(ValueError, match="current_max_attn_logits is None"): + attention.clip_qk() + + def test_clip_qk_below_threshold_no_update(self): + """Test that weights are not updated when max logits are below threshold.""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + qk_clip=True, + qk_clip_threshold=100.0, + qk_clip_alpha=0.5, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + attention.cuda() + + # Save original weights + original_weight = attention.linear_qkv.weight.data.clone() + + # Set current_max_attn_logits below threshold + attention.core_attention.current_max_attn_logits = torch.tensor( + [50.0, 60.0, 70.0, 80.0], device='cuda' + ) + + # Call clip_qk + attention.clip_qk() + + # Weights should not be updated + assert torch.equal(attention.linear_qkv.weight.data, original_weight) + # current_max_attn_logits should be reset + assert attention.core_attention.current_max_attn_logits is None + + def test_clip_qk_above_threshold_updates_weights(self): + """Test that weights are updated when max logits exceed threshold.""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + qk_clip=True, + qk_clip_threshold=100.0, + qk_clip_alpha=0.5, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + attention.cuda() + + # Save original weights + original_weight = attention.linear_qkv.weight.data.clone() + + # Set current_max_attn_logits above threshold + attention.core_attention.current_max_attn_logits = torch.tensor( + [150.0, 160.0, 170.0, 180.0], device='cuda' + ) + + # Call clip_qk + attention.clip_qk() + + # Weights should be updated + assert not torch.equal(attention.linear_qkv.weight.data, original_weight) + # current_max_attn_logits should be reset + assert attention.core_attention.current_max_attn_logits is None + + def test_clip_qk_gqa_configuration(self): + """Test clip_qk with GQA (Grouped Query Attention) configuration.""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=8, + num_query_groups=4, # GQA with 2 heads per group + use_cpu_initialization=True, + qk_clip=True, + qk_clip_threshold=100.0, + qk_clip_alpha=0.5, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + attention.cuda() + + # Save original weights + original_weight = attention.linear_qkv.weight.data.clone() + + # Set current_max_attn_logits for all heads (8 heads) + attention.core_attention.current_max_attn_logits = torch.tensor( + [150.0, 160.0, 170.0, 180.0, 190.0, 200.0, 210.0, 220.0], device='cuda' + ) + + # Call clip_qk + attention.clip_qk() + + # Weights should be updated + assert not torch.equal(attention.linear_qkv.weight.data, original_weight) + # current_max_attn_logits should be reset + assert attention.core_attention.current_max_attn_logits is None + + def test_clip_qk_mixed_logits(self): + """Test clip_qk with mixed logits (some above, some below threshold).""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + qk_clip=True, + qk_clip_threshold=100.0, + qk_clip_alpha=0.5, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + attention.cuda() + + # Save original weights + original_weight = attention.linear_qkv.weight.data.clone() + + # Set mixed current_max_attn_logits (some above, some below threshold) + attention.core_attention.current_max_attn_logits = torch.tensor( + [80.0, 150.0, 90.0, 200.0], device='cuda' + ) + + # Call clip_qk + attention.clip_qk() + + # Weights should be updated since at least one head exceeds threshold + assert not torch.equal(attention.linear_qkv.weight.data, original_weight) + # current_max_attn_logits should be reset + assert attention.core_attention.current_max_attn_logits is None + + +@pytest.mark.parametrize("output_gate", [False, True]) +@pytest.mark.parametrize("transformer_impl", ["transformer_engine", "native"]) +class TestSelfAttention: + + @pytest.fixture(scope='function', autouse=True) + def setup_method(self, output_gate, transformer_impl): + if transformer_impl == "native": + if output_gate: + pytest.skip("Native implementation does not support output gate.") + self.transformer_impl = transformer_impl self.output_gate = output_gate Utils.destroy_model_parallel() @@ -389,10 +613,17 @@ def run_self_attention(self, pg_collection): attention_output_gate=self.output_gate, tensor_model_parallel_size=tensor_model_parallel_size, use_cpu_initialization=False, + transformer_impl=self.transformer_impl, ) + if self.transformer_impl == "transformer_engine": + attn_layer_spec = ( + get_gpt_layer_with_transformer_engine_submodules().self_attention.submodules + ) + else: + attn_layer_spec = get_gpt_layer_local_spec().submodules.self_attention.submodules self.self_attention = SelfAttention( self.transformer_config, - get_gpt_layer_with_transformer_engine_submodules().self_attention.submodules, + attn_layer_spec, layer_number=1, attn_mask_type=AttnMaskType.causal, pg_collection=pg_collection, @@ -473,6 +704,8 @@ def _test_parallel_attention_correctness( tmp_path_dist_ckpt, atol, rtol, + cosine_similarity_threshold=None, + relative_l2_threshold=0.1, tp=1, sp=False, cp=1, @@ -480,6 +713,10 @@ def _test_parallel_attention_correctness( sequence_length=256, micro_batch_size=4, sequence_packing=False, + cp_partition_mode="zigzag", + compare_param_grads=False, + input_grad_atol=None, + input_grad_rtol=None, ): # Model initialization function def initialize_gpt_model( @@ -515,7 +752,6 @@ def initialize_gpt_model( with TempNamedDir(tmp_path_dist_ckpt / 'test_parallel_attn', sync=True) as ckpt_dir: # Set argument mock_args = parse_args(ignore_unknown_args=True) - mock_args.save_tokenizer_assets = False set_args(mock_args) # Initialize baseline model @@ -532,8 +768,24 @@ def initialize_gpt_model( mock_args.no_load_rng = True save_checkpoint(10, gpt_model, None, None, 0) + def get_param_grad(param): + grad = param.grad + if grad is None: + grad = getattr(param, "main_grad", None) + if grad is None: + return torch.zeros_like(param, dtype=torch.float32) + return grad + + def zero_param_grads(module): + for param in module.parameters(): + param.grad = None + main_grad = getattr(param, "main_grad", None) + if main_grad is not None: + main_grad.zero_() + # Calculate baseline output attention = gpt_model[0].decoder.layers[0].self_attention + zero_param_grads(attention) output_hidden_states_baseline, bias_hidden_states_baseline = attention( input_hidden_states, attention_mask=None ) @@ -541,6 +793,13 @@ def initialize_gpt_model( # Save baseline output input_grad_baseline = input_hidden_states.grad.detach() + param_grads_baseline = None + if compare_param_grads: + param_grads_baseline = { + name: get_param_grad(param).detach().float().clone() + for name, param in attention.named_parameters() + if param.requires_grad + } output_hidden_states_baseline = output_hidden_states_baseline.detach() bias_hidden_states_baseline = bias_hidden_states_baseline if bias_hidden_states_baseline is not None: @@ -559,9 +818,11 @@ def initialize_gpt_model( transformer_config.context_parallel_size = cp transformer_config.tensor_model_parallel_size = tp transformer_config.sequence_parallel = sp + transformer_config.cp_partition_mode = cp_partition_mode init_basic_mock_args(mock_args, tp, 1, bf16=True) mock_args.context_parallel_size = cp mock_args.sequence_parallel = sp + mock_args.cp_partition_mode = cp_partition_mode gpt_model = unwrap_model(get_model(initialize_gpt_model, config=transformer_config)) with mock.patch('megatron.training.checkpointing.check_checkpoint_args'): with mock.patch('megatron.training.checkpointing.update_num_microbatches'): @@ -572,10 +833,31 @@ def initialize_gpt_model( tp_rank = parallel_state.get_tensor_model_parallel_rank() def get_tensor_on_this_rank(tensor): - if cp > 1: - tensor = get_tensor_on_this_cp_rank(tensor, 0, cp_group) if sequence_packing: tensor = tensor.transpose(0, 1).contiguous().view(-1, 1, *tensor.shape[2:]) + if cp > 1: + cu_seqlens_tensor = torch.tensor( + [i * sequence_length for i in range(micro_batch_size + 1)], + device=tensor.device, + ) + batch = {"hidden_states": tensor, "cu_seqlens_padded": cu_seqlens_tensor} + get_cp_slice_for_thd( + batch, + cp_group, + keys=("hidden_states",), + cp_partition_mode=cp_partition_mode, + ) + tensor = batch["hidden_states"] + elif cp > 1: + cp_rank = torch.distributed.get_rank(cp_group) + if cp_partition_mode == "zigzag": + cp_idx = torch.tensor([cp_rank, 2 * cp - cp_rank - 1], device=tensor.device) + elif cp_partition_mode == "contiguous": + cp_idx = torch.tensor([2 * cp_rank, 2 * cp_rank + 1], device=tensor.device) + else: + raise ValueError(f"Unsupported test CP partition mode {cp_partition_mode!r}.") + tensor = tensor.view(2 * cp, -1, *tensor.shape[1:]) + tensor = tensor.index_select(0, cp_idx).view(-1, *tensor.shape[2:]) if tp > 1 and sp: sp_seg = tensor.shape[0] // tp tensor = tensor[tp_rank * sp_seg : (tp_rank + 1) * sp_seg] @@ -585,15 +867,27 @@ def get_tensor_on_this_rank(tensor): if sequence_packing: cu_seqlens = [i * sequence_length for i in range(micro_batch_size + 1)] packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) + packed_seq_params.cp_partition_mode = cp_partition_mode else: packed_seq_params = None input_hidden_states = get_tensor_on_this_rank(input_hidden_states) input_hidden_states = input_hidden_states.detach().requires_grad_(True) parallel_attention = gpt_model[0].decoder.layers[0].self_attention + zero_param_grads(parallel_attention) output_hidden_states_parallel, bias_hidden_states_parallel = parallel_attention( input_hidden_states, attention_mask=None, packed_seq_params=packed_seq_params ) output_hidden_states_parallel.sum().backward() + param_grads_parallel = None + if compare_param_grads: + assert tp == 1, "Parameter gradient parity only supports unsharded parameters." + param_grads_parallel = {} + for name, param in parallel_attention.named_parameters(): + if param.requires_grad: + grad = get_param_grad(param) + if cp > 1: + torch.distributed.all_reduce(grad, group=cp_group) + param_grads_parallel[name] = grad.detach().float().clone() input_grad_parallel = input_hidden_states.grad.detach() # Check if the output is close @@ -630,32 +924,65 @@ def get_tensor_on_this_rank(tensor): ~torch.isinf(bias_hidden_states_parallel) ), "bias_hidden_states_parallel contains inf" - torch.testing.assert_close( - output_hidden_states_baseline, - output_hidden_states_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in output_hidden_states: {msg}", + def assert_close_or_cosine_similarity( + baseline, parallel, tensor_name, tensor_atol=atol, tensor_rtol=rtol + ): + try: + torch.testing.assert_close( + baseline, + parallel, + atol=tensor_atol, + rtol=tensor_rtol, + msg=lambda msg: f"Mismatch in {tensor_name}: {msg}", + ) + return + except AssertionError as close_error: + if cosine_similarity_threshold is None: + raise close_error + baseline_flat = baseline.flatten().float() + parallel_flat = parallel.flatten().float() + cosine_sim = torch.nn.functional.cosine_similarity( + baseline_flat.unsqueeze(0), parallel_flat.unsqueeze(0) + ).item() + diff_norm = torch.linalg.vector_norm((parallel_flat - baseline_flat).float()) + baseline_norm = torch.linalg.vector_norm(baseline_flat.float()).clamp_min(1e-12) + relative_l2 = (diff_norm / baseline_norm).item() + assert cosine_sim >= cosine_similarity_threshold, ( + f"Mismatch in {tensor_name}: cosine similarity " + f"{cosine_sim} < {cosine_similarity_threshold}, " + f"while assert_close failed: {close_error}" + ) + assert relative_l2 <= relative_l2_threshold, ( + f"Mismatch in {tensor_name}: relative L2 " + f"{relative_l2} > {relative_l2_threshold}, " + f"while assert_close failed: {close_error}" + ) + + assert_close_or_cosine_similarity( + output_hidden_states_baseline, output_hidden_states_parallel, "output_hidden_states" ) - torch.testing.assert_close( + assert_close_or_cosine_similarity( input_grad_baseline, input_grad_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in input_grad: {msg}", + "input_grad", + tensor_atol=atol if input_grad_atol is None else input_grad_atol, + tensor_rtol=rtol if input_grad_rtol is None else input_grad_rtol, ) if has_bias: - torch.testing.assert_close( - bias_hidden_states_baseline, - bias_hidden_states_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in bias_hidden_states: {msg}", + assert_close_or_cosine_similarity( + bias_hidden_states_baseline, bias_hidden_states_parallel, "bias_hidden_states" ) + if compare_param_grads: + assert param_grads_baseline.keys() == param_grads_parallel.keys() + for name, grad_baseline in param_grads_baseline.items(): + assert_close_or_cosine_similarity( + grad_baseline, param_grads_parallel[name], f"param_grad[{name}]" + ) Utils.destroy_model_parallel() +# TODO(yuzhongw): Add test case for fallback_to_eager_attn @pytest.mark.parametrize("sequence_packing", [False, True]) @pytest.mark.parametrize("apply_rope_fusion", [False, True]) @pytest.mark.parametrize( diff --git a/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py b/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py index 0655e865c86..e591d7cd2d8 100644 --- a/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py +++ b/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py @@ -29,13 +29,10 @@ """ import re -from types import SimpleNamespace import pytest import torch -from megatron.core.dist_checkpointing.core import CheckpointingException -from megatron.core.dist_checkpointing.validation import StrictHandling from megatron.core.distributed.fsdp.src.megatron_fsdp.utils import ( get_mcore_tensor_parallel_partition_dim, is_mcore_tensor_model_parallel, @@ -45,80 +42,12 @@ copy_tensor_model_parallel_attributes, set_tensor_model_parallel_attributes, ) -from megatron.core.transformer import fsdp_dtensor_checkpoint from megatron.core.transformer.fsdp_dtensor_checkpoint import ( - MLA_UNFUSED_DOWN_PROJS, - _intersect_slice, - _shift_slice, - _strip_wrapper_prefixes, - absorbed_input_layernorm_key, flatten_state_dict, get_expert_index_from_key, - get_mla_fused_down_proj_splits, - get_mtp_inner_layer_paths, - get_unexpected_model_keys, - handle_mla_down_proj_in_state_dict, - handle_mtp_in_state_dict, - match_mla_fused_down_proj_key, - rename_mtp_inner_layer_keys, - validate_fsdp_dtensor_model_load, ) -# ============================================================================ -# Test the slice/path helpers shared by the SwiGLU, GDN and MLA handlers -# ============================================================================ -class TestSharedSliceHelpers: - """These three are used by every fused-parameter handler, so pin their semantics.""" - - def test_intersect_overlapping(self): - assert _intersect_slice(slice(0, 10), slice(4, 20)) == slice(4, 10) - - def test_intersect_contained(self): - assert _intersect_slice(slice(0, 10), slice(2, 5)) == slice(2, 5) - - @pytest.mark.parametrize( - "s1,s2", - [ - (slice(0, 4), slice(4, 8)), # touching but not overlapping - (slice(0, 4), slice(6, 8)), # disjoint - (slice(6, 8), slice(0, 4)), # disjoint, reversed - ], - ) - def test_intersect_without_overlap_is_empty(self, s1, s2): - """An FSDP shard that holds none of a section must yield a zero-length slice, not a - negative-length one, otherwise the section's reshape silently takes the wrong data.""" - result = _intersect_slice(s1, s2) - assert result.stop - result.start == 0 - assert list(range(10))[result] == [] - - def test_shift_slice(self): - assert _shift_slice(slice(6, 10), -4) == slice(2, 6) - - def test_shift_rebases_onto_local_storage(self): - """The handlers shift by -fsdp_slice.start to index into this rank's local shard.""" - fsdp_slice = slice(8, 16) - section = _intersect_slice(fsdp_slice, slice(12, 20)) - assert _shift_slice(section, -fsdp_slice.start) == slice(4, 8) - - @pytest.mark.parametrize( - "path,expected", - [ - ("module.decoder.layers.0.weight", "decoder.layers.0.weight"), - ("model.module.decoder.layers.0.weight", "decoder.layers.0.weight"), - ("module.module.decoder.weight", "decoder.weight"), - ("decoder.layers.0.weight", "decoder.layers.0.weight"), - ("", ""), - ], - ) - def test_strip_wrapper_prefixes(self, path, expected): - assert _strip_wrapper_prefixes(path) == expected - - def test_strip_keeps_inner_module_segments(self): - """Only leading wrapper segments go; 'module' deeper in the path is a real submodule.""" - assert _strip_wrapper_prefixes("module.decoder.module.weight") == "decoder.module.weight" - - # ============================================================================ # Test get_expert_index_from_key # ============================================================================ @@ -730,633 +659,6 @@ def test_vlm_scenario(self): assert all("vision_encoder" in k for k in should_skip) -# ============================================================================ -# Fixtures for the fused-MLA / MTP handlers -# -# These build real nn.Modules so that named_modules() behaves exactly as it does -# in production; only the leaf parameters are placeholders. -# ============================================================================ -Q_LORA_RANK = 24 -KV_LORA_RANK = 16 -QK_POS_EMB_HEAD_DIM = 8 -MLA_SPLIT_SIZES = [Q_LORA_RANK, KV_LORA_RANK + QK_POS_EMB_HEAD_DIM] - - -def _mla_config(): - return SimpleNamespace( - q_lora_rank=Q_LORA_RANK, kv_lora_rank=KV_LORA_RANK, qk_pos_emb_head_dim=QK_POS_EMB_HEAD_DIM - ) - - -class _FusedMLAAttention(torch.nn.Module): - """Stands in for FusedMLASelfAttention: one fused down-proj plus a config.""" - - def __init__(self, absorb_layernorm=False): - super().__init__() - self.config = _mla_config() - self.linear_qkv_down_proj = torch.nn.Module() - self.linear_qkv_down_proj.weight = torch.nn.Parameter(torch.zeros(sum(MLA_SPLIT_SIZES), 4)) - if absorb_layernorm: - self.linear_qkv_down_proj.layer_norm_weight = torch.nn.Parameter(torch.zeros(4)) - - -class _UnfusedMLAAttention(torch.nn.Module): - """Stands in for MLASelfAttention: separate q/kv down-projections.""" - - def __init__(self): - super().__init__() - self.config = _mla_config() - self.linear_q_down_proj = torch.nn.Module() - self.linear_kv_down_proj = torch.nn.Module() - - -class _Layer(torch.nn.Module): - def __init__(self, attention): - super().__init__() - self.self_attention = attention - self.input_layernorm = torch.nn.Module() - - -class _MTPLayer(torch.nn.Module): - """Stands in for MultiTokenPredictionLayer.""" - - def __init__(self, mtp_layer_pattern=None): - super().__init__() - self.mtp_layer_pattern = mtp_layer_pattern - self.mtp_model_layer = torch.nn.Module() - - -class _Decoder(torch.nn.Module): - def __init__(self, layers): - super().__init__() - self.layers = torch.nn.ModuleList(layers) - - -class _Model(torch.nn.Module): - """Wrapper mirroring the ``module.`` indirection the handlers expect.""" - - def __init__(self, decoder_layers=(), mtp_layers=()): - super().__init__() - inner = torch.nn.Module() - inner.decoder = _Decoder(list(decoder_layers)) - if mtp_layers: - inner.mtp = _Decoder(list(mtp_layers)) - self.module = inner - - def get_parameter(self, name): - # The handlers only need an object carrying Megatron-FSDP shard metadata; the - # split itself is exercised separately. - return SimpleNamespace(name=name) - - -def _fused_model(num_layers=1, absorb_layernorm=False, with_mtp=False): - layers = [_Layer(_FusedMLAAttention(absorb_layernorm)) for _ in range(num_layers)] - mtp = [_MTPLayer()] if with_mtp else [] - return _Model(decoder_layers=layers, mtp_layers=mtp) - - -# ============================================================================ -# Test get_mla_fused_down_proj_splits -# ============================================================================ -class TestGetMLAFusedDownProjSplits: - """The fused down-proj split sizes must match FusedMLASelfAttention.sharded_state_dict.""" - - def test_fused_model_reports_q_and_kv_sizes(self): - splits = get_mla_fused_down_proj_splits(_fused_model()) - assert splits == {"decoder.layers.0.self_attention": MLA_SPLIT_SIZES} - - def test_kv_section_includes_pos_emb_head_dim(self): - """kv is kv_lora_rank + qk_pos_emb_head_dim, not kv_lora_rank alone.""" - splits = get_mla_fused_down_proj_splits(_fused_model()) - _, kv = splits["decoder.layers.0.self_attention"] - assert kv == KV_LORA_RANK + QK_POS_EMB_HEAD_DIM - - def test_split_sizes_sum_to_fused_width(self): - model = _fused_model() - splits = get_mla_fused_down_proj_splits(model) - fused_weight = model.module.decoder.layers[0].self_attention.linear_qkv_down_proj.weight - assert sum(splits["decoder.layers.0.self_attention"]) == fused_weight.shape[0] - - def test_unfused_model_reports_nothing(self): - model = _Model(decoder_layers=[_Layer(_UnfusedMLAAttention())]) - assert get_mla_fused_down_proj_splits(model) == {} - - def test_multiple_layers(self): - splits = get_mla_fused_down_proj_splits(_fused_model(num_layers=3)) - assert set(splits) == { - "decoder.layers.0.self_attention", - "decoder.layers.1.self_attention", - "decoder.layers.2.self_attention", - } - - def test_mtp_attention_included(self): - """The MTP layer's attention is fused too and must be rewritten alongside the decoder.""" - model = _Model(decoder_layers=[_Layer(_FusedMLAAttention())], mtp_layers=[_MTPLayer()]) - model.module.mtp.layers[0].mtp_model_layer = _Layer(_FusedMLAAttention()) - splits = get_mla_fused_down_proj_splits(model) - assert "mtp.layers.0.mtp_model_layer.self_attention" in splits - - -# ============================================================================ -# Test match_mla_fused_down_proj_key -# ============================================================================ -class TestMatchMLAFusedDownProjKey: - - SPLITS = {"decoder.layers.0.self_attention": MLA_SPLIT_SIZES} - - def test_weight_key_matched(self): - key = "decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - wrapper, path, sizes, leaf = match_mla_fused_down_proj_key(key, self.SPLITS) - assert wrapper == "" - assert path == "decoder.layers.0.self_attention" - assert sizes == MLA_SPLIT_SIZES - assert leaf == "weight" - - def test_wrapper_prefix_preserved(self): - key = "module.decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - wrapper, path, _, leaf = match_mla_fused_down_proj_key(key, self.SPLITS) - assert wrapper == "module." - assert path == "decoder.layers.0.self_attention" - assert leaf == "weight" - - def test_nested_wrapper_prefix_preserved(self): - key = "model.module.decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - wrapper, _, _, _ = match_mla_fused_down_proj_key(key, self.SPLITS) - assert wrapper == "model.module." - - def test_layer_norm_leaf_matched(self): - key = "decoder.layers.0.self_attention.linear_qkv_down_proj.layer_norm_weight" - _, _, _, leaf = match_mla_fused_down_proj_key(key, self.SPLITS) - assert leaf == "layer_norm_weight" - - def test_extra_state_leaf_matched(self): - key = "decoder.layers.0.self_attention.linear_qkv_down_proj._extra_state" - _, _, _, leaf = match_mla_fused_down_proj_key(key, self.SPLITS) - assert leaf == "_extra_state" - - def test_already_unfused_key_not_matched(self): - """An unfused checkpoint-style key must not be rewritten a second time.""" - key = "decoder.layers.0.self_attention.linear_q_down_proj.weight" - assert match_mla_fused_down_proj_key(key, self.SPLITS) is None - - def test_unrelated_key_not_matched(self): - key = "decoder.layers.0.mlp.linear_fc1.weight" - assert match_mla_fused_down_proj_key(key, self.SPLITS) is None - - def test_other_layer_not_matched(self): - key = "decoder.layers.1.self_attention.linear_qkv_down_proj.weight" - assert match_mla_fused_down_proj_key(key, self.SPLITS) is None - - def test_empty_splits_match_nothing(self): - key = "decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - assert match_mla_fused_down_proj_key(key, {}) is None - - -# ============================================================================ -# Test absorbed_input_layernorm_key -# ============================================================================ -class TestAbsorbedInputLayernormKey: - """With fuse_input_layernorm the layernorm moves up one level, per the layer spec's - sharded_state_dict_keys_map: self_attention.linear_qkv_down_proj.layer_norm_ -> - input_layernorm.""" - - def test_weight(self): - assert ( - absorbed_input_layernorm_key("", "decoder.layers.0.self_attention", "layer_norm_weight") - == "decoder.layers.0.input_layernorm.weight" - ) - - def test_bias(self): - assert ( - absorbed_input_layernorm_key("", "decoder.layers.3.self_attention", "layer_norm_bias") - == "decoder.layers.3.input_layernorm.bias" - ) - - def test_wrapper_preserved(self): - assert ( - absorbed_input_layernorm_key( - "module.", "decoder.layers.0.self_attention", "layer_norm_weight" - ) - == "module.decoder.layers.0.input_layernorm.weight" - ) - - def test_top_level_attention_rejected(self): - with pytest.raises(AssertionError): - absorbed_input_layernorm_key("", "self_attention", "layer_norm_weight") - - -# ============================================================================ -# Test handle_mla_down_proj_in_state_dict -# -# split_fused_fsdp_param needs real Megatron-FSDP DTensors, so it is replaced with a -# recorder; everything around it (key rewriting, split sizes, optimizer state) is real. -# ============================================================================ -@pytest.fixture -def recorded_splits(monkeypatch): - calls = [] - - def fake_split(data, dist_param, split_sizes, is_expert_param=False, split_dim=0): - calls.append({"data": data, "split_sizes": list(split_sizes)}) - return [f"{data}::q", f"{data}::kv"] - - monkeypatch.setattr(fsdp_dtensor_checkpoint, "split_fused_fsdp_param", fake_split) - return calls - - -class TestHandleMLADownProjInStateDict: - - WEIGHT_KEY = "decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - Q_KEY = "decoder.layers.0.self_attention.linear_q_down_proj.weight" - KV_KEY = "decoder.layers.0.self_attention.linear_kv_down_proj.weight" - - def test_fused_weight_becomes_q_and_kv(self, recorded_splits): - model_sd, _ = handle_mla_down_proj_in_state_dict( - _fused_model(), {self.WEIGHT_KEY: "W"}, None - ) - assert set(model_sd) == {self.Q_KEY, self.KV_KEY} - assert model_sd[self.Q_KEY] == "W::q" - assert model_sd[self.KV_KEY] == "W::kv" - - def test_unfused_projection_names_match_mcore(self): - assert MLA_UNFUSED_DOWN_PROJS == ("linear_q_down_proj", "linear_kv_down_proj") - - def test_config_split_sizes_are_forwarded(self, recorded_splits): - handle_mla_down_proj_in_state_dict(_fused_model(), {self.WEIGHT_KEY: "W"}, None) - assert recorded_splits[0]["split_sizes"] == MLA_SPLIT_SIZES - - def test_absorbed_layernorm_is_moved_not_split(self, recorded_splits): - ln_key = "decoder.layers.0.self_attention.linear_qkv_down_proj.layer_norm_weight" - model_sd, _ = handle_mla_down_proj_in_state_dict( - _fused_model(absorb_layernorm=True), {ln_key: "LN"}, None - ) - assert model_sd == {"decoder.layers.0.input_layernorm.weight": "LN"} - assert recorded_splits == [] - - def test_extra_state_left_alone(self, recorded_splits): - key = "decoder.layers.0.self_attention.linear_qkv_down_proj._extra_state" - model_sd, _ = handle_mla_down_proj_in_state_dict(_fused_model(), {key: "X"}, None) - assert model_sd == {key: "X"} - - def test_unrelated_keys_untouched(self, recorded_splits): - other = "decoder.layers.0.mlp.linear_fc1.weight" - model_sd, _ = handle_mla_down_proj_in_state_dict( - _fused_model(), {self.WEIGHT_KEY: "W", other: "M"}, None - ) - assert model_sd[other] == "M" - - def test_unfused_model_is_noop(self, recorded_splits): - model = _Model(decoder_layers=[_Layer(_UnfusedMLAAttention())]) - original = {self.Q_KEY: "Q", self.KV_KEY: "KV"} - model_sd, opt_sd = handle_mla_down_proj_in_state_dict(model, original, None) - assert model_sd == original - assert opt_sd is None - assert recorded_splits == [] - - def test_input_state_dict_not_mutated(self, recorded_splits): - original = {self.WEIGHT_KEY: "W"} - handle_mla_down_proj_in_state_dict(_fused_model(), original, None) - assert original == {self.WEIGHT_KEY: "W"} - - def test_all_layers_rewritten(self, recorded_splits): - model_sd, _ = handle_mla_down_proj_in_state_dict( - _fused_model(num_layers=3), - { - f"decoder.layers.{i}.self_attention.linear_qkv_down_proj.weight": f"W{i}" - for i in range(3) - }, - None, - ) - assert len(model_sd) == 6 - assert all("linear_qkv_down_proj" not in k for k in model_sd) - - def test_optimizer_state_split(self, recorded_splits): - opt_key = f"module.{self.WEIGHT_KEY}" - optimizer_state_dict = {"state": {opt_key: {"exp_avg": "A", "exp_avg_sq": "B", "step": 7}}} - _, opt_sd = handle_mla_down_proj_in_state_dict(_fused_model(), {}, optimizer_state_dict) - state = opt_sd["state"] - assert set(state) == {f"module.{self.Q_KEY}", f"module.{self.KV_KEY}"} - assert state[f"module.{self.Q_KEY}"]["exp_avg"] == "A::q" - assert state[f"module.{self.KV_KEY}"]["exp_avg"] == "A::kv" - assert state[f"module.{self.Q_KEY}"]["exp_avg_sq"] == "B::q" - # Non-tensor entries are carried over to both halves. - assert state[f"module.{self.Q_KEY}"]["step"] == 7 - - def test_optimizer_layernorm_moved_not_split(self, recorded_splits): - """Regression: the absorbed layernorm has no q/kv halves to split into.""" - ln_key = "module.decoder.layers.0.self_attention.linear_qkv_down_proj.layer_norm_weight" - _, opt_sd = handle_mla_down_proj_in_state_dict( - _fused_model(absorb_layernorm=True), {}, {"state": {ln_key: {"exp_avg": "A"}}} - ) - assert set(opt_sd["state"]) == {"module.decoder.layers.0.input_layernorm.weight"} - assert recorded_splits == [] - - def test_optimizer_unrelated_state_preserved(self, recorded_splits): - other = "module.decoder.layers.0.mlp.linear_fc1.weight" - _, opt_sd = handle_mla_down_proj_in_state_dict( - _fused_model(), {}, {"state": {other: {"exp_avg": "A"}}} - ) - assert opt_sd["state"][other] == {"exp_avg": "A"} - - def test_empty_optimizer_state_tolerated(self, recorded_splits): - _, opt_sd = handle_mla_down_proj_in_state_dict(_fused_model(), {}, {"state": {}}) - assert opt_sd["state"] == {} - - -# ============================================================================ -# Test get_mtp_inner_layer_paths -# ============================================================================ -class TestGetMTPInnerLayerPaths: - - def test_gpt_mtp_layer_included(self): - model = _Model(mtp_layers=[_MTPLayer(mtp_layer_pattern=None)]) - assert get_mtp_inner_layer_paths(model) == ["mtp.layers.0"] - - def test_mamba_mtp_layer_excluded(self): - """Mamba MTP keeps mtp_model_layer as its native checkpoint name.""" - model = _Model(mtp_layers=[_MTPLayer(mtp_layer_pattern="M-M-")]) - assert get_mtp_inner_layer_paths(model) == [] - - def test_model_without_mtp(self): - model = _Model(decoder_layers=[_Layer(_UnfusedMLAAttention())]) - assert get_mtp_inner_layer_paths(model) == [] - - def test_multiple_mtp_layers(self): - model = _Model(mtp_layers=[_MTPLayer(), _MTPLayer()]) - assert get_mtp_inner_layer_paths(model) == ["mtp.layers.0", "mtp.layers.1"] - - -# ============================================================================ -# Test rename_mtp_inner_layer_keys -# ============================================================================ -class TestRenameMTPInnerLayerKeys: - """The model uses mtp_model_layer; checkpoints use transformer_layer.""" - - PATHS = ["mtp.layers.0"] - - def test_key_renamed(self): - sd = {"mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight": 1} - out = rename_mtp_inner_layer_keys(sd, self.PATHS) - assert out == {"mtp.layers.0.transformer_layer.mlp.linear_fc1.weight": 1} - - def test_wrapper_prefix_preserved(self): - sd = {"module.mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight": 1} - out = rename_mtp_inner_layer_keys(sd, self.PATHS) - assert list(out) == ["module.mtp.layers.0.transformer_layer.mlp.linear_fc1.weight"] - - def test_expert_key_renamed(self): - """The MTP layer's expert weights were the bulk of the silently dropped keys.""" - sd = {"mtp.layers.0.mtp_model_layer.mlp.experts.linear_fc1.weight3": 1} - out = rename_mtp_inner_layer_keys(sd, self.PATHS) - assert list(out) == ["mtp.layers.0.transformer_layer.mlp.experts.linear_fc1.weight3"] - - def test_decoder_key_untouched(self): - sd = {"decoder.layers.0.mlp.linear_fc1.weight": 1} - assert rename_mtp_inner_layer_keys(sd, self.PATHS) == sd - - def test_mtp_key_outside_inner_layer_untouched(self): - sd = {"mtp.layers.0.enorm.weight": 1} - assert rename_mtp_inner_layer_keys(sd, self.PATHS) == sd - - def test_only_inner_layer_occurrence_renamed(self): - sd = {"mtp.layers.0.mtp_model_layer.sub.mtp_model_layer.weight": 1} - out = rename_mtp_inner_layer_keys(sd, self.PATHS) - assert list(out) == ["mtp.layers.0.transformer_layer.sub.mtp_model_layer.weight"] - - def test_unlisted_mtp_layer_untouched(self): - sd = {"mtp.layers.1.mtp_model_layer.mlp.linear_fc1.weight": 1} - assert rename_mtp_inner_layer_keys(sd, self.PATHS) == sd - - def test_input_not_mutated(self): - sd = {"mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight": 1} - rename_mtp_inner_layer_keys(sd, self.PATHS) - assert list(sd) == ["mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight"] - - -# ============================================================================ -# Test handle_mtp_in_state_dict -# ============================================================================ -class TestHandleMTPInStateDict: - - MODEL_KEY = "mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight" - RENAMED_KEY = "mtp.layers.0.transformer_layer.mlp.linear_fc1.weight" - - def test_model_state_renamed(self): - model_sd, _ = handle_mtp_in_state_dict( - _Model(mtp_layers=[_MTPLayer()]), {self.MODEL_KEY: 1}, None - ) - assert model_sd == {self.RENAMED_KEY: 1} - - def test_optimizer_state_renamed(self): - opt_key = f"module.{self.MODEL_KEY}" - _, opt_sd = handle_mtp_in_state_dict( - _Model(mtp_layers=[_MTPLayer()]), {}, {"state": {opt_key: {"exp_avg": "A"}}} - ) - assert list(opt_sd["state"]) == [f"module.{self.RENAMED_KEY}"] - - def test_noop_without_mtp(self): - model = _Model(decoder_layers=[_Layer(_UnfusedMLAAttention())]) - sd = {"decoder.layers.0.mlp.linear_fc1.weight": 1} - model_sd, opt_sd = handle_mtp_in_state_dict(model, sd, None) - assert model_sd == sd - assert opt_sd is None - - def test_noop_for_mamba_mtp(self): - model = _Model(mtp_layers=[_MTPLayer(mtp_layer_pattern="M-M-")]) - sd = {self.MODEL_KEY: 1} - model_sd, _ = handle_mtp_in_state_dict(model, sd, None) - assert model_sd == sd - - def test_empty_optimizer_state_tolerated(self): - _, opt_sd = handle_mtp_in_state_dict(_Model(mtp_layers=[_MTPLayer()]), {}, {"state": {}}) - assert opt_sd["state"] == {} - - -# ============================================================================ -# Test get_unexpected_model_keys -# ============================================================================ -class TestGetUnexpectedModelKeys: - """ "Unexpected" follows StrictHandling: requested by this rank, absent from the - checkpoint. This is the direction that silently zeroed weights.""" - - def test_missing_model_key_reported(self): - metadata = {"model": {"decoder.layers.0.mlp.linear_fc1.weight": object()}} - load = { - "model": { - "decoder.layers.0.mlp.linear_fc1.weight": 1, - "decoder.layers.0.self_attention.linear_qkv_down_proj.weight": 1, - } - } - assert get_unexpected_model_keys(metadata, load) == { - "model.decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - } - - def test_nothing_missing(self): - metadata = {"model": {"a.weight": object(), "b.weight": object()}} - load = {"model": {"a.weight": 1, "b.weight": 1}} - assert get_unexpected_model_keys(metadata, load) == set() - - def test_expert_parallel_extras_not_reported(self): - """Regression: under EP the checkpoint holds every expert but each rank owns a - slice. Reporting checkpoint-side extras would fail every EP run.""" - metadata = { - "model": { - f"decoder.layers.0.mlp.experts.linear_fc1.weight{i}": object() for i in range(8) - } - } - load = { - "model": {f"decoder.layers.0.mlp.experts.linear_fc1.weight{i}": 1 for i in range(2)} - } - assert get_unexpected_model_keys(metadata, load) == set() - - def test_optimizer_section_ignored(self): - """Optimizer state is legitimately absent when it is not being restored.""" - metadata = {"model": {"a.weight": object()}} - load = {"model": {"a.weight": 1}, "optimizer": {"state": {"a.weight": {"exp_avg": 1}}}} - assert get_unexpected_model_keys(metadata, load) == set() - - def test_rng_section_ignored(self): - metadata = {"model": {"a.weight": object()}} - load = {"model": {"a.weight": 1}, "rng_state": [1, 2, 3]} - assert get_unexpected_model_keys(metadata, load) == set() - - def test_extra_state_ignored(self): - """TE _extra_state is dropped from both sides before save and load.""" - metadata = {"model": {"a.weight": object()}} - load = {"model": {"a.weight": 1, "a._extra_state": 1}} - assert get_unexpected_model_keys(metadata, load) == set() - - def test_virtual_pipeline_model_sections_covered(self): - metadata = {"model0": {"a.weight": object()}, "model1": {}} - load = {"model0": {"a.weight": 1}, "model1": {"b.weight": 1}} - assert get_unexpected_model_keys(metadata, load) == {"model1.b.weight"} - - def test_non_model_top_level_key_ignored(self): - metadata = {"model": {}} - load = {"model": {}, "checkpoint_version": 3.0} - assert get_unexpected_model_keys(metadata, load) == set() - - -# ============================================================================ -# Test validate_fsdp_dtensor_model_load -# ============================================================================ -class TestValidateFsdpDtensorModelLoad: - - MISSING_KEY = "decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - - def _mismatched(self): - return {"model": {}}, {"model": {self.MISSING_KEY: 1}} - - def _matched(self): - return {"model": {"a.weight": object()}}, {"model": {"a.weight": 1}} - - def test_raise_unexpected_raises(self): - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException) as excinfo: - validate_fsdp_dtensor_model_load( - metadata, load, "/ckpt", strict=StrictHandling.RAISE_UNEXPECTED - ) - assert self.MISSING_KEY in str(excinfo.value) - assert "/ckpt" in str(excinfo.value) - - def test_error_explains_the_cause(self): - """The message should point at sharded_state_dict(), which is the actual trap.""" - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException, match="sharded_state_dict"): - validate_fsdp_dtensor_model_load(metadata, load, "/ckpt") - - def test_raises_by_default(self): - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException): - validate_fsdp_dtensor_model_load(metadata, load, "/ckpt") - - def test_no_raise_when_complete(self): - metadata, load = self._matched() - assert validate_fsdp_dtensor_model_load(metadata, load, "/ckpt") == set() - - def test_raise_all_behaves_like_raise_unexpected(self): - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException): - validate_fsdp_dtensor_model_load( - metadata, load, "/ckpt", strict=StrictHandling.RAISE_ALL - ) - - @pytest.mark.parametrize("strict", [StrictHandling.LOG_UNEXPECTED, StrictHandling.LOG_ALL]) - def test_log_variants_do_not_raise(self, strict, caplog): - metadata, load = self._mismatched() - with caplog.at_level("WARNING"): - result = validate_fsdp_dtensor_model_load(metadata, load, "/ckpt", strict=strict) - assert result == {f"model.{self.MISSING_KEY}"} - assert self.MISSING_KEY in caplog.text - - def test_ignore_all_disables_the_check(self): - metadata, load = self._mismatched() - assert ( - validate_fsdp_dtensor_model_load( - metadata, load, "/ckpt", strict=StrictHandling.IGNORE_ALL - ) - == set() - ) - - def test_assume_ok_unexpected_raises(self): - """ASSUME_OK_UNEXPECTED defers to the underlying strategy, which never raises on a - partial DCP load, so for fsdp_dtensor it has to raise here instead. This is the - default value of --dist-ckpt-strictness, so the check must be on out of the box.""" - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException): - validate_fsdp_dtensor_model_load( - metadata, load, "/ckpt", strict=StrictHandling.ASSUME_OK_UNEXPECTED - ) - - def test_assume_ok_unexpected_quiet_when_complete(self): - metadata, load = self._matched() - assert ( - validate_fsdp_dtensor_model_load( - metadata, load, "/ckpt", strict=StrictHandling.ASSUME_OK_UNEXPECTED - ) - == set() - ) - - @pytest.mark.parametrize( - "strict", [StrictHandling.RETURN_UNEXPECTED, StrictHandling.RETURN_ALL] - ) - def test_return_variants_report_without_raising(self, strict): - metadata, load = self._mismatched() - assert validate_fsdp_dtensor_model_load(metadata, load, "/ckpt", strict=strict) == { - f"model.{self.MISSING_KEY}" - } - - def test_string_strict_flag_accepted(self): - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException): - validate_fsdp_dtensor_model_load(metadata, load, "/ckpt", strict="raise_unexpected") - - def test_invalid_strict_flag_rejected(self): - metadata, load = self._matched() - with pytest.raises(ValueError): - validate_fsdp_dtensor_model_load(metadata, load, "/ckpt", strict="nonsense") - - def test_expert_parallel_load_does_not_raise(self): - """Regression: the validator must stay quiet on a healthy expert-parallel load.""" - metadata = { - "model": { - f"decoder.layers.0.mlp.experts.linear_fc1.weight{i}": object() for i in range(8) - } - } - load = { - "model": {f"decoder.layers.0.mlp.experts.linear_fc1.weight{i}": 1 for i in range(2)} - } - assert validate_fsdp_dtensor_model_load(metadata, load, "/ckpt") == set() - - def test_message_truncated(self): - metadata = {"model": {}} - load = {"model": {f"decoder.layers.{i}.weight": 1 for i in range(50)}} - with pytest.raises(CheckpointingException) as excinfo: - validate_fsdp_dtensor_model_load(metadata, load, "/ckpt", max_reported=5) - message = str(excinfo.value) - assert "and 45 more" in message - assert message.count("\n model.decoder.layers.") == 5 - - # ============================================================================ # Test checkpoint_inspector.py pretrained-only checkpoint handling # (hasattr mcore_data guard — part of PR #3912 but related) diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index 3131222f800..f4c07430400 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/tests/unit_tests/transformer/test_multi_latent_attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import os from inspect import signature @@ -741,6 +741,189 @@ def test_gpu_forward_thd(self): assert bias.shape[0] == config.hidden_size +@pytest.mark.parametrize("gate_granularity", ("elementwise", "headwise")) +class TestMLAOutputGate: + + @pytest.fixture(scope='function', autouse=True) + def setup_and_teardown(self, gate_granularity): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + self.transformer_config = MLATransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + use_cpu_initialization=True, + q_lora_rank=32, + kv_lora_rank=32, + qk_head_dim=128, + v_head_dim=128, + qk_pos_emb_head_dim=64, + rope_type="yarn", + rotary_base=10000, + original_max_position_embeddings=32, + attention_output_gate=True, + gated_attention_proj_granularity=gate_granularity, + ) + self.parallel_attention = MLASelfAttention( + self.transformer_config, + get_mla_self_attn_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + ) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_constructor_builds_configured_gate_projection(self, gate_granularity): + gate = self.parallel_attention.linear_gate + assert gate is not None + expected_projection_size = ( + self.parallel_attention.query_projection_size + if gate_granularity == "elementwise" + else self.transformer_config.num_attention_heads + ) + assert gate.weight.shape == (expected_projection_size, self.transformer_config.hidden_size) + assert gate.bias is None or gate.bias.numel() == 0 + + def test_gate_matches_independent_fp32_sigmoid_reference(self, gate_granularity): + output_size = self.parallel_attention.query_projection_size + core_attn_out = torch.linspace(-2.0, 2.0, 2 * output_size, dtype=torch.bfloat16).view( + 2, 1, output_size + ) + gate_size = ( + output_size + if gate_granularity == "elementwise" + else self.transformer_config.num_attention_heads + ) + gate = torch.linspace(-4.0, 4.0, 2 * gate_size, dtype=torch.bfloat16).view(2, 1, gate_size) + + output = self.parallel_attention._apply_mla_output_gate(core_attn_out, gate) + reference_scale = torch.sigmoid(gate.float()).to(core_attn_out.dtype) + if gate_granularity == "headwise": + reference_scale = torch.repeat_interleave( + reference_scale, self.transformer_config.v_head_dim, dim=-1 + ) + expected = core_attn_out * reference_scale + + assert output.dtype == core_attn_out.dtype + torch.testing.assert_close(output, expected, atol=0, rtol=0) + + def test_gate_matches_fp32_sigmoid_reference_vjp(self, gate_granularity): + generator = torch.Generator().manual_seed(20260814) + output_size = self.parallel_attention.query_projection_size + gate_size = ( + output_size + if gate_granularity == "elementwise" + else self.transformer_config.num_attention_heads + ) + gate_input = torch.randn((8, 1, gate_size), generator=generator, dtype=torch.bfloat16) + core_output = torch.randn((8, 1, output_size), generator=generator, dtype=torch.bfloat16) + output_gradient = torch.randn(core_output.shape, generator=generator, dtype=torch.bfloat16) + + actual_gate = gate_input.detach().clone().requires_grad_(True) + actual_core = core_output.detach().clone().requires_grad_(True) + actual_output = self.parallel_attention._apply_mla_output_gate(actual_core, actual_gate) + actual_gradients = torch.autograd.grad( + actual_output, (actual_gate, actual_core), output_gradient + ) + + sigmoid_fp32 = torch.sigmoid(gate_input.float()) + native_scale = sigmoid_fp32.to(core_output.dtype) + if gate_granularity == "headwise": + native_scale = torch.repeat_interleave( + native_scale, self.transformer_config.v_head_dim, dim=-1 + ) + expected_core_gradient = output_gradient * native_scale + head_shape = (*core_output.shape[:-1], gate_size, self.transformer_config.v_head_dim) + scale_gradient = ( + output_gradient.float().view(head_shape) * core_output.float().view(head_shape) + ).sum(dim=-1) + # The fused broadcast reduction returns a native-dtype gradient to the + # FP32-sigmoid cast before applying the sigmoid derivative. + scale_gradient = scale_gradient.to(gate_input.dtype).float() + else: + expected_core_gradient = output_gradient * native_scale + scale_gradient = (output_gradient * core_output).float() + + expected_output = core_output * native_scale + expected_gate_gradient = (scale_gradient * sigmoid_fp32 * (1.0 - sigmoid_fp32)).to( + gate_input.dtype + ) + + torch.testing.assert_close(actual_output, expected_output, atol=0, rtol=0) + torch.testing.assert_close(actual_gradients[0], expected_gate_gradient) + torch.testing.assert_close(actual_gradients[1], expected_core_gradient) + + def test_missing_gate_spec_raises(self): + submodules = get_mla_self_attn_submodules() + submodules.linear_gate = None + with pytest.raises(ValueError, match="linear_gate module spec"): + MLASelfAttention( + self.transformer_config, + submodules, + layer_number=1, + attn_mask_type=AttnMaskType.causal, + ) + + def test_fused_down_projection_is_rejected(self, gate_granularity): + with pytest.raises(ValueError, match="does not support fused down projections"): + MLATransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + q_lora_rank=32, + kv_lora_rank=32, + qk_head_dim=128, + v_head_dim=128, + qk_pos_emb_head_dim=64, + rope_type=self.transformer_config.rope_type, + rotary_base=10000, + original_max_position_embeddings=32, + attention_output_gate=True, + gated_attention_proj_granularity=gate_granularity, + mla_down_proj_fusion=True, + ) + + def test_gpu_forward_thd(self): + if not is_te_min_version("1.10.0"): + pytest.skip("MLA requires TransformerEngine >= 1.10.0") + + attention = self.parallel_attention.cuda().bfloat16() + call_order = [] + core_attention_hook = attention.core_attention.register_forward_pre_hook( + lambda *_: call_order.append("core_attention") + ) + gate_projection_hook = attention.linear_gate.register_forward_pre_hook( + lambda *_: call_order.append("linear_gate") + ) + sequence_length = 32 + hidden_states = torch.ones( + (sequence_length, 1, self.transformer_config.hidden_size), + device='cuda', + dtype=torch.bfloat16, + requires_grad=True, + ) + packed_seq_params = make_test_packed_seq_params(sequence_length=sequence_length) + + try: + with mock.patch.dict( + os.environ, {"NVTE_FUSED_ATTN": "1", "NVTE_FLASH_ATTN": "0"}, clear=False + ): + output, bias = attention( + hidden_states, attention_mask=None, packed_seq_params=packed_seq_params + ) + output.sum().backward() + finally: + core_attention_hook.remove() + gate_projection_hook.remove() + + assert call_order == ["core_attention", "linear_gate"] + assert output.shape == hidden_states.shape + assert bias.shape == (self.transformer_config.hidden_size,) + assert hidden_states.grad is not None + assert attention.linear_gate.weight.grad is not None + + @pytest.mark.parametrize("rope_type", ('yarn', 'rope')) class TestParallelMLAAttentionPrecision: @@ -1438,19 +1621,25 @@ def test_clip_qk_with_absorption_raises_error(self): ], ) @pytest.mark.parametrize( - ("tp", "sp", "cp"), + ("tp", "sp", "cp", "output_gate", "gate_granularity"), [ - (4, False, 1), # TP w/o SP - (4, True, 1), # TP w/ SP - (1, False, 4), # CP - (2, False, 2), # CP + TP w/o SP - (2, True, 2), # CP + TP w/ SP + (4, False, 1, False, "elementwise"), # TP w/o SP + (4, True, 1, False, "elementwise"), # TP w/ SP + (1, False, 4, False, "elementwise"), # CP + (2, False, 2, False, "elementwise"), # CP + TP w/o SP + (2, True, 2, False, "elementwise"), # CP + TP w/ SP + (4, True, 1, True, "elementwise"), # Elementwise gate with TP + SP + (4, True, 1, True, "headwise"), # Headwise gate with TP + SP + (1, False, 2, True, "elementwise"), # Elementwise gate with CP + (1, False, 2, True, "headwise"), # Headwise gate with CP ], ) @pytest.mark.skipif(not is_te_min_version("1.10.0"), reason="Requires TransformerEngine >= 1.10.0") def test_parallel_multi_latent_attention_correctness( - tmp_path_dist_ckpt, rope_type, apply_rope_fusion, tp, sp, cp + tmp_path_dist_ckpt, rope_type, apply_rope_fusion, tp, sp, cp, output_gate, gate_granularity ): + if output_gate and (rope_type != "yarn" or apply_rope_fusion): + pytest.skip("Gated MLA parallel coverage uses one representative YARN configuration.") if cp > 1 and not is_te_min_version("2.5.0", check_equality=True): pytest.skip("MLA CP requires TransformerEngine >= 2.5.0") if rope_type == "yarn" and apply_rope_fusion and not is_torch_min_version("2.5.0"): @@ -1524,6 +1713,8 @@ def initialize_gpt_model( bf16=True, rope_type=rope_type, apply_rope_fusion=apply_rope_fusion, + attention_output_gate=output_gate, + gated_attention_proj_granularity=gate_granularity, hidden_dropout=0.0, attention_dropout=0.0, )