From 7de59f570484a4767a3311ea90416fce732bd582 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Thu, 27 Aug 2026 03:52:01 +0000 Subject: [PATCH 1/4] cp: add Ling-V3 Tiny hybrid model training support Signed-off-by: svcnemo-autobot --- .../inference/contexts/dynamic_context.py | 7 +- ...rimental_attention_variant_module_specs.py | 49 +- megatron/core/models/gpt/gpt_layer_specs.py | 3 + megatron/core/models/hybrid/hybrid_block.py | 25 + .../models/hybrid/hybrid_layer_allocation.py | 30 +- .../core/models/hybrid/hybrid_layer_specs.py | 24 +- megatron/core/optimizer/qk_clip.py | 71 +- megatron/core/ssm/gated_delta_net/__init__.py | 8 + megatron/core/ssm/gated_delta_net/common.py | 82 +- megatron/core/ssm/gated_delta_net/gdn.py | 380 +++++--- megatron/core/ssm/gated_delta_net/kda.py | 681 ++++++++++++++ megatron/core/transformer/attention.py | 5 + .../transformer/fsdp_dtensor_checkpoint.py | 153 +++- megatron/core/transformer/moe/router.py | 31 +- .../transformer/multi_latent_attention.py | 99 +++ .../core/transformer/transformer_config.py | 526 ++++++++++- megatron/training/argument_utils.py | 93 +- megatron/training/arguments.py | 78 +- megatron/training/training.py | 422 +++++++-- .../get_test_results_from_tensorboard_logs.py | 2 + .../test_pretraining_regular_pipeline.py | 2 + .../model_config.yaml | 128 +++ tests/test_utils/recipes/h100/mamba.yaml | 7 + .../mfsdp_v1/test_gdn_fsdp_checkpoint.py | 176 ++++ .../contexts/test_dynamic_context.py | 16 + ...rimental_attention_variant_module_specs.py | 47 + .../models/test_hybrid_moe_model.py | 11 + tests/unit_tests/optimizer/test_qk_clip.py | 98 ++ .../ssm/gated_delta_net/test_gdn.py | 841 ++++++++++++++++++ .../ssm/gated_delta_net/test_kda.py | 424 +++++++++ tests/unit_tests/ssm/test_hybrid_block.py | 17 +- .../ssm/test_hybrid_layer_allocation.py | 78 +- tests/unit_tests/test_argument_utils.py | 18 + .../test_num_floating_point_operations.py | 178 +++- .../transformer/moe/test_routers.py | 22 + .../unit_tests/transformer/test_attention.py | 10 +- .../test_fsdp_dtensor_checkpoint.py | 89 +- .../test_multi_latent_attention.py | 207 ++++- 38 files changed, 4716 insertions(+), 422 deletions(-) create mode 100644 megatron/core/ssm/gated_delta_net/kda.py create mode 100644 tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp1_ep8_cp2_kda_gated_mla_1N8G/model_config.yaml create mode 100644 tests/unit_tests/distributed/mfsdp_v1/test_gdn_fsdp_checkpoint.py create mode 100644 tests/unit_tests/optimizer/test_qk_clip.py create mode 100644 tests/unit_tests/ssm/gated_delta_net/test_gdn.py create mode 100644 tests/unit_tests/ssm/gated_delta_net/test_kda.py diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index afa071631cb..806b9c11808 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -456,9 +456,14 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Mamba and GDN use the same slot-indexed recurrent-state cache contract. Build # one map in global layer order; independently generated per-symbol maps both # start at zero and would alias if they were simply unioned. + layer_maps = get_layer_maps_from_layer_type_list( + mamba_inference_state_config.layer_type_list + ) 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)) + )(layer_maps) + if layer_maps[Symbols.KDA]: + 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/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index 87c63a126bb..c794c4649fb 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -5,7 +5,13 @@ 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, + GatedDeltaNet2, + GatedDeltaNetSubmodules, + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, +) from megatron.core.transformer.enums import AttnMaskType, LayerType from megatron.core.transformer.experimental_attention_variant.absorbed_mla import ( AbsorbedMLASelfAttention, @@ -58,7 +64,7 @@ ########## # Canonical ``experimental_attention_variant`` names served by the gated delta net family. -GDN_ATTENTION_VARIANTS = ("gdn", "gdn2") +GDN_ATTENTION_VARIANTS = ("gdn", "gdn2", "kda") # Deprecated ``experimental_attention_variant`` spellings mapped to their canonical name. _DEPRECATED_ATTENTION_VARIANT_ALIASES = {"gated_delta_net": "gdn"} @@ -72,25 +78,36 @@ 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: + 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}, + ) return attention diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index a61c36bff5b..38dfcf9ceee 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 5c58bf24bae..772994d450d 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -45,9 +45,13 @@ 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 @@ -183,6 +187,18 @@ def __init__( name=(name + f".layers.{i}") if name is not None else None, ) elif layer_type == LayerSymbols.MLA: + layer = build_module( + submodules.mla_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, + 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.mla_layer, config=self.config, @@ -230,6 +246,15 @@ def __init__( 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") self.layers.append(layer) 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 e8b1ebc9a56..8bb7e878c37 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -15,7 +15,12 @@ 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_net import ( + GatedDeltaNet, + GatedDeltaNetSubmodules, + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, +) from megatron.core.ssm.gated_delta_product import ( GatedDeltaProductMixer, GatedDeltaProductMixerSubmodules, @@ -133,6 +138,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 @@ -200,6 +221,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, 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/ssm/gated_delta_net/__init__.py b/megatron/core/ssm/gated_delta_net/__init__.py index e5a5b4e6a9f..1046cd1ea8b 100644 --- a/megatron/core/ssm/gated_delta_net/__init__.py +++ b/megatron/core/ssm/gated_delta_net/__init__.py @@ -23,13 +23,21 @@ 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", "causal_conv1d", "chunk_gated_delta_rule", "chunk_gdn2", diff --git a/megatron/core/ssm/gated_delta_net/common.py b/megatron/core/ssm/gated_delta_net/common.py index 63f2d244840..af4ecad59ba 100644 --- a/megatron/core/ssm/gated_delta_net/common.py +++ b/megatron/core/ssm/gated_delta_net/common.py @@ -43,10 +43,12 @@ 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 @@ -58,9 +60,7 @@ @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 @@ -124,13 +124,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. @@ -143,6 +144,7 @@ def __init__( 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. + is_mtp_layer: Whether this module is inside an MTP prediction depth. """ if not HAVE_FLA: raise ImportError( @@ -154,6 +156,8 @@ def __init__( # Attributes from arguments self.layer_number = layer_number + self._pp_layer_offset = pp_layer_offset + self.is_mtp_layer = is_mtp_layer self.pp_layer_offset = pp_layer_offset self.bias = bias self.conv_bias = conv_bias @@ -161,7 +165,7 @@ def __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 self.cp_size = self.pg_collection.cp.size() @@ -183,6 +187,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 @@ -197,7 +216,10 @@ def __init__( ) 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" + 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" # 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 @@ -205,7 +227,7 @@ def __init__( 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 +267,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 +277,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 +294,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, @@ -289,9 +317,9 @@ def __init__( 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: + Called once from ``__init__``. Must set: - ``in_proj_dim`` - ``in_proj_split_names`` - ``in_proj_split_sections`` @@ -324,9 +352,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)) @@ -338,6 +364,7 @@ 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]: @@ -389,6 +416,7 @@ 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. @@ -402,11 +430,11 @@ def _prepare_input_for_gated_delta_rule( ``k``, ``v``, ``g``, plus the variant-specific gates), 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 +446,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) @@ -647,7 +675,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. @@ -655,7 +683,8 @@ def get_parameter_local_cp( Args: param (torch.Tensor): The entire parameter to get the local parameter for. dim (int): The dimension to split the parameter along. Usually the dimension of head. - cp_group (torch.distributed.ProcessGroup): The context parallel group. + cp_group (torch.distributed.ProcessGroup | None): The context parallel group, or None + when no headwise sharding is active. split_sections (Optional[list[int]]): If not None, first split the parameter along the dimension dim into sections, then get the local hidden parallel weights separately, @@ -665,6 +694,9 @@ def get_parameter_local_cp( torch.Tensor: The local parameter for the current context parallel rank. """ + if cp_group is None: + return param + cp_size = cp_group.size() cp_rank = cp_group.rank() @@ -842,6 +874,9 @@ def a2a_cp_to_hp( The hidden-parallel tensor and the sequence-dim inverse permutation to hand to :func:`a2a_hp_to_cp` (``None`` outside the ``thd`` + CP>1 case). """ + if cp_size == 1: + return qkvzba, None + if cp_size > 1: # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a. head_perm = _build_head_perm_for_split_sections( @@ -887,6 +922,9 @@ def a2a_hp_to_cp( Returns: The context-parallel tensor, matching the layout of the GDN module input. """ + if cp_size == 1: + return norm_out + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': if cp_size > 1: assert thd_cp_a2a_inv is not None diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py index c7a983841cb..f3d4f1df611 100644 --- a/megatron/core/ssm/gated_delta_net/gdn.py +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -5,22 +5,26 @@ # This source code is licensed under the Apache license found in the # LICENSE file in the root directory of this source tree. -from functools import partial +from dataclasses import replace from typing import Optional import torch import torch.nn.functional as F from megatron.core import tensor_parallel +from megatron.core.context_parallel import ContextParallelLayoutManager 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.jit import jit_fuser from megatron.core.packed_seq_params import PackedSeqParams +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, @@ -38,7 +42,8 @@ class GatedDeltaNet(SSMDynamicInferenceMixin, _GDNBase): - # pylint: disable=missing-class-docstring + """Gated DeltaNet with a head-wise scalar memory-decay gate.""" + def _setup_variant_attrs(self): """Set the GDN in_proj sizing, split tables, gate parameter dims, and kernel.""" # alpha, beta @@ -56,11 +61,8 @@ 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.feat_dim_split = self._get_feat_dim_split( + self.cp_size if self.config.linear_cp_mode == "headwise" else 1 ) self.dt_bias_dim = self.num_v_heads_local_tp @@ -72,6 +74,15 @@ def _setup_variant_attrs(self): 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 GDN projection split sizes for runtime headwise CP.""" + 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( self, @@ -96,20 +107,65 @@ 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. - - Return: - (tuple[torch.Tensor, torch.Tensor]) GDN output and bias. - """ - + """Perform a forward pass through the GDN module.""" + del attention_mask, sequence_len_offset, kwargs 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 + cp_group = ( + packed_seq_params.cp_group + if packed_seq_params is not None and packed_seq_params.cp_group is not None + else active_pg_collection.cp + ) + 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() + + input_layout_state = None + if ( + self.config.linear_cp_mode == "chunkwise" + and self.config.linear_cp_layout == "zigzag" + and cp_size_chunkwise > 1 + ): + layout_manager = ContextParallelLayoutManager( + layer_layouts=("contiguous",), + boundary_layout="zigzag", + sequence_parallel=self.config.sequence_parallel, + cp_group=cp_group_chunkwise, + tp_group=self.tp_group, + tp_cp_group=getattr(active_pg_collection, "tp_cp", None), + ) + layout_packed_seq_params = packed_seq_params + if packed_seq_params is not None and packed_seq_params.total_tokens is None: + layout_packed_seq_params = replace( + packed_seq_params, + total_tokens=hidden_states.shape[0] * cp_size_runtime * self.sp_size, + ) + input_layout_state = layout_manager.build_forward_state(layout_packed_seq_params) + if input_layout_state is not None: + hidden_states, packed_seq_params = input_layout_state.prepare_layer(0, hidden_states) + + 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(): @@ -120,7 +176,7 @@ def forward( not self.config.batch_invariant_mode ), "GDN dynamic inference does not support batch-invariant mode." assert ( - self.cp_size == 1 + cp_size_runtime == 1 ), "Context parallelism is not supported for GDN dynamic inference." assert ( inference_context.num_speculative_tokens == 0 @@ -133,159 +189,186 @@ def forward( assert not self.config.sequence_parallel raise NotImplementedError("GDN static-batching inference is not supported.") - 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" - assert ( - not self.config.deterministic_mode - ), "Packed sequence does not support deterministic mode." + 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", "zigzag") != "zigzag" + ) + or ( + (packed_seq_params is None or packed_seq_params.qkv_format != "thd") + and self.config.linear_cp_layout != "zigzag" + ) + ): + raise ValueError("GatedDeltaNet with headwise CP requires zigzag layout.") - # Resolve cu_seqlens with alignment padding handling. + 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" + assert not self.config.deterministic_mode, "Packed sequence does not support deterministic mode." 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, " f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" ) - num_packed_seqs = cu_seqlens_q.shape[0] - 1 - assert num_packed_seqs > 0, ( - "Number of packed sequences must be greater than 0, " - f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" - ) else: cu_seqlens_q = None - cu_seqlens_kv = None - - # Input projection - nvtx_range_push(suffix="in_proj") - qkvzba, _ = self.in_proj(hidden_states) - nvtx_range_pop(suffix="in_proj") - - qkvzba, thd_cp_a2a_inv = a2a_cp_to_hp( - qkvzba, - self.in_proj_split_sections, - self.cp_size, - self.pg_collection.cp, - cu_seqlens_q, - seq_len, - packed_seq_params, - ) - # Transpose: s b x --> b s x - # From sbhd to bshd format - qkvzba = qkvzba.transpose(0, 1) - - # 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) - - # Convolution on qkv - nvtx_range_push(suffix="conv1d") - 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=self.pg_collection.cp, - split_sections=qkv_channels_split_sections, - ) - conv1d_bias = ( - get_parameter_local_cp( - self.conv1d.bias, + if cp_size_chunkwise > 1: + if cu_seqlens_q is None: + cu_seqlens_q = ( + torch.arange(batch + 1, device=torch.cuda.current_device(), dtype=torch.long) + * seq_len_global + ) + 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 cp_size_chunkwise > 1 and packed_seq_params is None 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.") + + def _compute(hidden_states): + nvtx_range_push(suffix="in_proj") + projected, _ = self.in_proj(hidden_states) + nvtx_range_pop(suffix="in_proj") + projected, thd_cp_a2a_inv = a2a_cp_to_hp( + projected, + self.in_proj_split_sections, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + seq_len_post_headwise, + packed_seq_params, + ) + projected = projected.transpose(0, 1) + qkv, gate, beta, alpha = torch.split( + projected, self._get_feat_dim_split(cp_size_headwise), dim=-1 + ) + gate = gate.reshape(batch, seq_len_post_headwise, -1, self.value_head_dim) + + nvtx_range_push(suffix="conv1d") + qkv_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=self.pg_collection.cp, - split_sections=qkv_channels_split_sections, + cp_group=cp_group_headwise, + split_sections=qkv_sections, ) - if self.conv_bias - else None - ) - if self.config.deterministic_mode: - qkv = qkv.transpose(1, 2).contiguous() # b, s, d -> b, d, s - conv_out = F.conv1d( - input=qkv, # Torch-native only accept [b, d, s] format input - 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, + conv1d_bias = ( + get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_sections, + ) + if self.conv_bias + else None ) - qkv = self.act_fn(conv_out[..., :seq_len]) - qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d - else: - assert self.activation in ["silu", "swish"] - qkv, _ = causal_conv1d( - x=qkv, # FLA conv1d accepts [b, s, d] format input - weight=conv1d_weight.squeeze(1), # d, 1, w -> d, w - bias=conv1d_bias, - activation=self.activation, + kernel_seq_len = qkv.shape[1] + if self.config.deterministic_mode: + qkv = qkv.transpose(1, 2).contiguous() + qkv = self.act_fn( + 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, + )[..., :kernel_seq_len] + ).transpose(1, 2) + else: + assert self.activation in ["silu", "swish"] + qkv, _ = causal_conv1d( + x=qkv.contiguous(), + weight=conv1d_weight.squeeze(1), + bias=conv1d_bias, + activation=self.activation, + initial_state=None, + output_final_state=False, + cu_seqlens=cu_seqlens_q, + cp_context=chunkwise_cp_context, + ) + 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 + ) + kernel_inputs = self._prepare_input_for_gated_delta_rule( + 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") + 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="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 - ) - - # 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 - ) - 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) + def _norm_and_restore(core_attn_out, gate): + norm_out = self._apply_gated_norm(core_attn_out, gate) + norm_out = ( + norm_out.reshape(batch, kernel_seq_len, -1).transpose(0, 1).contiguous() + ) + return a2a_hp_to_cp( + norm_out, + cp_size_headwise, + cp_group_headwise, + packed_seq_params, + thd_cp_a2a_inv, + ) + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() + norm_out = self.norm_out_checkpoint.checkpoint( + _norm_and_restore, core_attn_out, gate + ) + else: + norm_out = _norm_and_restore(core_attn_out, gate) + out, out_bias = self.out_proj(norm_out) + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint.discard_output_and_register_recompute(out) + return out, out_bias + + if self.recompute_gdn and self.training: + out, out_bias = tensor_parallel.checkpoint(_compute, False, hidden_states) 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) - + out, out_bias = _compute(hidden_states) + if input_layout_state is not None: + out = input_layout_state.finalize_layer(0, out) return out, out_bias def _split_projection( @@ -425,22 +508,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,6 +529,9 @@ 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 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..15f8a8eaf5e --- /dev/null +++ b/megatron/core/ssm/gated_delta_net/kda.py @@ -0,0 +1,681 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Kimi Delta Attention, a channel-wise Gated DeltaNet variant.""" + +import math +from dataclasses import dataclass, replace +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 ContextParallelLayoutManager +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.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.feat_dim_split = self._get_feat_dim_split(self.cp_size) + 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( # DevSkim: ignore DS148264 -- non-security parameter initialization. + 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 = ( + packed_seq_params.cp_group + if packed_seq_params is not None and packed_seq_params.cp_group is not None + else base_cp_group + ) + 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() + + input_layout_state = None + if ( + self.config.linear_cp_mode == "chunkwise" + and self.config.linear_cp_layout == "zigzag" + and cp_size_chunkwise > 1 + ): + layout_manager = ContextParallelLayoutManager( + layer_layouts=("contiguous",), + boundary_layout="zigzag", + sequence_parallel=self.config.sequence_parallel, + cp_group=cp_group_chunkwise, + tp_group=self.tp_group, + tp_cp_group=getattr(active_pg_collection, "tp_cp", None), + ) + layout_packed_seq_params = packed_seq_params + if packed_seq_params is not None and packed_seq_params.total_tokens is None: + layout_packed_seq_params = replace( + packed_seq_params, + total_tokens=hidden_states.shape[0] * cp_size_runtime * self.sp_size, + ) + input_layout_state = layout_manager.build_forward_state(layout_packed_seq_params) + hidden_states, packed_seq_params = input_layout_state.prepare_layer(0, hidden_states) + + 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 self.config.linear_cp_layout != "zigzag" + ) + or ( + (packed_seq_params is None or packed_seq_params.qkv_format != "thd") + and self.config.linear_cp_layout != "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 input_layout_state is not None: + out = input_layout_state.finalize_layer(0, out) + + 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 20a586204ed..a55544b5d47 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -1674,6 +1674,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..d6f83b7e9ce 100644 --- a/megatron/core/transformer/fsdp_dtensor_checkpoint.py +++ b/megatron/core/transformer/fsdp_dtensor_checkpoint.py @@ -395,66 +395,153 @@ def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis, is_expert_param def handle_gdn_in_state_dict(model, model_state_dict, optimizer_state_dict): - """Handle GDN (Gated DeltaNet) fused projections in model and optimizer state dicts. + """Handle GDN-family fused projections in model and optimizer state dicts. GDN layers fuse query/key/value/gate/beta/alpha projections into a single - ``in_proj.weight`` ColumnParallelLinear, and query/key/value into ``conv1d`` - (weight + optional bias). For FSDP checkpoints these fused tensors must be - split back into their constituent sub-tensors so that each can be - independently TP-sharded — otherwise loading a checkpoint written at TP=M - into TP=N would slice across logical component boundaries. + ``in_proj.weight`` ColumnParallelLinear, while KDA uses its own five-way + query/key/value/g/gate split and a separate beta projection. Both variants + share the query/key/value ``conv1d`` split. For FSDP checkpoints these fused + tensors must be split back into their constituent sub-tensors so that each + can be independently TP-sharded — otherwise loading a checkpoint written at + TP=M into TP=N would slice across logical component boundaries. This is analogous to :func:`handle_swiglu_in_state_dict` which splits ``linear_fc1`` into ``weight_w`` / ``weight_v``. - Sub-key naming follows ``GatedDeltaNet.sharded_state_dict()``:: + Sub-key naming follows each module's ``in_proj_split_names`` and + ``in_proj_split_sections`` metadata:: - in_proj.weight → .query / .key / .value / .z / .beta / .alpha (6-way) + GDN in_proj.weight → .query / .key / .value / .z / .beta / .alpha (6-way) + KDA in_proj.weight → .query / .key / .value / .g / .gate (5-way) conv1d.weight → .query / .key / .value (3-way) conv1d.bias → .query / .key / .value (3-way) """ assert HAVE_MEGATRON_FSDP, "This function requires Megatron-FSDP to be installed." - GDN_IN_PROJ_NAMES = ["query", "key", "value", "z", "beta", "alpha"] GDN_CONV1D_NAMES = ["query", "key", "value"] # ------------------------------------------------------------------ - # Build per-GDN-module split-size map by walking the model tree. - # GDN modules are identified by the presence of qk_dim / v_dim / - # in_proj_dim attributes (set in GatedDeltaNet.__init__). + # Build the per-GDN-family-module split-size map by walking the model tree. + # Both GDN and KDA expose qk_dim / v_dim / in_proj_dim; the split metadata + # is authoritative because their in_proj layouts are different. # ------------------------------------------------------------------ - _gdn_info = {} # normalized_path → {in_proj_sizes, conv1d_sizes} - for name, mod in model.named_modules(): - if not (hasattr(mod, 'qk_dim') and hasattr(mod, 'v_dim') and hasattr(mod, 'in_proj_dim')): + projection_info = {} # normalized path → split sizes and names + for name, module in model.named_modules(): + if not all(hasattr(module, attr) for attr in ('qk_dim', 'v_dim', 'in_proj_dim')): + continue + has_names = hasattr(module, 'in_proj_split_names') + has_sections = hasattr(module, 'in_proj_split_sections') + if not has_names and not has_sections: continue - tp = getattr(mod, 'tp_size', 1) - qk = mod.qk_dim // tp - v = mod.v_dim // tp - nvh = mod.num_value_heads // tp - _gdn_info[_strip_wrapper_prefixes(name)] = { - 'in_proj_sizes': [qk, qk, v, v, nvh, nvh], - 'conv1d_sizes': [qk, qk, v], + if not (has_names and has_sections): + raise ValueError( + f"GDN-family module {name!r} must define both in_proj_split_names and " + "in_proj_split_sections." + ) + + in_proj_names = tuple(module.in_proj_split_names) + in_proj_sizes = tuple(module.in_proj_split_sections) + if not in_proj_names or len(in_proj_sizes) != len(in_proj_names): + raise ValueError( + f"GDN-family module {name!r} has mismatched in_proj split metadata: " + f"{len(in_proj_sizes)} sizes for {len(in_proj_names)} names." + ) + if any(not isinstance(split_name, str) or not split_name for split_name in in_proj_names): + raise ValueError( + f"GDN-family module {name!r} has invalid in_proj split names: {in_proj_names}." + ) + if len(set(in_proj_names)) != len(in_proj_names): + raise ValueError( + f"GDN-family module {name!r} has duplicate in_proj split names: " + f"{in_proj_names}." + ) + if any( + not isinstance(size, int) or isinstance(size, bool) or size <= 0 + for size in in_proj_sizes + ): + raise ValueError( + f"GDN-family module {name!r} has invalid in_proj split sizes: {in_proj_sizes}." + ) + + tp_size = getattr(module, 'tp_size', 1) + if not isinstance(tp_size, int) or isinstance(tp_size, bool) or tp_size <= 0: + raise ValueError(f"GDN-family module {name!r} has invalid tp_size={tp_size!r}.") + if module.in_proj_dim % tp_size != 0: + raise ValueError( + f"GDN-family module {name!r} has in_proj_dim={module.in_proj_dim} " + f"which is not divisible by tp_size={tp_size}." + ) + expected_size = module.in_proj_dim // tp_size + if sum(in_proj_sizes) != expected_size: + raise ValueError( + f"GDN-family module {name!r} has in_proj split sizes totaling " + f"{sum(in_proj_sizes)}, expected {expected_size}." + ) + + qk_size = getattr(module, 'qk_dim_local_tp', None) + value_size = getattr(module, 'v_dim_local_tp', None) + if qk_size is None or value_size is None: + if module.qk_dim % tp_size != 0 or module.v_dim % tp_size != 0: + raise ValueError( + f"GDN-family module {name!r} qk_dim/v_dim must be divisible by " + f"tp_size={tp_size}." + ) + qk_size = module.qk_dim // tp_size + value_size = module.v_dim // tp_size + + normalized_name = _strip_wrapper_prefixes(name) + if normalized_name in projection_info: + raise ValueError(f"Multiple GDN-family modules normalize to path {normalized_name!r}.") + projection_info[normalized_name] = { + 'in_proj_sizes': in_proj_sizes, + 'in_proj_names': in_proj_names, + 'conv1d_sizes': (qk_size, qk_size, value_size), } - if not _gdn_info: + if not projection_info: return model_state_dict, optimizer_state_dict + parameter_map = {} + for name, parameter in model.named_parameters(): + normalized_name = _strip_wrapper_prefixes(name) + if normalized_name in parameter_map and parameter_map[normalized_name] is not parameter: + raise ValueError( + f"Multiple parameters normalize to GDN-family key {normalized_name!r}." + ) + parameter_map[normalized_name] = parameter + + def _get_parameter(key): + normalized_key = _strip_wrapper_prefixes(key) + try: + return parameter_map[normalized_key] + except KeyError as error: + raise KeyError( + f"No model parameter matches GDN-family state-dict key {key!r} " + f"(normalized as {normalized_key!r})." + ) from error + def _match_gdn_key(key): - """Return (split_sizes, sub_names, split_dim) if *key* is a GDN fused - parameter that needs splitting, else ``None``.""" + """Return split metadata when ``key`` names a fused GDN-family parameter.""" norm = _strip_wrapper_prefixes(key) - for gdn_path, info in _gdn_info.items(): - if not norm.startswith(gdn_path + '.'): - continue - rel = norm[len(gdn_path) + 1 :] + for module_path, info in projection_info.items(): + if module_path: + if not norm.startswith(module_path + '.'): + continue + rel = norm[len(module_path) + 1 :] + else: + rel = norm if rel == 'in_proj.weight': - return info['in_proj_sizes'], GDN_IN_PROJ_NAMES, 0 + return info['in_proj_sizes'], info['in_proj_names'], 0 if rel in ('conv1d.weight', 'conv1d.bias'): return info['conv1d_sizes'], GDN_CONV1D_NAMES, 0 return None 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-family projection DTensor into component DTensors. + + The implementation handles both model-state DTensors and optimizer-state + tensors while preserving FSDP and tensor-parallel metadata. + """ total_split = sum(split_sizes) if isinstance(data, DTensor) and data.shape[split_dim] == total_split: return list( @@ -518,7 +605,7 @@ def split_gdn_fused(data, dist_param, split_sizes, split_dim): if match is None: continue sizes, names, dim = match - dist_param = model.get_parameter(f"module.{key}") + dist_param = _get_parameter(key) sub_tensors = split_gdn_fused(model_state_dict[key], dist_param, sizes, dim) for sub_name, tensor in zip(names, sub_tensors): model_state_dict[f"{key}.{sub_name}"] = tensor @@ -546,7 +633,7 @@ def split_gdn_fused(data, dist_param, split_sizes, split_dim): for sub_name in names: new_opt_state[f"{key}.{sub_name}"] = opt_state[key].copy() for subkey in ["exp_avg", "exp_avg_sq"]: - dist_param = model.get_parameter(key[len("module.") :]) + dist_param = _get_parameter(key) sub_tensors = split_gdn_fused(opt_state[key][subkey], dist_param, sizes, dim) for sub_name, tensor in zip(names, sub_tensors): new_opt_state[f"{key}.{sub_name}"][subkey] = tensor diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 6ce93c76837..807b2c4cdcd 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -554,6 +554,19 @@ def _apply_global_aux_loss( ) return probs + 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 = getattr(self, 'mtp_layer_number', None) 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, @@ -591,14 +604,7 @@ 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 - - if self.is_mtp_layer: - layer_number = self.layer_number + self.config.num_layers - else: - layer_number = self.layer_number + layer_number, num_layers = self._get_metric_layer_number() get_moe_metrics_tracker().record( aux_loss_name, @@ -694,14 +700,7 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = 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 - - if self.is_mtp_layer: - layer_number = self.layer_number + self.config.num_layers - else: - layer_number = self.layer_number + layer_number, num_layers = self._get_metric_layer_number() get_moe_metrics_tracker().record( "z_loss", diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index d821c14effe..62ad4d9096c 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -18,6 +18,7 @@ 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, @@ -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): @@ -161,6 +163,10 @@ def __init__( 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 @@ -217,6 +223,26 @@ def __init__( pg_collection=self.pg_collection, ) + 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 +351,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 ( @@ -457,6 +485,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] # ================= @@ -467,6 +498,65 @@ def forward( 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 @@ -1092,6 +1182,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): diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index a9df601bd1e..4646fa19a2f 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -277,7 +277,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 +326,17 @@ 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', 'gdn2', 'kda', 'gated_delta_net', 'dsa', 'dsv4_hybrid'] + ] = None + """Type of experimental attention variant to use. + + ``gdn``, ``gdn2``, and ``kda`` select Gated DeltaNet-family implementations. + ``gated_delta_net`` is a deprecated compatibility alias for ``gdn`` and is normalized in + ``__post_init__`` with a ``DeprecationWarning``. + """ + linear_attention_type: Optional[str] = None + """Deprecated alias for experimental_attention_variant.""" experimental_attention_variant_loss_scale_func: Optional[Callable[[torch.Tensor], None]] = None """Optional hook for experimental attention variants to receive the main loss scale.""" @@ -343,6 +372,9 @@ class TransformerConfig(ModelParallelConfig): ``none`` disables fused DSA kernels. Explicit ``tilelang`` or ``cudnn`` enables only that backend. Unsupported DSA layouts continue to use the PyTorch fallback.""" + apply_dsa_kernel_fusion: Optional[bool] = None + """Deprecated alias for selecting the DSA kernel backend.""" + dsa_indexer_rope_interleaved: bool = False """Whether DSA indexer RoPE should use MLA-style interleaving.""" @@ -383,6 +415,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 #################### @@ -561,7 +607,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", "gdn", "gdn_norm_out", "gdp_in_proj", "gdp_qkv", "mhc". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -570,7 +616,14 @@ 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. + "mhc": recompute HyperConnection intermediate activations via + CheckpointWithoutOutput + MHCCheckpointManager. Requires + 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 + GDN-family layers. "gdp_in_proj": recompute the GatedDeltaProduct input projection and its CP gather/split preprocessing. "gdp_qkv": recompute the GatedDeltaProduct causal conv and QKV preparation. @@ -578,8 +631,8 @@ class TransformerConfig(ModelParallelConfig): 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. + "mhc" use output-discarding checkpointing, "core_attn", "mlp", "moe", + "shared_experts", and "gdn" use normal checkpointing. """ #################### @@ -1050,6 +1103,9 @@ class TransformerConfig(ModelParallelConfig): linear_cp_layout: CPLayout = "zigzag" """CP layout for linear-attention layers.""" + linear_cp_mode: Literal["headwise", "chunkwise"] = "headwise" + """Context-parallel execution mode for GDN-family linear-attention layers.""" + attention_cp_layout: CPLayout = "zigzag" """CP layout for softmax-attention layers.""" @@ -1463,11 +1519,6 @@ def __post_init__(self): 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 @@ -1519,6 +1570,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 @@ -1528,47 +1590,179 @@ def __post_init__(self): f"tensor_model_parallel_size ({self.tensor_model_parallel_size})." ) + 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.linear_cp_layout not in ("zigzag", "contiguous"): + raise ValueError(f"Unsupported cp_partition_mode: {self.linear_cp_layout}") + + if self.linear_cp_layout == "contiguous" and ( + self.context_parallel_size > 1 or self.hybrid_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.linear_cp_layout == "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.linear_cp_layout == "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." - # 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=})." - ) - assert self.linear_num_value_heads % tp_cp_size == 0, ( - f"{self.linear_num_value_heads=} must be a multiple of " - f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." + 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( @@ -1584,6 +1778,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.hybrid_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 @@ -2011,6 +2307,7 @@ def __post_init__(self): "mlp", "moe", "shared_experts", + "gdn", "gdn_norm_out", "gdp_in_proj", "gdp_qkv", @@ -2033,14 +2330,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, " @@ -2856,6 +3170,122 @@ 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".' + use_mhc_recompute = ( + self.recompute_granularity == "selective" and "mhc" in self.recompute_modules + ) + + # 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.hybrid_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": @@ -3437,8 +3867,16 @@ def __post_init__(self): if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": 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." + ) if self.cache_mla_latents: assert ( diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py index 124083ead83..81c0b5f7339 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 @@ -274,6 +274,80 @@ 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!r}, 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!r}, expected {expected}." + ) + compact.append(ratio) + else: + assert ratio == 0, ( + f"csa_compress_ratios should not pad non-attention hybrid symbol " + f"{symbol!r} 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!r}." + ) + + 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 @@ -345,8 +419,21 @@ def core_transformer_config_from_args(args, config_class=None): 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) + kw_args['inference_sampling_seed'] = args.seed diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index d3b71ef0bb2..2b4d3fa7727 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.""" @@ -787,6 +787,8 @@ def validate_args(args, defaults={}): 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 +868,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: @@ -918,10 +919,12 @@ def validate_args(args, defaults={}): 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 @@ -3678,28 +3681,43 @@ def _add_heterogeneous_args(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', diff --git a/megatron/training/training.py b/megatron/training/training.py index e4b0afb56f3..4aeec0ecf0e 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Pretrain utilities.""" @@ -860,8 +860,15 @@ def moe_layer_flops(total_tokens, hidden_size, moe_ffn_hidden_size, return routed_flops + shared_flops def attn_layer_flops( - total_tokens, seqlen_squared_sum, hidden_size, num_heads, gqa=True, - gqa_groups=8, kv_channels=None, + total_tokens, + seqlen_squared_sum, + hidden_size, + num_heads, + gqa=True, + gqa_groups=8, + kv_channels=None, + attention_output_gate=False, + gated_attention_proj_granularity="elementwise", ): """Calculate FLOPs for an attention layer. @@ -874,16 +881,85 @@ def attn_layer_flops( """ p = (kv_channels * num_heads / hidden_size) if kv_channels else 1 g = gqa_groups if gqa else num_heads + gate_projection_size = 0 + if attention_output_gate: + if gated_attention_proj_granularity == "elementwise": + gate_projection_size = hidden_size * p + elif gated_attention_proj_granularity == "headwise": + gate_projection_size = num_heads + else: + raise ValueError( + "gated_attention_proj_granularity must be either 'elementwise' or " + f"'headwise', got {gated_attention_proj_granularity!r}." + ) # 4 * total_tokens * h * p * (h + h*(g/n)): QKV + output projections (fwd*3, with FMA*2). + # The output gate contributes one additional hidden_size -> gate_size projection. # 2 * sum(L^2) * h * p: core attention (causal mask -> /2 cancels with FMA *2). return ( - 4 * total_tokens * hidden_size * p - * (hidden_size + hidden_size * (g / num_heads)) + 4 * total_tokens * hidden_size * p * (hidden_size + hidden_size * (g / num_heads)) + + 2 * total_tokens * hidden_size * gate_projection_size + 2 * seqlen_squared_sum * hidden_size * p ) - def mamba_layer_flops(total_tokens, hidden_size, state_dim=16, - head_dim=64, num_groups=1, num_heads=128): + def mla_attn_layer_flops( + total_tokens, + seqlen_squared_sum, + hidden_size, + num_heads, + q_lora_rank, + kv_lora_rank, + qk_head_dim, + qk_pos_emb_head_dim, + v_head_dim, + attention_output_gate=False, + gated_attention_proj_granularity="elementwise", + ): + """Calculate FLOPs for a Multi-Latent Attention (MLA) layer. + + Mirrors the MLA term in ``transformer_flops`` so the hybrid estimate + matches the standard-model estimate per attention layer (the DSv4 + attention layers -- Window/CSA/HCA -- are all MLA-based). The generic + ``attn_layer_flops`` above assumes dense MHA/GQA QKV projections and + badly overcounts MLA, whose Q/KV go through low-rank ``lora`` ranks. + + Returns a forward-equivalent value (FMA factor 2 baked in, no + forward/backward factor): hybrid_flops applies the global factor 3. + """ + fma = 2 + if q_lora_rank is None: + q_term = hidden_size * num_heads * (qk_head_dim + qk_pos_emb_head_dim) + else: + q_term = q_lora_rank * ( + hidden_size + num_heads * (qk_head_dim + qk_pos_emb_head_dim) + 1 + ) + gate_projection_size = 0 + if attention_output_gate: + if gated_attention_proj_granularity == "elementwise": + gate_projection_size = num_heads * v_head_dim + elif gated_attention_proj_granularity == "headwise": + gate_projection_size = num_heads + else: + raise ValueError( + "gated_attention_proj_granularity must be either 'elementwise' or " + f"'headwise', got {gated_attention_proj_granularity!r}." + ) + # Token-linear part (q lora+rope+norm, kv lora+rope+norm, output proj). + token_linear = fma * ( + q_term + + kv_lora_rank * (hidden_size + num_heads * (qk_head_dim + v_head_dim) + 1) + + hidden_size * qk_pos_emb_head_dim + + (num_heads * v_head_dim) * hidden_size + + hidden_size * gate_projection_size + ) + # Core attention (L^2) part: QK^T and (softmax(QK^T))V. /2 (causal) cancels *2 (FMA). + core = fma * ( + num_heads * (qk_head_dim + qk_pos_emb_head_dim) / 2 + num_heads * v_head_dim / 2 + ) + return token_linear * total_tokens + core * seqlen_squared_sum + + def mamba_layer_flops( + total_tokens, hidden_size, state_dim=16, head_dim=64, num_groups=1, num_heads=128 + ): """Calculate FLOPs for a Mamba layer.""" # Note (rwaleffe): flops estimate for scan should be updated based on new SSD kernels, # but small percent of overall layer flops @@ -961,48 +1037,211 @@ def gdn_layer_flops(total_tokens, hidden_size, ) ) - def hybrid_flops(total_tokens, seqlen_squared_sum, hidden_size, - num_attn_layers, num_mamba_layers, num_mlp_layers, num_moe_layers, - gdp_num_householder, - num_gdn_layers=0, - mamba_state_dim=128, mamba_head_dim=64, - mamba_num_groups=8, mamba_num_heads=128, - num_attn_heads=32, gqa=True, - gqa_groups=8, kv_channels=None, - mlp_expansion=4.0, swiglu=False, - moe_latent_size=None, - moe_ffn_hidden_size=2048, shared_expert_ffn_hidden_size=2048, num_experts_routed_to=1, - use_gated_delta_product=False, - gdn_qk_head_dim=128, gdn_v_head_dim=128, - gdn_num_qk_heads=16, gdn_num_v_heads=32, - gdn_conv_kernel_dim=4, gdn_use_gdn2=False, - vocab_size=256000, mtp_num_layers=0): + def kda_layer_flops( + total_tokens, + hidden_size, + qk_head_dim=128, + v_head_dim=128, + num_qk_heads=16, + num_v_heads=16, + conv_kernel_dim=4, + ): + """Calculate FLOPs for a direct-projection Kimi Delta Attention layer.""" + if num_qk_heads != num_v_heads or qk_head_dim != v_head_dim: + raise ValueError( + "KDA FLOPs require the equal K/V head layout enforced by KimiDeltaAttention." + ) + qk_dim = qk_head_dim * num_qk_heads + v_dim = v_head_dim * num_v_heads + in_proj_dim = 3 * qk_dim + 2 * v_dim + non_core_flops = ( + 2 + * total_tokens + * ( + hidden_size * (in_proj_dim + num_qk_heads) + + conv_kernel_dim * (2 * qk_dim + v_dim) + + hidden_size * v_dim + ) + ) + state_update_flops = num_v_heads * (qk_head_dim**2 + 3 * qk_head_dim * v_head_dim) + core_flops = 2 * total_tokens * state_update_flops + return non_core_flops + core_flops + + def hybrid_flops( + total_tokens, + seqlen_squared_sum, + hidden_size, + num_attn_layers, + num_mla_layers, + num_kda_layers, + num_mamba_layers, + num_mlp_layers, + num_moe_layers, + num_gdn_layers=0, + mamba_state_dim=128, + mamba_head_dim=64, + mamba_num_groups=8, + mamba_num_heads=128, + num_attn_heads=32, + gqa=True, + gqa_groups=8, + kv_channels=None, + mlp_expansion=4.0, + swiglu=False, + moe_latent_size=None, + moe_ffn_hidden_size=2048, + shared_expert_ffn_hidden_size=2048, + num_experts_routed_to=1, + gdp_num_householder=0, + use_gated_delta_product=False, + gdn_qk_head_dim=128, + gdn_v_head_dim=128, + gdn_num_qk_heads=16, + gdn_num_v_heads=32, + gdn_conv_kernel_dim=4, + kda_qk_head_dim=128, + kda_v_head_dim=128, + kda_num_qk_heads=16, + kda_num_v_heads=16, + kda_conv_kernel_dim=4, + vocab_size=256000, + mtp_num_layers=0, + q_lora_rank=None, + kv_lora_rank=0, + qk_head_dim=0, + qk_pos_emb_head_dim=0, + v_head_dim=0, + attention_output_gate=False, + gated_attention_proj_granularity="elementwise", + experimental_attention_variant=None, + dsv4_n_layers_r0=0, + dsv4_n_layers_r4=0, + dsv4_n_layers_r128=0, + o_groups=None, + o_lora_rank=None, + csa_window_size=None, + seq_length=None, + dsa_indexer_n_heads=None, + dsa_indexer_head_dim=None, + dsa_indexer_topk=None, + ): """Calculate total FLOPs for the hybrid model.""" + # Self-attention (already summed over all attention layers, fwd-equivalent + # with the FMA factor baked in; the global ``* 3`` below adds fwd+bwd). + if experimental_attention_variant == "dsv4_hybrid": + # DSv4 uses sparse MLA attention. Keep the shared helper as the + # single source of truth for both HybridModel and GPTModel. + dsv4_token_term, dsv4_core_term = _dsv4_hybrid_self_attention_flops( + hidden_size=hidden_size, + num_attention_heads=num_attn_heads, + v_head_dim=v_head_dim, + q_lora_rank=q_lora_rank, + o_groups=o_groups, + o_lora_rank=o_lora_rank, + csa_window_size=csa_window_size, + seq_length=seq_length, + n_layers_r0=dsv4_n_layers_r0, + n_layers_r4=dsv4_n_layers_r4, + n_layers_r128=dsv4_n_layers_r128, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + ) + attn_flops_total = 2 * ( + dsv4_token_term * total_tokens + dsv4_core_term * seqlen_squared_sum + ) + else: + # HybridStack assigns '*' to regular attention and '+' to MLA independently of the + # model-wide multi_latent_attention flag. + plain_attn_layer_flops = attn_layer_flops( + total_tokens, + seqlen_squared_sum, + hidden_size, + num_attn_heads, + gqa, + gqa_groups, + kv_channels, + attention_output_gate, + gated_attention_proj_granularity, + ) + attn_flops_total = num_attn_layers * plain_attn_layer_flops + if num_mla_layers: + attn_flops_total += num_mla_layers * mla_attn_layer_flops( + total_tokens, + seqlen_squared_sum, + hidden_size, + num_attn_heads, + q_lora_rank, + kv_lora_rank, + qk_head_dim, + qk_pos_emb_head_dim, + v_head_dim, + attention_output_gate, + gated_attention_proj_granularity, + ) + + kda_flops_total = 0 + if num_kda_layers: + kda_flops_total = num_kda_layers * kda_layer_flops( + total_tokens, + hidden_size, + kda_qk_head_dim, + kda_v_head_dim, + kda_num_qk_heads, + kda_num_v_heads, + kda_conv_kernel_dim, + ) + mamba_flops = ( - gated_delta_product_layer_flops(total_tokens, hidden_size, - gdp_num_householder, - mamba_state_dim, mamba_head_dim, - mamba_num_groups, mamba_num_heads) + gated_delta_product_layer_flops( + total_tokens, + hidden_size, + gdp_num_householder, + mamba_state_dim, + mamba_head_dim, + mamba_num_groups, + mamba_num_heads, + ) if use_gated_delta_product - else mamba_layer_flops(total_tokens, hidden_size, - mamba_state_dim, mamba_head_dim, - mamba_num_groups, mamba_num_heads) + else mamba_layer_flops( + total_tokens, + hidden_size, + mamba_state_dim, + mamba_head_dim, + mamba_num_groups, + mamba_num_heads, + ) ) + flops_fwd = ( - num_attn_layers * attn_layer_flops(total_tokens, seqlen_squared_sum, - hidden_size, num_attn_heads, gqa, - gqa_groups, kv_channels) + - num_mlp_layers * mlp_layer_flops(total_tokens, hidden_size, - mlp_expansion, swiglu) + - num_mamba_layers * mamba_flops + - num_moe_layers * moe_layer_flops(total_tokens, hidden_size, moe_ffn_hidden_size, - shared_expert_ffn_hidden_size, num_experts_routed_to, - moe_latent_size, swiglu) + - num_gdn_layers * gdn_layer_flops(total_tokens, hidden_size, - gdn_qk_head_dim, gdn_v_head_dim, - gdn_num_qk_heads, gdn_num_v_heads, - gdn_conv_kernel_dim, gdn_use_gdn2) + - (2 * total_tokens * hidden_size * vocab_size * (1 + mtp_num_layers)) # logits computation + attn_flops_total + + kda_flops_total + + num_mlp_layers * mlp_layer_flops(total_tokens, hidden_size, mlp_expansion, swiglu) + + num_mamba_layers * mamba_flops + + num_moe_layers + * moe_layer_flops( + total_tokens, + hidden_size, + moe_ffn_hidden_size, + shared_expert_ffn_hidden_size, + num_experts_routed_to, + moe_latent_size, + swiglu, + ) + + num_gdn_layers + * gdn_layer_flops( + total_tokens, + hidden_size, + gdn_qk_head_dim, + gdn_v_head_dim, + gdn_num_qk_heads, + gdn_num_v_heads, + gdn_conv_kernel_dim, + ) + + + # MTP norms (eh_norm + final_norm) and eh projection (2 * h^2). + 2 * mtp_num_layers * (3 * hidden_size + 2 * hidden_size * hidden_size) * total_tokens + + 2 * total_tokens * hidden_size * vocab_size * (1 + mtp_num_layers) ) return flops_fwd * 3 @@ -1104,6 +1343,20 @@ def transformer_flops(): + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim) + 1 ) + gate_projection_size = 0 + if args.attention_output_gate: + gate_granularity = getattr( + args, "gated_attention_proj_granularity", "elementwise" + ) + if gate_granularity == "elementwise": + gate_projection_size = args.num_attention_heads * args.v_head_dim + elif gate_granularity == "headwise": + gate_projection_size = args.num_attention_heads + else: + raise ValueError( + "gated_attention_proj_granularity must be either 'elementwise' or " + f"'headwise', got {gate_granularity!r}." + ) # Token-linear part of MLA self-attention (lora projs, kv proj, RoPE, output proj). standard_self_attn_term = ( forward_backward_expansion_factor @@ -1121,6 +1374,8 @@ def transformer_flops(): + args.hidden_size * args.qk_pos_emb_head_dim ## o proj + (args.num_attention_heads * args.v_head_dim) * args.hidden_size + ## output gate proj + + gate_projection_size * args.hidden_size ) ) # Core-attention (L^2) part: ``QK^T`` and ``(softmax(QK^T)) V``. The @@ -1199,7 +1454,17 @@ def transformer_flops(): num_linear_attention_layers = sum(linear_attention_pattern) num_standard_attention_layers = num_layers - num_linear_attention_layers - if is_gated_delta_net_variant(args.experimental_attention_variant): + if args.experimental_attention_variant == "kda": + linear_self_attn_term = forward_backward_expansion_factor * kda_layer_flops( + total_tokens=1, + hidden_size=args.hidden_size, + qk_head_dim=args.linear_key_head_dim, + v_head_dim=args.linear_value_head_dim, + num_qk_heads=args.linear_num_key_heads, + num_v_heads=args.linear_num_value_heads, + conv_kernel_dim=args.linear_conv_kernel_dim, + ) + elif is_gated_delta_net_variant(args.experimental_attention_variant): # Calculate the FLOPs for the gated delta net attention. qk_head_dim = args.linear_key_head_dim v_head_dim = args.linear_value_head_dim @@ -1330,28 +1595,51 @@ def _split_spec_part(part): # Main entrypoint for FLOPs calculation. if is_hybrid_model(args): - # Calculate the number of each type of layer. - from operator import itemgetter - from megatron.core.models.hybrid.hybrid_layer_allocation import ( Symbols, get_hybrid_layer_counts, ) - num_mamba_layers, num_gdn_layers, num_attn_layers, num_mlp_layers, num_moe_layers = ( - itemgetter(Symbols.MAMBA, Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, Symbols.MOE)( - get_hybrid_layer_counts(args.hybrid_layer_pattern) + + layer_counts = get_hybrid_layer_counts(args.hybrid_layer_pattern) + num_mamba_layers = layer_counts[Symbols.MAMBA] + num_gdn_layers = layer_counts[Symbols.GDN] + num_kda_layers = layer_counts[Symbols.KDA] + num_mlp_layers = layer_counts[Symbols.MLP] + num_moe_layers = layer_counts[Symbols.MOE] + + # DSv4 D/C/H/W layers are MLA variants, but their + # sparse attention FLOPs are handled by the DSv4 branch. + dsv4_mla_symbols = Symbols.MLA_ATTENTION - {Symbols.MLA} + dsv4_mla_layers = sum(layer_counts[symbol] for symbol in dsv4_mla_symbols) + num_mla_layers = layer_counts[Symbols.MLA] + if args.experimental_attention_variant == "dsv4_hybrid": + num_attn_layers = layer_counts[Symbols.ATTENTION] + dsv4_mla_layers + else: + num_attn_layers = layer_counts[Symbols.ATTENTION] + num_mla_layers += dsv4_mla_layers + + dsv4_n_layers_r0 = layer_counts[Symbols.WINDOW] + dsv4_n_layers_r4 = layer_counts[Symbols.CSA] + dsv4_n_layers_r128 = layer_counts[Symbols.HCA] + if args.experimental_attention_variant == "dsv4_hybrid": + assert num_mla_layers == 0, "dsv4_hybrid does not support dense + MLA layers" + assert num_attn_layers == (dsv4_n_layers_r0 + dsv4_n_layers_r4 + dsv4_n_layers_r128), ( + "dsv4_hybrid expects all attention layers to be Window/CSA/HCA; " + f"got {num_attn_layers} attention layers but only " + f"{dsv4_n_layers_r0 + dsv4_n_layers_r4 + dsv4_n_layers_r128} are W/C/H." ) - ) mtp_num_layers = args.mtp_num_layers if mtp_num_layers is None: mtp_num_layers = 0 - # Compute hybrid model FLOPs. + return hybrid_flops( total_tokens=total_real_tokens_in_batch, seqlen_squared_sum=seqlen_squared_sum_in_batch, hidden_size=args.hidden_size, num_attn_layers=num_attn_layers, + num_mla_layers=num_mla_layers, + num_kda_layers=num_kda_layers, num_mamba_layers=num_mamba_layers, num_mlp_layers=num_mlp_layers, num_moe_layers=num_moe_layers, @@ -1360,7 +1648,7 @@ def _split_spec_part(part): mamba_head_dim=args.mamba_head_dim, mamba_num_groups=args.mamba_num_groups, mamba_num_heads=args.mamba_num_heads, - gdp_num_householder=args.gdp_num_householder, + gdp_num_householder=getattr(args, "gdp_num_householder", 0), num_attn_heads=args.num_attention_heads, gqa=args.group_query_attention, gqa_groups=args.num_query_groups, @@ -1379,9 +1667,33 @@ def _split_spec_part(part): gdn_num_qk_heads=args.linear_num_key_heads or 16, gdn_num_v_heads=args.linear_num_value_heads or 32, gdn_conv_kernel_dim=args.linear_conv_kernel_dim or 4, - gdn_use_gdn2=(args.experimental_attention_variant == "gdn2"), + kda_qk_head_dim=args.linear_key_head_dim or 128, + kda_v_head_dim=args.linear_value_head_dim or 128, + kda_num_qk_heads=args.linear_num_key_heads or 16, + kda_num_v_heads=args.linear_num_value_heads or 16, + kda_conv_kernel_dim=args.linear_conv_kernel_dim or 4, vocab_size=args.padded_vocab_size, mtp_num_layers=mtp_num_layers, + q_lora_rank=args.q_lora_rank, + kv_lora_rank=args.kv_lora_rank, + qk_head_dim=args.qk_head_dim, + qk_pos_emb_head_dim=args.qk_pos_emb_head_dim, + v_head_dim=args.v_head_dim, + attention_output_gate=getattr(args, "attention_output_gate", False), + gated_attention_proj_granularity=getattr( + args, "gated_attention_proj_granularity", "elementwise" + ), + experimental_attention_variant=args.experimental_attention_variant, + dsv4_n_layers_r0=dsv4_n_layers_r0, + dsv4_n_layers_r4=dsv4_n_layers_r4, + dsv4_n_layers_r128=dsv4_n_layers_r128, + o_groups=getattr(args, "o_groups", None), + o_lora_rank=getattr(args, "o_lora_rank", None), + csa_window_size=getattr(args, "csa_window_size", None), + seq_length=args.seq_length, + dsa_indexer_n_heads=getattr(args, "dsa_indexer_n_heads", None), + dsa_indexer_head_dim=getattr(args, "dsa_indexer_head_dim", None), + dsa_indexer_topk=getattr(args, "dsa_indexer_topk", None), ) else: # Compute standard Transformer model FLOPs. 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..ee5a22d0c53 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,6 +45,7 @@ def collect_train_test_metrics( if golden_value_key in [ "iteration-time", + "grad-norm", "mem-allocated-bytes", "mem-max-allocated-bytes", "lm loss", @@ -52,6 +53,7 @@ def collect_train_test_metrics( "mtp_1 loss", "mtp_2 loss", "total loss", + "z_loss", ] } 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..b6780258df1 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,11 +13,13 @@ 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)], 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..1cd5118e931 100644 --- a/tests/test_utils/recipes/h100/mamba.yaml +++ b/tests/test_utils/recipes/h100/mamba.yaml @@ -109,3 +109,10 @@ products: - environment: [dev] scope: [mr-broken, mr-github-broken] platforms: [dgx_h100] + + - test_case: [hybrid_mr_mcore_te_tp1_pp1_ep8_cp2_kda_gated_mla_1N8G] + products: + # Enable after the official H100 functional run produces golden values. + - environment: [dev] + scope: [mr, mr-github, mr-github-slim] + platforms: [dgx_h100] diff --git a/tests/unit_tests/distributed/mfsdp_v1/test_gdn_fsdp_checkpoint.py b/tests/unit_tests/distributed/mfsdp_v1/test_gdn_fsdp_checkpoint.py new file mode 100644 index 00000000000..e2c90f760c4 --- /dev/null +++ b/tests/unit_tests/distributed/mfsdp_v1/test_gdn_fsdp_checkpoint.py @@ -0,0 +1,176 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""Distributed regression tests for GDN-family FSDP DTensor checkpoint splitting.""" + +import pytest +import torch +from torch.distributed._tensor import DeviceMesh, Shard, distribute_tensor + +from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + uneven_dtensor_to_full_tensor, +) +from megatron.core.transformer.fsdp_dtensor_checkpoint import handle_gdn_in_state_dict + + +class _FakeGDN(torch.nn.Module): + def __init__(self, split_names, split_sections): + super().__init__() + self.qk_dim = 16 + self.v_dim = 24 + self.tp_size = 1 + self.in_proj_dim = sum(split_sections) + self.in_proj_split_names = split_names + self.in_proj_split_sections = split_sections + self.in_proj = torch.nn.Linear(4, self.in_proj_dim, bias=False) + self.conv1d = torch.nn.Conv1d(4, 2 * self.qk_dim + self.v_dim, 1, bias=True) + + +def _make_wrapped_model(split_names, split_sections): + model = torch.nn.Module() + model.module = torch.nn.Module() + model.module.gdn = _FakeGDN(split_names, split_sections) + return model + + +def _make_sharded_tensor(shape, device_mesh, device, offset=0): + tensor = torch.arange(torch.tensor(shape).prod().item(), dtype=torch.float32, device=device) + tensor = tensor.reshape(shape) + offset + return tensor, distribute_tensor(tensor, device_mesh, [Shard(0)]) + + +def _assert_dtensor_splits(state_dict, key, names, reference): + reference_splits = torch.split( + reference, [state_dict[f"{key}.{name}"].shape[0] for name in names] + ) + for name, expected in zip(names, reference_splits): + actual = uneven_dtensor_to_full_tensor(state_dict[f"{key}.{name}"]) + torch.testing.assert_close(actual, expected) + + +def _optimizer_entry(exp_avg_dtensor, exp_avg_sq_dtensor): + return {"exp_avg": exp_avg_dtensor, "exp_avg_sq": exp_avg_sq_dtensor, "step": torch.tensor(1.0)} + + +@pytest.mark.parametrize( + ("split_names", "split_sections"), + [ + (("query", "key", "value", "z", "beta", "alpha"), (16, 16, 24, 24, 16, 24)), + (("query", "key", "value", "z", "f", "b", "w"), (16, 16, 24, 24, 16, 16, 24)), + (("query", "key", "value", "g", "gate"), (16, 16, 24, 16, 24)), + ], +) +def test_gdn_model_and_optimizer_dtensors_follow_variant_metadata( + distributed_setup, split_names, split_sections +): + mesh = DeviceMesh(distributed_setup.device.type, list(range(distributed_setup.world_size))) + model = _make_wrapped_model(split_names, split_sections) + gdn = model.module.gdn + + weight, weight_dtensor = _make_sharded_tensor( + tuple(gdn.in_proj.weight.shape), mesh, distributed_setup.device + ) + conv_weight, conv_weight_dtensor = _make_sharded_tensor( + tuple(gdn.conv1d.weight.shape), mesh, distributed_setup.device, offset=1_000 + ) + conv_bias, conv_bias_dtensor = _make_sharded_tensor( + tuple(gdn.conv1d.bias.shape), mesh, distributed_setup.device, offset=2_000 + ) + exp_avg, exp_avg_dtensor = _make_sharded_tensor( + tuple(gdn.in_proj.weight.shape), mesh, distributed_setup.device, offset=3_000 + ) + exp_avg_sq, exp_avg_sq_dtensor = _make_sharded_tensor( + tuple(gdn.in_proj.weight.shape), mesh, distributed_setup.device, offset=4_000 + ) + conv_weight_exp_avg, conv_weight_exp_avg_dtensor = _make_sharded_tensor( + tuple(gdn.conv1d.weight.shape), mesh, distributed_setup.device, offset=5_000 + ) + conv_weight_exp_avg_sq, conv_weight_exp_avg_sq_dtensor = _make_sharded_tensor( + tuple(gdn.conv1d.weight.shape), mesh, distributed_setup.device, offset=6_000 + ) + conv_bias_exp_avg, conv_bias_exp_avg_dtensor = _make_sharded_tensor( + tuple(gdn.conv1d.bias.shape), mesh, distributed_setup.device, offset=7_000 + ) + conv_bias_exp_avg_sq, conv_bias_exp_avg_sq_dtensor = _make_sharded_tensor( + tuple(gdn.conv1d.bias.shape), mesh, distributed_setup.device, offset=8_000 + ) + untouched = torch.tensor([9_000.0], device=distributed_setup.device) + + model_state = { + "gdn.in_proj.weight": weight_dtensor, + "gdn.conv1d.weight": conv_weight_dtensor, + "gdn.conv1d.bias": conv_bias_dtensor, + } + optimizer_state = { + "state": { + "module.gdn.in_proj.weight": _optimizer_entry(exp_avg_dtensor, exp_avg_sq_dtensor), + "module.gdn.conv1d.weight": _optimizer_entry( + conv_weight_exp_avg_dtensor, conv_weight_exp_avg_sq_dtensor + ), + "module.gdn.conv1d.bias": _optimizer_entry( + conv_bias_exp_avg_dtensor, conv_bias_exp_avg_sq_dtensor + ), + "module.unrelated.weight": {"step": untouched}, + } + } + + split_model_state, split_optimizer_state = handle_gdn_in_state_dict( + model, model_state, optimizer_state + ) + + assert "gdn.in_proj.weight" not in split_model_state + assert "gdn.conv1d.weight" not in split_model_state + assert "gdn.conv1d.bias" not in split_model_state + _assert_dtensor_splits(split_model_state, "gdn.in_proj.weight", split_names, weight) + _assert_dtensor_splits( + split_model_state, "gdn.conv1d.weight", ("query", "key", "value"), conv_weight + ) + _assert_dtensor_splits( + split_model_state, "gdn.conv1d.bias", ("query", "key", "value"), conv_bias + ) + + split_optimizer = split_optimizer_state["state"] + assert "module.gdn.in_proj.weight" not in split_optimizer + _assert_dtensor_splits( + {key: value["exp_avg"] for key, value in split_optimizer.items() if "exp_avg" in value}, + "module.gdn.in_proj.weight", + split_names, + exp_avg, + ) + _assert_dtensor_splits( + { + key: value["exp_avg_sq"] + for key, value in split_optimizer.items() + if "exp_avg_sq" in value + }, + "module.gdn.in_proj.weight", + split_names, + exp_avg_sq, + ) + optimizer_references = { + "module.gdn.conv1d.weight": (conv_weight_exp_avg, conv_weight_exp_avg_sq), + "module.gdn.conv1d.bias": (conv_bias_exp_avg, conv_bias_exp_avg_sq), + } + for key, (expected_exp_avg, expected_exp_avg_sq) in optimizer_references.items(): + _assert_dtensor_splits( + { + name: value["exp_avg"] + for name, value in split_optimizer.items() + if "exp_avg" in value + }, + key, + ("query", "key", "value"), + expected_exp_avg, + ) + _assert_dtensor_splits( + { + name: value["exp_avg_sq"] + for name, value in split_optimizer.items() + if "exp_avg_sq" in value + }, + key, + ("query", "key", "value"), + expected_exp_avg_sq, + ) + + gdn_states = [value for key, value in split_optimizer.items() if key.startswith("module.gdn.")] + assert all(value["step"].item() == 1.0 for value in gdn_states) + assert split_optimizer["module.unrelated.weight"]["step"] is untouched 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..72cc829df55 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 @@ -109,8 +109,10 @@ def _fn(variant): [ ("gdn", True), ("gdn2", True), + ("kda", True), ("gated_delta_net", True), ("dsa", False), + ("dsv4_hybrid", False), (None, False), ("some_unknown_variant", False), ], @@ -120,6 +122,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 +273,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 ( @@ -421,6 +467,7 @@ class TestGetExperimentalAttentionVariantModuleSpec: [ ("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"), ], diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index c687d2bc7a6..808ad06a431 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -35,6 +35,7 @@ "activation_func_fp8_input_store": False, "add_bias_linear": False, "add_qkv_bias": False, + "apply_dsa_kernel_fusion": None, "apply_query_key_layer_scaling": False, "apply_residual_connection_post_layernorm": False, "apply_rope_fusion": False, @@ -137,7 +138,10 @@ "fp8_wgrad": True, "fused_residual_rmsnorm": False, "fused_single_qkv_rope": False, + "gated_attention_proj_granularity": "elementwise", "gated_linear_unit": False, + "gdn_conv_pad_alignment": None, + "gdn_pre_gated_delta_rule_fusion": False, "gdp_cutedsl_kernel": False, "gdp_num_chunk_states_to_recompute": 2, "gdp_num_householder": 3, @@ -152,6 +156,7 @@ "heterogeneous_block_specs": False, "hidden_dropout": 0.0, "hidden_size": 2688, + "head_wise_attn_gate": False, "hierarchical_context_parallel_sizes": None, "high_priority_a2a_comm_stream": False, "inference_fuse_tp_communication": False, @@ -161,12 +166,16 @@ "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, "layernorm_zero_centered_gamma": False, "linear_attention_freq": None, + "linear_attention_type": None, "linear_cp_layout": "zigzag", + "linear_cp_mode": "headwise", "linear_conv_kernel_dim": 4, "linear_key_head_dim": 128, "linear_num_key_heads": 16, @@ -296,7 +305,9 @@ "recompute_method": None, "recompute_modules": ["core_attn"], "recompute_num_layers": None, + "rotary_base_per_layer": None, "rotary_interleaved": False, + "sequence_packing_scheduler": None, "sequence_parallel": True, "softmax_scale": None, "softmax_type": "vanilla", 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..2281ac56290 --- /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( # DevSkim: ignore DS148264 -- synthetic test data. + (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( # DevSkim: ignore DS148264 -- synthetic test data. + (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..b0384147404 --- /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( # DevSkim: ignore DS148264 -- synthetic test data. + (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..d3c659ab12f 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -11,7 +11,7 @@ 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 @@ -274,6 +274,21 @@ def test_gdn_inference_spec(self): 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.""" layer_pattern = Symbols.GDN + Symbols.ATTENTION + Symbols.MAMBA diff --git a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py index 8b4c181ee30..2869b0c6bc0 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: @@ -317,75 +318,103 @@ class TestGetHybridLayerCounts: def test_simple_pattern(self): assert get_hybrid_layer_counts("M*M*") == { '*': 2, + 'C': 0, 'D': 0, 'G': 0, + 'H': 0, + 'K': 0, 'M': 2, '+': 0, '-': 0, 'E': 0, + 'W': 0, } 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") == { '*': 1, + 'C': 0, 'D': 0, 'G': 1, + 'H': 0, + 'K': 0, 'M': 1, '+': 0, '-': 1, 'E': 1, + 'W': 0, } assert get_hybrid_layer_counts("MGD-E") == { '*': 0, + 'C': 0, 'D': 1, 'G': 1, + 'H': 0, + 'K': 0, 'M': 1, '+': 0, '-': 1, 'E': 1, + 'W': 0, } assert get_hybrid_layer_counts("MG+-E") == { '*': 0, + 'C': 0, 'D': 0, 'G': 1, + 'H': 0, + 'K': 0, 'M': 1, '+': 1, '-': 1, 'E': 1, + 'W': 0, } def test_with_pipes(self): # Pipes should be skipped in counting assert get_hybrid_layer_counts("M*|M*") == { '*': 2, + 'C': 0, 'D': 0, 'G': 0, + 'H': 0, + 'K': 0, 'M': 2, '+': 0, '-': 0, 'E': 0, + 'W': 0, } assert get_hybrid_layer_counts("M-M-|M-M*-") == { '*': 1, + 'C': 0, 'D': 0, 'G': 0, + 'H': 0, + 'K': 0, 'M': 4, '+': 0, '-': 4, 'E': 0, + 'W': 0, } def test_with_mtp(self): # MTP pattern "MM" repeated 2 depths -> 4 extra mamba layers assert get_hybrid_layer_counts("M*M*/MM/MM") == { '*': 2, + 'C': 0, 'D': 0, 'G': 0, + 'H': 0, + 'K': 0, 'M': 6, '+': 0, '-': 0, 'E': 0, + 'W': 0, } def test_with_pipes_and_mtp(self): @@ -393,91 +422,138 @@ def test_with_pipes_and_mtp(self): # MTP: MM x 2 depths -> +4 mamba assert get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") == { '*': 1, + 'C': 0, 'D': 0, 'G': 0, + 'H': 0, + 'K': 0, 'M': 8, '+': 0, '-': 4, 'E': 0, + 'W': 0, } def test_moe_pattern(self): assert get_hybrid_layer_counts("MEME") == { '*': 0, + 'C': 0, 'D': 0, 'G': 0, + 'H': 0, + 'K': 0, 'M': 2, '+': 0, '-': 0, 'E': 2, + 'W': 0, } 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") == { '*': 3, + 'C': 0, 'D': 0, 'G': 0, + 'H': 0, + 'K': 0, 'M': 7, '+': 0, '-': 0, 'E': 0, + 'W': 0, } def test_gdn_pattern(self): assert get_hybrid_layer_counts("GMGM") == { '*': 0, + 'C': 0, 'D': 0, 'G': 2, + 'H': 0, + 'K': 0, 'M': 2, '+': 0, '-': 0, 'E': 0, + 'W': 0, } def test_gdn_hybrid_pattern(self): # GDN + Mamba + Attention assert get_hybrid_layer_counts("G*GM*") == { '*': 2, + 'C': 0, 'D': 0, 'G': 2, + 'H': 0, + 'K': 0, 'M': 1, '+': 0, '-': 0, 'E': 0, + 'W': 0, } def test_dsa_pattern(self): assert get_hybrid_layer_counts("DMDM") == { '*': 0, + 'C': 0, 'D': 2, 'G': 0, + 'H': 0, + 'K': 0, 'M': 2, '+': 0, '-': 0, 'E': 0, + 'W': 0, } def test_mla_pattern(self): assert get_hybrid_layer_counts("+M+M") == { '*': 0, + 'C': 0, 'D': 0, 'G': 0, + 'H': 0, + 'K': 0, 'M': 2, '+': 2, '-': 0, 'E': 0, + 'W': 0, + } + + def test_kda_pattern(self): + assert get_hybrid_layer_counts("K+K/K") == { + '*': 0, + 'C': 0, + 'D': 0, + 'G': 0, + 'H': 0, + 'K': 3, + 'M': 0, + '+': 1, + '-': 0, + 'E': 0, + 'W': 0, } def test_empty_pattern(self): assert get_hybrid_layer_counts("") == { '*': 0, + 'C': 0, 'D': 0, 'G': 0, + 'H': 0, + 'K': 0, 'M': 0, '+': 0, '-': 0, 'E': 0, + 'W': 0, } @@ -761,7 +837,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) diff --git a/tests/unit_tests/test_argument_utils.py b/tests/unit_tests/test_argument_utils.py index 57a327cddfd..80a034833c5 100644 --- a/tests/unit_tests/test_argument_utils.py +++ b/tests/unit_tests/test_argument_utils.py @@ -677,6 +677,24 @@ 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_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 TestMegatronMixedPrecisionArguments: """Test language-model logit dtype CLI choices.""" diff --git a/tests/unit_tests/test_num_floating_point_operations.py b/tests/unit_tests/test_num_floating_point_operations.py index df5e4191843..c409fb9603c 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 @@ -100,6 +101,32 @@ def _make_hybrid_args(*, num_layers=4, hidden_size=512, num_attention_heads=8, s 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 + + class TestBSHDBackwardCompat: """For unpacked BSHD, the new optional arg must not change the result.""" @@ -282,10 +309,159 @@ def test_householder_count_changes_flops(self): ) num_gdp_layers = 2 expected_delta = 3 * num_gdp_layers * forward_delta_per_layer - assert flops_m4 - flops_m3 == expected_delta +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. + """ + + @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 + + assert num_floating_point_operations(standard, batch_size) == num_floating_point_operations( + hybrid, batch_size + ) + + @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 + 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 * (in_proj_dim + args.linear_num_key_heads) + + args.linear_conv_kernel_dim * (2 * qk_dim + v_dim) + + args.hidden_size * v_dim + ) + ) + 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 + + +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: """``total_real_tokens_in_batch`` removes padding from token-linear 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..86b44f2b9a2 100644 --- a/tests/unit_tests/transformer/test_attention.py +++ b/tests/unit_tests/transformer/test_attention.py @@ -480,6 +480,9 @@ def _test_parallel_attention_correctness( sequence_length=256, micro_batch_size=4, sequence_packing=False, + cp_partition_mode="zigzag", + input_grad_atol=None, + input_grad_rtol=None, ): # Model initialization function def initialize_gpt_model( @@ -559,9 +562,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'): @@ -585,6 +590,7 @@ 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) @@ -640,8 +646,8 @@ def get_tensor_on_this_rank(tensor): torch.testing.assert_close( input_grad_baseline, input_grad_parallel, - atol=atol, - rtol=rtol, + atol=atol if input_grad_atol is None else input_grad_atol, + rtol=rtol if input_grad_rtol is None else input_grad_rtol, msg=lambda msg: f"Mismatch in input_grad: {msg}", ) if has_bias: diff --git a/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py b/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py index 0655e865c86..d9e4221a5ce 100644 --- a/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py +++ b/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py @@ -16,13 +16,14 @@ Tests cover: 1. Per-TransformerLayer gated_linear_unit check in handle_swiglu_in_state_dict - 2. GDN fused projection splitting in handle_gdn_in_state_dict + 2. GDN-family fused projection splitting in handle_gdn_in_state_dict 3. Helper functions: get_expert_index_from_key, flatten_state_dict, etc. Note: handle_swiglu_in_state_dict and handle_gdn_in_state_dict require HAVE_MEGATRON_FSDP=True and a real distributed environment with DTensors. We test their internal helper logic (is_swiglu_key, _key_in_glu_layer, -_match_gdn_key) by extracting the logic into standalone testable units. +projection metadata validation) directly and cover real DTensor splitting in the +distributed M-FSDP test suite. For the checkpoint_inspector.py conversion functions, we test the non-distributed components: SWiGLU key detection, SWiGLU regex splitting, MTP key renaming, and --swiglu-modules prefix filtering. @@ -57,6 +58,7 @@ get_mla_fused_down_proj_splits, get_mtp_inner_layer_paths, get_unexpected_model_keys, + handle_gdn_in_state_dict, handle_mla_down_proj_in_state_dict, handle_mtp_in_state_dict, match_mla_fused_down_proj_key, @@ -483,6 +485,89 @@ def test_strip_wrappers(self): # ============================================================================ # Test GDN key matching logic # ============================================================================ +class _FakeGDNMetadata(torch.nn.Module): + def __init__(self, names, sections): + super().__init__() + self.qk_dim = 4 + self.v_dim = 6 + self.tp_size = 1 + self.in_proj_dim = sum(sections) + self.in_proj_split_names = names + self.in_proj_split_sections = sections + + +class TestGDNProjectionMetadata: + @staticmethod + def _preprocess(module): + return handle_gdn_in_state_dict(module, {}, None) + + @pytest.mark.parametrize( + ("names", "sections"), + [ + (("query", "key", "value", "z", "beta", "alpha"), (4, 4, 6, 6, 1, 1)), + (("query", "key", "value", "g", "gate"), (4, 4, 6, 4, 6)), + (("query", "key", "value", "z", "f", "b", "w"), (4, 4, 6, 6, 4, 4, 6)), + ], + ) + def test_accepts_variant_specific_metadata(self, names, sections): + model_state, optimizer_state = self._preprocess(_FakeGDNMetadata(names, sections)) + + assert model_state == {} + assert optimizer_state is None + + def test_rejects_incomplete_metadata(self): + module = _FakeGDNMetadata(("query",), (4,)) + del module.in_proj_split_sections + + with pytest.raises(ValueError, match="must define both"): + self._preprocess(module) + + def test_rejects_duplicate_names(self): + module = _FakeGDNMetadata(("query", "query"), (4, 4)) + + with pytest.raises(ValueError, match="duplicate"): + self._preprocess(module) + + @pytest.mark.parametrize( + ("names", "sections", "error"), + [ + (("query", "key"), (4,), "sizes for 2 names"), + (("query", ""), (4, 4), "invalid in_proj split names"), + (("query", "key"), (4, 0), "invalid in_proj split sizes"), + (("query", "key"), (4, True), "invalid in_proj split sizes"), + ], + ) + def test_rejects_invalid_metadata(self, names, sections, error): + module = _FakeGDNMetadata(names, sections) + + with pytest.raises(ValueError, match=error): + self._preprocess(module) + + def test_rejects_incorrect_total_size(self): + module = _FakeGDNMetadata(("query", "key"), (4, 4)) + module.in_proj_dim += 1 + + with pytest.raises(ValueError, match="expected"): + self._preprocess(module) + + @pytest.mark.parametrize("tp_size", [0, True]) + def test_rejects_invalid_tp_size(self, tp_size): + module = _FakeGDNMetadata(("query", "key"), (4, 4)) + module.tp_size = tp_size + + with pytest.raises(ValueError, match="invalid tp_size"): + self._preprocess(module) + + def test_rejects_dimensions_not_divisible_by_tp(self): + module = _FakeGDNMetadata(("query", "key"), (2, 2)) + module.tp_size = 2 + module.in_proj_dim = 8 + module.qk_dim = 5 + + with pytest.raises(ValueError, match="qk_dim/v_dim must be divisible"): + self._preprocess(module) + + class TestGDNKeyMatching: """Test the _match_gdn_key logic from handle_gdn_in_state_dict. 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, ) From 0317e417356dc1afb9784aac4575e3f809072956 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Thu, 27 Aug 2026 04:08:37 +0000 Subject: [PATCH 2/4] style: format GDN context parallel flow Signed-off-by: svcnemo-autobot --- megatron/core/ssm/gated_delta_net/gdn.py | 28 +++++++++--------------- 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py index f3d4f1df611..33da90630c6 100644 --- a/megatron/core/ssm/gated_delta_net/gdn.py +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -161,7 +161,9 @@ def forward( ) input_layout_state = layout_manager.build_forward_state(layout_packed_seq_params) if input_layout_state is not None: - hidden_states, packed_seq_params = input_layout_state.prepare_layer(0, hidden_states) + hidden_states, packed_seq_params = input_layout_state.prepare_layer( + 0, hidden_states + ) seq_len_local, batch, _ = hidden_states.shape seq_len_post_headwise = seq_len_local * self.sp_size * cp_size_headwise @@ -204,7 +206,9 @@ def forward( 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" - assert not self.config.deterministic_mode, "Packed sequence does not support deterministic mode." + assert ( + not self.config.deterministic_mode + ), "Packed sequence does not support deterministic mode." cu_seqlens_q = self._resolve_cu_seqlens( packed_seq_params.cu_seqlens_q_padded, packed_seq_params.cu_seqlens_q, @@ -270,17 +274,11 @@ def _compute(hidden_states): nvtx_range_push(suffix="conv1d") qkv_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_sections, + self.conv1d.weight, dim=0, cp_group=cp_group_headwise, split_sections=qkv_sections ) conv1d_bias = ( get_parameter_local_cp( - self.conv1d.bias, - dim=0, - cp_group=cp_group_headwise, - split_sections=qkv_sections, + self.conv1d.bias, dim=0, cp_group=cp_group_headwise, split_sections=qkv_sections ) if self.conv_bias else None @@ -340,15 +338,9 @@ def _compute(hidden_states): def _norm_and_restore(core_attn_out, gate): norm_out = self._apply_gated_norm(core_attn_out, gate) - norm_out = ( - norm_out.reshape(batch, kernel_seq_len, -1).transpose(0, 1).contiguous() - ) + norm_out = norm_out.reshape(batch, kernel_seq_len, -1).transpose(0, 1).contiguous() return a2a_hp_to_cp( - norm_out, - cp_size_headwise, - cp_group_headwise, - packed_seq_params, - thd_cp_a2a_inv, + norm_out, cp_size_headwise, cp_group_headwise, packed_seq_params, thd_cp_a2a_inv ) if self.recompute_norm_out and self.training: From 9aa15bc265f00155a315c58c6ad91145270cace6 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Thu, 27 Aug 2026 08:51:27 +0000 Subject: [PATCH 3/4] fix: shard GDN2 kernel inputs for context parallelism Signed-off-by: svcnemo-autobot --- megatron/core/ssm/gated_delta_net/gdn2.py | 11 ++++++++++- tests/unit_tests/ssm/test_gated_delta_net.py | 9 ++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net/gdn2.py b/megatron/core/ssm/gated_delta_net/gdn2.py index 62ad3ef0f32..14d462095e1 100644 --- a/megatron/core/ssm/gated_delta_net/gdn2.py +++ b/megatron/core/ssm/gated_delta_net/gdn2.py @@ -292,7 +292,16 @@ def forward( # 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, f, b, w + qkv, + gate, + A_log_local_cp, + dt_bias_local_cp, + batch, + seq_len, + f, + b, + w, + cp_size_headwise=self.cp_size, ) gate = kernel_inputs.pop("gate") nvtx_range_pop(suffix="prepare_input_for_gated_delta_rule") diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index b63420b2fa9..bbff83fd15e 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -371,7 +371,14 @@ def test_jit_compiled_helpers(self): # which are normally wrapped by @jit_fuser (torch.compile). 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 + qkv, + gate, + A_log_mock, + dt_bias_mock, + batch, + seq_len, + *gate_feats, + cp_size_headwise=gdn.cp_size, ) # The output gate (z) rides along under "gate" and is popped by forward before From acbd503fd7fd3747b0af087c0b57caff20e78ade Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Thu, 27 Aug 2026 09:03:34 +0000 Subject: [PATCH 4/4] test: remove unsupported dev-only GDN cases Signed-off-by: svcnemo-autobot --- .../ssm/gated_delta_net/test_gdn.py | 92 ------------------- 1 file changed, 92 deletions(-) diff --git a/tests/unit_tests/ssm/gated_delta_net/test_gdn.py b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py index 2281ac56290..5fad11da21b 100644 --- a/tests/unit_tests/ssm/gated_delta_net/test_gdn.py +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py @@ -3,7 +3,6 @@ import copy import inspect import os -from unittest import mock import pytest import torch @@ -35,13 +34,6 @@ # 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 = { @@ -425,26 +417,6 @@ def test_gpu_forward_rejects_sbhd_chunkwise_cp_batch_gt_one(self): 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( @@ -606,70 +578,6 @@ def test_jit_compiled_helpers(self): 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.")