From 139b583801e67aa6a2645c8b95eb0e1db0ae8f6a Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Tue, 18 Aug 2026 00:55:01 +0800 Subject: [PATCH 01/34] Prepare the Gated DeltaNet base for channel-wise variants Keep the existing private GDN base as the shared implementation for head-wise GatedDeltaNet and channel-wise variants. Generalize its variant hooks, parameter dtype handling, diagnostics, and norm-output recompute contract without exposing a broader linear-attention API. Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 6ff2fcc3e4f068f0bb00b196db1ced0c8dd47542) --- megatron/core/ssm/gated_delta_net/common.py | 187 +- megatron/core/ssm/gated_delta_net/gdn.py | 720 ++++--- .../core/transformer/transformer_config.py | 1663 ++++++++++------- .../ssm/gated_delta_net/test_gdn.py | 830 ++++++++ 4 files changed, 2381 insertions(+), 1019 deletions(-) create mode 100644 tests/unit_tests/ssm/gated_delta_net/test_gdn.py diff --git a/megatron/core/ssm/gated_delta_net/common.py b/megatron/core/ssm/gated_delta_net/common.py index 63f2d244840..30e0cfa5d9b 100644 --- a/megatron/core/ssm/gated_delta_net/common.py +++ b/megatron/core/ssm/gated_delta_net/common.py @@ -1,20 +1,16 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh. # Some of this code was adopted from https://github.com/huggingface/transformers # This source code is licensed under the Apache license found in the # LICENSE file in the root directory of this source tree. -# pylint: disable=unused-import - -import logging from dataclasses import dataclass from functools import lru_cache -from typing import Callable, Optional, Protocol, Union +from typing import Optional, Protocol, Union import torch import torch.nn as nn -import torch.nn.functional as F from megatron.core.fp8_utils import get_fp8_align_size from megatron.core.inference.contexts import BaseInferenceContext @@ -38,29 +34,44 @@ make_sharded_tensors_for_checkpoint, sharded_state_dict_default, ) -from megatron.core.utils import nvtx_range_pop, nvtx_range_push try: from fla.modules.convolution import causal_conv1d from fla.modules.l2norm import l2norm + from fla.ops.cp import build_cp_context from fla.ops.gated_delta_rule import chunk_gated_delta_rule HAVE_FLA = True except ImportError: + build_cp_context = None causal_conv1d = None l2norm = None chunk_gated_delta_rule = None HAVE_FLA = False -logger = logging.getLogger(__name__) +__all__ = [ + "HAVE_FLA", + "GatedDeltaNetSubmodules", + "_GDNBase", + "_build_head_perm_for_split_sections", + "_build_thd_cp_a2a_perm", + "_split_tensor_factory", + "a2a_cp_to_hp", + "a2a_hp_to_cp", + "build_cp_context", + "causal_conv1d", + "chunk_gated_delta_rule", + "get_parameter_local_cp", + "l2norm", + "tensor_a2a_cp2hp", + "tensor_a2a_hp2cp", +] @dataclass class GatedDeltaNetSubmodules: - """ - Contains the module specs for the input linear, output norm, and output linear layers. - """ + """Module specs shared by GDN-family layers.""" in_proj: Union[ModuleSpec, type] = IdentityOp out_norm: Union[ModuleSpec, type] = IdentityOp @@ -68,12 +79,7 @@ class GatedDeltaNetSubmodules: class GatedDeltaRuleInterface(Protocol): - """ - Unified typing protocol for linear attention interfaces, compliant to upstream FLA interfaces. - - Only ``q``/``k``/``v``/``g`` are common to every kernel, and only as keywords: each - variant inserts its own gates after ``g`` (e.g., ``beta`` for GDN, ``b``/``w`` for GDN2). - """ + """Callable interface shared by GDN-family kernels.""" def __call__( self, @@ -81,6 +87,7 @@ def __call__( k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, + beta: torch.Tensor, *, scale: float | None = None, initial_state: torch.Tensor | None = None, @@ -92,16 +99,16 @@ def __call__( class _GDNBase(MegatronModule): - """Common base class for the Gated Delta Net (GDN) family of layers. + """Shared implementation for the GDN-family layers. - Hosts everything the GDN variants share: the fused input projection, causal - convolution on q/k/v, the CP all-to-all plumbing, the kernel-input preparation - skeleton, the gated output norm + projection, and sharded checkpointing. + Provides the projection, Q/K/V causal convolution, gated delta-rule parameters, + context-parallel layout handling, gated output normalization, output projection, + and sharded-checkpoint plumbing shared by head-wise GatedDeltaNet and channel-wise + Kimi Delta Attention. """ dt_bias_dim: int a_log_dim: int - in_proj_qkvg_dim: int in_proj_extra_dim: int in_proj_dim: int @@ -124,13 +131,14 @@ def __init__( *, name: str | None = None, cp_comm_type: str | None = None, - pp_layer_offset: int = 0, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, ): """ Args: config: The config of the model. submodules: Contains the module specs for the input and output linear layers. - layer_number: The layer number of this GDN layer. + layer_number: The layer number of this GDN-family layer. bias: Whether to use bias in the linear layers. conv_bias: Whether to use bias in the causal convolution. conv_init: The initialization range for the causal convolution weights. @@ -138,11 +146,13 @@ def __init__( A_init_range: The initialization range for the attention weights. pg_collection: The required process groups to use for tensor model parallel and context parallel. - name (str | None): module instance name passed top-down from its paranet module + name (str | None): Optional module path prefix used for child module names. cp_comm_type (Optional[str]): Accepted for TransformerLayer compatibility and ignored; GDN implements context parallelism with its own all-to-alls rather than the attention CP communication schemes. - pp_layer_offset: Offset of this pipeline stage's first global layer. + pp_layer_offset (Optional[int]): Pipeline layer offset forwarded by + TransformerLayer. Stored for MTP/TransformerLayer API compatibility. + is_mtp_layer (bool): Whether this module is inside an MTP prediction depth. """ if not HAVE_FLA: raise ImportError( @@ -154,16 +164,19 @@ def __init__( # Attributes from arguments self.layer_number = layer_number - self.pp_layer_offset = pp_layer_offset + self._pp_layer_offset = pp_layer_offset + self.is_mtp_layer = is_mtp_layer self.bias = bias self.conv_bias = conv_bias self.conv_init = conv_init assert A_init_range[0] >= 0 and A_init_range[1] >= A_init_range[0] self.A_init_range = A_init_range self.use_qk_l2norm = use_qk_l2norm - assert pg_collection is not None, "pg_collection must be provided for GatedDeltaNet" + assert pg_collection is not None, "pg_collection must be provided for a GDN-family layer" self.pg_collection = pg_collection self.tp_group = pg_collection.tp + # Static/max CP size from model construction. Runtime dynamic CP paths must resolve + # the effective group from packed_seq_params instead of using this value. self.cp_size = self.pg_collection.cp.size() self.tp_size = self.pg_collection.tp.size() self.sp_size = self.tp_size if config.sequence_parallel else 1 @@ -183,6 +196,21 @@ def __init__( self.qk_dim_local_tp = self.qk_dim // self.tp_size self.v_dim_local_tp = self.v_dim // self.tp_size + # Headwise CP shards heads over the CP group; chunkwise CP keeps heads local. + if self.config.linear_cp_mode == "headwise": + num_key_heads_per_tp = self.num_key_heads // self.tp_size + num_value_heads_per_tp = self.num_value_heads // self.tp_size + assert num_key_heads_per_tp % self.cp_size == 0, ( + f"GDN-family head-parallel CP requires the static (max) cp_size ({self.cp_size}) " + f"to evenly divide num_key_heads per TP rank ({num_key_heads_per_tp}); " + f"all runtime dynamic cp_size values divide the static one and so will also divide." + ) + assert num_value_heads_per_tp % self.cp_size == 0, ( + f"GDN-family head-parallel CP requires the static (max) cp_size ({self.cp_size}) " + f"to evenly divide num_value_heads per TP rank ({num_value_heads_per_tp}); " + f"all runtime dynamic cp_size values divide the static one and so will also divide." + ) + self.num_v_heads_local_tp = self.num_value_heads // self.tp_size self.num_k_heads_local_tp = self.num_key_heads // self.tp_size @@ -192,20 +220,21 @@ def __init__( "in_proj_extra_dim", "in_proj_split_names", "in_proj_split_sections", - "feat_dim_split", "gated_delta_rule", ) self._setup_variant_attrs() for attr in attrs_to_check: - assert getattr(self, attr, None) is not None, f"Attribute {attr} for GDN is not set" - # QK, V, gate, shared across all variants - self.in_proj_qkvg_dim = self.qk_dim * 2 + self.v_dim * 2 - self.in_proj_dim = self.in_proj_qkvg_dim + self.in_proj_extra_dim + assert hasattr(self, attr), f"Attribute {attr} for the GDN-family variant is not set" + assert ( + getattr(self, attr) is not None + ), f"Attribute {attr} for the GDN-family variant is not set" + # Full input projection width: q, k, v, output gate, and variant-specific gate features. + self.in_proj_dim = self.qk_dim * 2 + self.v_dim * 2 + self.in_proj_extra_dim if self.config.fp8: fp8_align_size = get_fp8_align_size(self.config.fp8_recipe) assert self.in_proj_dim % fp8_align_size == 0, ( - "For FP8, the innermost dimension of the GDN layer " + "For FP8, the innermost dimension of the GDN-family layer " "input projection output tensor must be a multiple of 16." ) self.in_proj = build_module( @@ -245,7 +274,9 @@ def __init__( self.dt_bias = nn.Parameter( torch.empty( - self.dt_bias_dim, dtype=self.config.params_dtype, device=torch.cuda.current_device() + self.dt_bias_dim, + dtype=getattr(self, "gate_params_dtype", self.config.params_dtype), + device=torch.cuda.current_device(), ) ) setattr(self.dt_bias, "tensor_model_parallel", True) @@ -253,7 +284,9 @@ def __init__( self.A_log = nn.Parameter( torch.empty( - self.a_log_dim, dtype=self.config.params_dtype, device=torch.cuda.current_device() + self.a_log_dim, + dtype=getattr(self, "gate_params_dtype", self.config.params_dtype), + device=torch.cuda.current_device(), ) ) setattr(self.A_log, "tensor_model_parallel", True) @@ -268,8 +301,10 @@ def __init__( ) self.recompute_norm_out = False self.norm_out_checkpoint = None - if self.config.recompute_granularity == "selective": + self.recompute_gdn = False + if self.config.recompute_granularity == "selective" and self.config.recompute_modules: self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules + self.recompute_gdn = "gdn" in self.config.recompute_modules self.out_proj = build_module( submodules.out_proj, @@ -285,17 +320,20 @@ def __init__( tp_group=self.pg_collection.tp, name=(name + ".out_proj") if name is not None else None, ) + # TODO: Packed sequence cu_seqlens can vary per batch; cache only static SBHD + # cp_context entries here and revisit routing metadata lifetime in the CP layout refactor. + self._chunkwise_cp_context_cache: dict[tuple[int, int], tuple[torch.Tensor, object]] = {} self.reset_parameters() def _setup_variant_attrs(self): - """Set variant specifics on the module. Called once from ``__init__``. + """Set variant projection sections, gate parameter sizes, and kernel callable. Must set: - - ``in_proj_dim`` + - ``in_proj_extra_dim`` (the in_proj sections beyond q/k/v/z; the base + class derives ``in_proj_dim`` from it) - ``in_proj_split_names`` - ``in_proj_split_sections`` - - ``feat_dim_split`` - ``dt_bias_dim`` / ``a_log_dim`` (sizes of the gate parameters, which the base class creates after the conv1d module to preserve the original parameter registration order) @@ -306,8 +344,8 @@ def _setup_variant_attrs(self): def _reset_dt_bias(self): """Initialize ``dt_bias``. Called from ``reset_parameters`` under the RNG tracker. - Defaults to ones; variants whose kernel expects a different step-size - parametrization override this. + Defaults to ones; subclasses can override this if their kernel expects a + different step-size parametrization. """ torch.ones( self.dt_bias_dim, @@ -344,29 +382,6 @@ def forward( # pylint: disable=missing-function-docstring raise NotImplementedError - def _gated_norm_and_a2a( - self, - core_attn_out: torch.Tensor, - gate: torch.Tensor, - thd_cp_a2a_inv: torch.Tensor | None, - batch: int, - seq_len: int, - packed_seq_params: PackedSeqParams | None = None, - ) -> torch.Tensor: - # RMSNorm - nvtx_range_push(suffix="gated_norm") - norm_out_hp = self._apply_gated_norm(core_attn_out, gate) - nvtx_range_pop(suffix="gated_norm") - - # Transpose: b s x --> s b x - # From bshd back to sbhd format - norm_out_hp = norm_out_hp.reshape(batch, seq_len, -1) - norm_out_hp = norm_out_hp.transpose(0, 1).contiguous() - - return a2a_hp_to_cp( - norm_out_hp, self.cp_size, self.pg_collection.cp, packed_seq_params, thd_cp_a2a_inv - ) - @jit_fuser def _apply_gated_norm(self, x, gate): # Output Norm @@ -389,24 +404,25 @@ def _prepare_input_for_gated_delta_rule( batch: int, seq_len: int, *gate_feats: tuple[torch.Tensor], + cp_size_headwise: int | None = None, ) -> dict[str, torch.Tensor]: """ Prepare all gated delta rule kernel inputs. Fuses split, reshape, L2 norm, decay/gate activations, repeat_interleave, and - contiguous operations. ``gate_feats`` holds the variant-specific in_proj - sections, which ``_compute_gates`` turns into the decay and gating tensors. + contiguous operations. ``gate_feats`` holds the in_proj sections after qkv + and gate, which ``_compute_gates`` turns into the decay and gating tensors. Returns: (dict[str, Tensor]): Kernel inputs keyed by kernel argument name (``q``, - ``k``, ``v``, ``g``, plus the variant-specific gates), and the output + ``k``, ``v``, ``g``, and ``beta``), and the output gate (z) tensor under the ``gate`` key, which is not a kernel input. """ + cp_size = 1 if cp_size_headwise is None else cp_size_headwise + # Split qkv into query_key and value query_key, value = torch.split( - qkv, - [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size], - dim=-1, + qkv, [2 * self.qk_dim_local_tp // cp_size, self.v_dim_local_tp // cp_size], dim=-1 ) # Reshape query_key and value @@ -418,7 +434,7 @@ def _prepare_input_for_gated_delta_rule( query_key = l2norm(query_key.contiguous()) # Split query and key - split_size = self.qk_dim_local_tp // self.key_head_dim // self.cp_size + split_size = self.qk_dim_local_tp // self.key_head_dim // cp_size query, key = torch.split(query_key, [split_size, split_size], dim=2) # Expand query and key if needed (grouped query attention) @@ -450,19 +466,18 @@ def _compute_gates( *gate_feats: tuple[torch.Tensor], ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """ - Compute the log-decay ``g`` and the variant-specific kernel inputs. + Compute the log-decay ``g`` and remaining kernel inputs. Args: A_log_local_cp: CP-local slice of ``A_log``. dt_bias_local_cp: CP-local slice of ``dt_bias``. batch: Batch size. seq_len: Sequence length. - gate_feats: The variant-specific in_proj output sections (everything after - the qkv and output-gate sections, in ``feat_dim_split`` order). + gate_feats: The in_proj output sections after the qkv and output-gate sections. Returns: (tuple[Tensor, dict[str, Tensor]]): The log-decay ``g`` and a dict of the - remaining variant-specific kernel inputs keyed by kernel argument name. + remaining kernel inputs keyed by kernel argument name. """ raise NotImplementedError @@ -647,7 +662,7 @@ def _build_head_perm_for_split_sections( def get_parameter_local_cp( param: torch.Tensor, dim: int, - cp_group: torch.distributed.ProcessGroup, + cp_group: torch.distributed.ProcessGroup | None, split_sections: Optional[list[int]] = None, ) -> torch.Tensor: """Get the local parameter for the current context parallel rank. @@ -665,13 +680,15 @@ def get_parameter_local_cp( torch.Tensor: The local parameter for the current context parallel rank. """ - cp_size = cp_group.size() - cp_rank = cp_group.rank() + cp_size = cp_group.size() if cp_group is not None else 1 # No need to split if CP size is 1. if cp_size == 1: return param + assert cp_group is not None + cp_rank = cp_group.rank() + # Split first if needed. if split_sections is not None: inputs = torch.split(param, split_sections, dim=dim) @@ -693,7 +710,7 @@ def tensor_a2a_cp2hp( tensor: torch.Tensor, seq_dim: int, head_dim: int, - cp_group: torch.distributed.ProcessGroup, + cp_group: torch.distributed.ProcessGroup | None, split_sections: Optional[list[int]] = None, undo_attention_load_balancing: bool = True, ): @@ -714,12 +731,14 @@ def tensor_a2a_cp2hp( torch.Tensor: The all-to-all tensor. """ - cp_size = cp_group.size() + cp_size = cp_group.size() if cp_group is not None else 1 # No need to all-to-all if CP size is 1. if cp_size == 1: return tensor + assert cp_group is not None + # Limitations of mamba_context_parallel._all_to_all_cp2hp. assert seq_dim == 0, f"tensor_a2a_cp2hp only supports seq_dim == 0 for now, but got {seq_dim=}" assert ( @@ -756,7 +775,7 @@ def tensor_a2a_hp2cp( tensor: torch.Tensor, seq_dim: int, head_dim: int, - cp_group: torch.distributed.ProcessGroup, + cp_group: torch.distributed.ProcessGroup | None, split_sections: Optional[list[int]] = None, redo_attention_load_balancing: bool = True, ): @@ -777,12 +796,14 @@ def tensor_a2a_hp2cp( torch.Tensor: The all-to-all tensor. """ - cp_size = cp_group.size() + cp_size = cp_group.size() if cp_group is not None else 1 # No need to all-to-all if CP size is 1. if cp_size == 1: return tensor + assert cp_group is not None + # Limitations of mamba_context_parallel._all_to_all_hp2cp. assert seq_dim == 0, f"tensor_a2a_hp2cp only supports seq_dim == 0 for now, but got {seq_dim=}" assert ( diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py index c7a983841cb..15ecc4485c9 100644 --- a/megatron/core/ssm/gated_delta_net/gdn.py +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh. # Some of this code was adopted from https://github.com/huggingface/transformers @@ -12,35 +12,36 @@ import torch.nn.functional as F from megatron.core import tensor_parallel -from megatron.core.inference.contexts import BaseInferenceContext, DynamicInferenceContext -from megatron.core.inference.contexts.attention_context.triton.tensor_ops import ( - tensor_masked_update, -) +from megatron.core.context_parallel_layout import convert_module_input_tensors_cp_partition_mode +from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.jit import jit_fuser -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group +from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gated_delta_net.common import ( _GDNBase, a2a_cp_to_hp, + a2a_hp_to_cp, + build_cp_context, causal_conv1d, chunk_gated_delta_rule, get_parameter_local_cp, l2norm, ) -from megatron.core.ssm.ssm_inference import SSMDynamicInferenceMixin from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push -try: - from fla.modules.convolution import causal_conv1d_update - from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule -except ImportError: - causal_conv1d_update = None - fused_recurrent_gated_delta_rule = None +class GatedDeltaNet(_GDNBase): + """Gated DeltaNet with a head-wise scalar memory-decay gate.""" -class GatedDeltaNet(SSMDynamicInferenceMixin, _GDNBase): - # pylint: disable=missing-class-docstring def _setup_variant_attrs(self): """Set the GDN in_proj sizing, split tables, gate parameter dims, and kernel.""" + self.gdn_pre_gated_delta_rule_fusion = self.config.gdn_pre_gated_delta_rule_fusion + if self.config.deterministic_mode and self.gdn_pre_gated_delta_rule_fusion: + raise ValueError( + "Pre-GDR fusion is non-deterministic, but deterministic_mode=True. " + "Disable gdn_pre_gated_delta_rule_fusion or deterministic_mode." + ) + # alpha, beta self.in_proj_extra_dim = self.num_value_heads * 2 @@ -56,13 +57,6 @@ def _setup_variant_attrs(self): self.num_value_heads // self.tp_size, # beta self.num_value_heads // self.tp_size, # alpha ) - self.feat_dim_split = ( - (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // self.cp_size, # qkv - self.v_dim_local_tp // self.cp_size, # gate (z) - self.num_value_heads // self.tp_size // self.cp_size, # beta - self.num_value_heads // self.tp_size // self.cp_size, # alpha - ) - self.dt_bias_dim = self.num_v_heads_local_tp self.a_log_dim = self.num_v_heads_local_tp @@ -70,7 +64,15 @@ def _setup_variant_attrs(self): self.gated_delta_rule = torch_chunk_gated_delta_rule else: self.gated_delta_rule = chunk_gated_delta_rule - self.chunk_size = 64 + + def _get_feat_dim_split(self, cp_size_headwise: int) -> tuple[int, int, int, int]: + """Return GDN1 qkv/z/beta/alpha split sizes for a runtime headwise CP size.""" + return ( + (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size_headwise, + self.v_dim_local_tp // cp_size_headwise, + self.num_value_heads // self.tp_size // cp_size_headwise, + self.num_value_heads // self.tp_size // cp_size_headwise, + ) @jit_fuser def _compute_gates( @@ -96,9 +98,10 @@ def forward( packed_seq_params: Optional[PackedSeqParams] = None, sequence_len_offset: Optional[int] = None, *, + pg_collection: Optional[ProcessGroupCollection] = None, inference_params: Optional[BaseInferenceContext] = None, **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor | None]: + ): """ Perform a forward pass through the GDN module. @@ -108,30 +111,66 @@ def forward( inference_context = deprecate_inference_params(inference_context, inference_params) - seq_len, batch, _ = hidden_states.shape - seq_len = seq_len * self.sp_size * self.cp_size + active_pg_collection = pg_collection if pg_collection is not None else self.pg_collection + base_cp_group = active_pg_collection.cp + cp_group = resolve_cp_group(base_cp_group, packed_seq_params) + if self.config.linear_cp_mode == "chunkwise": + cp_group_chunkwise = cp_group + cp_group_headwise = None + elif self.config.linear_cp_mode == "headwise": + cp_group_chunkwise = None + cp_group_headwise = cp_group + elif cp_group.size() == 1: + cp_group_chunkwise = None + cp_group_headwise = None + else: + raise ValueError( + f"Unsupported linear_cp_mode {self.config.linear_cp_mode!r}; " + "expected 'headwise' or 'chunkwise'." + ) + cp_size_chunkwise = cp_group_chunkwise.size() if cp_group_chunkwise is not None else 1 + cp_size_headwise = cp_group_headwise.size() if cp_group_headwise is not None else 1 + cp_size_runtime = cp_group.size() + back_to_input_converter = None + if self.config.linear_cp_mode == "chunkwise": + hidden_states, back_to_input_converter = convert_module_input_tensors_cp_partition_mode( + hidden_states=hidden_states, + packed_seq_params=packed_seq_params, + cp_group=cp_group_chunkwise, + tp_group=self.tp_group, + tp_cp_group=getattr(active_pg_collection, "tp_cp", None), + target_partition_mode="contiguous", + sequence_parallel=self.config.sequence_parallel, + config=self.config, + ) + + seq_len_local, batch, _ = hidden_states.shape + seq_len_post_headwise = seq_len_local * self.sp_size * cp_size_headwise + seq_len_global = seq_len_post_headwise * cp_size_chunkwise if inference_context is not None: - if inference_context.is_dynamic_batching(): - assert ( - not self.config.deterministic_mode - ), "GDN dynamic inference requires the FLA recurrent kernels." - assert ( - not self.config.batch_invariant_mode - ), "GDN dynamic inference does not support batch-invariant mode." - assert ( - self.cp_size == 1 - ), "Context parallelism is not supported for GDN dynamic inference." - assert ( - inference_context.num_speculative_tokens == 0 - ), "GDN dynamic inference does not support speculative decoding." - assert ( - not inference_context.enable_prefix_caching - ), "GDN dynamic inference does not support prefix caching." - return self.ssm_dynamic_inference(hidden_states, inference_context) - assert inference_context.is_static_batching() + assert ( + inference_context.is_static_batching() + ), "GDN does not currently support dynamic inference batching." assert not self.config.sequence_parallel - raise NotImplementedError("GDN static-batching inference is not supported.") + # TODO: support inference + raise NotImplementedError("GDN does not support inference for now.") + + if cp_size_headwise > 1 and ( + ( + packed_seq_params is not None + and packed_seq_params.qkv_format == "thd" + and packed_seq_params.cp_partition_mode != "zigzag" + ) + or ( + (packed_seq_params is None or packed_seq_params.qkv_format != "thd") + and self.config.cp_partition_mode != "zigzag" + ) + ): + raise ValueError( + "GatedDeltaNet with headwise CP requires zigzag layout. CP partition " + "conversion must be handled before calling GatedDeltaNet." + ) if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': assert batch == 1, "Packed sequence expects batch dimension to be 1" @@ -143,16 +182,16 @@ def forward( cu_seqlens_q = self._resolve_cu_seqlens( packed_seq_params.cu_seqlens_q_padded, packed_seq_params.cu_seqlens_q, - seq_len, + seq_len_global, "cu_seqlens_q", - cp_size=self.cp_size, + cp_size=cp_size_runtime, ) cu_seqlens_kv = self._resolve_cu_seqlens( packed_seq_params.cu_seqlens_kv_padded, packed_seq_params.cu_seqlens_kv, - seq_len, + seq_len_global, "cu_seqlens_kv", - cp_size=self.cp_size, + cp_size=cp_size_runtime, ) assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " @@ -167,6 +206,86 @@ def forward( cu_seqlens_q = None cu_seqlens_kv = None + if cp_size_chunkwise > 1: + if cu_seqlens_q is None: + cache_key = (seq_len_global, batch) + cached = self._chunkwise_cp_context_cache.get(cache_key) + if cached is None: + cached_cu_seqlens = ( + torch.arange( + batch + 1, device=torch.cuda.current_device(), dtype=torch.long + ) + * seq_len_global + ) + cached_ctx = build_cp_context( + cu_seqlens=cached_cu_seqlens, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + cached = (cached_cu_seqlens, cached_ctx) + self._chunkwise_cp_context_cache[cache_key] = cached + cu_seqlens_q, chunkwise_cp_context = cached + else: + chunkwise_cp_context = build_cp_context( + cu_seqlens=cu_seqlens_q, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + else: + chunkwise_cp_context = None + + if self.recompute_gdn and self.training: + + def _checkpointed_compute(hidden_states): + return self._forward_compute( + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ) + + out, out_bias = tensor_parallel.checkpoint(_checkpointed_compute, False, hidden_states) + else: + out, out_bias = self._forward_compute( + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ) + + if back_to_input_converter is not None: + out = back_to_input_converter.convert( + out, seq_dim=0, sequence_parallel=self.config.sequence_parallel + ) + + return out, out_bias + + def _forward_compute( + self, + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ): + """Core GDN computation (in_proj -> conv1d -> gated_delta_rule -> norm -> out_proj).""" # Input projection nvtx_range_push(suffix="in_proj") qkvzba, _ = self.in_proj(hidden_states) @@ -175,24 +294,171 @@ def forward( qkvzba, thd_cp_a2a_inv = a2a_cp_to_hp( qkvzba, self.in_proj_split_sections, - self.cp_size, - self.pg_collection.cp, + cp_size_headwise, + cp_group_headwise, cu_seqlens_q, - seq_len, + seq_len_post_headwise, packed_seq_params, ) - # Transpose: s b x --> b s x - # From sbhd to bshd format - qkvzba = qkvzba.transpose(0, 1) + if self.gdn_pre_gated_delta_rule_fusion: + if cp_size_chunkwise > 1 and batch > 1: + raise ValueError( + "GDN chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " + "when cp_context is used. Use packed THD input or micro_batch_size=1." + ) + if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + nvtx_range_push(suffix="fused_streamed_pre_gated_delta_rule") + seq_idx = ( + packed_seq_params.seq_idx + if packed_seq_params is not None + and packed_seq_params.qkv_format == 'thd' + and cp_size_chunkwise == 1 + else None + ) + fused_cu_seqlens_q = ( + cu_seqlens_q + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + else None + ) + query, key, value, gate, beta, g = self._fused_streamed_pre_gated_delta_rule( + qkvzba, + cu_seqlens_q=fused_cu_seqlens_q, + seq_idx=seq_idx, + cp_group=cp_group_chunkwise if cp_size_chunkwise > 1 else None, + cp_group_headwise=cp_group_headwise, + ) + kernel_inputs = {"q": query, "k": key, "v": value, "g": g, "beta": beta} + nvtx_range_pop(suffix="fused_streamed_pre_gated_delta_rule") + else: + nvtx_range_push(suffix="pre_gated_delta_rule") + if cp_size_chunkwise > 1 and packed_seq_params is None and batch > 1: + # TODO: If additional gated delta rule backends are added, handle this + # SBHD + chunkwise CP + batch>1 case per backend instead of + # unconditionally rejecting it. + raise ValueError( + "GDN chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " + "when cp_context is used. Use packed THD input or micro_batch_size=1." + ) + if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + query, key, value, gate, beta, g = self.pre_gated_delta_rule( + qkvzba, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + chunkwise_cp_context, + packed_seq_params=packed_seq_params, + ) + kernel_inputs = {"q": query, "k": key, "v": value, "g": g, "beta": beta} + nvtx_range_pop(suffix="pre_gated_delta_rule") + + nvtx_range_push(suffix="gated_delta_rule") + core_attn_out, _ = self.gated_delta_rule( + **kernel_inputs, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens_q, + cp_context=chunkwise_cp_context, + ) + nvtx_range_pop(suffix="gated_delta_rule") + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() + norm_func = partial( + self._gated_norm_and_layout_restore, + thd_cp_a2a_inv=thd_cp_a2a_inv, + batch=batch, + seq_len=seq_len_post_headwise, + packed_seq_params=packed_seq_params, + cp_size_headwise=cp_size_headwise, + cp_group_headwise=cp_group_headwise, + cp_size_chunkwise=cp_size_chunkwise, + cp_group_chunkwise=cp_group_chunkwise, + cu_seqlens_q=cu_seqlens_q, + ) + norm_out = self.norm_out_checkpoint.checkpoint(norm_func, core_attn_out, gate) + else: + norm_out = self._gated_norm_and_layout_restore( + core_attn_out, + gate, + thd_cp_a2a_inv, + batch, + seq_len_post_headwise, + packed_seq_params, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + ) + + # Output projection + nvtx_range_push(suffix="out_proj") + out, out_bias = self.out_proj(norm_out) + nvtx_range_pop(suffix="out_proj") + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint.discard_output_and_register_recompute(out) + + return out, out_bias + + def _gated_norm_and_layout_restore( + self, + core_attn_out: torch.Tensor, + gate: torch.Tensor, + thd_cp_a2a_inv: torch.Tensor | None, + batch: int, + seq_len: int, + packed_seq_params: PackedSeqParams | None, + cp_size_headwise: int, + cp_group_headwise: torch.distributed.ProcessGroup | None, + cp_size_chunkwise: int, + cp_group_chunkwise: torch.distributed.ProcessGroup | None, + cu_seqlens_q: torch.Tensor | None, + ) -> torch.Tensor: + nvtx_range_push(suffix="gated_norm") + norm_out = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + norm_out = norm_out.reshape(batch, seq_len, -1) + norm_out = norm_out.transpose(0, 1).contiguous() + + return a2a_hp_to_cp( + norm_out, cp_size_headwise, cp_group_headwise, packed_seq_params, thd_cp_a2a_inv + ) - # Split the tensor into q, k, v, gate (z), and the variant-specific gate features - # (beta, alpha for GDN; f, b, w for GDN2) - qkv, gate, beta, alpha = self._split_projection(qkvzba, batch, seq_len) + def pre_gated_delta_rule( + self, + qkvzba, + batch, + seq_len, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q=None, + chunkwise_cp_context=None, + packed_seq_params=None, + ): + """Prepare QKV, gate, beta, and decay tensors before the gated delta rule.""" + + qkvzba = qkvzba.transpose(0, 1) + qkv, gate, beta, alpha = torch.split( + qkvzba, self._get_feat_dim_split(cp_size_headwise), dim=-1 + ) + gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) - # Convolution on qkv nvtx_range_push(suffix="conv1d") - seq_len = qkv.shape[1] + kernel_seq_len = qkv.shape[1] qkv_channels_split_sections = [ self.qk_dim_local_tp, self.qk_dim_local_tp, @@ -201,219 +467,155 @@ def forward( conv1d_weight = get_parameter_local_cp( self.conv1d.weight, dim=0, - cp_group=self.pg_collection.cp, + cp_group=cp_group_headwise, split_sections=qkv_channels_split_sections, ) conv1d_bias = ( get_parameter_local_cp( self.conv1d.bias, dim=0, - cp_group=self.pg_collection.cp, + cp_group=cp_group_headwise, split_sections=qkv_channels_split_sections, ) if self.conv_bias else None ) if self.config.deterministic_mode: - qkv = qkv.transpose(1, 2).contiguous() # b, s, d -> b, d, s + qkv = qkv.transpose(1, 2).contiguous() conv_out = F.conv1d( - input=qkv, # Torch-native only accept [b, d, s] format input + input=qkv, weight=conv1d_weight, bias=conv1d_bias, stride=self.conv1d.stride, padding=self.conv1d.padding, dilation=self.conv1d.dilation, - groups=self.conv_dim_local_tp // self.cp_size, + groups=self.conv_dim_local_tp // cp_size_headwise, ) - qkv = self.act_fn(conv_out[..., :seq_len]) - qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d + qkv = self.act_fn(conv_out[..., :kernel_seq_len]) + qkv = qkv.transpose(1, 2) else: assert self.activation in ["silu", "swish"] + orig_seq = qkv.shape[1] + pad_n = 0 + conv_input = qkv.contiguous() + conv_cu_seqlens = cu_seqlens_q + conv_cp_context = chunkwise_cp_context + if self.config.gdn_conv_pad_alignment is not None: + if packed_seq_params is None or cu_seqlens_q is None: + raise ValueError( + "gdn_conv_pad_alignment is only supported with packed sequence " + "parameters in THD format. SBHD inputs do not need causal-conv padding." + ) + if chunkwise_cp_context is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + pad_n = -orig_seq % self.config.gdn_conv_pad_alignment + if pad_n > 0: + conv_input = torch.nn.functional.pad(conv_input, (0, 0, 0, pad_n)) + conv_cu_seqlens = cu_seqlens_q.clone() + conv_cu_seqlens[-1] += pad_n qkv, _ = causal_conv1d( - x=qkv, # FLA conv1d accepts [b, s, d] format input - weight=conv1d_weight.squeeze(1), # d, 1, w -> d, w + x=conv_input, + weight=conv1d_weight.squeeze(1), bias=conv1d_bias, activation=self.activation, initial_state=None, output_final_state=False, - cu_seqlens=cu_seqlens_q, + cu_seqlens=conv_cu_seqlens, + cp_context=conv_cp_context, ) + if pad_n > 0: + qkv = qkv[:, :orig_seq, :] nvtx_range_pop(suffix="conv1d") - A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp) - dt_bias_local_cp = get_parameter_local_cp( - self.dt_bias, dim=0, cp_group=self.pg_collection.cp - ) + A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group_headwise) + dt_bias_local_cp = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group_headwise) - # Prepare all kernel inputs (split, reshape, L2 norm, gates, contiguous) nvtx_range_push(suffix="prepare_input_for_gated_delta_rule") kernel_inputs = self._prepare_input_for_gated_delta_rule( - qkv, gate, A_log_local_cp, dt_bias_local_cp, batch, seq_len, beta, alpha + qkv, + gate, + A_log_local_cp, + dt_bias_local_cp, + batch, + kernel_seq_len, + beta, + alpha, + cp_size_headwise=cp_size_headwise, ) - gate = kernel_inputs.pop("gate") nvtx_range_pop(suffix="prepare_input_for_gated_delta_rule") - nvtx_range_push(suffix="gated_delta_rule") - core_attn_out, _ = self.gated_delta_rule( - **kernel_inputs, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - cu_seqlens=cu_seqlens_q, - ) - nvtx_range_pop(suffix="gated_delta_rule") - - if self.recompute_norm_out: - self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() - norm_func = partial( - self._gated_norm_and_a2a, - thd_cp_a2a_inv=thd_cp_a2a_inv, - batch=batch, - seq_len=seq_len, - packed_seq_params=packed_seq_params, - ) - norm_out = self.norm_out_checkpoint.checkpoint(norm_func, core_attn_out, gate) - else: - norm_out = self._gated_norm_and_a2a( - core_attn_out, gate, thd_cp_a2a_inv, batch, seq_len, packed_seq_params - ) - - # Output projection - nvtx_range_push(suffix="out_proj") - out, out_bias = self.out_proj(norm_out) - nvtx_range_pop(suffix="out_proj") - - if self.recompute_norm_out: - self.norm_out_checkpoint.discard_output_and_register_recompute(out) - - return out, out_bias + gate = kernel_inputs.pop("gate") - def _split_projection( - self, projected: torch.Tensor, batch: int, seq_len: int - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Split the fused projection into qkv, output gate, beta, and alpha.""" - qkv, gate, beta, alpha = torch.split(projected, self.feat_dim_split, dim=-1) - gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) - return qkv, gate, beta, alpha - - def _prepare_inference_inputs( - self, qkv: torch.Tensor, beta: torch.Tensor, alpha: torch.Tensor, batch: int, seq_len: int - ) -> dict[str, torch.Tensor]: - """Prepare raw FLA inputs while leaving normalization and gates fused in-kernel.""" - query_key, value = torch.split(qkv, [2 * self.qk_dim_local_tp, self.v_dim_local_tp], dim=-1) - query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) - query, key = torch.chunk(query_key, 2, dim=2) - value = value.reshape(batch, seq_len, -1, self.value_head_dim) - return { - "q": query.contiguous(), - "k": key.contiguous(), - "v": value.contiguous(), - "g": alpha.contiguous(), - "beta": beta.contiguous(), - } - - def mamba_state_shapes_per_request(self) -> tuple[tuple[int, ...], tuple[int, ...]]: - """Return the TP-local convolution and delta-rule cache shapes.""" return ( - (self.conv_dim_local_tp, self.conv_kernel_dim), - (self.num_v_heads_local_tp, self.key_head_dim, self.value_head_dim), + kernel_inputs["q"], + kernel_inputs["k"], + kernel_inputs["v"], + gate, + kernel_inputs["beta"], + kernel_inputs["g"], ) - def ssm_decode( - self, - projected: torch.Tensor, - conv_state: torch.Tensor, - ssm_state: torch.Tensor, - batch_indices: torch.Tensor, - intermediate_conv_state: torch.Tensor | None = None, - intermediate_ssm_state: torch.Tensor | None = None, - ) -> torch.Tensor: - """Run one CUDA-graph-compatible GDN decode token per request.""" - batch, seq_len, _ = projected.shape - assert seq_len == 1, "GDN speculative decoding is not supported." - assert ( - intermediate_conv_state is None and intermediate_ssm_state is None - ), "GDN speculative decoding state capture is not supported." - assert causal_conv1d_update is not None and fused_recurrent_gated_delta_rule is not None - - qkv, gate, beta, alpha = self._split_projection(projected, batch, seq_len) - read_indices = batch_indices.clamp(min=0) - - active_conv_state = conv_state[read_indices].contiguous() - qkv_dtype = qkv.dtype - qkv, active_conv_state = causal_conv1d_update( - x=qkv.to(conv_state.dtype), - cache=active_conv_state, - weight=self.conv1d.weight.squeeze(1).to(conv_state.dtype), - bias=self.conv1d.bias.to(conv_state.dtype) if self.conv1d.bias is not None else None, - activation=self.activation, - ) - qkv = qkv.to(qkv_dtype) - tensor_masked_update(conv_state, batch_indices, active_conv_state) + def _fused_streamed_pre_gated_delta_rule( + self, qkvzba, cu_seqlens_q=None, seq_idx=None, cp_group=None, cp_group_headwise=None + ): + """Call the streamed fused pre-GDR wrapper.""" - kernel_inputs = self._prepare_inference_inputs(qkv, beta, alpha, batch, seq_len) - active_ssm_state = ssm_state[read_indices].contiguous() - core_attn_out, final_ssm_state = fused_recurrent_gated_delta_rule( - **kernel_inputs, - A_log=self.A_log, - dt_bias=self.dt_bias, - initial_state=active_ssm_state, - output_final_state=True, - use_qk_l2norm_in_kernel=self.use_qk_l2norm, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=True, - ) - tensor_masked_update(ssm_state, batch_indices, final_ssm_state) - return self._apply_gated_norm(core_attn_out, gate).reshape(batch, seq_len, -1) + try: + from megatron.core.fusions.fused_pre_gated_delta_rule import ( + fused_streamed_pre_gated_delta_rule, + ) + except ImportError as exc: + raise ImportError( + "gdn_pre_gated_delta_rule_fusion requires the streamed pre-GDR fusion " + "dependencies, including causal-conv1d." + ) from exc - def ssm_prefill( - self, - projected: torch.Tensor, - conv_state: torch.Tensor, - ssm_state: torch.Tensor, - context: DynamicInferenceContext, - ) -> torch.Tensor: - """Run packed variable-length GDN prefill and populate request states.""" - assert ( - not context.is_chunked_prefill_enabled() - ), "GDN dynamic inference does not support chunked prefill." - metadata = context.mamba_metadata - cu_seqlens = metadata.cu_seqlens - batch_indices = metadata.batch_indices_prefill - token_count = projected.shape[0] - - projected = projected.transpose(0, 1).contiguous() - qkv, gate, beta, alpha = self._split_projection(projected, 1, token_count) - read_indices = batch_indices.clamp(min=0) - - qkv_dtype = qkv.dtype - qkv, final_conv_state = causal_conv1d( - x=qkv.to(conv_state.dtype), - weight=self.conv1d.weight.squeeze(1).to(conv_state.dtype), - bias=self.conv1d.bias.to(conv_state.dtype) if self.conv1d.bias is not None else None, - activation=self.activation, - initial_state=conv_state[read_indices].contiguous(), - output_final_state=True, - cu_seqlens=cu_seqlens, + qkv_channels_split_sections = [ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ] + conv1d_weight = get_parameter_local_cp( + self.conv1d.weight, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + conv1d_bias = ( + get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + if self.conv_bias + else None + ) + A_log = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group_headwise) + dt_bias = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group_headwise) + num_value_heads = A_log.numel() + num_key_heads = (conv1d_weight.shape[0] - num_value_heads * self.value_head_dim) // ( + 2 * self.key_head_dim ) - qkv = qkv.to(qkv_dtype) - tensor_masked_update(conv_state, batch_indices, final_conv_state) - kernel_inputs = self._prepare_inference_inputs(qkv, beta, alpha, 1, token_count) - core_attn_out, final_ssm_state = chunk_gated_delta_rule( - **kernel_inputs, - A_log=self.A_log, - dt_bias=self.dt_bias, - initial_state=ssm_state[read_indices].contiguous(), - output_final_state=True, - use_qk_l2norm_in_kernel=self.use_qk_l2norm, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=True, - cu_seqlens=cu_seqlens, + return fused_streamed_pre_gated_delta_rule( + qkvzba, + conv1d_weight, + conv1d_bias, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=self.key_head_dim, + value_head_dim=self.value_head_dim, + use_qk_l2norm=self.use_qk_l2norm, + cu_seqlens=cu_seqlens_q, + seq_idx=seq_idx, + cp_group=cp_group, ) - tensor_masked_update(ssm_state, batch_indices, final_ssm_state) - y = self._apply_gated_norm(core_attn_out, gate) - return y.reshape(1, token_count, -1).transpose(0, 1).contiguous() #################### @@ -425,22 +627,20 @@ def torch_chunk_gated_delta_rule( v, g, beta, - scale=None, chunk_size=64, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False, cu_seqlens=None, - **kwargs, + cp_context=None, + scale=None, ) -> tuple[torch.Tensor, torch.Tensor | None]: # pylint: disable=line-too-long ''' Torch-native implementation of chunked gated delta rule for deterministic mode. Need this because FLA is not deterministic. - ``scale`` defaults to ``1 / sqrt(K)``, matching the FLA kernel. Extra keyword - arguments are accepted and ignored so this stays interchangeable with the FLA - kernel, which takes several options this implementation does not model. + ``scale`` defaults to ``1 / sqrt(K)``, matching the FLA kernel. Reference: https://github.com/huggingface/transformers/blob/144c8ce2809a2e21914017652700e1ecb450501e/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L470-L547 ''' @@ -448,65 +648,67 @@ def torch_chunk_gated_delta_rule( assert ( cu_seqlens is None ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now." + assert ( + cp_context is None + ), "cp_context is not supported for torch_chunk_gated_delta_rule for now." - query, key, value = q, k, v - initial_dtype = query.dtype + initial_dtype = q.dtype if use_qk_l2norm_in_kernel: - query = l2norm(query, dim=-1, eps=1e-6) - key = l2norm(key, dim=-1, eps=1e-6) - query, key, value, beta, g = [ - x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) + q = l2norm(q, dim=-1, eps=1e-6) + k = l2norm(k, dim=-1, eps=1e-6) + q, k, v, beta, g = [ + x.transpose(1, 2).contiguous().to(torch.float32) for x in (q, k, v, beta, g) ] - batch_size, num_heads, sequence_length, k_head_dim = key.shape - v_head_dim = value.shape[-1] + batch_size, num_heads, sequence_length, k_head_dim = k.shape + v_head_dim = v.shape[-1] pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size - query = F.pad(query, (0, 0, 0, pad_size)) - key = F.pad(key, (0, 0, 0, pad_size)) - value = F.pad(value, (0, 0, 0, pad_size)) + q = F.pad(q, (0, 0, 0, pad_size)) + k = F.pad(k, (0, 0, 0, pad_size)) + v = F.pad(v, (0, 0, 0, pad_size)) beta = F.pad(beta, (0, pad_size)) g = F.pad(g, (0, pad_size)) total_sequence_length = sequence_length + pad_size if scale is None: - scale = 1 / (query.shape[-1] ** 0.5) - query = query * scale + scale = 1 / (q.shape[-1] ** 0.5) + q = q * scale - v_beta = value * beta.unsqueeze(-1) - k_beta = key * beta.unsqueeze(-1) + v_beta = v * beta.unsqueeze(-1) + k_beta = k * beta.unsqueeze(-1) # reshape to chunks - query, key, value, k_beta, v_beta = [ + q, k, v, k_beta, v_beta = [ x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) - for x in (query, key, value, k_beta, v_beta) + for x in (q, k, v, k_beta, v_beta) ] g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) mask = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0 + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0 ) # chunk decay g = g.cumsum(dim=-1) decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() - attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) + attn = -((k_beta @ k.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) for i in range(1, chunk_size): row = attn[..., i, :i].clone() sub = attn[..., :i, :i].clone() attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) - value = attn @ v_beta + v = attn @ v_beta k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) last_recurrent_state = ( - torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value) + torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(v) if initial_state is None - else initial_state.to(value) + else initial_state.to(v) ) - core_attn_out = torch.zeros_like(value) + core_attn_out = torch.zeros_like(v) mask = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1 + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1 ) # for each chunk for i in range(0, total_sequence_length // chunk_size): - q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] + q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state v_new = v_i - v_prime diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 89bc09f64b3..a5933021700 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1,5 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import inspect import logging import math import warnings @@ -9,7 +10,6 @@ import torch import torch.nn.functional as F -from megatron.core.context_parallel import CPLayout from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.inference.moe import InferenceGroupedGemmBackend from megatron.core.quantization.quant_config import RecipeConfig @@ -25,8 +25,10 @@ CudaGraphModule, CudaGraphScope, InferenceCudaGraphScope, + LayerType, ) from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout +from megatron.core.utils import experimental_api from .._rank_utils import log_single_rank from ..fusions.fused_bias_geglu import quick_gelu @@ -52,6 +54,7 @@ @dataclass +@experimental_api class TransformerConfig(ModelParallelConfig): """Configuration object for megatron-core transformers. @@ -88,13 +91,6 @@ class TransformerConfig(ModelParallelConfig): This prevents MTP loss gradients from flowing back to the main model, only training the MTP heads themselves.""" - mtp_hsm: bool = False - """Enable uniform per-token Hidden State Mixing (HSM) for MTP layers. - At every MTP depth, each token independently draws its input from the main model - hidden state and the outputs of the earlier depths, all aligned on the same target - token. Only takes effect during training and requires at least two MTP layers, - since a single depth has nothing to mix.""" - mtp_hybrid_override_pattern: Optional[str] = None """DEPRECATED: Use unified hybrid_layer_pattern instead. Legacy argument for loading old checkpoints. @@ -155,16 +151,6 @@ class TransformerConfig(ModelParallelConfig): If attention backend is local we use the local pytorch implementation in mcore. Users can specify exact backend by changing this config. """ - flash_attention_version: Optional[Literal[2, 3, 4]] = None - """Pin the FlashAttention generation (2, 3, or 4) used by both the training - (TransformerEngine) and inference (mcore dynamic-batching) attention paths. When - None, each path selects a version automatically based on what is installed. Pinning - is required for batch-invariant mode: the training-side logprob recompute and the - inference engine must run the same kernel, since different FlashAttention - generations use different tile sizes and softmax accumulation orders and therefore - differ bitwise. On the training side this is enforced via TransformerEngine's - NVTE_FLASH_ATTN_V2/V3/V4 selection environment variables.""" - softmax_scale: Optional[float] = None """Softmax scale for attention scaling.""" @@ -235,7 +221,7 @@ class TransformerConfig(ModelParallelConfig): activation_func_clamp_value: Optional[float] = None """Clamp the output of the linear_fc1 in the activation function. Only used when activation_func - is quick_gelu or SwiGLU (MoE only).""" + is quick_gelu or weighted SwiGLU (MoE only).""" num_moe_experts: Optional[int] = None """Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None @@ -277,7 +263,26 @@ 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.""" + """Apply a full head_dim output gate (num_attention_heads * head_dim rows + fused inline per-group [q, gate, k, v] into linear_qkv). Mutually + exclusive with `head_wise_attn_gate` (per-head scalar gate).""" + + 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 +304,17 @@ class TransformerConfig(ModelParallelConfig): #################### # attention variant #################### - experimental_attention_variant: Optional[Literal['gdn', 'gdn2', 'dsa', 'gated_delta_net']] = ( + experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = ( 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.""" + """Type of attention variant to use. Currently support gated_delta_net, dsa, and dsv4_hybrid.""" + + cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag" + """How THD sequence rows are partitioned across context-parallel ranks. + + Contiguous partitioning is defined only for THD inputs; BSHD context parallelism uses the + standard zigzag layout. + """ experimental_attention_variant_loss_scale_func: Optional[Callable[[torch.Tensor], None]] = None """Optional hook for experimental attention variants to receive the main loss scale.""" @@ -339,18 +346,16 @@ class TransformerConfig(ModelParallelConfig): top-k indices.""" dsa_kernel_backend: Literal["none", "tilelang", "cudnn"] = "none" - """Optional fused DSA kernel backend. - ``none`` disables fused DSA kernels. Explicit ``tilelang`` or ``cudnn`` enables only that - backend. Unsupported DSA layouts continue to use the PyTorch fallback.""" + """Optional fused ordinary-DSA kernel backend. Unsupported layouts use PyTorch fallback.""" dsa_indexer_rope_interleaved: bool = False """Whether DSA indexer RoPE should use MLA-style interleaving.""" dsa_indexer_rotate_activation: bool = True - """Whether DSA indexer should apply Hadamard rotate_activation to q/k before scoring.""" + """Whether DSA indexer should apply Hadamard rotation before scoring.""" dsa_indexer_scoring_relu: bool = True - """Whether DSA indexer should apply ReLU to q@k^T scores before weighting.""" + """Whether DSA indexer should apply ReLU to q@k scores before weighting.""" dsa_indexer_k_norm_epsilon: Optional[float] = None """Optional epsilon override for the DSA indexer key LayerNorm.""" @@ -359,16 +364,14 @@ class TransformerConfig(ModelParallelConfig): """Whether DSA indexer key LayerNorm should run on fp32 inputs.""" #################### - # Compressed sparse attention + # DeepSeek-v4 hybrid attention #################### csa_window_size: int = 128 """Sliding window size for compressed sparse attention.""" - # TODO(#6402): consumed by DSv4 Hybrid attention orchestration, which selects the - # per-layer compression ratio and builds the compressed-KV rotary embedding. - # Neither field has a production reader in this primitive-only PR. csa_compress_ratios: Optional[List[int]] = None - """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...].""" + """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...]. A value of 0 is a + sliding-window-only layer (no compressor / no top-k indexer; the 'W' hybrid layer symbol).""" csa_compress_rotary_base: float = 40000.0 """RoPE base for compressed KV positions in compressed sparse attention.""" @@ -377,9 +380,16 @@ class TransformerConfig(ModelParallelConfig): """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be disabled.""" + apply_dsa_kernel_fusion: Optional[bool] = None + """Deprecated DSv4 fused-kernel switch. Use ``dsa_kernel_backend`` instead. For + ``dsv4_hybrid`` only, True maps to ``cudnn`` and False maps to ``none``.""" + #################### # linear attention #################### + linear_attention_type: Optional[str] = None + """Type of linear attention to use. + Deprecated. Use experimental_attention_variant instead.""" linear_attention_freq: Optional[Union[int, List[int]]] = None """Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) layers. @@ -402,6 +412,14 @@ class TransformerConfig(ModelParallelConfig): linear_num_value_heads: Optional[int] = 32 """Number of value and gate heads for the gated delta net.""" + gdn_pre_gated_delta_rule_fusion: bool = False + """Whether to use the streamed Triton fusion for GatedDeltaNet pre-GDR preprocessing.""" + + gdn_conv_pad_alignment: Optional[int] = None + """When set, pad packed GDN causal-conv inputs to this token alignment. + This is only valid without chunkwise CP: padding a chunk-local causal-conv input changes + the sequence seen by later chunks and therefore changes the GDN recurrence numerics.""" + #################### # initialization #################### @@ -539,13 +557,13 @@ class TransformerConfig(ModelParallelConfig): fused_single_qkv_rope: bool = False """If set, avoid splitting QKV before ROPE forward and avoid concatenating ROPE dgrads.""" - fused_residual_rmsnorm: bool = False - """If True, fuses residual connection and RMSNorm backward pass when TE is used.""" - use_transformer_engine_op_fuser: bool = False """If True, submodules may use Transformer Engine's operation fuser API to enable advanced fusions.""" + fused_residual_rmsnorm: bool = False + """If True, fuses residual connection and RMSNorm backward pass when TE is used.""" + #################### # activation recomputation #################### @@ -580,7 +598,7 @@ class TransformerConfig(ModelParallelConfig): recompute_modules: Optional[List[str]] = None """The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", - "shared_experts", "gdn_norm_out", "gdp_in_proj", "gdp_qkv", "mhc". + "shared_experts", "mhc", "gdn", "gdn_norm_out". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -589,16 +607,17 @@ class TransformerConfig(ModelParallelConfig): "mlp": recompute the dense MLP submodule. "moe": recompute the MoE layer. "shared_experts": recompute the shared experts in the MoE layer. - "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all. - "gdp_in_proj": recompute the GatedDeltaProduct input projection and its CP gather/split - preprocessing. - "gdp_qkv": recompute the GatedDeltaProduct causal conv and QKV preparation. "mhc": recompute HyperConnection intermediate activations via - CheckpointWithoutOutput + CheckpointWithoutOutputManager. Requires - enable_mhc_connections=True. Cannot be used with "mlp". - "moe_act", "layernorm", "mla_up_proj", "gdn_norm_out", "gdp_in_proj", "gdp_qkv", and - "mhc" use output-discarding checkpointing, "core_attn", "mlp", "moe", and - "shared_experts" use normal checkpointing. + CheckpointWithoutOutput + MHCCheckpointManager. Requires + enable_hyper_connections=True. Cannot be used with "mlp". + "gdn": recompute the entire GatedDeltaNet module (in_proj, conv1d, gated delta rule, + gated norm, CP all-to-all and out_proj). Requires + experimental_attention_variant="gated_delta_net". + "gdn_norm_out": recompute gated output normalization and layout restoration for + Gated DeltaNet-family layers, including GatedDeltaNet and KDA. + "moe_act", "layernorm", "mla_up_proj", "mhc", and "gdn_norm_out" use + output-discarding checkpointing, + "core_attn", "mlp", "moe", "shared_experts", and "gdn" use normal checkpointing. """ #################### @@ -731,7 +750,7 @@ class TransformerConfig(ModelParallelConfig): in the hidden_states gradient.""" moe_shared_expert_gate: bool = False - """Enable gate for shared expert. Only effective when + """Enable gate for shared expert. Only effective when moe-shared-expert-intermediate-size is set.""" moe_shared_expert_overlap: bool = False @@ -769,9 +788,6 @@ class TransformerConfig(ModelParallelConfig): for each individual sample. - "global_aux_loss": Load balancing loss calculated at global batch level. - "sinkhorn": Balancing algorithm used in S-BASE. - - "quantile_balancing": Dual coordinate-descent quantile balancing (QB). Load balance is - handled entirely by an internal per-expert bias update; auxiliary losses must be disabled - (`moe_aux_loss_coeff` = 0) when QB is selected. - "none": No load balancing. A list of strings can be provided to combine multiple aux-loss load balancing types. The default is "aux_loss". @@ -822,8 +838,8 @@ class TransformerConfig(ModelParallelConfig): By default, softmax is done after top-k.""" moe_router_topk_scaling_factor: Optional[float] = None - """Scaling factor for routing score in top-k selection. Defaults to None, which means no - scaling.""" + """Scaling factor for routing score in top-k selection, only works when moe_router_pre_softmax + enabled. Defaults to None, which means no scaling.""" moe_router_score_function: Literal['softmax', 'sigmoid', 'sqrtsoftplus'] = "softmax" """Score function for MoE routing. Can be "softmax", "sigmoid" or "sqrtsoftplus".""" @@ -833,9 +849,6 @@ class TransformerConfig(ModelParallelConfig): improve stability especially when the number of experts is large (e.g. finegrained-moe). None means no changes for dtype.""" - moe_router_skip_muon: bool = False - """Use the scalar optimizer instead of Muon for MoE router parameters.""" - moe_router_enable_expert_bias: bool = False """TopK routing with dynamic per-expert bias in the aux-loss-free load balancing strategy. The routing decision is based on the sum of the routing scores and the expert bias. @@ -847,12 +860,6 @@ class TransformerConfig(ModelParallelConfig): and decreased for the experts with more assigned tokens. The default value 1e-3 is same as that used in DeepSeekV3.""" - moe_router_quantile_balancing_ema: float = 0.0 - """EMA coefficient for the quantile-balancing per-expert bias (`qb_beta`), used only when - `moe_router_load_balancing_type` is "quantile_balancing". At each global batch the bias is - updated as `qb_beta = ema * qb_beta + (1 - ema) * local_quantile`. The default 0.0 means - no memory: the bias is replaced by the latest global-batch quantile estimate each step.""" - moe_router_force_load_balancing: bool = False """[Experimental] Force load balancing with random logits for MoE router, supports naive topk and group-limited topk. This is an experimental feature and only for benchmark.""" @@ -864,36 +871,49 @@ class TransformerConfig(ModelParallelConfig): If negative, generates bias once per layer and reuses it (abs value is std). This is an experimental feature for benchmarking purposes.""" - use_grouped_gemm_for_dense_mlp: bool = False + moe_n_hash_layers: int = 0 + """Number of leading transformer layers that use hash-based MoE routing. + Layers with layer_number <= moe_n_hash_layers use a pre-computed tid2eid + lookup table for expert selection instead of learned top-k routing.""" + + actual_vocab_size: Optional[int] = None + """Padded actual vocabulary size. Required when moe_n_hash_layers > 0 for the + tid2eid lookup buffer in hash-based MoE routing.""" + + dense_grouped_gemm: bool = False """Use GroupedLinear(num_groups=1) for dense MLP to trigger the ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 fusion on SM100+ with MXFP8 recipe. Requires ``use_te_op_fuser=True`` and SwiGLU activation. """ - moe_grouped_gemm: bool = False - """Use grouped GEMM to execute multiple local MoE experts together. - - The concrete implementation is selected by Transformer Engine. Set - ``moe_use_grouped_tensor=True`` to use its CUDA-graph-safe GroupedTensor path. + use_grouped_gemm_for_dense_mlp: bool = False + """Alias of ``dense_grouped_gemm``. Use GroupedLinear(num_groups=1) for dense MLP to + trigger the ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 fusion on SM100+ with MXFP8 recipe. + Requires ``use_te_op_fuser=True`` and SwiGLU activation. """ - moe_use_grouped_tensor: bool = False - """Use Transformer Engine's native GroupedTensor path for grouped MoE GEMMs. + log_moe_overload_factor: bool = False + """When True, log MoE overload metrics (avg/max vs balanced token count per step; max cum + overload = peak cumulative actual tokens / peak cumulative balanced count over interleaved + fwd/bwd) to TensorBoard/W&B and console. Records tokens_per_expert.sum() after dispatch; + use for debugging.""" - This path uses padded expert segments and CUDA split metadata so it can be captured in CUDA - graphs. Enabling the Transformer Engine operation fuser also enables this option. + moe_grouped_gemm: bool = False + """When there are multiple experts per rank, compress multiple local (potentially small) gemms + in a single kernel launch to improve the utilization and performance by leveraging the Grouped + GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm). """ moe_single_grouped_weight: bool = False """When using TE GroupedLinear for MoE experts, store expert weights as a single grouped parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` and - ``moe_use_grouped_tensor=True``. + ``use_transformer_engine_op_fuser=True``. """ moe_single_grouped_bias: bool = False """When using TE GroupedLinear for MoE experts, store expert biases as a single grouped - parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``, - ``moe_use_grouped_tensor=True``, and ``add_bias_linear=True``.""" + parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` + and ``add_bias_linear=True``.""" moe_aux_loss_coeff: Union[float, List[float]] = 0.0 """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended. @@ -918,20 +938,22 @@ class TransformerConfig(ModelParallelConfig): moe_enable_deepep: bool = False """[Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models.""" - moe_flex_dispatcher_backend: Literal['deepep', 'hybridep', 'ncclep'] = "deepep" + moe_flex_dispatcher_backend: Literal['deepep', 'deepepv2', 'hybridep', 'ncclep'] = "deepep" """[Experimental] The backend to use for flex token dispatcher. The default is "deepep". - Options are "deepep", "hybridep", and "ncclep". Currently only "hybridep" backend supports - the MNNVL case. "ncclep" uses NVIDIA NCCL Expert Parallelism via TransformerEngine's + Options are "deepep", "deepepv2", "hybridep", and "ncclep". Currently only "hybridep" backend + supports the MNNVL case. "ncclep" uses NVIDIA NCCL Expert Parallelism via TransformerEngine's transformer_engine.pytorch.ep API.""" moe_permute_fusion_into_hybridep: bool = False """Fuse token rearrangement ops during token dispatching for HybridEP.""" - moe_hybridep_pad_uneven_dispatch_inputs: bool = False - """Pad uneven HybridEP dispatch inputs to the group maximum before dispatch. - Enable when local HybridEP input token counts can differ across ranks, for example - with dynamically packed THD inputs. Leave disabled when dispatcher inputs are - already padded to equal token counts. + moe_hybridep_pad_variable_tokens: bool = False + """Dynamically pad uneven local token counts to the HybridEP group maximum before dispatch. + + Enable this when local token counts can differ across ranks. When disabled, the caller must + guarantee equal token counts across the HybridEP communication group, for example by padding + THD inputs to a fixed maximum before dispatch. CUDA Graph inputs should be statically padded + upstream and leave this option disabled. """ moe_per_layer_logging: bool = False @@ -941,6 +963,11 @@ class TransformerConfig(ModelParallelConfig): """moe_expert_capacity_factor (float): The capacity factor for each expert, None means no token will be dropped. The default is None.""" + moe_expert_rank_capacity_factor: Optional[float] = None + """moe_expert_rank_capacity_factor (float): The capacity factor for each expert rank. Tokens + exceeding this budget will be dropped. None means no token will be dropped. + The default is None.""" + moe_pad_expert_input_to_capacity: bool = False """moe_pad_expert_input_to_capacity (bool): If True, pads the input for each expert to match the expert capacity length, effective only after the moe_expert_capacity_factor is set. The @@ -975,19 +1002,6 @@ class TransformerConfig(ModelParallelConfig): moe_latent_size: Optional[int] = None """Latent projection dimension for MoE. If None, MoE latent projections are not used.""" - moe_use_norm_before_up_proj: bool = False - """Apply normalization before the latent-to-hidden MoE projection. Requires - ``moe_latent_size`` to be set.""" - - gtp_remat_opt_in_modules: list[str] = field(default_factory=list) - """Extra modules to apply GTP_remat weight sharding to, beyond the default set (attention, - Mamba, MLP, expert linears, embeddings). Allowed values: - - - ``"moe_latent_proj"`` — shard ``fc1_latent_proj`` / ``fc2_latent_proj`` (MoE latent - projections, ``parallel_mode="duplicated"``). Only beneficial when ``moe_latent_size`` - is large enough for the all-gather to amortize. - """ - moe_flex_dispatcher_num_sms: Optional[int] = None """Number of SMs for the flex token dispatcher's dispatch/combine communication, for all backends (deepep, hybridep, ncclep). None lets each backend use its own default. Unifies the @@ -1014,26 +1028,22 @@ class TransformerConfig(ModelParallelConfig): moe_hybridep_num_sms_preprocessing: int = 108 """Number of SMs to use for HybridEP preprocessing (metadata scan kernel).""" - moe_ncclep_zero_copy: bool = False + moe_ncclep_static_shape: bool = False + """For the 'ncclep' flex dispatcher: feed the experts the full fixed-size receive buffer + instead of narrowing to the (data-dependent) number of received tokens, removing the D2H sync + and dynamic shapes from the dispatch (required for CUDA-graph capture of the MoE A2A and for the + 1F1B EP comm overlap). The fused grouped GEMM consumes the ragged per-expert counts on device + and walks only the received tokens (no slack GEMM, no last-expert padding). This requires the + CuTe DSL / device-offset grouped GEMM, so it is only supported with the fused op + (use_transformer_engine_op_fuser, NVTE_CUTEDSL_FUSED_GROUPED_MLP=1) on sm100+ (Blackwell or + later); the dispatcher asserts this. On older GPUs leave it False (dynamic shape). Defaults to + False (narrow to the received tokens).""" + + moe_ncclep_use_symm_mem: bool = False """For the 'ncclep' flex dispatcher: use the NCCL symmetric-memory zero-copy IO path (ep_bootstrap zero_copy + symm-mem-backed receive/combine buffers) instead of the default HBM - staged-copy path, saving one copy on the wire. Requires moe_expert_rank_capacity_factor (the - symm-mem buffers are a fixed [recv_capacity, hidden] and cannot be resized per step) and the - fused op (use_transformer_engine_op_fuser). Defaults to False.""" - - moe_dispatch_fwd_dtype: Literal['bf16', 'mxfp8'] = 'bf16' - """Wire dtype of the MoE dispatch forward payload ('ncclep' flex dispatcher only). With - 'mxfp8', TransformerEngine quantizes the payload before the all-to-all and the receive - buffer comes back as a per-expert MXFP8 GroupedTensor that the grouped GEMM consumes - directly. Requires moe_grouped_gemm and use_transformer_engine_op_fuser. Defaults to - 'bf16' (no quantization on the wire).""" - - moe_combine_bwd_dtype: Literal['bf16', 'mxfp8'] = 'bf16' - """Wire dtype of the MoE combine backward gradient ('ncclep' flex dispatcher only). With - 'mxfp8', TransformerEngine quantizes the gradient before the all-to-all and the - expert-output gradient comes back as a per-expert MXFP8 GroupedTensor that the grouped - GEMM backward consumes directly. Same requirements as moe_dispatch_fwd_dtype. Defaults to - 'bf16' (no quantization on the wire).""" + staged-copy path. NOT SUPPORTED YET -- the dispatcher rejects this if set; the cross-stream + reuse ordering for the persistent symm-mem buffer is not implemented. Leave False.""" moe_mlp_glu_interleave_size: Optional[int] = None """When set, GLU activations in the MoE grouped MLP layer will use a @@ -1043,14 +1053,6 @@ class TransformerConfig(ModelParallelConfig): This data format is experimental and primarily intended to enable advanced fused kernels.""" - moe_expert_rank_capacity_factor: Optional[float] = None - """moe_expert_rank_capacity_factor (float): The capacity factor for each expert rank, i.e. the - per-rank token budget. None means no token will be dropped. The default is None. - With the 'hybridep' backend, tokens exceeding this budget are dropped. With the 'ncclep' - backend, setting it selects static shapes (fixed receive buffer, CUDA-graph capturable) and - leaving it None selects eager mode (receive buffer sized per step from the received-token - count).""" - ################## # Context Parallel ################## @@ -1059,6 +1061,8 @@ class TransformerConfig(ModelParallelConfig): str: all layers share same communication type. List[str]: each layer has its separate communication type. cp_comm_type of each layer can be "p2p" or "all_gather" or "a2a" or "a2a+p2p". + This option controls standard attention layers. Linear-attention layers use + `linear_cp_mode` instead. "p2p": Exchange KV chunks with P2P communications in ring topology. P2P is async and can be overlapped with attention compute. "all_gather": All-gather to get full sequence of KV before attention. The all-gather is not @@ -1070,11 +1074,18 @@ class TransformerConfig(ModelParallelConfig): and P2P communications in high-level CP groups (e.g., via IBLink). """ - linear_cp_layout: CPLayout = "zigzag" - """CP layout for linear-attention layers.""" - - attention_cp_layout: CPLayout = "zigzag" - """CP layout for softmax-attention layers.""" + linear_cp_mode: Optional[str] = "chunkwise" + """Context-parallel execution mode for linear-attention layers + (e.g. Gated Delta Net). Independent of `cp_comm_type`, which only controls standard attention. + Can be "chunkwise" or "headwise": + "chunkwise": Keep sequence chunks sharded across CP ranks and use CP-aware linear kernels + (e.g. chunk_gated_delta_rule + causal_conv1d with a CP context). This follows the chunkwise + DeltaNet idea of storing state at chunk boundaries and doing chunk-local matrix work, avoiding + a full per-token recurrent state materialization while keeping tensor-core-friendly matmuls. + See https://sustcsonglin.github.io/blog/2024/deltanet-2/#a-chunkwise-algorithm-for-deltanet. + "headwise": Scatter heads across the CP group with all-to-all (Ulysses-style); each rank runs + the linear-attention kernel on the full sequence for a shard of heads. Correct but memory-heavy. + """ ################## # Cuda Graphs @@ -1086,7 +1097,7 @@ class TransformerConfig(ModelParallelConfig): determines the scope of graph capture.""" cuda_graph_use_single_mempool: bool = True - """For cuda_graph_impl "local" with cuda_graph_scope "full_iteration" only. + """For cuda_graph_impl "full_iteration" only. When True, full-iteration graph replay (training and evaluation) and optimizer graph capture/replay share the same CUDA graph memory pool.""" @@ -1097,9 +1108,7 @@ class TransformerConfig(ModelParallelConfig): more details, see: https://pytorch.org/docs/stable/generated/torch.Tensor.backward.html.""" cuda_graph_warmup_steps: int = 3 - """Number of warmup steps for CUDA graphs. Note: GTP (``gtp_weight_remat_size > 1``) forces a - minimum of 2 per-graph warmup steps regardless of this value, because the first warmup builds - the weight-prefetch chain and the second exercises the prefetch path before capture.""" + """Number of warmup steps for CUDA graphs""" external_cuda_graph: bool = False """DEPRECATED and replaced by cuda_graph_impl. @@ -1138,6 +1147,14 @@ class TransformerConfig(ModelParallelConfig): transformed to an empty list in __post_init__. The deprecated values "full_iteration" and "full_iteration_inference" are also accepted and migrated to the new API in __post_init__.""" + create_attention_mask_in_dataloader: bool = True + """Whether training data loaders create and pass an attention_mask tensor. + + This mirrors the training argument of the same name so Transformer Engine CUDA graph capture + can preserve the actual forward signature. When disabled for causal/no-mask attention, + the graph capture path should not allocate a synthetic global-sequence attention mask. + """ + inference_cuda_graph_scope: Optional[InferenceCudaGraphScope] = field( default=None, metadata={ @@ -1169,13 +1186,33 @@ class TransformerConfig(ModelParallelConfig): CudaGraphScope instances deserialized from pre-refactor checkpoints are converted to their string names before normalization so existing CUDA_GRAPH_MODULES_DEPRECATIONS handles them.""" + thd_max_packed_sequences: Optional[int] = field( + default=None, metadata={"argparse_meta": {"arg_names": ["--thd-max-packed-sequences"]}} + ) + """Optional maximum number of THD packed sequences per microbatch, including any + dummy sequence appended for a padding tail. The dp_balanced packing scheduler reserves + that dummy slot when THD padding appends one. When set, cu_seqlens tensors are padded + to this size + 1 in both eager and CUDA Graph modes. THD CUDA Graph requires this value. + + Sizing guidance: choose a value that comfortably covers the worst-case packing, + roughly ceil(max_seqlen_per_dp_cp_rank * cp_size / min_seq_len_after_filter). + Setting it too small results in more microbatches with smaller packs (wasted + token budget); setting it too large just allocates a slightly larger cu_seqlens + buffer.""" + + cuda_graph_dynamic_microbatches: bool = False + """Allow CUDA graph replay when runtime microbatch count varies across iterations. + This option is only meaningful for cuda_graph_impl=transformer_engine. For THD sequence + packing, capture uses a conservative upper bound on the packed microbatch count so graph + replay can cover iterations whose real packed microbatch count changes.""" + #################### # Hyper-Connection Configuration #################### - enable_mhc_connections: bool = False + enable_hyper_connections: bool = False """Enable mHC residual connections.""" - mhc_num_residual_streams: int = 4 + num_residual_streams: int = 4 """Number of residual streams (n in paper).""" mhc_sinkhorn_iterations: int = 20 @@ -1185,35 +1222,34 @@ class TransformerConfig(ModelParallelConfig): """Initial value of Gating Factor (alpha in paper).""" use_fused_mhc: bool = False - """Use fused kernels for mHC operations when supported. - - With the default ``auto`` backend policy, selection is operation-specific - and unavailable accelerated implementations fall back to native torch. - Set ``mhc_fused_backend`` to request an explicit backend policy. - """ - - mhc_fused_backend: Literal["auto", "native", "triton", "cutile"] = "auto" - """Backend policy for fused mHC operations. - - ``auto`` selects the fastest available implementation for each operation. - Explicit policies require the requested dependency and device support, and - never select a different accelerated backend. Operations without an - implementation in the selected backend retain their native implementation. + """Use unified fused kernels for mHC operations. + + When True, attempts to replace the reference mHC modules (SinkhornKnopp, + H_aggregate, H_post_bda, ProjRms) with fused/autograd implementations for + better performance on supported GPUs. Backend selection is internal and + op-specific: Triton for Sinkhorn and H_post_bda backward when available, + cuTile for the remaining fused kernels when available, then native torch + fallback. If every mHC operation uses the native torch fallback, + use_fused_mhc remains enabled and a rank-0 warning is emitted. The all-native + fallback is functionally equivalent, but may not provide fused backend + performance benefits. """ mhc_recompute_layer_num: Optional[int] = None - """Number of layers per MHC recompute block. + """Number of layers in each mHC recompute group. - When set, every `mhc_recompute_layer_num` layers form a recompute block. The last layer - in each recompute block (i.e., layer_number % mhc_recompute_layer_num == 0 or the final - layer in the transformer block) will: - - NOT checkpoint its final MLP BDA - - Register the unified recompute hook on its MLP BDA output - - A new CheckpointWithoutOutputManager is created for subsequent layers + Layers are grouped in their local transformer-block order. The last layer in each group leaves + its final MLP BDA output uncheckpointed and closes the current ``MHCCheckpointManager``; a new + manager is created for the next group. - If None, all layers in the transformer block share a single recompute block. + In the standard forward path, the group-ending layer registers a unified recompute hook on its + MLP BDA output. In the fine-grained expert-parallel overlap schedule, the group outputs are + discarded explicitly and a compute-stream schedule node replays the group before its backward + computation. - Must be a positive integer when set.""" + If ``None``, all local layers in the transformer block share one recompute group. The value must + be a positive integer when set. + """ #################### # miscellaneous @@ -1238,24 +1274,6 @@ class TransformerConfig(ModelParallelConfig): training and inference as the kernels are not full optimized. Defaults to False.""" - batch_invariant_backend: Literal["te_native", "deepgemm", "triton"] = "te_native" - """Which batch-invariant GEMM backend to use when batch_invariant_mode is - enabled: "te_native" (default: keep the native cuBLASLt kernels and obtain - invariance via workspace starvation — lowest overhead, no extra - dependencies, and the configuration verified bitwise-identical to the TE - training forward), "deepgemm" (DeepGEMM bf16 kernels), or "triton" - (persistent Triton matmul; any dtype).""" - - batch_invariant_collective: Literal["ordered", "multimem"] = "ordered" - """Cross-rank EP combine collective under batch_invariant_mode. "ordered" - (default) reduces with an explicit fixed rank-order fp32 Triton kernel — - deterministic by construction on any hardware. "multimem" keeps the native - NVLS in-switch reduce-scatter: measured to return the correctly-rounded - exact fp32 sum (bitwise-equal to an fp64 reference over 16.7M adversarial - channels on B200), deterministic and batch-invariant, with better scaling - at large NVLink domains; software paths that must match it bitwise should - accumulate in fp64.""" - use_te_activation_func: bool = False """Whether to use ffn activation functions implemented by TransformerEngine""" @@ -1313,6 +1331,12 @@ class TransformerConfig(ModelParallelConfig): """ Multimodal rope section is for channel dimension of temporal, height and width in rope calculation. """ + mrope_interleaved: bool = False + """When True, use the interleaved T/H/W MRoPE layout (Qwen3.5-VL style) where + H freqs occupy stride-3 positions {1,4,7,...} and W freqs occupy {2,5,8,...}. + When False (default), use the original section-based layout (Qwen2-VL style) + that cycles through T/H/W per mrope_section chunk.""" + is_hybrid_model: bool = False """ Indicates whether this is a hybrid model. """ @@ -1329,9 +1353,6 @@ class TransformerConfig(ModelParallelConfig): """The number of heads used in Mamba layers. If None, the number of heads will be hidden_size * expand // mamba_head_dim.""" - gdp_num_householder: int = 3 - """The number of Householder reflections used in Gated Delta Product layers.""" - mamba_training_ssm_states_dtype: Optional[torch.dtype] = None """dtype of the materialized inter-chunk SSM states in Mamba training forwards and backwards. None causes the states to follow the activation dtype.""" @@ -1341,17 +1362,6 @@ class TransformerConfig(ModelParallelConfig): ) """Controls usage of the memory efficient path for Mamba layers.""" - gdp_cutedsl_kernel: bool = False - """Whether to use the CuTeDSL kernel for the GatedDeltaProduct mixer.""" - - gdp_num_chunk_states_to_recompute: int = 2 - """Checkpoint-coarsening ratio N in [0, 64] for the CuTeDSL GatedDeltaProduct kernel. - N=0 checkpoints every chunk state for the backward pass (dense, no recompute). N>=1 - stores one checkpoint per group of N+1 chunks (1/(N+1) the checkpoint memory) and the - backward recomputes each group's N missing chunk states, trading recompute time for - activation memory monotonically in N (sweet spot N=2..3). Only honored by kernel - builds that expose the num_chunk_states_to_recompute argument.""" - mlp_chunks_for_prefill: int = 1 """The number of chunks along the sequence dimension to use for MLP computation during prefill.""" @@ -1377,6 +1387,9 @@ class TransformerConfig(ModelParallelConfig): """Transformer implementation to use. Options are 'transformer_engine' for Transformer Engine and 'local' for MCore.""" + fallback_to_eager_attn: bool = False + """Whether to fallback to eager attention in TE implementation. + Suggested for when desired features are not available in TE implementation.""" ##################################### # Fine-grained Activation Offloading ##################################### @@ -1388,7 +1401,7 @@ class TransformerConfig(ModelParallelConfig): offload_modules: Optional[list[str]] = field(default_factory=list) """The submodules to offload its input. choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj", - "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp", "gdp_qkv". + "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp". "attn_norm": offload the input of the normalization in the attention part. "qkv_linear": offload the input of the qkv linear part. "core_attn": offload the input of the core attention part. @@ -1397,8 +1410,6 @@ class TransformerConfig(ModelParallelConfig): "expert_fc1": offload the input of the expert fc1 part. "moe_act": offload the input of the moe act part. "fused_group_mlp": offload the input of the whole fused grouped MLP. - "gdp_qkv": offload the input of the causal conv and QKV preparation in the - GatedDeltaProduct mixer. """ min_offloaded_tensor_size: int = 1024 * 1024 """The minimum size of the tensor to be offloaded.""" @@ -1421,6 +1432,14 @@ class TransformerConfig(ModelParallelConfig): https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#activation-offload-fraction. """ + fine_grained_offloading_max_inflight_offloads: Optional[int] = None + """Per fine-grained offloading group name, max number of inflight offloads for that name not + yet joined on the main stream (wait_event on D2H). The same cap applies to every name (e.g., + ``moe_act`` and ``qkv_linear`` each have their own pending queue). 0 = wait after every + offload for that name. 1 = at most one not-yet-waited offload per name, etc. None = do not + insert these joins. This feature is particularly useful when using with full-iteration CUDA + graphs""" + moe_paged_stash: bool = False """If True, enable paged stash for all routed-expert activations needed for backward""" @@ -1431,7 +1450,7 @@ class TransformerConfig(ModelParallelConfig): """Scale factor for paged stash CUDA buffer allocation. Sign selects sizing: positive = avg-based, negative = actual-max. Magnitude is headroom - (e.g. 1.10 = 10%%).""" + (e.g. 1.10 = 10%).""" moe_paged_stash_buffer_size_factor_cpu: float = 0.0 """Scale factor for paged stash host buffer. 0 disables host buffer. @@ -1439,83 +1458,12 @@ class TransformerConfig(ModelParallelConfig): Same sign convention as moe_paged_stash_buffer_size_factor_cuda: positive = avg-based, negative = actual-max; scale = abs(factor).""" - fine_grained_offloading_max_inflight_offloads: Optional[int] = None - """Per fine-grained offloading group name, max number of inflight offloads for that name not - yet joined on the main stream (wait_event on D2H). The same cap applies to every name (e.g., - ``moe_act`` and ``qkv_linear`` each have their own pending queue). 0 = wait after every - offload for that name. 1 = at most one not-yet-waited offload per name, etc. None = do not - insert these joins. This feature is particularly useful when using with full-iteration CUDA - graphs""" - - def _validate_cp_layouts(self) -> None: - """Validate context-parallel layout settings.""" - if self.linear_cp_layout not in ("contiguous", "zigzag"): - raise ValueError( - "linear_cp_layout must be either 'contiguous' or 'zigzag', " - f"got {self.linear_cp_layout!r}" - ) - if self.attention_cp_layout not in ("contiguous", "zigzag"): - raise ValueError( - "attention_cp_layout must be either 'contiguous' or 'zigzag', " - f"got {self.attention_cp_layout!r}" - ) - if self.context_parallel_size > 1 and self.attention_cp_layout == "contiguous": - raise ValueError( - "attention_cp_layout='contiguous' is not yet supported with context parallelism." - ) - if self.linear_cp_layout == "contiguous" and self.hybrid_context_parallel: - raise ValueError( - "hybrid_context_parallel is not supported with linear_cp_layout='contiguous'." - ) - if ( - self.context_parallel_size > 1 - and self.linear_cp_layout != self.attention_cp_layout - and self.sequence_parallel - and self.tensor_model_parallel_size > 1 - and self.tensor_model_parallel_size % 2 != 0 - ): - raise ValueError( - "Sequence-parallel CP layout conversion requires an even " - f"tensor-parallel size, got {self.tensor_model_parallel_size}." - ) - if ( - self.linear_cp_layout == "contiguous" - and self.context_parallel_size > 1 - and (self.mtp_num_layers or 0) > 0 - ): - raise ValueError( - "linear_cp_layout='contiguous' with context parallelism does not yet support MTP." - ) - def __post_init__(self): """Python dataclass method that is used to modify attributes after initialization. See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more details. """ super().__post_init__() - self._validate_cp_layouts() - - # Resolve deprecated attention variant spellings up front so that every consumer - # downstream only has to handle the canonical names. Imported lazily because the - # spec module imports this one. - from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( - is_gated_delta_net_variant, - normalize_experimental_attention_variant, - ) - - if self.experimental_attention_variant is not None: - self.experimental_attention_variant = normalize_experimental_attention_variant( - self.experimental_attention_variant - ) - - if self.use_transformer_engine_op_fuser and self.moe_grouped_gemm: - self.moe_use_grouped_tensor = True - - if self.moe_use_grouped_tensor and not self.moe_grouped_gemm: - raise ValueError("moe_use_grouped_tensor=True requires moe_grouped_gemm=True.") - - if self.mtp_hsm and (self.mtp_num_layers is None or self.mtp_num_layers < 2): - raise ValueError("mtp_hsm=True requires mtp_num_layers >= 2.") # When fp32 residual connections are enabled, pipeline parallel communication must # use fp32 to match the dtype of the residual stream between pipeline stages. @@ -1535,11 +1483,6 @@ def __post_init__(self): f"Only one of self.fp16: {self.fp16} and self.bf16 {self.bf16} should be True." ) - if self.gdp_num_householder < 1: - raise ValueError( - f"gdp_num_householder must be positive, got {self.gdp_num_householder}." - ) - # Apply BF16 matmul precision setting if needed if self.bf16 and self.disable_bf16_reduced_precision_matmul: torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False @@ -1568,45 +1511,203 @@ def __post_init__(self): f"tensor_model_parallel_size ({self.tensor_model_parallel_size})." ) - if is_gated_delta_net_variant(self.experimental_attention_variant): - # gdn2 may also be enabled for GDN layers built via the hybrid layer pattern - # symbol 'G', where linear_attention_freq is unused; the GPT experimental - # attention route raises a clear error downstream if it is missing. - if self.experimental_attention_variant == "gdn": - assert ( - self.linear_attention_freq is not None - ), "linear_attention_freq must be set for linear gdn." + if self.head_wise_attn_gate: + tp = self.tensor_model_parallel_size + if self.attention_output_gate: + raise ValueError( + "head_wise_attn_gate and attention_output_gate cannot both be enabled " + "(incompatible linear_qkv row layouts)." + ) + if self.num_attention_heads % tp != 0: + raise ValueError( + f"head_wise_attn_gate requires num_attention_heads " + f"({self.num_attention_heads}) divisible by tp ({tp})." + ) + # The gate peel-off in SelfAttention runs before the + # num_query_groups < world_size AllGather+reslice fallback, so the + # gate cannot ride that path without over-taking V/K rows. + if self.num_query_groups < tp: + raise ValueError( + f"head_wise_attn_gate requires num_query_groups " + f"({self.num_query_groups}) >= tp ({tp})." + ) + # TE FP8/FP4 GEMM requires per-partition output dim aligned to + # 16/32. attention_output_gate's wider gate gets this for free; + # head_wise's num_attention_heads-row increment can mis-align. + if self.fp8 is not None or self.fp4 is not None: + align = 32 if self.fp4 is not None else 16 + linear_qkv_out_dim = ( + self.kv_channels * self.num_attention_heads + + 2 * self.kv_channels * self.num_query_groups + + self.num_attention_heads + ) + per_partition = linear_qkv_out_dim // tp + if per_partition % align != 0: + fp_name = "fp4" if self.fp4 is not None else "fp8" + raise ValueError( + f"head_wise_attn_gate under {fp_name} requires per-partition " + f"linear_qkv output dim ({per_partition}) to be a multiple of " + f"{align}; got num_attention_heads={self.num_attention_heads}, " + f"num_query_groups={self.num_query_groups}, " + f"kv_channels={self.kv_channels}, tp={tp}." + ) - # Check required parameters - assert ( - self.linear_conv_kernel_dim is not None - ), "linear_conv_kernel_dim must be set for gated delta net." - assert ( - self.linear_key_head_dim is not None - ), "linear_key_head_dim must be set for gated delta net." - assert ( - self.linear_value_head_dim is not None - ), "linear_value_head_dim must be set for gated delta net." - assert ( - self.linear_num_key_heads is not None - ), "linear_num_key_heads must be set for gated delta net." - assert ( - self.linear_num_value_heads is not None - ), "linear_num_value_heads must be set for gated delta net." - assert self.linear_num_value_heads % self.linear_num_key_heads == 0, ( - f"linear_num_value_heads ({self.linear_num_value_heads}) must be a multiple of " - f"linear_num_key_heads ({self.linear_num_key_heads})." + 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.cp_partition_mode not in ("zigzag", "contiguous"): + raise ValueError(f"Unsupported cp_partition_mode: {self.cp_partition_mode}") + + if self.cp_partition_mode == "contiguous" and ( + self.context_parallel_size > 1 or self.dynamic_context_parallel + ): + if self.sequence_packing_scheduler is None: + raise ValueError( + "cp_partition_mode='contiguous' with context parallelism requires THD " + "inputs from a sequence_packing_scheduler; BSHD inputs are not supported. " + "The legacy non-scheduler CP batch slicing path only supports zigzag layout." + ) + if ( + (self.mtp_num_layers or 0) > 0 + and self.tensor_model_parallel_size > 1 + and self.sequence_parallel + ): + raise ValueError( + "MTP with tensor_model_parallel_size > 1, sequence_parallel=True, and " + "cp_partition_mode='contiguous' has a known token-side padding-mask layout " + "bug. This combination is temporarily unsupported and will be fixed in a " + "follow-up change." + ) + + if self.context_parallel_size > 1: + if self.cp_partition_mode == "contiguous": + if ( + self.multi_latent_attention + and self.experimental_attention_variant != "dsv4_hybrid" + ): + raise ValueError( + "cp_partition_mode='contiguous' is not supported with " + "multi_latent_attention outside dsv4_hybrid." + ) + if self.experimental_attention_variant not in ("dsv4_hybrid", "gated_delta_net"): + raise ValueError( + "cp_partition_mode='contiguous' with context parallelism currently " + "requires experimental_attention_variant to be either 'dsv4_hybrid' " + "or 'gated_delta_net'." + ) + if ( + self.experimental_attention_variant == "gated_delta_net" + and self.linear_cp_mode == "headwise" + ): + raise ValueError( + "cp_partition_mode='contiguous' is incompatible with " + "gated_delta_net linear_cp_mode='headwise'." + ) + elif self.cp_partition_mode == "zigzag": + if self.experimental_attention_variant == "dsv4_hybrid": + raise ValueError( + "DSv4 Hybrid with context parallelism requires " + "cp_partition_mode='contiguous'." + ) + + # Normalize the deprecated DSv4 kernel switch only after all deprecated attention + # selectors have been folded into experimental_attention_variant, and immediately + # before the centralized attention-variant validation consumes dsa_kernel_backend. + if self.apply_dsa_kernel_fusion is not None: + if self.experimental_attention_variant != "dsv4_hybrid": + log_single_rank( + logger, + logging.WARNING, + "apply_dsa_kernel_fusion is deprecated and ignored outside " + "experimental_attention_variant='dsv4_hybrid'; use dsa_kernel_backend " + "instead.", + ) + else: + legacy_backend = "cudnn" if self.apply_dsa_kernel_fusion else "none" + if self.dsa_kernel_backend not in ("none", legacy_backend): + raise ValueError( + "Conflicting DSA kernel controls: " + f"apply_dsa_kernel_fusion={self.apply_dsa_kernel_fusion} maps to " + f"dsa_kernel_backend={legacy_backend!r}, but " + f"dsa_kernel_backend={self.dsa_kernel_backend!r} was also selected. " + "Remove apply_dsa_kernel_fusion and use dsa_kernel_backend only." + ) + log_single_rank( + logger, + logging.WARNING, + "apply_dsa_kernel_fusion is deprecated and will be removed in a future " + f"release; use dsa_kernel_backend={legacy_backend!r} instead.", + ) + self.dsa_kernel_backend = legacy_backend + + if self.experimental_attention_variant in ["gated_delta_net"]: + assert ( + self.linear_attention_freq is not None + ), f"linear_attention_freq must be set for linear attention." + + if self.experimental_attention_variant == "gated_delta_net": + if self.pad_packed_seq_alignment is not None: + tail_policy = self.thd_tail_padding_policy or 'append_dummy_seq' + assert tail_policy == 'append_dummy_seq', ( + "gated_delta_net with pad_packed_seq_alignment requires " + "thd_tail_padding_policy='append_dummy_seq'." + ) + + # Check required parameters + assert ( + self.linear_conv_kernel_dim is not None + ), "linear_conv_kernel_dim must be set for gated delta net." + assert ( + self.linear_key_head_dim is not None + ), "linear_key_head_dim must be set for gated delta net." + assert ( + self.linear_value_head_dim is not None + ), "linear_value_head_dim must be set for gated delta net." + assert ( + self.linear_num_key_heads is not None + ), "linear_num_key_heads must be set for gated delta net." + assert ( + self.linear_num_value_heads is not None + ), "linear_num_value_heads must be set for gated delta net." + assert self.linear_num_value_heads % self.linear_num_key_heads == 0, ( + f"linear_num_value_heads ({self.linear_num_value_heads}) must be a multiple of " + f"linear_num_key_heads ({self.linear_num_key_heads})." + ) + 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: + assert self.linear_cp_mode in ("headwise", "chunkwise"), ( + f"linear_cp_mode must be one of 'headwise' or 'chunkwise', " + f"got {self.linear_cp_mode!r}." + ) + if self.gdn_conv_pad_alignment is not None: + assert self.linear_cp_mode != "chunkwise", ( + "gdn_conv_pad_alignment is incompatible with " + "linear_cp_mode='chunkwise' when context_parallel_size > 1. " + "Padding chunk-local GDN causal-conv inputs can change later " + "chunk numerics." + ) + # Check tensor parallelism compatibility. Headwise CP splits linear-attention heads + # across CP ranks; chunkwise CP keeps all TP-local heads on each CP rank. + linear_head_parallel_size = self.tensor_model_parallel_size + if self.context_parallel_size > 1 and self.linear_cp_mode == "headwise": + linear_head_parallel_size *= self.context_parallel_size + assert self.linear_num_key_heads % linear_head_parallel_size == 0, ( f"{self.linear_num_key_heads=} must be a multiple of " - f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." + f"{linear_head_parallel_size=} for {self.linear_cp_mode=}." ) - assert self.linear_num_value_heads % tp_cp_size == 0, ( + assert self.linear_num_value_heads % linear_head_parallel_size == 0, ( f"{self.linear_num_value_heads=} must be a multiple of " - f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." + f"{linear_head_parallel_size=} for {self.linear_cp_mode=}." ) elif self.experimental_attention_variant == "dsa": _validate_dsa_kernel_backend_dependencies(self.dsa_kernel_backend) @@ -1624,6 +1725,106 @@ def __post_init__(self): "dsa_indexer_skip_topk_offset must be non-negative, got " f"{self.dsa_indexer_skip_topk_offset}." ) + assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention" + if self.context_parallel_size > 1: + cp_comm_types = ( + self.cp_comm_type + if isinstance(self.cp_comm_type, list) + else [self.cp_comm_type] + ) + assert all( + cp_comm_type is not None + and cp_comm_type.replace("_", "").lower() == "allgather" + for cp_comm_type in cp_comm_types + ), ( + "DSAttention context parallelism currently supports " + "cp_comm_type=allgather only." + ) + elif self.experimental_attention_variant == "dsv4_hybrid": + assert self.multi_latent_attention, "DSv4 Hybrid requires multi_latent_attention." + assert self.csa_compress_ratios is not None, "csa_compress_ratios must be set" + mtp_layers = self.mtp_num_layers or 0 + # Minimum length is num_layers + mtp_num_layers (the GPT path uses exactly this, + # where mtp_num_layers == #MTP transformer layers). On HybridModel an MTP "depth" + # can expand to MULTIPLE hybrid layers, so mtp_num_layers (= depth count) undercounts + # the real MTP layers and csa_compress_ratios must be at least long enough to index + # every MTP attention layer (num_layers + layer_number - 1). Hence ">=", not "==". + expected_len = self.num_layers + mtp_layers + assert len(self.csa_compress_ratios) >= expected_len, ( + f"csa_compress_ratios length ({len(self.csa_compress_ratios)}) must be at least " + f"num_layers + mtp_num_layers ({self.num_layers} + {mtp_layers} = {expected_len})" + ) + assert all( + ratio in [0, 4, 128] for ratio in self.csa_compress_ratios + ), "csa_compress_ratios must be 0, 4, or 128" + assert ( + self.tensor_model_parallel_size == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + assert not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention." + self.hetereogenous_dist_checkpoint = True + + if self.dsa_kernel_backend == "tilelang": + raise ValueError( + "dsv4_hybrid does not support dsa_kernel_backend='tilelang'; use 'cudnn' " + "for fused CSA kernels or 'none' for the PyTorch fallback." + ) + _validate_dsa_kernel_backend_dependencies(self.dsa_kernel_backend) + + if self.dsa_kernel_backend == "cudnn": + sm = torch.cuda.get_device_capability() + assert sm[0] >= 9, ( + f"dsa_kernel_backend='cudnn' requires SM90+ (Hopper or later), " + f"but current device has compute capability {sm[0]}.{sm[1]}." + ) + uses_ratio4_indexer = 4 in self.csa_compress_ratios and not self.csa_dense_mode + indexer_loss_enabled = (self.dsa_indexer_loss_coeff or 0.0) > 0 + if ( + sm[0] == 9 + and uses_ratio4_indexer + and indexer_loss_enabled + and not self.dsa_indexer_use_sparse_loss + ): + raise ValueError( + "DSv4 with fused DSA and dense indexer loss is not supported on SM90 " + "because the cuDNN Frontend SM90 dense DSA kernels are not reliable for " + "this path. Use sparse indexer loss or set dsa_kernel_backend='none'." + ) + + from cudnn import DSA + + if ( + self.context_parallel_size > 1 or self.dynamic_context_parallel + ) and uses_ratio4_indexer: + required_wrappers = [DSA.indexer_forward_wrapper] + if indexer_loss_enabled and not self.dsa_indexer_use_sparse_loss: + required_wrappers.extend( + [ + DSA.dense_indexer_score_recompute_wrapper, + DSA.dense_attn_score_recompute_wrapper, + DSA.dense_indexer_backward_wrapper, + ] + ) + missing_offsets = [ + wrapper.__name__ + for wrapper in required_wrappers + if "q_causal_offsets" not in inspect.signature(wrapper).parameters + ] + if missing_offsets: + raise ValueError( + "DSv4 CP with ratio-4 fused DSA requires cuDNN Frontend wrappers " + "with q_causal_offsets support; missing from: " + f"{', '.join(missing_offsets)}. Install a compatible cuDNN Frontend " + "build or set dsa_kernel_backend='none'." + ) + + if ( + self.gdn_pre_gated_delta_rule_fusion + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "gdn_pre_gated_delta_rule_fusion is only supported with " + "experimental_attention_variant='gated_delta_net'." + ) if self.fp8: # cannot support first last layer bf16 with delayed scaling @@ -1714,16 +1915,10 @@ def __post_init__(self): "to avoid costly dtype conversions during decode." ) - # Gated linear units (SwiGLU/GeGLU) are supported by the torch and vllm - # grouped-GEMM backends only. - if self.gated_linear_unit and self.inference_grouped_gemm_backend not in ( - "torch", - "vllm", - ): + if self.gated_linear_unit: raise ValueError( - "--transformer-impl='inference_optimized' supports gated linear units " - "(SwiGLU/GeGLU) only with --inference-grouped-gemm-backend torch or vllm, " - f"got '{self.inference_grouped_gemm_backend}'." + "--transformer-impl='inference_optimized' does not yet support " + "gated linear units (SwiGLU/GeGLU)." ) if self.fp8 == "mxfp8": @@ -1762,24 +1957,6 @@ def __post_init__(self): "Set inference_grouped_gemm_backend to 'torch' for MXFP8." ) - if self.batch_invariant_mode: - if self.inference_grouped_gemm_backend not in ( - InferenceGroupedGemmBackend.TORCH, - InferenceGroupedGemmBackend.VLLM, - ): - raise ValueError( - "batch_invariant_mode requires inference_grouped_gemm_backend " - "'torch' or 'vllm'." - ) - if ( - self.expert_model_parallel_size > 1 - and self.inference_moe_token_dispatcher_type != "nvls" - ): - raise ValueError( - "batch_invariant_mode with inference-optimized MoE and expert " - "parallelism requires inference_moe_token_dispatcher_type='nvls'." - ) - if self.num_moe_experts is not None and self.num_moe_experts <= 0: raise ValueError("num_moe_experts must be non-negative.") @@ -1828,23 +2005,20 @@ def __post_init__(self): "moe_single_grouped_weight is currently supported with high-precision " "primary weights, fp8_recipe='mxfp8', or fp4_recipe='nvfp4'." ) - if self.fp4 and not self.fp4_param: + if not self.use_transformer_engine_op_fuser: raise ValueError( - "moe_single_grouped_weight with FP4 compute requires fp4_param=True " - "(--fp4-param-gather). Without FP4 parameter gather, Transformer Engine " - "uses a split-quantize fallback that is being deprecated." + "moe_single_grouped_weight requires " + "use_transformer_engine_op_fuser=True. The non-op-fuser TE GroupedLinear " + "path splits the grouped parameter into per-expert tensors and does not " + "support single-grouped-weight training." ) - if not self.moe_use_grouped_tensor: - raise ValueError("moe_single_grouped_weight requires moe_use_grouped_tensor=True.") if self.moe_single_grouped_bias and not self.add_bias_linear: raise ValueError("moe_single_grouped_bias requires add_bias_linear=True.") - if self.moe_single_grouped_bias and not self.moe_use_grouped_tensor: - raise ValueError("moe_single_grouped_bias requires moe_use_grouped_tensor=True.") if self.moe_enable_deepep: if self.moe_token_dispatcher_type != "flex": raise ValueError("DeepEP backend is only supported with flex token dispatcher.") - if self.moe_flex_dispatcher_backend == "hybridep": + if self.moe_flex_dispatcher_backend in ("deepepv2", "hybridep"): raise ValueError("Only one backend is supported for flex token dispatcher.") self.moe_flex_dispatcher_backend = "deepep" warnings.warn( @@ -1854,10 +2028,10 @@ def __post_init__(self): if self.moe_token_dispatcher_type == "flex": if self.moe_pad_expert_input_to_capacity and ( - self.moe_enable_deepep or self.moe_flex_dispatcher_backend == "deepep" + self.moe_enable_deepep or self.moe_flex_dispatcher_backend in ("deepep", "deepepv2") ): raise ValueError( - "Flex token dispatcher with deepep backend does not support " + "Flex token dispatcher with deepep/deepepv2 backend does not support " "moe_pad_expert_input_to_capacity" ) @@ -1867,31 +2041,6 @@ def __post_init__(self): "moe_flex_dispatcher_backend='ncclep' requires " "moe_token_dispatcher_type='flex'." ) - if self.moe_use_grouped_tensor and not self.use_transformer_engine_op_fuser: - raise ValueError( - "moe_use_grouped_tensor=True without use_transformer_engine_op_fuser is " - "not yet supported with the NCCL-EP dispatcher. Use the TE op-fuser path " - "or select the alltoall, DeepEP, or HybridEP dispatcher." - ) - - if self.moe_dispatch_fwd_dtype != 'bf16' or self.moe_combine_bwd_dtype != 'bf16': - if ( - self.moe_token_dispatcher_type != "flex" - or self.moe_flex_dispatcher_backend != "ncclep" - ): - raise ValueError( - "moe_dispatch_fwd_dtype / moe_combine_bwd_dtype require the 'ncclep' flex " - "dispatcher backend." - ) - if not (self.use_transformer_engine_op_fuser and self.moe_grouped_gemm): - raise ValueError( - "moe_dispatch_fwd_dtype / moe_combine_bwd_dtype = 'mxfp8' require BOTH " - "use_transformer_engine_op_fuser and moe_grouped_gemm: only the fused " - "grouped GEMM path consumes the pre-quantized MXFP8 GroupedTensor payload." - ) - - if self.moe_use_norm_before_up_proj and self.moe_latent_size is None: - raise ValueError("moe_use_norm_before_up_proj requires moe_latent_size to be set.") # moe_deepep_num_sms / moe_hybridep_num_sms are deprecated and unified into # moe_flex_dispatcher_num_sms. If either is set, route it (an explicit @@ -1978,11 +2127,10 @@ def __post_init__(self): if ( self.moe_flex_dispatcher_backend == "hybridep" and not self.use_transformer_engine_op_fuser - and not self.moe_use_grouped_tensor ): raise ValueError( "moe_expert_rank_capacity_factor with the 'hybridep' backend requires " - "use_transformer_engine_op_fuser=True or moe_use_grouped_tensor=True." + "use_transformer_engine_op_fuser to be enabled." ) if self.cpu_offloading and ( @@ -2054,10 +2202,9 @@ def __post_init__(self): "mlp", "moe", "shared_experts", - "gdn_norm_out", - "gdp_in_proj", - "gdp_qkv", "mhc", + "gdn", + "gdn_norm_out", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -2076,14 +2223,20 @@ 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" in self.recompute_modules + and self.experimental_attention_variant != "gated_delta_net" ): raise ValueError( - "gdn_norm_out in recompute_modules is only supported with " - "experimental_attention_variant='gdn' or 'gdn2'." + "gdn in recompute_modules is only supported with " + "experimental_attention_variant='gated_delta_net'." ) + 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 GatedDeltaNet module, including gated norm." + ) if "core_attn" in self.recompute_modules: warnings.warn( "If you are using transformer_engine as the transformer implementation, " @@ -2130,8 +2283,10 @@ def __post_init__(self): # Validation for "mhc" in recompute_modules if self.recompute_granularity == "selective" and "mhc" in self.recompute_modules: - if not self.enable_mhc_connections: - raise ValueError("'mhc' in recompute_modules requires enable_mhc_connections=True.") + if not self.enable_hyper_connections: + raise ValueError( + "'mhc' in recompute_modules requires enable_hyper_connections=True." + ) if "mlp" in self.recompute_modules: raise ValueError( "'mhc' and 'mlp' in recompute_modules cannot be used together. " @@ -2146,92 +2301,94 @@ def __post_init__(self): "mhc_recompute_layer_num must be a positive integer when " "'mhc' is in recompute_modules." ) - if self.fine_grained_activation_offloading: - raise NotImplementedError( - "'mhc' in recompute_modules + fine_grained_activation_offloading is " - "not yet supported. The mHC recompute hook currently fires before " - "the offloading backward chunk is initialized, causing tensor_pop " - "on a None chunk. Disable one of them." - ) + if self.fine_grained_activation_offloading and self.offload_modules: + # mHC checkpoints wrap input_layernorm (inside attn_norm offload context) + # and pre_mlp_layernorm (inside mlp_norm offload context). The unified + # recompute trigger runs before GroupCommitFunction.backward() initializes + # the backward chunk, so tensor_pop hits a None chunk for these modules. + # Other offload modules (qkv_linear, core_attn, attn_proj, expert_fc1, + # moe_act) live inside self_attention/MLP which are NOT wrapped by mHC + # checkpoints, so they are safe to use with mHC recompute -- except + # under the attention-only CUDA Graph split, which narrows the three + # attention-scope modules further down in __post_init__ (the split + # predicate needs normalized cuda_graph_modules, which are not + # available yet here). + _MHC_CONFLICTING_OFFLOAD_MODULES = {"attn_norm", "mlp_norm"} + conflicting = _MHC_CONFLICTING_OFFLOAD_MODULES & set(self.offload_modules) + if conflicting: + raise ValueError( + f"'mhc' in recompute_modules is incompatible with " + f"offload_modules {conflicting}. The mHC recompute replay starts " + f"before the offloading backward chunk is initialized for these " + f"modules, causing tensor_pop on a None chunk. Remove " + f"{conflicting} from offload_modules or remove 'mhc' from " + f"recompute_modules." + ) - if self.enable_mhc_connections and not ( + use_mhc_recompute = ( self.recompute_granularity == "selective" and "mhc" in self.recompute_modules + ) + # NOTE: the mHC-recompute/CUDA-graph capability gate lives further down in + # __post_init__, after cuda_graph_modules normalization and the deprecated + # enable_cuda_graph/external_cuda_graph migration, so that string module + # forms and legacy flags reach the same gate. + + # The attention-only Transformer Engine split keeps the mHC producer and + # its checkpoint registration eager per microbatch (only input-layernorm + # + self-attention are captured), so it composes with EP a2a overlap: + # the captured attention backward reads the fixed arena slot that the + # group-end recompute node repopulates by direct kernel output. + # Full-iteration capture is also valid because it records the explicit + # schedule barrier and its fixed-address recompute kernels as part of + # the iteration graph. Other per-layer graph forms remain rejected. + # Normalize locally: cuda_graph_modules is only normalized in bulk further + # down in __post_init__, so a config built directly (rather than through + # megatron/training/arguments.py) still holds string forms like ["attn"] + # here. Comparing them raw would reject a combination the later gate + # accepts, and the string form is part of this field's declared type. + normalized_graph_modules, _, _ = normalize_cuda_graph_modules(self.cuda_graph_modules) + is_te_attn_split = ( + self.cuda_graph_impl == "transformer_engine" + and not self.enable_cuda_graph + and not self.external_cuda_graph + and list(normalized_graph_modules or []) == [CudaGraphModule.attn] + and list(self.recompute_modules) == ["mhc"] + ) + is_full_iteration = ( + self.cuda_graph_impl == "full_iteration" + and not self.enable_cuda_graph + and not self.external_cuda_graph + ) + if ( + self.overlap_moe_expert_parallel_comm + and use_mhc_recompute + and not is_te_attn_split + and not is_full_iteration + and ( + self.cuda_graph_impl != "none" or self.enable_cuda_graph or self.external_cuda_graph + ) ): + raise ValueError( + "mHC recompute with overlap_moe_expert_parallel_comm requires CUDA graphs " + "to be disabled, except for the attention-only Transformer Engine split " + "(cuda_graph_impl='transformer_engine', cuda_graph_modules=[attn], " + "recompute_modules=[mhc]) or full-iteration capture " + "(cuda_graph_impl='full_iteration', hidden_dropout=0, " + "attention_dropout=0), because other graph scopes cannot represent the " + "explicit schedule-owned recompute barrier." + ) + + if self.enable_hyper_connections and not use_mhc_recompute: warnings.warn( "HyperConnections are enabled but 'mhc' is not in " "recompute_modules with selective recompute. Consider adding 'mhc' to " "recompute_modules with selective recompute to reduce activation memory." ) - if self.use_fused_mhc and not self.enable_mhc_connections: - raise ValueError("use_fused_mhc requires enable_mhc_connections=True.") - - valid_mhc_fused_backends = ("auto", "native", "triton", "cutile") - if self.mhc_fused_backend not in valid_mhc_fused_backends: - raise ValueError( - f"Unknown mhc_fused_backend {self.mhc_fused_backend!r}; expected one of " - f"{valid_mhc_fused_backends}." - ) - if self.mhc_fused_backend != "auto" and not self.use_fused_mhc: - raise ValueError("mhc_fused_backend requires use_fused_mhc=True when set explicitly.") - - if self.enable_mhc_connections and self.recompute_granularity == "full": - raise NotImplementedError( - "enable_mhc_connections is not yet compatible with full activation recompute. " - "Use selective recompute with 'mhc' in recompute_modules, or disable " - "activation recompute." - ) - - if self.enable_mhc_connections and self.inference_fuse_tp_communication: - raise NotImplementedError( - "enable_mhc_connections is not compatible with inference_fuse_tp_communication. " - "The fused inference TP path assumes single-stream residual tensors." - ) - - # Note: mHC + MoE is deliberately NOT rejected here. HyperConnectionTransformerLayer - # raises for a MoE MLP submodule at build time, which is the precise check; a - # config-level `num_moe_experts` guard would also block the documented composition - # (wrapping MoE as a HybridStack layer via HyperConnectionHybridLayer). - - if self.enable_mhc_connections: - # TransformerBlock expands to n-stream at `pre_process` and contracts back at - # the stage holding the final layernorm, so every intermediate pipeline stage - # exchanges [s, b, n*C] while the p2p buffers are still sized from hidden_size. - # Pipeline support must resize the p2p buffers before this guard can be lifted. - if self.pipeline_model_parallel_size > 1: - raise NotImplementedError( - "enable_mhc_connections does not support pipeline_model_parallel_size > 1 " - "yet. Inter-stage activations are n-stream ([s, b, n*C]) while pipeline " - "p2p buffers are sized from hidden_size, so the shapes disagree. Use " - "pipeline_model_parallel_size=1 until mHC pipeline support lands." - ) - - # The residual carried across an mHC layer is the n-stream tensor consumed by - # `fused_h_res_h_post_bda` (a bmm against h_res), not the single-stream residual - # the base layer casts. Upcasting it alone would mismatch the h_res dtype. - if self.fp32_residual_connection: - raise NotImplementedError( - "enable_mhc_connections is not compatible with fp32_residual_connection. " - "The mHC residual is the n-stream tensor consumed by the H_res batched " - "matmul, which must share the activation dtype." - ) - - # SinkhornKnopp.backward's analytic chain rule (grad = grad_M_init * M_init) - # is exact only because the per-row max shift is annihilated by the first - # row-normalization T_r. With zero iterations that invariance does not hold - # and the dropped argmax term makes the gradient silently wrong. - if self.mhc_sinkhorn_iterations < 1: - raise ValueError( - "mhc_sinkhorn_iterations must be >= 1; the Sinkhorn-Knopp backward " - f"assumes at least one row-normalization pass, got " - f"{self.mhc_sinkhorn_iterations}." - ) - - if self.mhc_init_gating_factor < 0: - raise ValueError( - "mhc_init_gating_factor must be non-negative, got " - f"{self.mhc_init_gating_factor}." - ) + # Validation for use_fused_mhc + if self.use_fused_mhc: + if not self.enable_hyper_connections: + raise ValueError("use_fused_mhc requires enable_hyper_connections=True.") if self.fine_grained_activation_offloading: assert ( @@ -2247,7 +2404,6 @@ def __post_init__(self): "attn_norm", "mlp_norm", "qkv_linear", - "gdp_qkv", } invalid_modules = set(self.offload_modules) - allowed_modules assert not invalid_modules, ( @@ -2260,17 +2416,6 @@ def __post_init__(self): "because the input of attn_proj is the output of core_attn, " "which is needed in core_attn.backward()." ) - if ( - "gdp_qkv" in self.offload_modules - and self.recompute_granularity == "selective" - and "gdp_in_proj" in self.recompute_modules - ): - raise ValueError( - "gdp_qkv cannot be set in offload_modules together with gdp_in_proj in " - "recompute_modules, because gdp_in_proj discards the input of the causal " - "conv and rematerializes it at the start of the mixer backward, leaving " - "nothing for the gdp_qkv offload group to keep on the host." - ) if self.recompute_granularity == "selective" and "moe" in self.recompute_modules: offload_inside_moe = {"moe_act", "expert_fc1", "fused_group_mlp"} & set( self.offload_modules @@ -2291,7 +2436,10 @@ def __post_init__(self): assert ( self.delta_offload_bytes_across_pp_ranks >= 0 ), "delta_offload_bytes_across_pp_ranks must be non-negative." - + if self.fine_grained_offloading_max_inflight_offloads is not None: + assert ( + self.fine_grained_offloading_max_inflight_offloads >= 0 + ), "fine_grained_offloading_max_inflight_offloads must be non-negative when set." if "fused_group_mlp" in self.offload_modules: if not self.use_transformer_engine_op_fuser: raise ValueError("fused_group_mlp requires use_transformer_engine_op_fuser.") @@ -2301,39 +2449,21 @@ def __post_init__(self): "fused_group_mlp offloads the whole fused grouped MLP and cannot be " f"combined with expert_fc1 or moe_act. Remove: {moe_partial_offload}" ) - - if self.gdp_cutedsl_kernel: - assert 0 <= self.gdp_num_chunk_states_to_recompute <= 64, ( - "gdp_num_chunk_states_to_recompute must be in range [0, 64], got " - f"{self.gdp_num_chunk_states_to_recompute}." - ) - - if self.gtp_remat_opt_in_modules: - _allowed_gtp_remat_opt_in_modules = {"moe_latent_proj"} - invalid = set(self.gtp_remat_opt_in_modules) - _allowed_gtp_remat_opt_in_modules - assert not invalid, ( - f"Invalid choices for gtp_remat_opt_in_modules: {invalid}. " - f"Allowed modules are: {_allowed_gtp_remat_opt_in_modules}" - ) - if self.moe_paged_stash: - if self.cpu_offloading: - raise ValueError("moe_paged_stash cannot be enabled with cpu_offloading.") - if self.moe_expert_rank_capacity_factor is None: - raise ValueError( - "moe_paged_stash requires moe_expert_rank_capacity_factor to be set; " - "there is no need to use paged stashing without it." - ) + assert not self.cpu_offloading, "moe_paged_stash cannot be enabled with cpu_offloading." + assert self.moe_expert_rank_capacity_factor is not None, ( + "moe_paged_stash requires moe_expert_rank_capacity_factor to be set; " + "there is no need to use paged stashing without it." + ) moe_offload_conflict = {"expert_fc1", "moe_act", "fused_group_mlp"} & set( self.offload_modules ) - if moe_offload_conflict: - raise ValueError( - "When moe_paged_stash is enabled, offload_modules must not include " - f"expert_fc1, moe_act, or fused_group_mlp " - f"(paged stash covers those activations). " - f"Remove: {moe_offload_conflict}" - ) + assert not moe_offload_conflict, ( + "When moe_paged_stash is enabled, offload_modules must not include " + f"expert_fc1, moe_act, or fused_group_mlp " + f"(paged stash covers those activations). " + f"Remove: {moe_offload_conflict}" + ) if ( self.num_layers_in_first_pipeline_stage is not None @@ -2577,32 +2707,18 @@ def __post_init__(self): if self.activation_func != F.silu or not self.gated_linear_unit: raise ValueError("Storing activation input in FP8 is supported only for SwiGLU.") - if ( - self.activation_func_clamp_value is not None - and self.activation_func == F.silu - and self.gated_linear_unit - ): - if ( - not math.isfinite(self.activation_func_clamp_value) - or self.activation_func_clamp_value <= 0 - ): - raise ValueError( - "activation_func_clamp_value for SwiGLU must be finite and greater than zero." - ) - if self.num_moe_experts is None: - raise ValueError( - "activation_func_clamp_value for SwiGLU is only supported with MoE." - ) - if self.glu_linear_offset != 0.0: - raise ValueError( - "glu_linear_offset must be zero when activation_func_clamp_value " - "is set for SwiGLU." - ) - if self.use_te_activation_func: - raise ValueError( - "use_te_activation_func must be False " - "when activation_func_clamp_value is not None for SwiGLU" - ) + if self.activation_func_clamp_value is not None: + # swiglu + if self.activation_func == F.silu and self.gated_linear_unit: + if self.num_moe_experts is None: + raise ValueError( + "activation_func_clamp_value for SwiGLU is only supported with MoE." + ) + if self.use_te_activation_func: + raise ValueError( + "use_te_activation_func must be False " + "when activation_func_clamp_value is not None for SwiGLU" + ) if self.apply_rope_fusion: if self.multi_latent_attention: @@ -2611,7 +2727,18 @@ def __post_init__(self): "It is experimental and may change in future versions." ) else: - if self.rotary_interleaved: + fused_mrope_available = False + # Triton fused mRoPE supports split-half RoPE only. Keep rotary_interleaved + # configs on the TE validation path so the TE >= 2.3 check still applies. + if self.mrope_section is not None and not self.rotary_interleaved: + try: + from megatron.core.fusions.fused_mrope import is_fused_mrope_available + + fused_mrope_available = is_fused_mrope_available() + except ImportError: + fused_mrope_available = False + + if self.rotary_interleaved and not fused_mrope_available: if not is_te_min_version("2.3.0"): raise ValueError( "rotary_interleaved does not work with apply_rope_fusion for " @@ -2623,9 +2750,14 @@ def __post_init__(self): fused_apply_rotary_pos_emb_thd, ) - if fused_apply_rotary_pos_emb is None and fused_apply_rotary_pos_emb_thd is None: + if ( + fused_apply_rotary_pos_emb is None + and fused_apply_rotary_pos_emb_thd is None + and not fused_mrope_available + ): raise ValueError( - "apply_rope_fusion is not available. Please install TE >= 1.4." + "apply_rope_fusion is not available. Please install TE >= 1.4 " + "or Triton for fused mRoPE." ) if self.fused_single_qkv_rope: @@ -2733,6 +2865,38 @@ def __post_init__(self): "'sqrtsoftplus', or unset --moe-router-enable-expert-bias." ) + if self.moe_n_hash_layers > 0: + assert ( + self.actual_vocab_size is not None + ), "actual_vocab_size must be set when moe_n_hash_layers > 0." + if self.pipeline_model_parallel_size > 1 and not self.is_hybrid_model: + assert self.pipeline_model_parallel_layout is not None, ( + "pipeline_model_parallel_layout must be set when using hash MoE " + "layers with pipeline parallelism (PP > 1)." + ) + # The embedding is always in layout[0][0] (PP rank 0, VPP rank 0). + # All hash MoE layers must be in the same virtual pipeline stage. + embedding_stage = self.pipeline_model_parallel_layout.layout[0][0] + n_decoders_with_embedding = embedding_stage.count(LayerType.decoder) + assert self.moe_n_hash_layers <= n_decoders_with_embedding, ( + f"Currently, All hash MoE layers must be in the same virtual pipeline stage " + f"as the embedding. The embedding stage has " + f"{n_decoders_with_embedding} decoder layers, but " + f"moe_n_hash_layers={self.moe_n_hash_layers}." + ) + assert ( + not self.overlap_moe_expert_parallel_comm + ), "overlap_moe_expert_parallel_comm does not support moe_n_hash_layers > 0 for now." + log_single_rank( + logger, + logging.WARNING, + f"Hash MoE layer initialized with placeholder round-robin tid2eid. " + f"For real training, you MUST either (a) load tid2eid from a " + f"pre-trained DSv4 checkpoint, or (b) provide a frequency-aware " + f"initialization (e.g., Sinkhorn-balanced over token frequency). " + f"Round-robin will cause severe expert imbalance.", + ) + if self.num_moe_experts and self.fp8: # TE version below 1.7.0 will raise Error when handle zeros tokens for expert if not is_te_min_version("1.7.0.dev0"): @@ -2904,6 +3068,119 @@ def _scope_to_str(s): self.cuda_graph_impl == "full_iteration" and self.cuda_graph_modules ), 'cuda_graph_modules must be empty when cuda_graph_impl="full_iteration".' + # mHC selective recompute couples with CUDA graphs only through the guarded + # attention-only Transformer Engine split. This gate must stay below the + # cuda_graph_modules normalization and the deprecated flag migration above: + # earlier placement would compare unnormalized string module forms and let + # enable_cuda_graph/external_cuda_graph bypass the gate entirely. + if use_mhc_recompute and self.cuda_graph_impl != "none": + if self.cuda_graph_impl == "local": + # Intentionally fail-closed even for inference-only local-graph + # configs that carry leftover training recompute args: mHC + # recompute is inert outside training, but silently accepting + # the combination would mask misconfigured training runs. + raise ValueError( + "mHC recompute is not supported with cuda_graph_impl='local': " + "eager mHC recompute and its per-microbatch checkpoint " + "registration need host execution between captured segments, " + "which the local per-layer implementation does not provide. Use " + "cuda_graph_impl='transformer_engine' with " + "cuda_graph_modules=['attn'], cuda_graph_impl='full_iteration' " + "with dropout disabled, or disable CUDA graphs." + ) + if self.cuda_graph_impl == "full_iteration": + # Full-iteration capture records the whole eager iteration — + # including mHC checkpoint registration, recompute kernels, and + # storage rebinding — into one graph, so replays re-execute the + # recompute at fixed addresses by construction (no partial-graph + # bridge involved). The one mechanical hazard is RNG-consuming + # ops inside a checkpointed region: the recompute-time RNG rewind + # cannot run under stream capture, so a captured recompute would + # replay a different dropout mask than its captured forward. + if self.hidden_dropout != 0.0 or self.attention_dropout != 0.0: + raise ValueError( + "mHC recompute with cuda_graph_impl='full_iteration' requires " + "hidden_dropout=0 and attention_dropout=0: RNG state cannot be " + "rewound inside CUDA graph capture, so a captured recompute " + "would replay a different dropout mask than its forward pass." + ) + elif list(self.cuda_graph_modules or []) != [CudaGraphModule.attn] or list( + self.recompute_modules + ) != ["mhc"]: + raise ValueError( + "mHC recompute with Transformer Engine CUDA Graphs currently supports " + "only the initial attention-only split: cuda_graph_modules=[attn] and " + "recompute_modules=[mhc]. The eager mHC producer must remain outside " + "the captured consumer." + ) + if ( + self.cuda_graph_impl == "transformer_engine" + and self.fine_grained_activation_offloading + ): + # HyperConnectionTransformerLayer._te_cuda_graph_capture replaces + # TransformerLayer's implementation rather than extending it, so it + # never plants the offload synchronization edges the full-layer + # capture plants -- backward_record() on the graph input and + # forward_record() after capture. _set_offload_modules plants those + # exactly for the attention-scope modules under an attn-scope graph, + # so without them the offload copies race the captured attention. + attn_scope_offload = {"qkv_linear", "core_attn", "attn_proj"} & set( + self.offload_modules or [] + ) + if attn_scope_offload: + raise ValueError( + f"mHC recompute with attention-only TE CUDA Graphs is incompatible " + f"with offload_modules {sorted(attn_scope_offload)}. The split " + f"capture path omits the offload stream synchronization the " + f"full-layer capture path performs, so the offload copies can race " + f"the captured attention. Remove {sorted(attn_scope_offload)} from " + f"offload_modules, or drop 'attn' from cuda_graph_modules." + ) + # The replay half of the same gap is fixed, not rejected: see + # _replay_mhc_attention_consumer. + if self.virtual_pipeline_model_parallel_size is not None: + # VPP is admitted only together with EP overlap. Interleaving used + # to diverge here (grad norms ~1e8 from the first iteration) on a + # caching-allocator use-after-free: mHC post-processing ran inside + # the communication-stream combine node, so the recompute subgraph + # was allocated on the compute stream and read from another, a + # window the allocator cannot track. The fix -- the post node owning + # a compute-stream schedule node -- lives in the overlap schedule, + # so the non-overlap VPP path has never carried it and stays + # unvalidated. StaticBufferLoader itself is VPP-safe, since only the + # pre_process chunk carries a data iterator. + if not self.overlap_moe_expert_parallel_comm: + raise ValueError( + "mHC recompute supports interleaved pipeline (VPP) " + "schedules only together with " + "overlap_moe_expert_parallel_comm: the non-overlap VPP " + "path is unvalidated." + ) + + cuda_graph_captures_attention = self.cuda_graph_impl == "full_iteration" or ( + self.cuda_graph_impl in ("local", "transformer_engine") + and (not self.cuda_graph_modules or CudaGraphModule.attn in self.cuda_graph_modules) + ) + + cp_layout_conversion_required = self.experimental_attention_variant == "gated_delta_net" + # TODO: Extend this predicate as GDN2/KDA are introduced, and for DSv4 when + # dsa_cp_balance_indexer is introduced; those paths will also require module-local THD CP + # layout conversion. + if ( + (self.context_parallel_size > 1 or self.dynamic_context_parallel) + and self.sequence_packing_scheduler is not None + and cuda_graph_captures_attention + and cp_layout_conversion_required + ): + raise ValueError( + "THD context parallel layout conversion is required for this model " + "configuration, but it is not supported by CUDA graph capture that includes " + "attention " + f"(experimental_attention_variant={self.experimental_attention_variant!r}, " + f"cuda_graph_impl={self.cuda_graph_impl!r}, " + f"cuda_graph_modules={self.cuda_graph_modules!r})." + ) + if self.cuda_graph_impl != "none": if self.cpu_offloading and self.cuda_graph_impl != "full_iteration": @@ -2923,6 +3200,12 @@ def _scope_to_str(s): if CudaGraphModule.moe_preprocess not in self.cuda_graph_modules: self.cuda_graph_modules.append(CudaGraphModule.moe_preprocess) + if self.cuda_graph_impl != "transformer_engine": + assert not self.cuda_graph_dynamic_microbatches, ( + "cuda_graph_dynamic_microbatches is only supported with " + "cuda_graph_impl=transformer_engine." + ) + assert ( CudaGraphModule.moe not in self.cuda_graph_modules or CudaGraphModule.moe_router not in self.cuda_graph_modules @@ -3014,27 +3297,6 @@ def _scope_to_str(s): "moe_input_jitter_eps is not supported with graphed moe recomputation." ) - if ( - self.gtp_weight_remat_size > 1 - and self.cuda_graph_impl == "local" - and (self.fp8 is not None or self.fp4 is not None) - and self.moe_shared_expert_intermediate_size is not None - and not self.moe_shared_expert_overlap - and ( - full_cudagraph - or CudaGraphModule.moe in self.cuda_graph_modules - or CudaGraphModule.moe_router in self.cuda_graph_modules - ) - ): - assert "shared_experts" not in self.recompute_modules, ( - "GTP + local CUDA graphs that capture shared_experts " - "(moe_router/moe scope) cannot recompute it under fp8/fp4: " - "te_checkpoint requires .backward(), but the local fwd-graph " - "warmup uses .grad(). Drop 'shared_experts' from " - "--recompute-modules (GTP-shard + offload instead), or use " - "--cuda-graph-impl full_iteration." - ) - if self.fine_grained_activation_offloading: offload_modules = set(self.offload_modules or []) if self.cuda_graph_impl == "local": @@ -3074,17 +3336,11 @@ def _scope_to_str(s): "fine-grained activation offloading." ) if self.cuda_graph_impl == "full_iteration": - assert ( - self.fine_grained_offloading_max_inflight_offloads is not None - and self.fine_grained_offloading_max_inflight_offloads >= 0 - ), ( - "fine_grained_offloading_max_inflight_offloads must be a non-negative " - "integer when using fine-grained activation offloading with " - "full-iteration CUDA graphs" + assert self.fine_grained_offloading_max_inflight_offloads is not None, ( + "fine_grained_offloading_max_inflight_offloads must be set when using " + "fine-grained activation offloading with full-iteration CUDA graphs " ) - # Only meaningful for MoE models; dense models never dispatch tokens, - # so the (unused) dispatcher default must not fail validation. if self.num_moe_experts is not None and self.moe_token_dispatcher_type in ["allgather"]: if self.variable_seq_lengths is True: raise ValueError( @@ -3129,15 +3385,60 @@ def _scope_to_str(s): 'flex', ], 'overlap_moe_expert_parallel_comm is supported with alltoall/flex token dispatcher' - assert ( - self.recompute_granularity != 'full' - ), 'disable full recomputation when enabling overlap_moe_expert_parallel_comm' - assert ( - self.recompute_method is None - ), 'disable recomputation method when enabling overlap_moe_expert_parallel_comm' - assert ( - self.recompute_num_layers is None - ), 'recompute_num_layers must be None when enabling overlap_moe_expert_parallel_comm' + if self.recompute_granularity == 'full': + # Full recompute runs per layer segment; recompute_method / + # recompute_num_layers control the decoder grouping and are validated by + # the shared checks above. MTP under 'block' and the FP8/FP4 block window + # intentionally differ from the non-overlap checkpoint path. See + # megatron/core/models/common/model_chunk_schedule_plan.py. + # The replay is hand-rolled rather than a checkpoint primitive, so it + # cannot shard the retained segment input across TP ranks. + assert not self.distribute_saved_activations, ( + 'overlap_moe_expert_parallel_comm full recompute does not support ' + 'distribute_saved_activations: the segment replay does not go through ' + 'a checkpoint primitive, so the retained input is not sharded.' + ) + # MTP is segmented under both methods here, so unlike + # MultiTokenPredictionLayer._checkpointed_forward there is no 'block' + # exemption. The 'uniform' restriction below is config parity with that + # function's own assert, not a constraint of this schedule - multi-layer + # decoder segments hand the mHC bridge over correctly. + if self.mtp_num_layers and self.recompute_method == 'uniform': + assert ( + self.recompute_num_layers == 1 + ), 'recompute_num_layers must be 1 for MTP recompute' + # With dropout the forward RNG stream is interleaved across microbatches + # in the 1F1B schedule and cannot yet be faithfully replayed. + assert self.attention_dropout == 0.0 and self.hidden_dropout == 0.0, ( + 'overlap_moe_expert_parallel_comm full recompute currently requires ' + 'attention_dropout==0 and hidden_dropout==0 (RNG replay across the ' + 'interleaved 1F1B schedule is not yet supported for nonzero dropout)' + ) + # Delayed-scaling FP8 accumulates amax into a persistent + # amax_history across microbatches. Both the initial no_grad forward + # and the backward-time recompute run under the delayed FP8 autocast, + # and this path does not enter TE's activation-recompute + # (recompute_phase) machinery that te_checkpoint uses to snapshot and + # restore FP8 metadata, so the recompute would update amax_history a + # second time and silently diverge. Current-scaling recipes + # (tensorwise / mxfp8 / blockwise) derive scales from the current + # tensor each pass, so the replayed forward re-derives identical + # scales (validated bitwise-identical to recompute-off); only delayed + # scaling is rejected here. + assert not (self.fp8 and self.fp8_recipe == Fp8Recipe.delayed), ( + 'overlap_moe_expert_parallel_comm full recompute is not yet supported ' + 'with delayed-scaling FP8 (fp8_recipe="delayed"): the backward-time ' + 'recompute would double-update the persistent amax_history. Use a ' + 'current-scaling recipe (tensorwise / mxfp8 / blockwise) instead.' + ) + else: + assert self.recompute_method is None, ( + 'disable recomputation method when enabling ' 'overlap_moe_expert_parallel_comm' + ) + assert self.recompute_num_layers is None, ( + 'recompute_num_layers must be None when enabling ' + 'overlap_moe_expert_parallel_comm' + ) assert ( "moe" not in self.recompute_modules ), 'disable moe in recompute_modules when enabling overlap_moe_expert_parallel_comm' @@ -3150,11 +3451,9 @@ def _scope_to_str(s): assert ( not self.moe_shared_expert_overlap ), 'disable moe_shared_expert_overlap when enabling overlap_moe_expert_parallel_comm' - assert self.mtp_num_layers in ( - None, - 0, - 1, - ), 'MTP supports at most one layer when enabling overlap_moe_expert_parallel_comm.' + assert ( + self.mtp_num_layers is None or self.mtp_num_layers == 1 + ), 'MTP layernum only supports 1 when enabling overlap_moe_expert_parallel_comm.' # NCCL EP (ncclep flex backend) mirrors hybridep's comm/compute overlap, but a few # configs are not yet safe under the 1F1B split and are gated here. @@ -3162,6 +3461,14 @@ def _scope_to_str(s): self.moe_token_dispatcher_type == 'flex' and self.moe_flex_dispatcher_backend == 'ncclep' ): + if not self.moe_ncclep_static_shape: + warnings.warn( + 'overlap_moe_expert_parallel_comm with ncclep and ' + 'moe_ncclep_static_shape=False: get_permuted_hidden_states_by_experts ' + 'does a device-to-host sync that serializes the 1F1B overlap (correct, ' + 'but loses the overlap benefit). Set moe_ncclep_static_shape=True for ' + 'the overlapped path (needs the fused op on sm100+).' + ) assert not ( self.fine_grained_activation_offloading and 'expert_fc1' in (self.offload_modules or []) @@ -3174,7 +3481,8 @@ def _scope_to_str(s): if self.cuda_graph_impl != "none": if self.cuda_graph_impl == "transformer_engine": assert ( - CudaGraphModule.moe not in self.cuda_graph_modules + self.cuda_graph_impl == "transformer_engine" + and CudaGraphModule.moe not in self.cuda_graph_modules and CudaGraphModule.mlp not in self.cuda_graph_modules ), ( 'CUDA graph scope on moe and mlp is not ' @@ -3246,6 +3554,16 @@ def _scope_to_str(s): "2.3.0.dev0+39c0e70" ), "Must have at least TE version 2.3 or higher to use symmetric memory all reduce" + if self.rotary_base_per_layer is not None: + if self.mtp_num_layers is not None: + total_num_layers = self.num_layers + self.mtp_num_layers + else: + total_num_layers = self.num_layers + assert len(self.rotary_base_per_layer) == total_num_layers, ( + f"rotary_base_per_layer length ({len(self.rotary_base_per_layer)}) " + f"must equal num_layers ({total_num_layers})" + ) + if self.no_rope_freq: assert not self.flash_decode, "flash_decode cannot be used with no_rope." if isinstance(self.no_rope_freq, int): @@ -3263,6 +3581,25 @@ def _scope_to_str(s): f"the number of layers ({self.num_layers})" ) + if self.fallback_to_eager_attn: + assert self.transformer_impl == "transformer_engine", ( + f"fallback_to_eager_attn is only available with transformer_engine implementation," + f" but got {self.transformer_impl=}." + ) + + if self.fallback_to_eager_attn or self.transformer_impl == "local": + if self.context_parallel_size > 1 and self.cp_comm_type is not None: + all_cp_comm_types_are_all_gather = ( + all(item == "all_gather" for item in self.cp_comm_type) + if isinstance(self.cp_comm_type, list) + else self.cp_comm_type == "all_gather" + ) + if not all_cp_comm_types_are_all_gather: + raise ValueError( + f"fallback_to_eager_attn only supports all_gather communication type " + f"for context parallelism, but got {self.cp_comm_type=} instead." + ) + if self.transformer_impl == "inference_optimized": assert self.normalization == "RMSNorm" assert not self.layernorm_zero_centered_gamma @@ -3270,23 +3607,6 @@ def _scope_to_str(s): assert not self.add_qkv_bias assert not self.use_kitchen - if self.experimental_attention_variant == "dsa": - assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention" - if self.context_parallel_size > 1: - cp_comm_types = ( - self.cp_comm_type - if isinstance(self.cp_comm_type, list) - else [self.cp_comm_type] - ) - assert all( - cp_comm_type is not None - and cp_comm_type.replace("_", "").lower() == "allgather" - for cp_comm_type in cp_comm_types - ), ( - "DSAttention context parallelism currently supports " - "cp_comm_type=allgather only." - ) - if self.inference_fuse_tp_communication: assert self.transformer_impl == "inference_optimized", ( "inference_fuse_tp_communication is only supported " @@ -3302,91 +3622,43 @@ def _scope_to_str(s): "for inference_optimized transformer implementation." ) - if self.flash_attention_version is not None: - assert self.flash_attention_version in (2, 3, 4), ( - "flash_attention_version must be one of 2, 3, or 4, got " - f"{self.flash_attention_version}" - ) - if self.batch_invariant_mode: - from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( - _BATCH_INVARIANT_BACKENDS, - ) - - # argparse validates via the Literal annotation; guard here too so - # programmatic TransformerConfig construction fails at build time - # rather than inside enable_batch_invariant_mode() after model init. - assert self.batch_invariant_backend in _BATCH_INVARIANT_BACKENDS, ( - f"Unknown batch_invariant_backend {self.batch_invariant_backend!r}; " - f"expected one of {_BATCH_INVARIANT_BACKENDS}." - ) - assert self.params_dtype == torch.bfloat16, ( - "Batch invariant mode supports BF16 model parameters only; " - f"got {self.params_dtype}." - ) assert ( self.attention_backend == AttnBackend.flash - ), "Batch invariant mode only supports FlashAttention (--attention-backend flash)" - # The training (TransformerEngine) and inference attention paths must run - # the same FlashAttention kernel, so the version cannot be left to each - # path's autodetection. FlashAttention-2 is excluded because it does not - # expose the fixed num_splits schedule the batch-invariant kernels require. - assert self.flash_attention_version in (3, 4), ( - "Batch invariant mode requires --flash-attention-version 3 or 4 so the " - "training and inference attention paths run the same batch-invariant " - f"FlashAttention kernel (got {self.flash_attention_version})." - ) - # Context parallelism routes through TE's FA2 fwd/bwd kernels directly, which - # cannot be pinned to another version; dropout is not batch-invariant. + ), "Batch invariant mode only supports FlashAttention" + + if self.cuda_graph_impl != "none" and ( + self.sequence_packing_scheduler is not None or self.dynamic_context_parallel + ): + if self.thd_max_packed_sequences is None: + raise ValueError("THD CUDA Graph requires --thd-max-packed-sequences to be set.") assert ( - self.context_parallel_size == 1 - ), "Batch invariant mode does not support context parallelism" + self.pad_packed_seq_alignment is not None + ), "THD CUDA Graph requires --pad-packed-seq-alignment to be set." assert ( - self.attention_dropout == 0.0 - ), "Batch invariant mode does not support attention dropout" - if (self.num_moe_experts or 0) > 0: - from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( - HAVE_DEEPGEMM_BF16, - ) + self.pad_packed_seq_alignment == "max" + or self.pad_packed_seq_alignment == self.max_seqlen_per_dp_cp_rank + ), ( + "THD CUDA Graph requires --pad-packed-seq-alignment='max' " + "or --pad-packed-seq-alignment equal to max_seqlen_per_dp_cp_rank " + f"({self.max_seqlen_per_dp_cp_rank}), got {self.pad_packed_seq_alignment}." + ) - if self.transformer_impl != "inference_optimized": - assert self.moe_token_dispatcher_type == "alltoall", ( - "Batch-invariant MoE training requires " - "moe_token_dispatcher_type='alltoall'." - ) - # DeepGEMM is used by the "deepgemm"/"triton" backends, and by - # the torch inference grouped-GEMM path under any backend. The - # "te_native" backend with the vLLM inference backend (or the - # training path, where TE grouped GEMM stays native) does not - # need it. - needs_deepgemm = self.batch_invariant_backend in ("deepgemm", "triton") or ( - self.transformer_impl == "inference_optimized" - and self.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.TORCH - ) - assert not needs_deepgemm or HAVE_DEEPGEMM_BF16, ( - "batch_invariant_mode=True with MoE requires DeepGEMM with bf16 " - "grouped-GEMM bindings (m_grouped_bf16_gemm_nt_contiguous) for " - "this backend combination. " - "Install via `uv pip install -e .[batch_invariant]`." - ) - assert not ( - self.fp8 or self.fp4 - ), "Batch-invariant MoE is bf16-only. Disable fp8/fp4 to use it." - assert not (self.moe_permute_fusion or self.moe_permute_fusion_into_hybridep), ( - "Batch-invariant MoE requires the unfused permute/unpermute path so " - "top-k reductions use the fixed batch-invariant add tree." - ) - assert not ( - self.moe_pad_expert_input_to_capacity - or self.moe_pad_experts_for_cuda_graph_inference - ), ( - "Batch-invariant MoE supports dynamic dropless routing only. " - "Disable MoE capacity/expert padding." - ) + # 'extend_last' THD tail padding with context parallelism requires the + # global metadata to be extended before CP slicing, which only the + # sequence-packing scheduler path performs. Restrict this combination. + if ( + self.sequence_packing_scheduler is None + and self.pad_packed_seq_alignment is not None + and (self.thd_tail_padding_policy or 'append_dummy_seq') == 'extend_last' + and (self.context_parallel_size > 1 or self.dynamic_context_parallel) + ): + raise ValueError( + "thd_tail_padding_policy='extend_last' with context parallelism is only " + "supported together with a sequence_packing_scheduler. Either enable a " + "sequence_packing_scheduler, or use thd_tail_padding_policy='append_dummy_seq'." + ) - # Scheduler-value, max-seqlen, and variable_seq_lengths handling live in - # ModelParallelConfig.__post_init__ next to the field definitions; only the - # transformer-stack requirements are validated here. if self.sequence_packing_scheduler is not None: # Check TE version. if not HAVE_PACKAGING: @@ -3402,19 +3674,28 @@ def _scope_to_str(s): f"but got {get_te_version()} (TE < 2.9.0 may have convergence issues)." ) - # TODO(tailaim): add support for other dispatcher types - # Only relevant for MoE models; dense models never dispatch tokens, - # so the (unused) dispatcher default must not fail validation. For - # allgather specifically, the general variable_seq_lengths check - # above raises first (packing derives variable_seq_lengths=True). + # Needed for passing variable sequences between pp stages. + self.variable_seq_lengths = True + if self.num_moe_experts is not None: - assert self.moe_token_dispatcher_type == "alltoall", ( - f"sequence_packing only supports moe_token_dispatcher_type='alltoall', " - f"got '{self.moe_token_dispatcher_type}'" + assert self.moe_token_dispatcher_type in ("alltoall", "flex"), ( + f"sequence_packing only supports moe_token_dispatcher_type in " + f"('alltoall', 'flex'), got '{self.moe_token_dispatcher_type}'" + ) + + supported_schedulers = ['dp_balanced', 'default_dynamic_cp'] + if ( + self.sequence_packing_scheduler is not None + and self.sequence_packing_scheduler not in supported_schedulers + ): + raise ValueError( + f"Unsupported scheduler: {self.sequence_packing_scheduler}. " + f"Available schedulers: {supported_schedulers}" ) @dataclass +@experimental_api class MLATransformerConfig(TransformerConfig): """Configuration object for megatron-core Multi-Latent Attention (MLA) transformers. @@ -3429,10 +3710,12 @@ class MLATransformerConfig(TransformerConfig): """Rank of Query tensor's low rank representation.""" kv_lora_rank: int = 512 - """Rank of Key and Value tensors' low rank representation.""" + """Rank of Key and Value tensors' low rank representation. + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_head_dim: int = 128 - """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim""" + """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_pos_emb_head_dim: int = 64 """Dimension of the position embedding in the QK projection.""" @@ -3470,6 +3753,12 @@ class MLATransformerConfig(TransformerConfig): mscale_all_dim: float = 0.0 """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" + o_groups: int = 8 + """Number of groups for grouped low-rank output projection (wo_a).""" + + o_lora_rank: int = 1024 + """Low-rank dimension per group for grouped output (wo_a). Used when o_groups > 0.""" + cache_mla_latents: bool = False """Cache the low dimensional tensors for MLA rather than full KV cache. This is only for the dynamic inference backend and requires that @@ -3482,12 +3771,32 @@ class MLATransformerConfig(TransformerConfig): def __post_init__(self): super().__post_init__() - if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": + if ( + self.multi_latent_attention + and self.apply_rope_fusion + and self.rope_type != "yarn" + and self.experimental_attention_variant != "dsv4_hybrid" + ): raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.") if self.attention_output_gate: raise NotImplementedError("Output gate is not supported for MLA yet.") + # DSv4 hybrid: derive qk_head_dim and kv_lora_rank from v_head_dim and qk_pos_emb_head_dim + if self.experimental_attention_variant == "dsv4_hybrid": + assert ( + not self.mla_down_proj_fusion + ), "MLA down projection fusion must be disabled for DSv4 hybrid mode." + log_single_rank( + logger, + logging.WARNING, + f"DSv4 hybrid mode is enabled, deriving qk_head_dim and kv_lora_rank from " + f"v_head_dim and qk_pos_emb_head_dim", + ) + derived = self.v_head_dim - self.qk_pos_emb_head_dim + self.qk_head_dim = derived + self.kv_lora_rank = derived + if self.cache_mla_latents: assert ( self.apply_rope_fusion is False 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..ec53ca490d6 --- /dev/null +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py @@ -0,0 +1,830 @@ +# 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": "gated_delta_net", + "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='gated_delta_net'"): + _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_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, + linear_attention_freq=None, + recompute_granularity="selective", + recompute_modules=["gdn_norm_out"], + ) + + assert config.recompute_modules == ["gdn_norm_out"] + + +def test_gdn_conv_pad_alignment_rejects_chunkwise_cp(): + with pytest.raises(AssertionError, match="gdn_conv_pad_alignment is incompatible"): + _make_gdn_config( + context_parallel_size=2, linear_cp_mode="chunkwise", gdn_conv_pad_alignment=4096 + ) + + +def test_gdn_chunkwise_cp_head_divisibility_ignores_cp_size(): + config = _make_gdn_config( + tensor_model_parallel_size=2, + context_parallel_size=4, + linear_cp_mode="chunkwise", + linear_num_key_heads=4, + linear_num_value_heads=8, + ) + assert config.linear_cp_mode == "chunkwise" + + +def test_torch_chunk_gated_delta_rule_preserves_public_signature(): + signature = inspect.signature(torch_chunk_gated_delta_rule) + assert tuple(signature.parameters) == ( + "q", + "k", + "v", + "g", + "beta", + "chunk_size", + "initial_state", + "output_final_state", + "use_qk_l2norm_in_kernel", + "cu_seqlens", + "cp_context", + "scale", + ) + + +def test_gdn_headwise_cp_head_divisibility_includes_cp_size(): + with pytest.raises(AssertionError, match="linear_head_parallel_size"): + _make_gdn_config( + tensor_model_parallel_size=2, + context_parallel_size=4, + linear_cp_mode="headwise", + linear_num_key_heads=4, + linear_num_value_heads=8, + ) + + +@pytest.mark.parametrize( + ("tp_size", "sp", "cp_size", "linear_cp_mode"), + [ + # cp_size=1: the CP path is inactive, so linear_cp_mode choice is irrelevant. + # Cover the "chunkwise" default and skip the "headwise" variants for brevity. + (1, False, 1, None), + (2, False, 1, None), + (2, True, 1, None), + # cp_size=2: exercise both CP paths. + (1, False, 2, "headwise"), + (2, False, 2, "headwise"), + (2, True, 2, "headwise"), + (1, False, 2, "chunkwise"), + (2, False, 2, "chunkwise"), + (2, True, 2, "chunkwise"), + ], +) +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +class TestGatedDeltaNet: + + @pytest.fixture(scope='function', autouse=True) + def setup_method(self, tp_size, sp, cp_size, linear_cp_mode): + # Initialize parallel and random seed + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp_size, + pipeline_model_parallel_size=1, + context_parallel_size=cp_size, + ) + model_parallel_cuda_manual_seed(123) + self.tp_size = tp_size + self.cp_size = cp_size + self.sp_size = tp_size if sp else 1 + self.linear_cp_mode = linear_cp_mode + if self.linear_cp_mode == "headwise": + self.cp_size_chunkwise = 1 + self.cp_size_headwise = self.cp_size + elif self.linear_cp_mode == "chunkwise": + self.cp_size_chunkwise = self.cp_size + self.cp_size_headwise = 1 + elif self.cp_size == 1: + self.cp_size_chunkwise = 1 + self.cp_size_headwise = 1 + else: + raise ValueError(f"Invalid linear CP mode: {self.linear_cp_mode}") + + # Get TP and CP process groups from device mesh + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + tp_cp_group = parallel_state.get_tensor_and_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group, tp_cp=tp_cp_group) + + # Initialize model, with the same config as Qwen Next except `num_layers` + self.transformer_config = TransformerConfig( + hidden_size=2048, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + num_layers=1, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + num_attention_heads=16, + num_query_groups=2, + activation_func=F.silu, + bf16=True, + tensor_model_parallel_size=tp_size, + sequence_parallel=sp, + context_parallel_size=cp_size, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + linear_cp_mode=self.linear_cp_mode, + transformer_impl="transformer_engine", + ) + gdn_submodules = get_experimental_attention_variant_module_spec( + config=self.transformer_config + ).submodules + + self.gdn = GatedDeltaNet( + self.transformer_config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + self.gdn = self.gdn.cuda().bfloat16() + + def teardown_method(self): + Utils.destroy_model_parallel() + + def test_gpu_forward(self): + gdn = self.gdn + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + attention_mask = None + + output, bias = gdn(hidden_states, attention_mask) + + assert output.dim() == 3, f"Output too many dimensions ({output.shape=})" + assert output.shape[0] == seq_length // self.sp_size // self.cp_size, ( + f"Output shape {output.shape[0]=} mismatch with " + f" {seq_length=} // {self.sp_size=} // {self.cp_size=}." + ) + assert ( + output.shape[1] == micro_batch_size + ), f"Output shape {output.shape[1]=} mismatch with {micro_batch_size=}" + assert ( + output.shape[2] == gdn.config.hidden_size + ), f"Output shape {output.shape[2]=} mismatch with {gdn.config.hidden_size=}" + assert ( + output.dtype == hidden_states.dtype + ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + + @pytest.mark.flaky_in_dev # Issue #5473 + def test_selective_recompute_gdn(self): + """Whole-module 'gdn' recompute must match the non-recompute forward and gradients. + + The same module/input is run twice (recompute off, then on); the forward output and + all parameter / input gradients must agree within a tight tolerance (rtol/atol=1e-4). + The recompute path is run-to-run deterministic on these kernels (empirically bitwise), + so a tolerance well below the bf16 floor is expected to hold. + """ + gdn = self.gdn + gdn.train() + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + torch.manual_seed(1234) + base_input = torch.randn( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + def run(recompute): + gdn.recompute_gdn = recompute + gdn.zero_grad(set_to_none=True) + hidden_states = base_input.clone().detach().requires_grad_(True) + output, _ = gdn(hidden_states, None) + output.float().square().mean().backward() + param_grads = { + name: param.grad.detach().clone() + for name, param in gdn.named_parameters() + if param.grad is not None + } + return output.detach().clone(), hidden_states.grad.detach().clone(), param_grads + + try: + out_ref, dinput_ref, pgrad_ref = run(recompute=False) + out_rc, dinput_rc, pgrad_rc = run(recompute=True) + finally: + gdn.recompute_gdn = False + + torch.testing.assert_close(out_rc, out_ref, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(dinput_rc, dinput_ref, rtol=1e-4, atol=1e-4) + assert pgrad_ref.keys() == pgrad_rc.keys(), "recompute changed the set of grad params" + assert len(pgrad_ref) > 0, "expected at least one parameter gradient" + for name in pgrad_ref: + torch.testing.assert_close( + pgrad_rc[name], + pgrad_ref[name], + rtol=1e-4, + atol=1e-4, + msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}", + ) + + @pytest.mark.flaky_in_dev # Issue #5473 + def test_selective_recompute_gdn_norm_out(self): + """Output-discarding 'gdn_norm_out' recompute must preserve outputs and gradients.""" + gdn = self.gdn + gdn.train() + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + torch.manual_seed(1234) + base_input = torch.randn( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + def run(recompute_norm_out): + gdn.recompute_gdn = False + gdn.recompute_norm_out = recompute_norm_out + gdn.norm_out_checkpoint = None + gdn.zero_grad(set_to_none=True) + hidden_states = base_input.clone().detach().requires_grad_(True) + output, _ = gdn(hidden_states, None) + output.float().square().mean().backward() + param_grads = { + name: param.grad.detach().clone() + for name, param in gdn.named_parameters() + if param.grad is not None + } + return output.detach().clone(), hidden_states.grad.detach().clone(), param_grads + + try: + out_ref, dinput_ref, pgrad_ref = run(recompute_norm_out=False) + out_rc, dinput_rc, pgrad_rc = run(recompute_norm_out=True) + finally: + gdn.recompute_norm_out = False + gdn.norm_out_checkpoint = None + + torch.testing.assert_close(out_rc, out_ref, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(dinput_rc, dinput_ref, rtol=1e-4, atol=1e-4) + assert pgrad_ref.keys() == pgrad_rc.keys(), "recompute changed the set of grad params" + assert len(pgrad_ref) > 0, "expected at least one parameter gradient" + for name in pgrad_ref: + torch.testing.assert_close( + pgrad_rc[name], + pgrad_ref[name], + rtol=1e-4, + atol=1e-4, + msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}", + ) + + def test_gpu_forward_rejects_sbhd_chunkwise_cp_batch_gt_one(self): + if not (self.linear_cp_mode == "chunkwise" and self.cp_size > 1): + pytest.skip("Only chunkwise CP with CP>1 uses the FLA CP batch guard.") + + gdn = self.gdn + + micro_batch_size = 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + with pytest.raises(ValueError, match="requires micro_batch_size == 1"): + gdn(hidden_states, None) + + def test_gpu_forward_rejects_sbhd_conv_padding(self): + gdn = self.gdn + gdn.config.gdn_conv_pad_alignment = 4096 + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + expected_error = ( + "incompatible with GDN chunkwise CP" + if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 + else "only supported with packed sequence" + ) + with pytest.raises(ValueError, match=expected_error): + gdn(hidden_states, None) + + def test_deterministic_mode(self): + if self.cp_size > 1: + pytest.skip( + "deterministic_mode uses torch_chunk_gated_delta_rule, which does not support CP." + ) + + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + det_config = copy.deepcopy(self.transformer_config) + det_config.deterministic_mode = True + + gdn_submodules = get_experimental_attention_variant_module_spec( + config=det_config + ).submodules + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + gdn = ( + GatedDeltaNet( + det_config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + .cuda() + .bfloat16() + ) + + # deterministic_mode must select the torch-native kernel, not FLA. + assert gdn.gated_delta_rule is torch_chunk_gated_delta_rule + + micro_batch_size = 2 + seq_length = 64 + torch.manual_seed(0) + base_input = torch.randn( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + def run(): + hidden_states = base_input.clone().requires_grad_(True) + output, _ = gdn(hidden_states, None) + output.float().sum().backward() + grads = { + name: param.grad.detach().clone() + for name, param in gdn.named_parameters() + if param.grad is not None + } + gdn.zero_grad(set_to_none=True) + return output.detach().clone(), grads, hidden_states.grad.detach().clone() + + out1, grads1, input_grad1 = run() + out2, grads2, input_grad2 = run() + + rank = torch.distributed.get_rank() + assert torch.equal(out1, out2), f"Output not reproducible ({rank=})" + assert torch.equal(input_grad1, input_grad2), f"Input grad not reproducible ({rank=})" + assert set(grads1.keys()) == set(grads2.keys()) + for name in grads1: + assert torch.equal( + grads1[name], grads2[name] + ), f"Grad not reproducible for {name} ({rank=})" + + def test_module_construction(self): + gdn = self.gdn + assert gdn.in_proj_dim == 2 * gdn.qk_dim + 2 * gdn.v_dim + 2 * gdn.num_value_heads + assert gdn.A_log.shape == (gdn.num_value_heads // self.tp_size,) + assert gdn.dt_bias.shape == (gdn.num_value_heads // self.tp_size,) + + def test_sharded_state_dict_splits_gdn_parameters(self): + sharded_sd = self.gdn.sharded_state_dict(prefix="gdn.") + + in_proj_weight = sharded_sd["gdn.in_proj.weight"] + conv1d_weight = sharded_sd["gdn.conv1d.weight"] + assert isinstance(in_proj_weight, ShardedTensorFactory) + assert isinstance(conv1d_weight, ShardedTensorFactory) + + in_proj_chunks = in_proj_weight.build() + assert tuple(chunk.key for chunk in in_proj_chunks) == tuple( + f"gdn.in_proj.weight.{name}" for name in self.gdn.in_proj_split_names + ) + assert sum(chunk.data.numel() for chunk in in_proj_chunks) == in_proj_weight.data.numel() + + conv1d_chunks = conv1d_weight.build() + assert tuple(chunk.key for chunk in conv1d_chunks) == ( + "gdn.conv1d.weight.query", + "gdn.conv1d.weight.key", + "gdn.conv1d.weight.value", + ) + assert sum(chunk.data.numel() for chunk in conv1d_chunks) == conv1d_weight.data.numel() + + def test_jit_compiled_helpers(self): + import torch._dynamo + + gdn = self.gdn + batch = 2 + seq_len = 16 + + device = torch.cuda.current_device() + num_v_heads_local = gdn.num_value_heads // gdn.tp_size // self.cp_size_headwise + qk_dim_local = gdn.qk_dim_local_tp // self.cp_size_headwise + v_dim_local = gdn.v_dim_local_tp // self.cp_size_headwise + + qkv = torch.randn( + batch, seq_len, 2 * qk_dim_local + v_dim_local, device=device, dtype=torch.bfloat16 + ) + gate = torch.randn( + batch, + seq_len, + num_v_heads_local, + gdn.value_head_dim, + device=device, + dtype=torch.bfloat16, + ) + gate_feats = ( + torch.randn(batch, seq_len, num_v_heads_local, device=device, dtype=torch.bfloat16), + torch.randn(batch, seq_len, num_v_heads_local, device=device, dtype=torch.bfloat16), + ) # beta, alpha + + # Disable dynamo so coverage.py can trace through the method bodies, + # which are normally wrapped by @jit_fuser (torch.compile). + A_log_mock = torch.randn(num_v_heads_local, device=device, dtype=torch.bfloat16) + dt_bias_mock = torch.randn(num_v_heads_local, device=device, dtype=torch.bfloat16) + + with torch._dynamo.config.patch(disable=True): + kernel_inputs = gdn._prepare_input_for_gated_delta_rule( + qkv, + gate, + A_log_mock, + dt_bias_mock, + batch, + seq_len, + *gate_feats, + cp_size_headwise=self.cp_size_headwise, + ) + + query = kernel_inputs["q"] + key = kernel_inputs["k"] + value = kernel_inputs["v"] + g = kernel_inputs["g"] + gate_out = kernel_inputs["gate"] + beta_out = kernel_inputs["beta"] + + assert query.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) + assert key.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) + assert value.shape == (batch, seq_len, num_v_heads_local, gdn.value_head_dim) + for t in (query, key, value, gate_out, beta_out): + assert t.is_contiguous() + + assert g.dtype == torch.float32 + assert g.shape == (batch, seq_len, num_v_heads_local) + assert beta_out.shape == (batch, seq_len, num_v_heads_local) + + def test_fused_pre_gated_delta_rule_headwise_cp_uses_cp_local_parameters(self): + if not HAVE_FUSED_PRE_GDR: + pytest.skip("causal-conv1d fused backward is not installed.") + if not (self.linear_cp_mode == "headwise" and self.cp_size > 1): + pytest.skip("Only headwise CP with CP>1 needs CP-local fused pre-GDR params.") + + gdn = self.gdn + batch = 2 + seq_len = 16 + qk_channels = gdn.qk_dim_local_tp // self.cp_size_headwise + v_channels = gdn.v_dim_local_tp // self.cp_size_headwise + num_key_heads = qk_channels // gdn.key_head_dim + num_value_heads = v_channels // gdn.value_head_dim + qkvzba_dim = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + qkvzba = torch.randn( + seq_len, batch, qkvzba_dim, device=torch.cuda.current_device(), dtype=torch.bfloat16 + ) + captured = {} + + def fake_fused_streamed_pre_gated_delta_rule( + qkvzba_arg, + conv1d_weight, + conv1d_bias, + A_log, + dt_bias, + *, + num_key_heads, + num_value_heads, + **kwargs, + ): + captured.update( + { + "qkvzba": qkvzba_arg, + "conv1d_weight": conv1d_weight, + "conv1d_bias": conv1d_bias, + "A_log": A_log, + "dt_bias": dt_bias, + "num_key_heads": num_key_heads, + "num_value_heads": num_value_heads, + "cp_group": kwargs["cp_group"], + } + ) + return tuple(torch.empty(0, device=qkvzba_arg.device) for _ in range(6)) + + with mock.patch( + "megatron.core.fusions.fused_pre_gated_delta_rule." + "fused_streamed_pre_gated_delta_rule", + side_effect=fake_fused_streamed_pre_gated_delta_rule, + ): + gdn._fused_streamed_pre_gated_delta_rule(qkvzba, cp_group_headwise=gdn.pg_collection.cp) + + assert captured["qkvzba"] is qkvzba + assert captured["conv1d_weight"].shape == ( + 2 * qk_channels + v_channels, + 1, + gdn.conv_kernel_dim, + ) + assert captured["conv1d_bias"] is None + assert captured["A_log"].shape == (num_value_heads,) + assert captured["dt_bias"].shape == (num_value_heads,) + assert captured["num_key_heads"] == num_key_heads + assert captured["num_value_heads"] == num_value_heads + assert captured["cp_group"] is None + + def test_gpu_forward_thd_correctness(self): + if self.sp_size > 1: + pytest.skip("Sequence parallel is not supported for this test case.") + if self.cp_size > 1 and self.linear_cp_mode == "chunkwise": + pytest.skip("Chunkwise CP is not supported for this test case.") + + atol, rtol = 3e-4, 3e-4 + + # Input shape + sequence_length = 32 + micro_batch_size = 4 + cu_seqlens = [0, 32, 64, 96, 128] + # sbhd input shape: [sequence length, batch size, hidden size] + sub_sequence_length = sequence_length // self.cp_size + hidden_states_sbhd = torch.rand( + (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size) + ) + attention_mask_sbhd = None + hidden_states_sbhd = hidden_states_sbhd.cuda().bfloat16() + # thd input shape: [sequence length * batch size, 1, hidden size] + hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() + hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) + attention_mask_thd = None + packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) + _set_gdn_test_cp_partition_mode(packed_seq_params, self.cp_size, self.linear_cp_mode) + + # THD format + output_thd, _ = self.gdn( + hidden_states_thd, attention_mask_thd, packed_seq_params=packed_seq_params + ) + # SBHD format + output_sbhd, _ = self.gdn(hidden_states_sbhd, attention_mask_sbhd) + output_sbhd_T = output_sbhd.transpose(0, 1).contiguous().view(*output_thd.shape) + + rank = torch.distributed.get_rank() + assert output_thd.shape[0] == sub_sequence_length * micro_batch_size + assert output_thd.shape[1] == 1 + assert output_thd.shape[2] == self.gdn.config.hidden_size + torch.testing.assert_close( + output_sbhd_T, + output_thd, + atol=atol, + rtol=rtol, + msg=lambda msg: f"Output mismatch ({rank=}): {msg}", + ) + + def test_gpu_forward_thd_padding_correctness(self): + if self.sp_size > 1: + pytest.skip("Sequence parallel is not supported for this test case.") + if self.cp_size > 1 and self.linear_cp_mode == "chunkwise": + pytest.skip("Chunkwise CP is not supported for this test case.") + + atol, rtol = 3e-4, 3e-4 + sequence_length = 32 + micro_batch_size = 4 + + # sbhd input shape: [sequence length, batch size, hidden size] + sub_sequence_length = sequence_length // self.cp_size + hidden_states_sbhd = torch.rand( + (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + output_sbhd, _ = self.gdn(hidden_states_sbhd, None) + + # thd input shape: [sequence length * batch size, 1, hidden size] + hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() + hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) + output_bshd = output_sbhd.transpose(0, 1).contiguous() + + rank = torch.distributed.get_rank() + + # A) padded branch: prefer *_padded when available. + padded_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 128] + ) + _set_gdn_test_cp_partition_mode(padded_params, self.cp_size, self.linear_cp_mode) + output_thd_padded, _ = self.gdn(hidden_states_thd, None, packed_seq_params=padded_params) + output_thd2bshd = output_thd_padded.view(*output_bshd.shape) + torch.testing.assert_close( + output_bshd[:, :30, :], + output_thd2bshd[:, :30, :], + atol=atol, + rtol=rtol, + msg=lambda msg: f"THD padded output mismatch ({rank=}): {msg}", + ) + + # B) no-padded branch: use actual cu_seqlens when it matches total_sequence_length. + no_padding_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 128]) + _set_gdn_test_cp_partition_mode(no_padding_params, self.cp_size, self.linear_cp_mode) + output_thd_no_padding, _ = self.gdn( + hidden_states_thd, None, packed_seq_params=no_padding_params + ) + assert output_thd_no_padding.shape == output_thd_padded.shape + + # C) explicit causal-conv padding is only applied to packed inputs and + # should not affect the original unpadded token outputs. + self.gdn.config.gdn_conv_pad_alignment = 48 + output_thd_conv_pad, _ = self.gdn( + hidden_states_thd, None, packed_seq_params=no_padding_params + ) + self.gdn.config.gdn_conv_pad_alignment = None + assert output_thd_conv_pad.shape == output_thd_no_padding.shape + torch.testing.assert_close( + output_thd_conv_pad, + output_thd_no_padding, + atol=atol, + rtol=rtol, + msg=lambda msg: f"THD conv-padded output mismatch ({rank=}): {msg}", + ) + + # D) padded mismatch branch: if *_padded[-1] mismatches total_sequence_length, should raise. + padded_mismatch_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 126] + ) + _set_gdn_test_cp_partition_mode(padded_mismatch_params, self.cp_size, self.linear_cp_mode) + with pytest.raises(ValueError, match="does not match"): + self.gdn(hidden_states_thd, None, packed_seq_params=padded_mismatch_params) + + # E) actual mismatch branch without *_padded: should raise. + actual_mismatch_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 129]) + _set_gdn_test_cp_partition_mode(actual_mismatch_params, self.cp_size, self.linear_cp_mode) + with pytest.raises(ValueError, match="does not match"): + self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) + + +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +class TestGDNCuSeqlensResolve: + + @pytest.fixture + def mock_gdn(self): + class MockGDN: + _resolve_cu_seqlens = GatedDeltaNet._resolve_cu_seqlens + + return MockGDN() + + def test_padded_preferred_when_available(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + padded = torch.tensor([0, 504, 1008], dtype=torch.int32) + result = mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q", cp_size=2) + assert torch.equal(result, padded) + + def test_actual_used_when_no_padding(self, mock_gdn): + actual = torch.tensor([0, 504, 1008], dtype=torch.int32) + result = mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + assert torch.equal(result, actual) + + def test_raises_when_padding_mismatch(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_raises_when_padded_mismatches_total(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + padded = torch.tensor([0, 504, 1004], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_raises_when_not_divisible_by_cp_size(self, mock_gdn): + actual = torch.tensor([0, 505, 1008], dtype=torch.int32) + with pytest.raises(ValueError, match="must be divisible by cp_size"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_cp1_still_validates_total(self, mock_gdn): + mock_gdn.cp_size = 1 + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=1) From da9d99067f257832329d842cde5bec8a98a904c6 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Tue, 18 Aug 2026 00:56:15 +0800 Subject: [PATCH 02/34] Add Kimi Delta Attention to hybrid models Add KDA as the channel-wise Gated DeltaNet variant in HybridModel. Reuse the existing private GDN-family base and gdn_norm_out recompute selector while keeping KDA projection, kernel, and context-parallel behavior variant-specific. Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 16f448cf5280f7514b05590906c5d56198ce1f07) --- .../inference/contexts/dynamic_context.py | 1716 +---- megatron/core/models/hybrid/hybrid_block.py | 10 + .../models/hybrid/hybrid_layer_allocation.py | 33 +- .../core/models/hybrid/hybrid_layer_specs.py | 193 +- megatron/core/ssm/gated_delta_net/__init__.py | 24 +- megatron/core/ssm/gated_delta_net/kda.py | 445 ++ .../core/transformer/transformer_config.py | 6 + megatron/training/arguments.py | 5796 +++++++++++------ .../contexts/test_dynamic_context.py | 16 + tests/unit_tests/ssm/test_hybrid_block.py | 253 +- .../ssm/test_hybrid_layer_allocation.py | 165 +- .../ssm/test_kimi_delta_attention.py | 383 ++ tests/unit_tests/test_argument_utils.py | 237 +- .../unit_tests/transformer/test_attention.py | 389 +- 14 files changed, 5767 insertions(+), 3899 deletions(-) create mode 100644 megatron/core/ssm/gated_delta_net/kda.py create mode 100644 tests/unit_tests/ssm/test_kimi_delta_attention.py diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 794c109c962..bc699392830 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1,18 +1,17 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging import math import operator import warnings from contextlib import nullcontext -from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple +from typing import Dict, List, Optional, Sequence, Tuple import torch # type: ignore import torch.nn.functional as F # type: ignore from torch import Tensor # type: ignore from megatron.core import parallel_state -from megatron.core.inference.batch_dimensions_utils import TOKEN_ROUNDER as _TOKEN_ROUNDER from megatron.core.inference.batch_dimensions_utils import ( CUDAGraphBatchDimensionBuilder, InferenceBatchDimensions, @@ -23,8 +22,6 @@ PrefixCachingEvictionPolicy, ) from megatron.core.inference.inference_request import DynamicInferenceRequest -from megatron.core.inference.moe import InferenceGroupedGemmBackend -from megatron.core.inference.moe.vllm_fused_moe import VllmFusedMoeBuffers from megatron.core.inference.sampling.base import Sampling from megatron.core.inference.sampling_params import SamplingParams from megatron.core.inference.unified_memory import ( @@ -38,11 +35,8 @@ get_layer_maps_from_layer_type_list, ) from megatron.core.package_info import __version__ as mcore_version -from megatron.core.ssm.ops.gdp.metadata import max_gdp_chunk_counts from megatron.core.transformer import MLATransformerConfig, TransformerConfig -from megatron.core.transformer.enums import InferenceCudaGraphScope from megatron.core.transformer.moe.token_dispatcher_inference import ( - InferenceAllGatherDispatcherBase, NCCLAllGatherDispatcher, NVLSAllGatherVDispatcher, ) @@ -55,12 +49,9 @@ from .base_context import BaseInferenceContext from .gpu_view import ContextGPUView from .kv_block_allocator import KVBlockAllocator -from .mamba_slot_allocator import MAX_INTERMEDIATE_OFFSETS_PER_REQUEST, MambaSlotAllocator +from .mamba_slot_allocator import MambaSlotAllocator from .routing_metadata import RoutingMetadata -# These callbacks are currently consumed only by the Dynamo frontend. -KVEventListener = Callable[[str, dict[str, Any]], None] - try: from .fused_kv_append_kernel import triton_append_key_value_cache except ImportError: @@ -233,69 +224,6 @@ def get_mem_size_str(n_bytes: int) -> str: raise Exception(f"something went wrong, n_bytes={n_bytes}.") -class DynamoHelper: - """Manage KV-cache lifecycle events consumed by the Dynamo frontend.""" - - def __init__(self) -> None: - self._kv_event_listeners: list[KVEventListener] = [] - self._pending_kv_stored_events: list[dict[str, Any]] = [] - - @property - def has_kv_event_listeners(self) -> bool: - """Return whether any KV-event listeners are registered.""" - return bool(self._kv_event_listeners) - - def add_kv_event_listener(self, listener: KVEventListener) -> None: - """Register a KV-cache lifecycle listener. - - Args: - listener: Callback invoked with the event kind and payload. - """ - self._kv_event_listeners.append(listener) - - def queue_kv_stored_event(self, payload: dict[str, Any]) -> None: - """Queue a stored event for publication after a successful forward pass. - - Args: - payload: Stored-event payload. - """ - self._pending_kv_stored_events.append(payload) - - def publish_pending_kv_stored_events(self) -> None: - """Publish blocks whose KV contents were produced by a successful forward pass.""" - pending, self._pending_kv_stored_events = self._pending_kv_stored_events, [] - for payload in pending: - self._emit_kv_event("stored", payload) - - def discard_pending_kv_stored_events(self) -> None: - """Discard registrations left by an interrupted or failed forward pass.""" - self._pending_kv_stored_events.clear() - - def on_kv_blocks_deregistered(self, _block_ids: list[int], hashes: set[int]) -> None: - """Publish removal events for deregistered KV blocks. - - Args: - _block_ids: Deregistered block IDs, unused by Dynamo. - hashes: Hashes of the deregistered blocks. - """ - if hashes: - self._emit_kv_event("removed", {"block_hashes": list(hashes)}) - - def notify_kv_cache_cleared(self) -> None: - """Notify listeners that no previously advertised block is routable.""" - self.discard_pending_kv_stored_events() - if self._kv_event_listeners: - self._emit_kv_event("cleared", {}) - - def _emit_kv_event(self, kind: str, payload: dict[str, Any]) -> None: - """Notify Dynamo listeners without allowing frontend failures to stop inference.""" - for listener in tuple(self._kv_event_listeners): - try: - listener(kind, payload) - except Exception: # pylint: disable=broad-exception-caught - logging.exception("KV-event listener failed while handling %r", kind) - - @internal_api # pylint: disable=line-too-long class DynamicInferenceContext(BaseInferenceContext): @@ -318,7 +246,7 @@ class DynamicInferenceContext(BaseInferenceContext): """ DEFAULT_MAX_TOKENS = 16384 - TOKEN_ROUNDER = _TOKEN_ROUNDER + TOKEN_ROUNDER = 64 REQUEST_ROUNDER = 4 TMS_TAG = "inference_context" @@ -339,8 +267,6 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Hyperparameter for choosing to prioritize prefix hit matches vs minimizing idle load self.prefix_caching_routing_alpha = inference_config.prefix_caching_routing_alpha - self.media_cache_coordinator_policy = inference_config.media_cache_coordinator_policy - self.media_cache_routing_weight = inference_config.media_cache_routing_weight # Monotonic clock for prefix caching LRU eviction ordering. # Incremented each engine step but kept independent so the engine step @@ -350,24 +276,12 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Prefix caching hit tracking (accumulated, reset by engine after logging). self.prefix_cache_hits = 0 # requests that matched at least one cached block self.prefix_cache_blocks_matched = 0 # total matched blocks across all requests - # Prefill compute accounting (drained into engine accumulators each step). - # computed = prompt tokens actually run through the model this step; - # skipped = prompt tokens whose prefill was skipped via a prefix-cache hit. - # A high skipped fraction confirms prefix caching is saving prefill compute - # (so any per-step latency growth is attention-over-context, not re-prefill). - self.prefix_cache_prefill_computed_tokens = 0 - self.prefix_cache_prefill_skipped_tokens = 0 # Engine step counter (used for logging, metrics, and event tracking) self.step_count = 0 self.async_sched_step_count = 0 self.async_sched_compaction_step_count = 0 - # Per-request VLM data (empty when not using multimodal). - self._request_to_image_embeddings: Dict[int, Optional[Tensor]] = {} - self._request_to_image_token_mask: Dict[int, Optional[Tensor]] = {} - self._request_to_image_token_count: Dict[int, int] = {} - self.cache_mla_latent = ( isinstance(model_config, MLATransformerConfig) and model_config.cache_mla_latents ) @@ -392,13 +306,11 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC else: self.num_attention_heads_per_partition = 1 - self.batch_invariant_mode = model_config.batch_invariant_mode self.num_speculative_tokens = inference_config.num_speculative_tokens assert self.num_speculative_tokens < inference_config.block_size_tokens, ( f"num_speculative_tokens ({self.num_speculative_tokens}) must be < " f"block_size_tokens ({inference_config.block_size_tokens})" ) - self._async_sched_token_offsets = None # Cache the PP group we should use for PP collectives inside the context. # If the model provides a pg_collection with a pp group, prefer it. @@ -428,51 +340,30 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Mamba states. mamba_inference_state_config = inference_config.mamba_inference_state_config self.is_hybrid_model = mamba_inference_state_config is not None - self.gdp_num_householder = 0 if self.is_hybrid_model: self.mamba_conv_states_shape = mamba_inference_state_config.conv_states_shape self.mamba_ssm_states_shape = mamba_inference_state_config.ssm_states_shape self.mamba_conv_states_dtype = mamba_inference_state_config.conv_states_dtype self.mamba_ssm_states_dtype = mamba_inference_state_config.ssm_states_dtype self.mamba_chunk_size = mamba_inference_state_config.mamba_chunk_size - self.ssm_chunk_alignment = mamba_inference_state_config.ssm_chunk_alignment - self.gdp_num_householder = mamba_inference_state_config.gdp_num_householder - - if self.batch_invariant_mode: - # Gated Delta Product does not implement batch-invariant mode yet. - assert self.gdp_num_householder == 0, ( - "batch_invariant_mode does not support Gated Delta Product layers; " - "set batch_invariant_mode=False." - ) - assert not self.enable_prefix_caching, ( - "batch_invariant_mode does not support Mamba prefix caching; " - "set enable_prefix_caching=False." - ) - assert self.num_speculative_tokens == 0, ( - "batch_invariant_mode for Mamba dynamic inference only supports " - "one-token decode; set num_speculative_tokens=0." - ) - assert self.mamba_ssm_states_dtype == torch.float32, ( - "batch_invariant_mode requires FP32 Mamba SSM states so state-passing " - "boundaries are not rounded between decode chunks." - ) - # 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. - 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)) - recurrent_layer_map = {} - for global_layer_idx, layer_type in enumerate( - mamba_inference_state_config.layer_type_list - ): - if layer_type in (Symbols.MAMBA, Symbols.GDN): - recurrent_layer_map[global_layer_idx] = len(recurrent_layer_map) + # For hybrid models, the layer map converts the global layer index to the + # corresponding attention layer index or Mamba layer index depending on the + # layer type. + attention_layer_map, dsa_layer_map, gdn_layer_map, kda_layer_map, mamba_layer_map = ( + operator.itemgetter( + Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.GDN, Symbols.KDA, Symbols.MAMBA + )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)) + ) + + if len(gdn_layer_map) > 0: + raise NotImplementedError("GDN layers are not supported for inference.") + if len(kda_layer_map) > 0: + raise NotImplementedError("KDA layers are not supported for inference.") self.num_attention_layers = len(attention_layer_map) + len(dsa_layer_map) - self.num_mamba_layers = len(recurrent_layer_map) - self.layer_map = attention_layer_map | dsa_layer_map | recurrent_layer_map + self.num_mamba_layers = len(mamba_layer_map) + self.layer_map = attention_layer_map | dsa_layer_map | mamba_layer_map else: # The layer map is the identity function for pure Transformer models. # Use the same per-PP-rank layer count as TransformerBlock (handles @@ -601,7 +492,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC paused_buffer_size_bytes = int(paused_buffer_size_bytes * (1.0 - mamba_memory_ratio)) block_count = buffer_size_bytes // self.block_size_bytes - block_count = max(2, block_count) # need >= 1 usable block + 1 dummy block + block_count = max(2, block_count) # need >= 1 active block + 1 dummy block paused_block_count = paused_buffer_size_bytes // self.block_size_bytes elif self.is_hybrid_model and inference_config.max_requests is not None: # Auto-derive mamba/KV split from max_requests. Allocate exactly enough @@ -615,19 +506,19 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC ) mamba_max_requests = inference_config.max_requests - # Subtract Mamba memory proportionally from both configured buffer inputs. + # Subtract mamba memory proportionally from active and paused buffers. mamba_memory_ratio = mamba_memory_needed / total_memory buffer_size_bytes = int(buffer_size_bytes * (1.0 - mamba_memory_ratio)) paused_buffer_size_bytes = int(paused_buffer_size_bytes * (1.0 - mamba_memory_ratio)) block_count = buffer_size_bytes // self.block_size_bytes - block_count = max(2, block_count) # need >= 1 usable block + 1 dummy block + block_count = max(2, block_count) # need >= 1 active block + 1 dummy block paused_block_count = paused_buffer_size_bytes // self.block_size_bytes else: block_count = buffer_size_bytes // ( self.block_size_bytes + mamba_states_memory_per_request ) - block_count = max(2, block_count) # need >= 1 usable block + 1 dummy block + block_count = max(2, block_count) # need >= 1 active block + 1 dummy block paused_block_count = paused_buffer_size_bytes // ( self.block_size_bytes + mamba_states_memory_per_request ) @@ -652,17 +543,13 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.kv_block_allocator = KVBlockAllocator( context=self, - pool_size=( + total_count=( block_count if self.unified_memory_level == 0 else block_count + paused_block_count ), - paused_limit=paused_block_count, + paused_count=paused_block_count, enable_prefix_caching=self.enable_prefix_caching, prefix_caching_eviction_policy=self.prefix_caching_eviction_policy, ) - self.dynamo_helper = DynamoHelper() - self.kv_block_allocator.add_blocks_deregistered_observer( - self.dynamo_helper.on_kv_blocks_deregistered - ) # Track request metadata. request_metadata_types = inference_config.request_metadata_types @@ -672,8 +559,6 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Initialize context state. self.params_dtype = model_config.params_dtype - self.hidden_size = model_config.hidden_size - self.inference_cuda_graph_scope = model_config.inference_cuda_graph_scope self.max_sequence_length = inference_config.max_sequence_length # Block ids. With speculative decoding, blocks are pre-allocated when the @@ -686,7 +571,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Set max_requests, max_tokens. if inference_config.max_requests is None: # Maximize compute utilization by defaulting to 1 block per request. - self.max_requests = self.kv_block_allocator.pool_size - 1 # -1 for dummy block + self.max_requests = self.kv_block_allocator.total_count - 1 # -1 for dummy block # Adjust max_requests for Mamba memory constraints if necessary if self.is_hybrid_model and mamba_max_requests < self.max_requests: @@ -704,15 +589,6 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.max_tokens = inference_config.max_tokens or self.DEFAULT_MAX_TOKENS - # Per-step upper bound on Mamba intermediate-state extractions, shared with - # MambaMetadata and MambaSlotAllocator so scratch/metadata buffers and the - # budget accounting agree. Bounded both by the token budget (one block - # boundary per block_size_tokens) and by the request budget - # (MAX_INTERMEDIATE_OFFSETS_PER_REQUEST per request); - token_based_count = math.ceil(self.max_tokens / self.block_size_tokens) - request_based_count = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * self.max_requests - self.max_mamba_intermediate_states_per_step = min(token_based_count, request_based_count) - assert self.max_tokens >= self.max_requests, ( f"max_tokens ({self.max_tokens}) must be >= " f"max_requests ({self.max_requests}), " @@ -725,7 +601,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.non_graph_attn_metadata = {} self.graph_attn_metadata["mha_metadata"] = GraphedMHAMetadata( - block_count_total=self.kv_block_allocator.pool_size, + block_count_total=self.kv_block_allocator.total_count, max_kv_block_count=self.max_kv_block_count, max_requests=self.max_requests, block_size_tokens=self.block_size_tokens, @@ -733,7 +609,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC ) self.non_graph_attn_metadata["mha_metadata"] = NonGraphedMHAMetadata( - block_count_total=self.kv_block_allocator.pool_size, + block_count_total=self.kv_block_allocator.total_count, max_kv_block_count=self.max_kv_block_count, max_requests=self.max_requests, block_size_tokens=self.block_size_tokens, @@ -753,12 +629,6 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC and model_config.inference_moe_token_dispatcher_type == 'nccl' ) - # are we using the inference_optimized nvls ep dispatcher for MoEs? - self._nvls_dispatcher = ( - get_pg_size(self.expert_model_parallel_group) > 1 - and model_config.inference_moe_token_dispatcher_type == 'nvls' - ) - # are we using the training a2a dispatcher for MoEs? # Note that this is not optimal for speed. self._training_ep_dispatcher = ( @@ -777,15 +647,12 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC ) # CUDA graph token budget for prefill/mixed graphs. Decode graphs are always - # capped at max_requests * (num_speculative_tokens + 1) inside the helper. By - # default the prefill/mixed range is bounded by `cuda_graph_max_tokens`, clamped - # to never fall below that decode bound nor exceed `max_tokens`; setting - # `cuda_graph_all_prefills` widens the range to the full `max_tokens`. - decode_bound = self.max_requests * (self.num_speculative_tokens + 1) + # capped at max_requests * (num_speculative_tokens + 1) inside the helper; this + # only widens the prefill/mixed range when `cuda_graph_all_prefills` is set. cuda_graph_max_tokens = ( self.max_tokens if inference_config.cuda_graph_all_prefills - else min(max(inference_config.cuda_graph_max_tokens, decode_bound), self.max_tokens) + else self.max_requests * (self.num_speculative_tokens + 1) ) # CUDA graph config list. @@ -806,58 +673,22 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Allocate per-step dispatcher buffers upfront so update_metadata never # triggers an allocation inside a captured CUDA graph. - # - # The shared _valid_tokens_tensor scalar is read as a pointer by both fused - # MoE backends (mcore_fused_moe and vllm_fused_moe) regardless of EP size, so - # allocate it unconditionally (covers EP=1, where no dispatcher comm buffers - # exist). The EP>1 dispatchers below reallocate it as part of their own buffer - # setup, which is harmless. - InferenceAllGatherDispatcherBase.allocate_valid_tokens_tensor() - if self._nccl_ep_dispatcher: - NCCLAllGatherDispatcher.allocate_buffers() - elif self._nvls_dispatcher: - # Use moe_latent_size if set (latent MoE: SuperV3, UltraV3), else hidden_size. - moe_hidden_size = model_config.moe_latent_size or model_config.hidden_size - NVLSAllGatherVDispatcher.allocate_buffers( - per_rank_worst_case_token_count=self.round_up_tokens(self.max_tokens) // tp_size, - topk=model_config.moe_router_topk, - hidden_size=moe_hidden_size, - ep_group=self.expert_model_parallel_group, - ) - - # Pre-allocate the vLLM fused-MoE intermediates so no allocation happens - # inside CUDA graph capture; one buffer set is shared by all MoE layers - # and graphs. Like the dispatcher buffers above, these persist across - # engine suspend/resume. - if ( - model_config.num_moe_experts - and model_config.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.VLLM - ): - ep_size = get_pg_size(self.expert_model_parallel_group) - moe_hidden_size = model_config.moe_latent_size or model_config.hidden_size - # Worst-case rows entering the MoE: the fixed NVLS AGV buffer height - # (per-rank worst case * ep_size); max_tokens covers the EP=1 / NCCL paths. - moe_max_rows = max( - self.max_tokens, self.round_up_tokens(self.max_tokens) // tp_size * ep_size - ) - VllmFusedMoeBuffers.allocate_buffers( - max_tokens=moe_max_rows, - topk=model_config.moe_router_topk, - fc1_output_size=model_config.moe_ffn_hidden_size - * (2 if model_config.gated_linear_unit else 1), - hidden_size=moe_hidden_size, - num_local_experts=model_config.num_moe_experts // ep_size, - ) + if get_pg_size(self.expert_model_parallel_group) > 1: + if self._nccl_ep_dispatcher: + NCCLAllGatherDispatcher.allocate_buffers() + else: + # Use moe_latent_size if set (latent MoE: SuperV3, UltraV3), else hidden_size. + moe_hidden_size = model_config.moe_latent_size or model_config.hidden_size + NVLSAllGatherVDispatcher.allocate_buffers( + per_rank_worst_case_token_count=self.round_up_tokens(self.max_tokens) + // tp_size, + topk=model_config.moe_router_topk, + hidden_size=moe_hidden_size, + ep_group=self.expert_model_parallel_group, + ) # Deal with chunked prefill self.enable_chunked_prefill = inference_config.enable_chunked_prefill - if self.batch_invariant_mode and self.is_hybrid_model and self.enable_chunked_prefill: - # A chunk plus its final token must fit in one step; otherwise a prompt - # of that length can never advance without an invalid one-token tail. - assert self.max_tokens > self.ssm_chunk_alignment, ( - "batch-invariant SSM chunked prefill requires max_tokens > " - f"ssm_chunk_alignment ({self.ssm_chunk_alignment})." - ) # FlashInfer. if inference_config.use_flashinfer_fused_rope is True: @@ -867,30 +698,17 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.use_flashinfer_fused_rope = inference_config.use_flashinfer_fused_rope self.inference_grouped_gemm_backend = model_config.inference_grouped_gemm_backend - # Placeholder for the MTP decoder hidden-states buffer; allocated inside - # initialize_all_tensors() when num_speculative_tokens > 0. - self.mtp_decoder_hidden_states = None - # Allocate GPU state. self.is_tensor_state_allocated = False - self._bookkeeping_no_real_work = False self.initialize_all_tensors() - # Bind the GPU real-token-count tensor onto the NVLS dispatcher class - # so it can mask out CUDA-graph padding tokens during routing. The - # tensor lives inside gpu_view._buf (fixed address) and is refreshed - # each step by transfer_bookkeeping_to_gpu(). NVLS-only — the NCCL - # dispatcher requires equal token counts across ranks already. - if self._nvls_dispatcher: - NVLSAllGatherVDispatcher.set_real_token_count_tensor(self.gpu_view.real_token_count) - # Print info. - pool_size = self.kv_block_allocator.pool_size - usable_blocks = pool_size - 1 - paused_limit = self.kv_block_allocator.paused_limit - pool_size_bytes = pool_size * self.block_size_bytes - usable_bytes = usable_blocks * self.block_size_bytes - paused_limit_bytes = paused_limit * self.block_size_bytes + active_blocks = self.kv_block_allocator.active_count + total_blocks = self.kv_block_allocator.total_count + paused_blocks = self.kv_block_allocator.paused_count + active_kv_bytes = active_blocks * self.block_size_bytes + total_kv_bytes = total_blocks * self.block_size_bytes + paused_kv_bytes = paused_blocks * self.block_size_bytes log_lines = [ "DynamicInferenceContext: configuration summary", @@ -901,10 +719,9 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC f" max_kv_blocks_per_req: {self.max_kv_block_count}", f" KV cache:", f" block_size_bytes: {get_mem_size_str(self.block_size_bytes)}", - f" pool_size: {pool_size} ({get_mem_size_str(pool_size_bytes)})", - f" usable_blocks: {usable_blocks} ({get_mem_size_str(usable_bytes)})", - f" paused_limit: {paused_limit} " - f"({get_mem_size_str(paused_limit_bytes)})", + f" active_blocks: {active_blocks} ({get_mem_size_str(active_kv_bytes)})", + f" paused_blocks: {paused_blocks} ({get_mem_size_str(paused_kv_bytes)})", + f" total_blocks: {total_blocks} ({get_mem_size_str(total_kv_bytes)})", ] if self.is_hybrid_model: @@ -945,21 +762,12 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC and prefix_caching_mamba_gb is not None and prefix_caching_mamba_gb > 0 ): - assert self.mamba_slot_allocator is not None prefix_cache_bytes = int(prefix_caching_mamba_gb * 1024**3) - # The allocator contains the PP-synchronized durable capacity. - # Scratch remains stage-local because it is temporary forward - # storage and does not participate in prefix-cache matching. - scratch_slots = self.max_mamba_intermediate_states_per_step - scratch_bytes = scratch_slots * mamba_bytes_per_req - durable_slots = self.mamba_slot_allocator.max_slots + prefix_cache_slots = prefix_cache_bytes // mamba_bytes_per_req log_lines += [ f" Mamba prefix cache:", f" budget: {get_mem_size_str(prefix_cache_bytes)}", - f" extraction_scratch: {scratch_slots} slots " - f"({get_mem_size_str(scratch_bytes)})", - f" durable_slots: {durable_slots} " - f"({get_mem_size_str(durable_slots * mamba_bytes_per_req)})", + f" slots: {prefix_cache_slots}", f" per_slot: {get_mem_size_str(mamba_bytes_per_req)}", ] @@ -972,7 +780,7 @@ def _allocate_memory_buffer(self): self.memory_buffer = torch.empty( ( self.num_attention_layers, - self.kv_block_allocator.pool_size, + self.kv_block_allocator.total_count, self.block_size_tokens, self.kv_reduced_dim, ), @@ -984,7 +792,7 @@ def _allocate_memory_buffer(self): ( 2, # key and value self.num_attention_layers, - self.kv_block_allocator.pool_size, + self.kv_block_allocator.total_count, self.block_size_tokens, self.num_attention_heads_per_partition, self.hidden_size_per_attention_head, @@ -1008,35 +816,25 @@ def _allocate_mamba_states(self): self.mamba_metadata = MambaMetadata( max_requests=self.max_requests, max_tokens=self.max_tokens, - max_intermediate_count=self.max_mamba_intermediate_states_per_step, mamba_chunk_size=self.mamba_chunk_size, d_conv=self.mamba_conv_states_shape[-1], - decode_indices_dtype=self._mamba_decode_indices_dtype, - gdp_num_householder=self.gdp_num_householder, ) # Bind the unified CPU/GPU buffers so the per-step Mamba metadata # fields ride along with the single coalesced H2D in # transfer_bookkeeping_to_gpu(). - _mamba_cpu_bufs = { - "batch_indices_decode": self._cpu_mamba_batch_indices_decode, - "batch_indices_prefill": self._cpu_mamba_batch_indices_prefill, - "seq_idx": self._cpu_mamba_seq_idx, - "cu_seqlens": self._cpu_mamba_cu_seqlens, - "cu_chunk_seqlens": self._cpu_mamba_cu_chunk_seqlens, - "last_chunk_indices": self._cpu_mamba_last_chunk_indices, - "seq_idx_for_varlen": self._cpu_mamba_seq_idx_for_varlen, - "conv_seq_idx": self._cpu_mamba_conv_seq_idx, - "conv_seq_start": self._cpu_mamba_conv_seq_start, - } - if self.gdp_num_householder > 0: - _mamba_cpu_bufs.update( - { - "gdp_chunk_indices": self._cpu_gdp_chunk_indices, - "gdp_chunk_indices_dp": self._cpu_gdp_chunk_indices_dp, - "gdp_chunk_offsets": self._cpu_gdp_chunk_offsets, - } - ) - self.mamba_metadata.bind_cpu_buffers(_mamba_cpu_bufs) + self.mamba_metadata.bind_cpu_buffers( + { + "batch_indices_decode": self._cpu_mamba_batch_indices_decode, + "batch_indices_prefill": self._cpu_mamba_batch_indices_prefill, + "seq_idx": self._cpu_mamba_seq_idx, + "cu_seqlens": self._cpu_mamba_cu_seqlens, + "cu_chunk_seqlens": self._cpu_mamba_cu_chunk_seqlens, + "last_chunk_indices": self._cpu_mamba_last_chunk_indices, + "seq_idx_for_varlen": self._cpu_mamba_seq_idx_for_varlen, + "conv_seq_idx": self._cpu_mamba_conv_seq_idx, + "conv_seq_start": self._cpu_mamba_conv_seq_start, + } + ) self.mamba_metadata.bind_gpu_buffers(self.gpu_view) self.mamba_conv_states = torch.empty( (self.num_mamba_layers, self.max_requests) + self.mamba_conv_states_shape, @@ -1150,9 +948,6 @@ def initialize_all_tensors(self) -> None: device='cpu', pin_memory=True, ) - self._async_sched_token_offsets = torch.arange( - self.num_speculative_tokens + 1, device='cpu' - ) # Track request metadata. Backed by pinned CPU memory: bookkeeping is # CPU-resident; GPU consumers read from the active-slice mirror in @@ -1200,10 +995,6 @@ def initialize_all_tensors(self) -> None: _tok_int32_bytes = self.max_tokens * 4 # Request-level fields are all 4 bytes wide (5 int32 + 2 float32 = 7 fields). _req_4byte_bytes = self.max_requests * 4 - # Scalar: real (unpadded) token count for the current step. Refreshed - # in transfer_bookkeeping_to_gpu(); read on GPU via - # `gpu_view.real_token_count` (MoE routing masks padding tokens). - _real_token_count_bytes = 4 # MHA section: 5 fields (int32) shared between GraphedMHAMetadata and # NonGraphedMHAMetadata. max_bs == max_requests. _mha_query_lengths_bytes = self.max_requests * 4 @@ -1215,7 +1006,6 @@ def initialize_all_tensors(self) -> None: 3 * _tok_int64_bytes + 3 * _tok_int32_bytes + 7 * _req_4byte_bytes - + _real_token_count_bytes + _mha_query_lengths_bytes + _mha_cu_query_seq_lengths_bytes + _mha_kv_seq_lengths_bytes @@ -1224,18 +1014,12 @@ def initialize_all_tensors(self) -> None: ) # Mamba section (hybrid models only). Must match the MambaMetadata # shapes (mirrors the layout documented in ContextGPUView). - # batch_indices_decode is int32 in batch-invariant mode and int64 otherwise; - # all other fields are int32. + # batch_indices_decode is int64; all other fields are int32. if self.is_hybrid_model: - self._mamba_decode_indices_dtype = ( - torch.int32 if self.batch_invariant_mode else torch.int64 - ) - _decode_index_bytes = 4 if self.batch_invariant_mode else 8 - _mamba_align_pad = ( - _decode_index_bytes - _pre_mamba_bytes % _decode_index_bytes - ) % _decode_index_bytes + # mamba_batch_indices_decode is int64; pad to 8-byte alignment. + _mamba_align_pad = (8 - _pre_mamba_bytes % 8) % 8 self._max_mamba_chunks = self.max_tokens // self.mamba_chunk_size + self.max_requests - _mamba_batch_indices_decode_bytes = self.max_requests * _decode_index_bytes + _mamba_batch_indices_decode_bytes = self.max_requests * 8 _mamba_batch_indices_prefill_bytes = self.max_requests * 4 _mamba_seq_idx_bytes = self.max_tokens * 4 _mamba_cu_seqlens_bytes = (self.max_requests + 1) * 4 @@ -1269,27 +1053,6 @@ def initialize_all_tensors(self) -> None: + _mamba_conv_seq_idx_bytes + _mamba_conv_seq_start_bytes ) - # GDP chunk descriptor section, only present for GDP models. All int32, - # so no alignment padding is needed after the Mamba section. - # gdp_chunk_indices int32 (max_gdp_chunks, 2) - # gdp_chunk_indices_dp int32 (max_gdp_chunks_dp, 2) - # gdp_chunk_offsets int32 (max_requests + 1,) - if self.gdp_num_householder > 0: - self._max_gdp_chunks, self._max_gdp_chunks_dp = max_gdp_chunk_counts( - self.max_tokens, self.max_requests, self.gdp_num_householder - ) - _gdp_chunk_indices_bytes = self._max_gdp_chunks * 2 * 4 - _gdp_chunk_indices_dp_bytes = self._max_gdp_chunks_dp * 2 * 4 - _gdp_chunk_offsets_bytes = (self.max_requests + 1) * 4 - else: - self._max_gdp_chunks = 0 - self._max_gdp_chunks_dp = 0 - _gdp_chunk_indices_bytes = 0 - _gdp_chunk_indices_dp_bytes = 0 - _gdp_chunk_offsets_bytes = 0 - _total_bytes += ( - _gdp_chunk_indices_bytes + _gdp_chunk_indices_dp_bytes + _gdp_chunk_offsets_bytes - ) self._cpu_bookkeeping_buf = torch.empty( _total_bytes, dtype=torch.uint8, device='cpu', pin_memory=True ) @@ -1369,14 +1132,6 @@ def initialize_all_tensors(self) -> None: ].view(torch.int32) _off += _req_4byte_bytes - # Scalar staging slot for the real (unpadded) token count. Refreshed - # from `self.batch_dimensions.token_count` in transfer_bookkeeping_to_gpu() - # and read on GPU via `gpu_view.real_token_count`. - self._staging_real_token_count = self._cpu_bookkeeping_buf[ - _off : _off + _real_token_count_bytes - ].view(torch.int32) - _off += _real_token_count_bytes - # Static tensor addresses to make `last_token_logits` graphable with speculative decoding. max_logit_idxs = self.max_requests * (self.num_speculative_tokens + 1) self.active_logit_idxs = torch.zeros( @@ -1420,7 +1175,7 @@ def initialize_all_tensors(self) -> None: _off += _mamba_align_pad self._cpu_mamba_batch_indices_decode = self._cpu_bookkeeping_buf[ _off : _off + _mamba_batch_indices_decode_bytes - ].view(self._mamba_decode_indices_dtype) + ].view(torch.int64) _off += _mamba_batch_indices_decode_bytes self._cpu_mamba_batch_indices_prefill = self._cpu_bookkeeping_buf[ _off : _off + _mamba_batch_indices_prefill_bytes @@ -1457,24 +1212,6 @@ def initialize_all_tensors(self) -> None: ].view(torch.int32) _off += _mamba_conv_seq_start_bytes - if self.gdp_num_householder > 0: - self._cpu_gdp_chunk_indices = ( - self._cpu_bookkeeping_buf[_off : _off + _gdp_chunk_indices_bytes] - .view(torch.int32) - .view(self._max_gdp_chunks, 2) - ) - _off += _gdp_chunk_indices_bytes - self._cpu_gdp_chunk_indices_dp = ( - self._cpu_bookkeeping_buf[_off : _off + _gdp_chunk_indices_dp_bytes] - .view(torch.int32) - .view(self._max_gdp_chunks_dp, 2) - ) - _off += _gdp_chunk_indices_dp_bytes - self._cpu_gdp_chunk_offsets = self._cpu_bookkeeping_buf[ - _off : _off + _gdp_chunk_offsets_bytes - ].view(torch.int32) - _off += _gdp_chunk_offsets_bytes - assert _off == _total_bytes, f"layout bug: wrote {_off} of {_total_bytes} bytes" # GPU view: the single interface for GPU code to read context state. @@ -1485,16 +1222,7 @@ def initialize_all_tensors(self) -> None: max_kv_blocks=self.max_kv_block_count, device=torch.cuda.current_device(), max_mamba_chunks=self._max_mamba_chunks, - max_gdp_chunks=self._max_gdp_chunks, - max_gdp_chunks_dp=self._max_gdp_chunks_dp, - # Same presence predicate as the CPU-side GDP section above; the two - # layouts must byte-match, so they gate on the same flag. - has_gdp=self.gdp_num_householder > 0, - mamba_decode_indices_dtype=( - self._mamba_decode_indices_dtype if self.is_hybrid_model else torch.int64 - ), ) - self._bookkeeping_h2d_done_event = torch.cuda.Event() # Cache of (input_ids_view, pos_ids_view) keyed by num_tokens. Instead of slicing and # unsqueezing on every new inference step (constructing new TensorImpls at 30-60 us), @@ -1541,40 +1269,6 @@ def initialize_all_tensors(self) -> None: and self.config.enable_prefix_caching ): self._allocate_mamba_cache(self.config.prefix_caching_mamba_gb) - elif self.is_hybrid_model and self.config.enable_prefix_caching: - # Memory-only mode: prefix caching on a hybrid model without a Mamba - # cache budget deduplicates identical KV prefixes for memory savings, - # but does NOT cache Mamba recurrent state. Prefill skipping is - # therefore disabled (prefix_skip_tokens is forced to 0) and every - # token is recomputed, so results stay correct -- but the main latency - # benefit of prefix caching is forgone. Warn so a user who expected - # full caching knows to set prefix_caching_mamba_gb. - logging.warning( - "enable_prefix_caching is set on a hybrid (Mamba) model but " - "prefix_caching_mamba_gb is not configured (got %r). Running in " - "memory-only mode: identical KV prefixes are deduplicated for " - "memory savings, but Mamba state caching and prefill skipping are " - "disabled (every token is recomputed). Set prefix_caching_mamba_gb " - "> 0 to enable full prefix caching.", - self.config.prefix_caching_mamba_gb, - ) - - # MTP speculative decoding: persistent buffer for decoder hidden states. - # Only needed for block-scope CUDA graphs, where the Python assignment in - # forward() runs only during graph capture. Using copy_() into a fixed - # buffer ensures every batch-size graph replay writes to the same GPU - # address. Sized to max_tokens; only [:actual_tokens] is valid each step. - if ( - self.num_speculative_tokens > 0 - and self.inference_cuda_graph_scope == InferenceCudaGraphScope.block - ): - self.mtp_decoder_hidden_states = torch.empty( - self.max_tokens, - 1, - self.hidden_size, - device=torch.cuda.current_device(), - dtype=self.params_dtype, - ) # Reset tensor-related metadata. self.reset_metadata() @@ -1795,10 +1489,6 @@ def pad_active_slices(self): # Padded gather indices fan in to row 0 harmlessly when used by FlashInfer. self.active_request_last_token_idxs[padding_request_slice].fill_(0) - def _clear_input_id_padding(self) -> None: - """Restore zero-filled padding without touching active VLM ``-1`` placeholders.""" - self.token_to_input_ids[self.padding_slice].zero_() - def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None: """Append to KV cache. @@ -1819,7 +1509,6 @@ def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) padded_active_token_count=self.padded_active_token_count, token_to_block_idx=self.gpu_view.token_to_block_idx, token_to_local_position_within_kv_block=self.gpu_view.token_to_local_position_within_kv_block, - dummy_block_idx=self.kv_block_allocator.dummy_block_idx, ) block_idx = self.gpu_view.token_to_block_idx[: self.padded_active_token_count] @@ -1908,58 +1597,15 @@ def _allocate_mamba_cache(self, mamba_gb: float) -> None: ssm_size = _math.prod(self.mamba_ssm_states_shape) * self.mamba_ssm_states_dtype.itemsize per_slot_bytes = self.num_mamba_layers * (conv_size + ssm_size) total_bytes = int(mamba_gb * 1024**3) - - # MambaSlotAllocator allocates two GPU buffer families with the same - # per-slot footprint, both of which must fit in this budget: - # - "durable" cache: self.ssm_states / self.conv_states, sized to - # `max_slots` slots (computed below). - # - "scratch" buffers: self.intermediate_ssm_out / self.intermediate_conv_out, - # fixed CUDA-graph-safe staging for intermediate-state - # extraction, sized to the per-step token-budget cap - # `scratch_slots` = max_mamba_intermediate_states_per_step. - # The scratch is not part of the durable cache but consumes the same - # per-slot bytes, so reserve it from the budget up front before sizing the - # durable cache; otherwise total usage silently exceeds mamba_gb (and can - # OOM) when scratch_slots > max_slots. - scratch_slots = self.max_mamba_intermediate_states_per_step - scratch_bytes = scratch_slots * per_slot_bytes - max_slots = (total_bytes - scratch_bytes) // per_slot_bytes # durable slots - local_max_slots = max_slots - - # Prefix-cache state is replicated across pipeline stages. Use the - # smallest stage-local capacity so identical cache operations produce - # identical block-to-slot mappings on every stage. A recurrent prefix - # is executable only when all stages retain its state, so additional - # slots on an individual stage would not increase usable cache capacity. - if get_pg_size(self.pipeline_parallel_group) > 1: - max_slots_tensor = torch.tensor( - max_slots, dtype=torch.int64, device=torch.cuda.current_device() - ) - torch.distributed.all_reduce( - max_slots_tensor, - op=torch.distributed.ReduceOp.MIN, - group=self.pipeline_parallel_group, - ) - max_slots = int(max_slots_tensor.item()) - + max_slots = total_bytes // per_slot_bytes if max_slots < 1: - if local_max_slots >= 1: - raise ValueError( - f"Mamba prefix cache budget (prefix_caching_mamba_gb={mamba_gb:.4g} GB) " - f"has room for {local_max_slots} durable slots on this pipeline stage, but " - f"another stage has room for fewer than one. Cache capacity is mirrored to " - f"the minimum across stages; increase prefix_caching_mamba_gb, reduce " - f"max_tokens, or rebalance the Mamba layers." - ) - raise ValueError( - f"Mamba prefix cache budget (prefix_caching_mamba_gb={mamba_gb:.4g} GB) " - f"is too small. The CUDA-graph extraction scratch reserves " - f"{scratch_bytes / 1024**3:.4g} GB ({scratch_slots} slots x " - f"{per_slot_bytes / 1024:.1f} KB/slot), leaving room for " - f"fewer than one durable cache slot. Increase prefix_caching_mamba_gb " - f"to at least {(scratch_bytes + per_slot_bytes) / 1024**3:.4g} GB, or " - f"reduce max_tokens." + logging.warning( + "Mamba cache budget (%.3f GB) too small for even 1 slot " + "(need %.3f GB per slot). Mamba caching disabled.", + mamba_gb, + per_slot_bytes / 1024**3, ) + return self.mamba_slot_allocator = MambaSlotAllocator( context=self, @@ -1975,14 +1621,9 @@ def _allocate_mamba_cache(self, mamba_gb: float) -> None: ) logging.info( - "Mamba prefix cache: %d durable slots (%.3f GB) + %d scratch slots " - "(%.3f GB) = %.3f GB total within %.3f GB budget, per-slot %.1f KB", + "Mamba prefix cache: %d slots (%.3f GB), per-slot %.1f KB", max_slots, max_slots * per_slot_bytes / 1024**3, - scratch_slots, - scratch_bytes / 1024**3, - (max_slots + scratch_slots) * per_slot_bytes / 1024**3, - mamba_gb, per_slot_bytes / 1024, ) @@ -2050,6 +1691,7 @@ def apply_rotary_emb_query( cp_group=cp_group, mscale=mscale, mla_rotary_interleaved=config.multi_latent_attention, + max_seqlen=query_emb.size(0), ) return query @@ -2376,8 +2018,6 @@ def add_dummy_requests_for_expert_parallel_step( self.request_to_kv_block_ids[0:N, 0] = dummy_block_idx # 3. Token-level state consumed by the triton KV append kernel. - # Dummy slots are active, so initialize their embedding IDs explicitly. - self.token_to_input_ids[0:T].zero_() self.token_to_block_idx[0:T] = dummy_block_idx # Compute per-request token positions: e.g. query_lengths [3,2] -> [0,1,2,0,1] query_lengths = self.request_query_lengths[0:N] @@ -2411,9 +2051,7 @@ def initialize_attention_state( *, construct_graph_dimensions: Optional[InferenceBatchDimensions] = None, is_expert_parallel_dummy_cuda_graph_step: bool = False, - transfer_bookkeeping_to_gpu: bool = True, - record_bookkeeping_done_event: bool = False, - ) -> Optional[torch.cuda.Event]: + ) -> None: """Initialize attention state so that every layer can use it. Args: @@ -2421,17 +2059,8 @@ def initialize_attention_state( The graph config to use for constructing the cuda graphs. is_expert_parallel_dummy_cuda_graph_step (bool): Whether this is a dummy expert model parallel step. - transfer_bookkeeping_to_gpu (bool): Whether to publish the prepared - CPU bookkeeping snapshot to GPU before returning. Legacy - callers publish immediately; async scheduling binds the GPU - views here and publishes their values later. - record_bookkeeping_done_event (bool): Whether to record an event - after the bookkeeping H2D transfer. - - Returns: - Optional[torch.cuda.Event]: Event marking bookkeeping H2D - completion, or `None` when no event was requested or no - transfer was performed. + Return: + None. """ # Launch deferred Mamba GPU ops first (state zeroing/restore) so they # overlap with the CPU work below. These are non-blocking GPU kernels. @@ -2509,7 +2138,6 @@ def initialize_attention_state( self.padded_active_token_count = self.padded_batch_dimensions.token_count self.padded_active_request_count = self.padded_batch_dimensions.req_count self.padding_slice = slice(self.active_token_count, self.padded_active_token_count) - self._clear_input_id_padding() self.build_active_slices( min(self.padded_active_request_count, self.max_requests - self.paused_request_count) @@ -2597,19 +2225,10 @@ def initialize_attention_state( self._cpu_mha_cu_kv_seq_lengths[real_bs] ) - # Block table: [0:real_bs] real, [real_bs:padded_bs] = dummy block. + # Block table: [0:real_bs] real, [real_bs:padded_bs] = -1 sentinel. self._cpu_mha_block_table[:real_bs] = request_to_kv_block_ids_view[:real_bs] if real_bs < padded_bs: - self._cpu_mha_block_table[real_bs:padded_bs] = self.kv_block_allocator.dummy_block_idx - # Real rows must avoid having a -1 sentinel in their trailing columns, - # because the kernel treats the -1 sentinel as a real value. - # We cannot avoid writing -1 into `request_to_kv_block_ids`; other logic needs it. - # The only option is to overwrite the -1 with a dummy block index via `masked_fill`. - if real_bs > 0: - _real_rows = self._cpu_mha_block_table[:real_bs] - # masked_fill_ over the pinned int32 view: no index_put/nonzero - # temporaries on the per-step CPU path. - _real_rows.masked_fill_(_real_rows < 0, self.kv_block_allocator.dummy_block_idx) + self._cpu_mha_block_table[real_bs:padded_bs] = -1 # Max sequence lengths (Python scalars; consumed as kernel launch args). if not self.using_cuda_graph_this_step() and real_bs > 0: @@ -2679,17 +2298,12 @@ def initialize_attention_state( # No-op when the queue is already empty (regular non-warmup steps). self._execute_pending_mamba_ops() - # Record whether this step produces real output — false on CUDA-graph - # capture (warmup) or dummy EP steps. Used by transfer_bookkeeping_to_gpu - # to publish real_token_count=0 so MoE routing masks all padding tokens. - self._bookkeeping_no_real_work = ( - construct_graph_dimensions is not None or is_expert_parallel_dummy_cuda_graph_step - ) - - # Preserve the existing behavior for callers that do not publish explicitly. - if transfer_bookkeeping_to_gpu: - return self.transfer_bookkeeping_to_gpu(record_done_event=record_bookkeeping_done_event) - return None + # Run the H2D transfer here so callers that bypass the controller + # (e.g. unit tests that call `model.forward()` directly after + # `initialize_attention_state()`) see populated GPU bookkeeping. The + # text-generation controller still calls `transfer_bookkeeping_to_gpu` + # explicitly; that second call is a cheap idempotent re-copy. + self.transfer_bookkeeping_to_gpu() def _execute_pending_mamba_ops(self) -> None: """Execute Mamba GPU operations deferred from add_request() / update_requests(). @@ -2715,33 +2329,18 @@ def _execute_pending_mamba_ops(self) -> None: self.mamba_ssm_states[:, indices] = 0.0 self._pending_mamba_zeros.clear() - def transfer_bookkeeping_to_gpu( - self, skip_token_input_ids: bool = False, record_done_event: bool = False - ) -> Optional[torch.cuda.Event]: + def transfer_bookkeeping_to_gpu(self) -> None: """Batch transfer CPU bookkeeping state to GPU staging buffers. - Legacy steps call this from initialize_attention_state(). Async - scheduling instead delays publication until after preparation and the - GPU sample-to-input copy. Legacy transfers block because the pinned CPU - source is re-staged in place. Async scheduling requests an event-tracked - non-blocking copy and synchronizes that event before reusing the source. + Called after initialize_attention_state() and before the forward pass. + All copies use non_blocking=True with pinned CPU memory. CUDA stream + ordering guarantees the forward pass sees completed transfers. The bookkeeping fields are backed by one contiguous pinned CPU buffer - and one contiguous GPU buffer; a single memcpy covers the whole - transfer. Request-level staging slots are refreshed from the persistent - CPU tensors immediately before the H2D (GPU reads them at `[:n_active]` + and one contiguous GPU buffer; a single cudaMemcpyAsync suffices. + Request-level staging slots are refreshed from the persistent CPU + tensors immediately before the H2D (GPU reads them at `[:n_active]` while CPU bookkeeping keeps them at `[paused_count:total_count)`). - - Args: - skip_token_input_ids (bool): If true, leave - `gpu_view.token_to_input_ids` unchanged while copying the rest - of the bookkeeping buffer. - record_done_event (bool): Whether to record and return an event after - an asynchronous bookkeeping transfer. - - Returns: - Optional[torch.cuda.Event]: Event marking H2D completion, or `None` - when no event was requested. """ n_active = self.total_request_count - self.paused_request_count active_slice = slice(self.paused_request_count, self.total_request_count) @@ -2775,32 +2374,11 @@ def transfer_bookkeeping_to_gpu( self._staging_request_query_lengths[n_active:padded_active] = 0 self._staging_request_kv_length_offsets[n_active:padded_active] = 0 - # Real (unpadded) token count for this step. CUDA-graph replay pads - # the token dim to a captured size; MoE routing reads this on GPU and - # rewrites padding rows' routing entries to -1 so they don't go to - # any expert. Set to 0 on CUDA-graph capture / dummy EP steps so - # every row gets masked out. - self._staging_real_token_count[0] = ( - 0 if self._bookkeeping_no_real_work else self.batch_dimensions.token_count - ) - - # Coalesced H2D: one copy for the entire bookkeeping buffer. + # Coalesced H2D: one cudaMemcpyAsync for the entire bookkeeping buffer. # Copying the whole (max_tokens + max_requests)-sized buffer including # unused slots is cheap (~71 KB total, ~3-5 us on PCIe Gen4) and saves - # redundant launch overheads vs. per-field copies. Async scheduling - # decode steps skip token_to_input_ids here because sampled tokens are - # already GPU-resident and copied directly into the GPU input buffer. - if skip_token_input_ids: - token_to_input_ids_offset = ( - self.token_to_input_ids.numel() * self.token_to_input_ids.element_size() - ) - else: - token_to_input_ids_offset = 0 - # Only event-tracked callers may leave the copy in flight; legacy callers - # block before the pinned CPU source can be re-staged. - self.gpu_view._buf[token_to_input_ids_offset:].copy_( - self._cpu_bookkeeping_buf[token_to_input_ids_offset:], non_blocking=record_done_event - ) + # 8 redundant launch overheads vs. the prior per-field copies. + self.gpu_view._buf.copy_(self._cpu_bookkeeping_buf, non_blocking=True) # MHA metadata GPU views were already bound to state_data in # initialize_attention_state(); the H2D above populates the underlying @@ -2811,54 +2389,6 @@ def transfer_bookkeeping_to_gpu( self.mamba_metadata.load_from_cpu(self._pending_mamba_transfer) self._pending_mamba_transfer = None - done_event = None - if record_done_event: - done_event = self._bookkeeping_h2d_done_event - done_event.record(torch.cuda.current_stream()) - - return done_event - - def copy_async_sched_sample_to_forward( - self, sampled_tokens_cuda: Tensor, sampled_mtp_tokens_cuda: Optional[Tensor] = None - ) -> None: - """Populate GPU input token IDs from sampled CUDA tokens for async scheduling. - - Async scheduling keeps sampled tokens GPU-resident for the next decode - forward. CPU bookkeeping is prepared independently and published later; - this direct GPU copy populates the live input-ID view without waiting - for the sample's CPU copy. - - Args: - sampled_tokens_cuda (Tensor): 1D CUDA tensor containing one sampled - token per active decode request. - sampled_mtp_tokens_cuda (Optional[Tensor]): MTP draft tokens with shape - ``[num_speculative_tokens, active_request_count]``. - """ - active_request_count = self.total_request_count - self.paused_request_count - - if self.num_speculative_tokens > 0: - expected_shape = (self.num_speculative_tokens, active_request_count) - if sampled_mtp_tokens_cuda is None or tuple(sampled_mtp_tokens_cuda.shape) != ( - expected_shape - ): - actual_shape = ( - None if sampled_mtp_tokens_cuda is None else sampled_mtp_tokens_cuda.shape - ) - raise RuntimeError( - f"Expected MTP draft token shape {expected_shape}, got {actual_shape}." - ) - - tokens_per_request = self.num_speculative_tokens + 1 - token_count = active_request_count * tokens_per_request - grouped_tokens = self.gpu_view.token_to_input_ids[:token_count].view( - active_request_count, tokens_per_request - ) - grouped_tokens[:, 0].copy_(sampled_tokens_cuda, non_blocking=True) - if sampled_mtp_tokens_cuda is not None: - grouped_tokens[:, 1:].copy_(sampled_mtp_tokens_cuda.transpose(0, 1), non_blocking=True) - if token_count < self.padded_active_token_count: - self.gpu_view.token_to_input_ids[token_count : self.padded_active_token_count].zero_() - def reset_tensors(self) -> None: """Fill all bookkeeping tensors with sentinel values.""" @@ -2885,34 +2415,19 @@ def reset_tensors(self) -> None: self.token_to_block_idx.fill_(-1) self.token_to_local_position_within_kv_block.fill_(0) - def reset_metadata( - self, preserve_prefix_cache: bool = False, *, preserve_counters: bool = False - ) -> None: + def reset_metadata(self) -> None: """Reset all bookkeeping state: counters, block allocator, attention/mamba state. This must be called after ``initialize_all_tensors()`` and after any suspend/resume cycle to bring the context back to a clean state. - - Args: - preserve_prefix_cache: When True, keep the KV block allocator's prefix-cache - state (hash index, ref counts, cached blocks) intact. Used by the idle - ``dummy_forward`` path, which only needs to clear the transient one-token - step state -- wiping the allocator there would destroy cross-request prefix - reuse for any subsequent request (the engine idles between requests at low - concurrency, especially with EP > 1). - preserve_counters: When True, keep engine-step, prefix-cache clock, - prefill-token, and async-scheduling counters intact. """ - # There is no prefix-cache state to preserve when caching is disabled. - preserve_prefix_cache = preserve_prefix_cache and self.enable_prefix_caching # Reset request/token counts. self.total_request_count = 0 self.active_token_count = 0 - if not preserve_counters: - self.lifetime_prefill_token_count = 0 - self.async_sched_step_count = 0 - self.async_sched_compaction_step_count = 0 + self.lifetime_prefill_token_count = 0 + self.async_sched_step_count = 0 + self.async_sched_compaction_step_count = 0 self.paused_request_count = 0 self.batch_dimensions = InferenceBatchDimensions( token_count=0, prefill_req_count=0, decode_req_count=0 @@ -2928,8 +2443,7 @@ def reset_metadata( # Reset attention, mamba, and block allocator state. self.reset_attention_state() self.reset_mamba_state() - if not preserve_prefix_cache: - self.kv_block_allocator.reset() + self.kv_block_allocator.reset() self.request_to_kv_block_ids.fill_(-1) # Reset chunked prefill state @@ -2941,15 +2455,7 @@ def reset_metadata( token_count=0, prefill_req_count=0, decode_req_count=0 ) - # Reset VLM data. - # Guarded because tests can construct via __new__ and skip __init__. - getattr(self, "_request_to_image_embeddings", {}).clear() - getattr(self, "_request_to_image_token_mask", {}).clear() - getattr(self, "_request_to_image_token_count", {}).clear() - - def reset( - self, preserve_prefix_cache: bool = False, *, preserve_counters: bool = False - ) -> None: + def reset(self) -> None: """Reset entire context. This method does: @@ -2960,32 +2466,18 @@ def reset( This method is useful after cuda graph warmup iterations, where the context's memory buffer is referenced by the cuda graph system and cannot be deallocated. - - Args: - preserve_prefix_cache: When True, keep the KV and Mamba prefix-cache - state (hash indices and cached blocks/slots) intact. Used by - the idle ``dummy_forward`` path so an idle step between requests does - not destroy cross-request prefix reuse. - preserve_counters: When True, keep engine-step, prefix-cache clock, - prefill-token, and async-scheduling counters intact. """ - # There is no prefix-cache state to preserve when caching is disabled. - preserve_prefix_cache = preserve_prefix_cache and self.enable_prefix_caching - self.dynamo_helper.discard_pending_kv_stored_events() self.reset_tensors() - self.reset_metadata( - preserve_prefix_cache=preserve_prefix_cache, preserve_counters=preserve_counters - ) + self.reset_metadata() - if not preserve_counters: - self.step_count = 0 - self.prefix_cache_lru_clock = 0 + # Reset lifetime counters (not reset in reset_metadata, which is also + # called during suspend/resume where these must persist). + self.step_count = 0 + self.prefix_cache_lru_clock = 0 - # Reset Mamba cache state. - if not preserve_prefix_cache and self.mamba_slot_allocator is not None: + # Reset Mamba cache state + if self.mamba_slot_allocator is not None: self.mamba_slot_allocator.reset() - if not preserve_prefix_cache: - self.dynamo_helper.notify_kv_cache_cleared() def current_input_and_position_ids( self, *, num_warmup_tokens: Optional[int] = None @@ -3074,46 +2566,13 @@ def last_token_logits(self, logits: Tensor) -> Tensor: ) return logits.squeeze(0)[self.active_logit_idxs[: self.num_last_token_logits], :] - def _find_mamba_match_count( - self, req: DynamicInferenceRequest, start_block: int, end_block: int - ) -> int: - """Find the farthest cached Mamba state within a chunk-local block range. - - Mamba state restore is only valid for blocks that the current chunk also - assigns from the KV cache. Chunked prefill can schedule a prompt prefix - that is shorter than the farthest cached full-prompt Mamba boundary, so - this helper intentionally uses the same block domain as KV matching. - """ - if self.mamba_slot_allocator is None or not req.precomputed_block_hashes: - return 0 - - end_block = min(end_block, len(req.precomputed_block_hashes)) - if start_block >= end_block: - return 0 - - mamba_map = self.mamba_slot_allocator.hash_to_block_id - hashes = req.precomputed_block_hashes[start_block:end_block] - for i in range(len(hashes) - 1, -1, -1): - if hashes[i] in mamba_map: - return i + 1 - return 0 - def _compute_prefix_match( - self, - req: DynamicInferenceRequest, - prefill_chunk_length: int, - record_mamba_match: bool = False, + self, req: DynamicInferenceRequest, prefill_chunk_length: int ) -> Tuple[list, int, int, int, int, int]: """Compute prefix match results and skip counts for a request chunk. Shared by check_availability (budget checks) and add_request (execution). - Args: - req: Request being scheduled. - prefill_chunk_length: Number of prompt tokens considered in this chunk. - record_mamba_match: If True, store the chunk-local executable Mamba - match count on the request for diagnostics/tests. - Returns: Tuple of (matched_block_ids, num_blocks_from_pool, already_allocated_blocks, overall_required_blocks, @@ -3152,11 +2611,7 @@ def _compute_prefix_match( # Only applies to the first chunk (finished == 0); continuation chunks # already had Mamba state restored during the first chunk. if self.is_hybrid_model and self.mamba_slot_allocator is not None and finished == 0: - num_mamba_matched = self._find_mamba_match_count( - req, already_allocated_blocks, already_allocated_blocks + num_matched - ) - if record_mamba_match: - req._mamba_num_matched_blocks = num_mamba_matched + num_mamba_matched = getattr(req, '_mamba_num_matched_blocks', 0) assert ( num_mamba_matched <= num_matched ), f"Mamba match ({num_mamba_matched}) > KV match ({num_matched})" @@ -3176,8 +2631,6 @@ def _compute_prefix_match( else: prefix_skip_tokens = 0 elif self.is_hybrid_model and finished == 0: - if record_mamba_match: - req._mamba_num_matched_blocks = 0 prefix_skip_tokens = 0 # Clamp so that effective_prefill_chunk_length >= 2 when possible. @@ -3207,36 +2660,19 @@ def check_availability(self, req: DynamicInferenceRequest) -> Tuple[bool, bool, Check if the request can be added to the context. """ # Note that for hybrid models checking the total request count is sufficient - # because we allocate a single set of Mamba state tensors for each request. + # because we allocate a single set of Mamba state tensors for each request request_can_be_added = ( self.total_request_count < self.max_requests and self.paused_request_count == 0 ) - # A handoff is the exception: it keeps its live recurrent-state slot after - # its request row is removed, so both capacities must be checked independently. - if self.is_hybrid_model and self.kv_block_allocator.enable_handoff_pinning: - request_can_be_added &= self.mamba_metadata.mamba_state_free_slot_count > 0 - matched_block_ids, num_blocks_from_pool, _, _, _, effective_prefill_chunk_length = ( + _, num_blocks_from_pool, _, _, _, effective_prefill_chunk_length = ( self._compute_prefix_match(req, req.remaining_prompt_length) ) request_tokens_can_be_added = ( self.active_token_count + effective_prefill_chunk_length <= self.max_tokens ) - # add_request pins the matched blocks before allocating. Only matches that - # are currently evictable (ref_count == 0) count against the evictable - # pool; matches already pinned by another in-flight request are not in - # get_evictable_block_count() and pinning them frees nothing. Reserve only - # the ref_count == 0 matches so availability is not under-reported. - potential_matched_count = 0 - if matched_block_ids: - matched_tensor = torch.tensor(matched_block_ids, dtype=torch.int32, device='cpu') - potential_matched_count = int( - (self.kv_block_allocator.block_ref_counts[matched_tensor] == 0).sum() - ) - kv_cache_available = self.kv_block_allocator.is_memory_available( - num_blocks_from_pool, potential_matched_count=potential_matched_count - ) + kv_cache_available = self.kv_block_allocator.is_memory_available(num_blocks_from_pool) return request_can_be_added, request_tokens_can_be_added, kv_cache_available def _find_kv_match_count( @@ -3321,47 +2757,31 @@ def add_request( overall_required_blocks, prefix_skip_tokens, effective_prefill_chunk_length, - ) = self._compute_prefix_match(req, prefill_chunk_length, record_mamba_match=True) + ) = self._compute_prefix_match(req, prefill_chunk_length) num_matched_blocks = len(matched_block_ids) effective_kv_offset = req.finished_chunk_token_count + prefix_skip_tokens - # Track prefix cache hits. num_cached_tokens accumulates across prefill - # chunks: each chunk matches a disjoint block range (start advances with - # finished_chunk_token_count), so a long cached prefix is discovered - # incrementally and must be summed, not overwritten. + # Track prefix cache hits. if num_matched_blocks > 0: self.prefix_cache_hits += 1 self.prefix_cache_blocks_matched += num_matched_blocks - req.num_cached_tokens += num_matched_blocks * self.block_size_tokens # Slice tokens to skip matched prefix this_round_tokens = req.remaining_prompt_tokens[prefix_skip_tokens:prefill_chunk_length] - # Pin matched (shared) blocks BEFORE allocation. allocate_memory_blocks() - # may trigger LRU eviction, and a matched block still at ref_count == 0 - # would be an eviction candidate — descendant-first LRU could evict a - # matched leaf and immediately reuse its ID for a new block, leaving the - # block table with a duplicate ID and a dangling parent. Incrementing ref - # counts first removes matched blocks from the evictable set (see - # evict_lru_blocks / get_evictable_block_count), so eviction falls back to - # genuinely unused cached blocks. - matched_tensor = None - if num_matched_blocks > 0: - matched_tensor = torch.tensor(matched_block_ids, dtype=torch.int32, device='cpu') - self.kv_block_allocator.block_ref_counts[matched_tensor] += 1 - if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU: - self.kv_block_allocator.update_timestamps(matched_tensor) - new_block_ids = None if num_blocks_from_pool > 0: new_block_ids = self.kv_block_allocator.allocate_memory_blocks(num_blocks_from_pool) if new_block_ids is None or len(new_block_ids) != num_blocks_from_pool: - # Roll back the pin so a failed add does not leak ref counts on - # the matched blocks (which would make them permanently unevictable). - if matched_tensor is not None: - self.kv_block_allocator.block_ref_counts[matched_tensor] -= 1 raise BlockOverflowError(req.request_id) + # Increment ref counts and update timestamps for matched (shared) blocks + if num_matched_blocks > 0: + matched_tensor = torch.tensor(matched_block_ids, dtype=torch.int32, device='cpu') + self.kv_block_allocator.block_ref_counts[matched_tensor] += 1 + if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU: + self.kv_block_allocator.update_timestamps(matched_tensor) + # Note that we decremented the total_request_count for the chunked prefill request # in update_requests, so setting current_id to the total_request_count will again # make the last request the continuing chunked prefill request if one exists. @@ -3443,63 +2863,6 @@ def add_request( self.token_to_block_idx[ self.active_token_count : self.active_token_count + effective_prefill_chunk_length ] = self.request_to_kv_block_ids[current_id][token_offset_range // self.block_size_tokens] - if num_matched_blocks > 0: - # Some tokens we are about to compute may land inside a block we matched - # by hash. That block already holds the correct KV for exactly these - # tokens and is shared with whoever cached it, so send those writes to the - # dummy block rather than perturb a concurrent reader's values. - # - # Only the write mapping moves. `request_to_kv_block_ids` still points at - # the real block, so attention reads the cached KV through the block table, - # and `token_to_block_idx` is rebuilt every step so nothing needs restoring. - # - # All positions below are absolute token positions within the request. - # A chunk resuming mid-block, with one fresh block past the match: - # - # token 0 B 2B 3B 4B - # blocks |--0--|--1--|--2--|--3--| - # matched [=================] blocks 0-2 were hash-matched - # chunk [==============] tokens computed this step - # redirect [========] the overlap -> dummy block - # write [=====] the rest -> real block 3 - - # 1. The blocks matched here sit right after the blocks this request - # already owns. The partial block carried over from the previous - # chunk joins them when it was itself matched earlier -- the tokens - # completing it are cached too, and it is still shared with whoever - # cached it. Everything before that is a prefix of matched blocks, so - # starting the span at 0 is safe; the intersection below clips it. - matched_prefix_start_token = ( - 0 - if req.num_matched_prefix_blocks >= already_allocated_blocks - else already_allocated_blocks * self.block_size_tokens - ) - matched_prefix_end_token = ( - already_allocated_blocks + num_matched_blocks - ) * self.block_size_tokens - - # 2. This chunk computes the tokens starting at `effective_kv_offset` - # (the prompt offset left over after skipping the cached prefix). - chunk_start_token = effective_kv_offset - chunk_end_token = effective_kv_offset + effective_prefill_chunk_length - - # 3. The redundant tokens are where those two spans overlap. Non-empty - # whenever we matched more blocks than we skipped tokens for: the - # ">= 2 computed tokens" clamp, the Mamba back-off, or memory-only - # hybrid mode where nothing is skipped but blocks are still shared. - overlap_start_token = max(matched_prefix_start_token, chunk_start_token) - overlap_end_token = min(matched_prefix_end_token, chunk_end_token) - - # 4. Rebase the overlap onto this chunk's slots and redirect it. The - # overlap starts past the chunk start only when this chunk's own - # matched blocks begin later than the tokens it computes, i.e. the - # chunk opens inside a freshly allocated block that it must write. - if overlap_end_token > overlap_start_token: - redirect_start = self.active_token_count + overlap_start_token - chunk_start_token - redirect_end = self.active_token_count + overlap_end_token - chunk_start_token - self.token_to_block_idx[redirect_start:redirect_end] = ( - self.kv_block_allocator.dummy_block_idx - ) self.token_to_local_position_within_kv_block[ self.active_token_count : self.active_token_count + effective_prefill_chunk_length ] = (token_offset_range % self.block_size_tokens) @@ -3520,29 +2883,9 @@ def _register_range(start: int, end: int): return block_ids_to_hash = self.request_to_kv_block_ids[current_id][start:end].tolist() block_hashes_slice = req.precomputed_block_hashes[start:end] - # Parent hash of block k is the hash of block k-1 in the chain; - # block 0 is a root (parent hash 0). Enables LRU eviction to keep - # parents cached until their children are gone. - parent_hashes_slice = [ - req.precomputed_block_hashes[k - 1] if k > 0 else 0 for k in range(start, end) - ] self.kv_block_allocator.register_kv_block_hashes( - block_ids_to_hash, block_hashes_slice, parent_hashes_slice + block_ids_to_hash, block_hashes_slice ) - if self.dynamo_helper.has_kv_event_listeners: - token_start = start * self.block_size_tokens - token_end = end * self.block_size_tokens - token_ids = req.prompt_tokens[token_start:token_end].tolist() - self.dynamo_helper.queue_kv_stored_event( - { - "block_hashes": list(block_hashes_slice), - "token_ids": token_ids, - "num_block_tokens": [self.block_size_tokens] * (end - start), - "parent_hash": ( - int(req.precomputed_block_hashes[start - 1]) if start > 0 else None - ), - } - ) # Range 1: prior-chunk partial block that this chunk just completed _register_range(previously_complete, min(already_allocated_blocks, num_complete_blocks)) @@ -3566,35 +2909,23 @@ def _register_range(start: int, end: int): else: self._pending_mamba_zeros.append(mamba_idx) - # compute_and_store_offsets sets CPU state + GPU staging buffers that - # commit_intermediate_states() consumes after the forward pass. Run it for - # EVERY prefill chunk (not just the first): the last complete block of a - # multi-chunk prompt falls in a continuation chunk, and caching its Mamba - # state is precisely what lets a later turn skip prefill on a hybrid model. - # Mamba slot allocation / state restore above stays first-chunk-only. - if self.is_hybrid_model and self.mamba_slot_allocator is not None: - self.mamba_slot_allocator.compute_and_store_offsets( - req, - current_id, - prefix_skip_tokens, - prefill_chunk_length, - num_matched_blocks, - matched_block_ids, - overall_required_blocks, - ) - - # Extend the request's run of matched blocks, so a later chunk knows not to - # rewrite the partial block it inherits from this one. Only extend when this - # chunk's matches continue the run: a gap means the blocks in between were - # computed by this request and the tokens completing them must be written. - if num_matched_blocks > 0 and req.num_matched_prefix_blocks >= already_allocated_blocks: - req.num_matched_prefix_blocks = already_allocated_blocks + num_matched_blocks + # compute_and_store_offsets sets both CPU state (hash_to_block_id, + # _eos_cache_block_id_gpu) and GPU staging buffers. Runs immediately + # because commit_intermediate_states() reads the CPU state after the + # forward pass. + if self.mamba_slot_allocator is not None: + self.mamba_slot_allocator.compute_and_store_offsets( + req, + current_id, + prefix_skip_tokens, + prefill_chunk_length, + num_matched_blocks, + matched_block_ids, + overall_required_blocks, + ) self.active_token_count += effective_prefill_chunk_length self.lifetime_prefill_token_count += effective_prefill_chunk_length - if self.enable_prefix_caching: - self.prefix_cache_prefill_computed_tokens += effective_prefill_chunk_length - self.prefix_cache_prefill_skipped_tokens += prefix_skip_tokens self.total_request_count += 1 self.num_prefill_requests += 1 @@ -3713,117 +3044,53 @@ def release_memory_blocks_from_request_indexes(self, request_indexes) -> None: sa._intermediate_block_ids_cpu[request_indexes] = -1 sa._eos_cache_block_id_cpu[request_indexes] = -1 - def _get_paused_request_count_within_block_budget(self) -> int: - """Count the left-most paused requests whose blocks fit the paused budget.""" - block_budget = self.kv_block_allocator.paused_limit - if self.paused_request_count == 0 or block_budget == 0: - return 0 - - if not self.enable_prefix_caching: - paused_block_counts = self.request_kv_block_counts[: self.paused_request_count] - cumulative_block_counts = paused_block_counts.cumsum(dim=0) - return int(torch.count_nonzero(cumulative_block_counts <= block_budget).item()) - - seen_block_ids: set[int] = set() - retained_request_count = 0 - for request_idx in range(self.paused_request_count): - block_ids = self.request_to_kv_block_ids[request_idx] - seen_block_ids.update(block_ids[block_ids >= 0].tolist()) - if len(seen_block_ids) > block_budget: - break - retained_request_count += 1 - return retained_request_count - - def _get_releasable_block_counts( - self, request_start_idx: int, request_end_idx: int - ) -> list[int]: - """Count blocks made allocatable by each right-most request suffix. - - The returned list has one entry for every possible suffix length, including - zero. For prefix caching, a shared block is credited only once all selected - request references account for its current allocator reference count. - """ - suffix_request_count = request_end_idx - request_start_idx - releasable_counts = [0] * (suffix_request_count + 1) - if suffix_request_count == 0: - return releasable_counts - - if not self.enable_prefix_caching: - suffix_block_counts = ( - self.request_kv_block_counts[request_start_idx:request_end_idx] - .flip(dims=[0]) - .cumsum(dim=0) - ) - return [0, *suffix_block_counts.tolist()] - - block_rows = self.request_to_kv_block_ids[request_start_idx:request_end_idx] - request_offsets, block_offsets = torch.where(block_rows >= 0) - if request_offsets.numel() == 0: - return releasable_counts - - block_ids = block_rows[request_offsets, block_offsets].long() - unique_block_ids, inverse, selected_reference_counts = torch.unique( - block_ids, return_inverse=True, return_counts=True - ) - current_reference_counts = self.kv_block_allocator.block_ref_counts[unique_block_ids] - assert torch.all( - selected_reference_counts <= current_reference_counts - ), "selected more KV block references than the allocator owns" - - # A fully selected block becomes releasable when the suffix reaches its - # left-most reference. - suffix_depths = suffix_request_count - request_offsets - release_suffix_depths = torch.zeros_like(unique_block_ids) - release_suffix_depths.scatter_reduce_( - 0, inverse, suffix_depths, reduce="amax", include_self=True - ) - fully_selected = selected_reference_counts == current_reference_counts - releases_by_suffix = torch.bincount( - release_suffix_depths[fully_selected], minlength=suffix_request_count + 1 - ) - return releases_by_suffix.cumsum(dim=0).tolist() - def resume_paused_requests( - self, active_request_count: int, newly_paused_request_ids: Optional[torch.Tensor] - ) -> tuple[int, Optional[torch.Tensor]]: - """Resume as many paused requests as compute and KV capacity permit. + self, active_request_count: int, newly_paused_request_ids: torch.Tensor + ) -> tuple[int, torch.Tensor]: + """Resume as many paused requests as we have space for in the active buffer. Args: active_request_count (int): Number of active requests. - newly_paused_request_ids (Optional[torch.Tensor]): Newly paused request ids. + newly_paused_request_ids (torch.Tensor): List of newly paused request ids. + next_tokens (torch.Tensor): Sampled tokens. Returns: - (tuple[int, Optional[torch.Tensor]]) Updated active count and newly paused ids. + (tuple[int, torch.Tensor]) active_request_count, newly_paused_request_ids. """ # Assign released blocks to paused requests. # todo: @shanmugamr, un-pause requests using FIFO, rather than LIFO. resume_request_count = 0 if self.paused_request_count > 0: - # Check which paused requests will actually need a new block upon - # resuming. Flip before cumsum because requests resume from the right. + active_block_count_avail = self.kv_block_allocator.get_active_avail() + # Clone not needed: flip() makes a copy. + paused_block_counts = self.request_kv_block_counts[: self.paused_request_count] + # Flip counts before cumsum, since paused requests are resumed from + # the right-most index, so we must count resumed blocks starting from + # the right side. + paused_block_counts = paused_block_counts.flip(dims=[0]) + + # Check which paused requests will actually need a new block upon resuming offsets = self.request_last_kv_block_offset[: self.paused_request_count] needs_new_block = ( offsets >= self.block_size_tokens - 1 - self.num_speculative_tokens - ).to(self.request_kv_block_counts.dtype) + ).to(paused_block_counts.dtype) needs_new_block = needs_new_block.flip(dims=[0]) + # Add +1 ONLY to the block counts of requests that finished their previous memory block + paused_block_counts += needs_new_block + paused_block_counts_cumsum = paused_block_counts.cumsum(dim=0) + resume_request_count = min( + torch.nonzero(paused_block_counts_cumsum <= active_block_count_avail).numel(), + self.kv_block_allocator.total_avail, + ) + # Constrain resumptions by the maximum allowed active requests and tokens max_allowed_active = min( self.max_requests, self.max_tokens // (self.num_speculative_tokens + 1) ) allowed_to_resume = max(0, max_allowed_active - active_request_count) - candidate_count = min(self.paused_request_count, allowed_to_resume) - - if candidate_count > 0: - new_block_counts_cumsum = needs_new_block[:candidate_count].cumsum(dim=0) - max_new_block_count = int(new_block_counts_cumsum[-1].item()) - block_count_avail = self.kv_block_allocator.pool_avail - if max_new_block_count > block_count_avail: - block_count_avail = self.kv_block_allocator.get_allocatable_count() - resume_request_count = int( - torch.count_nonzero(new_block_counts_cumsum <= block_count_avail).item() - ) + resume_request_count = min(resume_request_count, allowed_to_resume) self.paused_request_count -= resume_request_count active_request_count += resume_request_count @@ -3839,10 +3106,8 @@ def resume_paused_requests( num_new_blocks = needs_new_block.sum().item() if num_new_blocks > 0: + assert num_new_blocks <= self.kv_block_allocator.total_avail block_ids = self.kv_block_allocator.allocate_memory_blocks(num_new_blocks) - assert ( - block_ids is not None and block_ids.numel() == num_new_blocks - ), f"failed to allocate {num_new_blocks} blocks for resumed requests" # Apply updates only to the requests that required a new block relative_row_idx = torch.nonzero(needs_new_block).squeeze(1) @@ -3868,66 +3133,54 @@ def evict_overflow_paused_requests( active_request_count: int, next_tokens: torch.Tensor, new_speculative_tokens: Optional[torch.Tensor] = None, - ) -> Optional[torch.Tensor]: - """Evict requests that overflow the paused-block retention budget. + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """Evict requests that overflow the paused buffer. Args: active_request_count (int): Number of active requests. next_tokens (torch.Tensor): Sampled tokens. - new_speculative_tokens (Optional[torch.Tensor]): Speculative tokens. Returns: - (Optional[torch.Tensor]) Evicted request ids. + (torch.Tensor) Evicted request ids. """ - # Keep the largest left-most paused prefix that fits the configured - # paused-block retention budget. Requests to its right must either resume - # or be evicted. - retained_paused_request_count = self._get_paused_request_count_within_block_budget() - overflow_paused_request_count = self.paused_request_count - retained_paused_request_count - if overflow_paused_request_count == 0: - return None - - max_allowed_active = min( - self.max_requests, self.max_tokens // (self.num_speculative_tokens + 1) + # Overflow paused block count. + overflow_paused_block_count = ( + self.kv_block_allocator.get_paused_used() - self.kv_block_allocator.paused_count ) - allowed_to_resume = max(0, max_allowed_active - active_request_count) - needs_new_block = ( - self.request_last_kv_block_offset[: self.paused_request_count] - >= self.block_size_tokens - 1 - self.num_speculative_tokens - ) - overflow_needs_new_block = needs_new_block[ - retained_paused_request_count : self.paused_request_count - ] - releasable_block_counts = self._get_releasable_block_counts( - retained_paused_request_count, self.paused_request_count - ) - new_block_counts_by_survivor_count = overflow_needs_new_block.cumsum(dim=0) - new_block_counts_by_survivor_count = torch.cat( - (new_block_counts_by_survivor_count.new_zeros(1), new_block_counts_by_survivor_count) - ) - block_count_avail = self.kv_block_allocator.get_allocatable_count() - - # Preserve current ordering: evict the smallest right-most suffix whose - # physical releases let every remaining overflow request resume. - evict_request_count = None - for candidate_evict_count in range(overflow_paused_request_count + 1): - survivor_count = overflow_paused_request_count - candidate_evict_count - if survivor_count > allowed_to_resume: - continue - survivor_new_block_count = int( - new_block_counts_by_survivor_count[survivor_count].item() - ) - candidate_avail = block_count_avail + releasable_block_counts[candidate_evict_count] - if survivor_new_block_count <= candidate_avail: - evict_request_count = candidate_evict_count - break + # Nothing to evict? + if overflow_paused_block_count <= 0: + return None - assert evict_request_count is not None - if evict_request_count == 0: + # Overflow paused block count. + paused_block_counts = self.request_kv_block_counts[: self.paused_request_count] + paused_block_counts_cumsum = paused_block_counts.cumsum(dim=0) + valid_paused_request_count = torch.nonzero( + paused_block_counts_cumsum <= self.kv_block_allocator.paused_count + ).numel() + overflow_paused_request_count = self.paused_request_count - valid_paused_request_count + + # Nothing to evict? (Similar to checking overflow_paused_block_count + # above, but here we allow up to one paused request to overflow into the + # active buffer. + if overflow_paused_request_count == 0: return None + # Evict request count. (Flip paused_block_counts because evictions are + # counted from the right-most paused requests. + paused_block_counts = paused_block_counts[-overflow_paused_request_count:].flip(dims=[0]) + paused_block_counts_cumsum = paused_block_counts.cumsum(dim=0) + remaining_paused_request_counts = torch.arange( + overflow_paused_request_count - 1, + -1, + -1, + dtype=paused_block_counts_cumsum.dtype, + device='cpu', + ) + net_block_counts = paused_block_counts_cumsum - remaining_paused_request_counts + evict_request_count = torch.nonzero(net_block_counts >= 0)[0].item() + 1 + # Eviction index range. evict_start_idx = self.paused_request_count - evict_request_count evict_end_idx = self.paused_request_count @@ -3988,184 +3241,97 @@ def evict_overflow_paused_requests( return evict_request_ids - def _get_async_sched_rows_requiring_new_block(self) -> Tensor: - """Return active request rows that need a block during the next prepare. - - Returns: - Tensor: Boolean mask over active request rows. - """ - active_slice = slice(self.paused_request_count, self.total_request_count) - tokens_per_request = self.num_speculative_tokens + 1 - return ( - self.request_last_kv_block_offset[active_slice] + tokens_per_request - >= self.block_size_tokens - ) - - def can_prepare_requests(self) -> bool: - """Return whether requests can be prepared without lifecycle changes. - - Returns: - bool: Whether all requests are active decode requests and the shared - KV-block pool can satisfy the exact next-step allocation demand. - """ - if self.num_prefill_requests != 0 or self.paused_request_count != 0: - return False - - rows_requiring_new_block = self._get_async_sched_rows_requiring_new_block() - num_new_blocks = int(rows_requiring_new_block.sum().item()) - return num_new_blocks <= self.kv_block_allocator.get_allocatable_count() - - def prepare_requests(self) -> None: + def prepare_requests(self, new_tokens: Tensor) -> None: """Speculatively prepare active decode requests for the next forward pass. Async scheduling only supports decode-only steps with no pause, - evict, or resume lifecycle changes. If preparation cannot allocate the - required KV blocks without a lifecycle change, this method raises. The - prepared decode layout establishes the active token count. + evict, or resume lifecycle changes. If preparing the next token would + require one of those lifecycle changes, this method raises and the caller + should treat async scheduling as unsupported for that workload. + + Args: + new_tokens (Tensor): Newly sampled token for each active request. """ + if new_tokens.is_cuda: + new_tokens = new_tokens.cpu() + active_request_count = self.total_request_count - self.paused_request_count + if self.num_speculative_tokens != 0: + raise RuntimeError("Async scheduling does not support speculative tokens.") if self.num_prefill_requests != 0: raise RuntimeError("Async scheduling only supports decode-only steps.") if self.paused_request_count != 0: raise RuntimeError("Async scheduling does not support paused requests.") + if new_tokens.numel() != active_request_count: + raise RuntimeError( + f"Expected {active_request_count} new tokens, got {new_tokens.numel()}." + ) if active_request_count == 0: self.active_token_count = 0 return active_slice = slice(0, active_request_count) - tokens_per_request = self.num_speculative_tokens + 1 - last_block_offsets = self.request_last_kv_block_offset[active_slice] - token_offsets = self._async_sched_token_offsets - rows_requiring_new_block = self._get_async_sched_rows_requiring_new_block() - num_new_blocks = int(rows_requiring_new_block.sum().item()) - - block_ids = None + rows_requiring_new_block = ( + self.request_last_kv_block_offset[active_slice] >= self.block_size_tokens - 1 + ) + num_new_blocks = rows_requiring_new_block.sum().item() if num_new_blocks > 0: - if num_new_blocks > self.kv_block_allocator.get_allocatable_count(): + active_block_count_avail = self.kv_block_allocator.get_active_avail() + if num_new_blocks > active_block_count_avail: raise RuntimeError("Async scheduling cannot pause requests to allocate new blocks.") block_ids = self.kv_block_allocator.allocate_memory_blocks(num_new_blocks) if block_ids is None: raise RuntimeError("Async scheduling cannot evict requests to allocate new blocks.") - self.active_token_count = active_request_count * tokens_per_request - active_token_slice = slice(0, self.active_token_count) - grouped_token_block_ids = self.token_to_block_idx[active_token_slice].view( - active_request_count, tokens_per_request - ) - grouped_token_block_ids.copy_(self.request_last_kv_block_id[active_slice, None]) - - if block_ids is not None: row_idx = torch.nonzero(rows_requiring_new_block, as_tuple=True)[0] col_idx = self.request_kv_block_counts[row_idx] self.request_to_kv_block_ids[row_idx, col_idx] = block_ids self.request_kv_block_counts[row_idx] += 1 self.request_last_kv_block_id[row_idx] = block_ids - grouped_token_block_ids[row_idx] = torch.where( - last_block_offsets[row_idx, None] + 1 + token_offsets[None, :] - >= self.block_size_tokens, - block_ids[:, None], - grouped_token_block_ids[row_idx], - ) self.request_kv_length_offsets[active_slice].add_(self.request_query_lengths[active_slice]) - self.request_query_lengths[active_slice].fill_(tokens_per_request) + self.request_query_lengths[active_slice].fill_(1) + self.request_last_kv_block_offset[active_slice] = ( - last_block_offsets + tokens_per_request + self.request_last_kv_block_offset[active_slice] + 1 ) % self.block_size_tokens - token_positions = ( - self.request_kv_length_offsets[active_slice, None] + token_offsets[None, :] - ) - token_request_idxs = torch.arange(active_request_count, device='cpu').repeat_interleave( - tokens_per_request + self.active_token_count = active_request_count + self.token_to_input_ids[:active_request_count] = new_tokens + self.token_to_pos_ids[:active_request_count] = self.request_kv_length_offsets[active_slice] + self.token_to_request_idx[:active_request_count] = torch.arange( + active_request_count, device='cpu' ) - self.token_to_pos_ids[active_token_slice] = token_positions.flatten() - self.token_to_request_idx[active_token_slice] = token_request_idxs - self.token_to_position_in_request[active_token_slice] = self.token_to_pos_ids[ - active_token_slice + self.token_to_position_in_request[:active_request_count] = self.token_to_pos_ids[ + :active_request_count ] - self.token_to_local_position_within_kv_block[active_token_slice] = ( - self.token_to_pos_ids[active_token_slice] % self.block_size_tokens - ) - - def commit_sampled_tokens( - self, sampled_tokens_cpu: Tensor, sampled_mtp_tokens_cpu: Optional[Tensor] = None - ) -> None: - """Commit sampled CPU token IDs to the prepared request state. - - This establishes the post-resolution active token count and populates - the CPU input-ID staging rows in survivor order. Overlapped async - scheduling has already copied the same samples into the live GPU input - view for the speculative forward. - - Args: - sampled_tokens_cpu (Tensor): Sampled CPU token for each active request. - sampled_mtp_tokens_cpu (Optional[Tensor]): MTP draft tokens with shape - ``[num_speculative_tokens, active_request_count]``. - """ - assert sampled_tokens_cpu.device == torch.device( - 'cpu' - ), "Sampled tokens must be on the CPU before they are committed." - - if sampled_mtp_tokens_cpu is not None: - assert sampled_mtp_tokens_cpu.device == torch.device( - 'cpu' - ), "MTP draft tokens must be on the CPU before they are committed." - - active_request_count = self.total_request_count - self.paused_request_count - if sampled_tokens_cpu.numel() != active_request_count: - raise RuntimeError( - f"Expected {active_request_count} new tokens, got {sampled_tokens_cpu.numel()}." - ) - - expected_mtp_shape = (self.num_speculative_tokens, active_request_count) - if self.num_speculative_tokens == 0: - if sampled_mtp_tokens_cpu is not None and sampled_mtp_tokens_cpu.numel() != 0: - raise RuntimeError( - "Received MTP draft tokens when speculative decoding is disabled." - ) - else: - if sampled_mtp_tokens_cpu is None or tuple(sampled_mtp_tokens_cpu.shape) != ( - expected_mtp_shape - ): - actual_shape = ( - None if sampled_mtp_tokens_cpu is None else sampled_mtp_tokens_cpu.shape - ) - raise RuntimeError( - f"Expected MTP draft token shape {expected_mtp_shape}, got {actual_shape}." - ) - - tokens_per_request = self.num_speculative_tokens + 1 - active_token_count = active_request_count * tokens_per_request - self.active_token_count = active_token_count - grouped_tokens = self.token_to_input_ids[:active_token_count].view( - active_request_count, tokens_per_request + self.token_to_local_position_within_kv_block[:active_request_count] = ( + self.token_to_pos_ids[:active_request_count] % self.block_size_tokens ) - grouped_tokens[:, 0] = sampled_tokens_cpu - if sampled_mtp_tokens_cpu is not None: - grouped_tokens[:, 1:] = sampled_mtp_tokens_cpu.transpose(0, 1) + self.token_to_block_idx[:active_request_count] = self.request_last_kv_block_id[active_slice] - def resolve_requests(self, active_requests_mask: Tensor) -> Tuple[Tensor, Tensor]: + def resolve_requests(self, active_requests_mask: Tensor) -> Tensor: """Resolve finished requests after an async scheduling forward pass. - Prefill requests transition to decode during resolution. Request rows use - the same hole-filling order as ``update_requests`` so seeded sampling stays - consistent with legacy scheduling. Token tensors and the active token - count are left untouched; prepare rebuilds derived token metadata and the - controller commits sampled input IDs after resolution. + Async scheduling supports only request completion. The active request rows + and current decode-token rows are compacted in survivor order so any + following legacy or async scheduling step sees a consistent context. Args: active_requests_mask (Tensor): 1D mask marking requests that remain active. Returns: - Tuple[Tensor, Tensor]: Request IDs that finished and source row indices - for surviving requests in their resolved destination order. + Tensor: Request IDs for requests that finished during resolution. """ if active_requests_mask.is_cuda: active_requests_mask = active_requests_mask.cpu() + if self.num_speculative_tokens != 0: + raise RuntimeError("Async scheduling does not support speculative tokens.") + if self.num_prefill_requests != 0: + raise RuntimeError("Async scheduling only supports decode-only steps.") if self.paused_request_count != 0: raise RuntimeError("Async scheduling does not support paused requests.") @@ -4176,38 +3342,29 @@ def resolve_requests(self, active_requests_mask: Tensor) -> Tuple[Tensor, Tensor f"got {active_requests_mask.numel()}." ) - self.num_prefill_requests = 0 - self.request_in_prefill_status_tensor[self.request_in_prefill_status_tensor == 1] = 0 - + survivor_idxs = torch.nonzero(active_requests_mask == 1, as_tuple=True)[0] finished_idxs = torch.nonzero(active_requests_mask == 0, as_tuple=True)[0] finished_request_ids = self.request_ids[finished_idxs].clone() - active_request_count = int(active_requests_mask.sum().item()) - survivor_idxs = torch.arange(active_request_count, device='cpu') - finished_idxs_on_left = torch.nonzero( - active_requests_mask[:active_request_count] == 0, as_tuple=True - )[0] - active_idxs_on_right = ( - torch.nonzero(active_requests_mask[active_request_count:] == 1, as_tuple=True)[0] - + active_request_count - ) - assert finished_idxs_on_left.numel() == active_idxs_on_right.numel() - survivor_idxs[finished_idxs_on_left] = active_idxs_on_right - self.reset_attention_state() if finished_idxs.numel() > 0: self.release_memory_blocks_from_request_indexes(finished_idxs) + active_request_count = survivor_idxs.numel() if active_request_count == 0: self.request_to_kv_block_ids.fill_(-1) self.total_request_count = 0 + self.active_token_count = 0 self.reset_mamba_state() - return finished_request_ids, survivor_idxs + return finished_request_ids dst_idxs = torch.arange(active_request_count, device='cpu') if not torch.equal(survivor_idxs, dst_idxs): self.request_kv_length_offsets[dst_idxs] = self.request_kv_length_offsets[survivor_idxs] + self.request_in_prefill_status_tensor[dst_idxs] = self.request_in_prefill_status_tensor[ + survivor_idxs + ] self.request_query_lengths[dst_idxs] = self.request_query_lengths[survivor_idxs] self.request_output_lengths[dst_idxs] = self.request_output_lengths[survivor_idxs] self.request_ids[dst_idxs] = self.request_ids[survivor_idxs] @@ -4217,18 +3374,25 @@ def resolve_requests(self, active_requests_mask: Tensor) -> Tuple[Tensor, Tensor self.request_last_kv_block_offset[dst_idxs] = self.request_last_kv_block_offset[ survivor_idxs ] - if self.is_hybrid_model: - self.mamba_metadata.request_to_mamba_state_idx[dst_idxs] = ( - self.mamba_metadata.request_to_mamba_state_idx[survivor_idxs] - ) for metadata_tensor in self.request_metadata.values(): metadata_tensor[dst_idxs] = metadata_tensor[survivor_idxs] + + self.token_to_input_ids[dst_idxs] = self.token_to_input_ids[survivor_idxs] + self.token_to_pos_ids[dst_idxs] = self.token_to_pos_ids[survivor_idxs] + self.token_to_block_idx[dst_idxs] = self.token_to_block_idx[survivor_idxs] + self.token_to_local_position_within_kv_block[dst_idxs] = ( + self.token_to_local_position_within_kv_block[survivor_idxs] + ) + self.token_to_position_in_request[dst_idxs] = self.token_to_position_in_request[ + survivor_idxs + ] + + self.token_to_request_idx[:active_request_count] = dst_idxs stale_slice = slice(active_request_count, old_active_request_count) self.request_to_kv_block_ids[stale_slice] = -1 - if self.is_hybrid_model: - self.mamba_metadata.request_to_mamba_state_idx[stale_slice] = -1 self.total_request_count = active_request_count - return finished_request_ids, survivor_idxs + self.active_token_count = active_request_count + return finished_request_ids def update_requests( self, @@ -4478,7 +3642,7 @@ def update_requests( active_request_count, newly_paused_request_ids ) - # 6.b. Evict requests that overflow the paused-block retention budget. + # 6.b. Evict requests that overflow the paused buffer. evict_request_ids = self.evict_overflow_paused_requests( active_request_count, next_tokens, new_speculative_tokens ) @@ -4488,13 +3652,8 @@ def update_requests( active_request_count, newly_paused_request_ids ) - assert ( - active_request_count > 0 - or self.total_request_count == 0 - or self.chunked_prefill_request_id != -1 - ), "active_request_count == %d with %d paused requests and no hidden chunked prefill." % ( - active_request_count, - self.paused_request_count, + assert active_request_count > 0 or self.chunked_prefill_request_id != -1, ( + "active_request_count == %d with no hidden chunked prefill." % active_request_count ) # 6.d. Swap the chunked prefill request to the end of the active requests @@ -4698,83 +3857,25 @@ def _processed_log_probs( n_active: int, active_query_lengths: Optional[Tensor], sampling: Optional[Sampling], - row_to_request: Optional[Tensor] = None, ) -> Tensor: - """Calculate raw or sampling-processed per-row log probabilities. - - Args: - logits (Tensor): Raw logits with shape `[num_rows, vocab_size]`. - n_active (int): Number of active requests represented by the rows. - active_query_lengths (Optional[Tensor]): CPU token counts used to - map prefill rows to active requests, or `None` for decode. - sampling (Optional[Sampling]): Backend providing processed logprobs. - row_to_request (Optional[Tensor]): Explicit CPU mapping from each - logit row to an active request. - - Returns: - Tensor: Per-row log probabilities over the vocabulary. - """ + """Sample the logprobs if desired.""" if self.config.logprobs_mode == "raw_logprobs": return F.log_softmax(logits, dim=-1) assert sampling is not None, "processed_logprobs requires a sampling backend" # Map each logits row to its active request. - if row_to_request is None and active_query_lengths is not None: - row_to_request = torch.arange(n_active).repeat_interleave(active_query_lengths) - return sampling.log_probs_kernel(logits, self, token_to_request_index=row_to_request) - - def calculate_log_probs_tensors( - self, - logits: Tensor, - new_tokens: Tensor, - only_last_token_logits: Optional[bool] = False, - sampling: Optional[Sampling] = None, - row_to_request: Optional[Tensor] = None, - ) -> Tuple[Tensor, Tensor]: - """Calculate selected-token and full-distribution log probabilities. - - Args: - logits (Tensor): Raw model output logits with shape - `[1, sequence_length, vocab_size]`. - new_tokens (Tensor): Newly sampled tokens for active requests. - only_last_token_logits (Optional[bool]): Whether logits contain only - each request's final token row. - sampling (Optional[Sampling]): Sampling backend used for processed - log probabilities. - row_to_request (Optional[Tensor]): Explicit CPU mapping from each - logit row to an active request. - - Returns: - Tuple[Tensor, Tensor]: Selected-token log probabilities flattened in - active-token order and the full per-row log-probability tensor. - """ - logits_squeezed = logits.squeeze(0) - n_active = self.total_request_count - self.paused_request_count - - if only_last_token_logits or self.is_decode_only(): - seq_idx = torch.arange(len(new_tokens), dtype=torch.int32, device=logits.device) - active_logits = logits_squeezed[: len(new_tokens)].float() - log_probs = self._processed_log_probs( - active_logits, n_active, None, sampling, row_to_request - ) - return log_probs[seq_idx, new_tokens], log_probs - - logits_squeezed = logits_squeezed.float() - active_slice = slice(self.paused_request_count, self.total_request_count) - active_query_lengths_cpu = self.request_query_lengths[active_slice] - active_query_lengths_gpu = self.gpu_view.request_query_lengths[:n_active] - - # Shift away each request's first prompt token, then insert its sampled token. - active_token_ids = self.gpu_view.token_to_input_ids[: self.active_token_count].roll(-1, 0) - new_token_idx = active_query_lengths_gpu.cumsum(0) - 1 - active_token_ids[new_token_idx] = new_tokens - - log_probs = self._processed_log_probs( - logits_squeezed, n_active, active_query_lengths_cpu, sampling - ) - seq_idx = torch.arange(self.active_token_count, device=log_probs.device) - return log_probs[seq_idx, active_token_ids], log_probs + request_idx = torch.arange(n_active, device=logits.device) + row_to_request = ( + request_idx + if active_query_lengths is None + else request_idx.repeat_interleave(active_query_lengths) + ) + md = self.active_request_metadata + temperature = md["temperature"][:n_active].to(logits.device, torch.float32)[row_to_request] + top_k = md["top_k"][:n_active].to(logits.device, torch.long)[row_to_request] + top_p = md["top_p"][:n_active].to(logits.device, torch.float32)[row_to_request] + return sampling.log_probs_kernel(logits, temperature, top_k, top_p) def calculate_log_probs( self, @@ -4799,15 +3900,58 @@ def calculate_log_probs( log_probs (Tensor): Used to compute top n logprobs later if required. """ - selected_log_probs, log_probs = self.calculate_log_probs_tensors( - logits, new_tokens, only_last_token_logits=only_last_token_logits, sampling=sampling - ) + # Calculate log_probs (sequence_length x vocab_size) + logits_squeezed = logits.squeeze(0).float() + n_active = self.total_request_count - self.paused_request_count if only_last_token_logits or self.is_decode_only(): + seq_idx = torch.arange(len(new_tokens), dtype=torch.int32, device=logits.device) + log_probs = self._processed_log_probs( + logits_squeezed[seq_idx], n_active, None, sampling + ) + selected_log_probs = log_probs[seq_idx, new_tokens] return [[lp] for lp in selected_log_probs.tolist()], log_probs - active_slice = slice(self.paused_request_count, self.total_request_count) - active_query_lengths = self.request_query_lengths[active_slice] + # Get the selected token ids for all tokens. + # We shift the active token window left by one to remove the first prompt token for + # prefill requests and then set the token ids explicitly for the newly generated tokens. + # This is necessary because we calculate the log probs *before* updating the request metadata. + # + # Example (decode & prefill mix): + # + # active_query_lengths: [ 1 | 1 | 2 | 5 ] + # + # new_tokens : [ 52 | 12 | 3 | 86 ] + # + # seq_idx : [ 0 | 1 | 2 3 | 4 5 6 7 8 ] + # + # new_token_idx : [ 0 | 1 | 3 | 8 ] + # + # active_token_ids before left shift: + # : [ 31 | 75 | 45 16 | 90 12 72 24 88 ] + # + # active_token_ids after shift: + # : [ XX | XX | 16 XX | 12 72 24 88 XX ] (XX = undefined) + # + # active_token_ids[new_token_idx] = new_tokens + # : [ 52 | 12 | 16 3 | 12 72 24 88 86 ] + active_token_ids = self.gpu_view.token_to_input_ids[: self.active_token_count].roll(-1, 0) + active_query_lengths = self.gpu_view.request_query_lengths[:n_active] + + new_token_idx = active_query_lengths.cumsum(0) - 1 + active_token_ids[new_token_idx] = new_tokens + + # Compute (possibly processed) log-probs over all active-token rows. + log_probs = self._processed_log_probs( + logits_squeezed, n_active, active_query_lengths, sampling + ) + + # Extract the log probs for only the selected tokens. + # (sequence_length x vocab_size) -> (sequence_length) + seq_idx = torch.arange(self.active_token_count, device=log_probs.device) + selected_log_probs = log_probs[seq_idx, active_token_ids] + + # Split the log probs across request boundaries selected_log_probs_list = selected_log_probs.cpu().split( active_query_lengths.tolist(), dim=0 ) @@ -4835,11 +3979,11 @@ def get_kvcache_utilization_stats(self) -> dict: } """ # Total usable blocks exclude the reserved dummy block. - total_blocks = max(self.kv_block_allocator.pool_size - 1, 1) - block_count_avail = int(self.kv_block_allocator.pool_avail) + total_blocks = max(self.kv_block_allocator.total_count - 1, 1) + block_count_avail = int(self.kv_block_allocator.total_avail) # Overall allocated blocks in the buffer right now. - allocated_blocks = (self.kv_block_allocator.pool_size - 1) - block_count_avail + allocated_blocks = (self.kv_block_allocator.total_count - 1) - block_count_avail allocated_blocks = int(max(0, allocated_blocks)) # Active unique blocks referenced by current active requests only. @@ -4875,185 +4019,3 @@ def get_kvcache_utilization_stats(self) -> dict: 'total_request_count': int(total_request_count), 'max_requests': int(self.max_requests), } - - # ----- VLM per-request data management ----- - - def add_vlm_request_data( - self, - request_id: int, - image_embeddings: Optional[Tensor] = None, - image_token_mask: Optional[Tensor] = None, - ) -> None: - """Attach per-request image data to the context. - - Called at add_request time when a multimodal request has images. - No-op overhead for text-only requests that never call this. - - Args: - request_id: The request identifier. - image_embeddings: Tensor of shape [seq_img, 1, hidden] or None. - image_token_mask: 1-D tensor with -1 for text positions and - non-negative indices for image embedding positions; or None. - """ - self._request_to_image_embeddings[request_id] = image_embeddings - self._request_to_image_token_mask[request_id] = image_token_mask - if image_embeddings is not None: - self._request_to_image_token_count[request_id] = int( - image_embeddings.shape[0] * image_embeddings.shape[1] - ) - else: - self._request_to_image_token_count[request_id] = 0 - - def remove_vlm_request_data(self, request_id: int) -> None: - """Remove image data for a finished request.""" - self._request_to_image_embeddings.pop(request_id, None) - self._request_to_image_token_mask.pop(request_id, None) - self._request_to_image_token_count.pop(request_id, None) - - @property - def has_vlm_data(self) -> bool: - """True iff any active request has attached image data. - - Used by callers on the hot path (``TextGenerationController._ - dynamic_step_forward_logits``) to skip ``current_image_token_mask`` - / ``current_image_embeddings`` entirely on text-only workloads. - """ - return bool(self._request_to_image_token_mask) - - def current_image_token_mask(self) -> Optional[Tensor]: - """Flattened image-token mask aligned with current_input_ids. - - Returns a [1, padded_active_token_count] tensor with -1 for non-image - positions and non-negative indices into the concatenated image - embeddings, or None when there are no active tokens. During - decode-only steps (no image tokens consumed) returns all -1. - """ - # Text-only workloads never populate the per-request VLM dicts; short- - # circuit so the decode hot path doesn't pay for .tolist() / torch.cat. - if not self._request_to_image_token_mask: - return None - - if self.padded_active_token_count is None or self.padded_active_token_count == 0: - return None - - if self.is_decode_only(): - # Return the same [1, padded_active_token_count] shape the prefill - # branch does. Callers advanced-index a batch-first [b, seq, h] - # embedding tensor with this mask; the 2D form is the one that - # broadcasts correctly across the batch axis, a 1D mask would - # attempt to select along dim 0 and either error or silently - # index wrong. - return torch.full( - (1, self.padded_active_token_count), - -1, - dtype=torch.long, - device=torch.cuda.current_device(), - ) - - # Batch the three device→host copies into a single sync to keep the - # prefill hot path from stalling three separate times per step. A full - # vectorization (build the flat mask on-device via cumsum + gather) - # would remove even this remaining sync; TODO before this path takes - # heavy multimodal traffic. - active_slice = slice(self.paused_request_count, self.total_request_count) - _sync_batch = torch.stack( - [ - self.request_ids[active_slice], - self.request_query_lengths[active_slice], - # KV offset — how far into the stored prompt this step begins. - # Non-zero under chunked prefill and after pause/resume, so - # slicing per_request_mask must start at kv_offset, not 0. - self.request_kv_length_offsets[active_slice], - ], - dim=0, - ).tolist() - active_request_ids, active_query_lengths, active_kv_offsets = ( - _sync_batch[0], - _sync_batch[1], - _sync_batch[2], - ) - - segments: List[Tensor] = [] - cumulative_offset = 0 - for request_id, query_len, kv_offset in zip( - active_request_ids, active_query_lengths, active_kv_offsets - ): - per_request_mask = self._request_to_image_token_mask.get(request_id, None) - if per_request_mask is None or kv_offset >= per_request_mask.numel(): - # Past the end of the stored prompt mask — decode step, or a - # text-only request. Return an all-text (-1) segment; the - # stored mask is prompt-length and doesn't cover generated - # positions. - seg = torch.full( - (query_len,), -1, dtype=torch.long, device=torch.cuda.current_device() - ) - else: - # kv_offset is normally 0 for a fresh prefill and grows with - # chunked prefill / pause-and-resume — slice into the stored - # mask from where this step actually resumes. - seg = per_request_mask[kv_offset : kv_offset + query_len].clone() - # Pad if the request is being prefilled beyond the mask (mixed - # prefill+decode chunk). - if seg.numel() < query_len: - pad = torch.full( - (query_len - seg.numel(),), -1, dtype=seg.dtype, device=seg.device - ) - seg = torch.cat([seg, pad]) - positive = seg >= 0 - if positive.any(): - seg[positive] += cumulative_offset - - segments.append(seg) - cumulative_offset += int(self._request_to_image_token_count.get(request_id, 0)) - - if len(segments) == 0: - return None - - mask = torch.cat(segments, dim=0) - if mask.numel() < int(self.padded_active_token_count): - pad = torch.full( - (int(self.padded_active_token_count) - mask.numel(),), - -1, - dtype=torch.long, - device=mask.device, - ) - mask = torch.cat([mask, pad], dim=0) - return mask.unsqueeze(0) - - def current_image_embeddings(self) -> Optional[Tensor]: - """Concatenate image embeddings for active requests. - - Each per-request embedding has shape [seq_img_i, 1, hidden] where - seq_img_i may vary across requests under dynamic resolution. Each is - flattened to [seq_img_i, hidden], concatenated along dim 0, and - unsqueezed to [sum(seq_img_i), 1, hidden]. The downstream consumer - does ``permute(1, 0, 2).reshape(-1, hidden)`` to yield the flat - [total_image_tokens, hidden] array indexed by - :meth:`current_image_token_mask`. - - Returns tensor of shape [total_image_tokens, 1, hidden] or None. - """ - # Text-only workloads never populate the per-request VLM dicts; short- - # circuit so the decode hot path doesn't pay for .tolist() / torch.cat. - if not self._request_to_image_embeddings: - return None - - # Decode steps consume no image tokens, so the concatenated embedding - # is unused. Short-circuit before the .tolist() sync + torch.cat to - # keep the decode critical path free of the D2H stall that would - # otherwise fire on every step, mirroring the mask helper above. - if self.is_decode_only(): - return None - - active_request_ids = self.request_ids[ - self.paused_request_count : self.total_request_count - ].tolist() - - parts: List[Tensor] = [] - for request_id in active_request_ids: - emb = self._request_to_image_embeddings.get(request_id, None) - if emb is not None: - parts.append(emb.reshape(-1, emb.shape[-1])) - if not parts: - return None - return torch.cat(parts, dim=0).unsqueeze(1) diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 20a85839658..fb455ad9ede 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -55,6 +55,7 @@ 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 @@ -242,6 +243,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") if self.config.enable_mhc_connections: diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index bcd429dbfd7..f0750064936 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -16,15 +16,19 @@ 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, CSA, HCA, WINDOW, MLP, MOE} + # MLA-based attention layers (incompatible with standard '*' attention in one model). + MLA_ATTENTION = {DS_ATTENTION, CSA, HCA, WINDOW} @classmethod def name_sorted_valid_layer_symbols(cls) -> list[str]: @@ -175,10 +179,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, 'E': 0, 'W': 0} - >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") - {'*': 1, 'G': 0, 'D': 0, 'M': 8, '-': 4, 'E': 0} + >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") + {'*': 1, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'K': 0, 'M': 8, '-': 4, 'E': 0, 'W': 0} """ parsed = parse_hybrid_pattern(pattern) counts = {symbol: 0 for symbol in Symbols.name_sorted_valid_layer_symbols()} @@ -294,9 +298,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 +329,11 @@ def validate_segment_layers(segment: str) -> List[str]: f"one of {Symbols.VALID_LAYERS}" ) - # Disallow Attention + MLA/DSA hybridity. - if Symbols.ATTENTION in segment and (Symbols.DS_ATTENTION in segment or Symbols.MLA in segment): - raise ValueError("Not supported to have both Attention and MLA/DSA in one model") + # Disallow standard Attention ('*') mixed with any MLA-based attention (D/C/H/W). + if Symbols.ATTENTION in segment and any(s in segment for s in Symbols.MLA_ATTENTION): + raise ValueError( + "Not supported to have both Attention and MLA/DSA/CSA/HCA/Window in one model" + ) return layer_type_list diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index f0be7ef891b..1d422d32dc6 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from functools import partial from megatron.core.extensions.transformer_engine import ( @@ -15,10 +15,11 @@ get_moe_module_spec, ) from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules -from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules -from megatron.core.ssm.gated_delta_product import ( - GatedDeltaProductMixer, - GatedDeltaProductMixerSubmodules, +from megatron.core.ssm.gated_delta_net import ( + GatedDeltaNet, + GatedDeltaNetSubmodules, + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, ) from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules @@ -81,8 +82,9 @@ submodules=MultiTokenPredictionLayerSubmodules( enorm=TENorm, hnorm=TENorm, - # Hybrid MTP selects the combined projection normally and - # per-stream projections when mHC is enabled. + # Both projection forms are populated so the layer can switch on + # `config.enable_hyper_connections` at runtime: mHC=True uses + # `e_proj`+`h_proj` per-stream, mHC=False uses fused `eh_proj`. eh_proj=TEColumnParallelLinear, e_proj=TEColumnParallelLinear, h_proj=TEColumnParallelLinear, @@ -95,19 +97,6 @@ ) -def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): - return ModuleSpec( - module=MambaLayer, - submodules=MambaLayerSubmodules( - mixer=ModuleSpec( - module=GatedDeltaProductMixer, - submodules=GatedDeltaProductMixerSubmodules(in_proj=in_proj, out_proj=out_proj), - ), - mamba_bda=get_bias_dropout_add, - ), - ) - - hybrid_stack_spec = ModuleSpec( module=HybridStack, submodules=HybridStackSubmodules( @@ -137,6 +126,22 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): self_attn_bda=get_bias_dropout_add, ), ), + kda_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, + self_attention=ModuleSpec( + module=KimiDeltaAttention, + submodules=KimiDeltaAttentionSubmodules( + in_proj=TEColumnParallelLinear, + beta_proj=TEColumnParallelLinear, + out_norm=TENorm, + out_proj=TERowParallelLinear, + ), + ), + self_attn_bda=get_bias_dropout_add, + ), + ), # Started with spec from gpt_layer_specs.py (with MLP removed) # Using the TE spec because we had problems getting the non-TE spec # working @@ -190,28 +195,6 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): self_attn_bda=get_bias_dropout_add, ), ), - mla_layer=ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - input_layernorm=TENorm, - self_attention=ModuleSpec( - module=MLASelfAttention, - params={"attn_mask_type": AttnMaskType.causal}, - submodules=MLASelfAttentionSubmodules( - linear_q_proj=TEColumnParallelLinear, - linear_q_down_proj=TELinear, - linear_q_up_proj=TEColumnParallelLinear, - linear_kv_down_proj=TELinear, - linear_kv_up_proj=TEColumnParallelLinear, - core_attention=TEDotProductAttention, - linear_proj=TERowParallelLinear, - q_layernorm=IdentityOp, - kv_layernorm=IdentityOp, - ), - ), - self_attn_bda=get_bias_dropout_add, - ), - ), # Started with spec from gpt_layer_specs.py # Using the TE spec because we had problems getting the non-TE spec # working @@ -238,22 +221,6 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): ) -gated_delta_product_stack_spec = ModuleSpec( - module=HybridStack, - submodules=HybridStackSubmodules( - mamba_layer=_get_gated_delta_product_mamba_layer_spec( - TELayerNormColumnParallelLinear, TERowParallelLinear - ), - gdn_layer=hybrid_stack_spec.submodules.gdn_layer, - attention_layer=hybrid_stack_spec.submodules.attention_layer, - dsa_layer=hybrid_stack_spec.submodules.dsa_layer, - mlp_layer=hybrid_stack_spec.submodules.mlp_layer, - moe_layer=hybrid_stack_spec.submodules.moe_layer, - mtp_block_spec=hybrid_stack_spec.submodules.mtp_block_spec, - ), -) - - hybrid_inference_stack_spec = ModuleSpec( module=HybridStack, submodules=HybridStackSubmodules( @@ -270,20 +237,6 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): mamba_bda=get_bias_dropout_add, ), ), - gdn_layer=ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - self_attention=ModuleSpec( - module=GatedDeltaNet, - submodules=GatedDeltaNetSubmodules( - in_proj=InferenceLayerNormColumnParallelLinear, - out_norm=TENorm, - out_proj=InferenceRowParallelLinear, - ), - ), - self_attn_bda=get_bias_dropout_add, - ), - ), # Started with spec from gpt_layer_specs.py (with MLP removed) # Using the TE spec because we had problems getting the non-TE spec # working @@ -337,28 +290,6 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): self_attn_bda=get_bias_dropout_add, ), ), - mla_layer=ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - input_layernorm=TENorm, - self_attention=ModuleSpec( - module=MLASelfAttention, - params={"attn_mask_type": AttnMaskType.causal}, - submodules=MLASelfAttentionSubmodules( - linear_q_proj=TEColumnParallelLinear, - linear_q_down_proj=TELinear, - linear_q_up_proj=TEColumnParallelLinear, - linear_kv_down_proj=TELinear, - linear_kv_up_proj=TEColumnParallelLinear, - core_attention=TEDotProductAttention, - linear_proj=InferenceRowParallelLinear, - q_layernorm=IdentityOp, - kv_layernorm=IdentityOp, - ), - ), - self_attn_bda=get_bias_dropout_add, - ), - ), # Started with spec from gpt_layer_specs.py # Using the TE spec because we had problems getting the non-TE spec # working @@ -391,7 +322,8 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): submodules=MultiTokenPredictionLayerSubmodules( enorm=TENorm, hnorm=TENorm, - # Keep both projection forms available for Hybrid MTP. + # Populate both projection forms so the layer can switch on + # `config.enable_hyper_connections` at runtime. eh_proj=InferenceColumnParallelLinear, e_proj=InferenceColumnParallelLinear, h_proj=InferenceColumnParallelLinear, @@ -406,24 +338,59 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): ) -gated_delta_product_inference_stack_spec = ModuleSpec( - module=HybridStack, - submodules=HybridStackSubmodules( - mamba_layer=_get_gated_delta_product_mamba_layer_spec( - InferenceLayerNormColumnParallelLinear, InferenceRowParallelLinear - ), - gdn_layer=hybrid_inference_stack_spec.submodules.gdn_layer, - attention_layer=hybrid_inference_stack_spec.submodules.attention_layer, - dsa_layer=hybrid_inference_stack_spec.submodules.dsa_layer, - mlp_layer=hybrid_inference_stack_spec.submodules.mlp_layer, - moe_layer=hybrid_inference_stack_spec.submodules.moe_layer, - mtp_block_spec=hybrid_inference_stack_spec.submodules.mtp_block_spec, - ), -) - - # Backward-compatible aliases mamba_stack_spec = hybrid_stack_spec mamba_inference_stack_spec = hybrid_inference_stack_spec -gdp_stack_spec = gated_delta_product_stack_spec -gdp_inference_stack_spec = gated_delta_product_inference_stack_spec + + +def hybrid_dsv4_stack_spec(config): + """Config-aware hybrid stack spec whose ``D`` (DS_ATTENTION) layer runs the DSv4 + ``CompressedSparseAttention`` (CSA/HCA + ``CSAIndexer``), identical to the GPT + ``dsv4_hybrid`` path, instead of the legacy ``DSAttention``. + + The default ``hybrid_stack_spec`` wires the ``D`` layer to ``AbsorbedMLASelfAttention + + DSAttention`` (DS3.2-style sparse attention with no CSA/HCA compression). To run real + DSv4 on HybridModel — and to be numerically equivalent to a GPT ``dsv4_hybrid`` + attention layer — we reuse GPT's own ``get_dsv4_hybrid_module_spec_for_backend`` + (which is config-aware, e.g. picks the qk-layernorm form from ``config``) so the two + model paths build the *same* attention module. Selected via + ``--spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_dsv4_stack_spec``; + ``hybrid_builder`` invokes this function with ``config`` when the spec is callable. + """ + import dataclasses + + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + _get_backend_spec_provider, + get_dsv4_hybrid_module_spec_for_backend, + ) + + backend = _get_backend_spec_provider(config) + dsv4_attention = get_dsv4_hybrid_module_spec_for_backend(config, backend) + + def _wrap_dsv4_layer(compress_ratio=None): + # Wrap the DSv4 attention in a hybrid TransformerLayer. When compress_ratio is given + # (the 'C'/'H'/'W' layer symbols), bake it into the attention params so the layer uses a + # fixed CSA(4)/HCA(128)/window-only(0) ratio regardless of csa_compress_ratios; otherwise + # the layer reads its ratio from config.csa_compress_ratios (array-driven 'D' / GPT-parity + # path). compress_ratio=0 builds neither the compressor nor the top-k indexer, so the + # 'W' layer reuses the entire CSA/HCA code path as pure sliding-window attention. + attn = dsv4_attention + if compress_ratio is not None: + attn = dataclasses.replace( + dsv4_attention, params={**dsv4_attention.params, "compress_ratio": compress_ratio} + ) + return ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, self_attention=attn, self_attn_bda=get_bias_dropout_add + ), + ) + + submodules = dataclasses.replace( + hybrid_stack_spec.submodules, + dsa_layer=_wrap_dsv4_layer(), # 'D': array-driven (or window) DSv4 attention + csa_layer=_wrap_dsv4_layer(compress_ratio=4), # 'C': CSA + hca_layer=_wrap_dsv4_layer(compress_ratio=128), # 'H': HCA + window_layer=_wrap_dsv4_layer(compress_ratio=0), # 'W': sliding-window-only + ) + return ModuleSpec(module=HybridStack, submodules=submodules) diff --git a/megatron/core/ssm/gated_delta_net/__init__.py b/megatron/core/ssm/gated_delta_net/__init__.py index e5a5b4e6a9f..ba139891c36 100644 --- a/megatron/core/ssm/gated_delta_net/__init__.py +++ b/megatron/core/ssm/gated_delta_net/__init__.py @@ -1,6 +1,6 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Gated Delta Net (GDN) family of layers. +"""Gated Delta Net (GDN) layer exports. This package replaces the former ``megatron/core/ssm/gated_delta_net.py`` module at the same import path; the names below preserve that module's public surface. @@ -9,6 +9,9 @@ from megatron.core.ssm.gated_delta_net.common import ( HAVE_FLA, GatedDeltaNetSubmodules, + _build_head_perm_for_split_sections, + _build_thd_cp_a2a_perm, + _split_tensor_factory, causal_conv1d, chunk_gated_delta_rule, get_parameter_local_cp, @@ -17,26 +20,25 @@ tensor_a2a_hp2cp, ) from megatron.core.ssm.gated_delta_net.gdn import GatedDeltaNet, torch_chunk_gated_delta_rule -from megatron.core.ssm.gated_delta_net.gdn2 import ( - HAVE_FLA_GDN2, - GatedDeltaNet2, - chunk_gdn2, - torch_chunk_gdn2, +from megatron.core.ssm.gated_delta_net.kda import ( + HAVE_FLA_KDA, + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, ) __all__ = [ "HAVE_FLA", - "HAVE_FLA_GDN2", + "HAVE_FLA_KDA", "GatedDeltaNet", - "GatedDeltaNet2", "GatedDeltaNetSubmodules", + "KimiDeltaAttention", + "KimiDeltaAttentionSubmodules", + "_split_tensor_factory", "causal_conv1d", "chunk_gated_delta_rule", - "chunk_gdn2", "get_parameter_local_cp", "l2norm", "tensor_a2a_cp2hp", "tensor_a2a_hp2cp", "torch_chunk_gated_delta_rule", - "torch_chunk_gdn2", ] diff --git a/megatron/core/ssm/gated_delta_net/kda.py b/megatron/core/ssm/gated_delta_net/kda.py new file mode 100644 index 00000000000..d278ce003c3 --- /dev/null +++ b/megatron/core/ssm/gated_delta_net/kda.py @@ -0,0 +1,445 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh. + +"""Kimi Delta Attention, a channel-wise Gated DeltaNet variant.""" + +import math +from dataclasses import dataclass +from functools import partial + +import torch +import torch.nn.functional as F + +from megatron.core import tensor_parallel +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, + causal_conv1d, + get_parameter_local_cp, + l2norm, +) +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, + ) -> None: + self._validate_config(config) + 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, + ) + + 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) + self.dt_bias._no_reinit = True + self.dt_bias._no_weight_decay = True + self.A_log._no_weight_decay = True + + # Pass raw g to FLA and let the KDA kernel apply A_log and dt_bias. + self.use_gate_in_kernel = True + + @staticmethod + def _validate_config(config: TransformerConfig) -> None: + """Validate the direct-projection KDA layout supported by this implementation.""" + + if config.linear_num_key_heads != config.linear_num_value_heads: + raise ValueError( + "KimiDeltaAttention currently requires equal key and value head counts." + ) + if config.linear_key_head_dim != config.linear_value_head_dim: + raise ValueError( + "KimiDeltaAttention currently requires equal key and value head dimensions." + ) + tp_cp_size = config.tensor_model_parallel_size * config.context_parallel_size + if config.linear_num_key_heads % tp_cp_size != 0: + raise ValueError( + "KimiDeltaAttention key heads must be divisible by tensor parallel size times " + "context parallel size." + ) + + def _setup_variant_attrs(self) -> None: + """Set KDA dimensions and projection checkpoint metadata.""" + + self.in_proj_extra_dim = self.qk_dim + 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 = ( + (2 * self.qk_dim_local_tp + self.v_dim_local_tp) // self.cp_size, + self.qk_dim_local_tp // self.cp_size, + self.v_dim_local_tp // 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 reset_parameters(self) -> None: + """Initialize convolution and KDA gate parameters.""" + + if not self.config.perform_initialization: + return + + with tensor_parallel.get_cuda_rng_tracker().fork(): + if self.conv_init is not None: + torch.nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init) + + dt_min, dt_max, dt_init_floor = 0.001, 0.1, 1e-4 + dt = torch.exp( + torch.rand(self.dt_bias_dim, device=self.dt_bias.device, dtype=torch.float32) + * (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))) + + A = torch.empty(self.a_log_dim, device=self.A_log.device, dtype=torch.float32).uniform_( + *self.A_init_range + ) + self.A_log.data.copy_(torch.log(A)) + + @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 _prepare_kda_inputs( + self, + qkv: torch.Tensor, + raw_g: torch.Tensor, + beta: torch.Tensor, + output_gate: torch.Tensor, + batch: int, + seq_len: int, + ) -> tuple[torch.Tensor, ...]: + """Split, normalize, and reshape tensors for the FLA KDA kernel.""" + + 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, + ) + query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) + value = value.reshape(batch, seq_len, -1, self.value_head_dim) + if self.use_qk_l2norm: + query_key = l2norm(query_key.contiguous()) + + num_key_heads = self.num_k_heads_local_tp // self.cp_size + query, key = torch.split(query_key, [num_key_heads, num_key_heads], dim=2) + 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() + output_gate = output_gate.reshape( + batch, seq_len, self.num_v_heads_local_tp // self.cp_size, self.value_head_dim + ) + return ( + query.contiguous(), + key.contiguous(), + value.contiguous(), + raw_g.contiguous(), + beta.contiguous(), + output_gate.contiguous(), + ) + + @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 _gated_norm_and_a2a( + self, + core_attn_out: torch.Tensor, + gate: torch.Tensor, + thd_cp_a2a_inv: torch.Tensor | None, + batch: int, + seq_len: int, + packed_seq_params: PackedSeqParams | None = None, + ) -> torch.Tensor: + """Apply KDA's gated output norm and restore its context-parallel layout.""" + + nvtx_range_push(suffix="gated_norm") + norm_out_hp = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + norm_out_hp = norm_out_hp.reshape(batch, seq_len, -1) + norm_out_hp = norm_out_hp.transpose(0, 1).contiguous() + + return a2a_hp_to_cp( + norm_out_hp, self.cp_size, self.pg_collection.cp, packed_seq_params, thd_cp_a2a_inv + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + inference_context: BaseInferenceContext | None = None, + packed_seq_params: PackedSeqParams | None = None, + sequence_len_offset: int | None = None, + *, + inference_params: BaseInferenceContext | None = 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) + if inference_context is not None: # pragma: no cover + raise NotImplementedError("KDA recurrent inference is not implemented.") + + seq_len, batch, _ = hidden_states.shape + total_seq_len = seq_len * self.sp_size * self.cp_size + + cu_seqlens_q = None + 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, + total_seq_len, + "cu_seqlens_q", + cp_size=self.cp_size, + ) + cu_seqlens_kv = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_kv_padded, + packed_seq_params.cu_seqlens_kv, + total_seq_len, + "cu_seqlens_kv", + cp_size=self.cp_size, + ) + self._validate_packed_cu_seqlens(cu_seqlens_q, cu_seqlens_kv) + + nvtx_range_push(suffix="kda_in_proj") + qkvfg, _ = self.in_proj(hidden_states) + beta, _ = self.beta_proj(hidden_states) + nvtx_range_pop(suffix="kda_in_proj") + + qkvfg, thd_cp_a2a_inv = a2a_cp_to_hp( + qkvfg, + self.in_proj_split_sections, + self.cp_size, + self.pg_collection.cp, + cu_seqlens_q, + total_seq_len, + packed_seq_params, + ) + beta, _ = a2a_cp_to_hp( + beta, + (self.num_k_heads_local_tp,), + self.cp_size, + self.pg_collection.cp, + cu_seqlens_q, + total_seq_len, + packed_seq_params, + ) + + qkvfg = qkvfg.transpose(0, 1) + beta = beta.transpose(0, 1) + seq_len = qkvfg.shape[1] + qkv, raw_g, output_gate = torch.split(qkvfg, self.feat_dim_split, dim=-1) + + nvtx_range_push(suffix="kda_conv1d") + 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, + dim=0, + cp_group=self.pg_collection.cp, + 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( + qkv, + 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, + ) + qkv = self.act_fn(conv_out[..., :seq_len]).transpose(1, 2).contiguous() + else: + if self.activation not in ("silu", "swish"): + raise ValueError(f"FLA causal convolution requires SiLU, got {self.activation}.") + qkv, _ = causal_conv1d( + x=qkv, + weight=conv1d_weight.squeeze(1), + bias=conv1d_bias, + activation=self.activation, + initial_state=None, + output_final_state=False, + cu_seqlens=cu_seqlens_q, + ) + nvtx_range_pop(suffix="kda_conv1d") + + query, key, value, raw_g, beta, output_gate = self._prepare_kda_inputs( + qkv, raw_g, beta, output_gate, batch, seq_len + ) + 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 + ) + + nvtx_range_push(suffix="kda_chunk") + core_attn_out, _ = self.gated_delta_rule( + q=query, + k=key, + v=value, + g=raw_g, + beta=beta, + A_log=A_log_local_cp, + dt_bias=dt_bias_local_cp, + 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, + ) + nvtx_range_pop(suffix="kda_chunk") + + 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, output_gate) + else: + norm_out = self._gated_norm_and_a2a( + core_attn_out, output_gate, thd_cp_a2a_inv, batch, seq_len, packed_seq_params + ) + + nvtx_range_push(suffix="kda_out_proj") + out, out_bias = self.out_proj(norm_out) + nvtx_range_pop(suffix="kda_out_proj") + + if self.recompute_norm_out: + self.norm_out_checkpoint.discard_output_and_register_recompute(out) + + return out, out_bias + + 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/transformer_config.py b/megatron/core/transformer/transformer_config.py index a5933021700..2d30c38e6da 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -412,6 +412,12 @@ 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.""" diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index dc093508775..095dce88cb6 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -11,10 +11,20 @@ from pathlib import Path import torch +import torch.nn.functional as F +from packaging.version import Version as PkgVersion +from megatron.core.activations import squared_relu +from megatron.core.dist_checkpointing.validation import StrictHandling +from megatron.core.fusions.fused_bias_geglu import quick_gelu +from megatron.core.model_parallel_config import _parse_pad_packed_seq_alignment from megatron.core.msc_utils import MultiStorageClientFeature +from megatron.core.quantization.utils import ( + kitchen_quantization_recipe_config, + load_quantization_recipe, +) from megatron.core.rerun_state_machine import RerunStateMachine -from megatron.core.transformer import TransformerConfig +from megatron.core.transformer import MLATransformerConfig, TransformerConfig from megatron.core.transformer.cuda_graph_config import ( ALLOWED_INFERENCE_SCOPES, get_deprecated_cuda_graph_modules_migration, @@ -23,6 +33,10 @@ validate_deprecated_cuda_graph_modules_migration_inputs, ) from megatron.core.transformer.enums import AttnBackend, CudaGraphModule, InferenceCudaGraphScope +from megatron.core.transformer.heterogeneous.heterogeneous_config import ( + HeterogeneousTransformerConfig, + MLPConfig, +) from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.utils import ( get_torch_version, @@ -30,9 +44,10 @@ is_te_min_version, is_torch_min_version, ) -from megatron.training.argument_utils import ( # noqa: F401 # pylint: disable=unused-import +from megatron.training.argument_utils import ( ArgumentGroupFactory, - core_transformer_config_from_args, + _normalize_dsv4_hybrid_csa_compress_ratios, + _resolve_dsa_kernel_backend_cli_default, ) from megatron.training.global_vars import set_global_variables from megatron.training.utils import ( @@ -44,7 +59,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): - """"Add Megatron-LM arguments to the given parser.""" + """ "Add Megatron-LM arguments to the given parser.""" # Standard arguments. parser = _add_network_size_args(parser) @@ -67,7 +82,6 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): parser = _add_experimental_attention_variant_args(parser) parser = _add_heterogeneous_args(parser) parser = _add_logging_args(parser) - parser = _add_logits_distillation_args(parser) parser = _add_straggler_detector_args(parser) parser = _add_workload_inspector_server_args(parser) parser = _add_inference_args(parser) @@ -80,18 +94,23 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): parser = _add_msc_args(parser) parser = _add_kitchen_quantization_arguments(parser) parser = _add_sft_args(parser) + parser = _add_varlen_dataset_args(parser) + parser = _add_logits_distillation_args(parser) parser = _add_fault_injector_args(parser) return parser + def parse_and_validate_args(extra_args_provider=None, ignore_unknown_args=False, args_defaults={}): args = parse_args(extra_args_provider, ignore_unknown_args) if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False): from megatron.training.checkpointing import load_args_from_checkpoint - assert args.load is not None or args.pretrained_checkpoint is not None, "--use-checkpoint-args requires --load or --pretrained-checkpoint argument" + assert ( + args.load is not None or args.pretrained_checkpoint is not None + ), "--use-checkpoint-args requires --load or --pretrained-checkpoint argument" assert args.non_persistent_ckpt_type != "local", ( "--use-checkpoint-args is not supported with --non_persistent_ckpt_type=local. " "Two-stage checkpoint loading is not implemented, and all arguments must be defined " @@ -116,8 +135,7 @@ def parse_and_validate_args(extra_args_provider=None, ignore_unknown_args=False, def parse_args(extra_args_provider=None, ignore_unknown_args=False): """Parse all arguments.""" - parser = argparse.ArgumentParser(description='Megatron-LM Arguments', - allow_abbrev=False) + parser = argparse.ArgumentParser(description='Megatron-LM Arguments', allow_abbrev=False) parser = add_megatron_arguments(parser) @@ -139,7 +157,6 @@ def parse_args(extra_args_provider=None, ignore_unknown_args=False): args = load_yaml(args.yaml_cfg) - # Args from environment args.rank = int(os.getenv('RANK', '0')) args.world_size = int(os.getenv("WORLD_SIZE", '1')) @@ -198,17 +215,22 @@ def validate_model_config_args_from_heterogeneous_config(args): args.heterogeneous_layers_config_path ).read_text() - hf_config_dict = types.SimpleNamespace(**json.loads(args.heterogeneous_layers_config_encoded_json)) - - assert hf_config_dict.hidden_act == "silu", ( - f"hidden_act in heterogeneous config is {hf_config_dict.hidden_act}, should be silu" + hf_config_dict = types.SimpleNamespace( + **json.loads(args.heterogeneous_layers_config_encoded_json) ) + assert ( + hf_config_dict.hidden_act == "silu" + ), f"hidden_act in heterogeneous config is {hf_config_dict.hidden_act}, should be silu" + n_kv_heads_in_group = [ - config["attention"]["n_heads_in_group"] for config in hf_config_dict.block_configs + config["attention"]["n_heads_in_group"] + for config in hf_config_dict.block_configs if config["attention"]["n_heads_in_group"] is not None ] - assert all(num == n_kv_heads_in_group[0] for num in n_kv_heads_in_group), "num query head must be consistent across all layers" + assert all( + num == n_kv_heads_in_group[0] for num in n_kv_heads_in_group + ), "num query head must be consistent across all layers" args_to_validate = { "swiglu": True, @@ -234,16 +256,17 @@ def validate_model_config_args_from_heterogeneous_config(args): incompatible_args[key] = (provided_value, value) if incompatible_args: - incompatible_args_str = ', '.join([ - f"{k}: {provided_value} (provided) != {value} (expected)" - for k, (provided_value, value) in incompatible_args.items() - ]) - raise ValueError( - f"Arguments differ from heterogeneous config: {incompatible_args_str}" + incompatible_args_str = ', '.join( + [ + f"{k}: {provided_value} (provided) != {value} (expected)" + for k, (provided_value, value) in incompatible_args.items() + ] ) + raise ValueError(f"Arguments differ from heterogeneous config: {incompatible_args_str}") + def _eval_pattern(pattern): - """ Validate and evaluate a string containing a Python list expression """ + """Validate and evaluate a string containing a Python list expression""" assert isinstance(pattern, str) # validate input, only allow comma, digits, [, ], (, ), +, and * @@ -266,12 +289,12 @@ def _normalize_cuda_graph_modules_args(args): args.cuda_graph_modules ) validate_deprecated_cuda_graph_modules_migration_inputs( - deprecated_scopes, - args.cuda_graph_impl, - args.inference_cuda_graph_scope, + deprecated_scopes, args.cuda_graph_impl, args.inference_cuda_graph_scope ) if used_full_scope: - warn_rank_0('full scope is deprecated. Use empty cuda_graph_modules to capture the whole layer.') + warn_rank_0( + 'full scope is deprecated. Use empty cuda_graph_modules to capture the whole layer.' + ) for scope, attr, value in deprecated_scopes: migration = get_deprecated_cuda_graph_modules_migration( @@ -302,7 +325,7 @@ def _normalize_inference_cuda_graph_scope_arg(args): def no_rope_freq_type(x): - """ Controls which layers to skip performing Rotary Position Embedding. + """Controls which layers to skip performing Rotary Position Embedding. - An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. - A string "N": Same as above, but provided as a string - A string containing a Python list expression that defines a custom pattern, e.g.: @@ -324,6 +347,22 @@ def no_rope_freq_type(x): # it's a single int but in str return int(x) + +def compress_ratios_type(x): + """Per-layer compress ratios for compressed sparse attention. + + Accepts a string containing a Python list expression, e.g.: + "[0,0,4,128,4,128]" + "([0]+[4,128]*2)*3" + The result must be a list of integers. Each value represents the + compression ratio for the corresponding transformer layer. + """ + if isinstance(x, list): + return x + assert isinstance(x, str) + return _eval_pattern(x) + + def moe_freq_type(x): """Frequency between MoE layers and Dense layers. @@ -349,6 +388,7 @@ def moe_freq_type(x): # it's a single int but in str return int(x) + def la_freq_type(x): """Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) layers. @@ -374,6 +414,7 @@ def la_freq_type(x): # it's a single int but in str return int(x) + def tuple_type(x): """ Convert a string to a tuple of integers. @@ -396,8 +437,11 @@ def validate_args(args, defaults={}): args.exit_on_missing_checkpoint = True # Temporary - assert args.non_persistent_ckpt_type in ['global', 'local', None], \ - 'Currently only global and local checkpoints are supported' + assert args.non_persistent_ckpt_type in [ + 'global', + 'local', + None, + ], 'Currently only global and local checkpoints are supported' if args.non_persistent_ckpt_type == 'local': try: from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import ( @@ -411,31 +455,21 @@ def validate_args(args, defaults={}): update_use_dist_ckpt(args) - # GTP_remat counts toward total_model_size (an independent weight-shard axis), so the - # args.data_parallel_size below is the replicate degree (matches - # parallel_state). gtp_weight_remat_size is derived from --tensor-parallel-num-weight-shards. - from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards - (args.tensor_parallel_num_weight_shards, args.gtp_weight_remat_size) = ( - resolve_tensor_parallel_weight_shards( - args.tensor_model_parallel_size, - args.tensor_parallel_num_weight_shards, - getattr(args, "gtp_weight_remat_size", 1), - ) - ) total_model_size = ( args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size - * args.gtp_weight_remat_size ) # Total model size. - assert args.world_size % total_model_size == 0, ( - f"world size ({args.world_size}) is not divisible by total_model_size ({total_model_size=})" - ) + assert ( + args.world_size % total_model_size == 0 + ), f"world size ({args.world_size}) is not divisible by total_model_size ({total_model_size=})" if args.attention_backend == AttnBackend.local: - assert args.spec[0] == 'local' , '--attention-backend local is only supported with --spec local' + assert ( + args.spec[0] == 'local' + ), '--attention-backend local is only supported with --spec local' # Pipeline model parallel size. args.transformer_pipeline_model_parallel_size = args.pipeline_model_parallel_size @@ -444,13 +478,10 @@ def validate_args(args, defaults={}): args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size - * args.gtp_weight_remat_size ) args.data_parallel_size = args.world_size // total_model_size if args.perform_rl_step: - assert args.refit_method != 'nccl_m2n', 'nccl_m2n is unsupported by the built-in RL loop' - # ---------------------------------------------------------------- # CUDA graphs # @@ -481,13 +512,13 @@ def validate_args(args, defaults={}): # ---------------------------------------------------------------- # Persisting CGs only makes sense if we build any CGs. - assert not args.rl_persist_cuda_graphs or args.cuda_graph_impl != "none", ( - "--rl-persist-cuda-graphs is set but no CUDA graphs are being built." - ) + assert ( + not args.rl_persist_cuda_graphs or args.cuda_graph_impl != "none" + ), "--rl-persist-cuda-graphs is set but no CUDA graphs are being built." # Training CGs only makes sense if we build any CGs. - assert not args.rl_training_cuda_graphs or args.cuda_graph_impl != "none", ( - "--rl-training-cuda-graphs is set but no CUDA graphs are being built." - ) + assert ( + not args.rl_training_cuda_graphs or args.cuda_graph_impl != "none" + ), "--rl-training-cuda-graphs is set but no CUDA graphs are being built." # If CUDA graphs persist and KV cache memory address is not static, we need # either UVM or torch_memory_saver to maintain memory address stability for CGs. if args.rl_persist_cuda_graphs and args.rl_kv_cache_management_mode != "persist": @@ -501,8 +532,8 @@ def validate_args(args, defaults={}): # Offload mode requires CG persistence: CG recapture runs dummy forward # passes that corrupt the preserved KV data. - assert ( - (not args.rl_kv_cache_management_mode == "offload") or (args.rl_persist_cuda_graphs) + assert (not args.rl_kv_cache_management_mode == "offload") or ( + args.rl_persist_cuda_graphs ), "--rl-kv-cache-management-mode=offload requires --rl-persist-cuda-graphs" # There's no need to manually offload the KV cache with UVM. @@ -511,7 +542,7 @@ def validate_args(args, defaults={}): and args.rl_kv_cache_management_mode == "offload" ), "--rl-kv-cache-management-mode=offload is incompatible with UVM" # We currently cannot recapture CGs in offload mode. - assert not( + assert not ( not args.rl_persist_cuda_graphs and args.rl_kv_cache_management_mode == "offload" ), "Cannot recapture CUDA graphs while offloading KV cache." @@ -528,82 +559,75 @@ def validate_args(args, defaults={}): "installed. See https://github.com/fzyzcjy/torch_memory_saver." ) - if args.rl_max_inflight_requests is not None: - requests_per_batch = args.grpo_prompts_per_step * args.grpo_group_size - assert args.rl_generation_lag is None, \ - "--rl-generation-lag and --rl-max-inflight-requests are mutually exclusive." - assert args.rl_max_inflight_requests >= 1, \ - f"--rl-max-inflight-requests ({args.rl_max_inflight_requests}) must be >= 1." - if args.rl_max_inflight_requests > requests_per_batch: - assert args.rl_partial_rollouts, \ - f"--rl-max-inflight-requests above one training batch " \ - f"({requests_per_batch} requests) requires --rl-partial-rollouts." - # Total in-flight requests = (lag + 1) trainer batches of P * G requests each. - args.rl_generation_lag = args.rl_max_inflight_requests / requests_per_batch - 1 - if args.rl_generation_lag is None: - # With --rl-partial-rollouts the lag is autotuned from engine capacity - # at inference launch; otherwise generation is fully synchronous. - if not args.rl_partial_rollouts: - args.rl_generation_lag = 0 - else: - assert args.rl_generation_lag >= -1, \ - f"--rl-generation-lag ({args.rl_generation_lag}) must be >= -1." - if args.rl_generation_lag > 0: - assert args.rl_partial_rollouts, \ - "--rl-generation-lag requires --rl-partial-rollouts." - assert args.rl_submission_granularity == "B" or args.rl_partial_rollouts, \ - f"--rl-submission-granularity {args.rl_submission_granularity} requires " \ - "--rl-partial-rollouts." - assert args.rl_consumption_granularity != "R", \ - "--rl-consumption-granularity R is not currently supported." + submit_rollouts_at_rollout_granularity = args.rl_submission_granularity == "R" + if args.rl_generation_lag > 0: + assert args.rl_partial_rollouts, "--rl-generation-lag requires --rl-partial-rollouts." + if submit_rollouts_at_rollout_granularity: + assert ( + args.rl_partial_rollouts + ), "Rollout submission granularity requires streaming grouped rollouts." + assert ( + args.rl_consumption_granularity != "R" + ), "--rl-consumption-granularity R is not currently supported." assert not ( - args.rl_submission_granularity == "B" - and args.rl_consumption_granularity == "G" + args.rl_submission_granularity == "B" and args.rl_consumption_granularity == "G" ), "--rl-submission-granularity B with --rl-consumption-granularity G is not supported." args.grpo_samples_per_iteration = args.grpo_prompts_per_step * args.grpo_group_size if args.rl_use_sequence_packing: - assert args.micro_batch_size == 1, \ - "micro_batch_size must be 1 when using sequence packing. To increase compute per micro batch increase the sequence length." - - print_rank_0('using world size: {}, data-parallel size: {}, ' - 'context-parallel size: {}, ' - 'hierarchical context-parallel sizes: {}, ' - 'tensor-model-parallel size: {}, ' - 'pipeline-model-parallel size: {}'.format( - args.world_size, args.data_parallel_size, - args.context_parallel_size, - args.hierarchical_context_parallel_sizes, - args.tensor_model_parallel_size, - args.pipeline_model_parallel_size)) + assert ( + args.micro_batch_size == 1 + ), "micro_batch_size must be 1 when using sequence packing. To increase compute per micro batch increase the sequence length." + + print_rank_0( + 'using world size: {}, data-parallel size: {}, ' + 'context-parallel size: {}, ' + 'hierarchical context-parallel sizes: {}, ' + 'tensor-model-parallel size: {}, ' + 'pipeline-model-parallel size: {}'.format( + args.world_size, + args.data_parallel_size, + args.context_parallel_size, + args.hierarchical_context_parallel_sizes, + args.tensor_model_parallel_size, + args.pipeline_model_parallel_size, + ) + ) # Checks. if args.hierarchical_context_parallel_sizes: from numpy import prod + assert args.context_parallel_size == prod(args.hierarchical_context_parallel_sizes) if "a2a+p2p" in args.cp_comm_type: - assert args.hierarchical_context_parallel_sizes is not None, \ - "--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm" + assert ( + args.hierarchical_context_parallel_sizes is not None + ), "--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm" if args.expert_tensor_parallel_size is None: args.expert_tensor_parallel_size = args.tensor_model_parallel_size # Deprecated arguments. - assert args.batch_size is None, '--batch-size argument is no longer ' \ - 'valid, use --micro-batch-size instead' + assert args.batch_size is None, ( + '--batch-size argument is no longer ' 'valid, use --micro-batch-size instead' + ) del args.batch_size - assert args.warmup is None, '--warmup argument is no longer valid, use ' \ - '--lr-warmup-fraction instead' + assert args.warmup is None, ( + '--warmup argument is no longer valid, use ' '--lr-warmup-fraction instead' + ) del args.warmup - assert args.model_parallel_size is None, '--model-parallel-size is no ' \ - 'longer valid, use --tensor-model-parallel-size instead' + assert args.model_parallel_size is None, ( + '--model-parallel-size is no ' 'longer valid, use --tensor-model-parallel-size instead' + ) del args.model_parallel_size if args.checkpoint_activations: - print_rank_0('--checkpoint-activations is no longer valid, use --recompute-activations, ' - 'or, for more control, --recompute-granularity and --recompute-method.') + print_rank_0( + '--checkpoint-activations is no longer valid, use --recompute-activations, ' + 'or, for more control, --recompute-granularity and --recompute-method.' + ) exit() del args.checkpoint_activations @@ -634,12 +658,10 @@ def validate_args(args, defaults={}): del args.external_cuda_graph if getattr(args, 'cuda_graph_scope_deprecated', None) is not None: - assert not args.cuda_graph_modules, ( - "--cuda-graph-scope and --cuda-graph-modules cannot be used together." - ) - warn_rank_0( - '--cuda-graph-scope is deprecated, use --cuda-graph-modules instead.' - ) + assert ( + not args.cuda_graph_modules + ), "--cuda-graph-scope and --cuda-graph-modules cannot be used together." + warn_rank_0('--cuda-graph-scope is deprecated, use --cuda-graph-modules instead.') args.cuda_graph_modules = args.cuda_graph_scope_deprecated del args.cuda_graph_scope_deprecated @@ -647,10 +669,7 @@ def validate_args(args, defaults={}): # all subsequent validation sees fully-typed enum values. _normalize_cuda_graph_modules_args(args) _normalize_inference_cuda_graph_scope_arg(args) - assert ( - args.inference_cuda_graph_scope - in ALLOWED_INFERENCE_SCOPES[args.cuda_graph_impl] - ), ( + assert args.inference_cuda_graph_scope in ALLOWED_INFERENCE_SCOPES[args.cuda_graph_impl], ( "Invalid inference CUDA graph scope " f"{args.inference_cuda_graph_scope.name!r} for " f"--cuda-graph-impl={args.cuda_graph_impl!r}." @@ -662,31 +681,42 @@ def validate_args(args, defaults={}): # arguments that are passed to the program. We check this by # ensuring the arg is set to None. if getattr(args, key, None) is not None: - warn_rank_0('Overriding default arguments for {key}:{v} ' - 'with {key}:{v2}'.format(key=key, v=defaults[key], - v2=getattr(args, key))) + warn_rank_0( + 'Overriding default arguments for {key}:{v} ' + 'with {key}:{v2}'.format(key=key, v=defaults[key], v2=getattr(args, key)) + ) else: setattr(args, key, defaults[key]) if args.data_path is not None and args.split is None: legacy_default_split_value = '969, 30, 1' - warn_rank_0('Please specify --split when using --data-path. Using legacy default value ' - f'of "{legacy_default_split_value}"') + warn_rank_0( + 'Please specify --split when using --data-path. Using legacy default value ' + f'of "{legacy_default_split_value}"' + ) args.split = legacy_default_split_value use_data_path = (args.data_path is not None) or (args.data_args_path is not None) if use_data_path: # Exactly one of the two has to be None if we use it. assert (args.data_path is None) or (args.data_args_path is None) - use_per_split_data_path = any( - elt is not None - for elt in [args.train_data_path, args.valid_data_path, args.test_data_path]) or \ - args.per_split_data_args_path is not None + use_per_split_data_path = ( + any( + elt is not None + for elt in [args.train_data_path, args.valid_data_path, args.test_data_path] + ) + or args.per_split_data_args_path is not None + ) if use_per_split_data_path: - # Exactly one of the two has to be None if we use it. - assert any(elt is not None - for elt in [args.train_data_path, args.valid_data_path, args.test_data_path]) is False or \ - args.per_split_data_args_path is None + # Exactly one of the two has to be None if we use it. + assert ( + any( + elt is not None + for elt in [args.train_data_path, args.valid_data_path, args.test_data_path] + ) + is False + or args.per_split_data_args_path is None + ) if args.phase_transition_iterations: args.phase_transition_iterations = sorted( @@ -699,9 +729,7 @@ def validate_args(args, defaults={}): args, '_is_global_batch_size_explicitly_specified', args.global_batch_size is not None ) if args.step_batch_size_schedule is not None and is_global_batch_size_explicitly_specified: - raise ValueError( - 'Cannot specify both --step-batch-size-schedule and --global-batch-size' - ) + raise ValueError('Cannot specify both --step-batch-size-schedule and --global-batch-size') if args.global_batch_size is None: args.global_batch_size = args.micro_batch_size * args.data_parallel_size print_rank_0('setting global batch size to {}'.format(args.global_batch_size)) @@ -712,34 +740,36 @@ def validate_args(args, defaults={}): args.eval_global_batch_size = args.global_batch_size if args.eval_micro_batch_size is None: args.eval_micro_batch_size = args.micro_batch_size - # data_parallel_size is the replicate degree, so multiply the GTP-remat axis back in: evaluate() - # divides eval_global_batch_size by the same product to get its microbatch count, and without - # gtp_weight_remat_size here that division can silently floor (down to zero microbatches). - assert args.eval_global_batch_size % ( - args.eval_micro_batch_size * args.data_parallel_size * args.gtp_weight_remat_size - ) == 0, \ - f"eval_global_batch_size ({args.eval_global_batch_size}) must be divisible by " \ - f"eval_micro_batch_size ({args.eval_micro_batch_size}) * data_parallel_size ({args.data_parallel_size})" \ - f" * gtp_weight_remat_size ({args.gtp_weight_remat_size})" + assert ( + args.eval_global_batch_size % (args.eval_micro_batch_size * args.data_parallel_size) == 0 + ), ( + f"eval_global_batch_size ({args.eval_global_batch_size}) must be divisible by " + f"eval_micro_batch_size ({args.eval_micro_batch_size}) * data_parallel_size ({args.data_parallel_size})" + ) if args.perform_rl_step: num_generated_samples_per_inference_iteration = ( - args.grpo_samples_per_iteration * args.grpo_iterations) + args.grpo_samples_per_iteration * args.grpo_iterations + ) # Ensure that the number of prompts we collect is a multiple of the global batch size. - assert num_generated_samples_per_inference_iteration % args.global_batch_size == 0, \ - f"grpo_group_size * grpo_prompts_per_step * grpo_iterations should be divisible by global_batch_size" + assert ( + num_generated_samples_per_inference_iteration % args.global_batch_size == 0 + ), f"grpo_group_size * grpo_prompts_per_step * grpo_iterations should be divisible by global_batch_size" # For now only exit/checkpoint on iterations where we generate data. We don't currently # have a way to checkpoint the generated data. num_training_iterations_per_inference_iteration = ( - num_generated_samples_per_inference_iteration // args.global_batch_size) + num_generated_samples_per_inference_iteration // args.global_batch_size + ) if args.exit_interval is not None: - assert args.exit_interval % num_training_iterations_per_inference_iteration == 0, \ - f"exit_interval should be divisible by number of global batches per inference iteration." + assert ( + args.exit_interval % num_training_iterations_per_inference_iteration == 0 + ), f"exit_interval should be divisible by number of global batches per inference iteration." if args.save_interval is not None: - assert args.save_interval % num_training_iterations_per_inference_iteration == 0, \ - f"save_interval should be divisible by number of global batches per inference iteration." + assert ( + args.save_interval % num_training_iterations_per_inference_iteration == 0 + ), f"save_interval should be divisible by number of global batches per inference iteration." # === Hybrid layer pattern: deprecation handling and validation === @@ -772,6 +802,7 @@ def validate_args(args, defaults={}): get_hybrid_total_pipeline_segment_count, parse_hybrid_pattern, ) + sep = Symbols.MTP_SEPARATOR # Backward compat: convert legacy mtp_hybrid_override_pattern to unified format @@ -783,7 +814,9 @@ def validate_args(args, defaults={}): ): main_pattern = args.hybrid_layer_pattern or '' mtp_pattern = args.mtp_hybrid_override_pattern - args.hybrid_layer_pattern = main_pattern + sep + sep.join([mtp_pattern] * args.mtp_num_layers) + args.hybrid_layer_pattern = ( + main_pattern + sep + sep.join([mtp_pattern] * args.mtp_num_layers) + ) args.mtp_hybrid_override_pattern = None print_rank_0(f"Converted legacy MTP pattern to unified: {args.hybrid_layer_pattern}") @@ -877,7 +910,7 @@ def validate_args(args, defaults={}): f"MTP depth count ({inferred_mtp_num_layers}) in pattern " f"'{args.hybrid_layer_pattern}'. " f"Using the inferred value ({inferred_mtp_num_layers}).", - args.rank + args.rank, ) args.mtp_num_layers = inferred_mtp_num_layers @@ -888,25 +921,19 @@ def validate_args(args, defaults={}): + f"The supported position embedding types are rope and none." ) - if args.mtp_hsm and not (args.mtp_num_layers and args.mtp_num_layers >= 2): - warn_rank_0( - "--mtp-hsm needs at least two MTP layers to mix anything, but " - f"--mtp-num-layers is {args.mtp_num_layers}. Disabling Hidden State Mixing.", - args.rank, - ) - args.mtp_hsm = False - # Validate MTP args for hybrid vs non-hybrid models if args.hybrid_layer_pattern is not None: # Mamba/hybrid model MTP validation - if args.mtp_num_layers and not (args.hybrid_layer_pattern and sep in args.hybrid_layer_pattern): + if args.mtp_num_layers and not ( + args.hybrid_layer_pattern and sep in args.hybrid_layer_pattern + ): # Hybrid model wants MTP but no unified pattern - check for legacy args if args.mtp_hybrid_override_pattern is None: warn_rank_0( "Hybrid model with --mtp-num-layers but no MTP pattern. " "Use unified --hybrid-layer-pattern with '/' separator (e.g., 'M*M*/MM/MM') " "or legacy --mtp-hybrid-override-pattern for old checkpoints.", - args.rank + args.rank, ) else: # Non-hybrid (GPT) model MTP validation @@ -915,14 +942,11 @@ def validate_args(args, defaults={}): "--mtp-hybrid-override-pattern is for Mamba/hybrid models only. " "For GPT models, MTP replicates the main transformer layer structure. " "This argument will be ignored.", - args.rank + args.rank, ) # Infer use of MLA from unified pattern - if args.hybrid_layer_pattern and ( - Symbols.MLA in args.hybrid_layer_pattern - or Symbols.DS_ATTENTION in args.hybrid_layer_pattern - ): + if args.hybrid_layer_pattern and Symbols.DS_ATTENTION in args.hybrid_layer_pattern: args.multi_latent_attention = True # === End of hybrid layer pattern: deprecation handling and validation === @@ -944,7 +968,9 @@ def validate_args(args, defaults={}): if args.pipeline_model_parallel_layout is not None: # Parse the input flattened layout to a list and get the vpp size. # We will validate the layout more carefully in the TransformerConfig constructor. - num_stages = PipelineParallelLayerLayout.get_num_stages_from_str(args.pipeline_model_parallel_layout) + num_stages = PipelineParallelLayerLayout.get_num_stages_from_str( + args.pipeline_model_parallel_layout + ) assert num_stages % args.pipeline_model_parallel_size == 0, ( f"The length of pipeline_model_parallel_layout must be divisible" f" by pipeline_model_parallel_size ({num_stages=}," @@ -953,10 +979,15 @@ def validate_args(args, defaults={}): args.virtual_pipeline_model_parallel_size = num_stages // args.pipeline_model_parallel_size if args.virtual_pipeline_model_parallel_size == 1: args.virtual_pipeline_model_parallel_size = None - elif args.num_layers_per_virtual_pipeline_stage is not None or args.num_virtual_stages_per_pipeline_rank is not None: + elif ( + args.num_layers_per_virtual_pipeline_stage is not None + or args.num_virtual_stages_per_pipeline_rank is not None + ): if args.num_virtual_stages_per_pipeline_rank is None: - assert args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None, \ - 'please use --num-virtual-stages-per-pipeline-rank to specify virtual pipeline parallel degree when enable uneven pipeline parallelism' + assert ( + args.decoder_first_pipeline_num_layers is None + and args.decoder_last_pipeline_num_layers is None + ), 'please use --num-virtual-stages-per-pipeline-rank to specify virtual pipeline parallel degree when enable uneven pipeline parallelism' if args.num_layers is not None: num_layers = args.num_layers else: @@ -968,14 +999,19 @@ def validate_args(args, defaults={}): if args.account_for_loss_in_pipeline_split: num_layers += 1 - assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \ - 'number of layers of the model must be divisible pipeline model parallel size' - num_layers_per_pipeline_stage = num_layers // args.transformer_pipeline_model_parallel_size + assert ( + num_layers % args.transformer_pipeline_model_parallel_size == 0 + ), 'number of layers of the model must be divisible pipeline model parallel size' + num_layers_per_pipeline_stage = ( + num_layers // args.transformer_pipeline_model_parallel_size + ) - assert num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0, \ - 'number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage' - args.virtual_pipeline_model_parallel_size = num_layers_per_pipeline_stage // \ - args.num_layers_per_virtual_pipeline_stage + assert ( + num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0 + ), 'number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage' + args.virtual_pipeline_model_parallel_size = ( + num_layers_per_pipeline_stage // args.num_layers_per_virtual_pipeline_stage + ) else: args.virtual_pipeline_model_parallel_size = args.num_virtual_stages_per_pipeline_rank if args.virtual_pipeline_model_parallel_size == 1: @@ -985,7 +1021,10 @@ def validate_args(args, defaults={}): if args.hybrid_layer_pattern is None: args.virtual_pipeline_model_parallel_size = None - if args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None: + if ( + args.decoder_first_pipeline_num_layers is None + and args.decoder_last_pipeline_num_layers is None + ): # Divisibility check not applicable for T5 models which specify encoder_num_layers # and decoder_num_layers, or for hybrid models using --hybrid-layer-pattern. if args.num_layers is not None and args.hybrid_layer_pattern is None: @@ -997,59 +1036,73 @@ def validate_args(args, defaults={}): if args.account_for_loss_in_pipeline_split: num_layers += 1 - assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \ - 'Number of layers should be divisible by the pipeline-model-parallel size' + assert ( + num_layers % args.transformer_pipeline_model_parallel_size == 0 + ), 'Number of layers should be divisible by the pipeline-model-parallel size' if args.virtual_pipeline_model_parallel_size is not None: if args.overlap_p2p_comm: - assert args.pipeline_model_parallel_size > 1, \ - 'When interleaved schedule is used, pipeline-model-parallel size '\ + assert args.pipeline_model_parallel_size > 1, ( + 'When interleaved schedule is used, pipeline-model-parallel size ' 'should be greater than 1' + ) else: - assert args.pipeline_model_parallel_size > 2, \ - 'When interleaved schedule is used and p2p communication overlap is disabled, '\ - 'pipeline-model-parallel size should be greater than 2 to avoid having multiple '\ + assert args.pipeline_model_parallel_size > 2, ( + 'When interleaved schedule is used and p2p communication overlap is disabled, ' + 'pipeline-model-parallel size should be greater than 2 to avoid having multiple ' 'p2p sends and recvs between same 2 ranks per communication batch' + ) else: # Overlap P2P communication is disabled if not using the interleaved schedule. args.overlap_p2p_comm = False args.align_param_gather = False # Only print warning if PP size > 1. if args.rank == 0 and args.pipeline_model_parallel_size > 1: - print('WARNING: Setting args.overlap_p2p_comm and args.align_param_gather to False ' + print( + 'WARNING: Setting args.overlap_p2p_comm and args.align_param_gather to False ' 'since non-interleaved schedule does not support overlapping p2p communication ' - 'and aligned param AG') + 'and aligned param AG' + ) print_rank_0( f"Number of virtual stages per pipeline stage: {args.virtual_pipeline_model_parallel_size}" ) if args.overlap_param_gather: - assert args.use_distributed_optimizer or args.use_megatron_fsdp \ - or args.optimizer == 'dist_muon', \ - '--overlap-param-gather only supported with distributed optimizer, megatron fsdp, or dist_muon' - assert args.overlap_grad_reduce, \ - 'Must use --overlap-param-gather with --overlap-grad-reduce' + assert ( + args.use_distributed_optimizer + or args.use_megatron_fsdp + or args.optimizer == 'dist_muon' + ), '--overlap-param-gather only supported with distributed optimizer, megatron fsdp, or dist_muon' + assert ( + args.overlap_grad_reduce + ), 'Must use --overlap-param-gather with --overlap-grad-reduce' if args.use_torch_fsdp2: - assert is_torch_min_version("2.4.0"), \ - 'FSDP2 requires PyTorch >= 2.4.0 with FSDP 2 support.' - assert args.pipeline_model_parallel_size == 1, \ - '--use-torch-fsdp2 is not supported with pipeline parallelism' - assert args.expert_model_parallel_size == 1, \ - '--use-torch-fsdp2 is not supported with expert parallelism' - assert not args.use_distributed_optimizer, \ - "--use-torch-fsdp2 is not supported with MCore's distributed optimizer" - assert not args.gradient_accumulation_fusion, \ - '--use-torch-fsdp2 is not supported with gradient accumulation fusion' - assert args.ckpt_format in ('torch_dist', 'torch_dcp'), \ - '--use-torch-fsdp2 requires --ckpt-format torch_dist or torch_dcp' - assert args.untie_embeddings_and_output_weights, \ - '--use-torch-fsdp2 requires --untie-embeddings-and-output-weights' - assert not args.fp16, \ - '--use-torch-fsdp2 not supported with fp16 yet' - assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1", \ - 'FSDP always requires CUDA_DEVICE_MAX_CONNECTIONS value large than one' + assert is_torch_min_version("2.4.0"), 'FSDP2 requires PyTorch >= 2.4.0 with FSDP 2 support.' + assert ( + args.pipeline_model_parallel_size == 1 + ), '--use-torch-fsdp2 is not supported with pipeline parallelism' + assert ( + args.expert_model_parallel_size == 1 + ), '--use-torch-fsdp2 is not supported with expert parallelism' + assert ( + not args.use_distributed_optimizer + ), "--use-torch-fsdp2 is not supported with MCore's distributed optimizer" + assert ( + not args.gradient_accumulation_fusion + ), '--use-torch-fsdp2 is not supported with gradient accumulation fusion' + assert args.ckpt_format in ( + 'torch_dist', + 'torch_dcp', + ), '--use-torch-fsdp2 requires --ckpt-format torch_dist or torch_dcp' + assert ( + args.untie_embeddings_and_output_weights + ), '--use-torch-fsdp2 requires --untie-embeddings-and-output-weights' + assert not args.fp16, '--use-torch-fsdp2 not supported with fp16 yet' + assert ( + os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1" + ), 'FSDP always requires CUDA_DEVICE_MAX_CONNECTIONS value large than one' if args.fp8_param_gather and is_te_min_version("2.0.0"): args.fp8_param_gather = False @@ -1067,18 +1120,27 @@ def validate_args(args, defaults={}): ) if args.overlap_param_gather_with_optimizer_step: - assert args.use_distributed_optimizer, \ - '--overlap-param-gather-with-optimizer-step only supported with distributed optimizer' - assert args.overlap_param_gather, \ - 'Must use --overlap-param-gather-with-optimizer-step with --overlap-param-gather' - assert args.virtual_pipeline_model_parallel_size is not None, \ - '--overlap-param-gather-with-optimizer-step only supported with interleaved pipeline parallelism' - assert not args.use_dist_ckpt, \ - '--overlap-param-gather-with-optimizer-step not supported with distributed checkpointing yet' + assert ( + args.use_distributed_optimizer + ), '--overlap-param-gather-with-optimizer-step only supported with distributed optimizer' + assert ( + args.overlap_param_gather + ), 'Must use --overlap-param-gather-with-optimizer-step with --overlap-param-gather' + assert ( + args.virtual_pipeline_model_parallel_size is not None + ), '--overlap-param-gather-with-optimizer-step only supported with interleaved pipeline parallelism' + assert ( + not args.use_dist_ckpt + ), '--overlap-param-gather-with-optimizer-step not supported with distributed checkpointing yet' # Map string data-type to torch.dtype. dtype_map = { - 'fp32': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16, 'fp8': torch.uint8, 'auto': None, None: None, + 'fp32': torch.float32, + 'bf16': torch.bfloat16, + 'fp16': torch.float16, + 'fp8': torch.uint8, + 'auto': None, + None: None, } map_dtype = lambda d: d if isinstance(d, torch.dtype) else dtype_map[d] @@ -1089,19 +1151,24 @@ def validate_args(args, defaults={}): args.mamba_inference_conv_states_dtype = map_dtype(args.mamba_inference_conv_states_dtype) args.mamba_inference_ssm_states_dtype = map_dtype(args.mamba_inference_ssm_states_dtype) args.mamba_training_ssm_states_dtype = map_dtype(args.mamba_training_ssm_states_dtype) - args.logit_dtype = map_dtype(getattr(args, 'logit_dtype', None)) args.megatron_fsdp_main_params_dtype = map_dtype(args.megatron_fsdp_main_params_dtype) args.megatron_fsdp_main_grads_dtype = map_dtype(args.megatron_fsdp_main_grads_dtype) args.megatron_fsdp_grad_comm_dtype = map_dtype(args.megatron_fsdp_grad_comm_dtype) if args.fp8_param_gather: - assert args.use_distributed_optimizer or args.use_torch_fsdp2 or args.use_megatron_fsdp or not torch.is_grad_enabled(), \ - '--fp8-param-gather only supported with distributed optimizer, torch fsdp2, megatron fsdp, or inference mode' + assert ( + args.use_distributed_optimizer + or args.use_torch_fsdp2 + or args.use_megatron_fsdp + or not torch.is_grad_enabled() + ), '--fp8-param-gather only supported with distributed optimizer, torch fsdp2, megatron fsdp, or inference mode' # FP4 and FP8 are mutually exclusive if args.fp4 and args.fp8: - raise ValueError("--fp4-format and --fp8-format cannot be used simultaneously. Please choose one.") + raise ValueError( + "--fp4-format and --fp8-format cannot be used simultaneously. Please choose one." + ) # FP4 param requires FP4 mode if args.fp4_param_gather and not args.fp4: @@ -1109,7 +1176,9 @@ def validate_args(args, defaults={}): # FP4 requires TE >= 2.7.0.dev0 if args.fp4 and not is_te_min_version("2.7.0.dev0"): - raise ValueError("--fp4-format requires Transformer Engine >= 2.7.0.dev0 for NVFP4BlockScaling support.") + raise ValueError( + "--fp4-format requires Transformer Engine >= 2.7.0.dev0 for NVFP4BlockScaling support." + ) if ( args.fp8_recipe == 'mxfp8' @@ -1118,24 +1187,29 @@ def validate_args(args, defaults={}): ): raise ValueError("MXFP8 with inference optimized layers requires FlashInfer >= 0.6.4") + if args.inference_dynamic_batching_sampling_backend == 'flashinfer': + try: + import flashinfer # noqa: F401 + except ImportError as e: + raise ImportError( + "--inference-dynamic-batching-sampling-backend=flashinfer requires " + "the flashinfer package; install it or pass " + "--inference-dynamic-batching-sampling-backend=torch." + ) from e + if args.use_megatron_fsdp: # NOTE: The flag `use_custom_fsdp` is deprecated and will be removed in future versions. # Please use `use_megatron_fsdp` instead, as all functionality will be migrated there. # Future updates will drop support for `use_custom_fsdp` to avoid confusion. args.use_custom_fsdp = True - if args.megatron_fsdp_version == 2: - assert not args.use_distributed_optimizer, \ - '--megatron-fsdp-version 2 is not compatible with --use-distributed-optimizer' - else: - # Megatron-FSDP v1 requires the DistributedOptimizer. - if not args.use_distributed_optimizer: - warn_rank_0( - 'Megatron-FSDP v1 is only compatible with --use-distributed-optimizer. ' - 'Using DistributedOptimizer...', - args.rank, - ) - args.use_distributed_optimizer = True + # Megatron-FSDP requires the DistributedOptimizer. + if not args.use_distributed_optimizer: + warn_rank_0( + 'Megatron-FSDP is only compatible with --use-distributed-optimizer. Using DistributedOptimizer...', + args.rank, + ) + args.use_distributed_optimizer = True # Optimizer step MXFP8 buffer operation that is not relevant or supported for Megatron-FSDP. args.reuse_grad_buf_for_mxfp8_param_ag = False if args.moe_single_grouped_weight or args.moe_single_grouped_bias: @@ -1150,8 +1224,10 @@ def validate_args(args, defaults={}): "GroupedTensor param buffers." ) # Optimizer compatibility check. - assert args.optimizer in ('sgd', 'adam'), \ - f"Megatron-FSDP does not support the {args.optimizer} optimizer yet." + assert args.optimizer in ( + 'sgd', + 'adam', + ), f"Megatron-FSDP does not support the {args.optimizer} optimizer yet." if ( args.data_parallel_sharding_strategy in ["optim_grads_params", "optim_grads"] @@ -1163,21 +1239,40 @@ def validate_args(args, defaults={}): ) if args.data_parallel_sharding_strategy == "optim_grads_params": - assert args.check_weight_hash_across_dp_replicas_interval is None, \ - 'check_weight_hash_across_dp_replicas_interval is not supported with optim_grads_params' + assert ( + args.check_weight_hash_across_dp_replicas_interval is None + ), 'check_weight_hash_across_dp_replicas_interval is not supported with optim_grads_params' - assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1", \ - 'FSDP requires CUDA_DEVICE_MAX_CONNECTIONS > 1 or unset.' + assert ( + os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1" + ), 'FSDP requires CUDA_DEVICE_MAX_CONNECTIONS > 1 or unset.' + + assert ( + args.ckpt_format == "fsdp_dtensor" + ), "Megatron-FSDP requires the `fsdp_dtensor` checkpointing format." - assert args.ckpt_format == "fsdp_dtensor", \ - "Megatron-FSDP requires the `fsdp_dtensor` checkpointing format." + if args.megatron_fsdp_prefetch_recompute_forward_weights: + assert args.data_parallel_sharding_strategy == "optim_grads_params", ( + "--megatron-fsdp-prefetch-recompute-forward-weights is only supported " + 'with --data-parallel-sharding-strategy optim_grads_params.' + ) + assert args.recompute_granularity == "full", ( + "--megatron-fsdp-prefetch-recompute-forward-weights is only supported " + "with full activation recomputation." + ) + assert not args.overlap_moe_expert_parallel_comm, ( + "--megatron-fsdp-prefetch-recompute-forward-weights is not supported " + "with --overlap-moe-expert-parallel-comm." + ) if args.nccl_ub: # In Megatron-LM, required implementation for manual registration is already provided. # So we enable the manual registration by default when nccl-ub and use_megatron_fsdp is set. args.fsdp_manual_registration = True args.fsdp_double_buffer = True - warn_rank_0('FSDP double buffer and manual registration is enabled by default when --nccl-ub is enabled!') + warn_rank_0( + 'FSDP double buffer and manual registration is enabled by default when --nccl-ub is enabled!' + ) if args.megatron_fsdp_max_pool_double_buffer: # MaxPoolAllocator is a type of FSDP double buffer. @@ -1188,9 +1283,18 @@ def validate_args(args, defaults={}): "Meta device initialization (init_model_with_meta_device=True) is not " "supported or necessary for the 'no_shard' / 0 sharding strategy." ) + else: + assert not args.megatron_fsdp_prefetch_recompute_forward_weights, ( + "--megatron-fsdp-prefetch-recompute-forward-weights requires " "--use-megatron-fsdp." + ) + assert not args.megatron_fsdp_cache_param_bucket_views, ( + "--megatron-fsdp-cache-param-bucket-views requires " "--use-megatron-fsdp." + ) if args.fsdp_manual_registration: - assert args.use_megatron_fsdp, "FSDP manual registration is only supported with Megatron FSDP." + assert ( + args.use_megatron_fsdp + ), "FSDP manual registration is only supported with Megatron FSDP." assert args.nccl_ub, "FSDP manual registration is only supported with --nccl-ub argument." # Parameters dtype. @@ -1202,55 +1306,48 @@ def validate_args(args, defaults={}): # where NaNs in grads / loss are signal to the loss scaler. if not args.loss_scale: args.check_for_nan_in_loss_and_grad = False - warn_rank_0('Setting args.check_for_nan_in_loss_and_grad to False since ' - 'dynamic loss scaling is being used') + warn_rank_0( + 'Setting args.check_for_nan_in_loss_and_grad to False since ' + 'dynamic loss scaling is being used' + ) if args.bf16: assert not args.fp16 args.params_dtype = torch.bfloat16 # bfloat16 requires gradient accumulation and all-reduce to # be done in fp32. if args.accumulate_allreduce_grads_in_fp32: - assert args.main_grads_dtype == torch.float32, \ - "--main-grads-dtype can only be fp32 when --accumulate-allreduce-grads-in-fp32 is set" + assert ( + args.main_grads_dtype == torch.float32 + ), "--main-grads-dtype can only be fp32 when --accumulate-allreduce-grads-in-fp32 is set" if args.grad_reduce_in_bf16: args.accumulate_allreduce_grads_in_fp32 = False elif not args.accumulate_allreduce_grads_in_fp32 and args.main_grads_dtype == torch.float32: args.accumulate_allreduce_grads_in_fp32 = True print_rank_0('accumulate and all-reduce gradients in fp32 for bfloat16 data type.') - - if args.accumulate_allreduce_grads_in_fp32: - # The FP32-accumulation reduce-scatters only exist to recover FP32 accumulation from a - # lower-precision wire dtype. With FP32 main_grads the plain ring reduce-scatter already - # sums in FP32 on both axes, so they buy no precision and cost one extra - # unsharded-wgrad-sized scratch buffer per in-flight reduce-scatter. - for arg_name in ( - 'ddp_reduce_scatter_with_fp32_accumulation', - 'gtp_remat_reduce_scatter_with_fp32_accumulation', - ): - if getattr(args, arg_name, False): - setattr(args, arg_name, False) - warn_rank_0( - f"Setting args.{arg_name} to False since " - "--accumulate-allreduce-grads-in-fp32 already reduces in fp32" - ) - if args.cuda_graph_impl == "full_iteration": - assert not args.check_for_nan_in_loss_and_grad, \ - "--no-check-for-nan-in-loss-and-grad should be set with --cuda-graph-impl=full_iteration for training." + assert ( + not args.check_for_nan_in_loss_and_grad + ), "--no-check-for-nan-in-loss-and-grad should be set with --cuda-graph-impl=full_iteration for training." if args.inference_cuda_graph_scope == InferenceCudaGraphScope.block: if args.fp8 is not None: - assert args.transformer_impl == "inference_optimized", \ - "fp8 with --inference-cuda-graph-scope=block is only supported with " \ + assert args.transformer_impl == "inference_optimized", ( + "fp8 with --inference-cuda-graph-scope=block is only supported with " "--transformer-impl=inference_optimized" - assert args.fp8_recipe == "mxfp8", \ - "Only --fp8-recipe=mxfp8 is supported with --inference-cuda-graph-scope=block" + ) + assert ( + args.fp8_recipe == "mxfp8" + ), "Only --fp8-recipe=mxfp8 is supported with --inference-cuda-graph-scope=block" if args.cuda_graph_impl == 'local': - assert args.inference_dynamic_batching_num_cuda_graphs > 0 or args.inference_dynamic_batching_num_cuda_graphs == -1, \ - 'inference_dynamic_batching_num_cuda_graphs should be a positive integer or -1' \ + assert ( + args.inference_dynamic_batching_num_cuda_graphs > 0 + or args.inference_dynamic_batching_num_cuda_graphs == -1 + ), ( + 'inference_dynamic_batching_num_cuda_graphs should be a positive integer or -1' '-1 means that we will automatically determine the number of CUDA graphs to capture based on the `max_requests` value.' + ) print_rank_0('using {} for parameters ...'.format(args.params_dtype)) @@ -1272,43 +1369,39 @@ def validate_args(args, defaults={}): if args.train_iters and not args.skip_train: # If we use iteration-based training, make sure the # sample-based options are off. - assert args.train_samples is None, \ - 'expected iteration-based training' - assert args.lr_decay_samples is None, \ - 'expected iteration-based learning rate decay' - assert args.lr_warmup_samples == 0, \ - 'expected iteration-based learning rate warmup' + assert args.train_samples is None, 'expected iteration-based training' + assert args.lr_decay_samples is None, 'expected iteration-based learning rate decay' + assert args.lr_warmup_samples == 0, 'expected iteration-based learning rate warmup' if args.lr_warmup_fraction is not None: - assert args.lr_warmup_iters == 0, \ - 'can only specify one of lr-warmup-fraction and lr-warmup-iters' + assert ( + args.lr_warmup_iters == 0 + ), 'can only specify one of lr-warmup-fraction and lr-warmup-iters' # Sample-based training. if args.train_samples and not args.skip_train: # If we use sample-based training, make sure the # iteration-based options are off. - assert args.train_iters is None, \ - 'expected sample-based training' - assert args.lr_decay_iters is None, \ - 'expected sample-based learning rate decay' - assert args.lr_warmup_iters == 0, \ - 'expected sample-based learnig rate warmup' + assert args.train_iters is None, 'expected sample-based training' + assert args.lr_decay_iters is None, 'expected sample-based learning rate decay' + assert args.lr_warmup_iters == 0, 'expected sample-based learnig rate warmup' if args.lr_warmup_fraction is not None: - assert args.lr_warmup_samples == 0, \ - 'can only specify one of lr-warmup-fraction ' \ - 'and lr-warmup-samples' + assert args.lr_warmup_samples == 0, ( + 'can only specify one of lr-warmup-fraction ' 'and lr-warmup-samples' + ) if args.num_layers is not None: - assert args.encoder_num_layers is None, \ - 'cannot have both num-layers and encoder-num-layers specified' + assert ( + args.encoder_num_layers is None + ), 'cannot have both num-layers and encoder-num-layers specified' args.encoder_num_layers = args.num_layers else: - assert args.encoder_num_layers is not None, \ - 'either num-layers or encoder-num-layers should be specified' + assert ( + args.encoder_num_layers is not None + ), 'either num-layers or encoder-num-layers should be specified' args.num_layers = args.encoder_num_layers # Check required arguments. - required_args = ['num_layers', 'hidden_size', 'num_attention_heads', - 'max_position_embeddings'] + required_args = ['num_layers', 'hidden_size', 'num_attention_heads', 'max_position_embeddings'] for req_arg in required_args: _check_arg_is_not_none(args, req_arg) @@ -1329,9 +1422,10 @@ def validate_args(args, defaults={}): args.kv_channels = args.hidden_size // args.num_attention_heads if args.seq_length is not None and args.context_parallel_size > 1: - assert args.seq_length % (args.context_parallel_size * 2) == 0, \ - 'seq-length should be a multiple of 2 * context-parallel-size ' \ + assert args.seq_length % (args.context_parallel_size * 2) == 0, ( + 'seq-length should be a multiple of 2 * context-parallel-size ' 'if context-parallel-size > 1.' + ) if getattr(args, 'dataloader_inter_document_masking', False): # The dataset omits attention_mask when inter-document masking is @@ -1347,9 +1441,10 @@ def validate_args(args, defaults={}): args.seq_length = args.encoder_seq_length if args.seq_length is not None: - assert args.max_position_embeddings >= args.seq_length, \ - f"max_position_embeddings ({args.max_position_embeddings}) must be greater than " \ + assert args.max_position_embeddings >= args.seq_length, ( + f"max_position_embeddings ({args.max_position_embeddings}) must be greater than " f"or equal to seq_length ({args.seq_length})." + ) if args.decoder_seq_length is not None: assert args.max_position_embeddings >= args.decoder_seq_length if args.lr is not None: @@ -1368,8 +1463,9 @@ def validate_args(args, defaults={}): if args.fp16_lm_cross_entropy: assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.' if args.fp32_residual_connection: - assert args.fp16 or args.bf16, \ - 'residual connection in fp32 only supported when using fp16 or bf16.' + assert ( + args.fp16 or args.bf16 + ), 'residual connection in fp32 only supported when using fp16 or bf16.' if args.moe_grouped_gemm: dc = torch.cuda.get_device_capability() @@ -1398,30 +1494,33 @@ def validate_args(args, defaults={}): # Persistent fused layer norm. if not is_torch_min_version("1.11.0a0"): args.no_persist_layer_norm = True - print_rank_0('Persistent fused layer norm kernel is supported from ' - 'pytorch v1.11 (nvidia pytorch container paired with v1.11). ' - 'Defaulting to no_persist_layer_norm=True') + print_rank_0( + 'Persistent fused layer norm kernel is supported from ' + 'pytorch v1.11 (nvidia pytorch container paired with v1.11). ' + 'Defaulting to no_persist_layer_norm=True' + ) # Activation recomputing. if args.distribute_saved_activations: - assert args.tensor_model_parallel_size > 1, 'can distribute ' \ - 'recomputed activations only across tensor model ' \ - 'parallel groups' - assert args.recompute_granularity == 'full', \ - 'distributed recompute activations is only '\ - 'application to full recompute granularity' - assert args.recompute_method is not None, \ - 'for distributed recompute activations to work you '\ - 'need to use a recompute method ' - assert is_torch_min_version("1.10.0a0"), \ - 'distributed recompute activations are supported for pytorch ' \ - 'v1.10 and above (Nvidia Pytorch container >= 21.07). Current ' \ + assert args.tensor_model_parallel_size > 1, ( + 'can distribute ' 'recomputed activations only across tensor model ' 'parallel groups' + ) + assert args.recompute_granularity == 'full', ( + 'distributed recompute activations is only ' 'application to full recompute granularity' + ) + assert args.recompute_method is not None, ( + 'for distributed recompute activations to work you ' 'need to use a recompute method ' + ) + assert is_torch_min_version("1.10.0a0"), ( + 'distributed recompute activations are supported for pytorch ' + 'v1.10 and above (Nvidia Pytorch container >= 21.07). Current ' f'pytorch version is v{get_torch_version()}.' + ) if args.recompute_granularity == 'selective': - assert args.recompute_method is None, \ - 'recompute method is not yet supported for ' \ - 'selective recomputing granularity' + assert args.recompute_method is None, ( + 'recompute method is not yet supported for ' 'selective recomputing granularity' + ) # disable sequence parallelism when tp=1 # to avoid change in numerics when @@ -1435,36 +1534,90 @@ def validate_args(args, defaults={}): args.sequence_parallel = False if args.tp_comm_overlap: - assert args.sequence_parallel == True, 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' - - if args.hybrid_context_parallel: - assert not args.pipeline_model_parallel_size > 1, 'Hybrid context parallelism not supported with pipeline parallelism' - assert not args.enable_cuda_graph, 'Hybrid context parallelism not supported with CUDA Graph' - assert not args.use_megatron_fsdp, 'Hybrid context parallelism not supported with Megatron FSDP' - assert args.dataloader_type == 'single', 'Hybrid context parallelism only supported with single dataloader type' - assert args.calculate_per_token_loss, 'Hybrid context parallelism must be used with --calculate-per-token-loss' - - # Support for variable sequence lengths across batches/microbatches. - # set it if the dataloader supports generation of variable sequence lengths - # across batches/microbatches. Due to additional communication overhead - # during pipeline parallelism, it should not be set if sequence length - # is constant during training. - args.variable_seq_lengths = False - if args.mock_data and args.sft and args.sft_mock_dataset_config_json is None: - args.sft_mock_dataset_config_json = json.dumps( - { - "mode": "distribution", - "type": "lognormal", - "min_seq_len": args.seq_length // 2, - "max_seq_len": args.seq_length, - "mean_seq_len": args.seq_length // 4 * 3, - "lognormal_sigma": 1.1, - } + assert ( + args.sequence_parallel == True + ), 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' + + if args.dynamic_context_parallel: + assert ( + not args.enable_cuda_graph + ), 'Dynamic context parallelism not supported with CUDA Graph' + assert ( + not args.use_megatron_fsdp + ), 'Dynamic context parallelism not supported with Megatron FSDP' + assert ( + args.dataloader_type == 'single' + ), 'Dynamic context parallelism only supported with single dataloader type' + assert ( + args.calculate_per_token_loss + ), 'Dynamic context parallelism must be used with --calculate-per-token-loss' + if args.sequence_packing_scheduler is None: + args.sequence_packing_scheduler = 'default_dynamic_cp' + if args.sequence_packing_scheduler != 'default_dynamic_cp': + raise ValueError( + 'Dynamic context parallelism requires ' + 'sequence_packing_scheduler=default_dynamic_cp' + ) + + dp_cp_size = args.data_parallel_size * args.context_parallel_size + assert args.min_dynamic_context_parallel_size <= dp_cp_size, ( + f'min_dynamic_context_parallel_size ({args.min_dynamic_context_parallel_size}) ' + f'must be <= dp_size * cp_size ({dp_cp_size})' + ) + + import warnings + + warnings.warn( + f"Dynamic CP enabled: dp_size * context_parallel_size=" + f"{args.data_parallel_size * args.context_parallel_size} " + f"will be used as the maximum dynamic CP group size. " + f"Dynamic CP groups will range from " + f"min_dynamic_context_parallel_size={args.min_dynamic_context_parallel_size} " + f"to {args.data_parallel_size * args.context_parallel_size}." + ) + + if getattr(args, 'pad_packed_seq_alignment', None) is not None: + args.pad_packed_seq_alignment = _parse_pad_packed_seq_alignment( + args.pad_packed_seq_alignment ) + if args.max_seqlen_per_dp_cp_rank is None: + raise ValueError( + '--max-seqlen-per-dp-cp-rank must be set when ' + '--pad-packed-seq-alignment is enabled.' + ) + if args.pad_packed_seq_alignment != 'max': + if args.pad_packed_seq_alignment <= 0: + raise ValueError( + "--pad-packed-seq-alignment must be 'max' or a positive integer " "alignment." + ) + if args.pad_packed_seq_alignment > args.max_seqlen_per_dp_cp_rank: + raise ValueError( + '--pad-packed-seq-alignment must not exceed ' + f'--max-seqlen-per-dp-cp-rank ({args.max_seqlen_per_dp_cp_rank}), ' + f'got {args.pad_packed_seq_alignment}.' + ) + + if args.cuda_graph_impl != "none" and ( + args.sequence_packing_scheduler is not None or args.dynamic_context_parallel + ): + if getattr(args, 'pad_packed_seq_alignment', None) is None: + raise ValueError('THD CUDA Graph requires --pad-packed-seq-alignment to be set.') + if ( + args.pad_packed_seq_alignment != 'max' + and args.pad_packed_seq_alignment != args.max_seqlen_per_dp_cp_rank + ): + raise ValueError( + "THD CUDA Graph requires --pad-packed-seq-alignment='max' " + 'or --pad-packed-seq-alignment equal to ' + f'--max-seqlen-per-dp-cp-rank ({args.max_seqlen_per_dp_cp_rank}), ' + f'got {args.pad_packed_seq_alignment}.' + ) + # disable async_tensor_model_parallel_allreduce when # model parallel memory optimization is enabled - if (args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1) \ - and get_device_arch_version() < 10: + if ( + args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1 + ) and get_device_arch_version() < 10: # CUDA_DEVICE_MAX_CONNECTIONS requirement no longer exists since the Blackwell architecture if args.use_torch_fsdp2 or args.use_megatron_fsdp: fsdp_impl = "Torch-FSDP2" if args.use_torch_fsdp2 else "Megatron-FSDP" @@ -1487,9 +1640,10 @@ def validate_args(args, defaults={}): args.rank, ) else: - assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') == "1", \ - "Using tensor model parallelism or context parallelism require setting the environment variable " \ + assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') == "1", ( + "Using tensor model parallelism or context parallelism require setting the environment variable " "CUDA_DEVICE_MAX_CONNECTIONS to 1" + ) # Setting FSDP communication groups for high priority streams for Blackwell and later architectures # Assigning high priority to communication streams ensures that communication kernels are scheduled @@ -1497,123 +1651,9 @@ def validate_args(args, defaults={}): if args.use_torch_fsdp2 or args.use_megatron_fsdp and get_device_arch_version() >= 10: if 'dp_cp' not in args.high_priority_stream_groups: args.high_priority_stream_groups.append('dp_cp') - if args.expert_model_parallel_size > 1 and 'ep_dp' not in args.high_priority_stream_groups: + if args.expert_model_parallel_size > 1 and 'ep_dp' not in args.high_priority_stream_groups: args.high_priority_stream_groups.append('ep_dp') - - # Derive the internal gtp_weight_remat_size from the user-facing - # --tensor-parallel-num-weight-shards. gtp_weight_remat_size has no CLI flag (it is excluded - # from argument generation), so it is set here as a fresh attribute on args before it is - # consumed below (and in initialize/training, which read args.gtp_weight_remat_size directly). - # Mirrors ModelParallelConfig.__post_init__. - from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards - (args.tensor_parallel_num_weight_shards, args.gtp_weight_remat_size) = ( - resolve_tensor_parallel_weight_shards( - args.tensor_model_parallel_size, - args.tensor_parallel_num_weight_shards, - getattr(args, "gtp_weight_remat_size", 1), - ) - ) - # Same for the expert layers: derive the internal expert_gtp_weight_remat_size from the - # user-facing --expert-tensor-parallel-num-weight-shards (expert_tensor_parallel_size is - # defaulted earlier in validate_args). expert_gtp_weight_remat_size has no CLI flag. - (args.expert_tensor_parallel_num_weight_shards, args.expert_gtp_weight_remat_size) = ( - resolve_tensor_parallel_weight_shards( - args.expert_tensor_parallel_size, - args.expert_tensor_parallel_num_weight_shards, - getattr(args, "expert_gtp_weight_remat_size", 1), - ) - ) - - if args.gtp_weight_remat_size > 1 or args.expert_gtp_weight_remat_size > 1: - if args.fp4 and not args.fp4_param_gather: - raise ValueError( - "GTP (--tensor-parallel-num-weight-shards / " - "--expert-tensor-parallel-num-weight-shards > 1) with --fp4-format requires " - "--fp4-param-gather so NVFP4 weights are all-gathered as native NVFP4." - ) - gtp_weight_remat_size = args.gtp_weight_remat_size - egtp_weight_remat_size = args.expert_gtp_weight_remat_size - if get_device_arch_version() >= 10: - # Setting GTP communication groups for high priority streams for Blackwell and later - # architectures. Assigning high priority to communication streams ensures that - # communication kernels are scheduled with higher priority, minimizing the exposed - # communication when it is overlapped with other computation kernels. - if 'gtp_remat' not in args.high_priority_stream_groups: - args.high_priority_stream_groups.append('gtp_remat') - warn_rank_0("Setting 'gtp_remat' group for high priority streams.") - if ( - egtp_weight_remat_size > 1 - and 'expt_gtp_remat' not in args.high_priority_stream_groups - ): - args.high_priority_stream_groups.append('expt_gtp_remat') - warn_rank_0("Setting 'expt_gtp_remat' group for high priority streams.") - - # Sanity check for 'CUDA_GRAPHS_USE_NODE_PRIORITY'. - if args.cuda_graph_impl != "none": - assert os.environ.get('CUDA_GRAPHS_USE_NODE_PRIORITY') == "1", \ - 'GTP requires CUDA_GRAPHS_USE_NODE_PRIORITY=1 to make sure fine-grained GTP ' \ - 'comms can be well overlapped with GEMMs when CudaGraph is enabled for ' \ - 'Blackwell and later architecture.' - - # Sanity check for 'NCCL_PROTO'. - if os.environ.get('NCCL_PROTO', '').lower() == "simple": - warn_rank_0( - "Generally GTP prefers 'NCCL_PROTO=LL128 or LL' while get 'NCCL_PROTO=simple', " - "force setting NCCL_PROTO=Simple might introduce bad perf." - ) - - assert not args.ddp_average_in_collective, ( - "GTP requires --ddp-average-in-collective off (the default); averaged collectives " - "would need per-buffer 1/gtp_remat scaling." - ) - - assert args.ckpt_format in ('torch', 'torch_dist'), ( - f"GTP supports only --ckpt-format 'torch' (legacy) or 'torch_dist', got " - f"'{args.ckpt_format}'." - ) - assert not ( - getattr(args, 'dist_ckpt_optim_fully_reshardable', False) - and getattr(args, 'distrib_optim_fully_reshardable_mem_efficient', False) - ), ( - "GTP does not support the distributed-optimizer fully-reshardable + " - "mem-efficient checkpoint mode. Disable " - "--distrib-optim-fully-reshardable-mem-efficient (or " - "--dist-ckpt-optim-fully-reshardable)." - ) - - # GTP with the mxfp8 recipe requires --fp8-param-gather: GTP keeps no bf16 weight and - # relies on the optimizer maintaining the fp8 shard (the forward all-gathers fp8 and does - # not re-quantize). Without fp8-param-gather the fp8 forward weight would never be updated. - if getattr(args, 'fp8_recipe', None) == 'mxfp8': - assert getattr(args, 'fp8_param_gather', False), ( - "GTP + mxfp8 requires --fp8-param-gather (the optimizer maintains the fp8 shard; " - "GTP does not keep or re-quantize a bf16 weight)." - ) - # MXFP8 params cannot be mapped into the contiguous param buffer (TE's - # replace_raw_data does not support the MXFP8 tile-scaling layout), so the param - # all-gather must reuse the grad buffer instead. - assert getattr(args, 'reuse_grad_buf_for_mxfp8_param_ag', False), ( - "GTP + mxfp8 + --fp8-param-gather requires --reuse-grad-buf-for-mxfp8-param-ag " - "(MXFP8 params keep their own quantized storage; mapping them into the param " - "buffer via replace_raw_data is unsupported)." - ) - - # GTP symmetric memory registers pools with symmetric=True (NVLS needs symmetric - # windows), which contradicts --disable-symmetric-registration. - if getattr(args, 'gtp_remat_nccl_ub', False) or getattr(args, 'gtp_expert_remat_nccl_ub', False): - assert not getattr(args, 'disable_symmetric_registration', False), ( - "--gtp-remat-nccl-ub/--gtp-expert-remat-nccl-ub require symmetric window registration and " - "cannot be combined with --disable-symmetric-registration." - ) - if getattr(args, 'gtp_remat_reduce_scatter_with_fp32_accumulation', False): - print_rank_0( - "WARNING: --gtp-remat-nccl-ub/--gtp-expert-remat-nccl-ub take precedence over " - "--gtp-remat-reduce-scatter-with-fp32-accumulation on their groups: NVLS " - "symmetric reduce-scatters accumulate in fp32 in-switch " - "(NCCL multimem.ld_reduce .acc::f32)." - ) - # Disable bias gelu fusion if we are disabling bias altogether if not args.add_bias_linear: args.bias_gelu_fusion = False @@ -1623,22 +1663,24 @@ def validate_args(args, defaults={}): args.add_qkv_bias = True if args.qk_clip: - assert is_te_min_version("2.9.0"), \ - '--qk-clip is only supported with TE >= 2.9.0.' - assert 0.0 < args.qk_clip_alpha < 1.0, \ - '--qk-clip-alpha must be between 0.0 and 1.0 when using --qk-clip.' - assert args.qk_clip_threshold > 0, \ - '--qk-clip-threshold must be greater than 0 when using --qk-clip.' + assert is_te_min_version("2.9.0"), '--qk-clip is only supported with TE >= 2.9.0.' + assert ( + 0.0 < args.qk_clip_alpha < 1.0 + ), '--qk-clip-alpha must be between 0.0 and 1.0 when using --qk-clip.' + assert ( + args.qk_clip_threshold > 0 + ), '--qk-clip-threshold must be greater than 0 when using --qk-clip.' # decoupled log max attention logit check if args.log_max_attention_logit: - assert is_te_min_version("2.9.0"), \ - '--log-max-attention-logit is only supported with TE >= 2.9.0.' + assert is_te_min_version( + "2.9.0" + ), '--log-max-attention-logit is only supported with TE >= 2.9.0.' # Legacy RoPE arguments if args.use_rotary_position_embeddings: args.position_embedding_type = 'rope' - if args.position_embedding_type != 'rope': + if args.position_embedding_type not in ('rope', 'mrope'): args.apply_rope_fusion = False # Would just need to add 'NoPE' as a position_embedding_type to support this, but for now @@ -1654,8 +1696,9 @@ def validate_args(args, defaults={}): # MultiModal rotary embeddings arguments if args.position_embedding_type == "mrope": - assert args.mrope_section is not None, \ - '--mrope-section should be set when using --position-embedding-type mrope.' + assert ( + args.mrope_section is not None + ), '--mrope-section should be set when using --position-embedding-type mrope.' # MoE Spec check if args.num_experts == 0: @@ -1665,13 +1708,19 @@ def validate_args(args, defaults={}): warn_rank_0("moe_ffn_hidden_size is not set, using ffn_hidden_size for MoE instead.") # Expert parallelism check - if args.expert_model_parallel_size > 1: - assert args.num_experts is not None, "num_experts must be non None to use expert model parallelism" - assert args.num_experts % args.expert_model_parallel_size == 0, \ - "Number of experts should be a multiple of expert model parallel_size." + if args.expert_model_parallel_size > 1: + assert ( + args.num_experts is not None + ), "num_experts must be non None to use expert model parallelism" + assert ( + args.num_experts % args.expert_model_parallel_size == 0 + ), "Number of experts should be a multiple of expert model parallel_size." # MoE router check - if isinstance(args.moe_router_load_balancing_type, list) and len(args.moe_router_load_balancing_type) == 1: + if ( + isinstance(args.moe_router_load_balancing_type, list) + and len(args.moe_router_load_balancing_type) == 1 + ): args.moe_router_load_balancing_type = args.moe_router_load_balancing_type[0] if isinstance(args.moe_aux_loss_coeff, list) and len(args.moe_aux_loss_coeff) == 1: args.moe_aux_loss_coeff = args.moe_aux_loss_coeff[0] @@ -1679,40 +1728,81 @@ def validate_args(args, defaults={}): # torch_dcp (torch.distributed.checkpoint) checkpointing format checks. if args.ckpt_format == "torch_dcp": assert args.use_torch_fsdp2, "--ckpt-format torch_dcp is only tested with FSDP." - assert args.tensor_model_parallel_size <= 1, \ - "--ckpt-format torch_dcp is not tested with megatron tensor parallelism." - assert args.pipeline_model_parallel_size <= 1, \ - "--ckpt-format torch_dcp is not tested with megatron pipeline parallelism." + assert ( + args.tensor_model_parallel_size <= 1 + ), "--ckpt-format torch_dcp is not tested with megatron tensor parallelism." + assert ( + args.pipeline_model_parallel_size <= 1 + ), "--ckpt-format torch_dcp is not tested with megatron pipeline parallelism." # fsdp_dtensor checkpointing format checks. if args.ckpt_format == "fsdp_dtensor": - assert args.use_megatron_fsdp, "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP." - - # Scheduler-name and max-seqlen validation live in - # ModelParallelConfig.__post_init__; only the buffer-size check stays here - # because seq_length is not a core config field. The None case for - # max_seqlen_per_dp_cp_rank is rejected by the config check. - if args.sequence_packing_scheduler is not None: - args.variable_seq_lengths = True - # Packed microbatches carry different numbers of valid tokens, so the - # default per-microbatch loss averaging would weight tokens unevenly - # depending on how samples happened to be packed (same reasoning as - # the hybrid-context-parallel check above). - assert args.calculate_per_token_loss, ( - 'Sequence packing must be used with --calculate-per-token-loss' + assert ( + args.use_megatron_fsdp + ), "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP." + + # --use-varlen-dataset: independent of --sft. Cannot be combined with --sft + # because they are mutually-exclusive top-level dataset selectors that both + # drive the packed-sequence (THD) path. + if args.use_varlen_dataset: + assert not args.sft, ( + "--use-varlen-dataset and --sft are mutually exclusive; both " + "select the packed-sequence dataset family. Pick one." ) - if args.max_seqlen_per_dp_cp_rank is not None: - total_cp_ranks = args.context_parallel_size - assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, ( - f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) ' - f'must be >= single sequence max length ({args.seq_length})' + if args.varlen_sbhd_validation: + # ``--dynamic-context-parallel`` ⊥ ``--varlen-sbhd-validation`` is + # checked in ``GPTDatasetConfig.__post_init__``; only the + # scheduler check stays here, since ``sequence_packing_scheduler`` + # is a training-framework flag not stored on the dataset config. + assert args.sequence_packing_scheduler is None, ( + "--varlen-sbhd-validation does not use a sequence packing " + "scheduler; drop --sequence-packing-scheduler." ) + # SBHD validation is a real-data numerical-reference path only; + # MockVarlenDataset does not implement it. + assert not args.mock_data, ( + "--varlen-sbhd-validation is not supported with --mock-data; " + "SBHD validation requires a real dataset." + ) + else: + # VarlenDataset emits one unpacked sample per __getitem__; it + # relies on an upstream packing scheduler to group variable-length + # samples into THD batches. Auto-pick a default scheduler when + # the user did not request one explicitly: + # * ``--dynamic-context-parallel`` is already wired to + # ``default_dynamic_cp`` upstream (see the dynamic-cp block + # earlier in ``validate_args``). + # * Otherwise fall back to ``dp_balanced`` (static packing). + if args.sequence_packing_scheduler is None: + args.sequence_packing_scheduler = 'dp_balanced' + + # Packed-sequence buffer-size check. Placed after all scheduler auto-select + # logic (dynamic-cp and --use-varlen-dataset both set the scheduler above) + # so it validates the final resolved scheduler; the varlen path picks its + # default after the earlier generic validation has run. + if args.sequence_packing_scheduler is not None: + if args.sequence_packing_scheduler == 'dp_balanced': + total_cp_ranks = args.context_parallel_size + else: + total_cp_ranks = args.data_parallel_size * args.context_parallel_size + if args.max_seqlen_per_dp_cp_rank is None: + args.max_seqlen_per_dp_cp_rank = getattr(args, "max_seqlen_per_cp_rank", None) + if args.max_seqlen_per_dp_cp_rank is None: + args.max_seqlen_per_dp_cp_rank = ( + args.seq_length + total_cp_ranks - 1 + ) // total_cp_ranks + assert total_cp_ranks * args.max_seqlen_per_dp_cp_rank >= args.seq_length, ( + f'Packed sequence buffer size ({total_cp_ranks * args.max_seqlen_per_dp_cp_rank}) ' + f'must be >= single sequence max length ({args.seq_length})' + ) # Data blend checks - assert args.mock_data + \ - bool(args.data_path) + \ - any([args.train_data_path, args.valid_data_path, args.test_data_path]) \ - <= 1, "A single data source must be provided in training mode, else None" + assert ( + args.mock_data + + bool(args.data_path) + + any([args.train_data_path, args.valid_data_path, args.test_data_path]) + <= 1 + ), "A single data source must be provided in training mode, else None" if args.fim_data: extra_tokens = [ @@ -1725,11 +1815,11 @@ def validate_args(args, defaults={}): assert not args.mock_data, "Mock dataset is not supported with FIM dataset." assert args.fim_rate, "--fim-rate should be specified." assert args.fim_spm_rate, "--fim-spm-rate should be specified." - assert all(token is not None for token in extra_tokens), "FIM extra tokens should be specified." + assert all( + token is not None for token in extra_tokens + ), "FIM extra tokens should be specified." - assert not ( - args.cross_entropy_loss_fusion and args.cross_entropy_fusion_impl == 'te' - ), ( + assert not (args.cross_entropy_loss_fusion and args.cross_entropy_fusion_impl == 'te'), ( "Transformer Engine cross entropy loss fusion is disabled due to stability issues. " "Use --cross-entropy-fusion-impl native, or omit --cross-entropy-loss-fusion." ) @@ -1769,17 +1859,62 @@ def validate_args(args, defaults={}): args.use_layer_wise_distributed_optimizer = True args.use_distributed_optimizer = False + if ( + args.optimizer == 'muon' + and (args.chunked_optimizer_state_offload or args.offload_optimizer_states) + and args.optimizer_state_offload_fraction > 0.0 + ): + assert args.use_layer_wise_distributed_optimizer, ( + "Muon optimizer state offload requires the LayerWise distributed optimizer " + "(--use-distributed-optimizer) so non-Muon parameter groups are routed through " + "a sibling DistributedOptimizer" + ) + assert not args.use_torch_fsdp2, "Emerging optimizer does not support Torch-FSDP2 for now." - assert not args.use_megatron_fsdp, "Emerging optimizer does not support Megatron-FSDP for now." - assert args.ckpt_format in ["torch", "torch_dist"], "Emerging optimizer supports torch and torch_dist checkpoint format." - - assert not ( - args.use_layer_wise_distributed_optimizer and args.moe_single_grouped_weight - ), ( - "The LayerWise distributed optimizer does not support --moe-single-grouped-weight: " - "Muon semantics for a single grouped [E, N, K] expert weight are not defined. " - "Disable --moe-single-grouped-weight or use Adam/DistributedOptimizer." - ) + assert ( + not args.use_megatron_fsdp + ), "Emerging optimizer does not support Megatron-FSDP for now." + assert args.ckpt_format in [ + "torch", + "torch_dist", + ], "Emerging optimizer supports torch and torch_dist checkpoint format." + + if args.use_layer_wise_distributed_optimizer: + if not args.use_layer_wise_param_layout: + # Decoupled compact LayerWise: fp8 parameter gather is supported via the FP8-aware + # whole-param all-gather. Only mxfp8/blockwise (fp4 out of scope); mxfp8 needs + # reuse_grad_buf. fp4 is rejected unconditionally -- the LayerWise gather routes + # buckets by is_float8tensor, so an NVFP4 param would silently take the raw + # flatten path. + assert not getattr(args, 'fp4_param_gather', False), ( + "Decoupled compact LayerWise DDP layout supports fp8 parameter gather only " + "(mxfp8 or blockwise); fp4_param_gather is out of scope." + ) + if args.fp8_param_gather: + assert args.fp8_recipe in ('mxfp8', 'blockwise'), ( + "fp8 parameter gather on the decoupled compact LayerWise DDP layout requires " + f"fp8_recipe in {{'mxfp8', 'blockwise'}}; got {args.fp8_recipe!r}." + ) + if args.fp8_recipe == 'mxfp8': + assert args.reuse_grad_buf_for_mxfp8_param_ag, ( + "mxfp8 + --fp8-param-gather on the decoupled compact LayerWise DDP layout " + "requires --reuse-grad-buf-for-mxfp8-param-ag (or use fp8_recipe='blockwise')." + ) + assert args.num_distributed_optimizer_instances == 1, ( + "the decoupled compact LayerWise DDP layout (the default; pass " + "--use-layer-wise-param-layout for the padded layout) requires " + "num_distributed_optimizer_instances == 1: the non-DistOpt LayerWise (Muon) buffers " + "only all-reduce within a single optimizer instance, so partial DistOpt (>1 " + "instance) would under-reduce Muon gradients across the full data-parallel domain." + ) + else: + # Padded LayerWise param layout: fp8/fp4 parameter gather is not supported here. + assert not args.fp8_param_gather and not getattr(args, 'fp4_param_gather', False), ( + "Layer-wise (Muon) distributed optimizer with the padded param layout does not " + "support FP8/FP4 parameter gather. Use the default compact decoupled layout (do " + "not pass --use-layer-wise-param-layout) for fp8 parameter gather, or " + "fp8_param_gather=False." + ) # Make sure all functionality that requires Gloo process groups is disabled. if not args.use_gloo_process_groups: @@ -1790,11 +1925,14 @@ def validate_args(args, defaults={}): assert args.use_dist_ckpt if args.dist_ckpt_optim_fully_reshardable: - assert not args.distrib_optim_fully_reshardable_mem_efficient, \ - '--distrib-optim-fully-reshardable-mem-efficient requires -enable-gloo-process-groups' + assert ( + not args.distrib_optim_fully_reshardable_mem_efficient + ), '--distrib-optim-fully-reshardable-mem-efficient requires -enable-gloo-process-groups' if args.fake_process_group: - assert args.moe_token_dispatcher_type != "flex", "Fake process group is not supported with flex token dispatcher." + assert ( + args.moe_token_dispatcher_type != "flex" + ), "Fake process group is not supported with flex token dispatcher." # Disable nan check for fake process group args.check_for_nan_in_loss_and_grad = False warn_rank_0('check_for_nan_in_loss_and_grad is set to False for fake process group.') @@ -1804,23 +1942,31 @@ def validate_args(args, defaults={}): # Checkpointing if args.ckpt_fully_parallel_save_deprecated and args.rank == 0: - print('--ckpt-fully-parallel-save flag is deprecated and has no effect.' - ' Use --no-ckpt-fully-parallel-save to disable parallel save.') + print( + '--ckpt-fully-parallel-save flag is deprecated and has no effect.' + ' Use --no-ckpt-fully-parallel-save to disable parallel save.' + ) if ( args.use_dist_ckpt and not args.ckpt_fully_parallel_save and args.use_distributed_optimizer and args.rank == 0 ): - print('Warning: With non-parallel ckpt save and DistributedOptimizer,' - ' it will be impossible to resume training with different parallelism.' - ' Consider removing flag --no-ckpt-fully-parallel-save.') + print( + 'Warning: With non-parallel ckpt save and DistributedOptimizer,' + ' it will be impossible to resume training with different parallelism.' + ' Consider removing flag --no-ckpt-fully-parallel-save.' + ) if args.use_dist_ckpt_deprecated and args.rank == 0: - print('--use-dist-ckpt is deprecated and has no effect.' - ' Use --ckpt-format to select the checkpoint format.') + print( + '--use-dist-ckpt is deprecated and has no effect.' + ' Use --ckpt-format to select the checkpoint format.' + ) if args.dist_ckpt_format_deprecated and args.rank == 0: - print('--dist-ckpt-format is deprecated and has no effect.' - ' Use --ckpt-format to select the checkpoint format.') + print( + '--dist-ckpt-format is deprecated and has no effect.' + ' Use --ckpt-format to select the checkpoint format.' + ) if args.use_dist_ckpt and args.ckpt_fully_parallel_load: if args.ckpt_fully_parallel_load_exchange_algo != "broadcast": @@ -1829,25 +1975,6 @@ def validate_args(args, defaults={}): "Other algorithms cannot guarantee numerical stability yet." ) - # The PG-distribution cache stores one file per parallelization group, keyed by - # that group's minimum global rank, and each file holds both the save and the - # load distribution. If save and load used different parallelization groups, - # two distinct groups could map to the same file (they can share a minimum - # global rank), so one side would silently read a distribution computed for the - # other group's layout. Require both sides to use the same group. - if args.ckpt_pg_tensors_cache_path is not None: - assert ( - args.ckpt_fully_parallel_save_process_group - == args.ckpt_fully_parallel_load_process_group - ), ( - "--ckpt-pg-tensors-cache-path requires --ckpt-fully-parallel-save-process-group and " - "--ckpt-fully-parallel-load-process-group to be identical, but got " - f"save='{args.ckpt_fully_parallel_save_process_group}' and " - f"load='{args.ckpt_fully_parallel_load_process_group}'. The cache is keyed by the " - "parallelization group, so mixing groups can read a distribution computed for a " - "different layout." - ) - if args.load_main_params_from_ckpt: assert args.no_load_optim, '--load-main-params-from-ckpt must be used with --no-load-optim.' @@ -1862,39 +1989,11 @@ def validate_args(args, defaults={}): if not args.async_save: args.async_strategy = "mcore" - if args.logits_save_dir is not None: - assert args.logits_save_top_k is not None, '--logits-save-top-k is required when --logits-save-dir is set.' - assert args.async_save, ( - '--logits-save-dir requires --async-save (and --use-persistent-ckpt-worker). ' - 'Logits are flushed as an async request in the checkpoint queue.' - ) - if not args.freeze_all_layers: - warn_rank_0( - '--logits-save-dir without --freeze-all-layers: the LM loss is still computed and ' - 'gradients will update the model while logits are dumped. This is intended only ' - 'when dumping logits during active training; for a frozen-teacher dump pass ' - '--freeze-all-layers.' - ) - - if args.freeze_all_layers: - if args.use_distributed_optimizer: - warn_rank_0( - '--freeze-all-layers incompatible with use_distributed_optimizer. Disabling use_distributed_optimizer.' - ) - args.use_distributed_optimizer = False - if args.overlap_param_gather: - warn_rank_0( - '--freeze-all-layers incompatible with overlap_param_gather. Disabling overlap_param_gather.' - ) - args.overlap_param_gather = False - - if args.override_ckpt_iteration is not None: - assert not args.finetune, "Cannot override checkpoint iteration together with finetune flag." - # Inference args if args.inference_batch_times_seqlen_threshold > -1: - assert args.pipeline_model_parallel_size > 1, \ - "--inference-batch-times-seqlen-threshold requires setting --pipeline-model-parallel-size > 1." + assert ( + args.pipeline_model_parallel_size > 1 + ), "--inference-batch-times-seqlen-threshold requires setting --pipeline-model-parallel-size > 1." assert ( args.cuda_graph_impl == "none" ), "Pipeline-parallel microbatched inference is incompatible with CUDA graphs" @@ -1915,10 +2014,16 @@ def validate_args(args, defaults={}): if args.skip_train and not args.perform_rl_step and not args.no_load_optim: args.no_load_optim = True warn_rank_0('enabling --no-load-optim when skipping training.') - if args.skip_train and args.perform_rl_step and args.no_load_optim and args.rl_offload_optimizer_during_inference: - assert False, \ - '--no-load-optim with --skip-train --perform-rl-step skips the optimizer; ' \ + if ( + args.skip_train + and args.perform_rl_step + and args.no_load_optim + and args.rl_offload_optimizer_during_inference + ): + assert False, ( + '--no-load-optim with --skip-train --perform-rl-step skips the optimizer; ' '--rl-offload-optimizer-during-inference is incompatible (no optimizer to offload).' + ) # Optimizer CPU offload check if args.optimizer_cpu_offload: @@ -1931,22 +2036,68 @@ def validate_args(args, defaults={}): "must be used in conjunction with `--fp8-recipe delayed`." ) + # Optimizer configuration normalizes the deprecated spelling. Include it here so + # cross-config validation still runs before that configuration is constructed. + if ( + args.chunked_optimizer_state_offload or args.offload_optimizer_states + ) and args.optimizer_state_offload_fraction > 0.0: + if args.optimizer_state_offload_chunk_size_mb == 0: + warn_rank_0( + "optimizer state offload is enabled with chunk size 0, so all " + "selected optimizer tensor state is temporarily restored together. Set " + "--optimizer-state-offload-chunk-size-mb to a positive value to bound the " + "tensor-state GPU window. Selected master weights always use one full " + "restore window regardless of chunk size." + ) + assert ( + not args.use_megatron_fsdp + ), "chunked optimizer state offload does not support Megatron-FSDP" + if not args.no_save_optim or not args.no_load_optim: + assert args.ckpt_format == 'torch_dist', ( + "chunked optimizer state offload requires --ckpt-format torch_dist when " + "optimizer state is saved or loaded because its sharded load hooks preserve " + "CPU canonical storage" + ) + if not args.no_save_optim: + assert not args.async_save, ( + "chunked optimizer state offload does not support --async-save when optimizer " + "state is saved because the background writer can retain tensors backed by " + "reusable pinned CPU buffers" + ) + assert ( + args.cuda_graph_impl != 'full_iteration' + ), "chunked optimizer state offload does not support full-iteration CUDA graphs" + if args.optimizer == 'muon' and args.fp8_param_gather and args.overlap_param_gather: + warn_rank_0( + "Muon compact FP8 parameter gather with chunked optimizer state offload " + "forces LayerWise-owned pre-forward parameter buckets to complete before " + "masters are offloaded; sibling DistributedOptimizer gathers and " + "optimizer-state transfers remain overlapped." + ) + if args.non_persistent_ckpt_type == "local": - assert args.non_persistent_local_ckpt_dir is not None, "Tried to use local checkpointing without specifying --local-ckpt-dir!" + assert ( + args.non_persistent_local_ckpt_dir is not None + ), "Tried to use local checkpointing without specifying --local-ckpt-dir!" if args.replication: - assert args.replication_jump is not None, "--replication requires the value of --replication-jump!" - assert args.non_persistent_ckpt_type == "local", f"--replication requires args.non_persistent_ckpt_type == 'local', but got: {args.non_persistent_ckpt_type}" + assert ( + args.replication_jump is not None + ), "--replication requires the value of --replication-jump!" + assert ( + args.non_persistent_ckpt_type == "local" + ), f"--replication requires args.non_persistent_ckpt_type == 'local', but got: {args.non_persistent_ckpt_type}" elif args.replication_jump: warn_rank_0("--replication-jump was specified despite not using replication. Ignoring.") args.replication_jump = None if args.delay_wgrad_compute: - assert args.transformer_impl == 'transformer_engine', \ - "Delaying wgrad compute is only supported with transformer_engine implementation" + assert ( + args.transformer_impl == 'transformer_engine' + ), "Delaying wgrad compute is only supported with transformer_engine implementation" if args.overlap_grad_reduce: - assert is_te_min_version("2.8.0"), ( - "overlap_grad_reduce is only supported with TE >= 2.8.0 when enabling delay_wgrad_compute" - ) + assert is_te_min_version( + "2.8.0" + ), "overlap_grad_reduce is only supported with TE >= 2.8.0 when enabling delay_wgrad_compute" wgrad_in_graph_scope = CudaGraphModule.attn in args.cuda_graph_modules or ( CudaGraphModule.moe_router in args.cuda_graph_modules and args.moe_shared_expert_intermediate_size is not None @@ -1973,11 +2124,13 @@ def validate_args(args, defaults={}): ) if args.fine_grained_activation_offloading: - assert args.transformer_impl == 'transformer_engine', \ - "Fine-grained activation offloading is only supported with transformer_engine implementation" + assert ( + args.transformer_impl == 'transformer_engine' + ), "Fine-grained activation offloading is only supported with transformer_engine implementation" if is_te_min_version("2.10.0"): - assert os.getenv("NVTE_CPU_OFFLOAD_V1", "0") == "1", \ - "For fine-grained activation offloading with TE >= 2.10.0, NVTE_CPU_OFFLOAD_V1 should be set to 1 to avoid offloading weights." + assert ( + os.getenv("NVTE_CPU_OFFLOAD_V1", "0") == "1" + ), "For fine-grained activation offloading with TE >= 2.10.0, NVTE_CPU_OFFLOAD_V1 should be set to 1 to avoid offloading weights." if args.mtp_num_layers: # MTP is compatible with position embedding types that use position_ids. @@ -2011,19 +2164,23 @@ def validate_args(args, defaults={}): ), '--cuda-graph-modules must be empty when --cuda-graph-impl=full_iteration.' if args.multi_latent_attention: - assert not args.group_query_attention, "Group query attention is mutually exclusive with multi latent attention." + assert ( + not args.group_query_attention + ), "Group query attention is mutually exclusive with multi latent attention." if args.mla_down_proj_fusion: - assert args.multi_latent_attention, "--mla-down-proj-fusion requires --multi-latent-attention" - - assert ( - not args.moe_use_norm_before_up_proj or args.moe_latent_size is not None - ), "--moe-use-norm-before-up-proj requires --moe-latent-size to be set." + assert ( + args.multi_latent_attention + ), "--mla-down-proj-fusion requires --multi-latent-attention" # MoE latent projections if args.moe_latent_size is not None: - assert args.moe_latent_size > 0, "MoE latent projection dimension has to be greater than zero." - assert args.num_experts is not None, "MoE latent projections are applicable only for MoE models." + assert ( + args.moe_latent_size > 0 + ), "MoE latent projection dimension has to be greater than zero." + assert ( + args.num_experts is not None + ), "MoE latent projections are applicable only for MoE models." # Print arguments. _print_args("arguments", args) @@ -2034,6 +2191,7 @@ def validate_args(args, defaults={}): def _print_args(title, args): """Print arguments.""" from megatron.training.utils import is_rank0 + if is_rank0(): print(f'------------------------ {title} ------------------------', flush=True) str_list = [] @@ -2049,296 +2207,521 @@ def _check_arg_is_not_none(args, arg): assert getattr(args, arg) is not None, '{} argument is None'.format(arg) +def core_transformer_config_from_args(args, config_class=None): + + # Config class. + config_class = config_class or TransformerConfig + + if args.multi_latent_attention: + config_class = MLATransformerConfig + + if args.heterogeneous_layers_config_path is not None: + assert ( + not args.multi_latent_attention + ), "Multi latent attention with heterogeneous layers is not supported." + config_class = HeterogeneousTransformerConfig + + # Translate args to core transformer configuration + kw_args = {} + for f in dataclasses.fields(config_class): + if hasattr(args, f.name): + kw_args[f.name] = getattr(args, f.name) + kw_args['persist_layer_norm'] = not args.no_persist_layer_norm + kw_args['deallocate_pipeline_outputs'] = True + kw_args['pipeline_dtype'] = args.params_dtype + kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm + kw_args['num_moe_experts'] = args.num_experts + kw_args['actual_vocab_size'] = args.padded_vocab_size + kw_args['rotary_interleaved'] = args.rotary_interleaved + kw_args['num_layers_in_first_pipeline_stage'] = args.decoder_first_pipeline_num_layers + kw_args['num_layers_in_last_pipeline_stage'] = args.decoder_last_pipeline_num_layers + kw_args['fp8_param'] = args.fp8_param_gather + kw_args['fp4_param'] = args.fp4_param_gather + if args.swiglu: + kw_args['activation_func'] = F.silu + kw_args['gated_linear_unit'] = True + kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion + else: + kw_args['bias_activation_fusion'] = args.bias_gelu_fusion + if args.squared_relu: + assert not args.swiglu + kw_args['activation_func'] = squared_relu + elif args.quick_geglu: + assert not args.swiglu + kw_args['gated_linear_unit'] = True + kw_args['activation_func'] = quick_gelu + if args.init_method_xavier_uniform: + kw_args['init_method'] = torch.nn.init.xavier_uniform_ + kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_ + if args.group_query_attention: + kw_args['num_query_groups'] = args.num_query_groups + else: + kw_args['num_query_groups'] = None + kw_args['config_logger_dir'] = args.config_logger_dir + if args.rope_type is None: + # Pop 'rope_type' to let the config class use the default value. + kw_args.pop('rope_type', None) + else: + assert ( + args.multi_latent_attention or args.rope_type == 'rope' + ), f'Common attention only support rope_type="rope", but got {args.rope_type}.' + + if len(args.cp_comm_type) == 1: + kw_args['cp_comm_type'] = args.cp_comm_type[0] + if args.hybrid_layer_pattern is not None: + kw_args['is_hybrid_model'] = True + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + + _pat = args.hybrid_layer_pattern + _has_dsv4_csa = (Symbols.CSA in _pat) or (Symbols.HCA in _pat) or (Symbols.WINDOW in _pat) + _has_dsa = Symbols.DS_ATTENTION in _pat + if getattr(args, 'experimental_attention_variant', None) is None: + # 'C'/'H'/'W' run the DSv4 CompressedSparseAttention (CSA/HCA/window-only), which + # requires the full dsv4_hybrid contract (MLA, TP==1, no qk_clip, + # qk_head_dim/kv_lora_rank derivation). Set the variant so transformer_config runs + # that validation+derivation rather than silently skipping it. 'D' alone stays legacy + # DSv3 'dsa'. An explicit --experimental-attention-variant is always respected. + if _has_dsv4_csa: + kw_args['experimental_attention_variant'] = 'dsv4_hybrid' + elif _has_dsa: + kw_args['experimental_attention_variant'] = 'dsa' + # Normalize compact and legacy-padded ratios through the shared migration helper. + _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, _pat) + + _resolve_dsa_kernel_backend_cli_default(args, kw_args) + + kw_args['inference_sampling_seed'] = args.seed + + # handle quantization config + # NOTE: Kitchen arguments are only added to the namespace when + # Kitchen library is available. + if hasattr(args, "kitchen_config_file") and args.kitchen_config_file is not None: + kw_args['use_kitchen'] = True + kw_args['quant_recipe'] = load_quantization_recipe(args.kitchen_config_file) + elif hasattr(args, 'kitchen_recipe_number') and args.kitchen_recipe_number is not None: + kw_args['use_kitchen'] = True + kw_args['quant_recipe'] = kitchen_quantization_recipe_config(args.kitchen_recipe_number) + + kw_args['moe_latent_size'] = args.moe_latent_size + + if args.te_precision_config_file: + assert not 'quant_recipe' in kw_args, "Quantization recipe already configured." + # TODO(kwyss): Prohibit fp8_params or fp4_params with this flexibility + kw_args['quant_recipe'] = load_quantization_recipe(args.te_precision_config_file) + + if hasattr(args, "use_kitchen_attention"): + kw_args['use_kitchen_attention'] = args.use_kitchen_attention + if hasattr(args, "kitchen_attention_backend"): + kw_args['kitchen_attention_backend'] = args.kitchen_attention_backend + + # Return config. + return config_class(**kw_args) + + def _add_transformer_engine_args(parser): group = parser.add_argument_group(title='Transformer-Engine') # FP4 related arguments - group.add_argument('--fp4-param-gather', action='store_true', - help='Keep the compute param in fp4 (do not use any other intermediate ' - 'dtype) and perform the param all-gather in fp4.') + group.add_argument( + '--fp4-param-gather', + action='store_true', + help='Keep the compute param in fp4 (do not use any other intermediate ' + 'dtype) and perform the param all-gather in fp4.', + ) # FP8 related arguments - group.add_argument('--fp8-param-gather', action='store_true', - help='Keep the compute param in fp8 (do not use any other intermediate ' - 'dtype) and perform the param all-gather in fp8.') + group.add_argument( + '--fp8-param-gather', + action='store_true', + help='Keep the compute param in fp8 (do not use any other intermediate ' + 'dtype) and perform the param all-gather in fp8.', + ) # TE precision config file - group.add_argument('--te-precision-config-file', default=None, - help='Configuration file to select per-module precision overrides. ' - 'See TransformerEngineMixedPrecision.md') + group.add_argument( + '--te-precision-config-file', + default=None, + help='Configuration file to select per-module precision overrides. ' + 'See TransformerEngineMixedPrecision.md', + ) return parser + def _add_inference_args(parser): group = parser.add_argument_group(title='inference') - group.add_argument('--inference-batch-times-seqlen-threshold', - type=int, default=-1, - help='If (batch-size * sequence-length) is smaller than this threshold' - 'then batches will not be split up for pipelining.' - 'Requires setting --pipeline-model-parallel-size > 1.' - 'Setting this to -1 indicates that batch pipelining is not used.') - group.add_argument('--max-tokens-to-oom', - type=int, default=12000, - help='Maximum number of tokens during inference' - 'tokens here is # in prompt + # to generate' - 'Allows us to throw an error before OOM crashes server') - group.add_argument('--output-bert-embeddings', action='store_true', - help='Output Bert embeddings (via mean pooling) from ' - 'model, rather than its binary head output or entire ' - 'hidden batch.') - group.add_argument('--bert-embedder-type', default="megatron", - choices=["megatron", "huggingface"], - help='Select either Megatron or Huggingface as the ' - 'Bert embedder.') - group.add_argument('--cuda-graph-scope', nargs='+', type=_parse_cuda_graph_modules_arg, - default=None, dest='cuda_graph_scope_deprecated', - help=argparse.SUPPRESS) # hidden; use --cuda-graph-modules instead - group.add_argument('--cuda-graph-modules', nargs='+', type=_parse_cuda_graph_modules_arg, default=[], - help='Selects training capture coverage within per-layer CUDA graphs ' - '(local and transformer_engine implementations). ' - 'Valid values are "attn", "mlp", "moe", "moe_router", "moe_preprocess", and "mamba": ' - '"attn": captures operations in TransformerLayer._forward_attention(). ' - '"mlp": captures operations in TransformerLayer._forward_mlp() for a dense layer. ' - '"moe": captures operations in TransformerLayer._forward_mlp() for a MoE layer. ' - '"moe_router": captures operations in TransformerLayer._forward_mlp() up to MoELayer.router(), ' - 'including the shared experts if they are not overlapped with EP comm. ' - '"moe_preprocess": captures operations in MoELayer.preprocess(). Must be used together with "moe_router". ' - '"mamba": captures the mamba layer. ' - 'An empty list means capturing the whole Transformer layer. ' - 'This field is meaningless when --cuda-graph-impl=full_iteration and must be empty. ' - 'Backward compatibility: "full" is deprecated but kept for backward compatibility; ' - 'it is transformed to an empty list in validate_args. The deprecated values ' - '"full_iteration" and "full_iteration_inference" are also accepted and migrated ' - 'to the new API in validate_args.') - group.add_argument('--use-legacy-static-engine', action='store_true', default=False, - help='Use legacy static engine. (Current static engine uses dynamic engine under the hood)', - dest='use_legacy_static_engine') - group.add_argument('--inference-max-requests', type=int, default=8, - help='Maximum number of requests for inference.', - dest='inference_max_requests') - group.add_argument('--inference-max-seq-length', type=int, default=2560, - help='Maximum sequence length expected for inference (prefill + decode).', - dest='inference_max_seq_length') - group.add_argument('--inference-dynamic-batching', - action='store_true', default=False, - help='Enable dynamic batching mode.') - group.add_argument('--inference-dynamic-batching-buffer-size-gb', - type=float, default=40., - help='On-GPU portion of the shared KV cache block pool. ' - 'The total amount of memory allocated for the KV cache ' - '(CPU + GPU memory) depends on the value set for the ' - 'unified virtual memory (UVM) level (via ' - '`--inference-dynamic-batching-unified-memory-level`).' - 'If the UVM level is 0, then only GPU memory is used and ' - 'the total memory equals `buffer_size_gb`. If the UVM ' - 'level is 1, then additional memory is utilized on the ' - 'CPU and the total memory equals `buffer_size_gb + ' - 'paused_buffer_size_gb`.') - group.add_argument('--inference-dynamic-batching-paused-buffer-size-gb', - type=float, default=None, - help='Memory used to derive the paused-request block retention ' - 'budget. This does not reserve blocks from active requests: ' - 'active requests may use the entire shared pool of usable KV ' - 'cache blocks. Under allocation pressure, paused requests ' - 'retain blocks only within this budget and excess paused ' - 'requests may be evicted.') - group.add_argument('--inference-dynamic-batching-mamba-memory-ratio', type=float, default=None, - help='Percentage of memory buffer to allocate for Mamba states. ' - 'If not specified, allocates Mamba state tensors for each KV cache block. ' - 'Only used for hybrid models.') - group.add_argument('--inference-dynamic-batching-block-size', - type=int, default=256, - help='KV cache block size. ' - 'It should be a multiple of 256') - group.add_argument('--inference-dynamic-batching-max-requests', - type=int, default=None, - help='Override the inference context\'s `max_requests`. ' - 'By default, `max_requests` is set to the number of ' - 'blocks in the context\'s memory buffer.') - group.add_argument('--inference-dynamic-batching-max-tokens', - type=int, default=None, - help='Override the inference context\'s default `max_tokens`.') - group.add_argument('--inference-dynamic-batching-num-cuda-graphs', - type=int, default=16, - help='Maximum number of cuda graphs to capture, where the ' - 'cuda graph batch sizes range from 1 to `max_requests`. ' - '(See `dynamic_context.py` for details on how ' - '`max_requests` is computed). Due to rounding, the actual ' - 'number of cuda graphs may not equal this argument.' - 'The user can also pass -1, in which case we automatically determine the number of graphs ' \ - 'to capture based on the `max_requests`.') - group.add_argument('--inference-dynamic-batching-track-paused-request-events', - action='store_true', - help='Track paused request ids by adding \'paused\' events ' - 'to each request\'s event history. This has a very minor ' - 'impact on latency.') - group.add_argument('--inference-dynamic-batching-track-generated-token-events', - action='store_true', - help='Track per-token events with timestamps for each generated token. ' - 'When enabled, each generated token creates a GENERATED_TOKEN event ' - 'with a timestamp, useful for per-token latency analysis.') - group.add_argument('--decode-only-cuda-graphs', - action='store_true', default=False, - help='Only use cuda graphs for decode-only steps, not prefill and mixed steps.') - group.add_argument('--inference-cuda-graph-all-prefills', - action='store_true', default=False, - help='Extend prefill/mixed CUDA graph capture up to `max_tokens`. ' - 'By default, all graphs are limited by the decode limit of ' - '`max_requests * (num_speculative_tokens + 1)`.') - group.add_argument('--inference-cuda-graph-max-tokens', type=int, default=512, - dest='inference_cuda_graph_max_tokens', - help='Token ceiling for the largest captured prefill/mixed CUDA ' - 'graph (default: 512). Clamped to at least the decode limit ' - '`max_requests * (num_speculative_tokens + 1)` and at most ' - '`max_tokens`. Ignored when --inference-cuda-graph-all-prefills ' - 'is set (which extends capture to the full `max_tokens`).') - group.add_argument('--inference-dynamic-batching-unified-memory-level', - type=int, default=0, choices=[0, 1], - help='Set unified memory usage within the dynamic ' - 'inference context. The levels are: 0) no unified memory, ' - '1) allocate `memory_buffer` in unified memory. ' - 'Eventually, additional levels will be included to ' - 'control other tensors within the context.') - group.add_argument('--enable-chunked-prefill', dest='enable_chunked_prefill', - action='store_true', default=False, - help="Enable chunked prefill (disabled by default)") - group.add_argument('--num-speculative-tokens', type=int, default=0, - help='Number of speculative tokens generated during decode') - group.add_argument('--inference-dynamic-batching-prefix-caching', - dest='inference_dynamic_batching_enable_prefix_caching', - action=argparse.BooleanOptionalAction, - default=False, - help='Enable/disable prefix caching for dynamic batching inference. ' - 'When disabled, KV cache blocks cannot be shared between ' - 'requests with identical prompt prefixes.') - group.add_argument('--inference-dynamic-batching-prefix-caching-eviction-policy', - type=str, default='ref_zero', - choices=['ref_zero', 'lru'], - dest='inference_dynamic_batching_prefix_caching_eviction_policy', - help='Eviction policy for prefix caching blocks. ' - '"ref_zero" (default) immediately returns blocks to the ' - 'free pool when ref_count hits 0. "lru" keeps blocks ' - 'cached and evicts via LRU only when space is needed.') - group.add_argument('--inference-dynamic-batching-prefix-caching-coordinator-policy', - type=str, default='load_balanced', - choices=['longest_prefix', 'first_prefix_block', 'load_balanced'], - dest='inference_dynamic_batching_prefix_caching_coordinator_policy', - help='Coordinator routing policy for prefix caching. ' - '"load_balanced" (default) routes to the rank with the fewest ' - 'in-flight requests, ignoring prefix affinity. ' - '"first_prefix_block" routes based on the first block hash only. ' - '"longest_prefix" routes to the rank with the longest matching ' - 'prefix. "first_prefix_block" and "longest_prefix" both combine ' - 'prefix affinity with load balancing and fall back to ' - 'load-balanced routing when prefix caching is disabled or no ' - 'prefix match exists.') - group.add_argument('--inference-dynamic-batching-prefix-caching-routing-alpha', - type=float, default=0.5, - dest='inference_dynamic_batching_prefix_caching_routing_alpha', - help='Weight for prefix-aware routing score: ' - 'score = alpha * match + (1 - alpha) * normalized_load. ' - 'Higher alpha favors prefix cache hits; lower alpha ' - 'favors load balance. Default: 0.5.') - group.add_argument('--inference-dynamic-batching-media-cache-coordinator-policy', - type=str, default='affinity', - choices=['affinity', 'load_balanced'], - dest='inference_dynamic_batching_media_cache_coordinator_policy', - help='Coordinator routing policy for media caching. ' - '"affinity" prefers a rank assigned the same media; ' - '"load_balanced" ignores standalone media affinity.') - group.add_argument('--inference-dynamic-batching-media-cache-routing-weight', - type=float, default=1.0, - dest='inference_dynamic_batching_media_cache_routing_weight', - help='Media-cache hit weight in equivalent compact-prompt blocks. ' - 'Default: 1.0.') - group.add_argument('--inference-dynamic-batching-vision-embedding-cache-max-bytes', - type=int, default=0, - dest='inference_dynamic_batching_vision_embedding_cache_max_bytes', - help='Maximum GPU bytes retained per engine for reusable vision ' - 'embeddings. Zero disables the cache. Default: 0.') - group.add_argument('--inference-dynamic-batching-allow-stale-multimodal-embeddings', - action='store_true', - dest='inference_dynamic_batching_allow_stale_multimodal_embeddings', - help='Allow request-local and cached multimodal embeddings to survive ' - 'suspend/resume and generation-epoch changes. Use only when model ' - 'weights do not change across these boundaries.') - group.add_argument('--inference-dynamic-batching-prefix-caching-mamba-gb', - type=float, default=None, - dest='inference_dynamic_batching_prefix_caching_mamba_gb', - help='GPU memory budget (in GB) for the Mamba state cache ' - 'used by prefix caching on hybrid models. When set, Mamba ' - 'states at block boundaries are cached for reuse. This budget ' - 'covers both the durable cache (the ssm_states/conv_states ' - 'slots reused across requests) and the per-step extraction ' - 'scratch (the intermediate_ssm_out/intermediate_conv_out ' - 'buffers, sized to min(ceil(max_tokens / block_size), ' - '3 * max_requests) slots); the scratch is reserved first, so a ' - 'smaller max_tokens (or max_requests) shrinks the scratch and ' - 'leaves more durable slots.') - group.add_argument('--inference-dynamic-batching-cuda-graph-mixed-prefill-count', - type=int, default=16, - help='Number of mixed prefill requests to capture in a cuda graph.') - group.add_argument('--inference-dynamic-batching-cuda-graph-sizing-distribution', - type=str, default='exponential', - choices=['exponential', 'linear'], - dest='inference_dynamic_batching_cuda_graph_sizing_distribution', - help='Spacing of CUDA graph token counts. "exponential" (default) ' - 'halves from cuda_graph_max_tokens down to tp_size, giving a ' - 'log-spaced distribution with bounded relative padding. ' - '"linear" uses varying linear strides across the range.') - group.add_argument('--inference-dynamic-batching-sampling-backend', - type=str, default='torch', - choices=['torch', 'flashinfer'], - help='Which sampling kernels to use during inference. ' - 'Falls back to "torch" with a warning if "flashinfer" ' - 'is requested but the package is not installed.') - group.add_argument('--use-same-sampling-seed-across-dp-ranks', - action='store_false', dest='offset_sampling_seed_by_dp_rank', - default=True, - help='Use the same inference sampling seed on every data-parallel rank. ' - '--deterministic-mode also uses the same seed on every DP rank.') - group.add_argument('--inference-dynamic-batching-async-sched-mode', - type=str, default='legacy', - choices=['legacy', 'async'], - help='Async scheduling mode for dynamic batching. ' - '"legacy" (default) preserves the existing resolve-before-prepare ' - 'path. "async" overlaps asynchronous scheduling phases by reordering ' - 'them to prepare-before-resolve.') - group.add_argument('--inference-dynamic-batching-logprobs-mode', - type=str, default='raw_logprobs', - choices=['raw_logprobs', 'processed_logprobs'], - help='How returned inference log-probs are computed engine-wide. ' - '"raw_logprobs" (default) uses the unmodified model logits; ' - '"processed_logprobs" uses temperature and filters by top-k/top-p.') - group.add_argument('--inference-logging-step-interval', type=int, default=0, - help='Step interval for logging inference metrics. ' - 'Default to 0 to disable inference logging.') - group.add_argument('--inference-text-gen-server-logging', action=argparse.BooleanOptionalAction, - required=False, default=False, - help='Enable per-request logging in the inference text generation server.') - group.add_argument('--inference-wandb-logging', action=argparse.BooleanOptionalAction, - required=False, default=False, help='Enable inference wandb logging.') - group.add_argument("--inference-coordinator-port", type=int, - help="This port will be used to setup the inference coordinator on node-0") - group.add_argument('--mamba-inference-conv-states-dtype', type=str, - choices=['bf16', 'fp16', 'fp32'], default='bf16', - help='Dtype for the Mamba inference conv states tensor') - group.add_argument('--mamba-inference-ssm-states-dtype', type=str, - choices=['bf16', 'fp16', 'fp32'], default='bf16', - help='Dtype for the Mamba inference SSM states tensor') - group.add_argument('--inference-use-synchronous-zmq-collectives', action=argparse.BooleanOptionalAction, - required=False, default=False, help='Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for MoEs.') - group.add_argument('--inference-disable-ep-consensus', action=argparse.BooleanOptionalAction, - required=False, default=False, - help='Skip the EP-group consensus all-reduce in the inference engine control loop and step on local state only. ' - 'Pause/unpause take effect as soon as the signal is delivered to a rank. ' - 'Only safe when EP coordination is not required (e.g. ep_world_size == 1).') - group.add_argument('--inference-shards', type=str, default=None, metavar='SPEC', - help='Partition the world into independent inference models, each with ' - 'its own parallelism, e.g. "tp=2,role=prefill+tp=1,dp=2,role=decode". ' - 'Shards are separated by "+" or ";"; per-shard keys are ' - 'tp,pp,ep,expt_tp,dp (each defaults to 1) and must partition the full ' - 'world. Tagging shards role=prefill|decode enables disaggregated ' - 'inference (prefill hands KV to the decode pool); a dp>1 decode shard ' - 'is several independent decode instances.') + group.add_argument( + '--inference-batch-times-seqlen-threshold', + type=int, + default=-1, + help='If (batch-size * sequence-length) is smaller than this threshold' + 'then batches will not be split up for pipelining.' + 'Requires setting --pipeline-model-parallel-size > 1.' + 'Setting this to -1 indicates that batch pipelining is not used.', + ) + group.add_argument( + '--max-tokens-to-oom', + type=int, + default=12000, + help='Maximum number of tokens during inference' + 'tokens here is # in prompt + # to generate' + 'Allows us to throw an error before OOM crashes server', + ) + group.add_argument( + '--output-bert-embeddings', + action='store_true', + help='Output Bert embeddings (via mean pooling) from ' + 'model, rather than its binary head output or entire ' + 'hidden batch.', + ) + group.add_argument( + '--bert-embedder-type', + default="megatron", + choices=["megatron", "huggingface"], + help='Select either Megatron or Huggingface as the ' 'Bert embedder.', + ) + group.add_argument( + '--cuda-graph-scope', + nargs='+', + type=_parse_cuda_graph_modules_arg, + default=None, + dest='cuda_graph_scope_deprecated', + help=argparse.SUPPRESS, # hidden; use --cuda-graph-modules instead + ) + group.add_argument( + '--cuda-graph-modules', + nargs='+', + type=_parse_cuda_graph_modules_arg, + default=[], + help='Selects training capture coverage within per-layer CUDA graphs ' + '(local and transformer_engine implementations). ' + 'Valid values are "attn", "mlp", "moe", "moe_router", "moe_preprocess", and "mamba": ' + '"attn": captures operations in TransformerLayer._forward_attention(). ' + '"mlp": captures operations in TransformerLayer._forward_mlp() for a dense layer. ' + '"moe": captures operations in TransformerLayer._forward_mlp() for a MoE layer. ' + '"moe_router": captures operations in TransformerLayer._forward_mlp() up to MoELayer.router(), ' + 'including the shared experts if they are not overlapped with EP comm. ' + '"moe_preprocess": captures operations in MoELayer.preprocess(). Must be used together with "moe_router". ' + '"mamba": captures the mamba layer. ' + 'An empty list means capturing the whole Transformer layer. ' + 'This field is meaningless when --cuda-graph-impl=full_iteration and must be empty. ' + 'Backward compatibility: "full" is deprecated but kept for backward compatibility; ' + 'it is transformed to an empty list in validate_args. The deprecated values ' + '"full_iteration" and "full_iteration_inference" are also accepted and migrated ' + 'to the new API in validate_args.', + ) + group.add_argument( + '--use-legacy-static-engine', + action='store_true', + default=False, + help='Use legacy static engine. (Current static engine uses dynamic engine under the hood)', + dest='use_legacy_static_engine', + ) + group.add_argument( + '--inference-max-requests', + type=int, + default=8, + help='Maximum number of requests for inference.', + dest='inference_max_requests', + ) + group.add_argument( + '--inference-max-seq-length', + type=int, + default=2560, + help='Maximum sequence length expected for inference (prefill + decode).', + dest='inference_max_seq_length', + ) + group.add_argument( + '--inference-dynamic-batching', + action='store_true', + default=False, + help='Enable dynamic batching mode.', + ) + group.add_argument( + '--inference-dynamic-batching-buffer-size-gb', + type=float, + default=40.0, + help='Amount of on-GPU memory allocated for the KV cache. ' + 'The total amount of memory allocated for the KV cache ' + '(CPU + GPU memory) depends on the value set for the ' + 'unified virtual memory (UVM) level (via ' + '`--inference-dynamic-batching-unified-memory-level`).' + 'If the UVM level is 0, then only GPU memory is used and ' + 'the total memory equals `buffer_size_gb`. If the UVM ' + 'level is 1, then additional memory is utilized on the ' + 'CPU and the total memory equals `buffer_size_gb + ' + 'paused_buffer_size_gb`.', + ) + group.add_argument( + '--inference-dynamic-batching-paused-buffer-size-gb', + type=float, + default=None, + help='Amount of memory reserved for paused requests in ' + 'the dynamic inference context. Active requests are ' + 'paused when there are not enough active blocks available ' + 'to continue generating a request.', + ) + group.add_argument( + '--inference-dynamic-batching-mamba-memory-ratio', + type=float, + default=None, + help='Percentage of memory buffer to allocate for Mamba states. ' + 'If not specified, allocates Mamba state tensors for each KV cache block. ' + 'Only used for hybrid models.', + ) + group.add_argument( + '--inference-dynamic-batching-block-size', + type=int, + default=256, + help='KV cache block size. ' 'It should be a multiple of 256', + ) + group.add_argument( + '--inference-dynamic-batching-max-requests', + type=int, + default=None, + help='Override the inference context\'s `max_requests`. ' + 'By default, `max_requests` is set to the number of ' + 'blocks in the context\'s memory buffer.', + ) + group.add_argument( + '--inference-dynamic-batching-max-tokens', + type=int, + default=None, + help='Override the inference context\'s default `max_tokens`.', + ) + group.add_argument( + '--inference-dynamic-batching-num-cuda-graphs', + type=int, + default=16, + help='Maximum number of cuda graphs to capture, where the ' + 'cuda graph batch sizes range from 1 to `max_requests`. ' + '(See `dynamic_context.py` for details on how ' + '`max_requests` is computed). Due to rounding, the actual ' + 'number of cuda graphs may not equal this argument.' + 'The user can also pass -1, in which case we automatically determine the number of graphs ' + 'to capture based on the `max_requests`.', + ) + group.add_argument( + '--inference-dynamic-batching-track-paused-request-events', + action='store_true', + help='Track paused request ids by adding \'paused\' events ' + 'to each request\'s event history. This has a very minor ' + 'impact on latency.', + ) + group.add_argument( + '--inference-dynamic-batching-track-generated-token-events', + action='store_true', + help='Track per-token events with timestamps for each generated token. ' + 'When enabled, each generated token creates a GENERATED_TOKEN event ' + 'with a timestamp, useful for per-token latency analysis.', + ) + group.add_argument( + '--decode-only-cuda-graphs', + action='store_true', + default=False, + help='Only use cuda graphs for decode-only steps, not prefill and mixed steps.', + ) + group.add_argument( + '--inference-dynamic-batching-unified-memory-level', + type=int, + default=0, + choices=[0, 1], + help='Set unified memory usage within the dynamic ' + 'inference context. The levels are: 0) no unified memory, ' + '1) allocate `memory_buffer` in unified memory. ' + 'Eventually, additional levels will be included to ' + 'control other tensors within the context.', + ) + group.add_argument( + '--enable-chunked-prefill', + dest='enable_chunked_prefill', + action='store_true', + default=False, + help="Enable chunked prefill (disabled by default)", + ) + group.add_argument( + '--num-speculative-tokens', + type=int, + default=0, + help='Number of speculative tokens generated during decode', + ) + group.add_argument( + '--inference-dynamic-batching-prefix-caching', + dest='inference_dynamic_batching_enable_prefix_caching', + action=argparse.BooleanOptionalAction, + default=False, + help='Enable/disable prefix caching for dynamic batching inference. ' + 'When disabled, KV cache blocks cannot be shared between ' + 'requests with identical prompt prefixes.', + ) + group.add_argument( + '--inference-dynamic-batching-prefix-caching-eviction-policy', + type=str, + default='ref_zero', + choices=['ref_zero', 'lru'], + dest='inference_dynamic_batching_prefix_caching_eviction_policy', + help='Eviction policy for prefix caching blocks. ' + '"ref_zero" (default) immediately returns blocks to the ' + 'free pool when ref_count hits 0. "lru" keeps blocks ' + 'cached and evicts via LRU only when space is needed.', + ) + group.add_argument( + '--inference-dynamic-batching-prefix-caching-coordinator-policy', + type=str, + default='first_prefix_block', + choices=['longest_prefix', 'first_prefix_block', 'round_robin'], + dest='inference_dynamic_batching_prefix_caching_coordinator_policy', + help='Coordinator routing policy for prefix caching. ' + '"first_prefix_block" (default) routes based on the first ' + 'block hash only. "longest_prefix" routes to the rank with ' + 'the longest matching prefix. "round_robin" ignores prefix ' + 'affinity and cycles through ranks.', + ) + group.add_argument( + '--inference-dynamic-batching-prefix-caching-routing-alpha', + type=float, + default=0.5, + dest='inference_dynamic_batching_prefix_caching_routing_alpha', + help='Weight for prefix-aware routing score: ' + 'score = alpha * match + (1 - alpha) * normalized_load. ' + 'Higher alpha favors prefix cache hits; lower alpha ' + 'favors load balance. Default: 0.5.', + ) + group.add_argument( + '--inference-dynamic-batching-prefix-caching-mamba-gb', + type=float, + default=None, + dest='inference_dynamic_batching_prefix_caching_mamba_gb', + help='GPU memory budget (in GB) for the Mamba state cache ' + 'used by prefix caching on hybrid models. When set, Mamba ' + 'states at block boundaries are cached for reuse.', + ) + group.add_argument( + '--inference-dynamic-batching-cuda-graph-max-tokens', + type=int, + default=16384, + help='Maximum number of tokens to capture in a cuda graph.', + ) + group.add_argument( + '--inference-dynamic-batching-cuda-graph-mixed-prefill-count', + type=int, + default=16, + help='Number of mixed prefill requests to capture in a cuda graph.', + ) + group.add_argument( + '--inference-dynamic-batching-sampling-backend', + type=str, + default='torch', + choices=['torch', 'flashinfer'], + help='Which sampling kernels to use during inference. ' + 'Falls back to "torch" with a warning if "flashinfer" ' + 'is requested but the package is not installed.', + ) + group.add_argument( + '--inference-dynamic-batching-async-sched-mode', + type=str, + default='legacy', + choices=['legacy', 'serial'], + help='Async scheduling mode for dynamic batching. ' + '"legacy" (default) preserves the existing resolve-before-prepare ' + 'path. "serial" speculatively prepares and forwards decode-only ' + 'steps before resolving finished requests.', + ) + group.add_argument( + '--inference-dynamic-batching-logprobs-mode', + type=str, + default='raw_logprobs', + choices=['raw_logprobs', 'processed_logprobs'], + help='How returned inference log-probs are computed engine-wide. ' + '"raw_logprobs" (default) uses the unmodified model logits; ' + '"processed_logprobs" uses temperature and filters by top-k/top-p.', + ) + group.add_argument( + '--inference-logging-step-interval', + type=int, + default=0, + help='Step interval for logging inference metrics. ' + 'Default to 0 to disable inference logging.', + ) + group.add_argument( + '--inference-text-gen-server-logging', + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help='Enable per-request logging in the inference text generation server.', + ) + group.add_argument( + '--inference-wandb-logging', + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help='Enable inference wandb logging.', + ) + group.add_argument( + "--inference-coordinator-port", + type=int, + help="This port will be used to setup the inference coordinator on node-0", + ) + group.add_argument( + '--mamba-inference-conv-states-dtype', + type=str, + choices=['bf16', 'fp16', 'fp32'], + default='bf16', + help='Dtype for the Mamba inference conv states tensor', + ) + group.add_argument( + '--mamba-inference-ssm-states-dtype', + type=str, + choices=['bf16', 'fp16', 'fp32'], + default='bf16', + help='Dtype for the Mamba inference SSM states tensor', + ) + group.add_argument( + '--inference-use-synchronous-zmq-collectives', + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help='Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for MoEs.', + ) + group.add_argument( + '--inference-disable-ep-consensus', + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help='Skip the EP-group consensus all-reduce in the inference engine control loop and step on local state only. ' + 'Pause/unpause take effect as soon as the signal is delivered to a rank. ' + 'Only safe when EP coordination is not required (e.g. ep_world_size == 1).', + ) + group.add_argument( + '--inference-shards', + type=str, + default=None, + metavar='SPEC', + help='Partition the world into independent inference models, each with ' + 'its own parallelism, e.g. "tp=2,role=prefill+tp=1,dp=2,role=decode". ' + 'Shards are separated by "+" or ";"; per-shard keys are ' + 'tp,pp,ep,expt_tp,dp (each defaults to 1) and must partition the full ' + 'world. Tagging shards role=prefill|decode enables disaggregated ' + 'inference (prefill hands KV to the decode pool); a dp>1 decode shard ' + 'is several independent decode instances.', + ) + group.add_argument( + '--inference-cuda-graph-all-prefills', + action='store_true', + default=False, + help='Extend prefill/mixed CUDA graph capture up to `max_tokens`. ' + 'By default, all graphs are limited by the decode limit of ' + '`max_requests * (num_speculative_tokens + 1)`.', + ) + group.add_argument( + '--inference-dynamic-batching-cuda-graph-sizing-distribution', + type=str, + default='exponential', + choices=['exponential', 'linear'], + dest='inference_dynamic_batching_cuda_graph_sizing_distribution', + help='Spacing of CUDA graph token counts. "exponential" (default) ' + 'halves from cuda_graph_max_tokens down to tp_size, giving a ' + 'log-spaced distribution with bounded relative padding. ' + '"linear" uses varying linear strides across the range.', + ) return parser @@ -2366,6 +2749,7 @@ def _add_network_size_args(parser): "no_rope_freq", "moe_layer_freq", "linear_attention_freq", + "csa_compress_ratios", "moe_router_load_balancing_type", "moe_aux_loss_coeff", "cp_comm_type", @@ -2421,8 +2805,11 @@ def _add_network_size_args(parser): "barrier_with_L1_time", # args uses same var with a different name "num_moe_experts", + "actual_vocab_size", "fp8_param", "fp4_param", + # already generated by data args + "create_attention_mask_in_dataloader", # incompatible defaults in dataclass "gradient_accumulation_fusion", "overlap_p2p_comm", @@ -2431,249 +2818,418 @@ def _add_network_size_args(parser): "persist_layer_norm", "bias_dropout_fusion", "apply_rope_fusion", + "apply_dsa_kernel_fusion", + "dsa_kernel_backend", "mamba_training_ssm_states_dtype", - "max_seqlen_per_dp_cp_rank", - "hybrid_context_parallel", - "sequence_packing_scheduler", - # internal/derived: controlled only via --tensor-parallel-num-weight-shards - "gtp_weight_remat_size", - # internal/derived: controlled only via --expert-tensor-parallel-num-weight-shards - "expert_gtp_weight_remat_size", - "max_seqlen_per_dp_cp_rank", - "hybrid_context_parallel", - "sequence_packing_scheduler", ] transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude) transformer_group = transformer_factory.build_group(parser, "transformer configuration") group = parser.add_argument_group(title='network size') - group.add_argument('--encoder-num-layers', type=int, default=None, - help='Number of encoder transformer layers.') - group.add_argument('--decoder-num-layers', type=int, default=None, - help='Number of decoder transformer layers.') - group.add_argument('--group-query-attention', action='store_true', - help='Use group-query attention.') - group.add_argument('--window-size', type=tuple_type, default=None, - help='Window size for window attention. If not provided, ' - 'window attention will be disabled.') - group.add_argument('--window-attn-skip-freq', type=moe_freq_type, default=None, - help='Frequency of layers to skip window attention. Accepts either: ' - '- An integer N: Represents a (N-1):1 ratio, meaning one full attention layer ' - 'after (N-1) SWA layers. ' - '- A string containing a Python list expression that defines a custom pattern, ' - 'e.g.: "[1,1,1,0]*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' - 'where 1 indicates SWA and 0 indicates full attention. ') - group.add_argument('--max-position-embeddings', type=int, default=None, - help='Maximum number of position embeddings to use. ' - 'This is the size of position embedding.') - group.add_argument('--position-embedding-type', type=str, default='learned_absolute', - choices=['learned_absolute', 'rope', 'yarn', 'mrope', 'relative', 'none'], - help='Position embedding type.') - group.add_argument('--yarn-original-max-position-embeddings', type=int, default=None, - help='Original maximum position embeddings for YaRN RoPE frequency correction.') - group.add_argument('--yarn-beta-fast', type=float, default=None, - help='Beta fast for YaRN RoPE frequency correction.') - group.add_argument('--yarn-beta-slow', type=float, default=None, - help='Beta slow for YaRN RoPE frequency correction.') - group.add_argument('--yarn-correction-range-round-to-int', action='store_true', default=None, - help='Round YaRN correction range endpoints to integers.') - group.add_argument('--no-yarn-correction-range-round-to-int', action='store_false', - dest='yarn_correction_range_round_to_int', - help='Do not round YaRN correction range endpoints to integers.') - group.add_argument('--relative-attention-num-buckets', type=int, default=32, - help='Number of buckets for relative position embeddings.') - group.add_argument('--relative-attention-max-distance', type=int, default=128, - help='Maximum distance for relative position embeddings calculation.') - group.add_argument('--use-rotary-position-embeddings', action='store_true', - help='Use rotary positional embeddings or not. ' - 'Deprecated: use --position-embedding-type') - group.add_argument('--rotary-base', type=int, default=10000, - help='Base to use for rotary positional embeddings, default 10000') - group.add_argument('--rotary-percent', type=float, default=1.0, - help='Percent of rotary dimension to use, default 100%%') - group.add_argument('--rotary-seq-len-interpolation-factor', type=int, default=None, - help='Sequence length interpolation factor for rotary embeddings.') - group.add_argument('--use-rope-scaling', action='store_true', - help='Apply rope scaling as used in llama3.x') - group.add_argument('--rope-scaling-factor', type=float, default=8.0, - help='Rope scaling factor in llama3.x models') - group.add_argument('--no-rope-freq', type=no_rope_freq_type, default=None, - help='Controls which layers to skip performing Rotary Position Embedding. Accepts either: ' - '- An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. ' - '- A string containing a Python list expression that defines a custom pattern, e.g.: ' - '"([0]*3+[1]*1)*3" evaluates to [0,0,0,1,0,0,0,1,0,0,0,1] ' - 'where 1 indicates no-rope layer. This patten is equivalent to --no-rope-freq=4.' - 'By default this is disabled and set to None, indicating RoPE will be performed' - 'on every layer.' - ) - group.add_argument('--no-position-embedding', - action='store_false', - help='Disable position embedding. Deprecated: use --position-embedding-type', - dest='add_position_embedding') - group.add_argument('--make-vocab-size-divisible-by', type=int, default=128, - help='Pad the vocab size to be divisible by this value.' - 'This is added for computational efficieny reasons.') - group.add_argument('--openai-gelu', action='store_true', - help='Use OpenAIs GeLU implementation. This option' - 'should not be used unless for backward compatibility' - 'reasons.') - group.add_argument('--squared-relu', action='store_true', - help='Use squared relu activation instead of default gelu') - group.add_argument('--swiglu', action='store_true', - help='Use gated linear units and SiLU activation instead of default gelu') - group.add_argument('--quick-geglu', action='store_true', - help='Use quick geglu activation instead of default gelu') - group.add_argument('--onnx-safe', type=bool, required=False, - help='Use workarounds for known problems with ' - 'Torch ONNX exporter') - group.add_argument('--bert-no-binary-head', action='store_false', - help='Disable BERT binary head.', - dest='bert_binary_head') - group.add_argument('--untie-embeddings-and-output-weights', action='store_true', - help='Untie embeddings and output weights.') - return parser - -def _add_straggler_detector_args(parser): - from megatron.training.config import StragglerDetectionConfig - - straggler_factory = ArgumentGroupFactory(StragglerDetectionConfig) - group = straggler_factory.build_group(parser, "straggler") - - return parser - -def _add_workload_inspector_server_args(parser): - group = parser.add_argument_group(title='workload inspector') - group.add_argument('--run-workload-inspector-server', action='store_true', - help='If set, enables workload inspector server for on-demand profiling.') - return parser - -def _add_inprocess_restart_args(parser): - group = parser.add_argument_group(title='In-process restart') - - group.add_argument('--inprocess-restart', action='store_true', - help='Enables in-process restart.') - - group.add_argument('--inprocess-max-iterations', default=None, type=int, - help='Maximum number of in-process restart iterations.') - group.add_argument('--inprocess-monitor-thread-interval', default=1.0, type=float, - help='Monitoring interval (in seconds) for the monitoring thread.') - group.add_argument('--inprocess-monitor-process-interval', default=1.0, type=float, - help='Monitoring interval (in seconds) for the monitoring process.') - group.add_argument('--inprocess-progress-watchdog-interval', default=1.0, type=float, - help='Interval (in seconds) for automatic progress watchdog timestamp ' - 'updates.') - group.add_argument('--inprocess-heartbeat-interval', default=30, type=float, - help='Monitoring interval (in seconds) for detecting unresponsive ranks.') - - group.add_argument('--inprocess-soft-timeout', default=60, type=float, - help='Soft progress timeout (in seconds).') - group.add_argument('--inprocess-hard-timeout', default=90, type=float, - help='Hard progress timeout (in seconds).') - group.add_argument('--inprocess-heartbeat-timeout', default=60, type=float, - help='Timeout (in seconds) for a missing rank detection heartbeat.') - - group.add_argument('--inprocess-barrier-timeout', default=120, type=float, - help='Timeout (in seconds) for internal distributed barrier') - group.add_argument('--inprocess-completion-timeout', default=120, type=float, - help='Timeout (in seconds) for barrier on completion on all ranks') - - group.add_argument('--inprocess-last-call-wait', default=1, type=float, - help='Time interval (in seconds) for other ranks to report concurrent ' - 'terminal failures.') - group.add_argument('--inprocess-termination-grace-time', default=1, type=float, - help='Interval (in seconds) between SIGTERM and SIGKILL issued on hard ' - 'timeout') - - group.add_argument('--inprocess-granularity', default='node', type=str, - choices=['node', 'rank'], - help='Granularity for in-process restart.') - group.add_argument('--inprocess-active-world-size', - default=int(os.getenv('WORLD_SIZE', '1')), type=int, - help='The number of ranks initially executing the workload. ' - 'The remaining ranks from the allocation are set aside ' - 'as warm reserve.') - group.add_argument('--inprocess-empty-cuda-cache', action='store_true', - help='Release all unoccupied cached GPU memory on every in-process restart.') - return parser - -def _add_one_logger_args(parser): - group = parser.add_argument_group(title='one logger') - group.add_argument('--no-one-logger', action='store_false', - help='If set, disable using one_logger to track E2E metrics' - 'Note that one_logger is an internal tool and not ' - 'available externally. For installation, please go to ' - 'https://confluence.nvidia.com/display/MLWFO/Package+Repositories' - 'for more details', - dest='enable_one_logger') - group.add_argument('--one-logger-project', type=str, default='megatron-lm', - help='The one-logger project name. Will ignore if ' - '--no-one-logger is set') - group.add_argument('--one-logger-run-name', type=str, default=None, - help='The one-logger run name displayed. Will ignore if ' - '--no-one-logger is set') - group.add_argument('--one-logger-async', action='store_true', - help='If set, forces one_logger to use async mode.') - group.add_argument('--app-tag-run-name', type=str, default=None, - help='Jobs belonging to same training run, suppose to ' - 'have the same name. It will be used to track progress of ' - 'a training done over multiple different jobs') - group.add_argument('--app-tag-run-version', type=str, default='0.0.0', - help='The version of the training of which current job is ' - 'part of. It will be used to track the changes in the ' - 'application side which might change the performance ' - 'baseline') - return parser - - -def _add_ft_package_args(parser): - group = parser.add_argument_group(title='ft_package') - group.add_argument('--enable-ft-package', action='store_true', - help='If set, Fault Tolerance package is enabled. ' - 'Note: This feature is for Nvidia internal use only.') - group.add_argument('--calc-ft-timeouts', action='store_true', - help='If set, FT package will try to automatically compute the timeouts. ' - 'Note: This feature is for Nvidia internal use only.') - group.add_argument('--ft-num-warmup-iters', type=int, default=5, - help='Number of warmup iterations before monitoring step section and ' - 'out-of-section timeouts. The first N iterations are excluded from ' - 'timeout monitoring as they can be significantly slower than steady-state. ' - 'Default: 5. Note: This feature is for Nvidia internal use only.') - return parser - - -def _add_logging_args(parser): - from megatron.training.config import LoggerConfig - - log_factory = ArgumentGroupFactory(LoggerConfig, exclude = ["log_throughput_to_tensorboard", "throughput_window_size", "memory_keys", "log_l2_norm_grad_to_tensorboard", "log_runtime_to_tensorboard", "runtime_time_unit", "filter_warnings", "modules_to_filter", "set_level_for_all_loggers", "save_config_filepath"]) - group = log_factory.build_group(parser, title="logging") + group.add_argument( + '--encoder-num-layers', type=int, default=None, help='Number of encoder transformer layers.' + ) + group.add_argument( + '--decoder-num-layers', type=int, default=None, help='Number of decoder transformer layers.' + ) + group.add_argument( + '--group-query-attention', action='store_true', help='Use group-query attention.' + ) + group.add_argument( + '--window-size', + type=tuple_type, + default=None, + help='Window size for window attention. If not provided, ' + 'window attention will be disabled.', + ) + group.add_argument( + '--window-attn-skip-freq', + type=moe_freq_type, + default=None, + help='Frequency of layers to skip window attention. Accepts either: ' + '- An integer N: Represents a (N-1):1 ratio, meaning one full attention layer ' + 'after (N-1) SWA layers. ' + '- A string containing a Python list expression that defines a custom pattern, ' + 'e.g.: "[1,1,1,0]*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' + 'where 1 indicates SWA and 0 indicates full attention. ', + ) + group.add_argument( + '--max-position-embeddings', + type=int, + default=None, + help='Maximum number of position embeddings to use. ' + 'This is the size of position embedding.', + ) + group.add_argument( + '--position-embedding-type', + type=str, + default='learned_absolute', + choices=['learned_absolute', 'rope', 'yarn', 'mrope', 'relative', 'none'], + help='Position embedding type.', + ) + group.add_argument( + '--yarn-original-max-position-embeddings', + type=int, + default=None, + help='Original maximum position embeddings for YaRN RoPE frequency correction.', + ) + group.add_argument( + '--yarn-beta-fast', + type=float, + default=None, + help='Beta fast for YaRN RoPE frequency correction.', + ) + group.add_argument( + '--yarn-beta-slow', + type=float, + default=None, + help='Beta slow for YaRN RoPE frequency correction.', + ) + group.add_argument( + '--yarn-correction-range-round-to-int', + action='store_true', + default=None, + help='Round YaRN correction range endpoints to integers.', + ) + group.add_argument( + '--no-yarn-correction-range-round-to-int', + action='store_false', + dest='yarn_correction_range_round_to_int', + help='Do not round YaRN correction range endpoints to integers.', + ) + group.add_argument( + '--relative-attention-num-buckets', + type=int, + default=32, + help='Number of buckets for relative position embeddings.', + ) + group.add_argument( + '--relative-attention-max-distance', + type=int, + default=128, + help='Maximum distance for relative position embeddings calculation.', + ) + group.add_argument( + '--use-rotary-position-embeddings', + action='store_true', + help='Use rotary positional embeddings or not. ' + 'Deprecated: use --position-embedding-type', + ) + group.add_argument( + '--rotary-base', + type=int, + default=10000, + help='Base to use for rotary positional embeddings, default 10000', + ) + group.add_argument( + '--rotary-percent', + type=float, + default=1.0, + help='Percent of rotary dimension to use, default 100%%', + ) + group.add_argument( + '--rotary-seq-len-interpolation-factor', + type=int, + default=None, + help='Sequence length interpolation factor for rotary embeddings.', + ) + group.add_argument( + '--use-rope-scaling', action='store_true', help='Apply rope scaling as used in llama3.x' + ) + group.add_argument( + '--rope-scaling-factor', + type=float, + default=8.0, + help='Rope scaling factor in llama3.x models', + ) + group.add_argument( + '--no-rope-freq', + type=no_rope_freq_type, + default=None, + help='Controls which layers to skip performing Rotary Position Embedding. Accepts either: ' + '- An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. ' + '- A string containing a Python list expression that defines a custom pattern, e.g.: ' + '"([0]*3+[1]*1)*3" evaluates to [0,0,0,1,0,0,0,1,0,0,0,1] ' + 'where 1 indicates no-rope layer. This patten is equivalent to --no-rope-freq=4.' + 'By default this is disabled and set to None, indicating RoPE will be performed' + 'on every layer.', + ) + group.add_argument( + '--no-position-embedding', + action='store_false', + help='Disable position embedding. Deprecated: use --position-embedding-type', + dest='add_position_embedding', + ) + group.add_argument( + '--make-vocab-size-divisible-by', + type=int, + default=128, + help='Pad the vocab size to be divisible by this value.' + 'This is added for computational efficieny reasons.', + ) + group.add_argument( + '--openai-gelu', + action='store_true', + help='Use OpenAIs GeLU implementation. This option' + 'should not be used unless for backward compatibility' + 'reasons.', + ) + group.add_argument( + '--squared-relu', + action='store_true', + help='Use squared relu activation instead of default gelu', + ) + group.add_argument( + '--swiglu', + action='store_true', + help='Use gated linear units and SiLU activation instead of default gelu', + ) + group.add_argument( + '--quick-geglu', + action='store_true', + help='Use quick geglu activation instead of default gelu', + ) + group.add_argument( + '--onnx-safe', + type=bool, + required=False, + help='Use workarounds for known problems with ' 'Torch ONNX exporter', + ) + group.add_argument( + '--bert-no-binary-head', + action='store_false', + help='Disable BERT binary head.', + dest='bert_binary_head', + ) + group.add_argument( + '--untie-embeddings-and-output-weights', + action='store_true', + help='Untie embeddings and output weights.', + ) + return parser + + +def _add_straggler_detector_args(parser): + from megatron.training.config import StragglerDetectionConfig + + straggler_factory = ArgumentGroupFactory(StragglerDetectionConfig) + group = straggler_factory.build_group(parser, "straggler") + + return parser + - otel_group = parser.add_argument_group(title='opentelemetry') - otel_group.add_argument( - '--otel-enabled', +def _add_workload_inspector_server_args(parser): + group = parser.add_argument_group(title='workload inspector') + group.add_argument( + '--run-workload-inspector-server', action='store_true', - default=False, - help='Enable OpenTelemetry telemetry (traces and metrics). ' - 'See MEGATRON_OTEL_ENABLED env var for the env-var equivalent.', + help='If set, enables workload inspector server for on-demand profiling.', + ) + return parser + + +def _add_inprocess_restart_args(parser): + group = parser.add_argument_group(title='In-process restart') + + group.add_argument( + '--inprocess-restart', action='store_true', help='Enables in-process restart.' + ) + + group.add_argument( + '--inprocess-max-iterations', + default=None, + type=int, + help='Maximum number of in-process restart iterations.', + ) + group.add_argument( + '--inprocess-monitor-thread-interval', + default=1.0, + type=float, + help='Monitoring interval (in seconds) for the monitoring thread.', + ) + group.add_argument( + '--inprocess-monitor-process-interval', + default=1.0, + type=float, + help='Monitoring interval (in seconds) for the monitoring process.', + ) + group.add_argument( + '--inprocess-progress-watchdog-interval', + default=1.0, + type=float, + help='Interval (in seconds) for automatic progress watchdog timestamp ' 'updates.', + ) + group.add_argument( + '--inprocess-heartbeat-interval', + default=30, + type=float, + help='Monitoring interval (in seconds) for detecting unresponsive ranks.', + ) + + group.add_argument( + '--inprocess-soft-timeout', + default=60, + type=float, + help='Soft progress timeout (in seconds).', + ) + group.add_argument( + '--inprocess-hard-timeout', + default=90, + type=float, + help='Hard progress timeout (in seconds).', + ) + group.add_argument( + '--inprocess-heartbeat-timeout', + default=60, + type=float, + help='Timeout (in seconds) for a missing rank detection heartbeat.', ) - otel_group.add_argument( - '--otel-service-name', + + group.add_argument( + '--inprocess-barrier-timeout', + default=120, + type=float, + help='Timeout (in seconds) for internal distributed barrier', + ) + group.add_argument( + '--inprocess-completion-timeout', + default=120, + type=float, + help='Timeout (in seconds) for barrier on completion on all ranks', + ) + + group.add_argument( + '--inprocess-last-call-wait', + default=1, + type=float, + help='Time interval (in seconds) for other ranks to report concurrent ' + 'terminal failures.', + ) + group.add_argument( + '--inprocess-termination-grace-time', + default=1, + type=float, + help='Interval (in seconds) between SIGTERM and SIGKILL issued on hard ' 'timeout', + ) + + group.add_argument( + '--inprocess-granularity', + default='node', + type=str, + choices=['node', 'rank'], + help='Granularity for in-process restart.', + ) + group.add_argument( + '--inprocess-active-world-size', + default=int(os.getenv('WORLD_SIZE', '1')), + type=int, + help='The number of ranks initially executing the workload. ' + 'The remaining ranks from the allocation are set aside ' + 'as warm reserve.', + ) + group.add_argument( + '--inprocess-empty-cuda-cache', + action='store_true', + help='Release all unoccupied cached GPU memory on every in-process restart.', + ) + return parser + + +def _add_one_logger_args(parser): + group = parser.add_argument_group(title='one logger') + group.add_argument( + '--no-one-logger', + action='store_false', + help='If set, disable using one_logger to track E2E metrics' + 'Note that one_logger is an internal tool and not ' + 'available externally. For installation, please go to ' + 'https://confluence.nvidia.com/display/MLWFO/Package+Repositories' + 'for more details', + dest='enable_one_logger', + ) + group.add_argument( + '--one-logger-project', + type=str, + default='megatron-lm', + help='The one-logger project name. Will ignore if ' '--no-one-logger is set', + ) + group.add_argument( + '--one-logger-run-name', type=str, default=None, - help='Override OTEL_SERVICE_NAME for this training run.', + help='The one-logger run name displayed. Will ignore if ' '--no-one-logger is set', + ) + group.add_argument( + '--one-logger-async', + action='store_true', + help='If set, forces one_logger to use async mode.', ) - otel_group.add_argument( - '--otel-span-groups', + group.add_argument( + '--app-tag-run-name', type=str, default=None, - help='Comma-separated span-group spec controlling which OTel ' - 'instrumentation boundaries are active. Accepts preset keywords ' - '("default", "per_step", "full", "all") or individual group names ' - '("job", "checkpoint", "evaluate", "model_init", "load_checkpoint", ' - '"step", "forward_backward", "optimizer", "microbatch"), or a mix. ' - 'Defaults to "default" (coarse job/checkpoint/evaluate spans only). ' - 'Equivalent to MEGATRON_OTEL_SPAN_GROUPS env var.', + help='Jobs belonging to same training run, suppose to ' + 'have the same name. It will be used to track progress of ' + 'a training done over multiple different jobs', + ) + group.add_argument( + '--app-tag-run-version', + type=str, + default='0.0.0', + help='The version of the training of which current job is ' + 'part of. It will be used to track the changes in the ' + 'application side which might change the performance ' + 'baseline', + ) + return parser + + +def _add_ft_package_args(parser): + group = parser.add_argument_group(title='ft_package') + group.add_argument( + '--enable-ft-package', + action='store_true', + help='If set, Fault Tolerance package is enabled. ' + 'Note: This feature is for Nvidia internal use only.', + ) + group.add_argument( + '--calc-ft-timeouts', + action='store_true', + help='If set, FT package will try to automatically compute the timeouts. ' + 'Note: This feature is for Nvidia internal use only.', + ) + group.add_argument( + '--ft-num-warmup-iters', + type=int, + default=5, + help='Number of warmup iterations before monitoring step section and ' + 'out-of-section timeouts. The first N iterations are excluded from ' + 'timeout monitoring as they can be significantly slower than steady-state. ' + 'Default: 5. Note: This feature is for Nvidia internal use only.', + ) + return parser + + +def _add_logging_args(parser): + from megatron.training.config import LoggerConfig + + log_factory = ArgumentGroupFactory( + LoggerConfig, + exclude=[ + "log_throughput_to_tensorboard", + "throughput_window_size", + "memory_keys", + "log_l2_norm_grad_to_tensorboard", + "log_runtime_to_tensorboard", + "runtime_time_unit", + "filter_warnings", + "modules_to_filter", + "set_level_for_all_loggers", + "save_config_filepath", + ], ) + group = log_factory.build_group(parser, title="logging") return parser @@ -2681,639 +3237,1100 @@ def _add_logging_args(parser): def _add_regularization_args(parser): group = parser.add_argument_group(title='regularization') - group.add_argument('--weight-decay', type=float, default=0.01, - help='Weight decay coefficient for L2 regularization.') - group.add_argument('--apply-wd-to-qk-layernorm', action='store_true', - help='Apply weight decay to qk layernorm as a special case.') - group.add_argument('--clip-grad', type=float, default=1.0, - help='Gradient clipping based on global L2 norm.') - group.add_argument('--adam-beta1', type=float, default=0.9, - help='First coefficient for computing running averages ' - 'of gradient and its square') - group.add_argument('--adam-beta2', type=float, default=0.999, - help='Second coefficient for computing running averages ' - 'of gradient and its square') - group.add_argument('--adam-eps', type=float, default=1e-08, - help='Term added to the denominator to improve' - 'numerical stability') - group.add_argument('--sgd-momentum', type=float, default=0.9, - help='Momentum factor for sgd') - group.add_argument('--muon-momentum', type=float, default=0.9, - help='Momentum factor for Muon optimizer') - group.add_argument('--muon-no-split-qkv', action='store_false', default=True, - dest='muon_split_qkv', - help='Whether to split QKV parameters for Muon optimizer') - group.add_argument('--muon-nesterov', action='store_true', - help='Whether to use Nesterov-style momentum in the internal SGD') - group.add_argument('--muon-scale-mode', type=str, default='spectral', - choices=['spectral', 'unit_rms_norm', 'shape_scaling'], - help='Scale mode for Muon optimizer. With MuP, set ' - '--muon-scale-mode unit_rms_norm to use unit_rms_norm scaling, ' - 'or set --muon-scale-mode spectral to keep spectral scaling.') - group.add_argument('--muon-fp32-matmul-prec', type=str, default='medium', - choices=['low', 'medium', 'high'], - help='FP32 matmul precision for Newton-Schulz iteration') - group.add_argument('--muon-coefficient-type', type=str, default='quintic', - help='Newton-Schulz coefficient type for the Muon optimizer. ' - 'Valid types are discovered from the installed emerging_optimizers ' - 'package (e.g. simple, quintic, polar_express, aol). ' - 'Validated at optimizer creation time.') - group.add_argument('--muon-num-ns-steps', type=int, default=5, - help='Number of Newton-Schulz steps for Muon optimizer') - group.add_argument('--muon-tp-mode', type=str, default='duplicated', - choices=['blockwise', 'duplicated', 'distributed', 'auto'], - help='How to perform NS calculation for tensor model parallel weights. ' - 'blockwise orthogonalizes each shard on its own, so the update rule ' - 'depends on the parallelism config; duplicated and distributed both ' - 'orthogonalize the whole matrix and give TP-invariant results; auto ' - 'select between duplicated and distributed mode per-weight for ' - 'dense weights.') - group.add_argument('--muon-use-syrk', action='store_true', - help='Use the Triton SYRK kernel for the Gram matrix ' - 'in Newton-Schulz iteration.') - group.add_argument('--muon-extra-scale-factor', type=float, default=1.0, - help='Additional scale factor for the muon update') - group.add_argument('--muon-scalar-optimizer', type=str, default='adam', - choices=['adam', 'lion'], - help='Optimizer for scalar parameters (embeddings, biases, norms) ' - 'when using muon. Defaults to adam.') - group.add_argument('--lion-beta1', type=float, default=0.95, - help='First beta coefficient for Lion optimizer ' - '(used in sign update). Default: 0.95.') - group.add_argument('--lion-beta2', type=float, default=0.98, - help='Second beta coefficient for Lion optimizer ' - '(used in momentum EMA update). Default: 0.98.') - - group.add_argument('--no-weight-decay-cond-type', type=str, choices=['apply_wd_to_qk_layernorm'], - help='Type of no weight decay condition. Choices: ' - 'None (default): apply weight decay to 1D weights and biases.' - '"apply_wd_to_qk_layernorm": additionally apply weight decay to ' - 'qk layernorm as a special case.' - 'DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. ') + group.add_argument( + '--weight-decay', + type=float, + default=0.01, + help='Weight decay coefficient for L2 regularization.', + ) + group.add_argument( + '--apply-wd-to-qk-layernorm', + action='store_true', + help='Apply weight decay to qk layernorm as a special case.', + ) + group.add_argument( + '--clip-grad', type=float, default=1.0, help='Gradient clipping based on global L2 norm.' + ) + group.add_argument( + '--adam-beta1', + type=float, + default=0.9, + help='First coefficient for computing running averages ' 'of gradient and its square', + ) + group.add_argument( + '--adam-beta2', + type=float, + default=0.999, + help='Second coefficient for computing running averages ' 'of gradient and its square', + ) + group.add_argument( + '--adam-eps', + type=float, + default=1e-08, + help='Term added to the denominator to improve' 'numerical stability', + ) + group.add_argument('--sgd-momentum', type=float, default=0.9, help='Momentum factor for sgd') + group.add_argument( + '--muon-momentum', type=float, default=0.9, help='Momentum factor for Muon optimizer' + ) + group.add_argument( + '--muon-no-split-qkv', + action='store_false', + default=True, + dest='muon_split_qkv', + help='Whether to split QKV parameters for Muon optimizer', + ) + group.add_argument( + '--muon-nesterov', + action='store_true', + help='Whether to use Nesterov-style momentum in the internal SGD', + ) + group.add_argument( + '--muon-scale-mode', + type=str, + default='spectral', + choices=['spectral', 'unit_rms_norm', 'shape_scaling'], + help='Scale mode for Muon optimizer. With MuP, set ' + '--muon-scale-mode unit_rms_norm to use unit_rms_norm scaling, ' + 'or set --muon-scale-mode spectral to keep spectral scaling.', + ) + group.add_argument( + '--muon-fp32-matmul-prec', + type=str, + default='medium', + choices=['low', 'medium', 'high'], + help='FP32 matmul precision for Newton-Schulz iteration', + ) + group.add_argument( + '--muon-coefficient-type', + type=str, + default='quintic', + help='Newton-Schulz coefficient type for the Muon optimizer. ' + 'Valid types are discovered from the installed emerging_optimizers ' + 'package (e.g. simple, quintic, polar_express, aol). ' + 'Validated at optimizer creation time.', + ) + group.add_argument( + '--muon-num-ns-steps', + type=int, + default=5, + help='Number of Newton-Schulz steps for Muon optimizer', + ) + group.add_argument( + '--muon-tp-mode', + type=str, + default='blockwise', + choices=['blockwise', 'duplicated', 'distributed'], + help='How to perform NS calculation for tensor model parallel weights', + ) + group.add_argument( + '--muon-extra-scale-factor', + type=float, + default=1.0, + help='Additional scale factor for the muon update', + ) + group.add_argument( + '--muon-scalar-optimizer', + type=str, + default='adam', + choices=['adam', 'lion'], + help='Optimizer for scalar parameters (embeddings, biases, norms) ' + 'when using muon. Defaults to adam.', + ) + group.add_argument( + '--lion-beta1', + type=float, + default=0.95, + help='First beta coefficient for Lion optimizer ' '(used in sign update). Default: 0.95.', + ) + group.add_argument( + '--lion-beta2', + type=float, + default=0.98, + help='Second beta coefficient for Lion optimizer ' + '(used in momentum EMA update). Default: 0.98.', + ) + + group.add_argument( + '--no-weight-decay-cond-type', + type=str, + choices=['apply_wd_to_qk_layernorm'], + help='Type of no weight decay condition. Choices: ' + 'None (default): apply weight decay to 1D weights and biases.' + '"apply_wd_to_qk_layernorm": additionally apply weight decay to ' + 'qk layernorm as a special case.' + 'DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. ', + ) return parser def _add_rl_args(parser): group = parser.add_argument_group(title='rl') - group.add_argument('--perform-rl-step', action='store_true', - help="Use the RL training step.") - group.add_argument('--rl-prompts-per-eval', type=int, default=32, - help='Number of prompts to evaluate for for each RL task.' - 'This evaluation can be very expensive when using environments' - 'that evaluate pass@k so we default to a lower number.') + group.add_argument('--perform-rl-step', action='store_true', help="Use the RL training step.") + group.add_argument( + '--rl-prompts-per-eval', + type=int, + default=32, + help='Number of prompts to evaluate for for each RL task.' + 'This evaluation can be very expensive when using environments' + 'that evaluate pass@k so we default to a lower number.', + ) # TODO(rkirby): allow for "complete" evaluation when --rl-prompts-per-eval is set to -1 - group.add_argument('--grpo-prompts-per-step', type=int, default=32, - help="Number of GRPO groups (G in the paper).") - group.add_argument('--grpo-group-size', type=int, default=2, - help="Number of samples per a GRPO group.") - group.add_argument('--rl-generation-lag', type=float, default=None, - help='Number of trainer batches of rollout generation lag to allow ' - 'The number of in-flight trainer batches is this value plus one. ' - 'May be fractional or negative; the minimum of -1 keeps a single unit ' - 'of generation work in flight. If omitted, the lag is autotuned to the ' - 'inference engine\'s request capacity when --rl-partial-rollouts is ' - 'set, and is 0 otherwise. ' - 'Requires --rl-partial-rollouts when greater than 0. ' - 'Mutually exclusive with --rl-max-inflight-requests.') - group.add_argument('--rl-max-inflight-requests', type=int, default=None, - help='Maximum number of inference requests RL generation may keep inflight: ' - 'equivalent to (--rl-generation-lag + 1) training batches ' - 'of grpo_prompts_per_step * grpo_group_size requests each. ' - 'Requires --rl-partial-rollouts when above one training batch; ' - 'mutually exclusive with --rl-generation-lag.') + group.add_argument( + '--grpo-prompts-per-step', + type=int, + default=32, + help="Number of GRPO groups (G in the paper).", + ) + group.add_argument( + '--grpo-group-size', type=int, default=2, help="Number of samples per a GRPO group." + ) + group.add_argument( + '--rl-generation-lag', + type=int, + default=0, + help='Number of trainer batches of rollout generation lag to allow. ' + 'The number of in-flight trainer batches is this value plus one. ' + 'Requires --rl-partial-rollouts when greater than 0.', + ) # TODO: Refactor these string literals back to an enum after the megatron.training refactor. - group.add_argument('--rl-submission-granularity', type=str, - default="B", - choices=["R", "G", "B"], - help='Granularity for submitting rollout generation work. ' - 'R submits individual rollouts independently while still yielding ' - 'complete rollout groups to training. ' - 'G submits one rollout group at a time. ' - 'B submits grpo_prompts_per_step rollout groups together.') - group.add_argument('--rl-consumption-granularity', type=str, - default="B", - choices=["R", "G", "B"], - help='Granularity for consuming generated rollout groups. ' - 'G consumes groups as they complete. ' - 'B consumes complete trainer batches in submission order. ' - 'R is not currently supported.') - group.add_argument('--rl-durable-rollout-bank', action='store_true', - help='Persist completed rollout groups to a durable, write-through ' - 'ledger so they survive a SIGKILL (the SLURM time limit) and are ' - 'restored at restart instead of regenerated. No-op when unset.') - group.add_argument('--rl-rollout-bank-dir', type=str, default=None, - help='Directory for the durable rollout bank (on Lustre). Defaults to ' - '/rollout_bank so the bank stays coupled to the checkpoint.') - group.add_argument('--rl-rollout-bank-max-bytes', type=int, default=0, - help='Soft cap (bytes) on the rollout bank size; 0 = unbounded. On ' - 'exceed, a warning is logged. Compaction occurs at the next checkpoint ' - 'regardless of the cap and never blocks generation.') - group.add_argument('--grpo-iterations', type=int, default=2, - help="Number of iterations per a GRPO implementation.") - # As in DAPO, we keep upper/lower eps different. - # To have a vanilla GRPO, set them to be the same. - group.add_argument('--grpo-clamp-eps-lower', type=float, default=0.01, - help="Lower GRPO clipping bound.") - group.add_argument('--grpo-clamp-eps-upper', type=float, default=0.01, - help="Upper GRPO clipping bound. In vanilla implementation, equals to the lower one.") - group.add_argument('--grpo-kl-beta', type=float, default=0.001, - help="KL term weight in the GRPO loss.") - group.add_argument('--grpo-entropy-term-weight', type=float, default=0.0, - help="Entropy term weight in GRPO loss.") - group.add_argument('--grpo-filter-groups-with-same-reward', action='store_true', - help="Filter groups with same reward.") - group.add_argument('--langrl-env-config', type=str, default=None, - help="Path to YAML config file for RL environment configuration.") - group.add_argument('--rl-default-temperature', type=float, default=1.0, - help="Default temperature for model inference.") - group.add_argument('--rl-default-top-p', type=float, default=0, - help="Default top-p for model inference.") - group.add_argument('--rl-default-top-k', type=int, default=-1, - help="Default top-k for model inference.") - group.add_argument('--rl-offload-optimizer-during-inference', action='store_true', - help='Offload optimizer state to CPU during inference/rollout to save GPU memory') - group.add_argument('--rl-kv-cache-management-mode', type=str, default='persist', - choices=['persist', 'offload', 'recompute'], - help='KV cache management mode during RL training: ' - 'persist: leave KV cache in GPU memory (default), ' - 'offload: offload KV cache to CPU during training, ' - 'recompute: deallocate KV cache and recompute from scratch each cycle') - group.add_argument('--rl-persist-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='Persist CUDA graphs when the inference engine is suspended. ' - 'If False, CUDA graphs are deleted on suspend and re-captured on resume.') - group.add_argument('--rl-partial-rollouts', action=argparse.BooleanOptionalAction, default=False, - help='Allow inference to continue generating rollouts while training updates ' - 'the policy weights. This enables off-policy training where rollouts may ' - 'be generated with a stale version of the policy. Use ' - '--rl-generation-lag to control the degree of staleness.') - group.add_argument('--rl-inference-logprobs-is-correction', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='If set, use inference logprobs in importance sampling correction of the loss.') - group.add_argument('--rl-importance-sampling-truncation-coef', type=float, default=None, - help="If --inference-logprobs-is-correction is on and this coefficient is set, apply truncation for the IS correction at GRPO loss.") - group.add_argument('--rl-use-sequence-packing', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='Enable sequence packing') - group.add_argument('--rl-sequence-packing-max-sequences-per-bin', type=int, default=50, - help='Maximum number of sequences that can be packed into a single bin. ') - group.add_argument('--rl-sequence-packing-algo', type=str, default='fifo', - choices=['fifo', 'round-robin'], - help='Algorithm for distributing packed bins across ranks. ' - 'fifo: first-in-first-out sequential distribution, ' - 'round-robin: distribute bins cyclically across ranks for better load balancing') - group.add_argument('--rl-training-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, - default=False, - help='If set, do not toggle CUDA graphs on/off between inference and training phases.') - group.add_argument('--rl-inference-tensor-model-parallel-size', type=int, default=None, - help='Degree of tensor model parallelism for inference for RL.') group.add_argument( - '--rl-inference-pipeline-model-parallel-size', - type=int, - default=None, - help='Degree of pipeline model parallelism for inference for RL.', + '--rl-submission-granularity', + type=str, + default="B", + choices=["R", "G", "B"], + help='Granularity for submitting rollout generation work. ' + 'R submits individual rollouts independently while still yielding ' + 'complete rollout groups to training. ' + 'G submits one rollout group at a time. ' + 'B submits grpo_prompts_per_step rollout groups together.', + ) + group.add_argument( + '--rl-consumption-granularity', + type=str, + default="B", + choices=["R", "G", "B"], + help='Granularity for consuming generated rollout groups. ' + 'G consumes groups as they complete. ' + 'B consumes complete trainer batches in submission order. ' + 'R is not currently supported.', + ) + group.add_argument( + '--grpo-iterations', + type=int, + default=2, + help="Number of iterations per a GRPO implementation.", + ) + # As in DAPO, we keep upper/lower eps different. + # To have a vanilla GRPO, set them to be the same. + group.add_argument( + '--grpo-clamp-eps-lower', type=float, default=0.01, help="Lower GRPO clipping bound." + ) + group.add_argument( + '--grpo-clamp-eps-upper', + type=float, + default=0.01, + help="Upper GRPO clipping bound. In vanilla implementation, equals to the lower one.", + ) + group.add_argument( + '--grpo-kl-beta', type=float, default=0.001, help="KL term weight in the GRPO loss." + ) + group.add_argument( + '--grpo-entropy-term-weight', + type=float, + default=0.0, + help="Entropy term weight in GRPO loss.", + ) + group.add_argument( + '--grpo-filter-groups-with-same-reward', + action='store_true', + help="Filter groups with same reward.", + ) + group.add_argument( + '--langrl-env-config', + type=str, + default=None, + help="Path to YAML config file for RL environment configuration.", + ) + group.add_argument( + '--rl-default-temperature', + type=float, + default=1.0, + help="Default temperature for model inference.", + ) + group.add_argument( + '--rl-default-top-p', type=float, default=0, help="Default top-p for model inference." + ) + group.add_argument( + '--rl-default-top-k', type=int, default=-1, help="Default top-k for model inference." + ) + group.add_argument( + '--rl-offload-optimizer-during-inference', + action='store_true', + help='Offload optimizer state to CPU during inference/rollout to save GPU memory', + ) + group.add_argument( + '--rl-kv-cache-management-mode', + type=str, + default='persist', + choices=['persist', 'offload', 'recompute'], + help='KV cache management mode during RL training: ' + 'persist: leave KV cache in GPU memory (default), ' + 'offload: offload KV cache to CPU during training, ' + 'recompute: deallocate KV cache and recompute from scratch each cycle', + ) + group.add_argument( + '--rl-persist-cuda-graphs', + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help='Persist CUDA graphs when the inference engine is suspended. ' + 'If False, CUDA graphs are deleted on suspend and re-captured on resume.', + ) + group.add_argument( + '--rl-partial-rollouts', + action=argparse.BooleanOptionalAction, + default=False, + help='Allow inference to continue generating rollouts while training updates ' + 'the policy weights. This enables off-policy training where rollouts may ' + 'be generated with a stale version of the policy. Use --rl-generation-lag ' + 'to control the degree of staleness.', + ) + group.add_argument( + '--rl-inference-logprobs-is-correction', + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help='If set, use inference logprobs in importance sampling correction of the loss.', + ) + group.add_argument( + '--rl-importance-sampling-truncation-coef', + type=float, + default=None, + help="If --inference-logprobs-is-correction is on and this coefficient is set, apply truncation for the IS correction at GRPO loss.", + ) + group.add_argument( + '--rl-use-sequence-packing', + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help='Enable sequence packing', + ) + group.add_argument( + '--rl-sequence-packing-max-sequences-per-bin', + type=int, + default=50, + help='Maximum number of sequences that can be packed into a single bin. ', + ) + group.add_argument( + '--rl-sequence-packing-algo', + type=str, + default='fifo', + choices=['fifo', 'round-robin'], + help='Algorithm for distributing packed bins across ranks. ' + 'fifo: first-in-first-out sequential distribution, ' + 'round-robin: distribute bins cyclically across ranks for better load balancing', + ) + group.add_argument( + '--rl-training-cuda-graphs', + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help='If set, do not toggle CUDA graphs on/off between inference and training phases.', + ) + group.add_argument( + '--rl-inference-tensor-model-parallel-size', + type=int, + default=None, + help='Degree of tensor model parallelism for inference for RL.', + ) + group.add_argument( + '--rl-inference-pipeline-model-parallel-size', + type=int, + default=None, + help='Degree of pipeline model parallelism for inference for RL.', + ) + group.add_argument( + '--rl-inference-expert-model-parallel-size', + type=int, + default=None, + help='Degree of expert model parallelism for inference for RL.', + ) + group.add_argument( + '--rl-inference-expert-tensor-model-parallel-size', + type=int, + default=None, + help='Degree of expert tensor model parallelism for inference for RL. ' + 'For MoE models, this controls the TP size for expert layers specifically. ' + 'Defaults to training expert_tensor_parallel_size if not specified.', + ) + group.add_argument( + '--rl-inference-model-unified-memory-level', + type=int, + default=0, + choices=[0, 1], + help=( + 'Allocate the separate RL inference model parameters from a unified virtual memory (UVM) ' + 'CUDA mempool. Level 0 disables UVM (default). Level 1 enables UVM allocation so the ' + 'inference model weights can be prefetched to CPU when idle while keeping CUDA-graph-safe ' + 'device pointers.' + ), + ) + group.add_argument( + '--rl-offload-inference-model-weights-when-idle', + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help=( + 'When using a separate RL inference model, offload its weights to CPU when not doing rollout ' + 'inference, and restore to GPU right before inference. Works with two backends: ' + '1) UVM (when --rl-inference-model-unified-memory-level=1), or ' + '2) torch_memory_saver (when UVM is not enabled; requires torch_memory_saver to be installed).' + ), + ) + group.add_argument( + '--refit-method', + type=str, + default='gloo', + choices=['nccl', 'gloo', 'nvshmem'], + help=( + 'Method to refit the model weights between training and inference models during RL. ' + 'nccl: use NCCLCopyService to refit using NCCL; ' + 'gloo: use GlooCopyService over CPU; ' + 'nvshmem: use NVSHMEMCopyService to refit using the NVSHMEM.' + ), + ) + group.add_argument( + '--rl-verify-model-weights-swap', + action=argparse.BooleanOptionalAction, + default=False, + help='If set, verify that the model weights were correctly transferred by comparing forward pass outputs on' + 'the first swap of model weights.', + ) + + group.add_argument( + '--rl-skip-bos-token', + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help='Skip BOS token at the beginning of the sequences. Default is False.', + ) + group.add_argument( + '--rl-profile', + action='store_true', + default=False, + help='Enable RL profiling to collect detailed timer data (JSONL + CSV).', + ) + group.add_argument( + '--rl-profile-dir', + type=str, + default=None, + help='Directory to write RL profiling data. Defaults to {save}/profiles.', + ) + group.add_argument( + '--rl-inference-parsers', + nargs='*', + default=[], + help='List of response parsers to enable for RL inference ' + '(e.g. --rl-inference-parsers deepseek-r1-reasoning qwen3-coder-tool).', + ) + return parser + + +def _add_training_args(parser): + from megatron.training.config import ProfilingConfig, TrainingConfig + + prof_factory = ArgumentGroupFactory(ProfilingConfig) + prof_group = prof_factory.build_group(parser, "profiling") + + train_factory = ArgumentGroupFactory(TrainingConfig) + group = train_factory.build_group(parser, "training") + + group.add_argument( + '--batch-size', + type=int, + default=None, + help='Old batch size parameter, do not use. ' 'Use --micro-batch-size instead', + ) + group.add_argument( + '--recompute-activations', + action='store_true', + help='recompute activation to allow for training ' + 'with larger models, sequences, and batch sizes.', + ) + group.add_argument( + '--no-check-for-nan-in-loss-and-grad', + action='store_false', + help='Check for NaNs in loss and grad', + dest='check_for_nan_in_loss_and_grad', + ) + group.add_argument( + '--check-for-large-grads', + action='store_true', + help='Check for unexpectedly large grads', + dest='check_for_large_grads', + ) + group.add_argument( + '--result-rejected-tracker-filename', + type=str, + default=None, + help='Optional name of file tracking `result_rejected` events.', + ) + group.add_argument( + '--tp-comm-overlap-cfg', + type=str, + default=None, + help='Config file when tp_comm_overlap is enabled.', + ) + + # deprecated + group.add_argument( + '--checkpoint-activations', + action='store_true', + help='Checkpoint activation to allow for training ' + 'with larger models, sequences, and batch sizes.', + ) + group.add_argument( + '--no-masked-softmax-fusion', + action='store_false', + help='Disable fusion of query_key_value scaling, ' 'masking, and softmax.', + dest='masked_softmax_fusion', + ) + group.add_argument( + '--no-bias-gelu-fusion', + action='store_false', + help='Disable bias and gelu fusion.', + dest='bias_gelu_fusion', + ) + group.add_argument( + '--no-bias-swiglu-fusion', + action='store_false', + help='Disable bias and swiglu fusion, the fusion is ' + 'available only when using megatron-core.', + dest='bias_swiglu_fusion', + ) + group.add_argument( + '--no-bias-dropout-fusion', + action='store_false', + help='Disable bias and dropout fusion.', + dest='bias_dropout_fusion', + ) + group.add_argument( + '--no-rope-fusion', + action='store_false', + help='Disable rope fusion, the fusion is available ' 'only when using megatron-core.', + dest='apply_rope_fusion', + ) + group.add_argument( + '--rope-type', + type=str, + default=None, + choices=['rope', 'yarn'], + help='Type of rope to use. Note that MLA takes yarn by default, ' + 'and common attention takes rope by default.', + ) + group.add_argument( + '--use-flash-attn', + action='store_true', + help='use FlashAttention implementation of attention. ' 'https://arxiv.org/abs/2205.14135', + ) + group.add_argument( + '--optimizer', + type=str, + default='adam', + choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap', 'adaptive_muon'], + help='Optimizer function. ' + 'Note: dist_muon is deprecated; use --optimizer muon ' + 'with --use-distributed-optimizer instead.', + ) + group.add_argument( + '--optimizer-cpu-offload', action='store_true', help='Offload optimizer state to CPU' + ) + group.add_argument( + '--optimizer-cuda-graph', action='store_true', help='Enable CUDA graph for optimizer step' + ) + group.add_argument( + '--optimizer-offload-fraction', + type=float, + default=1.0, + help='Fraction used by --optimizer-cpu-offload. This is distinct from ' + '--optimizer-state-offload-fraction, and the two offload modes are mutually exclusive.', + ) + group.add_argument( + '--use-torch-optimizer-for-cpu-offload', + action='store_true', + help="Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu offload mode.", + ) + group.add_argument( + '--overlap-cpu-optimizer-d2h-h2d', + action='store_true', + default=False, + help='Overlap CPU optimizer step, gradients D2H and updated parameters H2D.', + ) + group.add_argument( + '--dump-param-to-param-group-map', + type=str, + default=None, + help="Path to a file containing parameter-to-parameter-group mapping. " + "Provide a JSON file that specifies which parameters belong to which " + "parameter group for global coordination.", + ) + group.add_argument( + '--no-pin-cpu-grads', + action='store_false', + dest='pin_cpu_grads', + help='Disable pinning of CPU memory for gradients.', + ) + group.add_argument( + '--no-pin-cpu-params', + action='store_false', + dest='pin_cpu_params', + help='Disable pinning of CPU memory for parameters.', + ) + group.add_argument( + '--chunked-optimizer-state-offload', + action='store_true', + help='Keep selected optimizer tensor states and master weights in CPU memory between ' + 'updates. GPU tensor-state updates run in bounded chunks, while selected master ' + 'weights use one full restore window.', + ) + group.add_argument( + '--optimizer-state-offload-chunk-size-mb', + type=int, + default=0, + help='Target size of each optimizer tensor-state staging buffer in MiB. Two buffers ' + 'can be live for overlap (roughly 2x this value per active offload manager). Zero ' + 'temporarily restores all selected tensor state for one full GPU update. A single ' + 'atomic parameter state may exceed this soft target and is reported at runtime. This ' + 'value does not bound the selected master-weight window, which is always restored in ' + 'full.', + ) + group.add_argument( + '--optimizer-state-offload-fraction', + type=float, + default=1.0, + help='Fraction of optimizer parameter bundles to offload. Selected bundles always ' + 'include both tensor state and any separate master weight. This is distinct from ' + '--optimizer-offload-fraction, which belongs to --optimizer-cpu-offload; the modes are ' + 'mutually exclusive. Zero disables optimizer state offload even when its flag is set.', + ) + group.add_argument( + '--offload-optimizer-states', + action='store_true', + dest='offload_optimizer_states', + help='Deprecated spelling for --chunked-optimizer-state-offload. It automatically ' + 'enables the replacement; the legacy-equivalent settings are chunk size 0 and offload ' + 'fraction 1.0 unless new tuning arguments are supplied. ' + 'With a nonzero fraction, optimizer-state checkpoint I/O requires --ckpt-format ' + 'torch_dist, optimizer-state saves must be synchronous, and optimizer or full-iteration ' + 'CUDA graphs are unsupported.', + ) + group.add_argument( + '--dataloader-type', + type=str, + default=None, + choices=['single', 'cyclic', 'external'], + help='Single pass vs multiple pass data loader', + ) + group.add_argument( + '--no-persist-layer-norm', + action='store_true', + help='Disable using persistent fused layer norm kernel. ' + 'This kernel supports only a set of hidden sizes. Please ' + 'check persist_ln_hidden_sizes if your hidden ' + 'size is supported.', + ) + group.add_argument( + '--no-gradient-accumulation-fusion', + action='store_false', + help='Disable fusing gradient accumulation to weight ' + 'gradient computation of linear layers', + dest='gradient_accumulation_fusion', + ) + group.add_argument( + '--use-mcore-models', + action='store_true', + dest='deprecated_use_mcore_models', + help='DEPRECATED. Use the implementation from megatron core.' + 'Now ignored and mcore models are the default.', + ) + + return parser + + +def _add_rerun_machine_args(parser): + from megatron.training.config import RerunStateMachineConfig + + rerun_factory = ArgumentGroupFactory(RerunStateMachineConfig, exclude=["check_for_nan_in_loss"]) + group = rerun_factory.build_group(parser, "rerun engine") + + return parser + + +def _add_initialization_args(parser): + from megatron.training.config import RNGConfig + + rng_factory = ArgumentGroupFactory(RNGConfig) + group = rng_factory.build_group(parser, "RNG and initialization") + + group.add_argument( + '--init-method-xavier-uniform', + action='store_true', + help='Enable Xavier uniform parameter initialization', + ) + + return parser + + +def _add_learning_rate_args(parser): + from megatron.training.config import SchedulerConfig + + sched_factory = ArgumentGroupFactory(SchedulerConfig, exclude=["no_weight_decay_cond_type"]) + group = sched_factory.build_group(parser, title="learning rate and weight decay") + + group.add_argument( + '--lr', + type=float, + default=None, + help='Initial learning rate. Depending on decay style ' + 'and initial warmup, the learning rate at each ' + 'iteration would be different.', + ) + group.add_argument( + '--warmup', + type=int, + default=None, + help='Old lr warmup argument, do not use. Use one of the' '--lr-warmup-* arguments above', + ) + group.add_argument( + '--min-lr', + type=float, + default=0.0, + help='Minimum value for learning rate. The scheduler' 'clip values below this threshold.', + ) + group.add_argument( + '--decoupled-lr', + type=float, + default=None, + help='Separate learning rate for the input and output layer', + ) + group.add_argument( + '--decoupled-min-lr', + type=float, + default=None, + help='Minimum value for learning rate for the input and output layer. The scheduler' + 'clip values below this threshold', + ) + group.add_argument( + '--freeze-all-layers', action='store_true', help='Freeze all layers of the model.' + ) + + return parser + + +def _add_checkpointing_args(parser): + from megatron.training.config import CheckpointConfig + + ckpt_factory = ArgumentGroupFactory( + CheckpointConfig, + exclude=[ + "most_recent_k", + "save_tokenizer_assets", + "save_optim", + "save_rng", + "load_optim", + "load_rng", + ], + ) + group = ckpt_factory.build_group(parser, "checkpointing") + + group.add_argument( + '--no-save-optim', action='store_true', default=None, help='Do not save current optimizer.' + ) + group.add_argument( + '--no-save-rng', action='store_true', default=None, help='Do not save current rng state.' + ) + group.add_argument( + '--no-load-optim', + action='store_true', + default=None, + help='Do not load optimizer when loading checkpoint.', + ) + group.add_argument( + '--no-load-rng', + action='store_true', + default=None, + help='Do not load rng state when loading checkpoint.', + ) + group.add_argument( + '--use-dist-ckpt', + action='store_true', + dest='use_dist_ckpt_deprecated', + help='Deprecated: see --ckpt-format.', + ) + group.add_argument( + '--dist-ckpt-format', + dest='dist_ckpt_format_deprecated', + help='Deprecated: see --ckpt-format.', + ) + group.add_argument( + '--dist-ckpt-workers', + type=int, + default=1, + help='Number of workers for distributed checkpointing. ' + 'Only used for async save. ' + 'If set to 1, the checkpointing is performed in a single process.', + ) + group.add_argument( + '--ckpt-fully-parallel-save', + action='store_true', + dest='ckpt_fully_parallel_save_deprecated', + help='Deprecated: see --no-ckpt-fully-parallel-save.', + ) + return parser + + +def _add_mixed_precision_args(parser): + group = parser.add_argument_group(title='mixed precision') + + group.add_argument( + '--grad-reduce-in-bf16', action='store_true', help='Reduce gradients in bfloat16.' + ) + group.add_argument( + '--loss-scale', + type=float, + default=None, + help='Static loss scaling, positive power of 2 ' + 'values can improve fp16 convergence. If None, dynamic' + 'loss scaling is used.', + ) + group.add_argument( + '--initial-loss-scale', + type=float, + default=2**32, + help='Initial loss-scale for dynamic loss scaling.', + ) + group.add_argument( + '--min-loss-scale', + type=float, + default=1.0, + help='Minimum loss scale for dynamic loss scaling.', + ) + group.add_argument( + '--loss-scale-window', + type=float, + default=1000, + help='Window over which to raise/lower dynamic scale.', + ) + group.add_argument( + '--hysteresis', type=int, default=2, help='hysteresis for dynamic loss scaling' + ) + group.add_argument( + '--attention-softmax-in-fp32', + action='store_true', + help='Run attention masking and softmax in fp32.', + ) + group.add_argument( + '--accumulate-allreduce-grads-in-fp32', + action='store_true', + help='Gradient accumulation and all-reduce in fp32.', + ) + group.add_argument( + '--fp16-lm-cross-entropy', + action='store_true', + help='Move the cross entropy unreduced loss calculation' 'for lm head to fp16.', + ) + group.add_argument( + '--reuse-grad-buf-for-mxfp8-param-ag', + action='store_true', + help='If True, reuse the grad buffer for MXFP8 parameter all-gather.', + ) + group.add_argument( + '--mamba-training-ssm-states-dtype', + type=str, + choices=['fp32', 'bf16'], + default=None, + help='Dtype of the materialized inter-chunk SSM states in Mamba training', + ) + + return parser + + +def _add_distributed_args(parser): + from megatron.training.config import DistributedInitConfig + + dist_init_factory = ArgumentGroupFactory(DistributedInitConfig) + group = dist_init_factory.build_group(parser, "distributed init") + + group.add_argument( + '--decoder-first-pipeline-num-layers', + type=int, + default=None, + help=( + 'The number of transformer layers on the first pipeline stage of the decoder. ' + 'Default None is even split of transformer layers across all pipeline stages' + ), + ) + group.add_argument( + '--decoder-last-pipeline-num-layers', + type=int, + default=None, + help=( + 'The number of transformer layers on the last pipeline stage of the decoder. ' + 'Default None is even split of transformer layers across all pipeline stages' + ), + ) + group.add_argument( + '--pipeline-model-parallel-layout', + type=str, + default=None, + help=( + 'A string that describes a custom pipeline model parallel layout. ' + 'e.g., "E|(t|)*3,m|m||L". E, L, t, m denotes embedding, loss, transformer ' + 'decoder layer, and mtp layer, respectively. Stages are split by "|". ' + 'Replicated stages or layers can be described with multiplication. ' + 'Commas can be used cosmetically. ' + 'Default None is not using this argument to set the layout.' + ), + ) + group.add_argument( + '--model-parallel-size', + type=int, + default=None, + help='Old model parallel argument, do not use. Use ' + '--tensor-model-parallel-size instead.', + ) + group.add_argument( + '--num-layers-per-virtual-pipeline-stage', + type=int, + default=None, + help='Number of layers per virtual pipeline stage', + ) + group.add_argument( + '--num-virtual-stages-per-pipeline-rank', + type=int, + default=None, + help='Number of virtual pipeline stages per pipeline parallelism rank', + ) + group.add_argument( + '--no-overlap-p2p-communication', + action='store_false', + help='overlap pipeline parallel communication with forward and backward chunks in 1F1B', + dest='overlap_p2p_comm', + ) + group.add_argument( + '--overlap-grad-reduce', + action='store_true', + default=False, + help='If set, overlap DDP grad reduce.', + ) + group.add_argument( + '--ddp-num-buckets', + type=int, + default=None, + help='Number of buckets for data-parallel communication', + ) + group.add_argument( + '--ddp-bucket-size', + type=int, + default=None, + help='Bucket size for data-parallel communication', + ) + group.add_argument( + '--ddp-pad-buckets-for-high-nccl-busbw', + action='store_true', + default=False, + help='If set, make sure the bucket size is divisible by a large power ' + 'of 2 (2^16) to ensure NCCL collectives have high bus bandwidth at large DP counts, ' + 'since NCCL message size (which for ring algorithms is bucket_size / dp_size) ' + 'apparently needs to be divisible by a power of 2 for high busbw.', + ) + group.add_argument( + '--ddp-reduce-scatter-with-fp32-accumulation', + action='store_true', + default=False, + help='If set, use a reduce-scatter implementation which sends lower-precision ' + 'values over the wire (using an all-to-all to keep total communication overhead in line ' + 'with the standard ring implementation) but performs accumulation locally in FP32.', + ) + group.add_argument( + '--ddp-param-name-patterns-for-fp32-local-accumulation', + nargs='+', + default=[], + help='List of param_name patterns (in Python\'s fnmatch format) to match against ' + 'to do local gradient accumulation in FP32. The special pattern \'all\' matches ' + 'every parameter.', + ) + group.add_argument( + '--ddp-average-in-collective', + action='store_true', + default=False, + help='If set, average directly in data-parallel communication collective.', + ) + group.add_argument( + '--disable-grad-buffers-cpu-backup', + action='store_true', + default=False, + help='If set, allocate DDP gradient buffers in a torch_memory_saver region ' + 'without CPU backup.', + ) + group.add_argument( + '--disable-param-buffers-cpu-backup', + action='store_true', + default=False, + help='If set, allocate DDP parameter buffers in a torch_memory_saver region ' + 'without CPU backup. Only applies when using the distributed optimizer.', + ) + group.add_argument( + '--overlap-param-gather', + action='store_true', + default=False, + help='If set, overlap param all-gather in distributed optimizer.', + ) + group.add_argument( + '--overlap-param-gather-with-optimizer-step', + action='store_true', + default=False, + help='If set, overlap param all-gather of first bucket with optimizer step.', + ) + group.add_argument( + '--no-align-param-gather', + action='store_false', + help='If not set, all PP stages will launch param all-gathers simultaneously. ' + 'Otherwise, each PP stage will independently launch as needed.', + dest='align_param_gather', + ) + group.add_argument( + '--use-distributed-optimizer', action='store_true', help='Use distributed optimizer.' + ) + group.add_argument( + '--use-nccl-ub', + action='store_true', + dest='nccl_ub', + help='Use the userbuffer registration for DP/FSDP communication buffers.' + 'This option will reduce GPU SM usage for the DP/FSDP communication,' + 'which is improving the performance of the overlapped computation.', + ) + group.add_argument( + '--disable-symmetric-registration', + action='store_true', + dest='disable_symmetric_registration', + default=False, + help='Disable symmetric (window) registration for NCCL userbuffer registration.' + 'This option will force to use conventional (local) userbuffer registration when use-nccl-ub is set.', + ) + group.add_argument( + '--fsdp-manual-registration', + action='store_true', + dest='fsdp_manual_registration', + default=False, + help='Manually register the FSDP communication buffers to NCCL user buffer.' + 'This option is only effective when use-megatron-fsdp and use-nccl-ub is set.', + ) + group.add_argument( + '--create-all-gather-group', + action='store_true', + help='Create a separate process group for all-gather operations ' + 'to overlap reduce-scatter and all-gather operations.', + ) + group.add_argument( + '--data-parallel-sharding-strategy', + type=str, + default='optim_grads_params', + choices=['no_shard', 'optim', 'optim_grads', 'optim_grads_params'], + help='Sharding strategy of data parallelism.', + ) + group.add_argument( + '--outer-dp-sharding-strategy', + type=str, + default='no_shard', + choices=['no_shard', 'optim'], + help='Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode. ' + 'Valid values are "no_shard" (DP Replication) and "optim" (Optimizer State Hybrid Sharding). ' + 'The "optim" option is only supported when --data-parallel-sharding-strategy is "optim_grads_params". ' + 'This option is only effective when Hybrid FSDP is enabled (i.e., when dp_outer_dim is not None). ' + 'Default: "no_shard".', + ) + group.add_argument( + '--no-gradient-reduce-div-fusion', + action='store_false', + dest='gradient_reduce_div_fusion', + help='If not set, fuse the division in gradient reduce.', ) group.add_argument( - '--rl-inference-expert-model-parallel-size', - type=int, - default=None, - help='Degree of expert model parallelism for inference for RL.', + '--fsdp-double-buffer', + action='store_true', + help="Enable double buffering for temporary memory needed for Megatron FSDP communications. " + "Double-buffering the communication memory improves memory management efficiency by " + "reusing previously allocated buffers, rather than creating new buffers for each FSDP communication. " + "This is required for user buffer registration and is enabled by default when using NCCL user buffers.", ) group.add_argument( - '--rl-inference-expert-tensor-model-parallel-size', + '--suggested-communication-unit-size', type=int, default=None, - help='Degree of expert tensor model parallelism for inference for RL. ' - 'For MoE models, this controls the TP size for expert layers specifically. ' - 'Defaults to training expert_tensor_parallel_size if not specified.', + help='Specifies the number of elements to communicate at once during FSDP (Fully Sharded Data Parallel) operations. ' + 'This flag also affects FSDP all-gather prefetch behavior. Setting a larger value increases the communication buffer size, ' + 'while a smaller value disables prefetching and may degrade performance. Adjust this value based on your system\'s memory ' + 'and performance requirements.', ) group.add_argument( - '--rl-inference-model-unified-memory-level', + '--keep-fp8-transpose-cache', + action='store_true', + help='If set, keep the fp8 transpose cache when using Megatron FSDP.', + ) + group.add_argument( + '--enable-full-sharding-in-hsdp', + action='store_true', + help='If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel (HSDP) mode.', + ) + group.add_argument( + '--num-distributed-optimizer-instances', type=int, - default=0, - choices=[0, 1], - help=( - 'Allocate the separate RL inference model parameters from a unified virtual memory (UVM) ' - 'CUDA mempool. Level 0 disables UVM (default). Level 1 enables UVM allocation so the ' - 'inference model weights can be prefetched to CPU when idle while keeping CUDA-graph-safe ' - 'device pointers.' - ), + default=1, + help='Number of Distributed Optimizer copies across Data Parallel domain.', ) group.add_argument( - '--rl-offload-inference-model-weights-when-idle', - action=argparse.BooleanOptionalAction, - required=False, - default=False, - help=( - 'When using a separate RL inference model, offload its weights to CPU when not doing rollout ' - 'inference, and restore to GPU right before inference. Works with two backends: ' - '1) UVM (when --rl-inference-model-unified-memory-level=1), or ' - '2) torch_memory_saver (when UVM is not enabled; requires torch_memory_saver to be installed).' - ), + '--torch-fsdp2-no-reshard-after-forward', + action='store_false', + dest='torch_fsdp2_reshard_after_forward', + help='Whether to reshard weights after forward pass when using PyTorch FSDP2. ' + 'Set to enable FSDP ZeRO-2.', ) - group.add_argument('--refit-method', type=str, default='gloo', - choices=['nccl', 'nccl_m2n', 'gloo', 'nvshmem', 'nixl'], - help=('Method to refit model weights. ' - 'nccl: use NCCLCopyService; ' - 'nccl_m2n: use the official NCCL M2N API from a non-RL ' - 'launcher such as the ReFIT benchmark; ' - 'gloo: use GlooCopyService over CPU; ' - 'nvshmem: use NVSHMEMCopyService; ' - 'nixl: use NixlCopyService.')) - group.add_argument('--rl-verify-model-weights-swap', action=argparse.BooleanOptionalAction, default=False, - help='If set, verify that the model weights were correctly transferred by comparing forward pass outputs on' - 'the first swap of model weights.') - - group.add_argument('--rl-skip-bos-token', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='Skip BOS token at the beginning of the sequences. Default is False.') - group.add_argument('--rl-profile', action='store_true', default=False, - help='Enable RL profiling to collect detailed timer data (JSONL + CSV).') - group.add_argument('--rl-profile-dir', type=str, default=None, - help='Directory to write RL profiling data. Defaults to {save}/profiles.') - group.add_argument('--rl-inference-parsers', nargs='*', default=[], - help='List of response parsers to enable for RL inference ' - '(e.g. --rl-inference-parsers deepseek-r1-reasoning qwen3-coder-tool).') - return parser - -def _add_training_args(parser): - from megatron.training.config import ProfilingConfig, TrainingConfig - - prof_factory = ArgumentGroupFactory(ProfilingConfig) - prof_group = prof_factory.build_group(parser, "profiling") - - train_factory = ArgumentGroupFactory(TrainingConfig) - group = train_factory.build_group(parser, "training") - - group.add_argument('--batch-size', type=int, default=None, - help='Old batch size parameter, do not use. ' - 'Use --micro-batch-size instead') - group.add_argument('--recompute-activations', action='store_true', - help='recompute activation to allow for training ' - 'with larger models, sequences, and batch sizes.') - group.add_argument('--no-check-for-nan-in-loss-and-grad', action='store_false', - help='Check for NaNs in loss and grad', - dest='check_for_nan_in_loss_and_grad') - group.add_argument('--check-for-large-grads', action='store_true', - help='Check for unexpectedly large grads', - dest='check_for_large_grads') - group.add_argument('--result-rejected-tracker-filename', type=str, default=None, - help='Optional name of file tracking `result_rejected` events.') - group.add_argument('--tp-comm-overlap-cfg', type=str, default=None, - help='Config file when tp_comm_overlap is enabled.') - - # deprecated - group.add_argument('--checkpoint-activations', action='store_true', - help='Checkpoint activation to allow for training ' - 'with larger models, sequences, and batch sizes.') - group.add_argument('--no-masked-softmax-fusion', - action='store_false', - help='Disable fusion of query_key_value scaling, ' - 'masking, and softmax.', - dest='masked_softmax_fusion') - group.add_argument('--no-bias-gelu-fusion', action='store_false', - help='Disable bias and gelu fusion.', - dest='bias_gelu_fusion') - group.add_argument('--no-bias-swiglu-fusion', action='store_false', - help='Disable bias and swiglu fusion, the fusion is ' - 'available only when using megatron-core.', - dest='bias_swiglu_fusion') - group.add_argument('--no-bias-dropout-fusion', action='store_false', - help='Disable bias and dropout fusion.', - dest='bias_dropout_fusion') - group.add_argument('--no-rope-fusion', action='store_false', - help='Disable rope fusion, the fusion is available ' - 'only when using megatron-core.', - dest='apply_rope_fusion') - group.add_argument('--rope-type', type=str, default=None, - choices=['rope', 'yarn'], - help='Type of rope to use. Note that MLA takes yarn by default, ' - 'and common attention takes rope by default.') - group.add_argument('--use-flash-attn', action='store_true', - help='use FlashAttention implementation of attention. ' - 'https://arxiv.org/abs/2205.14135') - group.add_argument('--optimizer', type=str, default='adam', - choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap', 'adaptive_muon'], - help='Optimizer function. ' - 'Note: dist_muon is deprecated; use --optimizer muon ' - 'with --use-distributed-optimizer instead.') - group.add_argument('--optimizer-cpu-offload', action='store_true', - help='Offload optimizer state to CPU') - group.add_argument('--optimizer-cuda-graph', action='store_true', - help='Enable CUDA graph for optimizer step') - group.add_argument('--optimizer-offload-fraction', type=float, default=1.0, - help='Ratio of optimizer state to offload to CPU') - group.add_argument('--use-torch-optimizer-for-cpu-offload', action='store_true', - help="Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu offload mode.") - group.add_argument('--overlap-cpu-optimizer-d2h-h2d', action='store_true', default=False, - help='Overlap CPU optimizer step, gradients D2H and updated parameters H2D.') - group.add_argument('--dump-param-to-param-group-map', type=str, default=None, - help="Path to a file containing parameter-to-parameter-group mapping. " - "Provide a JSON file that specifies which parameters belong to which " - "parameter group for global coordination.") - group.add_argument('--no-pin-cpu-grads', action='store_false', dest='pin_cpu_grads', - help='Disable pinning of CPU memory for gradients.') - group.add_argument('--no-pin-cpu-params', action='store_false', dest='pin_cpu_params', - help='Disable pinning of CPU memory for parameters.') - group.add_argument('--dataloader-type', type=str, default=None, - choices=['single', 'cyclic', 'external'], - help='Single pass vs multiple pass data loader') - group.add_argument('--no-persist-layer-norm', action='store_true', - help='Disable using persistent fused layer norm kernel. ' - 'This kernel supports only a set of hidden sizes. Please ' - 'check persist_ln_hidden_sizes if your hidden ' - 'size is supported.') - group.add_argument('--no-gradient-accumulation-fusion', - action='store_false', - help='Disable fusing gradient accumulation to weight ' - 'gradient computation of linear layers', - dest='gradient_accumulation_fusion') - group.add_argument('--use-mcore-models', action='store_true', - dest='deprecated_use_mcore_models', - help='DEPRECATED. Use the implementation from megatron core.' - 'Now ignored and mcore models are the default.') - - return parser - - -def _add_rerun_machine_args(parser): - from megatron.training.config import RerunStateMachineConfig - - rerun_factory = ArgumentGroupFactory(RerunStateMachineConfig, exclude=["check_for_nan_in_loss"]) - group = rerun_factory.build_group(parser, "rerun engine") - - return parser - - -def _add_initialization_args(parser): - from megatron.training.config import RNGConfig - - rng_factory = ArgumentGroupFactory(RNGConfig) - group = rng_factory.build_group(parser, "RNG and initialization") - - group.add_argument('--init-method-xavier-uniform', action='store_true', - help='Enable Xavier uniform parameter initialization') - - return parser - - -def _add_learning_rate_args(parser): - from megatron.training.config import SchedulerConfig - - sched_factory = ArgumentGroupFactory(SchedulerConfig, exclude=["no_weight_decay_cond_type"]) - group = sched_factory.build_group(parser, title="learning rate and weight decay") - - group.add_argument('--lr', type=float, default=None, - help='Initial learning rate. Depending on decay style ' - 'and initial warmup, the learning rate at each ' - 'iteration would be different.') - group.add_argument('--warmup', type=int, default=None, - help='Old lr warmup argument, do not use. Use one of the' - '--lr-warmup-* arguments above') - group.add_argument('--min-lr', type=float, default=0.0, - help='Minimum value for learning rate. The scheduler' - 'clip values below this threshold.') - group.add_argument('--decoupled-lr', type=float, default=None, - help='Separate learning rate for the input and output layer') - group.add_argument('--decoupled-min-lr', type=float, default=None, - help='Minimum value for learning rate for the input and output layer. The scheduler' - 'clip values below this threshold') - group.add_argument('--freeze-all-layers', action='store_true', - help='Freeze all layers of the model.') - - return parser - - -def _add_checkpointing_args(parser): - from megatron.training.config import CheckpointConfig - - ckpt_factory = ArgumentGroupFactory(CheckpointConfig, exclude=["most_recent_k", "save_optim", "save_rng", "load_optim", "load_rng"]) - group = ckpt_factory.build_group(parser, "checkpointing") - - group.add_argument('--no-save-optim', action='store_true', default=None, - help='Do not save current optimizer.') - group.add_argument('--no-save-rng', action='store_true', default=None, - help='Do not save current rng state.') - group.add_argument('--no-load-optim', action='store_true', default=None, - help='Do not load optimizer when loading checkpoint.') - group.add_argument('--no-load-rng', action='store_true', default=None, - help='Do not load rng state when loading checkpoint.') - group.add_argument('--override-ckpt-iteration', type=int, default=None, - help='Override the iteration stored in the loaded checkpoint. ' - 'Also resets consumed_train_samples accordingly so the ' - 'data loader replays samples from that iteration onward.') - group.add_argument('--use-dist-ckpt', action='store_true', - dest='use_dist_ckpt_deprecated', - help='Deprecated: see --ckpt-format.') - group.add_argument('--dist-ckpt-format', - dest='dist_ckpt_format_deprecated', - help='Deprecated: see --ckpt-format.') - group.add_argument('--dist-ckpt-workers', type=int, default=1, - help='Number of workers for distributed checkpointing. ' - 'Only used for async save. ' - 'If set to 1, the checkpointing is performed in a single process.') - group.add_argument('--ckpt-fully-parallel-save', action='store_true', - dest='ckpt_fully_parallel_save_deprecated', - help='Deprecated: see --no-ckpt-fully-parallel-save.') - group.add_argument('--ckpt-drop-redundant-extra-state', action='store_true', - default=False, - help='Drop TE `_extra_state` artifacts that carry no ' - 'irreplaceable state (no FP8, or block/current FP8 ' - 'scaling) from the distributed checkpoint, keeping them ' - 'local instead of writing them. Only delayed-scaling ' - '`_extra_state` (amax history + scale) is ever needed and ' - 'it is always persisted regardless of this flag. By ' - 'default (flag off) all `_extra_state` are persisted, ' - 'preserving the previous behavior. Older checkpoints load ' - 'unchanged either way.') - return parser - - -def _add_mixed_precision_args(parser): - group = parser.add_argument_group(title='mixed precision') - - group.add_argument('--grad-reduce-in-bf16', action='store_true', - help='Reduce gradients in bfloat16.') - group.add_argument('--loss-scale', type=float, default=None, - help='Static loss scaling, positive power of 2 ' - 'values can improve fp16 convergence. If None, dynamic' - 'loss scaling is used.') - group.add_argument('--initial-loss-scale', type=float, default=2**32, - help='Initial loss-scale for dynamic loss scaling.') - group.add_argument('--min-loss-scale', type=float, default=1.0, - help='Minimum loss scale for dynamic loss scaling.') - group.add_argument('--loss-scale-window', type=float, default=1000, - help='Window over which to raise/lower dynamic scale.') - group.add_argument('--hysteresis', type=int, default=2, - help='hysteresis for dynamic loss scaling') - group.add_argument('--attention-softmax-in-fp32', action='store_true', - help='Run attention masking and softmax in fp32.') - group.add_argument('--accumulate-allreduce-grads-in-fp32', - action='store_true', - help='Gradient accumulation and all-reduce in fp32.') - group.add_argument('--fp16-lm-cross-entropy', action='store_true', - help='Move the cross entropy unreduced loss calculation' - 'for lm head to fp16.') - group.add_argument('--output-logit-dtype', type=str, choices=['bf16', 'fp32'], default=None, - dest='logit_dtype', - help='Output dtype for the language-model output-layer GEMM. When the ' - 'requested dtype differs from the input dtype, Transformer Engine ' - 'general_gemm is used. By default, logits use the output-layer input dtype.') - group.add_argument('--reuse-grad-buf-for-mxfp8-param-ag', action='store_true', - help='If True, reuse the grad buffer for MXFP8 parameter all-gather.') - group.add_argument('--mamba-training-ssm-states-dtype', type=str, - choices=['fp32', 'bf16'], default=None, - help='Dtype of the materialized inter-chunk SSM states in Mamba training') - - return parser - - -def _add_distributed_args(parser): - from megatron.training.config import DistributedInitConfig - - dist_init_factory = ArgumentGroupFactory(DistributedInitConfig) - group = dist_init_factory.build_group(parser, "distributed init") - - group.add_argument('--decoder-first-pipeline-num-layers', - type=int, default=None, - help=('The number of transformer layers on the first pipeline stage of the decoder. ' - 'Default None is even split of transformer layers across all pipeline stages')) - group.add_argument('--decoder-last-pipeline-num-layers', - type=int, default=None, - help=('The number of transformer layers on the last pipeline stage of the decoder. ' - 'Default None is even split of transformer layers across all pipeline stages')) - group.add_argument('--pipeline-model-parallel-layout', - type=str, default=None, - help=('A string that describes a custom pipeline model parallel layout. ' - 'e.g., "E|(t|)*3,m|m||L". E, L, t, m denotes embedding, loss, transformer ' - 'decoder layer, and mtp layer, respectively. Stages are split by "|". ' - 'Replicated stages or layers can be described with multiplication. ' - 'Commas can be used cosmetically. ' - 'Default None is not using this argument to set the layout.')) - group.add_argument('--model-parallel-size', type=int, default=None, - help='Old model parallel argument, do not use. Use ' - '--tensor-model-parallel-size instead.') - group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None, - help='Number of layers per virtual pipeline stage') - group.add_argument('--num-virtual-stages-per-pipeline-rank', type=int, default=None, - help='Number of virtual pipeline stages per pipeline parallelism rank') - group.add_argument('--no-overlap-p2p-communication', action='store_false', - help='overlap pipeline parallel communication with forward and backward chunks in 1F1B', - dest='overlap_p2p_comm') - group.add_argument('--overlap-grad-reduce', action='store_true', - default=False, help='If set, overlap DDP grad reduce.') - group.add_argument('--ddp-num-buckets', type=int, default=None, - help='Number of buckets for data-parallel communication') - group.add_argument('--ddp-bucket-size', type=int, default=None, - help='Bucket size for data-parallel communication') - group.add_argument('--ddp-pad-buckets-for-high-nccl-busbw', action='store_true', - default=False, help='If set, make sure the bucket size is divisible by a large power ' - 'of 2 (2^16) to ensure NCCL collectives have high bus bandwidth at large DP counts, ' - 'since NCCL message size (which for ring algorithms is bucket_size / dp_size) ' - 'apparently needs to be divisible by a power of 2 for high busbw.') - group.add_argument('--ddp-reduce-scatter-with-fp32-accumulation', action='store_true', - default=False, help='If set, use a reduce-scatter implementation which sends lower-precision ' - 'values over the wire (using an all-to-all to keep total communication overhead in line ' - 'with the standard ring implementation) but performs accumulation locally in FP32.') - group.add_argument('--gtp-remat-reduce-scatter-with-fp32-accumulation', action='store_true', - default=False, help='Same trade as --ddp-reduce-scatter-with-fp32-accumulation, but for ' - 'the wgrad reduce-scatter GTP weight-remat performs over the gtp_remat axis: send ' - 'low-precision values over the wire via an all-to-all and accumulate locally in FP32. ' - 'Independent of the DDP flag (different collective, different process group). Costs one ' - 'extra unsharded-wgrad-sized scratch buffer per in-flight reduce-scatter.') - group.add_argument('--ddp-param-name-patterns-for-fp32-local-accumulation', nargs='+', default=[], - help='List of param_name patterns (in Python\'s fnmatch format) to match against ' - 'to do local gradient accumulation in FP32. The special pattern \'all\' matches ' - 'every parameter.') - group.add_argument('--ddp-average-in-collective', action='store_true', - default=False, help='If set, average directly in data-parallel communication collective.') - group.add_argument('--overlap-param-gather', action='store_true', - default=False, help='If set, overlap param all-gather in distributed optimizer.') - group.add_argument('--overlap-param-gather-with-optimizer-step', action='store_true', - default=False, help='If set, overlap param all-gather of first bucket with optimizer step.') - group.add_argument('--no-align-param-gather', action='store_false', - help='If not set, all PP stages will launch param all-gathers simultaneously. ' - 'Otherwise, each PP stage will independently launch as needed.', - dest='align_param_gather') - group.add_argument('--use-distributed-optimizer', action='store_true', - help='Use distributed optimizer.') - group.add_argument('--megatron-fsdp-version', type=int, default=1, choices=[1, 2], - help='Megatron-FSDP implementation version. Defaults to 1.') - group.add_argument('--no-use-layer-wise-param-layout', - action='store_false', - dest='use_layer_wise_param_layout', - help='Opt out of the precomputed LayerWise param layout. When set, ' - 'falls back to the legacy LayerWise ping-pong path: all params ' - '(including non-Muon embeddings, biases, layernorm) live in a single ' - 'LayerWise buffer and the optimizer uses the allgather_params() codepath. ' - 'The default (precomputed layout) routes non-Muon params through a ' - 'separate DistributedOptimizer with byte-level sharding, which is faster ' - 'and uses less padding but produces different bf16 reduction ordering ' - 'and so will not match legacy-path loss curves bit-for-bit.') - group.add_argument('--use-nccl-ub', action='store_true', dest='nccl_ub', - help='Use the userbuffer registration for DP/FSDP communication buffers.' - 'This option will reduce GPU SM usage for the DP/FSDP communication,' - 'which is improving the performance of the overlapped computation.') - group.add_argument('--disable-symmetric-registration', action='store_true', dest='disable_symmetric_registration', - default=False, help='Disable symmetric (window) registration for NCCL userbuffer registration.' - 'This option will force to use conventional (local) userbuffer registration when use-nccl-ub is set. ' - 'Cannot be combined with --gtp-remat-nccl-ub/--gtp-expert-remat-nccl-ub, which require symmetric windows.') - group.add_argument('--gtp-remat-nccl-ub', action='store_true', dest='gtp_remat_nccl_ub', - default=False, help='Register the wgrad reduce-scatter send buffers with NCCL symmetric ' - 'memory on the GTP group, independent of --use-nccl-ub (which covers the DP group).') - group.add_argument('--gtp-expert-remat-nccl-ub', action='store_true', dest='gtp_expert_remat_nccl_ub', - default=False, help='Like --gtp-remat-nccl-ub but for routed-expert (EGTP) groups.') - group.add_argument('--fsdp-manual-registration', action='store_true', dest='fsdp_manual_registration', - default=False, help='Manually register the FSDP communication buffers to NCCL user buffer.' - 'This option is only effective when use-megatron-fsdp and use-nccl-ub is set.') - group.add_argument('--create-all-gather-group', action='store_true', - help='Create a separate process group for all-gather operations ' - 'to overlap reduce-scatter and all-gather operations.') - group.add_argument('--data-parallel-sharding-strategy', type=str, default='optim_grads_params', - choices=['no_shard', 'optim', 'optim_grads', 'optim_grads_params'], - help='Sharding strategy of data parallelism.') - group.add_argument('--outer-dp-sharding-strategy', type=str, default='no_shard', - choices=['no_shard', 'optim'], - help='Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode. ' - 'Valid values are "no_shard" (DP Replication) and "optim" (Optimizer State Hybrid Sharding). ' - 'The "optim" option is only supported when --data-parallel-sharding-strategy is "optim_grads_params". ' - 'This option is only effective when Hybrid FSDP is enabled (i.e., when dp_outer_dim is not None). ' - 'Default: "no_shard".') - group.add_argument('--no-gradient-reduce-div-fusion', action='store_false', dest='gradient_reduce_div_fusion', - help='If not set, fuse the division in gradient reduce.') - group.add_argument('--fsdp-double-buffer', action='store_true', - help="Enable double buffering for temporary memory needed for Megatron FSDP communications. " - "Double-buffering the communication memory improves memory management efficiency by " - "reusing previously allocated buffers, rather than creating new buffers for each FSDP communication. " - "This is required for user buffer registration and is enabled by default when using NCCL user buffers.") - group.add_argument('--suggested-communication-unit-size', type=int, default=None, - help='Specifies the number of elements to communicate at once during FSDP (Fully Sharded Data Parallel) operations. ' - 'This flag also affects FSDP all-gather prefetch behavior. Setting a larger value increases the communication buffer size, ' - 'while a smaller value disables prefetching and may degrade performance. Adjust this value based on your system\'s memory ' - 'and performance requirements.') - group.add_argument('--keep-fp8-transpose-cache', action='store_true', - help='If set, keep the fp8 transpose cache when using Megatron FSDP.') - group.add_argument('--enable-full-sharding-in-hsdp', action='store_true', - help='If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel (HSDP) mode.') - group.add_argument('--num-distributed-optimizer-instances', type=int, default=1, - help='Number of Distributed Optimizer copies across Data Parallel domain.') - group.add_argument('--torch-fsdp2-no-reshard-after-forward', action='store_false', dest='torch_fsdp2_reshard_after_forward', - help='Whether to reshard weights after forward pass when using PyTorch FSDP2. ' - 'Set to enable FSDP ZeRO-2.') - group.add_argument('--cp-comm-type', nargs='+', type=str, default=["p2p"], - help='Inter-gpu communication type for context parallelism: ' - 'p2p, a2a, allgather or a2a+p2p. If a single string is provided, ' - 'all layers will share the same communication type. Users can also ' - 'specify separated types for each layer like ' - '--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p') - group.add_argument('--max-seqlen-per-dp-cp-rank', type=int, default=None, - help='Maximum sequence length per CP rank. This is used to calculate the ' - 'number of sub-samples assigned to each CP rank when using heterogeneous context parallel.') - group.add_argument('--hybrid-context-parallel', action='store_true', default=False, - help='Enables hybrid context parallel. This is used to balance the workload ' - 'of each CP rank when we use packed samples with variable sequence lengths. ' - 'Requires --max-seqlen-per-dp-cp-rank to be set.') - group.add_argument('--sequence-packing-scheduler', type=str, default=None, choices=['dp_balanced']) - group.add_argument('--fake-process-group', action='store_true', default=False, - help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ + group.add_argument( + '--cp-comm-type', + nargs='+', + type=str, + default=["p2p"], + help='Inter-gpu communication type for context parallelism: ' + 'p2p, a2a, allgather or a2a+p2p. If a single string is provided, ' + 'all layers will share the same communication type. Users can also ' + 'specify separated types for each layer like ' + '--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p', + ) + group.add_argument( + '--fake-process-group', + action='store_true', + default=False, + help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ This is quite useful for profiling memory usage of distributed training with just one GPU. \ - Setting WORLD_SIZE and RANK to the specific values for target distribtued scale.') + Setting WORLD_SIZE and RANK to the specific values for target distribtued scale.', + ) + group.add_argument( + '--use-layer-wise-param-layout', + action='store_true', + dest='use_layer_wise_param_layout', + default=False, + help='Opt INTO the padded shard-aligned LayerWise param layout. The default is the compact ' + 'decoupled layout, where LayerWise (Muon 2D) buffers use a no-padding DDP layout and locally ' + 'disable DistributedOptimizer (all-reduce grads + whole-param ping-pong + allgather_params), ' + 'while sibling buffers keep the byte-level DistributedOptimizer; this avoids the persistent ' + 'dp_size * max(shard_load) padding. The compact layout supports --fp8-param-gather (mxfp8 or ' + 'blockwise): the Muon buffers stay non-DistOpt and the param sync runs through the FP8-aware ' + 'allgather_params ping-pong path (stage fp32 master -> bf16, uneven all-gather, requantize ' + 'per rank); mxfp8 requires --reuse-grad-buf-for-mxfp8-param-ag. Pass this flag to restore the ' + 'padded layout (e.g. for bit-for-bit comparison; it uses a different bf16 reduction ordering).', + ) return parser @@ -3338,138 +4355,254 @@ def _add_tokenizer_args(parser): def _add_data_args(parser): group = parser.add_argument_group(title='data and dataloader') - group.add_argument('--data-path', nargs='*', default=None, - help='The weight and prefix list for a set of train, validation, and test' - 'datasets which split according to --split. The accepted formats are: ' - '(1) a single prefix, ' - '(2) a list of weight prefix pairs e.g. weight1 prefix1 weight2 prefix2, ' - '(3) a list of prefixes e.g. prefix1 prefix2. ' - 'For (3), weights are inferred from the lengths of the contributing datasets. ' - 'This argument is exclusive to the other independent --*-data-path arguments.') - group.add_argument('--phase-transition-iterations', type=str, default=None, - help='Comma-separated list of iterations where phase ' - 'transitions occur. Requires fixed global batch size across phases.') - group.add_argument('--split', type=str, default=None, - help='Comma-separated list of proportions for training,' - ' validation, and test split. For example the split ' - '`90,5,5` will use 90%% of data for training, 5%% for ' - 'validation and 5%% for test.') - group.add_argument('--train-data-path', nargs='*', default=None, - help='The weight and prefix list for an independent train dataset. ' - 'Follows the same pattern rules as --data-path.') - group.add_argument('--valid-data-path', nargs='*', default=None, - help='The weight and prefix list for an independent validation dataset. ' - 'Follows the same pattern rules as --data-path.') - group.add_argument('--test-data-path', nargs='*', default=None, - help='The weight and prefix list for an independent test dataset. ' - 'Follows the same pattern rules as --data-path.') - group.add_argument('--data-args-path', type=str, default=None, - help='Path to data-args. Instead of feeding `--data-path` ' - 'with weighted dataset, we pass in a file path from which ' - 'we read that argument. This is useful when the list of data is ' - 'too big.') - group.add_argument('--per-split-data-args-path', type=str, default=None, - help='Path to per-split-data-args. Instead of feeding ' - '`--(train|valid|test)-data-path` with weighted dataset, ' - 'we pass in a file path from which we read those arguments. ' - 'This is useful when the list of data is too big. Format is a ' - 'json file with `train`, `valid, `test` keys') - group.add_argument('--per-dataset-sequences-path', default=None, - help='Path to a json file with the sequences per dataset. Check the tools/build_sequences_per_dataset.py script to build this file.') - group.add_argument('--dataloader-fast-cache-load', action='store_true', - help='Option to use the fast cache loading path when building the datasets. Requires all the dataset caches to be built and stored in --data-cache-path.') - group.add_argument('--dataloader-defer-npy-index-mmap', action='store_true', - help='Defer the mmap of the dataset indexes (.npy files) until the first access. Requires all the dataset caches to be built and stored in --data-cache-path.') - group.add_argument('--data-cache-path', default=None, - help='Path to a directory to hold cached index files.') - group.add_argument('--no-mmap-bin-files', action='store_false', - help='Disable mmap-ing of .bin files.', - dest='mmap_bin_files') - group.add_argument('--mock-data', action='store_true', - help='Skip data loading and validation and opt for artificial ' - 'generation of mock data when an implementation is available.') - group.add_argument('--seq-length', type=int, default=None, - help='Maximum sequence length to process.') - group.add_argument('--encoder-seq-length', type=int, default=None, - help='Maximum encoder sequence length to process.' - 'This should be exclusive of --seq-length') - group.add_argument('--decoder-seq-length', type=int, default=None, - help="Maximum decoder sequence length to process.") - group.add_argument('--sample-rate', type=float, default=1.0, - help='sample rate for training data. Supposed to be 0 ' - ' < sample_rate < 1') - group.add_argument('--mask-prob', type=float, default=0.15, - help='Probability of replacing a token with mask.') - group.add_argument('--short-seq-prob', type=float, default=0.1, - help='Probability of producing a short sequence.') - group.add_argument('--num-workers', type=int, default=2, - help="Dataloader number of workers.") - group.add_argument('--reset-position-ids', action='store_true', - help='Reset posistion ids after end-of-document token.') - group.add_argument('--reset-attention-mask', action='store_true', - help='Reset self attention mask after ' - 'end-of-document token.') - group.add_argument('--dataloader-inter-document-masking', action='store_true', - help='Return cu_seqlens marking document boundaries ' - 'within each sample so that attention is restricted ' - 'to individual documents.') - group.add_argument('--eod-mask-loss', action='store_true', - help='Mask loss for the end of document tokens.') - group.add_argument('--no-create-attention-mask-in-dataloader', action='store_false', - help='If set, do not create attention_masks in dataloader.', - dest='create_attention_mask_in_dataloader') - group.add_argument('--num-dataset-builder-threads', type=int, default=1, - help='Number of parallel threads per rank for dataset builder') - group.add_argument('--object-storage-cache-path', type=str, default=None, - help='Path to cache index files when using s3 or msc dataloader') - group.add_argument('--mid-level-dataset-surplus', type=float, default=0.005, - help='The sample surplus to build for the mid-level datasets(s)') - group.add_argument('--allow-ambiguous-pad-tokens', action='store_true', - help='Whether to prevent pad tokens already present in the dataset ' - 'from being masked out when the pad token incorrectly shares the same id ' - 'with other special tokens in the tokenizer. Note that this argument has ' - 'no effect when the tokenizer correctly provides a unique id for the pad. ' - 'Masking out such ambiguous pad tokens results in training instability. ' - 'Such a scenario is best resolved by fixing the tokenizer; leaving this ' - 'option as False provides a workaround. ' - 'When left to the default of False, any token ids that collide with the ' - 'pad token id - as provided by the tokenizer - will not be masked out of ' - 'the loss calculation: it cannot be determined whether they are truly pad. ' - 'If instead this argument is set, the training flow will treat all tokens ' - 'that share the same id as the pad token as true pad tokens, potentially ' - 'causing severe training instability.') + group.add_argument( + '--data-path', + nargs='*', + default=None, + help='The weight and prefix list for a set of train, validation, and test' + 'datasets which split according to --split. The accepted formats are: ' + '(1) a single prefix, ' + '(2) a list of weight prefix pairs e.g. weight1 prefix1 weight2 prefix2, ' + '(3) a list of prefixes e.g. prefix1 prefix2. ' + 'For (3), weights are inferred from the lengths of the contributing datasets. ' + 'This argument is exclusive to the other independent --*-data-path arguments.', + ) + group.add_argument( + '--phase-transition-iterations', + type=str, + default=None, + help='Comma-separated list of iterations where phase ' + 'transitions occur. Requires fixed global batch size across phases.', + ) + group.add_argument( + '--split', + type=str, + default=None, + help='Comma-separated list of proportions for training,' + ' validation, and test split. For example the split ' + '`90,5,5` will use 90%% of data for training, 5%% for ' + 'validation and 5%% for test.', + ) + group.add_argument( + '--train-data-path', + nargs='*', + default=None, + help='The weight and prefix list for an independent train dataset. ' + 'Follows the same pattern rules as --data-path.', + ) + group.add_argument( + '--valid-data-path', + nargs='*', + default=None, + help='The weight and prefix list for an independent validation dataset. ' + 'Follows the same pattern rules as --data-path.', + ) + group.add_argument( + '--test-data-path', + nargs='*', + default=None, + help='The weight and prefix list for an independent test dataset. ' + 'Follows the same pattern rules as --data-path.', + ) + group.add_argument( + '--data-args-path', + type=str, + default=None, + help='Path to data-args. Instead of feeding `--data-path` ' + 'with weighted dataset, we pass in a file path from which ' + 'we read that argument. This is useful when the list of data is ' + 'too big.', + ) + group.add_argument( + '--per-split-data-args-path', + type=str, + default=None, + help='Path to per-split-data-args. Instead of feeding ' + '`--(train|valid|test)-data-path` with weighted dataset, ' + 'we pass in a file path from which we read those arguments. ' + 'This is useful when the list of data is too big. Format is a ' + 'json file with `train`, `valid, `test` keys', + ) + group.add_argument( + '--per-dataset-sequences-path', + default=None, + help='Path to a json file with the sequences per dataset. Check the tools/build_sequences_per_dataset.py script to build this file.', + ) + group.add_argument( + '--dataloader-fast-cache-load', + action='store_true', + help='Option to use the fast cache loading path when building the datasets. Requires all the dataset caches to be built and stored in --data-cache-path.', + ) + group.add_argument( + '--dataloader-defer-npy-index-mmap', + action='store_true', + help='Defer the mmap of the dataset indexes (.npy files) until the first access. Requires all the dataset caches to be built and stored in --data-cache-path.', + ) + group.add_argument( + '--data-cache-path', default=None, help='Path to a directory to hold cached index files.' + ) + group.add_argument( + '--no-mmap-bin-files', + action='store_false', + help='Disable mmap-ing of .bin files.', + dest='mmap_bin_files', + ) + group.add_argument( + '--mock-data', + action='store_true', + help='Skip data loading and validation and opt for artificial ' + 'generation of mock data when an implementation is available.', + ) + group.add_argument( + '--seq-length', type=int, default=None, help='Maximum sequence length to process.' + ) + group.add_argument( + '--encoder-seq-length', + type=int, + default=None, + help='Maximum encoder sequence length to process.' + 'This should be exclusive of --seq-length', + ) + group.add_argument( + '--decoder-seq-length', + type=int, + default=None, + help="Maximum decoder sequence length to process.", + ) + group.add_argument( + '--sample-rate', + type=float, + default=1.0, + help='sample rate for training data. Supposed to be 0 ' ' < sample_rate < 1', + ) + group.add_argument( + '--mask-prob', type=float, default=0.15, help='Probability of replacing a token with mask.' + ) + group.add_argument( + '--short-seq-prob', + type=float, + default=0.1, + help='Probability of producing a short sequence.', + ) + group.add_argument('--num-workers', type=int, default=2, help="Dataloader number of workers.") + group.add_argument( + '--reset-position-ids', + action='store_true', + help='Reset posistion ids after end-of-document token.', + ) + group.add_argument( + '--reset-attention-mask', + action='store_true', + help='Reset self attention mask after ' 'end-of-document token.', + ) + group.add_argument( + '--dataloader-inter-document-masking', + action='store_true', + help='Return cu_seqlens marking document boundaries ' + 'within each sample so that attention is restricted ' + 'to individual documents.', + ) + group.add_argument( + '--eod-mask-loss', action='store_true', help='Mask loss for the end of document tokens.' + ) + group.add_argument( + '--no-create-attention-mask-in-dataloader', + action='store_false', + help='If set, do not create attention_masks in dataloader.', + dest='create_attention_mask_in_dataloader', + ) + group.add_argument( + '--num-dataset-builder-threads', + type=int, + default=1, + help='Number of parallel threads per rank for dataset builder', + ) + group.add_argument( + '--object-storage-cache-path', + type=str, + default=None, + help='Path to cache index files when using s3 or msc dataloader', + ) + group.add_argument( + '--mid-level-dataset-surplus', + type=float, + default=0.005, + help='The sample surplus to build for the mid-level datasets(s)', + ) + group.add_argument( + '--allow-ambiguous-pad-tokens', + action='store_true', + help='Whether to prevent pad tokens already present in the dataset ' + 'from being masked out when the pad token incorrectly shares the same id ' + 'with other special tokens in the tokenizer. Note that this argument has ' + 'no effect when the tokenizer correctly provides a unique id for the pad. ' + 'Masking out such ambiguous pad tokens results in training instability. ' + 'Such a scenario is best resolved by fixing the tokenizer; leaving this ' + 'option as False provides a workaround. ' + 'When left to the default of False, any token ids that collide with the ' + 'pad token id - as provided by the tokenizer - will not be masked out of ' + 'the loss calculation: it cannot be determined whether they are truly pad. ' + 'If instead this argument is set, the training flow will treat all tokens ' + 'that share the same id as the pad token as true pad tokens, potentially ' + 'causing severe training instability.', + ) group.add_argument('--fim-data', action='store_true', help='Whether to use the FIM dataset.') - group.add_argument('--fim-rate', type=float, default=0.5, - help='Probability to convert a training sample into a FIM format.') - group.add_argument('--fim-spm-rate', type=float, default=0.5, - help='Probability that the a FIM sample uses the SPM format over the PSM format.') - group.add_argument('--fim-split-sample', type=str, default=None, - help='String around which to split the sample for FIM.') - group.add_argument('--fim-fragment-rate', type=float, default=None, - help='Rate of FIM on each fragment when --fim-split-sample is not None.') - group.add_argument('--fim-no-prefix', type=str, default=None, - help='Do not apply FIM to fragments that start with this prefix') - group.add_argument('--fim-prefix-token', type=str, default='', - help='FIM prefix token') - group.add_argument('--fim-middle-token', type=str, default='', - help='FIM middle token') - group.add_argument('--fim-suffix-token', type=str, default='', - help='FIM suffix token') - group.add_argument('--fim-pad-token', type=str, default='', - help='FIM PAD token') - group.add_argument('--fim-eod-token', type=str, default='<|endoftext|>', - help='FIM EOD token') + group.add_argument( + '--fim-rate', + type=float, + default=0.5, + help='Probability to convert a training sample into a FIM format.', + ) + group.add_argument( + '--fim-spm-rate', + type=float, + default=0.5, + help='Probability that the a FIM sample uses the SPM format over the PSM format.', + ) + group.add_argument( + '--fim-split-sample', + type=str, + default=None, + help='String around which to split the sample for FIM.', + ) + group.add_argument( + '--fim-fragment-rate', + type=float, + default=None, + help='Rate of FIM on each fragment when --fim-split-sample is not None.', + ) + group.add_argument( + '--fim-no-prefix', + type=str, + default=None, + help='Do not apply FIM to fragments that start with this prefix', + ) + group.add_argument( + '--fim-prefix-token', type=str, default='', help='FIM prefix token' + ) + group.add_argument( + '--fim-middle-token', type=str, default='', help='FIM middle token' + ) + group.add_argument( + '--fim-suffix-token', type=str, default='', help='FIM suffix token' + ) + group.add_argument('--fim-pad-token', type=str, default='', help='FIM PAD token') + group.add_argument('--fim-eod-token', type=str, default='<|endoftext|>', help='FIM EOD token') return parser def _add_autoresume_args(parser): group = parser.add_argument_group(title='autoresume') - group.add_argument('--adlr-autoresume', action='store_true', - help='Enable autoresume on adlr cluster.') - group.add_argument('--adlr-autoresume-interval', type=int, default=1000, - help='Intervals over which check for autoresume' - 'termination signal') + group.add_argument( + '--adlr-autoresume', action='store_true', help='Enable autoresume on adlr cluster.' + ) + group.add_argument( + '--adlr-autoresume-interval', + type=int, + default=1000, + help='Intervals over which check for autoresume' 'termination signal', + ) return parser @@ -3478,56 +4611,95 @@ def _add_biencoder_args(parser): group = parser.add_argument_group(title='biencoder') # network size - group.add_argument('--ict-head-size', type=int, default=None, - help='Size of block embeddings to be used in ICT and ' - 'REALM (paper default: 128)') - group.add_argument('--biencoder-projection-dim', type=int, default=0, - help='Size of projection head used in biencoder (paper' - ' default: 128)') - group.add_argument('--biencoder-shared-query-context-model', action='store_true', - help='Whether to share the parameters of the query ' - 'and context models or not') + group.add_argument( + '--ict-head-size', + type=int, + default=None, + help='Size of block embeddings to be used in ICT and ' 'REALM (paper default: 128)', + ) + group.add_argument( + '--biencoder-projection-dim', + type=int, + default=0, + help='Size of projection head used in biencoder (paper' ' default: 128)', + ) + group.add_argument( + '--biencoder-shared-query-context-model', + action='store_true', + help='Whether to share the parameters of the query ' 'and context models or not', + ) # checkpointing - group.add_argument('--ict-load', type=str, default=None, - help='Directory containing an ICTBertModel checkpoint') - group.add_argument('--bert-load', type=str, default=None, - help='Directory containing an BertModel checkpoint ' - '(needed to start ICT and REALM)') + group.add_argument( + '--ict-load', type=str, default=None, help='Directory containing an ICTBertModel checkpoint' + ) + group.add_argument( + '--bert-load', + type=str, + default=None, + help='Directory containing an BertModel checkpoint ' '(needed to start ICT and REALM)', + ) # data - group.add_argument('--titles-data-path', type=str, default=None, - help='Path to titles dataset used for ICT') - group.add_argument('--query-in-block-prob', type=float, default=0.1, - help='Probability of keeping query in block for ' - 'ICT dataset') - group.add_argument('--use-one-sent-docs', action='store_true', - help='Whether to use one sentence documents in ICT') - group.add_argument('--evidence-data-path', type=str, default=None, - help='Path to Wikipedia Evidence frm DPR paper') + group.add_argument( + '--titles-data-path', type=str, default=None, help='Path to titles dataset used for ICT' + ) + group.add_argument( + '--query-in-block-prob', + type=float, + default=0.1, + help='Probability of keeping query in block for ' 'ICT dataset', + ) + group.add_argument( + '--use-one-sent-docs', + action='store_true', + help='Whether to use one sentence documents in ICT', + ) + group.add_argument( + '--evidence-data-path', + type=str, + default=None, + help='Path to Wikipedia Evidence frm DPR paper', + ) # training - group.add_argument('--retriever-report-topk-accuracies', nargs='+', type=int, - default=[], help="Which top-k accuracies to report " - "(e.g. '1 5 20')") - group.add_argument('--retriever-score-scaling', action='store_true', - help='Whether to scale retriever scores by inverse ' - 'square root of hidden size') + group.add_argument( + '--retriever-report-topk-accuracies', + nargs='+', + type=int, + default=[], + help="Which top-k accuracies to report " "(e.g. '1 5 20')", + ) + group.add_argument( + '--retriever-score-scaling', + action='store_true', + help='Whether to scale retriever scores by inverse ' 'square root of hidden size', + ) # faiss index - group.add_argument('--block-data-path', type=str, default=None, - help='Where to save/load BlockData to/from') - group.add_argument('--embedding-path', type=str, default=None, - help='Where to save/load Open-Retrieval Embedding' - ' data to/from') + group.add_argument( + '--block-data-path', type=str, default=None, help='Where to save/load BlockData to/from' + ) + group.add_argument( + '--embedding-path', + type=str, + default=None, + help='Where to save/load Open-Retrieval Embedding' ' data to/from', + ) # indexer - group.add_argument('--indexer-batch-size', type=int, default=128, - help='How large of batches to use when doing indexing ' - 'jobs') - group.add_argument('--indexer-log-interval', type=int, default=1000, - help='After how many batches should the indexer ' - 'report progress') + group.add_argument( + '--indexer-batch-size', + type=int, + default=128, + help='How large of batches to use when doing indexing ' 'jobs', + ) + group.add_argument( + '--indexer-log-interval', + type=int, + default=1000, + help='After how many batches should the indexer ' 'report progress', + ) return parser @@ -3535,144 +4707,299 @@ def _add_vision_args(parser): group = parser.add_argument_group(title="vision") # general vision arguements - group.add_argument('--num-classes', type=int, default=1000, - help='num of classes in vision classificaiton task') - group.add_argument('--img-h', type=int, default=224, - help='Image height for vision classification task') - group.add_argument('--img-w', type=int, default=224, - help='Image height for vision classification task') - group.add_argument('--num-channels', type=int, default=3, - help='Number of channels in input image data') - group.add_argument('--patch-dim', type=int, default=16, - help='patch dimension') - group.add_argument('--classes-fraction', type=float, default=1.0, - help='training with fraction of classes.') - group.add_argument('--data-per-class-fraction', type=float, default=1.0, - help='training with fraction of data per class.') - group.add_argument('--no-data-sharding', action='store_false', - help='Disable data sharding.', - dest='data_sharding') - group.add_argument('--head-lr-mult', type=float, default=1.0, - help='learning rate multiplier for head during finetuning') + group.add_argument( + '--num-classes', type=int, default=1000, help='num of classes in vision classificaiton task' + ) + group.add_argument( + '--img-h', type=int, default=224, help='Image height for vision classification task' + ) + group.add_argument( + '--img-w', type=int, default=224, help='Image height for vision classification task' + ) + group.add_argument( + '--num-channels', type=int, default=3, help='Number of channels in input image data' + ) + group.add_argument('--patch-dim', type=int, default=16, help='patch dimension') + group.add_argument( + '--classes-fraction', type=float, default=1.0, help='training with fraction of classes.' + ) + group.add_argument( + '--data-per-class-fraction', + type=float, + default=1.0, + help='training with fraction of data per class.', + ) + group.add_argument( + '--no-data-sharding', + action='store_false', + help='Disable data sharding.', + dest='data_sharding', + ) + group.add_argument( + '--head-lr-mult', + type=float, + default=1.0, + help='learning rate multiplier for head during finetuning', + ) # pretraining type and backbone selection` - group.add_argument('--vision-pretraining', action='store_true', - help='flag to indicate vision pretraining') - group.add_argument('--vision-pretraining-type', type=str, default='classify', - choices=['classify', 'inpaint', 'dino'], - help='pretraining objectives') - group.add_argument('--vision-backbone-type', type=str, default='vit', - choices=['vit', 'mit', 'swin'], - help='backbone types types') - group.add_argument('--swin-backbone-type', type=str, default='tiny', - choices=['tiny', 'base', 'h3'], - help='pretraining objectives') + group.add_argument( + '--vision-pretraining', action='store_true', help='flag to indicate vision pretraining' + ) + group.add_argument( + '--vision-pretraining-type', + type=str, + default='classify', + choices=['classify', 'inpaint', 'dino'], + help='pretraining objectives', + ) + group.add_argument( + '--vision-backbone-type', + type=str, + default='vit', + choices=['vit', 'mit', 'swin'], + help='backbone types types', + ) + group.add_argument( + '--swin-backbone-type', + type=str, + default='tiny', + choices=['tiny', 'base', 'h3'], + help='pretraining objectives', + ) # inpainting arguments - group.add_argument('--mask-type', type=str, default='random', - choices=['random', 'row'], - help='mask types') - group.add_argument('--mask-factor', type=float, default=1.0, - help='mask size scaling parameter') + group.add_argument( + '--mask-type', type=str, default='random', choices=['random', 'row'], help='mask types' + ) + group.add_argument('--mask-factor', type=float, default=1.0, help='mask size scaling parameter') # dino arguments - group.add_argument('--iter-per-epoch', type=int, default=1250, - help='iterations per epoch') - group.add_argument('--dino-local-img-size', type=int, default=96, - help='Image size for vision classification task') - group.add_argument('--dino-local-crops-number', type=int, default=10, - help='Number of local crops') - group.add_argument('--dino-head-hidden-size', type=int, default=2048, - help='Hidden dimension size in dino head') - group.add_argument('--dino-bottleneck-size', type=int, default=256, - help='Bottle neck dimension in dino head ') - group.add_argument('--dino-freeze-last-layer', type=float, default=1, - help='Freezing last layer weights') - group.add_argument('--dino-norm-last-layer', action='store_true', - help='Disable Norm in last layer.') - group.add_argument('--dino-warmup-teacher-temp', type=float, default=0.04, - help='warump teacher temperature') - group.add_argument('--dino-teacher-temp', type=float, default=0.07, - help='teacher temperature') - group.add_argument('--dino-warmup-teacher-temp-epochs', type=int, default=30, - help='warmup teacher temperaure epochs') + group.add_argument('--iter-per-epoch', type=int, default=1250, help='iterations per epoch') + group.add_argument( + '--dino-local-img-size', + type=int, + default=96, + help='Image size for vision classification task', + ) + group.add_argument( + '--dino-local-crops-number', type=int, default=10, help='Number of local crops' + ) + group.add_argument( + '--dino-head-hidden-size', type=int, default=2048, help='Hidden dimension size in dino head' + ) + group.add_argument( + '--dino-bottleneck-size', type=int, default=256, help='Bottle neck dimension in dino head ' + ) + group.add_argument( + '--dino-freeze-last-layer', type=float, default=1, help='Freezing last layer weights' + ) + group.add_argument( + '--dino-norm-last-layer', action='store_true', help='Disable Norm in last layer.' + ) + group.add_argument( + '--dino-warmup-teacher-temp', type=float, default=0.04, help='warump teacher temperature' + ) + group.add_argument('--dino-teacher-temp', type=float, default=0.07, help='teacher temperature') + group.add_argument( + '--dino-warmup-teacher-temp-epochs', + type=int, + default=30, + help='warmup teacher temperaure epochs', + ) return parser + def _add_moe_args(parser): group = parser.add_argument_group(title="moe") # General arguments - group.add_argument('--num-experts', type=int, default=None, - help='Number of Experts in MoE (None means no MoE)') - group.add_argument('--moe-layer-freq', type=moe_freq_type, default=1, - help='Frequency between MoE layers and Dense layers. Accepts either: ' - '- An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers ' - '- A string containing a Python list expression that defines a custom pattern, e.g.: ' - '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' - 'where 1 indicates an expert layer and 0 indicates a dense layer. ' - 'Examples: "([0]+[1]*23)": 1 dense layer followed by 23 expert layers, ' - '"([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.') - group.add_argument('--moe-use-upcycling', action='store_true', - help='Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. ' - 'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.') + group.add_argument( + '--num-experts', type=int, default=None, help='Number of Experts in MoE (None means no MoE)' + ) + group.add_argument( + '--moe-layer-freq', + type=moe_freq_type, + default=1, + help='Frequency between MoE layers and Dense layers. Accepts either: ' + '- An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers ' + '- A string containing a Python list expression that defines a custom pattern, e.g.: ' + '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' + 'where 1 indicates an expert layer and 0 indicates a dense layer. ' + 'Examples: "([0]+[1]*23)": 1 dense layer followed by 23 expert layers, ' + '"([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.', + ) + group.add_argument( + '--moe-use-upcycling', + action='store_true', + help='Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. ' + 'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.', + ) # Router arguments - group.add_argument('--moe-router-load-balancing-type', nargs='+', type=str, - choices=['aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'quantile_balancing', 'none'], - default='aux_loss', - help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE; "quantile_balancing" (QB) uses dual coordinate descent on a per-expert bias to handle load balance internally; "none" implies no load balancing. The default is "aux_loss".') - group.add_argument('--moe-aux-loss-coeff', type=float, nargs='+', default=0.0, - help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.') + group.add_argument( + '--moe-router-load-balancing-type', + nargs='+', + type=str, + choices=['aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'none'], + default='aux_loss', + help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".', + ) + group.add_argument( + '--moe-aux-loss-coeff', + type=float, + nargs='+', + default=0.0, + help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.', + ) # Token dispatcher arguments # MoE communication overlap arguments - group.add_argument('--moe-upcycling-granularity', type=int, default=1, - help='This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. ' - 'For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy.') + group.add_argument( + '--moe-upcycling-granularity', + type=int, + default=1, + help='This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. ' + 'For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy.', + ) return parser + def _add_mla_args(parser): group = parser.add_argument_group(title="mla") - group.add_argument('--q-lora-rank', type=int, default=None, - help="Rank of Query tensor's low rank representation.") - group.add_argument('--kv-lora-rank', type=int, default=32, - help="Rank of Key and Value tensors' low rank representation.") - group.add_argument('--qk-head-dim', type=int, default=128, - help="Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim") - group.add_argument('--qk-pos-emb-head-dim', type=int, default=64, - help="Dimension of the position embedding in the QK projection.") - group.add_argument('--v-head-dim', type=int, default=128, - help="Dimension of the head in the V projection.") - group.add_argument('--rotary-scaling-factor', type=float, default=1.0, - help="Rotary scaling factor for the rotary embeddings.") - group.add_argument('--mscale', type=float, default=1.0, - help="Mscale for YaRN RoPE in multi-latent attention.") - group.add_argument('--mscale-all-dim', type=float, default=0.0, - help="Mscale all dimensions for YaRN RoPE in multi-latent attention.") - group.add_argument('--cache-mla-latents', action='store_true', default=False, - help="If set caches the mla down projected latents with mla flash decode.") + group.add_argument( + '--q-lora-rank', + type=int, + default=None, + help="Rank of Query tensor's low rank representation.", + ) + group.add_argument( + '--kv-lora-rank', + type=int, + default=32, + help="Rank of Key and Value tensors' low rank representation.", + ) + group.add_argument( + '--qk-head-dim', + type=int, + default=128, + help="Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim", + ) + group.add_argument( + '--qk-pos-emb-head-dim', + type=int, + default=64, + help="Dimension of the position embedding in the QK projection.", + ) + group.add_argument( + '--v-head-dim', type=int, default=128, help="Dimension of the head in the V projection." + ) + group.add_argument( + '--rotary-scaling-factor', + type=float, + default=1.0, + help="Rotary scaling factor for the rotary embeddings.", + ) + group.add_argument( + '--original-max-position-embeddings', + type=int, + default=4096, + help="Original maximum position embeddings for the original model, used by yarn.", + ) + group.add_argument( + '--mscale', type=float, default=1.0, help="Mscale for YaRN RoPE in multi-latent attention." + ) + group.add_argument( + '--mscale-all-dim', + type=float, + default=0.0, + help="Mscale all dimensions for YaRN RoPE in multi-latent attention.", + ) + group.add_argument( + '--o-groups', + type=int, + default=8, + help="Number of groups for grouped output (wo_a). 0 = single linear.", + ) + group.add_argument( + '--o-lora-rank', + type=int, + default=1024, + help="Low-rank dimension per group for grouped output (wo_a). Used when o-groups > 0.", + ) + group.add_argument( + '--cache-mla-latents', + action='store_true', + default=False, + help="If set caches the mla down projected latents with mla flash decode.", + ) group.add_argument( '--mla-down-proj-fusion', action='store_true', default=False, help="Enable fused q/kv down-projection and fused input layernorm when backend supports. " - "Otherwise fall back to the unfused MLA.", + "Otherwise fall back to the unfused MLA.", ) return parser + def _add_experimental_attention_variant_args(parser): group = parser.add_argument_group(title="experimental_attention_variant") # Linear attention - group.add_argument('--linear-attention-freq', type=la_freq_type, default=None, - help='Frequency between LA (linear attention) layers and' - ' SDPA (scaled dot-product attention) layers. Accepts either: ' - '- An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer ' - '- A string containing a Python list expression that defines a custom pattern, e.g.: ' - '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' - 'where 1 indicates an LA layer and 0 indicates a SDPA layer. ' - 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' - '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.') + group.add_argument( + '--linear-attention-freq', + type=la_freq_type, + default=None, + help='Frequency between LA (linear attention) layers and' + ' SDPA (scaled dot-product attention) layers. Accepts either: ' + '- An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer ' + '- A string containing a Python list expression that defines a custom pattern, e.g.: ' + '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' + 'where 1 indicates an LA layer and 0 indicates a SDPA layer. ' + 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' + '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.', + ) + group.add_argument( + '--csa-compress-ratios', + type=compress_ratios_type, + default=None, + help='Per-layer compress ratios for compressed sparse attention. ' + 'Accepts a string containing a Python list expression, e.g.: ' + '"[0,0,4,128,4,128]" or "([0]+[4,128]*2)*3". ' + 'Each value is the compression ratio for the corresponding ' + 'transformer layer (valid values: 0, 4, 128; 0 = sliding-window-only, the "W" ' + 'hybrid layer symbol). ' + 'For HybridModel with --hybrid-layer-pattern, the preferred compact form has one ' + 'entry per W/C/H attention symbol; legacy zero-padded one-entry-per-hybrid-layer ' + 'lists are also accepted.', + ) + group.add_argument( + '--no-dsa-kernel-fusion', + action='store_false', + default=None, + dest='apply_dsa_kernel_fusion', + help='Deprecated compatibility flag for DSv4 hybrid attention. ' + 'Use --dsa-kernel-backend none instead.', + ) + # Defined manually (and excluded from ArgumentGroupFactory) so that an omitted + # flag is distinguishable from an explicit "none": dsv4_hybrid launches resolve + # an omitted flag to "cudnn" in core_transformer_config_from_args, preserving + # the historical fused-by-default CLI behavior of --no-dsa-kernel-fusion. + group.add_argument( + '--dsa-kernel-backend', + type=str, + choices=['none', 'tilelang', 'cudnn'], + default=None, + help='Fused DSA kernel backend. When omitted, resolves to "cudnn" for ' + 'experimental_attention_variant="dsv4_hybrid" launches and to "none" ' + 'otherwise.', + ) + # Note: --dsa-indexer-{n-heads,head-dim,topk,loss-coeff,use-sparse-loss}, + # --csa-window-size, --csa-compress-rotary-base, --csa-dense-mode are + # auto-generated by ArgumentGroupFactory from TransformerConfig fields + # (none of them are in the exclude list at line 2500-2576). return parser + def _add_heterogeneous_args(parser): """ Heterogeneous models refer to transformer architectures where individual layers can differ @@ -3712,118 +5039,210 @@ def _add_heterogeneous_args(parser): ] """ group = parser.add_argument_group(title="heterogeneous architecture") - group.add_argument('--heterogeneous-layers-config-path', type=str, default=None, - help='Path to json file containing heterogeneous model configuration. ' - 'Use the format of the HuggingFace config files in llama nemotron ' - 'models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.') - group.add_argument('--heterogeneous-layers-config-encoded-json', type=str, default=None, - help='This is encoded json string of the heterogeneous model configuration. Used to keep the content ' - 'of the heterogeneous model specification in args when the model is loaded from a checkpoint. ' - 'Use the format of the HuggingFace config files in llama nemotron ' - 'models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.') + group.add_argument( + '--heterogeneous-layers-config-path', + type=str, + default=None, + help='Path to json file containing heterogeneous model configuration. ' + 'Use the format of the HuggingFace config files in llama nemotron ' + 'models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.', + ) + group.add_argument( + '--heterogeneous-layers-config-encoded-json', + type=str, + default=None, + help='This is encoded json string of the heterogeneous model configuration. Used to keep the content ' + 'of the heterogeneous model specification in args when the model is loaded from a checkpoint. ' + 'Use the format of the HuggingFace config files in llama nemotron ' + 'models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.', + ) return parser + def _add_experimental_args(parser): group = parser.add_argument_group(title='experimental') - group.add_argument('--enable-experimental', action='store_true', - help='Enable experimental features.') - group.add_argument('--spec', type=str, default=None, nargs='*', - help='Specify the pair ' - 'that returns a spec to customize a model, transformer ' - 'block, or transformer layer, depending on the use case.' - 'To use local spec specify local as the argument.' - 'For more details, see the model class, ' - '`transformer_block.py`, or `transformer_layer.py`') - group.add_argument('--hybrid-layer-pattern', type=str, default=None, - help='Specify a hybrid layer pattern using M (mamba), G (gdn), ' - '* (attention), D (dsa), - (mlp), E (moe). Use | to define pipeline ' - 'stage boundaries for flexible virtual pipeline parallel (fVPP). ' - 'Use / to separate MTP patterns. ' - 'Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". ' - 'When this flag is used, it is the sole indicator that a hybrid model ' - 'is being run.') - group.add_argument('--hybrid-override-pattern', type=str, default=None, - help='Deprecated. Use --hybrid-layer-pattern instead. ' - 'If specified, its value will be forwarded to --hybrid-layer-pattern.') - group.add_argument('--yaml-cfg', type=str, default=None, - help = 'Config file to add additional arguments') + group.add_argument( + '--enable-experimental', action='store_true', help='Enable experimental features.' + ) + group.add_argument( + '--spec', + type=str, + default=None, + nargs='*', + help='Specify the pair ' + 'that returns a spec to customize a model, transformer ' + 'block, or transformer layer, depending on the use case.' + 'To use local spec specify local as the argument.' + 'For more details, see the model class, ' + '`transformer_block.py`, or `transformer_layer.py`', + ) + group.add_argument( + '--hybrid-layer-pattern', + type=str, + default=None, + help='Specify a hybrid layer pattern using M (mamba), G (gdn), K (kda), ' + '* (attention), D (dsa), - (mlp), E (moe). Use | to define pipeline ' + 'stage boundaries for flexible virtual pipeline parallel (fVPP). ' + 'Use / to separate MTP patterns. ' + 'Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". ' + 'When this flag is used, it is the sole indicator that a hybrid model ' + 'is being run.', + ) + group.add_argument( + '--hybrid-override-pattern', + type=str, + default=None, + help='Deprecated. Use --hybrid-layer-pattern instead. ' + 'If specified, its value will be forwarded to --hybrid-layer-pattern.', + ) + group.add_argument( + '--yaml-cfg', type=str, default=None, help='Config file to add additional arguments' + ) # Args of precision-aware optimizer. - group.add_argument('--use-precision-aware-optimizer', action='store_true', - help='Use the precision-aware optimizer in TransformerEngine, which allows ' - 'setting the main params and optimizer states to lower precision, such as ' - 'fp16, bf16 and fp8.') - group.add_argument('--main-grads-dtype', default='fp32', choices=['fp32', 'bf16'], - help='Dtype of main grads when enabling precision-aware-optimizer.') - group.add_argument('--main-params-dtype', default='fp32', choices=['fp32', 'fp16'], - help='Dtype of main params when enabling precision-aware-optimizer.') - group.add_argument('--exp-avg-dtype', default='fp32', choices=['fp32', 'fp16', 'bf16', 'fp8'], - help='Dtype of exp_avg (1st moment in adam optimizer) when enabling ' - 'precision-aware-optimizer. This dtype is used for storing the ' - 'optimizer state in memory during training but does not affect ' - 'the precision in the kernel computation.') - group.add_argument('--exp-avg-sq-dtype', default='fp32', choices=['fp32', 'fp16', 'bf16', 'fp8'], - help='Dtype of exp_avg_sq (2nd moment in adam optimizer) when enabling ' - 'precision-aware-optimizer. This dtype is used for storing the ' - 'optimizer state in memory during training but does not affect ' - 'the precision in the kernel computation.') + group.add_argument( + '--use-precision-aware-optimizer', + action='store_true', + help='Use the precision-aware optimizer in TransformerEngine, which allows ' + 'setting the main params and optimizer states to lower precision, such as ' + 'fp16, bf16 and fp8.', + ) + group.add_argument( + '--main-grads-dtype', + default='fp32', + choices=['fp32', 'bf16'], + help='Dtype of main grads when enabling precision-aware-optimizer.', + ) + group.add_argument( + '--main-params-dtype', + default='fp32', + choices=['fp32', 'fp16'], + help='Dtype of main params when enabling precision-aware-optimizer.', + ) + group.add_argument( + '--exp-avg-dtype', + default='fp32', + choices=['fp32', 'fp16', 'bf16', 'fp8'], + help='Dtype of exp_avg (1st moment in adam optimizer) when enabling ' + 'precision-aware-optimizer. This dtype is used for storing the ' + 'optimizer state in memory during training but does not affect ' + 'the precision in the kernel computation.', + ) + group.add_argument( + '--exp-avg-sq-dtype', + default='fp32', + choices=['fp32', 'fp16', 'bf16', 'fp8'], + help='Dtype of exp_avg_sq (2nd moment in adam optimizer) when enabling ' + 'precision-aware-optimizer. This dtype is used for storing the ' + 'optimizer state in memory during training but does not affect ' + 'the precision in the kernel computation.', + ) # Megatron-FSDP Arguments - group.add_argument('--megatron-fsdp-main-params-dtype', default='fp32', choices=['fp32', 'bf16', 'fp16', 'auto'], - help="Data type for the main weight buffer utilized for distributed optimization " - "and quantization with Megatron-FSDP. If 'auto', then the native model parameter " - "data-type will be used for the main weight data-type. Replaces --main-params-dtype.") - group.add_argument('--megatron-fsdp-main-grads-dtype', default='auto', choices=['fp32', 'bf16', 'fp16', 'auto'], - help="Data type for the main gradient buffer utilized for distributed optimization " - "with Megatron-FSDP. If 'auto', then the native model gradient data-type will " - "be used for the main gradient / accumulation data-type. Replaces --main-grads-dtype.") - group.add_argument("--megatron-fsdp-grad-comm-dtype", default='auto', choices=['fp32', 'fp16', 'bf16', 'auto'], - help="When using Megatron-FSDP, this controls the data-type used when communicating " - "model gradients during FSDP. If 'auto', then the main gradient data-type will " - "be used for the gradient communication / reduction data-type. When using NCCL " - "v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels.") - group.add_argument( - '--megatron-fsdp-enable-fine-grained-param-gather', - action='store_true', - default=False, - dest='megatron_fsdp_enable_fine_grained_param_gather', - help=( - 'If set, enables fine-grained parameter gathering for Megatron-FSDP. ' - 'This allows greater overlap between parameter all-gather operations and ' - 'forward computation, at the cost of additional communication calls. ' - 'For MXFP8, this helps save memory during fine-grained activation ' - 'recomputation, because MXFP8 forward and backward passes use different ' - 'parameter representations (rowwise data for forward, colwise data for ' - 'backward). Only the rowwise parameters of modules involved in ' - 'recomputation will be unsharded.' - ), - ) + group.add_argument( + '--megatron-fsdp-main-params-dtype', + default='fp32', + choices=['fp32', 'bf16', 'fp16', 'auto'], + help="Data type for the main weight buffer utilized for distributed optimization " + "and quantization with Megatron-FSDP. If 'auto', then the native model parameter " + "data-type will be used for the main weight data-type.", + ) + group.add_argument( + '--megatron-fsdp-main-grads-dtype', + default='auto', + choices=['fp32', 'bf16', 'fp16', 'auto'], + help="Data type for the main gradient buffer utilized for distributed optimization " + "with Megatron-FSDP. If 'auto', then the native model gradient data-type will " + "be used for the main gradient / accumulation data-type.", + ) + group.add_argument( + "--megatron-fsdp-grad-comm-dtype", + default='auto', + choices=['fp32', 'fp16', 'bf16', 'auto'], + help="When using Megatron-FSDP, this controls the data-type used when communicating " + "model gradients during FSDP. If 'auto', then the main gradient data-type will " + "be used for the gradient communication / reduction data-type. When using NCCL " + "v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels.", + ) + group.add_argument( + '--megatron-fsdp-prefetch-recompute-forward-weights', + action='store_true', + default=False, + dest='megatron_fsdp_prefetch_recompute_forward_weights', + help=( + 'If set, Megatron-FSDP prefetches rowwise weights needed by activation ' + 'recomputation during backward before prefetching backward transpose ' + 'weights.' + ), + ) + group.add_argument( + '--megatron-fsdp-cache-param-bucket-views', + action='store_true', + default=False, + dest='megatron_fsdp_cache_param_bucket_views', + help=( + 'If set, Megatron-FSDP caches parameter bucket views to reduce repeated ' + 'Python-side view setup when attaching module parameters to all-gather buckets.' + ), + ) + group.add_argument( + '--megatron-fsdp-enable-fine-grained-param-gather', + action='store_true', + default=False, + dest='megatron_fsdp_enable_fine_grained_param_gather', + help=( + 'If set, enables fine-grained parameter gathering for Megatron-FSDP. ' + 'This allows greater overlap between parameter all-gather operations and ' + 'forward computation, at the cost of additional communication calls. ' + 'For MXFP8, this helps save memory during fine-grained activation ' + 'recomputation, because MXFP8 forward and backward passes use different ' + 'parameter representations (rowwise data for forward, colwise data for ' + 'backward). Only the rowwise parameters of modules involved in ' + 'recomputation will be unsharded.' + ), + ) - group.add_argument("--megatron-fsdp-max-pool-double-buffer", action='store_true', - help="When using Megatron-FSDP double buffering, use the MaxPoolAllocator instead of " - "the FixedPoolAllocator to support asymmetrical FSDP unit configurations. Will " - "increase memory overhead to recycle buffers that fit all FSDP units. Enables " - "NCCL user buffer registration and CUDA graph replay for mixed-arch models.") - group.add_argument("--fsdp-db-use-persist-buf-on-alloc-fail", action='store_true', - help="When using Megatron-FSDP double buffering, persist non-unit modules that " - "are not included in the symmetric buffer pool. May be necessary for NCCL " - "UBR or CUDA Graphs on hybrid architectures.") + group.add_argument( + "--megatron-fsdp-max-pool-double-buffer", + action='store_true', + help="When using Megatron-FSDP double buffering, use the MaxPoolAllocator instead of " + "the FixedPoolAllocator to support asymmetrical FSDP unit configurations. Will " + "increase memory overhead to recycle buffers that fit all FSDP units. Enables " + "NCCL user buffer registration and CUDA graph replay for mixed-arch models.", + ) + group.add_argument( + "--fsdp-db-use-persist-buf-on-alloc-fail", + action='store_true', + help="When using Megatron-FSDP double buffering, persist non-unit modules that " + "are not included in the symmetric buffer pool. May be necessary for NCCL " + "UBR or CUDA Graphs on hybrid architectures.", + ) return parser def _add_msc_args(parser): group = parser.add_argument_group(title="msc") - group.add_argument('--enable-msc', default=False, action='store_true', dest='enable_msc', - help='Enable the usage of Multi-Storage Client (MSC) in Megatron Core. ' - 'Disabled by default; pass this flag to opt in.') - group.add_argument('--disable-msc', default=False, action='store_true', - dest='disable_msc_deprecated', - help='[DEPRECATED] MSC is disabled by default; this flag is a no-op ' - 'and will be removed in a future release.') + group.add_argument( + '--enable-msc', + default=False, + action='store_true', + dest='enable_msc', + help='Enable the usage of Multi-Storage Client (MSC) in Megatron Core. ' + 'Disabled by default; pass this flag to opt in.', + ) + group.add_argument( + '--disable-msc', + default=False, + action='store_true', + dest='disable_msc_deprecated', + help='[DEPRECATED] MSC is disabled by default; this flag is a no-op ' + 'and will be removed in a future release.', + ) return parser + def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): """Add quant-specific arguments to the main parser @@ -3854,70 +5273,155 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): ) return parser + def _add_sft_args(parser): group = parser.add_argument_group(title='sft') group.add_argument('--sft', action="store_true", help='Megatron SFT training') - group.add_argument('--sft-tokenizer-prompt-format', type=str, default="nemotron-h-aligned", - help='SFT prompt format.') - group.add_argument('--sft-mock-dataset-config-json', type=str, default=None, - help='This config provides the necessary information for the mock ' - 'dataset. Accepts either an inline JSON literal or a path to a JSON ' - 'file containing the same schema. You can either specify a CSV file ' - 'that contains sequence lengths, where each line stores the length of ' - 'a sequence, for example: {"mode":"file","path":"/path/to/file"}. ' - 'Alternatively, you can specify a distribution (currently only ' - 'supporting lognormal distribution) along with the required ' - 'parameters, for example, {"mode":"distribution","type":"lognormal",' - '"min_seq_len":1024,"max_seq_len":2048,"mean_seq_len":1536,' - '"lognormal_sigma":1.1}, where sigma controls the variability of the ' - 'lognormal distribution. If not specified and --mock-data is set, ' - 'defaults to a lognormal distribution with min_seq_len=seq_length//2, ' - 'max_seq_len=seq_length, mean_seq_len=seq_length*3//4, ' - 'lognormal_sigma=1.1.') + group.add_argument( + '--sft-tokenizer-prompt-format', + type=str, + default="nemotron-h-aligned", + help='SFT prompt format.', + ) + group.add_argument( + '--sft-mock-dataset-config-json', + type=str, + default=None, + help='This config provides the necessary information for the mock dataset. ' + 'Accepts either an inline JSON literal or a path to a JSON file containing ' + 'the same schema. You can either specify a CSV file that contains sequence lengths, ' + 'where each line stores the length of a sequence, for example: ' + '{"mode":"file","path":"/path/to/file"}. Alternatively, you can specify a distribution ' + '(currently only supporting lognormal distribution) along with the required parameters, ' + 'for example, {"mode":"distribution","type":"lognormal","min_seq_len":1024,' + '"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, where sigma controls ' + 'the variability of the lognormal distribution. ' + 'If not specified and --mock-data is set, defaults to a lognormal distribution with ' + 'min_seq_len=seq_length//2, max_seq_len=seq_length, mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.', + ) + return parser + + +def _add_varlen_dataset_args(parser): + group = parser.add_argument_group(title='varlen dataset') + group.add_argument( + '--use-varlen-dataset', + action="store_true", + help='Train with VarlenDataset, a variable-length packed (THD) dataset ' + 'that consumes instruction-tuning data from a HuggingFace Hub repo id, ' + 'a local parquet file, or a local jsonl file. Schema (alpaca / sharegpt ' + '/ openai-messages) is auto-detected from the dataset columns. ' + 'Mutually exclusive with --sft. Auto-picks a sequence packing ' + 'scheduler when none is given: ``dp_balanced`` by default, ' + '``default_dynamic_cp`` when ``--dynamic-context-parallel`` is set. ' + 'Combine with --mock-data for a synthetic lognormal sequence-length ' + 'distribution; see --varlen-mock-dataset-config-json.', + ) + group.add_argument( + '--varlen-sbhd-validation', + action="store_true", + help='Reference SBHD mode for THD numerical verification. When set, ' + 'VarlenDataset emits SBHD-style samples right-padded to ' + '--seq-length (no cu_seqlens, no packing scheduler), so the run can ' + 'be compared against the THD path to validate correctness. ' + 'Incompatible with --dynamic-context-parallel and ' + '--sequence-packing-scheduler.', + ) + group.add_argument( + '--varlen-mock-dataset-config-json', + type=str, + default=None, + help='Mock-dataset config for --use-varlen-dataset --mock-data. ' + 'Accepts either an inline JSON literal or a path to a JSON file containing ' + 'the same schema as --sft-mock-dataset-config-json: either ' + '{"mode":"file","path":"/path/to/lengths.csv"}, ' + '{"mode":"distribution","type":"lognormal","min_seq_len":1024,' + '"max_seq_len":2048,"mean_seq_len":1536,"lognormal_sigma":1.1}, or ' + '{"mode":"verification","data_path":"/prefix/of/IndexedDataset"}. ' + 'If not specified, defaults to a lognormal distribution with ' + 'min_seq_len=seq_length//2, max_seq_len=seq_length, ' + 'mean_seq_len=seq_length*3//4, lognormal_sigma=1.1.', + ) return parser + def _add_logits_distillation_args(parser): group = parser.add_argument_group(title='Logits Distillation') - group.add_argument('--logits-save-top-k', type=int, default=None, - help='Number of top logits to save.') - group.add_argument('--logits-save-top-p', type=float, default=None, - help='Top-P (nucleus) threshold applied after top-K ' - 'selection when saving logits. Only the smallest ' - 'set of entries whose cumulative probability mass ' - 'reaches this threshold is kept. Must be in (0, 1].') - group.add_argument('--logits-save-top-p-min-k', type=int, default=1, - help='Minimum number of entries kept per token when ' - 'top-P masking is active, regardless of ' - 'cumulative mass. Default: 1.') - group.add_argument('--logits-save-dir', type=str, default=None, - help='Directory to save logits.') - group.add_argument('--logits-save-dtype', type=str, default='fp16', - choices=['fp16', 'bf16', 'fp32'], - help='Dtype for on-disk top-K log-probabilities.') - group.add_argument('--logits-load-dir', type=str, default=None, - help='Directory to load logits.') - group.add_argument('--logits-load-decode-threads', type=int, default=4, - help='Number of decode threads for cached-logits zstd ' - 'decompression and torch.load processing.') - group.add_argument('--logits-load-prefetch-factor', type=int, default=3, - help='PyTorch DataLoader prefetch factor for decoded ' - 'cached-logits iterations. (Non-MSC only)') - group.add_argument('--logits-load-msc-prefetch-depth', type=int, default=2, - help='For MSC/object-storage logits tar shards, number ' - 'of whole tar shards to prefetch into the MSC ' - 'cache ahead of sequential tar consumption.') - group.add_argument('--logits-load-kd-loss-alpha', type=float, default=1.0, - help='KD loss alpha for loading logits. Total loss is calculated as ' - 'alpha * kd_loss + (1 - alpha) * lm_loss.') - group.add_argument('--logits-load-ignore-errors', action='store_true', - default=False, - help='When set, KD loss errors are logged as warnings and ' - 'training falls back to LM-only loss instead of crashing.') + group.add_argument( + '--logits-save-top-k', type=int, default=None, help='Number of top logits to save.' + ) + group.add_argument( + '--logits-save-top-p', + type=float, + default=None, + help='Top-P (nucleus) threshold applied after top-K ' + 'selection when saving logits. Only the smallest ' + 'set of entries whose cumulative probability mass ' + 'reaches this threshold is kept. Must be in (0, 1].', + ) + group.add_argument( + '--logits-save-top-p-min-k', + type=int, + default=1, + help='Minimum number of entries kept per token when ' + 'top-P masking is active, regardless of ' + 'cumulative mass. Default: 1.', + ) + group.add_argument( + '--logits-save-dir', type=str, default=None, help='Directory to save logits.' + ) + group.add_argument( + '--logits-save-dtype', + type=str, + default='fp16', + choices=['fp16', 'bf16', 'fp32'], + help='Dtype for on-disk top-K log-probabilities.', + ) + group.add_argument( + '--logits-load-dir', type=str, default=None, help='Directory to load logits.' + ) + group.add_argument( + '--logits-load-decode-threads', + type=int, + default=4, + help='Number of decode threads for cached-logits zstd ' + 'decompression and torch.load processing.', + ) + group.add_argument( + '--logits-load-prefetch-factor', + type=int, + default=3, + help='PyTorch DataLoader prefetch factor for decoded ' + 'cached-logits iterations. (Non-MSC only)', + ) + group.add_argument( + '--logits-load-msc-prefetch-depth', + type=int, + default=2, + help='For MSC/object-storage logits tar shards, number ' + 'of whole tar shards to prefetch into the MSC ' + 'cache ahead of sequential tar consumption.', + ) + group.add_argument( + '--logits-load-kd-loss-alpha', + type=float, + default=1.0, + help='KD loss alpha for loading logits. Total loss is calculated as ' + 'alpha * kd_loss + (1 - alpha) * lm_loss.', + ) + group.add_argument( + '--logits-load-ignore-errors', + action='store_true', + default=False, + help='When set, KD loss errors are logged as warnings and ' + 'training falls back to LM-only loss instead of crashing.', + ) return parser def _add_fault_injector_args(parser): from megatron.training.config import FaultInjectorConfig + ArgumentGroupFactory(FaultInjectorConfig).build_group(parser, "fault injector") return parser diff --git a/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/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index 322f1d78911..42ec59b6cfa 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -1,17 +1,13 @@ -# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import pytest import torch -from megatron.core.extensions.transformer_engine import TEDotProductAttention -from megatron.core.models.hybrid.hybrid_block import HybridStack +from megatron.core.models.hybrid.hybrid_block import HybridStack, HyperConnectionHybridLayer from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols, validate_segment_layers -from megatron.core.models.hybrid.hybrid_layer_specs import ( - hybrid_inference_stack_spec, - hybrid_stack_spec, -) +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.gated_delta_net import GatedDeltaNet +from megatron.core.ssm.gated_delta_net import HAVE_FLA_KDA, GatedDeltaNet, KimiDeltaAttention from megatron.core.ssm.mamba_layer import MambaLayer from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig @@ -21,7 +17,6 @@ ) from megatron.core.transformer.experimental_attention_variant.dsa import DSAttention from megatron.core.transformer.mlp import MLP -from megatron.core.transformer.multi_latent_attention import MLASelfAttention from megatron.core.transformer.transformer_config import MLATransformerConfig from megatron.core.transformer.transformer_layer import TransformerLayer from tests.unit_tests.test_utilities import Utils @@ -37,8 +32,13 @@ def setup_method(self, method): def get_pg_collection(self): return ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'pp', 'cp']) - def get_hybrid_block(self, layer_pattern, **config_kwargs): + def get_mamba_block(self, layer_pattern, enable_hyper_connections=False): layer_type_list = validate_segment_layers(layer_pattern) + mhc_kwargs = ( + {"enable_hyper_connections": True, "hidden_dropout": 0.0, "mhc_sinkhorn_iterations": 5} + if enable_hyper_connections + else {} + ) transformer_config = TransformerConfig( hidden_size=256, # The Mamba layer places several constraints on this # Need to specify num_attention_heads and num_layers or TransformerConfig @@ -46,7 +46,7 @@ def get_hybrid_block(self, layer_pattern, **config_kwargs): num_layers=len(layer_type_list), num_attention_heads=4, use_cpu_initialization=True, - **config_kwargs, + **mhc_kwargs, ) modules = hybrid_stack_spec.submodules return HybridStack( @@ -57,29 +57,16 @@ def get_hybrid_block(self, layer_pattern, **config_kwargs): pg_collection=self.get_pg_collection(), ) - def get_dsa_hybrid_block(self, layer_pattern): + def get_hybrid_block(self, layer_pattern, **config_kwargs): layer_type_list = validate_segment_layers(layer_pattern) - transformer_config = MLATransformerConfig( + transformer_config = TransformerConfig( hidden_size=256, # The Mamba layer places several constraints on this # Need to specify num_attention_heads and num_layers or TransformerConfig # will generate errors. num_layers=len(layer_type_list), - num_attention_heads=16, + num_attention_heads=4, use_cpu_initialization=True, - bf16=True, - params_dtype=torch.bfloat16, - q_lora_rank=64, - kv_lora_rank=64, - qk_head_dim=64, - qk_pos_emb_head_dim=32, - v_head_dim=64, - rope_type='rope', - rotary_base=10000, - rotary_percent=1.0, - dsa_indexer_n_heads=8, - dsa_indexer_head_dim=64, - dsa_indexer_topk=32, - add_bias_linear=False, + **config_kwargs, ) modules = hybrid_stack_spec.submodules return HybridStack( @@ -90,8 +77,13 @@ def get_dsa_hybrid_block(self, layer_pattern): pg_collection=self.get_pg_collection(), ) - def get_mla_hybrid_block(self, layer_pattern): + def get_dsa_mamba_block(self, layer_pattern, enable_hyper_connections=False): layer_type_list = validate_segment_layers(layer_pattern) + mhc_kwargs = ( + {"enable_hyper_connections": True, "hidden_dropout": 0.0, "mhc_sinkhorn_iterations": 5} + if enable_hyper_connections + else {} + ) transformer_config = MLATransformerConfig( hidden_size=256, # The Mamba layer places several constraints on this # Need to specify num_attention_heads and num_layers or TransformerConfig @@ -109,6 +101,11 @@ def get_mla_hybrid_block(self, layer_pattern): rope_type='rope', rotary_base=10000, rotary_percent=1.0, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=32, + add_bias_linear=False, + **mhc_kwargs, ) modules = hybrid_stack_spec.submodules return HybridStack( @@ -246,8 +243,163 @@ def test_layer_types(self): assert isinstance(layers[2], TransformerLayer) assert isinstance(layers[2].mlp, MLP) + def test_hyper_connection_layer_wrappers(self): + """mHC wraps each hybrid layer while preserving the layer type underneath.""" + layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP + block = self.get_mamba_block(layer_pattern, enable_hyper_connections=True) + layers = block.layers + assert all(isinstance(layer, HyperConnectionHybridLayer) for layer in layers) + assert isinstance(layers[0].inner_layer, MambaLayer) + assert isinstance(layers[1].inner_layer, TransformerLayer) + assert isinstance(layers[1].inner_layer.self_attention, SelfAttention) + assert isinstance(layers[2].inner_layer, TransformerLayer) + assert isinstance(layers[2].inner_layer.mlp, MLP) + + def test_hyper_connection_recompute_plan_for_hybrid_layers(self): + """HybridStack creates per-layer mHC recompute managers when requested.""" + layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP + layer_type_list = validate_segment_layers(layer_pattern) + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=len(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + hidden_dropout=0.0, + mhc_sinkhorn_iterations=5, + recompute_granularity="selective", + recompute_modules=["core_attn", "mhc"], + ) + block = HybridStack( + transformer_config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + pg_collection=self.get_pg_collection(), + ) + + managers, block_ends = block._build_mhc_recompute_layer_plan(use_mhc_recompute=True) + assert len(managers) == len(block.layers) + assert all(manager is not None for manager in managers) + assert block_ends[-1] is True + + def test_hyper_connection_gpu_forward(self): + """mHC-enabled HybridStack expands internally and contracts back at the output.""" + layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP + block = self.get_mamba_block(layer_pattern, enable_hyper_connections=True) + block.cuda() + micro_batch_size = 2 + sequence_length = 32 + hidden_states = torch.ones((sequence_length, micro_batch_size, block.config.hidden_size)) + hidden_states = hidden_states.cuda() + attention_mask = torch.ones( + (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool + ) + attention_mask = attention_mask.cuda() + output = block(hidden_states, attention_mask=attention_mask) + assert output.shape[0] == sequence_length + assert output.shape[1] == micro_batch_size + assert output.shape[2] == block.config.hidden_size + assert output.dtype == torch.float32 + + def test_hyper_connection_gdn_gpu_forward(self): + """mHC runs through GDN, attention, and Mamba hybrid layers.""" + layer_pattern = Symbols.GDN + Symbols.ATTENTION + Symbols.MAMBA + layer_type_list = validate_segment_layers(layer_pattern) + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=len(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + activation_func=torch.nn.functional.silu, + enable_hyper_connections=True, + hidden_dropout=0.0, + mhc_sinkhorn_iterations=5, + ) + block = HybridStack( + transformer_config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + pg_collection=self.get_pg_collection(), + ) + block.cuda() + micro_batch_size = 2 + sequence_length = 32 + hidden_states = torch.ones((sequence_length, micro_batch_size, block.config.hidden_size)) + hidden_states = hidden_states.cuda() + attention_mask = torch.ones( + (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool + ).cuda() + output = block(hidden_states, attention_mask=attention_mask) + assert output.shape == (sequence_length, micro_batch_size, block.config.hidden_size) + + def test_hyper_connection_dsa_layer_wrappers(self): + """mHC wraps DeepSeek-style DSA and MLP split layers.""" + layer_pattern = Symbols.MAMBA + Symbols.DS_ATTENTION + Symbols.MLP + block = self.get_dsa_mamba_block(layer_pattern, enable_hyper_connections=True) + layers = block.layers + assert all(isinstance(layer, HyperConnectionHybridLayer) for layer in layers) + assert isinstance(layers[0].inner_layer, MambaLayer) + assert isinstance(layers[1].inner_layer, TransformerLayer) + assert isinstance(layers[1].inner_layer.self_attention, AbsorbedMLASelfAttention) + assert isinstance(layers[1].inner_layer.self_attention.core_attention, DSAttention) + assert isinstance(layers[2].inner_layer, TransformerLayer) + assert isinstance(layers[2].inner_layer.mlp, MLP) + + def test_hyper_connection_pipeline_boundary_shapes(self): + """HybridStack keeps n-stream tensors between PP stages and contracts at the end.""" + layer_type_list = validate_segment_layers(Symbols.MAMBA) + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=len(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + hidden_dropout=0.0, + mhc_sinkhorn_iterations=5, + ) + modules = hybrid_stack_spec.submodules + first_stage = HybridStack( + transformer_config, + modules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + post_process=False, + pg_collection=self.get_pg_collection(), + ).cuda() + last_stage = HybridStack( + transformer_config, + modules, + pre_process=False, + layer_type_list=layer_type_list, + pp_layer_offset=1, + post_process=True, + pg_collection=self.get_pg_collection(), + ).cuda() + + micro_batch_size = 2 + sequence_length = 32 + hidden_states = torch.ones( + (sequence_length, micro_batch_size, transformer_config.hidden_size), device='cuda' + ) + attention_mask = torch.ones( + (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool, device='cuda' + ) + + pp_hidden = first_stage(hidden_states, attention_mask=attention_mask) + assert pp_hidden.shape == ( + sequence_length, + micro_batch_size, + transformer_config.hidden_size * transformer_config.num_residual_streams, + ) + + last_stage.set_input_tensor(pp_hidden.detach()) + output = last_stage(hidden_states, attention_mask=attention_mask) + assert output.shape == (sequence_length, micro_batch_size, transformer_config.hidden_size) + def test_invalid_layer_types_cause_failure(self): - invalid_symbol = 'X' + invalid_symbol = '+' assert invalid_symbol not in Symbols.VALID_LAYERS # sanity check. layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP + invalid_symbol # validate_segment_layers() in hybrid_layer_allocation.py throws a ValueError. @@ -268,11 +420,20 @@ def test_gdn_layer_types(self): assert isinstance(layers[1].self_attention, SelfAttention) assert isinstance(layers[2], MambaLayer) - def test_gdn_inference_spec(self): - """The inference stack must materialize GDN rather than its IdentityOp default.""" - gdn_spec = hybrid_inference_stack_spec.submodules.gdn_layer - assert gdn_spec.module is TransformerLayer - assert gdn_spec.submodules.self_attention.module is GatedDeltaNet + @pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") + def test_kda_layer_type(self): + """K builds a TransformerLayer wrapping KimiDeltaAttention.""" + block = self.get_hybrid_block( + Symbols.KDA, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=4, + activation_func=torch.nn.functional.silu, + add_bias_linear=False, + ) + assert isinstance(block.layers[0], TransformerLayer) + assert isinstance(block.layers[0].self_attention, KimiDeltaAttention) def test_gdn_gpu_forward(self): """Test GPU forward pass with GDN, attention, and Mamba layers.""" @@ -311,7 +472,7 @@ def test_gdn_gpu_forward(self): def test_dsa_layer_types(self): """D symbol creates a TransformerLayer with absorbed MLA and DSA core attention.""" layer_pattern = Symbols.MAMBA + Symbols.DS_ATTENTION + Symbols.MAMBA - block = self.get_dsa_hybrid_block(layer_pattern) + block = self.get_dsa_mamba_block(layer_pattern) layers = block.layers assert isinstance(layers[0], MambaLayer) assert isinstance(layers[1], TransformerLayer) @@ -323,22 +484,4 @@ def test_mixed_attention_and_dsa_layer_types(self): """* and D in the same block fail.""" layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.DS_ATTENTION + Symbols.MAMBA with pytest.raises(ValueError): - block = self.get_dsa_hybrid_block(layer_pattern) - - def test_mla_layer_types(self): - """+ symbol creates a TransformerLayer with MLASelfAttention but - standard (non-DSA) core attention.""" - layer_pattern = Symbols.MAMBA + Symbols.MLA + Symbols.MAMBA - block = self.get_mla_hybrid_block(layer_pattern) - layers = block.layers - assert isinstance(layers[0], MambaLayer) - assert isinstance(layers[1], TransformerLayer) - assert isinstance(layers[1].self_attention, MLASelfAttention) - assert isinstance(layers[1].self_attention.core_attention, TEDotProductAttention) - assert isinstance(layers[2], MambaLayer) - - def test_mixed_attention_and_mla_layer_types(self): - """* and + in the same block fail (same reason as * and D).""" - layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLA + Symbols.MAMBA - with pytest.raises(ValueError): - block = self.get_mla_hybrid_block(layer_pattern) + block = self.get_dsa_mamba_block(layer_pattern) diff --git a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py index 8b4c181ee30..e33fe34082e 100644 --- a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py +++ b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py @@ -78,7 +78,8 @@ 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']), - ("M+M+", ['M', '+', 'M', '+']), + ("K+K", ['K', '+', 'K']), + ("M+M+", ['M', '+', 'M', '+']), ] for pattern, expected in test_cases: result = validate_segment_layers(pattern) @@ -102,11 +103,20 @@ def test_invalid_symbols_cause_failure(self): with pytest.raises(ValueError): # Not allowed to have both standard Attention and MLA/DSA validate_segment_layers("MDM*-") + + def test_window_symbol(self): + """'W' (sliding-window-only DSv4 attention) is a first-class MLA layer symbol.""" + assert Symbols.WINDOW == "W" + assert Symbols.WINDOW in Symbols.VALID_LAYERS + assert Symbols.WINDOW in Symbols.MLA_ATTENTION + # Parses and is counted like any other layer. + assert validate_segment_layers("WEWE") == ["W", "E", "W", "E"] + assert get_hybrid_layer_counts("WEWE")["W"] == 2 + # 'W' coexists with the other MLA-family attentions (C/H/D)... + validate_segment_layers("WEDECEHE") + # ...but not with standard '*' attention. with pytest.raises(ValueError): - # Not allowed to have both standard Attention and MLA (same reason - # as DSA: * uses the model-level rotary_pos_emb while + uses MLA's - # own decoupled RoPE). - validate_segment_layers("M+M*-") + validate_segment_layers("MWM*-") @pytest.mark.internal @@ -169,8 +179,6 @@ def test_main_pattern_only(self): ("GEGEGE*E", "GEGEGE*E"), ("MDMD", "MDMD"), ("DM", "DM"), - ("M+M+", "M+M+"), - ("+M", "+M"), ] for pattern, expected_main in test_cases: result = parse_hybrid_pattern(pattern) @@ -295,8 +303,6 @@ def test_complex_patterns(self): ("GEGEGE*E/GG/GG", "GEGEGE*E", "GG", 2), # DSA in main pattern with MTP ("MDMD/MD/MD", "MDMD", "MD", 2), - # MLA in main pattern with MTP - ("M+M+/M+/M+", "M+M+", "M+", 2), ] for pattern, expected_main, expected_mtp, expected_depths in test_cases: result = parse_hybrid_pattern(pattern) @@ -316,11 +322,14 @@ class TestGetHybridLayerCounts: def test_simple_pattern(self): assert get_hybrid_layer_counts("M*M*") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 2, 'D': 0, 'G': 0, + 'K': 0, 'M': 2, - '+': 0, '-': 0, 'E': 0, } @@ -328,50 +337,54 @@ def test_simple_pattern(self): def test_all_layer_types(self): # Not allowed to have both standard Attention and MLA/DSA, so we do separate asserts. assert get_hybrid_layer_counts("MG*-E") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 1, 'D': 0, 'G': 1, + 'K': 0, 'M': 1, - '+': 0, '-': 1, 'E': 1, } assert get_hybrid_layer_counts("MGD-E") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 1, 'G': 1, + 'K': 0, 'M': 1, '+': 0, '-': 1, 'E': 1, } - assert get_hybrid_layer_counts("MG+-E") == { - '*': 0, - 'D': 0, - 'G': 1, - 'M': 1, - '+': 1, - '-': 1, - 'E': 1, - } def test_with_pipes(self): # Pipes should be skipped in counting assert get_hybrid_layer_counts("M*|M*") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 2, 'D': 0, 'G': 0, + 'K': 0, 'M': 2, - '+': 0, '-': 0, 'E': 0, } assert get_hybrid_layer_counts("M-M-|M-M*-") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 1, 'D': 0, 'G': 0, + 'K': 0, 'M': 4, - '+': 0, '-': 4, 'E': 0, } @@ -379,11 +392,14 @@ def test_with_pipes(self): def test_with_mtp(self): # MTP pattern "MM" repeated 2 depths -> 4 extra mamba layers assert get_hybrid_layer_counts("M*M*/MM/MM") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 2, 'D': 0, 'G': 0, + 'K': 0, 'M': 6, - '+': 0, '-': 0, 'E': 0, } @@ -392,22 +408,28 @@ def test_with_pipes_and_mtp(self): # Main: M-M-|M-M*- -> 1 attn, 4 mamba, 4 mlp # MTP: MM x 2 depths -> +4 mamba assert get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 1, 'D': 0, 'G': 0, + 'K': 0, 'M': 8, - '+': 0, '-': 4, 'E': 0, } def test_moe_pattern(self): assert get_hybrid_layer_counts("MEME") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 0, 'G': 0, + 'K': 0, 'M': 2, - '+': 0, '-': 0, 'E': 2, } @@ -415,22 +437,28 @@ def test_moe_pattern(self): def test_mtp_with_attention(self): # MTP pattern "*M" repeated 3 depths -> 3 attn + 3 mamba from MTP assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 3, 'D': 0, 'G': 0, + 'K': 0, 'M': 7, - '+': 0, '-': 0, 'E': 0, } def test_gdn_pattern(self): assert get_hybrid_layer_counts("GMGM") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 0, 'G': 2, + 'K': 0, 'M': 2, - '+': 0, '-': 0, 'E': 0, } @@ -438,44 +466,70 @@ def test_gdn_pattern(self): def test_gdn_hybrid_pattern(self): # GDN + Mamba + Attention assert get_hybrid_layer_counts("G*GM*") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 2, 'D': 0, 'G': 2, + 'K': 0, 'M': 1, - '+': 0, '-': 0, 'E': 0, } def test_dsa_pattern(self): assert get_hybrid_layer_counts("DMDM") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 2, 'G': 0, + 'K': 0, 'M': 2, - '+': 0, - '-': 0, - 'E': 0, - } + '+': 0, + '-': 0, + 'E': 0, + } def test_mla_pattern(self): assert get_hybrid_layer_counts("+M+M") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 0, 'G': 0, + 'K': 0, 'M': 2, '+': 2, '-': 0, 'E': 0, } + def test_kda_pattern(self): + assert get_hybrid_layer_counts("K+K/K") == { + '*': 0, + 'D': 0, + 'G': 0, + 'K': 3, + 'M': 0, + '+': 1, + '-': 0, + 'E': 0, + } + def test_empty_pattern(self): assert get_hybrid_layer_counts("") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 0, 'G': 0, + 'K': 0, 'M': 0, - '+': 0, '-': 0, 'E': 0, } @@ -761,7 +815,7 @@ def test_standard_layer_types(self): """Standard symbols each produce a single-entry map at local index 0.""" maps = get_layer_maps_from_layer_type_list(["*", "M", "-", "E"]) # We always get all symbols returned, not only those contained in the pattern. - assert len(maps) == 7 + assert len(maps) == len(Symbols.VALID_LAYERS) attention_map, mamba_map, mlp_map, moe_map = operator.itemgetter( Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE )(maps) @@ -769,6 +823,13 @@ def test_standard_layer_types(self): assert mamba_map == {1: 0} assert mlp_map == {2: 0} assert moe_map == {3: 0} + assert maps[Symbols.KDA] == {} + + def test_kda(self): + """KDA layers are tracked independently from MLA and standard attention.""" + maps = get_layer_maps_from_layer_type_list(["K", "+", "K"]) + assert maps[Symbols.KDA] == {0: 0, 2: 1} + assert maps[Symbols.MLA] == {1: 0} def test_dsa(self): """D (DSA) layers are treated as separate layers for KV cache mapping.""" @@ -804,39 +865,3 @@ def test_all_mamba(self): assert mamba_map == {0: 0, 1: 1, 2: 2} assert mlp_map == {} assert moe_map == {} - - def test_mla(self): - """+ (MLA) layers are mapped independently of other attention types.""" - maps = get_layer_maps_from_layer_type_list(["+", "M", "+", "M"]) - attention_map, dsa_map, mamba_map, mla_map, mlp_map, moe_map = operator.itemgetter( - Symbols.ATTENTION, - Symbols.DS_ATTENTION, - Symbols.MAMBA, - Symbols.MLA, - Symbols.MLP, - Symbols.MOE, - )(maps) - assert attention_map == {} - assert dsa_map == {} - assert mla_map == {0: 0, 2: 1} - assert mamba_map == {1: 0, 3: 1} - assert mlp_map == {} - assert moe_map == {} - - def test_mixed_dsa_and_mla(self): - """D and + can coexist (both are MLA-based and use decoupled RoPE).""" - maps = get_layer_maps_from_layer_type_list(["D", "+", "M", "-"]) - attention_map, dsa_map, mamba_map, mla_map, mlp_map, moe_map = operator.itemgetter( - Symbols.ATTENTION, - Symbols.DS_ATTENTION, - Symbols.MAMBA, - Symbols.MLA, - Symbols.MLP, - Symbols.MOE, - )(maps) - assert attention_map == {} - assert dsa_map == {0: 0} - assert mla_map == {1: 0} - assert mamba_map == {2: 0} - assert mlp_map == {3: 0} - assert moe_map == {} diff --git a/tests/unit_tests/ssm/test_kimi_delta_attention.py b/tests/unit_tests/ssm/test_kimi_delta_attention.py new file mode 100644 index 00000000000..5e1bc44c9b9 --- /dev/null +++ b/tests/unit_tests/ssm/test_kimi_delta_attention.py @@ -0,0 +1,383 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. 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, GatedDeltaNet, KimiDeltaAttention +from megatron.core.ssm.gated_delta_net.common import _GDNBase +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 test_kda_and_gdn_are_sibling_gdn_family_variants(): + assert issubclass(GatedDeltaNet, _GDNBase) + assert issubclass(KimiDeltaAttention, _GDNBase) + assert not issubclass(KimiDeltaAttention, GatedDeltaNet) + + +def _make_config( + tp_size: int = 1, + cp_size: int = 1, + sequence_parallel: bool = False, + params_dtype: torch.dtype = torch.bfloat16, +) -> 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, + transformer_impl="transformer_engine", + ) + + +def test_kda_norm_out_recompute_config_accepts_gdn_family_selector(): + config = replace( + _make_config(), recompute_granularity="selective", recompute_modules=["gdn_norm_out"] + ) + + assert config.recompute_modules == ["gdn_norm_out"] + + +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_rejects_unsupported_head_layouts(): + config = _make_config() + config.linear_num_value_heads = 8 + with pytest.raises(ValueError, match="equal key and value head counts"): + KimiDeltaAttention._validate_config(config) + + config = _make_config() + config.linear_value_head_dim = 64 + with pytest.raises(ValueError, match="equal key and value head dimensions"): + KimiDeltaAttention._validate_config(config) + + config = _make_config() + config.context_parallel_size = 3 + with pytest.raises(ValueError, match="tensor parallel size times context parallel size"): + KimiDeltaAttention._validate_config(config) + + +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.") +def test_kda_selective_recompute_norm_out(): + 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 = ["gdn_norm_out"] + + 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.norm_out_checkpoint is None + assert recompute_kda.recompute_norm_out is True + assert recompute_kda.norm_out_checkpoint is not None + 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: + # The checkpointed gate/norm backward can choose a different BF16 reduction order. + # Keep output and input-gradient checks bitwise, while allowing only BF16-scale + # rounding in parameter gradients. + torch.testing.assert_close( + recompute_grads[name], + base_grads[name], + atol=2e-4, + rtol=5e-3, + msg=f"Selective recompute gradient mismatch for {name}", + ) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.internal +@pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") +def test_kda_cp2_packed_with_physical_padding_matches_cp1(): + """Match CP2 packed KDA against CP1 when physical and actual boundaries differ.""" + + sequence_length = 32 + micro_batch_size = 4 + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=1 + ) + try: + packed_seq_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 128] + ) + torch.manual_seed(2026) + hidden_states_sbhd = torch.rand( + (sequence_length, micro_batch_size, 128), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + model_parallel_cuda_manual_seed(2026) + reference_kda = _build_kda(_make_config(cp_size=1)) + reference_state = { + name: tensor.detach().clone() for name, tensor in reference_kda.state_dict().items() + } + reference_input = ( + hidden_states_sbhd.transpose(0, 1) + .contiguous() + .view(-1, 1, hidden_states_sbhd.shape[-1]) + .requires_grad_(True) + ) + reference_output, reference_bias = reference_kda( + reference_input, None, packed_seq_params=packed_seq_params + ) + assert reference_bias is None + reference_output.float().sum().backward() + reference_output_sbhd = ( + reference_output.detach() + .view(micro_batch_size, sequence_length, -1) + .transpose(0, 1) + .contiguous() + ) + reference_input_grad_sbhd = ( + reference_input.grad.detach() + .view(micro_batch_size, sequence_length, -1) + .transpose(0, 1) + .contiguous() + ) + finally: + Utils.destroy_model_parallel() + + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=2 + ) + try: + model_parallel_cuda_manual_seed(2026) + parallel_kda = _build_kda(_make_config(cp_size=2)) + parallel_kda.load_state_dict(reference_state) + cp_group = parallel_state.get_context_parallel_group() + + def get_packed_tensor_on_this_cp_rank(tensor): + local_tensor = get_tensor_on_this_cp_rank(tensor, 0, cp_group) + return local_tensor.transpose(0, 1).contiguous().view(-1, 1, tensor.shape[-1]) + + parallel_input = get_packed_tensor_on_this_cp_rank(hidden_states_sbhd) + parallel_input = parallel_input.detach().requires_grad_(True) + parallel_output, parallel_bias = parallel_kda( + parallel_input, None, packed_seq_params=packed_seq_params + ) + assert parallel_bias is None + parallel_output.float().sum().backward() + + expected_output = get_packed_tensor_on_this_cp_rank(reference_output_sbhd) + expected_input_grad = get_packed_tensor_on_this_cp_rank(reference_input_grad_sbhd) + torch.testing.assert_close(parallel_output, expected_output, atol=1e-2, rtol=1e-2) + grad_diff = (parallel_input.grad - expected_input_grad).float() + relative_l2 = torch.linalg.vector_norm(grad_diff) / torch.linalg.vector_norm( + expected_input_grad.float() + ).clamp_min(1e-12) + assert relative_l2 < _KDA_INPUT_GRAD_RELATIVE_L2 + # The fixed-seed CP comparison differs by at most one observed BF16 bin + # (0.03125) while aggregate relative L2 remains below 0.5%. + torch.testing.assert_close( + parallel_input.grad, expected_input_grad, atol=_KDA_INPUT_GRAD_ATOL, rtol=1e-2 + ) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.parametrize( + ("sequence_packing", "tp", "sp", "cp"), + [(False, 2, True, 1), (True, 1, False, 2), (True, 2, True, 2)], +) +@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): + config = _make_config() + _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/test_argument_utils.py b/tests/unit_tests/test_argument_utils.py index 8ff6719af72..b6492874abb 100644 --- a/tests/unit_tests/test_argument_utils.py +++ b/tests/unit_tests/test_argument_utils.py @@ -1,24 +1,23 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import signal -import sys from argparse import ArgumentError, ArgumentParser, Namespace from dataclasses import dataclass, field from typing import Callable, Literal, Optional, Union from unittest.mock import MagicMock, patch import pytest -import torch from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig from megatron.core.optimizer import OptimizerConfig +from megatron.core.transformer.spec_utils import ModuleSpec from megatron.training.argument_utils import ( ArgumentGroupFactory, TypeInferenceError, - core_transformer_config_from_args, + _normalize_dsv4_hybrid_csa_compress_ratios, + hybrid_config_from_args, pretrain_cfg_container_from_args, ) -from megatron.training.arguments import add_megatron_arguments, parse_args, validate_args from megatron.training.config import PretrainConfigContainer @@ -45,57 +44,6 @@ class DummyConfig: """Setting with enum type to test enum handling""" -@dataclass(init=False) -class CapturingTransformerConfig: - """Minimal config that records kwargs produced by core_transformer_config_from_args.""" - - moe_use_norm_before_up_proj: bool = False - - def __init__(self, **kwargs): - self.__dict__.update(kwargs) - - -def test_moe_norm_flag_reaches_transformer_config(): - """The generated LatentMoE norm flag should populate the model config.""" - parser = ArgumentParser() - add_megatron_arguments(parser) - - default_args = parser.parse_args([]) - assert default_args.moe_use_norm_before_up_proj is False - - enabled_args = parser.parse_args(["--moe-use-norm-before-up-proj"]) - # params_dtype has no CLI flag of its own; validate_args normally derives it. - enabled_args.params_dtype = torch.float32 - config = core_transformer_config_from_args( - enabled_args, config_class=CapturingTransformerConfig - ) - - assert config.moe_use_norm_before_up_proj is True - - -def test_moe_norm_flag_requires_latent_size(monkeypatch): - """validate_args should reject the LatentMoE norm flag without a latent size.""" - monkeypatch.setattr(sys, 'argv', ['test_argument_utils.py']) - args = parse_args() - args.num_layers = 2 - args.hidden_size = 128 - args.num_attention_heads = 4 - args.max_position_embeddings = 1024 - args.seq_length = 1024 - args.micro_batch_size = 1 - # Let validate_args derive a global batch size that is valid for the - # active data-parallel size in distributed unit-test jobs. - args.train_iters = 1 - args.lr = 1e-4 - args.tokenizer_type = 'NullTokenizer' - args.vocab_size = 1024 - args.moe_use_norm_before_up_proj = True - args.moe_latent_size = None - - with pytest.raises(AssertionError, match="--moe-use-norm-before-up-proj requires"): - validate_args(args) - - @dataclass class ConfigWithOptional: """Config with optional fields.""" @@ -709,12 +657,6 @@ def test_handled_unsupported_unions(self): class TestMegatronNetworkArgumentGeneration: """Test Megatron's TransformerConfig-derived argument group.""" - @staticmethod - def _parser() -> ArgumentParser: - from megatron.training.arguments import _add_network_size_args - - return _add_network_size_args(ArgumentParser(exit_on_error=False)) - def test_transformer_callback_fields_are_not_registered_as_cli_args(self): """Callback fields are runtime hooks, not CLI-provided values.""" from megatron.training.arguments import _add_network_size_args @@ -738,39 +680,112 @@ def test_transformer_callback_fields_are_not_registered_as_cli_args(self): for field_name in callback_fields: assert not hasattr(args, field_name) - def test_mhc_fused_backend_is_exposed_as_config_choice(self): - assert self._parser().parse_args([]).mhc_fused_backend == "auto" - - args = self._parser().parse_args(["--mhc-fused-backend", "native"]) + def test_hybrid_layer_pattern_help_lists_every_public_symbol(self): + """The public help must describe every symbol accepted by the pattern parser.""" + from megatron.training.arguments import _add_experimental_args + + parser = _add_experimental_args(ArgumentParser()) + action = next(action for action in parser._actions if action.dest == "hybrid_layer_pattern") + for symbol_and_name in ( + "M (mamba)", + "G (gdn)", + "K (kda)", + "* (attention)", + "D (dsa)", + "+ (mla)", + "- (mlp)", + "E (moe)", + ): + assert symbol_and_name in action.help + + +class TestDsv4HybridCsaCompressRatioNormalization: + """Test the shared DSv4 hybrid compression-ratio normalization contract.""" + + @pytest.mark.parametrize( + ("provided", "expected_config_ratios"), + [ + (None, [0, 0, 0, 4, 128, 0]), + ([0, 4, 128], [0, 0, 0, 4, 128, 0]), + ([0, 0, 0, 4, 128, 0], [0, 0, 0, 4, 128, 0]), + ], + ) + def test_normalizes_default_compact_and_padded_ratios(self, provided, expected_config_ratios): + args = Namespace(experimental_attention_variant='dsv4_hybrid', csa_compress_ratios=provided) + kw_args = {} + + _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, "-W|EC/H-") + + assert args.csa_compress_ratios == [0, 4, 128] + assert kw_args['csa_compress_ratios'] == expected_config_ratios + + @pytest.mark.parametrize( + ("provided", "message"), + [ + ([0, 8, 128], "ratio 8.*symbol 'C'.*expected 4"), + ([1, 0, 0, 4, 128, 0], "non-attention hybrid symbol '-'.*non-zero ratio 1"), + ([0, 4], r"length \(2\).*W/C/H attention symbols \(3\)"), + ], + ) + def test_rejects_invalid_ratios(self, provided, message): + args = Namespace(experimental_attention_variant='dsv4_hybrid', csa_compress_ratios=provided) - assert args.mhc_fused_backend == "native" + with pytest.raises(AssertionError, match=message): + _normalize_dsv4_hybrid_csa_compress_ratios(args, {}, "-W|EC/H-") - def test_mhc_fused_backend_rejects_unknown_choice(self): - with pytest.raises(ArgumentError, match="invalid choice"): - self._parser().parse_args(["--mhc-fused-backend", "cuda"]) +class TestHybridConfigFromArgs: + """Test static and config-aware hybrid stack spec resolution.""" -class TestMegatronMixedPrecisionArguments: - """Test language-model logit dtype CLI choices.""" + @staticmethod + def _args(): + return Namespace( + spec=["test_module", "test_spec"], + fp16_lm_cross_entropy=False, + hybrid_layer_pattern="M", + position_embedding_type="none", + rotary_percent=1.0, + rotary_base=10000, + make_vocab_size_divisible_by=128, + rotary_seq_len_interpolation_factor=None, + max_position_embeddings=1024, + untie_embeddings_and_output_weights=False, + padded_vocab_size=128, + ) @staticmethod - def _parser() -> ArgumentParser: - from megatron.training.arguments import _add_mixed_precision_args + def _transformer_config(): + config = MagicMock() + config.transformer_impl = "transformer_engine" + config.inference_fuse_tp_communication = False + return config + + @patch( + "megatron.training.argument_utils.HybridModelConfig", side_effect=lambda **kwargs: kwargs + ) + @patch("megatron.training.argument_utils.import_module") + def test_preserves_static_module_spec(self, mock_import_module, _mock_model_config): + static_spec = ModuleSpec(module=object) + mock_import_module.return_value = static_spec - return _add_mixed_precision_args(ArgumentParser(exit_on_error=False)) + config = hybrid_config_from_args(self._args(), config=self._transformer_config()) - def test_logit_dtype_defaults_to_input_dtype(self): - args = self._parser().parse_args([]) - assert args.logit_dtype is None + assert config["hybrid_stack_spec"] is static_spec - @pytest.mark.parametrize("dtype", ["bf16", "fp32"]) - def test_logit_dtype_accepts_supported_choices(self, dtype): - args = self._parser().parse_args(["--output-logit-dtype", dtype]) - assert args.logit_dtype == dtype + @patch( + "megatron.training.argument_utils.HybridModelConfig", side_effect=lambda **kwargs: kwargs + ) + @patch("megatron.training.argument_utils.import_module") + def test_resolves_config_aware_spec_factory(self, mock_import_module, _mock_model_config): + static_spec = ModuleSpec(module=object) + spec_factory = MagicMock(return_value=static_spec) + mock_import_module.return_value = spec_factory + transformer_config = self._transformer_config() - def test_logit_dtype_rejects_fp16(self): - with pytest.raises(ArgumentError, match="invalid choice"): - self._parser().parse_args(["--output-logit-dtype", "fp16"]) + config = hybrid_config_from_args(self._args(), config=transformer_config) + + spec_factory.assert_called_once_with(transformer_config) + assert config["hybrid_stack_spec"] is static_spec # --------------------------------------------------------------------------- @@ -802,6 +817,29 @@ def _make_args(**overrides): return Namespace(**defaults) +def _make_ddp_args(**overrides): + """Build the minimum Namespace required by get_megatron_ddp_config.""" + defaults = { + "accumulate_allreduce_grads_in_fp32": False, + "check_for_nan_in_loss_and_grad": False, + "check_for_large_grads": False, + "ddp_num_buckets": None, + "ddp_bucket_size": None, + "ddp_pad_buckets_for_high_nccl_busbw": False, + "ddp_reduce_scatter_with_fp32_accumulation": False, + "ddp_param_name_patterns_for_fp32_local_accumulation": [], + "ddp_average_in_collective": False, + "megatron_fsdp_main_params_dtype": None, + "megatron_fsdp_main_grads_dtype": None, + "megatron_fsdp_grad_comm_dtype": None, + "use_precision_aware_optimizer": False, + "use_megatron_fsdp": False, + "cuda_graph_impl": "none", + } + defaults.update(overrides) + return Namespace(**defaults) + + @pytest.fixture def mock_optimizer_config(): return MagicMock(spec=OptimizerConfig) @@ -857,6 +895,37 @@ def test_ddp_config_comes_from_helper(self, mock_opt, mock_ddp): assert result.ddp is mock_ddp_instance +class TestDDPConfigMapping: + """Test argparse Namespace → DistributedDataParallelConfig mapping.""" + + def test_cpu_backup_disable_fields_map_from_args(self): + from megatron.training.training import get_megatron_ddp_config + + args = _make_ddp_args( + disable_grad_buffers_cpu_backup=True, disable_param_buffers_cpu_backup=True + ) + + config = get_megatron_ddp_config(args) + + assert config.disable_grad_buffers_cpu_backup is True + assert config.disable_param_buffers_cpu_backup is True + + def test_cpu_backup_disable_cli_flags_map_to_ddp_config(self): + from megatron.training.arguments import _add_distributed_args + from megatron.training.training import get_megatron_ddp_config + + parser = _add_distributed_args(ArgumentParser()) + parsed_args = parser.parse_args( + ["--disable-grad-buffers-cpu-backup", "--disable-param-buffers-cpu-backup"] + ) + args = _make_ddp_args(**vars(parsed_args)) + + config = get_megatron_ddp_config(args) + + assert config.disable_grad_buffers_cpu_backup is True + assert config.disable_param_buffers_cpu_backup is True + + class TestCheckpointConfigMapping: """Test the boolean inversions and custom dest mappings for CheckpointConfig.""" diff --git a/tests/unit_tests/transformer/test_attention.py b/tests/unit_tests/transformer/test_attention.py index e82ae7e458a..9b2d3954a18 100644 --- a/tests/unit_tests/transformer/test_attention.py +++ b/tests/unit_tests/transformer/test_attention.py @@ -1,18 +1,19 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import copy from unittest import mock +import einops import pytest import torch from packaging import version +from torch.nn import functional as F import megatron.core.parallel_state as parallel_state +from megatron.core.datasets.data_schedule_utils import get_cp_slice_for_thd from megatron.core.hyper_comm_grid import HyperCommGrid -from megatron.core.models.common.embeddings.rope_utils import ( - get_pos_emb_on_this_cp_rank as get_tensor_on_this_cp_rank, -) from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_local_spec, get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_with_transformer_engine_submodules, ) @@ -21,6 +22,10 @@ from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig from megatron.core.transformer.attention import SelfAttention +from megatron.core.transformer.dot_product_attention_context_parallel import ( + AttentionFuncionWithContextParallel, + to_zz_mask_attn_bias, +) from megatron.core.transformer.enums import AttnMaskType from megatron.core.utils import is_te_min_version, unwrap_model from megatron.training.arguments import parse_args @@ -44,10 +49,19 @@ @pytest.mark.parametrize("output_gate", [False, True]) +@pytest.mark.parametrize( + ("transformer_impl", "fallback_to_eager_attn"), + [("transformer_engine", False), ("transformer_engine", True), ("native", False)], +) class TestParallelAttention: @pytest.fixture(scope='function', autouse=True) - def setup_method(self, output_gate): + def setup_method(self, output_gate, transformer_impl, fallback_to_eager_attn): + if output_gate: + if transformer_impl == "native": + pytest.skip("Native implementation does not support output gate.") + if fallback_to_eager_attn: + pytest.skip("No need to test output gate for fallback_to_eager_attn = True.") Utils.initialize_model_parallel(1, 1) model_parallel_cuda_manual_seed(123) self.transformer_config = TransformerConfig( @@ -58,11 +72,17 @@ def setup_method(self, output_gate): bf16=True, params_dtype=torch.bfloat16, attention_output_gate=output_gate, + transformer_impl=transformer_impl, + fallback_to_eager_attn=fallback_to_eager_attn, ) + if transformer_impl == "transformer_engine": + attn_layer_spec = ( + get_gpt_layer_with_transformer_engine_submodules().self_attention.submodules + ) + else: + attn_layer_spec = get_gpt_layer_local_spec().submodules.self_attention.submodules self.parallel_attention = SelfAttention( - self.transformer_config, - get_gpt_layer_with_transformer_engine_submodules().self_attention.submodules, - layer_number=1, + self.transformer_config, attn_layer_spec, layer_number=1 ) def teardown_method(self): @@ -73,10 +93,19 @@ def test_constructor(self): assert self.parallel_attention.layer_number == 1 num_weights = sum([p.numel() for p in self.parallel_attention.parameters()]) + + hidden_size = self.transformer_config.hidden_size + standard_num_weights = ( + hidden_size * hidden_size * 4 + hidden_size * 4 # QKVO weight # QKVO bias + ) if self.transformer_config.attention_output_gate: - assert num_weights == 82816 - else: - assert num_weights == 66304 + standard_num_weights += hidden_size * hidden_size + hidden_size # Gate weight and bias + if self.transformer_config.transformer_impl == "transformer_engine": + standard_num_weights += hidden_size * 2 # fused pre layernorm weight and bias + + assert ( + num_weights == standard_num_weights + ), f"{num_weights=} does not match {standard_num_weights=}." def test_cpu_forward(self): # we can't currently do this because the global memory buffer is on GPU @@ -111,6 +140,8 @@ def test_gpu_forward(self): @pytest.mark.parametrize("rotary_interleaved", [True, False]) @pytest.mark.parametrize("fused_qkv_rope", [True, False]) def test_fused_rope_gpu_forward(self, rotary_interleaved, fused_qkv_rope): + if self.transformer_config.fallback_to_eager_attn: + pytest.skip("No need to test fused RoPE for fallback_to_eager_attn = True.") self.parallel_attention.config.apply_rope_fusion = True if rotary_interleaved and not is_te_min_version("2.3.0"): pytest.skip("Only TE >= 2.3.0 supports interleaved fused RoPE.") @@ -374,6 +405,199 @@ class TestSelfAttention: @pytest.fixture(scope='function', autouse=True) def setup_method(self, output_gate): + self.output_gate = output_gate + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self): + Utils.destroy_model_parallel() + + def test_clip_qk_disabled_raises_error(self): + """Test that clip_qk raises ValueError when qk_clip is not enabled.""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + qk_clip=False, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + + with pytest.raises(ValueError, match="qk_clip option needs to be enabled"): + attention.clip_qk() + + def test_clip_qk_none_logits_raises_error(self): + """Test that clip_qk raises ValueError when current_max_attn_logits is None.""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + qk_clip=True, + qk_clip_threshold=100.0, + qk_clip_alpha=0.5, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + + with pytest.raises(ValueError, match="current_max_attn_logits is None"): + attention.clip_qk() + + def test_clip_qk_below_threshold_no_update(self): + """Test that weights are not updated when max logits are below threshold.""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + qk_clip=True, + qk_clip_threshold=100.0, + qk_clip_alpha=0.5, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + attention.cuda() + + # Save original weights + original_weight = attention.linear_qkv.weight.data.clone() + + # Set current_max_attn_logits below threshold + attention.core_attention.current_max_attn_logits = torch.tensor( + [50.0, 60.0, 70.0, 80.0], device='cuda' + ) + + # Call clip_qk + attention.clip_qk() + + # Weights should not be updated + assert torch.equal(attention.linear_qkv.weight.data, original_weight) + # current_max_attn_logits should be reset + assert attention.core_attention.current_max_attn_logits is None + + def test_clip_qk_above_threshold_updates_weights(self): + """Test that weights are updated when max logits exceed threshold.""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + qk_clip=True, + qk_clip_threshold=100.0, + qk_clip_alpha=0.5, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + attention.cuda() + + # Save original weights + original_weight = attention.linear_qkv.weight.data.clone() + + # Set current_max_attn_logits above threshold + attention.core_attention.current_max_attn_logits = torch.tensor( + [150.0, 160.0, 170.0, 180.0], device='cuda' + ) + + # Call clip_qk + attention.clip_qk() + + # Weights should be updated + assert not torch.equal(attention.linear_qkv.weight.data, original_weight) + # current_max_attn_logits should be reset + assert attention.core_attention.current_max_attn_logits is None + + def test_clip_qk_gqa_configuration(self): + """Test clip_qk with GQA (Grouped Query Attention) configuration.""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=8, + num_query_groups=4, # GQA with 2 heads per group + use_cpu_initialization=True, + qk_clip=True, + qk_clip_threshold=100.0, + qk_clip_alpha=0.5, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + attention.cuda() + + # Save original weights + original_weight = attention.linear_qkv.weight.data.clone() + + # Set current_max_attn_logits for all heads (8 heads) + attention.core_attention.current_max_attn_logits = torch.tensor( + [150.0, 160.0, 170.0, 180.0, 190.0, 200.0, 210.0, 220.0], device='cuda' + ) + + # Call clip_qk + attention.clip_qk() + + # Weights should be updated + assert not torch.equal(attention.linear_qkv.weight.data, original_weight) + # current_max_attn_logits should be reset + assert attention.core_attention.current_max_attn_logits is None + + def test_clip_qk_mixed_logits(self): + """Test clip_qk with mixed logits (some above, some below threshold).""" + transformer_config = TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + qk_clip=True, + qk_clip_threshold=100.0, + qk_clip_alpha=0.5, + ) + attention = SelfAttention( + transformer_config, + get_gpt_layer_with_transformer_engine_spec().submodules.self_attention.submodules, + layer_number=1, + ) + attention.cuda() + + # Save original weights + original_weight = attention.linear_qkv.weight.data.clone() + + # Set mixed current_max_attn_logits (some above, some below threshold) + attention.core_attention.current_max_attn_logits = torch.tensor( + [80.0, 150.0, 90.0, 200.0], device='cuda' + ) + + # Call clip_qk + attention.clip_qk() + + # Weights should be updated since at least one head exceeds threshold + assert not torch.equal(attention.linear_qkv.weight.data, original_weight) + # current_max_attn_logits should be reset + assert attention.core_attention.current_max_attn_logits is None + + +@pytest.mark.parametrize("output_gate", [False, True]) +@pytest.mark.parametrize("transformer_impl", ["transformer_engine", "native"]) +class TestSelfAttention: + + @pytest.fixture(scope='function', autouse=True) + def setup_method(self, output_gate, transformer_impl): + if transformer_impl == "native": + if output_gate: + pytest.skip("Native implementation does not support output gate.") + self.transformer_impl = transformer_impl self.output_gate = output_gate Utils.destroy_model_parallel() @@ -389,10 +613,17 @@ def run_self_attention(self, pg_collection): attention_output_gate=self.output_gate, tensor_model_parallel_size=tensor_model_parallel_size, use_cpu_initialization=False, + transformer_impl=self.transformer_impl, ) + if self.transformer_impl == "transformer_engine": + attn_layer_spec = ( + get_gpt_layer_with_transformer_engine_submodules().self_attention.submodules + ) + else: + attn_layer_spec = get_gpt_layer_local_spec().submodules.self_attention.submodules self.self_attention = SelfAttention( self.transformer_config, - get_gpt_layer_with_transformer_engine_submodules().self_attention.submodules, + attn_layer_spec, layer_number=1, attn_mask_type=AttnMaskType.causal, pg_collection=pg_collection, @@ -473,6 +704,8 @@ def _test_parallel_attention_correctness( tmp_path_dist_ckpt, atol, rtol, + cosine_similarity_threshold=None, + relative_l2_threshold=0.1, tp=1, sp=False, cp=1, @@ -480,6 +713,10 @@ def _test_parallel_attention_correctness( sequence_length=256, micro_batch_size=4, sequence_packing=False, + cp_partition_mode="zigzag", + compare_param_grads=False, + input_grad_atol=None, + input_grad_rtol=None, ): # Model initialization function def initialize_gpt_model( @@ -515,7 +752,6 @@ def initialize_gpt_model( with TempNamedDir(tmp_path_dist_ckpt / 'test_parallel_attn', sync=True) as ckpt_dir: # Set argument mock_args = parse_args(ignore_unknown_args=True) - mock_args.save_tokenizer_assets = False set_args(mock_args) # Initialize baseline model @@ -532,8 +768,24 @@ def initialize_gpt_model( mock_args.no_load_rng = True save_checkpoint(10, gpt_model, None, None, 0) + def get_param_grad(param): + grad = param.grad + if grad is None: + grad = getattr(param, "main_grad", None) + if grad is None: + return torch.zeros_like(param, dtype=torch.float32) + return grad + + def zero_param_grads(module): + for param in module.parameters(): + param.grad = None + main_grad = getattr(param, "main_grad", None) + if main_grad is not None: + main_grad.zero_() + # Calculate baseline output attention = gpt_model[0].decoder.layers[0].self_attention + zero_param_grads(attention) output_hidden_states_baseline, bias_hidden_states_baseline = attention( input_hidden_states, attention_mask=None ) @@ -541,6 +793,13 @@ def initialize_gpt_model( # Save baseline output input_grad_baseline = input_hidden_states.grad.detach() + param_grads_baseline = None + if compare_param_grads: + param_grads_baseline = { + name: get_param_grad(param).detach().float().clone() + for name, param in attention.named_parameters() + if param.requires_grad + } output_hidden_states_baseline = output_hidden_states_baseline.detach() bias_hidden_states_baseline = bias_hidden_states_baseline if bias_hidden_states_baseline is not None: @@ -559,9 +818,11 @@ def initialize_gpt_model( transformer_config.context_parallel_size = cp transformer_config.tensor_model_parallel_size = tp transformer_config.sequence_parallel = sp + transformer_config.cp_partition_mode = cp_partition_mode init_basic_mock_args(mock_args, tp, 1, bf16=True) mock_args.context_parallel_size = cp mock_args.sequence_parallel = sp + mock_args.cp_partition_mode = cp_partition_mode gpt_model = unwrap_model(get_model(initialize_gpt_model, config=transformer_config)) with mock.patch('megatron.training.checkpointing.check_checkpoint_args'): with mock.patch('megatron.training.checkpointing.update_num_microbatches'): @@ -572,10 +833,31 @@ def initialize_gpt_model( tp_rank = parallel_state.get_tensor_model_parallel_rank() def get_tensor_on_this_rank(tensor): - if cp > 1: - tensor = get_tensor_on_this_cp_rank(tensor, 0, cp_group) if sequence_packing: tensor = tensor.transpose(0, 1).contiguous().view(-1, 1, *tensor.shape[2:]) + if cp > 1: + cu_seqlens_tensor = torch.tensor( + [i * sequence_length for i in range(micro_batch_size + 1)], + device=tensor.device, + ) + batch = {"hidden_states": tensor, "cu_seqlens_padded": cu_seqlens_tensor} + get_cp_slice_for_thd( + batch, + cp_group, + keys=("hidden_states",), + cp_partition_mode=cp_partition_mode, + ) + tensor = batch["hidden_states"] + elif cp > 1: + cp_rank = torch.distributed.get_rank(cp_group) + if cp_partition_mode == "zigzag": + cp_idx = torch.tensor([cp_rank, 2 * cp - cp_rank - 1], device=tensor.device) + elif cp_partition_mode == "contiguous": + cp_idx = torch.tensor([2 * cp_rank, 2 * cp_rank + 1], device=tensor.device) + else: + raise ValueError(f"Unsupported test CP partition mode {cp_partition_mode!r}.") + tensor = tensor.view(2 * cp, -1, *tensor.shape[1:]) + tensor = tensor.index_select(0, cp_idx).view(-1, *tensor.shape[2:]) if tp > 1 and sp: sp_seg = tensor.shape[0] // tp tensor = tensor[tp_rank * sp_seg : (tp_rank + 1) * sp_seg] @@ -585,15 +867,27 @@ def get_tensor_on_this_rank(tensor): if sequence_packing: cu_seqlens = [i * sequence_length for i in range(micro_batch_size + 1)] packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) + packed_seq_params.cp_partition_mode = cp_partition_mode else: packed_seq_params = None input_hidden_states = get_tensor_on_this_rank(input_hidden_states) input_hidden_states = input_hidden_states.detach().requires_grad_(True) parallel_attention = gpt_model[0].decoder.layers[0].self_attention + zero_param_grads(parallel_attention) output_hidden_states_parallel, bias_hidden_states_parallel = parallel_attention( input_hidden_states, attention_mask=None, packed_seq_params=packed_seq_params ) output_hidden_states_parallel.sum().backward() + param_grads_parallel = None + if compare_param_grads: + assert tp == 1, "Parameter gradient parity only supports unsharded parameters." + param_grads_parallel = {} + for name, param in parallel_attention.named_parameters(): + if param.requires_grad: + grad = get_param_grad(param) + if cp > 1: + torch.distributed.all_reduce(grad, group=cp_group) + param_grads_parallel[name] = grad.detach().float().clone() input_grad_parallel = input_hidden_states.grad.detach() # Check if the output is close @@ -630,32 +924,65 @@ def get_tensor_on_this_rank(tensor): ~torch.isinf(bias_hidden_states_parallel) ), "bias_hidden_states_parallel contains inf" - torch.testing.assert_close( - output_hidden_states_baseline, - output_hidden_states_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in output_hidden_states: {msg}", + def assert_close_or_cosine_similarity( + baseline, parallel, tensor_name, tensor_atol=atol, tensor_rtol=rtol + ): + try: + torch.testing.assert_close( + baseline, + parallel, + atol=tensor_atol, + rtol=tensor_rtol, + msg=lambda msg: f"Mismatch in {tensor_name}: {msg}", + ) + return + except AssertionError as close_error: + if cosine_similarity_threshold is None: + raise close_error + baseline_flat = baseline.flatten().float() + parallel_flat = parallel.flatten().float() + cosine_sim = torch.nn.functional.cosine_similarity( + baseline_flat.unsqueeze(0), parallel_flat.unsqueeze(0) + ).item() + diff_norm = torch.linalg.vector_norm((parallel_flat - baseline_flat).float()) + baseline_norm = torch.linalg.vector_norm(baseline_flat.float()).clamp_min(1e-12) + relative_l2 = (diff_norm / baseline_norm).item() + assert cosine_sim >= cosine_similarity_threshold, ( + f"Mismatch in {tensor_name}: cosine similarity " + f"{cosine_sim} < {cosine_similarity_threshold}, " + f"while assert_close failed: {close_error}" + ) + assert relative_l2 <= relative_l2_threshold, ( + f"Mismatch in {tensor_name}: relative L2 " + f"{relative_l2} > {relative_l2_threshold}, " + f"while assert_close failed: {close_error}" + ) + + assert_close_or_cosine_similarity( + output_hidden_states_baseline, output_hidden_states_parallel, "output_hidden_states" ) - torch.testing.assert_close( + assert_close_or_cosine_similarity( input_grad_baseline, input_grad_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in input_grad: {msg}", + "input_grad", + tensor_atol=atol if input_grad_atol is None else input_grad_atol, + tensor_rtol=rtol if input_grad_rtol is None else input_grad_rtol, ) if has_bias: - torch.testing.assert_close( - bias_hidden_states_baseline, - bias_hidden_states_parallel, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Mismatch in bias_hidden_states: {msg}", + assert_close_or_cosine_similarity( + bias_hidden_states_baseline, bias_hidden_states_parallel, "bias_hidden_states" ) + if compare_param_grads: + assert param_grads_baseline.keys() == param_grads_parallel.keys() + for name, grad_baseline in param_grads_baseline.items(): + assert_close_or_cosine_similarity( + grad_baseline, param_grads_parallel[name], f"param_grad[{name}]" + ) Utils.destroy_model_parallel() +# TODO(yuzhongw): Add test case for fallback_to_eager_attn @pytest.mark.parametrize("sequence_packing", [False, True]) @pytest.mark.parametrize("apply_rope_fusion", [False, True]) @pytest.mark.parametrize( From b382bfe768fa81585fcd842d6cddcb63dd63bae8 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 16 Aug 2026 01:31:14 +0800 Subject: [PATCH 03/34] Honor GDN projection metadata in FSDP checkpoints Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 72cefc335f96c3f1137e34e18586338a9038ebd7) --- .../transformer/fsdp_dtensor_checkpoint.py | 682 +++----- megatron/training/checkpointing.py | 1409 +++++------------ .../mfsdp_v1/test_gdn_fsdp_checkpoint.py | 176 ++ tests/unit_tests/test_checkpointing.py | 381 ++--- .../test_fsdp_dtensor_checkpoint.py | 787 +-------- 5 files changed, 1066 insertions(+), 2369 deletions(-) create mode 100644 tests/unit_tests/distributed/mfsdp_v1/test_gdn_fsdp_checkpoint.py diff --git a/megatron/core/transformer/fsdp_dtensor_checkpoint.py b/megatron/core/transformer/fsdp_dtensor_checkpoint.py index 327a4f4da05..3bf2e4c1c10 100644 --- a/megatron/core/transformer/fsdp_dtensor_checkpoint.py +++ b/megatron/core/transformer/fsdp_dtensor_checkpoint.py @@ -44,33 +44,11 @@ HAVE_MEGATRON_FSDP = False from megatron.core import parallel_state -from megatron.core.dist_checkpointing.core import CheckpointingException -from megatron.core.dist_checkpointing.validation import StrictHandling, parse_strict_flag from megatron.core.tensor_parallel.layers import copy_tensor_model_parallel_attributes from megatron.core.transformer.transformer_layer import TransformerLayer from megatron.core.utils import get_attr_wrapped_model -def _strip_wrapper_prefixes(path): - """Strip DDP/FSDP wrapper prefixes (module., model.) from a module or state dict path.""" - parts = path.split('.') - while parts and parts[0] in ('module', 'model'): - parts = parts[1:] - return '.'.join(parts) - - -def _intersect_slice(s1, s2): - """Intersection of two step-1 slices, or an empty slice when they do not overlap.""" - start = max(s1.start, s2.start) - stop = min(s1.stop, s2.stop) - return slice(0, 0) if start >= stop else slice(start, stop) - - -def _shift_slice(s, offset): - """Move a step-1 slice by ``offset``, e.g. to rebase it onto a shard's local storage.""" - return slice(s.start + offset, s.stop + offset) - - def get_ep_layer_offset(num_experts: int | None = None) -> int: """ Get the expert layer offset for the current model. @@ -237,16 +215,21 @@ def handle_swiglu_in_state_dict(model, model_state_dict, optimizer_state_dict): # layers have gated_linear_unit=False while language decoder layers # have gated_linear_unit=True. # ------------------------------------------------------------------ + def _strip_wrappers(path): + """Strip DDP/FSDP wrapper prefixes (module., model.) from a path.""" + parts = path.split('.') + while parts and parts[0] in ('module', 'model'): + parts = parts[1:] + return '.'.join(parts) + _layer_glu = {} for name, module in model.named_modules(): if isinstance(module, TransformerLayer): - _layer_glu[_strip_wrapper_prefixes(name)] = getattr( - module.config, 'gated_linear_unit', False - ) + _layer_glu[_strip_wrappers(name)] = getattr(module.config, 'gated_linear_unit', False) def _key_in_glu_layer(key): """Return True if *key* belongs to a TransformerLayer with gated_linear_unit=True.""" - norm_key = _strip_wrapper_prefixes(key) + norm_key = _strip_wrappers(key) best_glu, best_len = None, -1 for layer_path, uses_glu in _layer_glu.items(): if norm_key.startswith(layer_path + '.') and len(layer_path) > best_len: @@ -255,6 +238,17 @@ def _key_in_glu_layer(key): return True # no TransformerLayer found — assume GLU for backward compat return best_glu + def intersection(s1, s2): + # Only works for step=1 + start = max(s1.start, s2.start) + stop = min(s1.stop, s2.stop) + if start >= stop: + return slice(0, 0) # Empty slice if no intersection + return slice(start, stop) + + def offset_slice(s, offset): + return slice(s.start + offset, s.stop + offset) + def is_swiglu_key(key): """ Check if this key should be handled as SwiGLU linear_fc1 weight or bias. @@ -280,10 +274,21 @@ def is_swiglu_key(key): def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis, is_expert_param): """ Split the SWiGLU linear_fc1 parameter into two parts: weight_w and weight_v. + + Args: + data: The tensor to split. May be a DTensor (model state dict) or a + plain Tensor (optimizer states from FusedAdam). + dist_param: The corresponding model parameter (always a DTensor). + Used for global shape, numel, FSDP slice, and dist index metadata. + swiglu_shard_axis: Axis along which to split W and V gates. + is_expert_param: Whether this is an expert parameter (affects TP mesh). """ - assert data.shape[swiglu_shard_axis] % 2 == 0, ( - f"SWiGLU weights must have an even size along the shard axis {swiglu_shard_axis}, " - f"got {data.shape[swiglu_shard_axis]}" + # Use dist_param (always a DTensor) for global shape/numel, + # as data may be a regular Tensor (e.g., optimizer states). + global_shape = dist_param.shape + assert global_shape[swiglu_shard_axis] % 2 == 0, ( + f"SWiGLU FC1 must have even global size along axis {swiglu_shard_axis}, " + f"got {global_shape[swiglu_shard_axis]} (global_shape={list(global_shape)})" ) fsdp_slice = dist_param.megatron_fsdp_slice @@ -292,25 +297,39 @@ def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis, is_expert_param tp_mesh = megatron_fsdp_dist_index.get_submesh( [megatron_fsdp_dist_index.tp_dim], is_expert_parallel=is_expert_param ) - data_size = data.numel() // tp_mesh.mesh.numel() + data_size = dist_param.numel() // tp_mesh.mesh.numel() w_slice = slice(0, data_size // 2) v_slice = slice(data_size // 2, data_size) - view_shape = list(data.shape) + view_shape = list(global_shape) view_shape[swiglu_shard_axis] = -1 - local_tensor = data.to_local() + if isinstance(data, DTensor): + assert data.shape == global_shape, ( + f"DTensor shape mismatch: data.shape={data.shape} vs " + f"dist_param.shape={global_shape}" + ) + local_tensor = data.to_local() + else: + # Plain Tensor must already be the FSDP-local shard; other layouts + # (TP-local / global) would silently misalign in the slice below. + expected = fsdp_slice.stop - fsdp_slice.start + assert data.numel() == expected, ( + f"Plain Tensor must be FSDP-local shard " + f"(expected numel={expected}, got {data.numel()}; fsdp_slice={fsdp_slice})" + ) + local_tensor = data weight_w = local_tensor.view(-1)[ - _shift_slice(_intersect_slice(fsdp_slice, w_slice), -fsdp_slice.start) + offset_slice(intersection(fsdp_slice, w_slice), -fsdp_slice.start) ] weight_v = local_tensor.view(-1)[ - _shift_slice(_intersect_slice(fsdp_slice, v_slice), -fsdp_slice.start) + offset_slice(intersection(fsdp_slice, v_slice), -fsdp_slice.start) ] weight_w = weight_w.reshape(view_shape) weight_v = weight_v.reshape(view_shape) # Fake parameters w and v are used to provide the correct parameter # shape and Tensor-Parallelism information. - per_tp_rank_shape = list(data.shape) + per_tp_rank_shape = list(global_shape) if is_mcore_tensor_model_parallel(dist_param): tp_dim = get_mcore_tensor_parallel_partition_dim(dist_param) assert tp_dim is not None, "Tensor model parallel dimension not found" @@ -395,85 +414,206 @@ 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"] + def _strip_wrappers(path): + """Strip DDP/FSDP wrapper prefixes (module., model.) from a path.""" + parts = path.split('.') + while parts and parts[0] in ('module', 'model'): + parts = parts[1:] + return '.'.join(parts) + # ------------------------------------------------------------------ - # Build per-GDN-module split-size map by walking the model tree. - # GDN modules are identified by the presence of qk_dim / v_dim / - # 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_wrappers(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_wrappers(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_wrappers(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``.""" - 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 :] + """Return split metadata when ``key`` names a fused GDN-family parameter.""" + norm = _strip_wrappers(key) + 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 intersection(s1, s2): + start = max(s1.start, s2.start) + stop = min(s1.stop, s2.stop) + return slice(0, 0) if start >= stop else slice(start, stop) + + def offset_slice(s, offset): + return slice(s.start + offset, s.stop + offset) + def split_gdn_fused(data, dist_param, split_sizes, split_dim): - """Split a fused GDN projection DTensor into per-component DTensors.""" + """Split a fused GDN-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: + # GDN tensors are already TP-local here (fast-path from #4799). return list( split_dtensor( data, split_sizes, dim=split_dim, update_uneven_dtensor_chunk_meta=True ) ) + # Use dist_param (always a DTensor) for global shape/numel, + # as data may be a regular Tensor (e.g., optimizer states). + global_shape = dist_param.shape + fsdp_slice = dist_param.megatron_fsdp_slice dist_index = dist_param.megatron_fsdp_dist_index tp_mesh = dist_index.get_submesh([dist_index.tp_dim], is_expert_parallel=False) - data_size = data.numel() // tp_mesh.mesh.numel() + data_size = dist_param.numel() // tp_mesh.mesh.numel() elems_per_unit = data_size // total_split - local_tensor = data.to_local() - view_shape = list(data.shape) + if isinstance(data, DTensor): + assert data.shape == global_shape, ( + f"DTensor shape mismatch: data.shape={data.shape} vs " + f"dist_param.shape={global_shape}" + ) + local_tensor = data.to_local() + else: + # Plain Tensor must already be the FSDP-local shard; other layouts + # (TP-local / global) would silently misalign in the slice below. + expected = fsdp_slice.stop - fsdp_slice.start + assert data.numel() == expected, ( + f"Plain Tensor must be FSDP-local shard " + f"(expected numel={expected}, got {data.numel()}; fsdp_slice={fsdp_slice})" + ) + local_tensor = data + view_shape = list(global_shape) - per_tp_rank_shape = list(data.shape) + per_tp_rank_shape = list(global_shape) if is_mcore_tensor_model_parallel(dist_param): tp_dim = get_mcore_tensor_parallel_partition_dim(dist_param) assert tp_dim is not None, "Tensor model parallel dimension not found" @@ -485,8 +625,8 @@ def split_gdn_fused(data, dist_param, split_sizes, split_dim): comp_flat = s * elems_per_unit comp_slice = slice(flat_offset, flat_offset + comp_flat) - shard = _intersect_slice(fsdp_slice, comp_slice) - comp_data = local_tensor.view(-1)[_shift_slice(shard, -fsdp_slice.start)] + shard = intersection(fsdp_slice, comp_slice) + comp_data = local_tensor.view(-1)[offset_slice(shard, -fsdp_slice.start)] comp_view = list(view_shape) comp_view[split_dim] = -1 @@ -518,7 +658,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 +686,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 @@ -555,296 +695,6 @@ def split_gdn_fused(data, dist_param, split_sizes, split_dim): return model_state_dict, optimizer_state_dict -def split_fused_fsdp_param(data, dist_param, split_sizes, is_expert_param=False, split_dim=0): - """Split a fused Megatron-FSDP parameter along ``split_dim`` into per-section DTensors. - - ``split_sizes`` are given in full, tensor-parallel-unsharded units along ``split_dim`` - (so they can be taken straight from the module config) and are scaled down to this - rank's TP shard internally. - - The returned DTensors alias the fused parameter's storage, so a save reads the fused - values and an in-place load writes back into the fused parameter. - - Same flat-slice arithmetic as :func:`handle_swiglu_in_state_dict`, generalized from an - even two-way split to arbitrary section sizes. - """ - assert HAVE_MEGATRON_FSDP, "This function requires Megatron-FSDP to be installed." - - fsdp_slice = dist_param.megatron_fsdp_slice - dist_index = dist_param.megatron_fsdp_dist_index - tp_mesh = dist_index.get_submesh([dist_index.tp_dim], is_expert_parallel=is_expert_param) - - per_tp_rank_shape = list(data.shape) - if is_mcore_tensor_model_parallel(dist_param): - tp_dim = get_mcore_tensor_parallel_partition_dim(dist_param) - assert tp_dim is not None, "Tensor model parallel dimension not found" - per_tp_rank_shape[tp_dim] //= tp_mesh.mesh.numel() - - total_full = sum(split_sizes) - assert data.shape[split_dim] == total_full, ( - f"Fused parameter is {data.shape[split_dim]} wide along dim {split_dim}, " - f"but the requested sections sum to {total_full}" - ) - - local_total = per_tp_rank_shape[split_dim] - local_sizes = [] - for size in split_sizes: - assert (size * local_total) % total_full == 0, ( - f"Section of size {size} does not divide evenly across the tensor-parallel " - f"group (fused dim {total_full} -> {local_total} on this rank)" - ) - local_sizes.append(size * local_total // total_full) - - data_size = 1 - for dim_size in per_tp_rank_shape: - data_size *= dim_size - elems_per_unit = data_size // local_total - - local_tensor = data.to_local() - view_shape = list(per_tp_rank_shape) - view_shape[split_dim] = -1 - - results = [] - flat_offset = 0 - for local_size in local_sizes: - section_numel = local_size * elems_per_unit - section_slice = slice(flat_offset, flat_offset + section_numel) - - shard = _intersect_slice(fsdp_slice, section_slice) - section_data = local_tensor.view(-1)[_shift_slice(shard, -fsdp_slice.start)] - section_data = section_data.reshape(view_shape) - - # A meta tensor carries the unsharded shape and TP attributes of the section. - meta_shape = list(per_tp_rank_shape) - meta_shape[split_dim] = local_size - section_meta = torch.empty(*meta_shape, device="meta") - copy_tensor_model_parallel_attributes(section_meta, dist_param) - - results.append( - make_fsdp_dtensor( - section_data.data, - section_meta, - dist_index=dist_index, - is_expert_param=is_expert_param, - run_check=True, - update_uneven_dtensor_chunk_meta=True, - ) - ) - flat_offset += section_numel - - return results - - -# The single down-projection registered by FusedMLASelfAttention, and the two separate -# projections registered by the unfused MLASelfAttention it replaces. -MLA_FUSED_DOWN_PROJ = 'linear_qkv_down_proj' -MLA_UNFUSED_DOWN_PROJS = ('linear_q_down_proj', 'linear_kv_down_proj') - - -def get_mla_fused_down_proj_splits(model): - """Map each FusedMLASelfAttention module path to its ``[q, kv]`` down-proj split sizes. - - Sizes match ``FusedMLASelfAttention.sharded_state_dict`` and are expressed in full, - tensor-parallel-unsharded units. Returns an empty dict for unfused models. - """ - splits = {} - for name, mod in model.named_modules(): - if not hasattr(mod, MLA_FUSED_DOWN_PROJ): - continue - config = mod.config - splits[_strip_wrapper_prefixes(name)] = [ - config.q_lora_rank, - config.kv_lora_rank + config.qk_pos_emb_head_dim, - ] - return splits - - -def match_mla_fused_down_proj_key(key, fused_splits): - """Match a state dict key against the fused MLA down-projections in ``fused_splits``. - - Returns ``(wrapper, attention_path, split_sizes, leaf)`` where ``wrapper`` is the - stripped ``module.``/``model.`` prefix and ``leaf`` is the part after the fused module - name, or ``None`` when the key is not a fused down-projection. - """ - norm = _strip_wrapper_prefixes(key) - for attention_path, split_sizes in fused_splits.items(): - fused_prefix = f'{attention_path}.{MLA_FUSED_DOWN_PROJ}.' - if norm.startswith(fused_prefix): - wrapper = key[: len(key) - len(norm)] - return wrapper, attention_path, split_sizes, norm[len(fused_prefix) :] - return None - - -def absorbed_input_layernorm_key(wrapper, attention_path, leaf): - """Rewrite a ``layer_norm_*`` leaf absorbed by the fused down-proj to ``input_layernorm.*``. - - With ``fuse_input_layernorm``, the fused module owns the layer's input layernorm as - ``layer_norm_weight``/``layer_norm_bias``; on disk it lives one level up under - ``input_layernorm``, per the layer spec's ``sharded_state_dict_keys_map``. - """ - assert '.' in attention_path, ( - f"Cannot locate the transformer layer owning {attention_path}.{leaf}; expected the " - f"attention module to be nested inside it." - ) - layer_path = attention_path.rsplit('.', 1)[0] - return f'{wrapper}{layer_path}.input_layernorm.{leaf[len("layer_norm_"):]}' - - -def handle_mla_down_proj_in_state_dict(model, model_state_dict, optimizer_state_dict): - """Rewrite a fused MLA down-projection into the unfused layout used on disk. - - ``mla_down_proj_fusion=True`` swaps ``MLASelfAttention`` for ``FusedMLASelfAttention``, - replacing ``linear_q_down_proj`` and ``linear_kv_down_proj`` with a single - ``linear_qkv_down_proj`` holding their row-concatenation, and on the Transformer Engine - backend also absorbing the layer's ``input_layernorm`` into it as ``layer_norm_*``. - - ``FusedMLASelfAttention.sharded_state_dict`` hides both changes from ``torch_dist`` - checkpoints so fused and unfused runs share one on-disk format. The Megatron-FSDP path - builds its keys from ``named_parameters()`` and never calls ``sharded_state_dict``, so - the same remapping is applied here. - - No-op for unfused models. - """ - assert HAVE_MEGATRON_FSDP, "This function requires Megatron-FSDP to be installed." - - fused_splits = get_mla_fused_down_proj_splits(model) - if not fused_splits: - return model_state_dict, optimizer_state_dict - - rewritten = 0 - - model_state_dict = model_state_dict.copy() - for key in list(model_state_dict.keys()): - match = match_mla_fused_down_proj_key(key, fused_splits) - if match is None: - continue - wrapper, attention_path, split_sizes, leaf = match - - if leaf.endswith('_extra_state'): - # Not a tensor; handle_fp8_extra_state_case drops these beforehand. - continue - - if leaf.startswith('layer_norm_'): - new_key = absorbed_input_layernorm_key(wrapper, attention_path, leaf) - model_state_dict[new_key] = model_state_dict.pop(key) - rewritten += 1 - continue - - sections = split_fused_fsdp_param( - model_state_dict[key], model.get_parameter(f'module.{key}'), split_sizes - ) - for proj_name, section in zip(MLA_UNFUSED_DOWN_PROJS, sections): - model_state_dict[f'{wrapper}{attention_path}.{proj_name}.{leaf}'] = section - del model_state_dict[key] - rewritten += 1 - - if rewritten: - logger.info( - f"[MLA] Rewrote {rewritten} fused {MLA_FUSED_DOWN_PROJ} key(s) across " - f"{len(fused_splits)} attention module(s) into the unfused layout." - ) - - if optimizer_state_dict is not None and len(optimizer_state_dict.get("state", {})) != 0: - optimizer_state_dict = optimizer_state_dict.copy() - optimizer_state = optimizer_state_dict["state"] - new_optimizer_state = {} - for key in list(optimizer_state.keys()): - match = match_mla_fused_down_proj_key(key, fused_splits) - if match is None: - new_optimizer_state[key] = optimizer_state[key] - continue - wrapper, attention_path, split_sizes, leaf = match - - if leaf.startswith('layer_norm_'): - # Absorbed input layernorm: moved, not split. - new_key = absorbed_input_layernorm_key(wrapper, attention_path, leaf) - new_optimizer_state[new_key] = optimizer_state[key] - continue - - new_keys = [ - f'{wrapper}{attention_path}.{proj_name}.{leaf}' - for proj_name in MLA_UNFUSED_DOWN_PROJS - ] - for new_key in new_keys: - new_optimizer_state[new_key] = optimizer_state[key].copy() - dist_param = model.get_parameter(key[len("module.") :]) - for subkey in ["exp_avg", "exp_avg_sq"]: - sections = split_fused_fsdp_param( - optimizer_state[key][subkey], dist_param, split_sizes - ) - for new_key, section in zip(new_keys, sections): - new_optimizer_state[new_key][subkey] = section - optimizer_state_dict["state"] = new_optimizer_state - - return model_state_dict, optimizer_state_dict - - -# MCore renamed MultiTokenPredictionLayer's inner transformer layer; checkpoints keep the -# original name. -MTP_INNER_LAYER = 'mtp_model_layer' -MTP_INNER_LAYER_CHECKPOINT_NAME = 'transformer_layer' - - -def get_mtp_inner_layer_paths(model): - """Paths of MTP layers whose inner layer is renamed on disk. - - Mamba MTP layers (``mtp_layer_pattern`` set) are excluded, matching - ``MultiTokenPredictionLayer.sharded_state_dict``: for them ``mtp_model_layer`` is - already the native checkpoint name. - """ - return [ - _strip_wrapper_prefixes(name) - for name, mod in model.named_modules() - if hasattr(mod, MTP_INNER_LAYER) and getattr(mod, 'mtp_layer_pattern', None) is None - ] - - -def rename_mtp_inner_layer_keys(state_dict, mtp_layer_paths): - """Rename ``mtp_model_layer.*`` keys to ``transformer_layer.*`` under ``mtp_layer_paths``.""" - state_dict = state_dict.copy() - for key in list(state_dict.keys()): - norm = _strip_wrapper_prefixes(key) - for mtp_path in mtp_layer_paths: - current_prefix = f'{mtp_path}.{MTP_INNER_LAYER}.' - if not norm.startswith(current_prefix): - continue - wrapper = key[: len(key) - len(norm)] - leaf = norm[len(current_prefix) :] - new_key = f'{wrapper}{mtp_path}.{MTP_INNER_LAYER_CHECKPOINT_NAME}.{leaf}' - state_dict[new_key] = state_dict.pop(key) - break - return state_dict - - -def handle_mtp_in_state_dict(model, model_state_dict, optimizer_state_dict): - """Rename MTP inner-layer keys to the ``transformer_layer`` name used on disk. - - ``MultiTokenPredictionLayer.sharded_state_dict`` remaps ``mtp_model_layer.*`` back to - ``transformer_layer.*`` so GPT MTP checkpoints written before the rename keep loading. - The Megatron-FSDP path never calls ``sharded_state_dict``, so the same mapping is - applied here. - - No-op for models without MTP layers. - """ - mtp_layer_paths = get_mtp_inner_layer_paths(model) - if not mtp_layer_paths: - return model_state_dict, optimizer_state_dict - - model_state_dict = rename_mtp_inner_layer_keys(model_state_dict, mtp_layer_paths) - logger.info( - f"[MTP] Renamed {MTP_INNER_LAYER} -> {MTP_INNER_LAYER_CHECKPOINT_NAME} for " - f"{len(mtp_layer_paths)} MTP layer(s)." - ) - - if optimizer_state_dict is not None and len(optimizer_state_dict.get("state", {})) != 0: - optimizer_state_dict = optimizer_state_dict.copy() - optimizer_state_dict["state"] = rename_mtp_inner_layer_keys( - optimizer_state_dict["state"], mtp_layer_paths - ) - - return model_state_dict, optimizer_state_dict - - def handle_fp8_extra_state_case(model_state_dict): """ Handle the case where FP8 extra state is present in the model state dict. @@ -907,108 +757,6 @@ def print_diff_in_state_dicts(state_dict_metadata, load_state_dict, limit=100): logger.info(f" {k}: meta shape={meta_shape}, load shape={load_shape}") -# Top-level sections holding model weights: "model" for a single chunk, "model0", "model1", -# ... with virtual pipeline parallelism. -_MODEL_SECTION_PATTERN = re.compile(r'^model\d*\.') - - -def get_unexpected_model_keys(state_dict_metadata, load_state_dict): - """Model weights this rank requests that the checkpoint does not contain. - - "Unexpected" follows :class:`~megatron.core.dist_checkpointing.validation.StrictHandling`: - keys accessed locally but absent from the checkpoint. The opposite direction (keys the - checkpoint holds that this rank does not request) is *not* reported, because that is - normal under pipeline, tensor and expert parallelism where every rank loads a subset of - the global checkpoint. - - Only the ``model`` sections are inspected, so optimizer and RNG state that is - intentionally absent does not count. Transformer Engine ``_extra_state`` entries are - skipped because they are dropped from both sides before saving and loading. - """ - metadata_keys = set(flatten_state_dict(state_dict_metadata).keys()) - return { - key - for key in flatten_state_dict(load_state_dict) - if _MODEL_SECTION_PATTERN.match(key) - and not key.endswith('._extra_state') - and key not in metadata_keys - } - - -def validate_fsdp_dtensor_model_load( - state_dict_metadata, - load_state_dict, - checkpoint_path, - strict=StrictHandling.RAISE_UNEXPECTED, - max_reported=20, -): - """Report model weights that an ``fsdp_dtensor`` load cannot supply. - - Torch DCP loads with ``allow_partial_load=True`` skip such keys silently, leaving those - parameters at their initialized values. That does not crash; it surfaces much later as - a convergence regression, so by default fail here instead. - - Only "unexpected" keys are determined (see :func:`get_unexpected_model_keys`), so the - ``*_ALL`` variants of ``strict`` behave like their ``*_UNEXPECTED`` counterparts; - identifying genuinely "missing" keys would require exchanging key sets across ranks. - - ``ASSUME_OK_UNEXPECTED`` is treated as ``RAISE_UNEXPECTED`` here. It means "rely on the - underlying strategy to raise", which a partial DCP load never does, and it exists to skip - the extra disk access an explicit check normally costs — but the caller already holds the - checkpoint metadata, so this check is free. Use ``IGNORE_ALL`` to actually skip it. - - Args: - state_dict_metadata: ``state_dict_metadata`` from the checkpoint's DCP metadata. - load_state_dict: state dict this rank is about to load into. - checkpoint_path: checkpoint location, for the error message. - strict (StrictHandling): how to handle a mismatch. Defaults to raising. - max_reported (int): cap on the number of keys named in the message. - - Returns: - Set[str]: the unexpected model keys, empty when ``strict`` disables the check. - - Raises: - CheckpointingException: if ``strict`` is a ``RAISE_*`` value and keys are missing. - """ - strict = parse_strict_flag(strict) - if strict is StrictHandling.IGNORE_ALL: - return set() - - unexpected_keys = get_unexpected_model_keys(state_dict_metadata, load_state_dict) - if not unexpected_keys or strict in ( - StrictHandling.RETURN_UNEXPECTED, - StrictHandling.RETURN_ALL, - ): - return unexpected_keys - - reported = sorted(unexpected_keys) - truncated = ( - f'\n ... and {len(reported) - max_reported} more' if len(reported) > max_reported else '' - ) - rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 - message = ( - f"Megatron-FSDP checkpoint load is incomplete: rank {rank} requests " - f"{len(reported)} model parameter(s) not present in the checkpoint at " - f"{checkpoint_path}, which would silently keep their initialized values.\n" - + '\n'.join(f' {key}' for key in reported[:max_reported]) - + truncated - + "\nThe model's parameter names have diverged from the checkpoint's, typically " - "because of a fused vs. unfused module variant or a renamed submodule. " - "fsdp_dtensor checkpoints bypass sharded_state_dict(), so key remapping " - "implemented there does not apply; add the mapping to " - "preprocess_fsdp_dtensor_state_dict instead." - ) - - if strict in ( - StrictHandling.ASSUME_OK_UNEXPECTED, - StrictHandling.RAISE_UNEXPECTED, - StrictHandling.RAISE_ALL, - ): - raise CheckpointingException(message) - logger.warning(message) - return unexpected_keys - - def validate_loaded_state_dict(state_dict, checkpoint_path): """ Validate the loaded state dict against the expected structure and types. diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index e0fe107302c..d092e82e54e 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -7,7 +7,6 @@ import multiprocessing import os import random -import re import shutil import sys import threading @@ -15,7 +14,7 @@ from argparse import Namespace from datetime import datetime from enum import Enum, auto -from logging import DEBUG, getLogger +from logging import getLogger from pathlib import Path from time import time from typing import Any, Dict, List, Optional, Union @@ -24,15 +23,8 @@ import torch from torch.distributed.checkpoint import FileSystemReader, default_planner -try: - from nemo.lens.helpers import managed_span as _otel_managed_span -except ImportError: - from megatron.core.telemetry.fallbacks import managed_span as _otel_managed_span - from megatron.core import dist_checkpointing, mpu, tensor_parallel -from megatron.core._rank_utils import safe_get_rank as get_rank_safe -from megatron.core.dist_checkpointing.dict_utils import dict_list_map_inplace -from megatron.core.dist_checkpointing.mapping import LocalNonpersistentObject, ShardedObject +from megatron.core.dist_checkpointing.mapping import ShardedObject from megatron.core.dist_checkpointing.strategies.async_utils import _disable_gc from megatron.core.dist_checkpointing.strategies.fully_parallel import ( FullyParallelLoadStrategyWrapper, @@ -43,23 +35,18 @@ TorchDistSaveShardedStrategy, get_async_strategy, ) -from megatron.core.msc_utils import MultiStorageClientFeature, maybe_msc +from megatron.core.msc_utils import MultiStorageClientFeature, open_file from megatron.core.num_microbatches_calculator import update_num_microbatches from megatron.core.optimizer import DistributedOptimizer -from megatron.core.post_training.modelopt.checkpointing import save_modelopt_state, save_sharded_modelopt_state from megatron.core.rerun_state_machine import get_rerun_state_machine -from megatron.core.tokenizers import MegatronTokenizer from megatron.core.utils import get_pg_rank, get_pg_size, unwrap_model -from megatron.training.argument_utils import _default_config_from_args -from megatron.training.config import TokenizerConfig -from megatron.training.global_vars import get_tokenizer from ..core.dist_checkpointing.utils import _clean_metadata_for_serialization from . import ft_integration, wandb_utils from .async_utils import get_save_and_finalize_callbacks, is_empty_async_queue, schedule_async_save from .global_vars import get_args from .one_logger_utils import on_save_checkpoint_start, on_save_checkpoint_success -from .utils import append_to_progress_log, is_last_rank, print_rank_0, print_rank_last, warn_rank_0 +from .utils import append_to_progress_log, is_last_rank, print_rank_0 try: from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( @@ -68,11 +55,9 @@ from megatron.core.transformer.fsdp_dtensor_checkpoint import ( handle_experts_in_state_dict, handle_fp8_extra_state_case, - handle_mla_down_proj_in_state_dict, - handle_mtp_in_state_dict, + handle_gdn_in_state_dict, handle_swiglu_in_state_dict, print_diff_in_state_dicts, - validate_fsdp_dtensor_model_load, ) HAVE_MEGATRON_FSDP = True @@ -82,6 +67,8 @@ # [ModelOpt]: Import try: + from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state + from megatron.post_training.utils import print_distributed_quant_summary has_nvidia_modelopt = True @@ -99,22 +86,6 @@ _deletion_processes = [] -def _maybe_compact_rollout_bank(iteration): - """Compact the rollout bank without making checkpointing depend eagerly on RL.""" - from megatron.rl.rl_utils import maybe_compact_rollout_bank - - maybe_compact_rollout_bank(iteration) - - -def _register_rollout_bank_compaction(async_save_request, iteration): - """Compact the rollout bank after an asynchronous checkpoint becomes durable.""" - - def rollout_bank_finalize_fn(iteration=iteration): - _maybe_compact_rollout_bank(iteration) - - async_save_request.add_finalize_fn(rollout_bank_finalize_fn) - - def finalize_deletion_processes(blocking=False): """Clean up deletion processes to prevent zombie processes. @@ -132,13 +103,9 @@ def finalize_deletion_processes(blocking=False): finished = [] for proc in _deletion_processes: if not proc.is_alive() or blocking: - if logger.isEnabledFor(DEBUG): - logger.debug( - "Joining deletion process %s (blocking=%s, is_alive=%s)", - proc.pid, - blocking, - proc.is_alive(), - ) + logger.debug( + f"Joining deletion process {proc.pid} (blocking={blocking}, is_alive={proc.is_alive()})" + ) proc.join() finished.append(proc) for proc in finished: @@ -148,7 +115,7 @@ def finalize_deletion_processes(blocking=False): def set_checkpoint_version(value): global _CHECKPOINT_VERSION if _CHECKPOINT_VERSION is not None: - assert _CHECKPOINT_VERSION == value, 'checkpoint versions do not match' + assert _CHECKPOINT_VERSION == value, "checkpoint versions do not match" _CHECKPOINT_VERSION = value @@ -173,15 +140,12 @@ def get_loaded_iteration(): return _LOADED_ITERATION -def check_checkpoint_args(checkpoint_args, skip_args: set[str] | None = None): +def check_checkpoint_args(checkpoint_args): """Ensure fixed arguments for a model are the same for the input arguments and the one retrieved from checkpoint.""" args = get_args() - skip_args = skip_args or set() def _compare(arg_name, old_arg_name=None, default=None): - if arg_name in skip_args: - return if old_arg_name is not None: ckpt_arg_name = old_arg_name else: @@ -192,17 +156,14 @@ def _compare(arg_name, old_arg_name=None, default=None): checkpoint_value = getattr(checkpoint_args, ckpt_arg_name) args_value = getattr(args, arg_name) error_message = ( - '{} value from checkpoint ({}) is not equal to the input argument value ({}).'.format( - arg_name, checkpoint_value, args_value - ) + '{} value from checkpoint ({}) is not equal to the ' + 'input argument value ({}).'.format(arg_name, checkpoint_value, args_value) ) assert checkpoint_value == args_value, error_message _compare('num_layers') _compare('hidden_size') _compare('num_attention_heads') - if hasattr(args, 'gdp_num_householder'): - _compare('gdp_num_householder', default=3) _compare('add_position_embedding', default=True) if args.vocab_file: _compare('max_position_embeddings') @@ -221,10 +182,22 @@ def _compare(arg_name, old_arg_name=None, default=None): _compare('pipeline_model_parallel_size') +def isfile(filename) -> bool: + if MultiStorageClientFeature.is_enabled(): + msc = MultiStorageClientFeature.import_package() + return msc.os.path.isfile(filename) + else: + return os.path.isfile(filename) + + def ensure_directory_exists(filename, check_parent=True): """Build filename's path if it does not already exists.""" dirname = os.path.dirname(filename) if check_parent else filename - maybe_msc.os.makedirs(dirname, exist_ok=True) + if MultiStorageClientFeature.is_enabled(): + msc = MultiStorageClientFeature.import_package() + msc.os.makedirs(dirname, exist_ok=True) + else: + os.makedirs(dirname, exist_ok=True) def get_checkpoint_name( @@ -237,7 +210,7 @@ def get_checkpoint_name( expert_parallel=None, expert_rank=None, return_base_dir=False, - basename='model_optim_rng.pt', + basename="model_optim_rng.pt", ): """Determine the directory name for this rank's checkpoint.""" if release: @@ -276,25 +249,25 @@ def get_checkpoint_name( return os.path.join(common_path, basename) -def get_load_checkpoint_path_by_args(args, load_arg='load'): +def get_load_checkpoint_path_by_args(args, load_arg="load"): """Get the checkpoint path based on the arguments.""" load_dir = getattr(args, load_arg) iteration, release = -1, False tracker_filename = 'because load directory is not defined' if load_dir is not None: tracker_filename = get_checkpoint_tracker_filename(load_dir) - if maybe_msc.os.path.isfile(tracker_filename): + if isfile(tracker_filename): iteration, release = read_metadata(tracker_filename) # Allow user to specify the loaded iteration. - if getattr(args, 'ckpt_step', None): + if getattr(args, "ckpt_step", None): iteration = args.ckpt_step return get_checkpoint_name(load_dir, iteration, release, return_base_dir=True) def get_distributed_optimizer_checkpoint_name(model_checkpoint_name): - return os.path.join(os.path.dirname(model_checkpoint_name), 'distrib_optim.pt') + return os.path.join(os.path.dirname(model_checkpoint_name), "distrib_optim.pt") def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): @@ -317,7 +290,7 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): expert_parallel=False, expert_rank=0, ) - if maybe_msc.os.path.isfile(filename): + if isfile(filename): return filename # Look for checkpoint with no pipelining and expert parallelism @@ -331,7 +304,7 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): expert_parallel=True, expert_rank=0, ) - if maybe_msc.os.path.isfile(filename): + if isfile(filename): return filename # Look for checkpoint with pipelining and no expert parallelism @@ -345,7 +318,7 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): expert_parallel=False, expert_rank=0, ) - if maybe_msc.os.path.isfile(filename): + if isfile(filename): return filename # Look for checkpoint with pipelining and expert parallelism @@ -359,7 +332,7 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): expert_parallel=True, expert_rank=0, ) - if maybe_msc.os.path.isfile(filename): + if isfile(filename): return filename # Look for a distributed checkpoint @@ -382,7 +355,7 @@ def checkpoint_exists(checkpoints_path): if checkpoints_path is None: return False path = get_checkpoint_tracker_filename(checkpoints_path) - return maybe_msc.os.path.isfile(path) + return isfile(path) def read_metadata(tracker_filename): @@ -391,7 +364,7 @@ def read_metadata(tracker_filename): iteration = -1 release = False - with maybe_msc.open(tracker_filename, 'r') as f: + with open_file(tracker_filename, 'r') as f: metastring = f.read().strip() try: iteration = int(metastring) @@ -430,24 +403,6 @@ def read_metadata(tracker_filename): return max_iter, release -def read_frozen_resume_iteration(save_dir): - """Resume iteration for a ``--freeze-all-layers`` run. - - Returns the integer recorded in ``save_dir``'s progress tracker - (``latest_checkpointed_iteration.txt``), or ``0`` when there is none -- i.e. a - fresh run, equivalent to ``--finetune`` on the first launch. Unlike - :func:`read_metadata` (which returns ``-1`` / errors on a missing or malformed - file), a missing tracker here simply means "start from the beginning". - """ - if save_dir is None: - return 0 - tracker_filename = get_checkpoint_tracker_filename(save_dir) - if not maybe_msc.os.path.isfile(tracker_filename): - return 0 - iteration, _release = read_metadata(tracker_filename) - return iteration - - def get_rng_state( ckpt_format: str, tp_group: torch.distributed.ProcessGroup, @@ -490,7 +445,7 @@ def get_rng_state( if dp_cp_group is not None else mpu.get_data_parallel_rank(with_context_parallel=True) ) - if ckpt_format == 'torch_dist': + if ckpt_format == "torch_dist": pp_rank = get_pg_rank(pp_group) pp_size = get_pg_size(pp_group) tp_rank = get_pg_rank(tp_group) @@ -502,10 +457,10 @@ def get_rng_state( (pp_rank, tp_rank), replica_id=dp_cp_rank, ) - elif ckpt_format == 'fsdp_dtensor': + elif ckpt_format == "fsdp_dtensor": pp_rank = get_pg_rank(pp_group) tp_rank = get_pg_rank(tp_group) - rng_state_list = {f'({pp_rank}, {tp_rank})': rng_state_list} + rng_state_list = {f"({pp_rank}, {tp_rank})": rng_state_list} return rng_state_list @@ -545,10 +500,10 @@ def _build_sharded_state_dict_metadata( args, 'use_layer_wise_distributed_optimizer', False ) - if has_distributed_optimizer and args.ckpt_format == 'fsdp_dtensor': + if has_distributed_optimizer and args.ckpt_format == "fsdp_dtensor": metadata['distrib_optim_sharding_type'] = 'fsdp_dtensor' - if has_distributed_optimizer and args.ckpt_format != 'fsdp_dtensor': + if has_distributed_optimizer and args.ckpt_format != "fsdp_dtensor": if args.dist_ckpt_optim_fully_reshardable: metadata['distrib_optim_sharding_type'] = 'fully_reshardable' metadata['distrib_optim_fully_reshardable_mem_efficient'] = ( @@ -574,7 +529,8 @@ def save_grads(save_dir, state_dict, iteration, grad_label): this can be handled in postprocessing.""" print_rank_0( - f' [{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] saving {grad_label} from iteration {iteration:7d}' + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving {grad_label} " + f"from iteration {iteration:7d}" ) if mpu.get_expert_data_parallel_rank() == 0: @@ -584,21 +540,22 @@ def save_grads(save_dir, state_dict, iteration, grad_label): tp_rank = mpu.get_tensor_model_parallel_rank() assert save_dir is not None assert iteration is not None - save_dir = os.path.join(save_dir, grad_label, f'iter_{iteration:07d}') + save_dir = os.path.join(save_dir, grad_label, f"iter_{iteration:07d}") os.makedirs(save_dir, exist_ok=True) # Save state_dict. - checkpoint_name = f'mp_rank_{tp_rank:02d}' + checkpoint_name = f"mp_rank_{tp_rank:02d}" if mpu.get_pipeline_model_parallel_world_size() > 1: - checkpoint_name += f'_{pp_rank:03d}' + checkpoint_name += f"_{pp_rank:03d}" if mpu.get_expert_model_parallel_world_size() > 1: - checkpoint_name += f'_{ep_rank:03d}' - full_save_path = os.path.join(save_dir, f'{checkpoint_name}.pth') + checkpoint_name += f"_{ep_rank:03d}" + full_save_path = os.path.join(save_dir, f"{checkpoint_name}.pth") # Convert back to dict (e.g., from collections.defaultdict) for easy loading later. torch.save(dict(state_dict), full_save_path) print_rank_0( - f' [{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] saved {grad_label} from iteration {iteration:7d}' + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saved {grad_label} " + f"from iteration {iteration:7d}" ) @@ -647,6 +604,27 @@ def save_checkpoint( start_ckpt = time() args = get_args() + # Keep the safety invariant even when arguments came through a validation path that does not + # perform the flat CLI cross-check, or were mutated after startup. + optimizer_config = getattr(optimizer, "config", None) if optimizer is not None else None + chunked_optimizer_state_offload = bool( + optimizer_config is not None + and getattr(optimizer_config, "chunked_optimizer_state_offload", False) + and getattr(optimizer_config, "optimizer_state_offload_fraction", 1.0) > 0.0 + ) + if chunked_optimizer_state_offload and not args.no_save_optim: + if args.async_save: + raise RuntimeError( + "chunked optimizer state offload does not support --async-save when optimizer " + "state is saved because reusable pinned CPU buffers cannot be handed to a " + "background writer" + ) + if getattr(args, "ckpt_format", None) != "torch_dist": + raise RuntimeError( + "chunked optimizer state offload requires --ckpt-format torch_dist; checkpoint " + "arguments may have overwritten the startup setting during resume" + ) + if args.async_save and not is_empty_async_queue(): print_rank_0( 'WARNING: Starting a checkpoint save before previous has finished. Consider increasing the checkpoint interval.' @@ -661,10 +639,6 @@ def save_checkpoint( # Only rank zero of the data parallel writes to the disk. model = unwrap_model(model) - # --freeze-all-layers: weights are frozen, so skip the weight write and its finalizers; dump - # progress is recorded separately via the post-logits finalize (see the async block below). - skip_weight_ckpt = getattr(args, 'freeze_all_layers', False) - # Handle non_persistent_ckpt flag. Besides overwriting `args.save` and # `args.use_dist_ckpt`, non-persistent global ckpt requires no additional logic ckpt_type = CheckpointType.GLOBAL if args.use_dist_ckpt else CheckpointType.LEGACY @@ -686,13 +660,13 @@ def save_checkpoint( save_dir = checkpointing_context['local_checkpoint_manager'].local_ckpt_dir else: raise NotImplementedError( - f'Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})' + f"Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})" ) ckpt_format = args.ckpt_format if ckpt_type == CheckpointType.GLOBAL else 'torch' print_rank_0( - f'[{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] saving checkpoint ' - f'at iteration {iteration:7d} to {save_dir} in {ckpt_format} format' + f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving checkpoint " + f"at iteration {iteration:7d} to {save_dir} in {ckpt_format} format" ) # Collect rng state across data parallel ranks. @@ -728,14 +702,9 @@ def save_checkpoint( return_base_dir=return_base_dir, ) - # Save dataloader state if the external dataloader supports it. + # Save dataloader state if the dataloader supports it (currently only Megatron Energon). maybe_save_dataloader_state( - train_data_iterator, - iteration, - getattr(args, 'dataloader_save', None), - tp_group=tp_group, - pp_group=pp_group, - dp_group=dp_group, + train_data_iterator, iteration, getattr(args, "dataloader_save", None) ) # Save distributed optimizer's custom parameter state. @@ -751,10 +720,14 @@ def save_checkpoint( optimizer.save_parameter_state(optim_checkpoint_name) # LayerWiseDistributedOptimizer save optimizer state to file on different ranks - if getattr(args, 'use_layer_wise_distributed_optimizer', False) and args.ckpt_format == 'torch': + if ( + getattr(args, "use_layer_wise_distributed_optimizer", False) + and not args.no_save_optim + and args.ckpt_format == 'torch' + ): dp_rank = mpu.get_data_parallel_rank() optim_checkpoint_name = os.path.join( - os.path.dirname(checkpoint_name), f'layer_wise_optimizer_{dp_rank}.pt' + os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt" ) ensure_directory_exists(optim_checkpoint_name) if not optimizer.is_stub_optimizer: @@ -791,7 +764,7 @@ def save_checkpoint( # exactly one rank. Neither dp_rank==0 nor edp_rank==0 alone covers all shards when # the dense and expert parallelism layouts disagree (e.g. TP > EP*ETP); the union # does, with at most one rank per (tp_rank, ep_rank) inside any DP group. - if not skip_weight_ckpt and ( + if ( not torch.distributed.is_initialized() or ckpt_type != CheckpointType.LEGACY or dp_rank == 0 @@ -806,21 +779,20 @@ def save_checkpoint( ) else: sharded_sd_metadata = None - with _otel_managed_span('checkpoint', 'megatron.checkpoint.save.state_dict', is_goodput_span=True): - state_dict = generate_state_dict( - args, - model, - optimizer, - opt_param_scheduler, - rng_state, - iteration=iteration, - optim_sd_kwargs=dict(metadata=sharded_sd_metadata), - model_sd_kwargs=dict(metadata=sharded_sd_metadata), - rerun_state=rerun_state, - ) + state_dict = generate_state_dict( + args, + model, + optimizer, + opt_param_scheduler, + rng_state, + iteration=iteration, + optim_sd_kwargs=dict(metadata=sharded_sd_metadata), + model_sd_kwargs=dict(metadata=sharded_sd_metadata), + rerun_state=rerun_state, + ) state_dict['num_floating_point_operations_so_far'] = num_floating_point_operations_so_far - if ckpt_type == CheckpointType.GLOBAL and ckpt_format == 'torch_dist': + if ckpt_type == CheckpointType.GLOBAL and ckpt_format == "torch_dist": if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: # TODO Handle non-empty directories (e.g., after a crash during saving). ensure_directory_exists(checkpoint_name, check_parent=False) @@ -852,11 +824,11 @@ def save_checkpoint( checkpointing_context['load_strategy'], 'cached_global_metadata', None ) if cached_global_metadata is not None: - logger.debug('Plugging in the read metadata from the load strategy...') + logger.debug("Plugging in the read metadata from the load strategy...") save_strategy.cached_global_metadata = cached_global_metadata else: logger.debug( - 'Failed to plug in the read metadata from the load strategy...' + "Failed to plug in the read metadata from the load strategy..." ) if args.ckpt_fully_parallel_save: @@ -873,67 +845,56 @@ def save_checkpoint( else mpu.get_expert_data_parallel_group() ) save_strategy = FullyParallelSaveStrategyWrapper( - save_strategy, - process_group, - args.ckpt_assume_constant_structure, - validate_access_integrity=args.ckpt_load_validate_sharding_integrity, - pg_cache_path=getattr(args, 'ckpt_pg_tensors_cache_path', None), - pg_cache_create=getattr(args, 'ckpt_pg_tensors_cache_create', False), + save_strategy, process_group, args.ckpt_assume_constant_structure ) - # Allow opting out of save-side sharding validation entirely (even on - # the first save) via --no-ckpt-load-validate-sharding-integrity. This - # skips the world-wide determine_global_metadata all_gather_object. - if not args.ckpt_load_validate_sharding_integrity: - validate_sharding_integrity = False # Store save strategy for future checkpoint saves if checkpointing_context is not None: checkpointing_context['save_strategy'] = save_strategy end_ckpt = time() logger.debug( - f'rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt ' + f"rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt " + ) + async_save_request = dist_checkpointing.save( + state_dict, + checkpoint_name, + save_strategy, + async_sharded_save=args.async_save, + validate_access_integrity=validate_sharding_integrity, + preprocess_common_before_consistancy_check=preprocess_common_state_dict_fn, + content_metadata=_clean_metadata_for_serialization(sharded_sd_metadata), + async_strategy=args.async_strategy, + verify_integrity=args.verify_integrity, ) - with _otel_managed_span('checkpoint', 'megatron.checkpoint.save.io_write', is_goodput_span=True): - async_save_request = dist_checkpointing.save( - state_dict, - checkpoint_name, - save_strategy, - async_sharded_save=args.async_save, - validate_access_integrity=validate_sharding_integrity, - preprocess_common_before_consistancy_check=preprocess_common_state_dict_fn, - content_metadata=_clean_metadata_for_serialization(sharded_sd_metadata), - async_strategy=args.async_strategy, - verify_integrity=args.verify_integrity, - ) # [ModelOpt]: save sharded modelopt_state if has_nvidia_modelopt: save_sharded_modelopt_state(model, checkpoint_name, (args.ckpt_format, 1)) - elif ckpt_type == CheckpointType.GLOBAL and ckpt_format in ['torch_dcp', 'fsdp_dtensor']: + elif ckpt_type == CheckpointType.GLOBAL and ckpt_format in ["torch_dcp", "fsdp_dtensor"]: if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: # TODO Handle non-empty directories (e.g., after a crash during saving). ensure_directory_exists(checkpoint_name, check_parent=False) - if ckpt_format == 'fsdp_dtensor': + if ckpt_format == "fsdp_dtensor": state_dict = preprocess_fsdp_dtensor_state_dict(args, state_dict, model[0]) if args.async_save: planner = torch.distributed.checkpoint.DefaultSavePlanner() coordinator_rank = 0 _, async_modules = get_async_strategy(args.async_strategy) - FileSystemWriterAsync = async_modules['FileSystemWriterAsync'] - save_state_dict_async_plan = async_modules['save_state_dict_async_plan'] + FileSystemWriterAsync = async_modules["FileSystemWriterAsync"] + save_state_dict_async_plan = async_modules["save_state_dict_async_plan"] _cpu_shm = getattr(args, 'async_ckpt_use_cpu_shm', False) _writer_kwargs = {} if _cpu_shm: if ( - 'use_cpu_shm_for_gpu_tensors' + "use_cpu_shm_for_gpu_tensors" in inspect.signature(FileSystemWriterAsync.__init__).parameters ): - _writer_kwargs['use_cpu_shm_for_gpu_tensors'] = True + _writer_kwargs["use_cpu_shm_for_gpu_tensors"] = True else: raise AssertionError( - 'Installed nvidia-resiliency-ext does not support ' - 'use_cpu_shm_for_gpu_tensors. Update nvidia-resiliency-ext ' - 'to use --async-ckpt-use-cpu-shm.' + "Installed nvidia-resiliency-ext does not support " + "use_cpu_shm_for_gpu_tensors. Update nvidia-resiliency-ext " + "to use --async-ckpt-use-cpu-shm." ) fs_storage_writer = FileSystemWriterAsync( checkpoint_name, @@ -950,6 +911,9 @@ def save_checkpoint( planner=planner, enable_cache=args.ckpt_assume_constant_structure, ) + async_save_request = get_save_and_finalize_callbacks( + fs_storage_writer, save_state_dict_ret + ) async_save_request = get_save_and_finalize_callbacks( fs_storage_writer, save_state_dict_ret, args.async_strategy ) @@ -968,7 +932,7 @@ def save_checkpoint( end_ckpt = time() logger.debug( - f'rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt ' + f"rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt " ) if ckpt_type == CheckpointType.LOCAL: try: @@ -978,7 +942,7 @@ def save_checkpoint( except ModuleNotFoundError: raise RuntimeError( "The 'nvidia_resiliency_ext' module is required for local " - 'checkpointing but was not found. Please ensure it is installed.' + "checkpointing but was not found. Please ensure it is installed." ) if (sharded_sd_metadata or {}).get('distrib_optim_sharding_type') in [ 'fully_reshardable', @@ -1015,7 +979,6 @@ def save_checkpoint( # Save. ensure_directory_exists(checkpoint_name) torch.save(state_dict, checkpoint_name) - start_misc = time() if ckpt_type != CheckpointType.LOCAL: if not args.async_save: @@ -1025,18 +988,15 @@ def save_checkpoint( torch.distributed.barrier() # And update the latest iteration - if not skip_weight_ckpt and ( - not torch.distributed.is_initialized() - or torch.distributed.get_rank() == 0 - ): + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: tracker_filename = get_checkpoint_tracker_filename(save_dir) if ckpt_type == CheckpointType.LOCAL: def iter_finalize_fn(): print_rank_0( - f' [{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] successfully ' - f'saved local checkpoint from iteration {iteration:7d}' + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully " + f"saved local checkpoint from iteration {iteration:7d}" ) if args.log_progress and args.async_save: append_to_progress_log( @@ -1068,8 +1028,6 @@ def _rank_and_size(explicit_rank, group, mpu_rank_fn, mpu_size_fn): mpu.get_pipeline_model_parallel_rank, mpu.get_pipeline_model_parallel_world_size, ) - gtp_remat_rank = mpu.get_gtp_weight_remat_rank() + 1 - gtp_remat_size_to_print = mpu.get_gtp_weight_remat_world_size() def iter_finalize_fn(): prev_iteration = 0 @@ -1077,29 +1035,19 @@ def iter_finalize_fn(): args, 'save_retain_interval', None ) # For backwards compatibility of tests. if save_retain_interval is not None: - if maybe_msc.os.path.exists(tracker_filename): - with maybe_msc.open(tracker_filename, 'r') as f: + if os.path.exists( + tracker_filename + ): # TODO: Make this work with MSC remote paths? + with open_file(tracker_filename, 'r') as f: prev_iteration = int(f.read().strip()) - with maybe_msc.open(tracker_filename, 'w') as f: - f.write('release' if release else str(iteration)) + with open_file(tracker_filename, 'w') as f: + f.write("release" if release else str(iteration)) print_rank_0( - f' [{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] successfully saved ' - f'checkpoint from iteration {int(iteration):7d} to {args.save} ' - f'[ t {tensor_mp_rank}/{tp_size_to_print}, ' - f'gtp_remat {gtp_remat_rank}/{gtp_remat_size_to_print}, ' - f'p {pipeline_mp_rank}/{pp_size_to_print} ]' + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully saved " + f"checkpoint from iteration {int(iteration):7d} to {args.save} " + f"[ t {tensor_mp_rank}/{tp_size_to_print}, " + f"p {pipeline_mp_rank}/{pp_size_to_print} ]" ) - # Save tokenizer files for torch_dist checkpoints (if enabled) - if ( - args.save_tokenizer_assets - and args.ckpt_format == 'torch_dist' - and iteration > 0 - ): - checkpoint_name = get_checkpoint_name( - save_dir, iteration=iteration, return_base_dir=True - ) - config = _default_config_from_args(TokenizerConfig, args) - save_tokenizer_assets(get_tokenizer(), config, checkpoint_name) if args.log_progress and args.async_save: append_to_progress_log( args.save, f'Saved async checkpoint\tIteration: {iteration}', barrier=False @@ -1158,9 +1106,7 @@ def iter_finalize_fn(): iter_finalize_fn() # Additional callback for one_logger (last rank) - if not skip_weight_ckpt and ( - not torch.distributed.is_initialized() or is_last_rank() - ): + if not torch.distributed.is_initialized() or is_last_rank(): def onelogger_finalize_fn(): on_save_checkpoint_success(productive_metrics, args.async_save) @@ -1172,9 +1118,7 @@ def onelogger_finalize_fn(): onelogger_finalize_fn() # Additional callback for wandb (last rank) - if not skip_weight_ckpt and ( - not torch.distributed.is_initialized() or is_last_rank() - ): + if not torch.distributed.is_initialized() or is_last_rank(): def wandb_finalize_fn(): wandb_utils.on_save_checkpoint_success( @@ -1188,64 +1132,14 @@ def wandb_finalize_fn(): wandb_finalize_fn() if args.async_save: - # Schedule logits flush AFTER the checkpoint request so the persistent - # worker processes checkpoint preload first (unblocking the main - # thread), then writes logits in the background. Finalize_fns are - # moved from the checkpoint request to the logits request so that - # "success" callbacks only fire after both writes are confirmed. - from megatron.training.distillation import get_logits_saver - - logits_saver = get_logits_saver() - if logits_saver is not None: - # In frozen-dump mode there is no checkpoint request (async_save_request is None); the - # logits request then owns its own finalize_fns. - if async_save_request is not None: - logits_finalize_fns = async_save_request.finalize_fns.copy() - async_save_request.finalize_fns.clear() - else: - logits_finalize_fns = [] - # Record run progress AFTER the logits tar is confirmed written, so a resumed job - # never skips or replays a window. Written by a single rank -- the last rank, which - # lives on the last pipeline stage where the logits saver is attached (get_logits_saver - # is None on earlier stages, including global rank 0 when PP > 1). - if skip_weight_ckpt and ( - not torch.distributed.is_initialized() or is_last_rank() - ): - - def progress_finalize_fn(): - tracker_filename = get_checkpoint_tracker_filename(args.save) - with maybe_msc.open(tracker_filename, 'w') as f: - f.write(str(iteration)) - print_rank_last(f" recorded logits-dump progress: iteration " - f"{iteration} to {tracker_filename}") - - logits_finalize_fns.append(progress_finalize_fn) - async_request_cls = get_async_strategy(args.async_strategy)[1]['AsyncRequest'] - async_logits_request = async_request_cls( - async_fn=logits_saver._write_batched_tar, - async_fn_args=logits_saver.take_pending_data(), - finalize_fns=logits_finalize_fns, - ) - - if ( - getattr(args, "rl_durable_rollout_bank", False) - and async_save_request is not None - and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0) - ): - # Compact the bank once this async checkpoint becomes durable, so the - # bank's compacted-through T tracks the model checkpoint. Rank-0 only. - _register_rollout_bank_compaction(async_save_request, iteration) - if async_save_request is not None: - schedule_async_save(async_save_request) - if logits_saver is not None: - schedule_async_save(async_logits_request) + schedule_async_save(async_save_request) print_rank_0( - f' [{datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}] scheduled ' - f'an async checkpoint save at iteration {iteration:7d} to {save_dir}' + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] scheduled " + f"an async checkpoint save at iteration {iteration:7d} to {save_dir}" ) end_misc = time() - logger.debug(f'rank: {rank}, takes {end_misc - start_misc} to finalize ckpt save ') + logger.debug(f"rank: {rank}, takes {end_misc - start_misc} to finalize ckpt save ") if not args.async_save: # Add a barrier so that all ranks wait for finalization to complete @@ -1253,151 +1147,9 @@ def progress_finalize_fn(): if torch.distributed.is_initialized(): torch.distributed.barrier() - # Durable rollout bank: compact at the checkpoint boundary so the bank's - # compacted-through T tracks this (now durable) checkpoint. Only on sync - # saves; async saves compact in their durability finalize callback above. - # Rank-0 no-op otherwise. - if getattr(args, "rl_durable_rollout_bank", False): - _maybe_compact_rollout_bank(iteration) - ft_integration.on_checkpointing_end(is_async_finalization=False) -def save_tokenizer_assets( - tokenizer: MegatronTokenizer, - config: TokenizerConfig, - checkpoint_path: str, -) -> None: - """Save tokenizer files to the checkpoint directory. - - Always saves tokenizer files to ensure checkpoints are self-contained - and portable. Handles both HuggingFace tokenizers and file-based tokenizers. - Compatible with MultiStorageClient for cloud storage support. - - Args: - tokenizer: The tokenizer instance to save. - config: Tokenizers config. - checkpoint_path: The checkpoint directory path. - """ - if tokenizer is None: - return - - # Only rank 0 saves tokenizer files - rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 - if rank != 0: - return - - def resolve_path(path_str: str) -> str: - """Resolve relative paths to absolute paths.""" - if not path_str: - return path_str - path_obj = Path(path_str) - if path_obj.is_absolute(): - return path_str - # Resolve relative to current working directory - return str(path_obj.resolve()) - - try: - # Check if MultiStorageClient is enabled - if MultiStorageClientFeature .is_enabled(): - msc = MultiStorageClientFeature .import_package() - checkpoint_path_obj = msc.Path(checkpoint_path) - tokenizer_dir = checkpoint_path_obj / "tokenizer" - tokenizer_dir.mkdir(parents=True, exist_ok=True) - use_msc = True - else: - tokenizer_dir = os.path.join(checkpoint_path, "tokenizer") - os.makedirs(tokenizer_dir, exist_ok=True) - use_msc = False - - tokenizer_type = config.tokenizer_type - - # Handle HuggingFace and Multimodal tokenizers - if tokenizer_type in ("HuggingFaceTokenizer", "MultimodalTokenizer"): - if use_msc: - import tempfile - - with tempfile.TemporaryDirectory() as tmp_dir: - if hasattr(tokenizer, "save_pretrained"): - tokenizer.save_pretrained(tmp_dir) - elif hasattr(tokenizer, "_tokenizer") and hasattr(tokenizer._tokenizer, "save_pretrained"): - tokenizer._tokenizer.save_pretrained(tmp_dir) - else: - logger.debug(f"{tokenizer_type} does not support save_pretrained(), skipping tokenizer save") - return - - logger.debug(f"Saving {tokenizer_type} files to {tokenizer_dir}") - for filename in os.listdir(tmp_dir): - src_path = os.path.join(tmp_dir, filename) - if os.path.isfile(src_path): - dest_path = tokenizer_dir / filename - with open(src_path, "rb") as src_f: - with msc.open(str(dest_path), "wb") as dest_f: - dest_f.write(src_f.read()) - else: - logger.debug(f"Saving {tokenizer_type} files to {tokenizer_dir}") - if hasattr(tokenizer, "save_pretrained"): - tokenizer.save_pretrained(tokenizer_dir) - elif hasattr(tokenizer, "_tokenizer") and hasattr(tokenizer._tokenizer, "save_pretrained"): - tokenizer._tokenizer.save_pretrained(tokenizer_dir) - return - - # Handle file-based tokenizers - resolve all paths - files_to_copy = [] - - if tokenizer_type in ("BertWordPieceLowerCase", "BertWordPieceCase"): - if config.vocab_file: - resolved_path = resolve_path(config.vocab_file) - files_to_copy.append(("vocab_file", resolved_path, "vocab.txt")) - - elif tokenizer_type == "GPT2BPETokenizer": - if config.vocab_file: - resolved_path = resolve_path(config.vocab_file) - files_to_copy.append(("vocab_file", resolved_path, "vocab.json")) - if config.merge_file: - resolved_path = resolve_path(config.merge_file) - files_to_copy.append(("merge_file", resolved_path, "merges.txt")) - - elif tokenizer_type in ("SentencePieceTokenizer", "GPTSentencePieceTokenizer", "Llama2Tokenizer"): - if config.tokenizer_model: - resolved_path = resolve_path(config.tokenizer_model) - files_to_copy.append(("tokenizer_model", resolved_path, "tokenizer.model")) - - elif tokenizer_type == "TikTokenizer": - if config.tokenizer_model: - resolved_path = resolve_path(config.tokenizer_model) - files_to_copy.append(("tokenizer_model", resolved_path, "tokenizer.json")) - - elif tokenizer_type == "NullTokenizer": - logger.debug(f"{tokenizer_type} requires no file artifacts") - return - - # Copy the files - if files_to_copy: - logger.debug(f"Saving {tokenizer_type} files to {tokenizer_dir}") - for config_attr, source_path, dest_filename in files_to_copy: - if source_path and os.path.exists(source_path): - if use_msc: - dest_path = tokenizer_dir / dest_filename - with open(source_path, "rb") as src_f: - with msc.open(str(dest_path), "wb") as dest_f: - dest_f.write(src_f.read()) - logger.debug(f"Copied {config_attr}: {source_path} -> {dest_path}") - else: - dest_path = os.path.join(tokenizer_dir, dest_filename) - shutil.copy2(source_path, dest_path) - logger.debug(f"Copied {config_attr}: {source_path} -> {dest_path}") - else: - logger.debug(f"{config_attr} not found at resolved path: {source_path}") - - except Exception as e: - if get_rank_safe() == 0: - logger.error(f"Failed to save tokenizer files: {e}") - import traceback - - logger.error(traceback.format_exc()) - - @_disable_gc() def _async_delete_checkpoint_impl( save_path, @@ -1432,7 +1184,8 @@ def _async_delete_checkpoint_impl( try: shutil.rmtree(checkpoint_name) # TODO: Make this work with MSC remote paths? print( - f' successfully deleted checkpoint from iteration {iteration_to_delete:7d} at {save_path}', + f' successfully deleted checkpoint from iteration {iteration_to_delete:7d} ' + f'at {save_path}', flush=True, ) if log_progress: @@ -1454,7 +1207,7 @@ def cleanup_old_non_persistent_checkpoint(save_dir, leave_ckpt_num=1, do_async=F return save_dir = Path(save_dir) - iter_prefix = 'iter_' + iter_prefix = "iter_" iter_ckpts = save_dir.rglob(f'{iter_prefix}*') sorted_iter_ckpts = sorted( iter_ckpts, key=lambda ckpt_name: int(ckpt_name.name[len(iter_prefix) :]) @@ -1475,13 +1228,12 @@ def remove_iter_ckpts(_iter_ckpts): remove_iter_ckpts(rm_iter_ckpts) -def maybe_save_dataloader_state( - train_iterator, iteration, dataloader_save_path, *, tp_group=None, pp_group=None, dp_group=None -): +def maybe_save_dataloader_state(train_iterator, iteration, dataloader_save_path): """Saves dataloader state if the dataloader supports it. - External dataloaders use this to store state at a specific iteration. The Megatron built-in - dataloader creates index files upfront to track its state. + Currently, this is only used by Megatron Energon dataloader (multimodal) to store its state at a + specific iteration. The Megatron built-in dataloader (text-only) creates index files upfront + to track its state. If the provided dataloader has `save_state` method, then it is called to save the state. Otherwise, no state is saved. @@ -1490,63 +1242,39 @@ def maybe_save_dataloader_state( train_iterator (iterable): Train dataloader. iteration (int): Current iteration. dataloader_save_path (str): Path where the dataloader state is saved. - tp_group (ProcessGroup): Tensor-parallel group, or MPU fallback when unset. - pp_group (ProcessGroup): Pipeline-parallel group, or MPU fallback when unset. - dp_group (ProcessGroup): Data-parallel group, or MPU fallback when unset. """ # If no dataloader or saving path is provided, exit early, otherwise, raise an error. - if train_iterator is None or dataloader_save_path is None or dataloader_save_path == '': + if train_iterator is None or dataloader_save_path is None or dataloader_save_path == "": return # If dataloader doesn't support saving state, raise an error. - if not hasattr(train_iterator.iterable, 'save_state'): + if not hasattr(train_iterator.iterable, "save_state"): raise RuntimeError( - f'Could not find a save_state for the train_iterator of type {type(train_iterator)}' + f"Could not find a save_state for the train_iterator of type {type(train_iterator)}" ) # Save dataloader state for each data parallel rank only once. first_rank = ( - get_pg_rank(pp_group) == 0 - if pp_group is not None - else mpu.is_pipeline_first_stage(ignore_virtual=True) - ) and ( - get_pg_rank(tp_group) == 0 - if tp_group is not None - else mpu.get_tensor_model_parallel_rank() == 0 + mpu.is_pipeline_first_stage(ignore_virtual=True) + and mpu.get_tensor_model_parallel_rank() == 0 ) if not first_rank: return - dp_rank = get_pg_rank(dp_group) if dp_group is not None else mpu.get_data_parallel_rank() - train_dataloader_state_dict = train_iterator.iterable.save_state() + dp_rank = mpu.get_data_parallel_rank() if dp_rank == 0: - print(f'saving dataloader checkpoint at iteration {iteration} to {dataloader_save_path}') + print(f"saving dataloader checkpoint at iteration {iteration} to {dataloader_save_path}") + train_dataloader_state_dict = train_iterator.iterable.save_state() data_state_save_path = get_checkpoint_name( - dataloader_save_path, - iteration, - pipeline_parallel=( - get_pg_size(pp_group) > 1 - if pp_group is not None - else mpu.get_pipeline_model_parallel_world_size() > 1 - ), - # Dataloader state is sharded only by DP rank. Keep it in the canonical TP0/PP0 directory. - tensor_rank=0, - pipeline_rank=0, - expert_parallel=False, - expert_rank=0, - basename=f'train_dataloader_dprank{dp_rank:03d}.pt', + dataloader_save_path, iteration, basename=f'train_dataloader_dprank{dp_rank:03d}.pt' ) - data_parallel_group = dp_group if dp_group is not None else mpu.get_data_parallel_group() - torch.distributed.barrier(group=data_parallel_group) + torch.distributed.barrier(group=mpu.get_data_parallel_group()) - if dp_rank == 0: + if mpu.get_data_parallel_rank() == 0: ensure_directory_exists(data_state_save_path) - torch.distributed.barrier(group=data_parallel_group) - - if train_dataloader_state_dict is None: - return + torch.distributed.barrier(group=mpu.get_data_parallel_group()) dataloader_save_dict = {} dataloader_save_dict['dataloader_state_dict'] = train_dataloader_state_dict @@ -1574,17 +1302,17 @@ def generate_state_dict( state_dict['iteration'] = iteration for i in range(len(model)): - key = 'model' + key = "model" if len(model) > 1: - key = f'model{i}' + key = f"model{i}" - if args.ckpt_format == 'torch_dist': + if args.ckpt_format == "torch_dist": model_sd = model[i].sharded_state_dict( **( model_sd_kwargs or { - 'metadata': { - 'dp_cp_group': mpu.get_data_parallel_group(with_context_parallel=True) + "metadata": { + "dp_cp_group": mpu.get_data_parallel_group(with_context_parallel=True) } } ) @@ -1597,25 +1325,26 @@ def generate_state_dict( # Optimizer stuff. if not args.no_save_optim: if optimizer is not None and not optimizer.is_stub_optimizer: - if args.ckpt_format == 'torch_dist': + + if args.ckpt_format == "torch_dist": optimizer_sd = optimizer.sharded_state_dict( state_dict, **( optim_sd_kwargs or { - 'metadata': { - 'dp_cp_group': mpu.get_data_parallel_group( + "metadata": { + "dp_cp_group": mpu.get_data_parallel_group( with_context_parallel=True ) } } ), ) - elif args.ckpt_format == 'fsdp_dtensor': + elif args.ckpt_format == "fsdp_dtensor": if optim_sd_kwargs is None: optim_sd_kwargs = {} - if 'metadata' not in optim_sd_kwargs: - optim_sd_kwargs['metadata'] = {} + if "metadata" not in optim_sd_kwargs: + optim_sd_kwargs["metadata"] = {} optim_sd_kwargs['metadata'].update(_build_sharded_state_dict_metadata(args)) optimizer_sd = optimizer.sharded_state_dict(state_dict, **optim_sd_kwargs) else: @@ -1632,124 +1361,32 @@ def generate_state_dict( # RNG states. if not args.no_save_rng and rng_state: - state_dict['rng_state'] = rng_state - - # Optionally avoid persisting TE `_extra_state` artifacts that carry no - # irreplaceable state (see `_localize_redundant_extra_states`). Opt-in via - # --ckpt-drop-redundant-extra-state (off by default, preserving prior behavior). - # Applied here so the SAME decision is used for both the save state dict and - # the load *request* (`generate_state_dict` is the shared chokepoint), which - # keeps the two symmetric and backward compatible. - if args.ckpt_format == "torch_dist" and getattr( - args, "ckpt_drop_redundant_extra_state", False - ): - _localize_redundant_extra_states(state_dict) + state_dict["rng_state"] = rng_state return state_dict -# Byte markers of the dict keys that TE `get_extra_state` writes ONLY under -# `if recipe.delayed():` (see TransformerEngine `module/base.py`). Their presence -# in the serialized payload is a robust, unpickle-free signal that the -# `_extra_state` carries persistent delayed-scaling state (amax history + scale). -_FP8_PERSISTENT_EXTRA_STATE_MARKERS = (b"amax_history_fwd", b"scale_fwd", b"scale_bwd") - - -def _fp8_extra_state_is_persistent(data) -> bool: - """Whether a TE `_extra_state` payload must be checkpointed. - - A `_extra_state` is worth persisting only when it carries state that cannot - be reconstructed from the run config at model-build time, i.e. the - delayed-scaling FP8 amax history + scale. Concretely: - - * ``None`` / empty ``uint8`` tensor -> FP8 disabled -> NOT persistent - * non-empty, no delayed-scaling key -> block/current scaling (NVFP4, MXFP8, - Float8CurrentScaling): recipe + scalars only, all config-derived - -> NOT persistent - * non-empty, has ``scale_fwd`` / ``amax_history_fwd`` / ``scale_bwd`` - -> delayed scaling -> persistent - - The check operates on the raw serialized bytes (the uint8 tensor that - ``get_extra_state`` returns) so it needs no unpickling and no TE import. - Anything unrecognized defaults to persistent, so real state is never - silently dropped. - """ - if data is None: - return False - if isinstance(data, torch.Tensor): - if data.numel() == 0: - return False - try: - raw = data.detach().cpu().contiguous().numpy().tobytes() - except Exception: - return True - return any(marker in raw for marker in _FP8_PERSISTENT_EXTRA_STATE_MARKERS) - # Unknown payload type: keep it persistent to be safe. - return True - - -def _localize_redundant_extra_states(state_dict): - """Rewrite non-persistent TE `_extra_state` ShardedObjects as local objects. - - For checkpoints with no FP8, or with block/current FP8 scaling (NVFP4, - MXFP8, Float8CurrentScaling), every `_extra_state` is either empty or a - recipe-only blob that the freshly-built model reproduces from config — see - `_fp8_extra_state_is_persistent`. Such artifacts are pure overhead: in a - large MoE model they account for the overwhelming majority of the - checkpoint's ShardedObjects (e.g. ~50k on the 55B hybrid-MoE run). - - Wrapping them in ``LocalNonpersistentObject`` (instead of ``ShardedObject``) - means, via the dist-checkpointing pipeline: - - * SAVE: ``save_preprocess`` drops them, so they are never written to disk. - * LOAD: ``load_preprocess`` unwraps them back into the loaded state dict with - the local (freshly-built) value, so the module's ``_extra_state`` key is - still present (strict ``load_state_dict`` stays happy) and - ``set_extra_state`` no-ops on the empty/recipe payload. - - Because this runs in ``generate_state_dict`` (shared by save and the load - request), the decision is symmetric: dropped keys are never *requested*, so - loading an older checkpoint that still stores them simply does not read them - (they are neither "missing" nor "unexpected" in the request) — backward - compatible. Delayed-scaling `_extra_state` stays a ``ShardedObject`` and is - saved/loaded exactly as before. - """ - - def _maybe_localize(value): - if ( - isinstance(value, ShardedObject) - and value.key.endswith("_extra_state") - and not _fp8_extra_state_is_persistent(value.data) - ): - return LocalNonpersistentObject(value.data) - return value - - dict_list_map_inplace(_maybe_localize, state_dict) - - def preprocess_fsdp_dtensor_state_dict(args, raw_state_dict, model): state_dict = raw_state_dict.copy() - handle_fp8_extra_state_case(state_dict['model']) - - def apply(handler): - """Run a state dict handler over the model and, when present, the optimizer state.""" - model_state_dict, optimizer_state_dict = handler( - model, state_dict['model'], state_dict.get('optimizer') - ) - state_dict['model'] = model_state_dict - if 'optimizer' in state_dict: - state_dict['optimizer'] = optimizer_state_dict - + handle_fp8_extra_state_case(state_dict["model"]) if args.swiglu: - apply(handle_swiglu_in_state_dict) - # Split a fused MLA q/kv down-projection (mla_down_proj_fusion) back into the unfused - # layout used on disk. No-op for unfused models. - apply(handle_mla_down_proj_in_state_dict) + if "optimizer" in state_dict: + model_state_dict, optimizer_state_dict = handle_swiglu_in_state_dict( + model, state_dict["model"], state_dict["optimizer"] + ) + state_dict["model"] = model_state_dict + state_dict["optimizer"] = optimizer_state_dict + else: + model_state_dict, _ = handle_swiglu_in_state_dict(model, state_dict["model"], None) + state_dict["model"] = model_state_dict + model_state_dict, optimizer_state_dict = handle_gdn_in_state_dict( + model, state_dict["model"], state_dict.get("optimizer") + ) + state_dict["model"] = model_state_dict + if optimizer_state_dict is not None: + state_dict["optimizer"] = optimizer_state_dict if args.num_experts: - state_dict['model'] = handle_experts_in_state_dict(state_dict['model'], args.num_experts) - # Rename the MTP inner layer to the name used on disk. Runs last because the handlers - # above resolve keys against live module paths, which still use the new name. - apply(handle_mtp_in_state_dict) + state_dict["model"] = handle_experts_in_state_dict(state_dict["model"], args.num_experts) preprocess_state_dict_for_uneven_dtensor(state_dict) return state_dict @@ -1812,7 +1449,7 @@ def fix_query_key_value_ordering(model, checkpoint_version): elif checkpoint_version == 1.0: fixed_param = _transpose_first_dim(param.data, 3, False, model) else: - print_rank_0(f'Invalid checkpoint version {checkpoint_version}.') + print_rank_0(f"Invalid checkpoint version {checkpoint_version}.") sys.exit() param.data.copy_(fixed_param) if name.endswith(('.key_value.weight', '.key_value.bias')): @@ -1821,36 +1458,36 @@ def fix_query_key_value_ordering(model, checkpoint_version): elif checkpoint_version == 1.0: fixed_param = _transpose_first_dim(param.data, 2, False, model) else: - print_rank_0(f'Invalid checkpoint version {checkpoint_version}.') + print_rank_0(f"Invalid checkpoint version {checkpoint_version}.") sys.exit() param.data.copy_(fixed_param) print_rank_0( - ' successfully fixed query-key-values ordering for checkpoint version {}'.format( - checkpoint_version - ) + " successfully fixed query-key-values ordering for" + " checkpoint version {}".format(checkpoint_version) ) def _get_non_persistent_iteration(non_persistent_global_dir, args, checkpointing_context=None): if args.non_persistent_ckpt_type is None: return -1 - elif args.non_persistent_ckpt_type == 'global': + elif args.non_persistent_ckpt_type == "global": tracker_filename = get_checkpoint_tracker_filename(non_persistent_global_dir) - if maybe_msc.os.path.isfile(tracker_filename): + if isfile(tracker_filename): iteration, release = read_metadata(tracker_filename) if release: - raise RuntimeError("Non-persistent checkpoint can't be a release checkpoint") + raise RuntimeError('Non-persistent checkpoint can\'t be a release checkpoint') else: iteration = -1 print_rank_0('WARNING: could not find the metadata file {}'.format(tracker_filename)) print_rank_0(' will not load any non-persistent checkpoint') return iteration - elif args.non_persistent_ckpt_type == 'local': + elif args.non_persistent_ckpt_type == "local": return checkpointing_context['local_checkpoint_manager'].find_latest() else: - assert ( - False - ), f'Please use local or global non-persistent checkpoints(got: {args.non_persistent_ckpt_type})' + assert False, ( + 'Please use local or global non-persistent checkpoints' + f'(got: {args.non_persistent_ckpt_type})' + ) def _load_non_persistent_base_checkpoint( @@ -1867,7 +1504,7 @@ def _load_non_persistent_base_checkpoint( Depending on the non_persistent_ckpt_type, different logic may be required. """ assert args.non_persistent_ckpt_type is not None - if args.non_persistent_ckpt_type == 'global': + if args.non_persistent_ckpt_type == "global": if not rank0: print_rank_0( f'Loading from a non-persistent checkpoint (non-persistent iter {non_persistent_iteration})' @@ -1883,7 +1520,7 @@ def _load_non_persistent_base_checkpoint( dp_cp_group=dp_cp_group, expt_dp_group=expt_dp_group, ) - elif args.non_persistent_ckpt_type == 'local': + elif args.non_persistent_ckpt_type == "local": intermediate_state_dict, checkpoint_name = checkpointing_context[ 'local_checkpoint_manager' ].load() @@ -1899,7 +1536,7 @@ def _load_non_persistent_base_checkpoint( return state_dict, checkpoint_name, False, CheckpointType.LOCAL else: raise NotImplementedError( - f'Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})' + f"Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})" ) @@ -1945,21 +1582,14 @@ def _load_global_dist_base_checkpoint( ) else: raise ValueError( - f'Invalid load process group: {args.ckpt_fully_parallel_load_process_group}' + f"Invalid load process group: {args.ckpt_fully_parallel_load_process_group}" ) load_strategy = FullyParallelLoadStrategyWrapper( - load_strategy, - process_group, - exchange_algo=args.ckpt_fully_parallel_load_exchange_algo, - per_rank_object_load=getattr( - args, 'ckpt_fully_parallel_load_per_rank_objects', False - ), - pg_cache_path=getattr(args, 'ckpt_pg_tensors_cache_path', None), - pg_cache_create=getattr(args, 'ckpt_pg_tensors_cache_create', False), + load_strategy, process_group, exchange_algo=args.ckpt_fully_parallel_load_exchange_algo ) if checkpointing_context is not None: - checkpointing_context['load_strategy'] = load_strategy + checkpointing_context["load_strategy"] = load_strategy state_dict = dist_checkpointing.load( sharded_state_dict, checkpoint_name, @@ -1973,21 +1603,26 @@ def _load_global_dist_base_checkpoint( def _get_checkpoint_format(checkpoint_name, args): """Get the format of an existing checkpoint.""" - checkpoint_dir = maybe_msc.Path(checkpoint_name) - is_torch_ckpt = any([f.name.startswith('mp_rank_0') for f in checkpoint_dir.iterdir()]) - is_torch_dcp = checkpoint_dir.joinpath('.metadata').exists() + if MultiStorageClientFeature.is_enabled(): + msc = MultiStorageClientFeature.import_package() + checkpoint_dir = msc.Path(checkpoint_name) + is_torch_ckpt = any([f.name.startswith("mp_rank_0") for f in checkpoint_dir.iterdir()]) + is_torch_dcp = checkpoint_dir.joinpath(".metadata").exists() + else: + is_torch_ckpt = any([f.startswith("mp_rank_0") for f in os.listdir(checkpoint_name)]) + is_torch_dcp = os.path.exists(os.path.join(checkpoint_name, ".metadata")) ckpt_format = None if dist_checkpointing.check_is_distributed_checkpoint(checkpoint_name): - ckpt_format = 'torch_dist' + ckpt_format = "torch_dist" elif is_torch_ckpt: - ckpt_format = 'torch' + ckpt_format = "torch" elif is_torch_dcp: - ckpt_format = 'torch_dcp' - if getattr(args, 'use_megatron_fsdp', False): - ckpt_format = 'fsdp_dtensor' + ckpt_format = "torch_dcp" + if getattr(args, "use_megatron_fsdp", False): + ckpt_format = "fsdp_dtensor" else: - raise NotImplementedError(f'unknown checkpoint format in {checkpoint_name}') + raise NotImplementedError(f"unknown checkpoint format in {checkpoint_name}") return ckpt_format @@ -2000,7 +1635,6 @@ def _load_base_checkpoint( checkpointing_context=None, dp_cp_group=None, expt_dp_group=None, - gpt_compat_layer_maps=None, ): """Load the base state_dict from the given directory @@ -2019,11 +1653,11 @@ def _load_base_checkpoint( tracker_filename = 'because load directory is not defined' if load_dir is not None: tracker_filename = get_checkpoint_tracker_filename(load_dir) - if maybe_msc.os.path.isfile(tracker_filename): + if isfile(tracker_filename): iteration, release = read_metadata(tracker_filename) # Allow user to specify the loaded iteration. - if getattr(args, 'ckpt_step', None): + if getattr(args, "ckpt_step", None): iteration = args.ckpt_step # Record the iteration loaded (stored separately from args to avoid @@ -2053,25 +1687,19 @@ def _load_base_checkpoint( print_rank_0(' will not load any checkpoints and will start from random') # Conditionally exit if checkpoint not found. if args.exit_on_missing_checkpoint: - print_rank_0( - ">> '--exit-on-missing-checkpoint' is set but no checkpoint was found under " - f"load directory '{load_dir}' (missing metadata/tracker file " - f"'{tracker_filename}'). Exiting with a non-zero status. <<" - ) + print_rank_0(">> '--exit-on-missing-checkpoint' set ... exiting. <<") if torch.distributed.is_initialized(): torch.distributed.barrier() - # Exit non-zero so that callers (e.g. CI harnesses) detect the missing - # checkpoint as a failure instead of silently treating exit code 0 as success. - sys.exit(1) + sys.exit() - return None, '', False, None + return None, "", False, None # Determine the type of the checkpoint on disk. checkpoint_name = get_checkpoint_name(load_dir, iteration, release, return_base_dir=True) ckpt_format = _get_checkpoint_format(checkpoint_name, args) if not rank0: - dist_infix = 'distributed ' if ckpt_format == 'torch_dist' else '' + dist_infix = "distributed " if ckpt_format == "torch_dist" else "" if release: print_rank_0(f' loading release {dist_infix}checkpoint from {load_dir}') else: @@ -2082,7 +1710,7 @@ def _load_base_checkpoint( ckpt_type = None # Handle global distributed checkpoint - if ckpt_format == 'torch_dist': + if ckpt_format == "torch_dist": return _load_global_dist_base_checkpoint( load_dir, args, @@ -2094,7 +1722,7 @@ def _load_base_checkpoint( dp_cp_group=dp_cp_group, expt_dp_group=expt_dp_group, ) - elif ckpt_format == 'torch': + elif ckpt_format == "torch": ckpt_type = CheckpointType.LEGACY # Handle global legacy checkpoint if rank0: @@ -2105,17 +1733,34 @@ def _load_base_checkpoint( ) try: state_dict = torch.load(checkpoint_name, map_location='cpu') + except ModuleNotFoundError: + from megatron.legacy.fp16_deprecated import loss_scaler + + # For backward compatibility. + if not rank0: + print_rank_0(' > deserializing using the old code structure ...') + sys.modules['fp16.loss_scaler'] = sys.modules[ + 'megatron.legacy.fp16_deprecated.loss_scaler' + ] + sys.modules['megatron.fp16.loss_scaler'] = sys.modules[ + 'megatron.legacy.fp16_deprecated.loss_scaler' + ] + sys.modules['megatron.model'] = sys.modules['megatron.legacy.model'] + state_dict = torch.load(checkpoint_name, map_location='cpu') + sys.modules.pop('fp16.loss_scaler', None) + sys.modules.pop('megatron.fp16.loss_scaler', None) + sys.modules.pop('megatron.model', None) except Exception as e: print('could not load the checkpoint') print(e) sys.exit() - elif ckpt_format == 'torch_dcp': + elif ckpt_format == "torch_dcp": ckpt_type = CheckpointType.TORCH_DCP if rank0: # _load_base_checkpoint is called from load_args_from_checkpoint. torch.distributed is not initialized. # Load only metadata. - state_dict = {'args': None, 'iteration': None} + state_dict = {"args": None, "iteration": None} torch.distributed.checkpoint.load(state_dict=state_dict, checkpoint_id=checkpoint_name) else: # _load_base_checkpoint is called from load_checkpoint with a proper state dict. @@ -2126,72 +1771,63 @@ def _load_base_checkpoint( torch.distributed.checkpoint.load_state_dict( state_dict=state_dict, storage_reader=fs_storage_reader ) - elif ckpt_format == 'fsdp_dtensor': - assert HAVE_MEGATRON_FSDP, 'Should not be called if Megatron-FSDP is not available.' + elif ckpt_format == "fsdp_dtensor": + assert HAVE_MEGATRON_FSDP, "Should not be called if Megatron-FSDP is not available." if rank0: - state_dict = {'args': None, 'iteration': None, 'checkpoint_version': None} - torch.distributed.checkpoint.load(state_dict=state_dict, checkpoint_id=checkpoint_name) - return state_dict, checkpoint_name, release, CheckpointType.FSDP_DTENSOR + return {}, checkpoint_name, release, CheckpointType.FSDP_DTENSOR state_dict = sharded_state_dict raw_optimizer_state_dict = ( - state_dict['optimizer'].copy() if 'optimizer' in state_dict else None + state_dict["optimizer"].copy() if "optimizer" in state_dict else None ) - raw_model_state_dict = state_dict['model'].copy() if 'model' in state_dict else None - model = state_dict.pop('_model') + raw_model_state_dict = state_dict["model"].copy() if "model" in state_dict else None + model = state_dict.pop("_model") state_dict = preprocess_fsdp_dtensor_state_dict(args, state_dict, model[0]) - fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(checkpoint_name) - state_dict_metadata = fs_storage_reader.read_metadata().state_dict_metadata - if gpt_compat_layer_maps is not None: - from megatron.core.dist_checkpointing.gpt_checkpoint_interop import ( - retarget_fsdp_state_dict_to_gpt_checkpoint, - ) - - state_dict['model'] = retarget_fsdp_state_dict_to_gpt_checkpoint( - state_dict['model'], - gpt_compat_layer_maps, - tuple(key for key in state_dict_metadata if key.startswith('model.')), - checkpoint_prefix='model', - ) - if 'optimizer' in state_dict: - state_dict['optimizer'] = retarget_fsdp_state_dict_to_gpt_checkpoint( - state_dict['optimizer'], - gpt_compat_layer_maps, - tuple(key for key in state_dict_metadata if key.startswith('optimizer.')), - checkpoint_prefix='optimizer', - ) ckpt_type = CheckpointType.FSDP_DTENSOR + fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(checkpoint_name) allow_partial_load = not getattr(args, 'strict_fsdp_dtensor_load', False) + + # Stash state_dict keys that have no presence in the checkpoint. + # For Megatron-saved checkpoints, DCP stores all entries (tensors + # as TensorStorageMetadata, non-tensors as BytesStorageMetadata), + # so nothing needs to be stashed. + # For Bridge-converted checkpoints, metadata keys like 'args', + # 'iteration', etc. don't exist — stash them so DCP doesn't + # error, then restore placeholder values after load. + state_dict_metadata = fs_storage_reader.read_metadata().state_dict_metadata + _stashed_keys = {} + for key in list(state_dict.keys()): + # Keep the key if it exists as a top-level metadata entry + # (e.g. 'args', 'iteration') OR if any flattened sub-key + # exists (e.g. 'model.layer.weight' for top-level 'model'). + if key not in state_dict_metadata and not any( + k.startswith(key + '.') for k in state_dict_metadata + ): + _stashed_keys[key] = state_dict.pop(key) + if allow_partial_load: rank = torch.distributed.get_rank() import time as _time _time.sleep(rank * 0.001) # Make that logs of different ranks do not overlap print_diff_in_state_dicts(state_dict_metadata, state_dict) - # A partial load silently skips model weights the checkpoint does not have, - # leaving them at their initialized values. Report that, unless this is a - # fine-tune where loading only part of the model is the intent. - if not getattr(args, 'finetune', False): - validate_fsdp_dtensor_model_load( - state_dict_metadata, - state_dict, - checkpoint_name, - strict=args.dist_ckpt_strictness, - ) planner = default_planner.DefaultLoadPlanner(allow_partial_load=allow_partial_load) torch.distributed.checkpoint.load_state_dict( state_dict=state_dict, storage_reader=fs_storage_reader, planner=planner ) + # Restore stashed keys (Bridge fallback placeholders). + state_dict.update(_stashed_keys) + if raw_optimizer_state_dict is not None: - state_dict['optimizer'] = raw_optimizer_state_dict + state_dict["optimizer"] = raw_optimizer_state_dict if raw_model_state_dict is not None: - state_dict['model'] = raw_model_state_dict + state_dict["model"] = raw_model_state_dict else: - raise NotImplementedError(f'checkpoint format {ckpt_format} not supported') + raise NotImplementedError(f"checkpoint format {ckpt_format} not supported") return state_dict, checkpoint_name, release, ckpt_type @@ -2263,10 +1899,10 @@ def _set_arg(arg_name, old_arg_name=None, force=False): checkpoint_value = getattr(checkpoint_args, arg_name, None) if checkpoint_value is not None: - print_rank_0(f'Setting {arg_name} to {checkpoint_value} from checkpoint') + print_rank_0(f"Setting {arg_name} to {checkpoint_value} from checkpoint") setattr(args, arg_name, checkpoint_value) else: - print_rank_0(f'Checkpoint did not provide arguments {arg_name}') + print_rank_0(f"Checkpoint did not provide arguments {arg_name}") # Model args. _set_arg('num_layers') @@ -2325,10 +1961,6 @@ def _set_arg(arg_name, old_arg_name=None, force=False): _set_arg('mamba_head_dim', force=True) _set_arg('mamba_num_groups', force=True) _set_arg('mamba_num_heads', force=True) - # GDP checkpoints created before this argument existed always used three reflections. - if not hasattr(checkpoint_args, 'gdp_num_householder'): - setattr(checkpoint_args, 'gdp_num_householder', 3) - _set_arg('gdp_num_householder', force=True) # We need to be able to override hybrid_layer_pattern from the command-line so that different # pipelining can be specified when re-loading a model (e.g. for inference or post-training). _set_arg('hybrid_layer_pattern') @@ -2339,7 +1971,6 @@ def _set_arg(arg_name, old_arg_name=None, force=False): # MoE latent projection. _set_arg('moe_latent_size', force=True) - _set_arg('moe_use_norm_before_up_proj', force=True) # Tokenizer args. if args.use_tokenizer_model_from_checkpoint_args: @@ -2366,95 +1997,6 @@ def _set_arg(arg_name, old_arg_name=None, force=False): return args, checkpoint_args -def _maybe_setup_gpt_to_hybrid_load(args, ckpt_args, model): - """Detect a GPT (pure transformer) checkpoint being loaded into a HybridModel run. - - Returns ``(layer_maps, load_optim)`` where ``layer_maps`` is used to retarget - the run's sharded state dict at the GPT checkpoint's keys (see - ``megatron.core.dist_checkpointing.gpt_checkpoint_interop``) and ``load_optim`` is - True when the GPT optimizer state should be translated and loaded as well. - Returns ``(None, False)`` when checkpoint and runtime already agree. Raises - RuntimeError for combinations that cannot be loaded. - """ - from megatron.core.dist_checkpointing.gpt_checkpoint_interop import gpt_compatible_layer_maps - from megatron.core.models.hybrid.hybrid_model import HybridModel - from megatron.core.models.mimo.model.base import MimoModel - - def _contains_hybrid_model(module): - # Megatron-FSDP and Float16Module both retain the wrapped module under - # ``module`` but are intentionally not handled by the regular - # ``unwrap_model`` helper. Multimodal wrappers (e.g. LLaVAModel) attach - # the language model under ``language_model`` instead. - while module is not None: - if isinstance(module, HybridModel): - return True - if isinstance(module, MimoModel): - language_module = module.mimo_config.language_model_spec.module - if isinstance(language_module, type) and issubclass(language_module, HybridModel): - return True - inner = getattr(module, 'language_model', None) - if inner is not None and isinstance(inner, HybridModel): - return True - module = getattr(module, 'module', None) - return False - - runtime_is_hybrid = any(_contains_hybrid_model(m) for m in model) - ckpt_pattern = getattr(ckpt_args, 'hybrid_layer_pattern', None) or getattr( - ckpt_args, 'hybrid_override_pattern', None - ) - if runtime_is_hybrid == bool(ckpt_pattern): - return None, False - if not runtime_is_hybrid: - raise RuntimeError( - f'The checkpoint was saved by a hybrid model run (hybrid layer pattern ' - f'{ckpt_pattern!r}) but the current run builds a non-hybrid model. Load it ' - f'with the hybrid training entrypoint instead.' - ) - - # GPT checkpoint feeding a HybridModel run: translate the sharded state - # dict at load time instead of converting the checkpoint on disk. - if not args.hybrid_layer_pattern: - raise RuntimeError( - 'Loading a GPT checkpoint into a hybrid model requires ' - '--hybrid-layer-pattern so checkpoint layers can be paired with ' - 'hybrid layer positions.' - ) - try: - layer_maps = gpt_compatible_layer_maps(args.hybrid_layer_pattern) - except ValueError as exc: - raise RuntimeError(f'Cannot load a GPT checkpoint into this hybrid model: {exc}') from exc - - ckpt_num_layers = getattr(ckpt_args, 'num_layers', None) - if ckpt_num_layers is not None and ckpt_num_layers != layer_maps.num_gpt_layers: - raise RuntimeError( - f'Hybrid layer pattern {args.hybrid_layer_pattern!r} pairs with a GPT ' - f'checkpoint of {layer_maps.num_gpt_layers} layers, but the checkpoint ' - f'has num_layers={ckpt_num_layers}.' - ) - - # The optimizer state is loaded unless the user opts out or the GPT run saved - # no optimizer state. Fresh layers (e.g. Mamba) have no counterpart in the GPT - # checkpoint, so their optimizer state stays freshly initialized; warn about it. - load_optim = not args.no_load_optim and not getattr(ckpt_args, 'no_save_optim', False) - if load_optim and layer_maps.fresh_init: - print_rank_0( - f'> WARNING: {len(layer_maps.fresh_init)} hybrid layer(s) have no GPT ' - f'counterpart (e.g. Mamba positions); their weights and optimizer state ' - f'start from a fresh initialization while the GPT-sourced attention and ' - f'MLP layers load their optimizer state. Pass --no-load-optim to start ' - f"every layer's optimizer state fresh." - ) - - print_rank_0( - f'> loading a GPT checkpoint into the hybrid model: {layer_maps.num_gpt_layers} ' - f'GPT layers feed {len(layer_maps.attention_to_gpt)} attention and ' - f'{len(layer_maps.mlp_to_gpt)} MLP positions; {len(layer_maps.fresh_init)} ' - f'hybrid layers keep their fresh initialization' - + ('; optimizer state will be loaded.' if load_optim else '; optimizer state starts fresh.') - ) - return layer_maps, load_optim - - def load_checkpoint( ddp_model, optimizer, @@ -2484,22 +2026,6 @@ def load_checkpoint( args = get_args() load_dir = getattr(args, load_arg) - # --freeze-all-layers: nothing trains, so load the model in --load weights-only (finetune-style) - # and auto-resume the data position by feeding this run's own progress tracker -- written to - # --save on the previous launch -- into the standard --override-ckpt-iteration path. Reading - # progress from --save (this job's output) rather than --load lets --load stay pinned to a - # fixed checkpoint across resubmits. An explicit --override-ckpt-iteration wins. (The main use - # today is offline-KD teacher-logit dumps.) - if getattr(args, 'freeze_all_layers', False): - # Weights only: don't adopt the loaded checkpoint's optimizer / LR-scheduler / rng, or run - # check_checkpoint_args against a checkpoint from a different run (finetune gates all of - # those; --freeze-all-layers alone would still load the scheduler and assert on arg drift). - args.finetune = True - if args.override_ckpt_iteration is None: - progress_iteration = read_frozen_resume_iteration(args.save) - if progress_iteration > 0: - args.override_ckpt_iteration = progress_iteration - # Finetuning directories pretrained_dir = getattr(args, 'pretrained_checkpoint', None) if pretrained_dir is not None and not checkpoint_exists(load_dir): @@ -2508,63 +2034,47 @@ def load_checkpoint( ) load_dir = pretrained_dir if not checkpoint_exists(load_dir): - raise FileNotFoundError('No checkpoint found in load directory or pretrained directory') + raise FileNotFoundError("No checkpoint found in load directory or pretrained directory") args.finetune = True model = unwrap_model(ddp_model) ckpt_format = args.ckpt_format - state_dict = None - release = False - if args.auto_detect_ckpt_format or ckpt_format in ('torch_dist', 'fsdp_dtensor'): - with _otel_managed_span('load_checkpoint', 'megatron.checkpoint.load.io_read', is_goodput_span=True): - state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( - load_dir, args, rank0=True, checkpointing_context=checkpointing_context - ) + if args.auto_detect_ckpt_format or ckpt_format == "torch_dist": + state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( + load_dir, args, rank0=True, checkpointing_context=checkpointing_context + ) ckpt_format = None if ckpt_type == CheckpointType.TORCH_DCP: - ckpt_format = 'torch_dcp' + ckpt_format = "torch_dcp" elif ckpt_type == CheckpointType.FSDP_DTENSOR: - ckpt_format = 'fsdp_dtensor' + ckpt_format = "fsdp_dtensor" elif ckpt_type == CheckpointType.LEGACY: - ckpt_format = 'torch' + ckpt_format = "torch" elif ckpt_type in [CheckpointType.LOCAL, CheckpointType.GLOBAL]: - ckpt_format = 'torch_dist' + ckpt_format = "torch_dist" elif ckpt_type == None: pass # Not loaded. else: - raise NotImplementedError(f'checkpoint format {ckpt_format} not supported') + raise NotImplementedError(f"checkpoint format {ckpt_format} not supported") load_kwargs = {} ignore_rng_state = False ignore_rerun_state = True - ckpt_args = types.SimpleNamespace() - if ( - ckpt_format in ('torch_dist', 'fsdp_dtensor') - and state_dict is not None - and 'args' in state_dict - ): - ckpt_args = state_dict.get('args') or types.SimpleNamespace() - - # Both model-space torch_dist and fsdp_dtensor checkpoints carry model-keyed - # optimizer state that can be retargeted from GPTModel to HybridModel. - gpt_compat_layer_maps, gpt_compat_load_optim = ( - _maybe_setup_gpt_to_hybrid_load(args, ckpt_args, model) - if ckpt_format in ('torch_dist', 'fsdp_dtensor') and state_dict is not None - else (None, False) - ) - gpt_compat_load_optim = gpt_compat_load_optim and not release + if ckpt_format == "torch_dist": + ckpt_args = types.SimpleNamespace() + if state_dict is not None and "args" in state_dict: + ckpt_args = state_dict.get("args") - if ckpt_format == 'torch_dist': - if not hasattr(ckpt_args, 'tensor_model_parallel_size'): - print_rank_0('WARNING: TP size not found in checkpoint args, using 1 as default.') - if not hasattr(ckpt_args, 'pipeline_model_parallel_size'): - print_rank_0('WARNING: PP size not found in checkpoint args, using 1 as default.') + if not hasattr(ckpt_args, "tensor_model_parallel_size"): + print_rank_0("WARNING: TP size not found in checkpoint args, using 1 as default.") + if not hasattr(ckpt_args, "pipeline_model_parallel_size"): + print_rank_0("WARNING: PP size not found in checkpoint args, using 1 as default.") ckpt_tp_pp = ( - getattr(ckpt_args, 'tensor_model_parallel_size', 1), - getattr(ckpt_args, 'pipeline_model_parallel_size', 1), + getattr(ckpt_args, "tensor_model_parallel_size", 1), + getattr(ckpt_args, "pipeline_model_parallel_size", 1), ) run_tp_pp = (args.tensor_model_parallel_size, args.pipeline_model_parallel_size) @@ -2572,7 +2082,7 @@ def load_checkpoint( run_world_size = getattr(args, 'world_size', 0) ckpt_dp = getattr(ckpt_args, 'data_parallel_size', 0) run_dp = getattr(args, 'data_parallel_size', 0) - mismatch_msg = '(TP, PP) mismatch after resume ({} vs {} from checkpoint)'.format( + mismatch_msg = "(TP, PP) mismatch after resume ({} vs {} from checkpoint)".format( run_tp_pp, ckpt_tp_pp ) @@ -2599,7 +2109,7 @@ def load_checkpoint( ignore_rng_state = True gen_sd_rng_state = None if ckpt_tp_pp != run_tp_pp: - print_rank_0('{}: RNG state will be ignored'.format(mismatch_msg)) + print_rank_0("{}: RNG state will be ignored".format(mismatch_msg)) if ckpt_type == CheckpointType.LOCAL: sharded_sd_metadata = _build_sharded_state_dict_metadata(args, dp_cp_group=dp_cp_group) @@ -2611,12 +2121,10 @@ def load_checkpoint( f'sharded_state_dict metadata loaded from the checkpoint: {sharded_sd_metadata}' ) - # Determine if optimizer state will be loaded. For a GPT->hybrid load the - # optimizer state is retargeted at the GPT checkpoint even under --finetune, - # which independently controls iteration and LR-schedule reset semantics. + # Determine if optimizer state will be loaded if ( not release - and (not args.finetune or gpt_compat_load_optim) + and not args.finetune and not args.no_load_optim and not getattr(ckpt_args, 'no_save_optim', False) ): @@ -2635,32 +2143,15 @@ def load_checkpoint( else 'dp_zero_gather_scatter' ) } - # Retargeting optimizer state onto a GPT checkpoint only works for the - # model-space formats, where each optimizer ShardedTensor carries the - # model param's key and sharding. Bucket-space formats key state by a - # flat buffer layout that the extra hybrid layers reshuffle. - if gpt_compat_load_optim and sharded_sd_metadata[ - 'distrib_optim_sharding_type' - ] not in ('fully_reshardable', 'fully_sharded_model_space'): - raise RuntimeError( - 'Loading optimizer state from a GPT checkpoint into a hybrid model ' - 'is only supported for model-space distributed-optimizer checkpoints ' - "(sharding type 'fully_reshardable' or 'fully_sharded_model_space'), " - f'but the checkpoint uses ' - f'{sharded_sd_metadata["distrib_optim_sharding_type"]!r}. Re-save the ' - 'GPT checkpoint with --dist-ckpt-optim-fully-reshardable, or pass ' - '--no-load-optim ' - 'to start from a fresh optimizer.' - ) if ( ckpt_tp_pp != run_tp_pp and sharded_sd_metadata['distrib_optim_sharding_type'] not in DistributedOptimizer.checkpoint_fully_reshardable_formats ): raise RuntimeError( - f'{mismatch_msg}: not supported for DistributedOptimizer with sharding type' - f' {sharded_sd_metadata["distrib_optim_sharding_type"]}.' - f' Please use `--ckpt-fully-parallel-save` flag during checkpoint saving.' + f"{mismatch_msg}: not supported for DistributedOptimizer with sharding type" + f" {sharded_sd_metadata['distrib_optim_sharding_type']}." + f" Please use `--ckpt-fully-parallel-save` flag during checkpoint saving." ) # Check if fully parallel load is compatible with sharding type @@ -2670,8 +2161,8 @@ def load_checkpoint( == 'dp_zero_gather_scatter' ): raise RuntimeError( - 'Fully parallel load is not supported for dp_zero_gather_scatter checkpoints. ' - 'Please remove --ckpt-fully-parallel-load flag' + "Fully parallel load is not supported for dp_zero_gather_scatter checkpoints. " + "Please remove --ckpt-fully-parallel-load flag" ) else: gen_sd_optim = None @@ -2684,7 +2175,7 @@ def load_checkpoint( # Ensure we have a dict before updating to avoid NoneType AttributeError. if sharded_sd_metadata is None: sharded_sd_metadata = {} - sharded_sd_metadata['dp_cp_group'] = dp_cp_group + sharded_sd_metadata["dp_cp_group"] = dp_cp_group optim_sd_kwargs = dict(metadata=sharded_sd_metadata, is_loading=True) model_sd_kwargs = dict(metadata=sharded_sd_metadata) @@ -2706,7 +2197,7 @@ def load_checkpoint( ) ignore_rerun_state = False if ckpt_world_size != run_world_size or ckpt_tp_pp != run_tp_pp or ckpt_dp != run_dp: - print_rank_0('Job sharding has changed: Rerun state will be ignored') + print_rank_0("Job sharding has changed: Rerun state will be ignored") # [ModelOpt]: Initial loading from non-resume sharded checkpoint to a Distillation Model # will result in key mismatch with loss modules potentially containing parameters, since @@ -2714,7 +2205,7 @@ def load_checkpoint( with ( contextlib.ExitStack() as stack ): # Allows multiple context managers for each model shard - if args.finetune and hasattr(model[0], 'hide_loss_modules'): + if args.finetune and hasattr(model[0], "hide_loss_modules"): for m in model: stack.enter_context(m.hide_loss_modules()) load_kwargs['sharded_state_dict'] = generate_state_dict( @@ -2727,42 +2218,26 @@ def load_checkpoint( model_sd_kwargs=model_sd_kwargs, rerun_state=gen_sd_rerun_state, ) - - if gpt_compat_layer_maps is not None: - from megatron.core.dist_checkpointing.gpt_checkpoint_interop import ( - retarget_sharded_state_dict_to_gpt_checkpoint, - ) - - # The optimizer sharded state dict is built from the (hybrid) model sharded - # state dict, so its entries carry the same ``decoder.layers..`` keys and - # sharding; the same retargeting points them at the GPT checkpoint too. - for sd_key, sub_sd in load_kwargs['sharded_state_dict'].items(): - is_model = sd_key == 'model' or re.fullmatch(r'model\d+', sd_key) - is_optim = gpt_compat_load_optim and ( - sd_key == 'optimizer' or re.fullmatch(r'optimizer\d+', sd_key) - ) - if is_model or is_optim: - retarget_sharded_state_dict_to_gpt_checkpoint(sub_sd, gpt_compat_layer_maps) - elif args.ckpt_format == 'torch_dcp': + elif args.ckpt_format == "torch_dcp": model_sd = model[0].state_dict() optimizer_sd = optimizer.state_dict(is_loading=True) if tp_group is None and pp_group is None: tp_group = mpu.get_tensor_model_parallel_group() pp_group = mpu.get_pipeline_model_parallel_group() sharded_state_dict = { - 'model': model_sd, - 'optimizer': optimizer_sd, - 'args': None, - 'iteration': 1, - 'rng_state': get_rng_state( + "model": model_sd, + "optimizer": optimizer_sd, + "args": None, + "iteration": 1, + "rng_state": get_rng_state( args.ckpt_format, tp_group, pp_group, dp_cp_group=dp_cp_group, dp_group=dp_group ), - 'checkpoint_version': None, - 'opt_param_scheduler': opt_param_scheduler.state_dict(), - 'num_floating_point_operations_so_far': 0, + "checkpoint_version": None, + "opt_param_scheduler": opt_param_scheduler.state_dict(), + "num_floating_point_operations_so_far": 0, } - load_kwargs['sharded_state_dict'] = sharded_state_dict - elif args.ckpt_format == 'fsdp_dtensor': + load_kwargs["sharded_state_dict"] = sharded_state_dict + elif args.ckpt_format == "fsdp_dtensor": reader = FileSystemReader(get_load_checkpoint_path_by_args(args)) try: state_dict_metadata = reader.read_metadata().state_dict_metadata @@ -2774,7 +2249,7 @@ def load_checkpoint( gen_sd_rng_state = None gen_sd_optim = None if not args.finetune: - if 'rerun_state_machine' in state_dict_metadata: + if "rerun_state_machine" in state_dict_metadata: gen_sd_rerun_state = get_rerun_state_machine().state_dict( data_iterator=None, ckpt_format=ckpt_format, force=True ) @@ -2782,9 +2257,8 @@ def load_checkpoint( gen_sd_rng_state = get_rng_state( args.ckpt_format, tp_group, pp_group, dp_cp_group=dp_cp_group, dp_group=dp_group ) - if (not args.finetune or gpt_compat_load_optim) and not args.no_load_optim: - gen_sd_optim = optimizer - if not args.finetune: + if not args.no_load_optim: + gen_sd_optim = optimizer gen_sd_opt_param_scheduler = opt_param_scheduler optim_sd_kwargs = dict( @@ -2792,43 +2266,18 @@ def load_checkpoint( is_loading=True, ) - # Megatron-FSDP materializes optimizer slots with a dummy zero-gradient - # step while building a loading state dict. A normal full resume - # overwrites every model parameter afterward, but GPT->Hybrid leaves - # fresh-only layers untouched. Temporarily zero the optimizer LR so the - # shape-materialization step cannot weight-decay or momentum-update - # those fresh model parameters. - optimizer_lrs = [] - if gpt_compat_layer_maps is not None and gen_sd_optim is not None: - pending_optimizers = [gen_sd_optim] - while pending_optimizers: - pending_optimizer = pending_optimizers.pop() - if hasattr(pending_optimizer, 'chained_optimizers'): - pending_optimizers.extend(pending_optimizer.chained_optimizers) - continue - inner_optimizer = getattr(pending_optimizer, 'optimizer', None) - if inner_optimizer is None: - continue - for param_group in inner_optimizer.param_groups: - optimizer_lrs.append((param_group, param_group.get('lr'))) - param_group['lr'] = 0.0 - - try: - state_dict = generate_state_dict( - args, - model=model, - optimizer=gen_sd_optim, - opt_param_scheduler=gen_sd_opt_param_scheduler, - rng_state=gen_sd_rng_state, - optim_sd_kwargs=optim_sd_kwargs, - rerun_state=gen_sd_rerun_state, - iteration=1, - ) - finally: - for param_group, lr in optimizer_lrs: - param_group['lr'] = lr - state_dict['_model'] = model - load_kwargs['sharded_state_dict'] = state_dict + state_dict = generate_state_dict( + args, + model=model, + optimizer=gen_sd_optim, + opt_param_scheduler=gen_sd_opt_param_scheduler, + rng_state=gen_sd_rng_state, + optim_sd_kwargs=optim_sd_kwargs, + rerun_state=gen_sd_rerun_state, + iteration=1, + ) + state_dict["_model"] = model + load_kwargs["sharded_state_dict"] = state_dict state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( load_dir, @@ -2837,7 +2286,6 @@ def load_checkpoint( checkpointing_context=checkpointing_context, dp_cp_group=dp_cp_group, expt_dp_group=expt_dp_group, - gpt_compat_layer_maps=gpt_compat_layer_maps, **load_kwargs, ) @@ -2850,9 +2298,9 @@ def load_checkpoint( set_checkpoint_version(state_dict.get('checkpoint_version', 0)) # Convert to regular torch tensor to DTensor. - if ckpt_type == CheckpointType.LEGACY and args.ckpt_format == 'torch_dcp': - dtensor_state_dict = _to_dtensor(ddp_model, state_dict['model']) - state_dict['model'] = dtensor_state_dict + if ckpt_type == CheckpointType.LEGACY and args.ckpt_format == "torch_dcp": + dtensor_state_dict = _to_dtensor(ddp_model, state_dict["model"]) + state_dict["model"] = dtensor_state_dict # Set iteration. if args.finetune or release: @@ -2865,9 +2313,8 @@ def load_checkpoint( iteration = state_dict['total_iters'] except KeyError: print_rank_0( - 'A metadata file exists but unable to load iteration from checkpoint {}, exiting'.format( - checkpoint_name - ) + 'A metadata file exists but unable to load ' + 'iteration from checkpoint {}, exiting'.format(checkpoint_name) ) sys.exit() num_floating_point_operations_so_far = state_dict.get('num_floating_point_operations_so_far', 0) @@ -2875,12 +2322,7 @@ def load_checkpoint( # Check arguments. if 'args' in state_dict and not args.finetune: checkpoint_args = state_dict['args'] - # A GPT block is split into separate attention and MLP positions in - # HybridModel, so num_layers intentionally differs even for an - # architecture-preserving load. Keep every other resume-time argument - # compatibility check. - skip_args = {'num_layers'} if gpt_compat_layer_maps is not None else None - check_checkpoint_args(checkpoint_args, skip_args=skip_args) + check_checkpoint_args(checkpoint_args) args.consumed_train_samples = getattr(checkpoint_args, 'consumed_train_samples', 0) args.skipped_train_samples = getattr(checkpoint_args, 'skipped_train_samples', 0) update_num_microbatches(consumed_samples=args.consumed_train_samples, verbose=True) @@ -2888,64 +2330,18 @@ def load_checkpoint( else: print_rank_0('could not find arguments in the checkpoint ...') - # --override-ckpt-iteration: rewind the data loader to this iteration, operating on `args` - # (not state_dict) so it also works on checkpoints with no saved `args` (release / HF). The - # GBS-match check applies only when adopting this checkpoint's own args (not a finetune load). - if getattr(args, 'override_ckpt_iteration', None) is not None: - if 'args' in state_dict and not args.finetune: - ckpt_global_batch_size = getattr(state_dict['args'], 'global_batch_size', None) - if ( - ckpt_global_batch_size is not None - and ckpt_global_batch_size != args.global_batch_size - ): - warn_rank_0( - '--override-ckpt-iteration recomputes consumed_train_samples = target_iter * ' - f'current global_batch_size ({args.global_batch_size}), but this checkpoint ' - f'was saved at global_batch_size {ckpt_global_batch_size}. If the target ' - "iteration is in the checkpoint run's units, the data loader resumes from a " - 'shifted sample offset.' - ) - iteration = args.override_ckpt_iteration - args.consumed_train_samples = iteration * args.global_batch_size - args.skipped_train_samples = 0 - update_num_microbatches(consumed_samples=args.consumed_train_samples, verbose=True) - print_rank_0(f'--override-ckpt-iteration: start at iteration {iteration} ' - f'(consumed_train_samples {args.consumed_train_samples})') - def load_model_state_dict(module, state_dict, strict: bool): """Helper function to load state dict with fallback for missing extra states.""" - # GTP native-FP8 weights: load_state_dict's copy_ re-quantizes into the FP8 param, which - # TE's IsMXFP8Tensor check rejects for our subclass. Present the base FP8 class for it. - from megatron.core.tensor_parallel.gtp_api import HAVE_GTP - - if HAVE_GTP: - from megatron.core.tensor_parallel.gtp_api import gtp_native_fp8_load_context - - load_ctx = lambda: gtp_native_fp8_load_context(module) - else: - from contextlib import nullcontext - - load_ctx = nullcontext try: - with load_ctx(): - module.load_state_dict(state_dict, strict=strict) + module.load_state_dict(state_dict, strict=strict) except Exception as e: if strict: # Fallback support for backward compatibility breaking changes in TransformerEngine - with load_ctx(): - load_return = module.load_state_dict(state_dict, strict=False) - print(f'load_return: {load_return}') - - # Megatron-FSDP DTensors are loaded into the model buffers in-place above. - # Replaying the translated raw state dict through ``load_state_dict`` would - # reapply HybridModel-local entries that were intentionally omitted (for - # example fresh Mamba layers). - gpt_fsdp_model_loaded_in_place = ( - ckpt_format == 'fsdp_dtensor' and gpt_compat_layer_maps is not None - ) + load_return = module.load_state_dict(state_dict, strict=False) + print(f"load_return: {load_return}") # Model. - if not skip_load_to_model_and_opt and not gpt_fsdp_model_loaded_in_place: + if not skip_load_to_model_and_opt: if len(ddp_model) == 1: load_model_state_dict(ddp_model[0], state_dict['model'], strict) else: @@ -2961,17 +2357,17 @@ def load_model_state_dict(module, state_dict, strict: bool): fix_query_key_value_ordering(model, checkpoint_version) # Optimizer. - if not release and (not args.finetune or gpt_compat_load_optim) and not args.no_load_optim: + if not release and not args.finetune and not args.no_load_optim: try: # Load state dict. if ( - getattr(args, 'use_layer_wise_distributed_optimizer', False) + getattr(args, "use_layer_wise_distributed_optimizer", False) and args.ckpt_format == 'torch' ): # LayerWiseDistributedOptimizer load optimizer state from file on different ranks dp_rank = mpu.get_data_parallel_rank() optim_checkpoint_name = os.path.join( - os.path.dirname(checkpoint_name), f'layer_wise_optimizer_{dp_rank}.pt' + os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt" ) optimizer.load_state_dict_from_file(optim_checkpoint_name) elif ( @@ -2983,11 +2379,11 @@ def load_model_state_dict(module, state_dict, strict: bool): # Load distributed optimizer's custom parameter state. # For distributed checkpoint it's already loaded in load_state_dict above - is_torch_dist = ckpt_format == 'torch_dist' + is_torch_dist = ckpt_format == "torch_dist" if ( args.use_distributed_optimizer and not is_torch_dist - and ckpt_format not in ['torch_dcp', 'fsdp_dtensor'] + and ckpt_format not in ["torch_dcp", "fsdp_dtensor"] ): # NOTE: this is a manual read of the tracker file. # This code should not be reached when reading from a non_persistent checkpoint @@ -3003,53 +2399,12 @@ def load_model_state_dict(module, state_dict, strict: bool): update_legacy_format=args.ckpt_convert_update_legacy_dist_opt_format, ) - # Load scheduler unless --finetune requests a fresh iteration and LR schedule. - if opt_param_scheduler is not None and not args.finetune: + # Load scheduler. + if opt_param_scheduler is not None: if 'lr_scheduler' in state_dict: # backward compatbility opt_param_scheduler.load_state_dict(state_dict['lr_scheduler']) else: opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler']) - - # Optimizer state dict can overwrite per-group max_lr/min_lr values. - # If scheduler override is requested, re-apply runtime lr bounds from - # args so scheduler math does not stay pinned to checkpoint lr settings. - if getattr(args, "override_opt_param_scheduler", False): - if ( - optimizer is None - or getattr(optimizer, "is_stub_optimizer", False) - or opt_param_scheduler is None - ): - print_rank_0( - " > WARNING: --override-opt-param-scheduler is set, but optimizer or " - "opt_param_scheduler is not available; skipping override of scheduler " - "and optimizer param_group max_lr/min_lr." - ) - else: - for param_group in optimizer.param_groups: - if param_group.get("is_decoupled_lr", False): - max_lr = getattr(args, "decoupled_lr", None) - min_lr = getattr(args, "decoupled_min_lr", None) - else: - max_lr = args.lr - min_lr = args.min_lr - if max_lr is not None: - param_group["max_lr"] = max_lr - if min_lr is not None: - param_group["min_lr"] = min_lr - # Synchronize scheduler num_steps with consumed_train_samples - # to ensure lr calculation is based on current training progress - if opt_param_scheduler.num_steps != args.consumed_train_samples: - print_rank_0( - f" > WARNING: scheduler num_steps ({opt_param_scheduler.num_steps}) " - f"differs from consumed_train_samples ({args.consumed_train_samples}). " - f"Resetting scheduler num_steps to match consumed_train_samples." - ) - opt_param_scheduler.num_steps = args.consumed_train_samples - opt_param_scheduler.step(increment=0) - print_rank_0( - " > restored optimizer param_group max_lr/min_lr from runtime config " - "because --override-opt-param-scheduler is set" - ) except KeyError as e: print_rank_0( 'Unable to load optimizer from checkpoint {}. ' @@ -3071,7 +2426,7 @@ def load_model_state_dict(module, state_dict, strict: bool): if 'rerun_state_machine' in state_dict: get_rerun_state_machine().load_state_dict(state_dict['rerun_state_machine']) except Exception as e: - print_rank_0(f'Unable to restore RerunMachine from checkpoint: {e}. Skipping.') + print_rank_0(f"Unable to restore RerunMachine from checkpoint: {e}. Skipping.") # rng states. if not release and not args.finetune and not args.no_load_rng and not ignore_rng_state: @@ -3079,7 +2434,7 @@ def load_model_state_dict(module, state_dict, strict: bool): cuda_rng_tracker = tensor_parallel.get_cuda_rng_tracker() graph_safe_rng = tensor_parallel.is_graph_safe_cuda_rng_tracker(cuda_rng_tracker) if 'rng_state' in state_dict: - if args.ckpt_format == 'fsdp_dtensor': + if args.ckpt_format == "fsdp_dtensor": # FSDP DTensor checkpoints store rng_state in a different format. tp_rank = ( get_pg_rank(tp_group) @@ -3091,10 +2446,10 @@ def load_model_state_dict(module, state_dict, strict: bool): if pp_group is not None else mpu.get_pipeline_model_parallel_rank() ) - if f'({pp_rank}, {tp_rank})' in state_dict['rng_state']: - rng_state = state_dict['rng_state'][f'({pp_rank}, {tp_rank})'] + if f"({pp_rank}, {tp_rank})" in state_dict['rng_state']: + rng_state = state_dict['rng_state'][f"({pp_rank}, {tp_rank})"] else: - print_rank_0('WARNING: RNG state not found for current TP/PP rank') + print_rank_0("WARNING: RNG state not found for current TP/PP rank") rng_state = next(iter(state_dict['rng_state'].values())) else: rng_state = state_dict['rng_state'] @@ -3160,12 +2515,9 @@ def load_model_state_dict(module, state_dict, strict: bool): if pp_group is not None else mpu.get_pipeline_model_parallel_world_size() ) - _gtp_remat_r = mpu.get_gtp_weight_remat_rank() - _gtp_remat_w = mpu.get_gtp_weight_remat_world_size() print_rank_0( f' successfully loaded checkpoint from {load_dir} ' f'[ t {_tp_r + 1}/{_tp_w}, ' - f'gtp_remat {_gtp_remat_r + 1}/{_gtp_remat_w}, ' f'p {_pp_r + 1}/{_pp_w} ] ' f'at iteration {iteration}' ) @@ -3194,7 +2546,28 @@ def load_model_state_dict(module, state_dict, strict: bool): log_printed = True if has_nvidia_modelopt: - print_distributed_quant_summary(model, msg='After loading checkpoint') + print_distributed_quant_summary(model, msg="After loading checkpoint") + + # Load teacher model in Distillation mode. + if getattr(args, "export_kd_teacher_load", None): + from megatron.post_training.checkpointing import load_modelopt_checkpoint + + unwrapped_model = unwrap_model(model)[0] + # Note: load_modelopt_checkpoint may call this function so we prevent infinite recursion. + if hasattr(unwrapped_model, 'teacher_model'): + teacher = unwrapped_model.teacher_model + print_rank_0( + f"Loading teacher as {type(teacher).__name__} from {args.export_kd_teacher_load} ..." + ) + # [WAR]: To avoid error out on loading teacher's checkpoint, we temporarily + # set args.finetune to True while loading the teacher checkpoint. + original_args_finetune, original_ckpt_format = args.finetune, args.ckpt_format + args.finetune = True + if args.export_kd_teacher_ckpt_format is not None: + args.ckpt_format = args.export_kd_teacher_ckpt_format + load_modelopt_checkpoint([teacher], load_arg='export_kd_teacher_load') + args.finetune, args.ckpt_format = original_args_finetune, original_ckpt_format + print_rank_0("... teacher loaded successfully.") return iteration, num_floating_point_operations_so_far @@ -3205,7 +2578,7 @@ def _to_dtensor(wrapped_model, model_state_dict): new_model_sd = dict() for k, v in model_state_dict.items(): # FP8 extra state cannot be converted to dtensor yet. - if '_extra_state' in k: + if "_extra_state" in k: new_model_sd[k] = v else: new_model_sd[k] = torch.distributed.tensor.distribute_tensor(v, device_mesh) @@ -3229,7 +2602,7 @@ def load_biencoder_checkpoint( tracker_filename = get_checkpoint_tracker_filename(load_path) - with maybe_msc.open(tracker_filename, 'r') as f: + with open_file(tracker_filename, 'r') as f: iteration = int(f.read().strip()) checkpoint_name = get_checkpoint_name( 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/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index d5c625ca8ca..d5edfa7d779 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -8,6 +8,7 @@ import pytest import torch import torch.distributed.checkpoint +from torch.distributed.checkpoint import FileSystemReader from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel @@ -24,9 +25,8 @@ _build_sharded_state_dict_metadata, _load_base_checkpoint, get_checkpoint_tracker_filename, - load_args_from_checkpoint, load_checkpoint, - maybe_save_dataloader_state, + preprocess_fsdp_dtensor_state_dict, read_metadata, save_checkpoint, ) @@ -50,6 +50,32 @@ def sharded_state_dict(self, *args, metadata: Optional[dict] = None, **kwargs): return self.state_dict() +class MockDeltaRule(torch.nn.Module): + """Minimal delta-rule module exposing production checkpoint metadata.""" + + def __init__(self, split_names, split_sections): + super().__init__() + self.qk_dim = 4 + self.v_dim = 6 + 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) + + +class MockDeltaRuleModel(MegatronModule): + """Minimal Megatron model used to exercise the real FSDP checkpoint path.""" + + def __init__(self, config, split_names, split_sections): + super().__init__(config=config) + self.delta_rule = MockDeltaRule(split_names, split_sections) + + def sharded_state_dict(self, *args, metadata: Optional[dict] = None, **kwargs): + return self.state_dict() + + class MockState: def __init__(self, state_dict): self._state_dict = state_dict @@ -76,128 +102,111 @@ def sharded_state_dict(self, *args, metadata: Optional[dict] = None, **kwargs): return self.state_dict() -def test_maybe_save_dataloader_state_uses_explicit_process_groups(tmp_path): - """Dataloader checkpoints use the supplied module groups and canonical model-parallel path.""" - groups = { - "tp": SimpleNamespace(rank=0, size=2), - "pp": SimpleNamespace(rank=0, size=2), - "dp": SimpleNamespace(rank=3, size=4), - } - barriers = [] - saved = [] - iterator = SimpleNamespace( - iterable=SimpleNamespace(save_state=lambda: {"global_sequence_id": 16}) - ) +def test_fsdp_dtensor_preprocessing_calls_delta_rule_handler(): + args = SimpleNamespace(swiglu=False, num_experts=None) + model = mock.sentinel.model + model_state = {"weight": mock.sentinel.model_weight} + optimizer_state = {"state": {}} + split_model_state = {"weight.query": mock.sentinel.query_weight} + split_optimizer_state = {"state": {"weight.query": {}}} with ( + mock.patch("megatron.training.checkpointing.handle_fp8_extra_state_case"), mock.patch( - "megatron.training.checkpointing.get_pg_rank", side_effect=lambda group: group.rank - ), - mock.patch( - "megatron.training.checkpointing.get_pg_size", side_effect=lambda group: group.size - ), - mock.patch( - "megatron.training.checkpointing.torch.distributed.barrier", - side_effect=lambda group: barriers.append(group), - ), - mock.patch( - "megatron.training.checkpointing.torch.save", - side_effect=lambda state, path: saved.append((state, path)), - ), - ): - maybe_save_dataloader_state( - iterator, - 2, - tmp_path, - tp_group=groups["tp"], - pp_group=groups["pp"], - dp_group=groups["dp"], - ) - - assert barriers == [groups["dp"], groups["dp"]] - assert saved[0][0] == {"dataloader_state_dict": {"global_sequence_id": 16}} - assert saved[0][1] == str( - tmp_path / "iter_0000002" / "mp_rank_00_000" / "train_dataloader_dprank003.pt" - ) - - -def test_maybe_save_dataloader_state_skips_empty_state_after_barriers(tmp_path): - """Ranks without dataloader state participate in barriers but do not write a file.""" - group = SimpleNamespace(rank=0, size=1) - iterator = SimpleNamespace(iterable=SimpleNamespace(save_state=lambda: None)) - barriers = [] - - with ( - mock.patch( - "megatron.training.checkpointing.get_pg_rank", - side_effect=lambda process_group: process_group.rank, - ), - mock.patch( - "megatron.training.checkpointing.get_pg_size", - side_effect=lambda process_group: process_group.size, - ), - mock.patch( - "megatron.training.checkpointing.torch.distributed.barrier", - side_effect=lambda group: barriers.append(group), - ), - mock.patch("megatron.training.checkpointing.torch.save") as save, + "megatron.training.checkpointing.handle_gdn_in_state_dict", + return_value=(split_model_state, split_optimizer_state), + ) as handle_gdn, + mock.patch("megatron.training.checkpointing.preprocess_state_dict_for_uneven_dtensor"), ): - maybe_save_dataloader_state( - iterator, 2, tmp_path, tp_group=group, pp_group=group, dp_group=group + state_dict = preprocess_fsdp_dtensor_state_dict( + args, {"model": model_state, "optimizer": optimizer_state}, model ) - assert barriers == [group, group] - save.assert_not_called() + handle_gdn.assert_called_once_with(model, model_state, optimizer_state) + assert state_dict["model"] is split_model_state + assert state_dict["optimizer"] is split_optimizer_state -class MockOptParamScheduler(MockState): - def __init__(self, state_dict): - super().__init__(state_dict) - self.num_steps = state_dict.get("num_steps", 0) - self.step_calls = [] +def _local_checkpoint_tensor(value): + if hasattr(value, "to_local"): + value = value.to_local() + return value.detach().clone() - def load_state_dict(self, state_dict): - super().load_state_dict(state_dict) - self.num_steps = state_dict.get("num_steps", self.num_steps) - def step(self, increment=1): - self.step_calls.append(increment) +@pytest.mark.parametrize( + ("split_names", "split_sections"), + [ + (("query", "key", "value", "z", "f", "b", "w"), (4, 4, 6, 6, 4, 4, 6)), + (("query", "key", "value", "g", "gate"), (4, 4, 6, 4, 6)), + ], +) +def test_fsdp_dtensor_delta_rule_save_load_roundtrip( + init_model_parallel, create_ckpt_load_args, tmp_path_dist_ckpt, split_names, split_sections +): + """Real Megatron-FSDP save/load preserves variant-specific fused projections.""" + args = create_ckpt_load_args + args.ckpt_format = "fsdp_dtensor" + args.use_megatron_fsdp = True + args.use_distributed_optimizer = True + args.use_dist_ckpt = True + args.swiglu = False + args.num_experts = None + args.save_tokenizer_assets = False -class MockOptimizer(MockState): - def state_dict(self, is_loading=False): - state_dict = super().state_dict(is_loading=is_loading).copy() - state_dict["param_groups"] = [param_group.copy() for param_group in self.param_groups] - return state_dict + config = TransformerConfig(num_layers=1, kv_channels=1) + ddp_config = DistributedDataParallelConfig( + use_distributed_optimizer=True, use_megatron_fsdp=True + ) - def load_state_dict(self, state_dict): - super().load_state_dict(state_dict) - self.param_groups = [param_group.copy() for param_group in state_dict["param_groups"]] + def build_model(seed): + torch.manual_seed(seed) + return FullyShardedDataParallel( + config=config, + ddp_config=ddp_config, + module=MockDeltaRuleModel(config, split_names, split_sections), + ) + source_model = build_model(123) + source_optimizer = MockState({"state": {}}) + source_scheduler = MockState({"scheduler": "source"}) -@pytest.mark.parametrize( - ("checkpoint_args", "configured_num_householder", "expected_num_householder"), - [(SimpleNamespace(gdp_num_householder=5), 3, 5), (SimpleNamespace(), 5, 3)], -) -def test_load_args_restores_gdp_num_householder_from_checkpoint( - checkpoint_args, configured_num_householder, expected_num_householder -): - args = SimpleNamespace( - load="checkpoint", - iteration=0, - gdp_num_householder=configured_num_householder, - use_tokenizer_model_from_checkpoint_args=False, - use_mp_args_from_checkpoint_args=False, - ) - state_dict = {"args": checkpoint_args, "iteration": 12} + with TempNamedDir(tmp_path_dist_ckpt / "test_delta_rule_fsdp_roundtrip", sync=True) as ckpt_dir: + args.save = ckpt_dir + args.load = ckpt_dir + set_args(args) + save_checkpoint(7, [source_model], source_optimizer, source_scheduler, 0) + + checkpoint_dir = ckpt_dir / "iter_0000007" + metadata_keys = FileSystemReader(checkpoint_dir).read_metadata().state_dict_metadata + metadata_prefix = "model.module.delta_rule" + for name in split_names: + assert f"{metadata_prefix}.in_proj.weight.{name}" in metadata_keys + for parameter in ("weight", "bias"): + for name in ("query", "key", "value"): + assert f"{metadata_prefix}.conv1d.{parameter}.{name}" in metadata_keys + assert f"{metadata_prefix}.in_proj.weight" not in metadata_keys + assert f"{metadata_prefix}.conv1d.weight" not in metadata_keys + assert f"{metadata_prefix}.conv1d.bias" not in metadata_keys + + source_state = { + key: _local_checkpoint_tensor(value) for key, value in source_model.state_dict().items() + } + target_model = build_model(456) + target_optimizer = MockState({"state": {}}) + target_scheduler = MockState({"scheduler": "target"}) - with mock.patch( - "megatron.training.checkpointing._load_base_checkpoint", - return_value=(state_dict, "checkpoint", False, CheckpointType.LEGACY), - ): - restored_args, _ = load_args_from_checkpoint(args) + loaded_iteration, _ = load_checkpoint( + [target_model], target_optimizer, target_scheduler, strict=True + ) - assert restored_args.gdp_num_householder == expected_num_householder + assert loaded_iteration == 7 + target_state = { + key: _local_checkpoint_tensor(value) for key, value in target_model.state_dict().items() + } + assert source_state.keys() == target_state.keys() + for key in source_state: + torch.testing.assert_close(target_state[key], source_state[key], rtol=0, atol=0) def create_checkpoint(load_path, ckpt_format): @@ -242,7 +251,6 @@ def create_args(): args.auto_detect_ckpt_format = False args.ckpt_convert_update_legacy_dist_opt_format = False args.ckpt_step = None - args.override_opt_param_scheduler = False args.swiglu = True args.num_experts = 1 args.verify_integrity = False @@ -250,6 +258,87 @@ def create_args(): yield args +def test_async_save_rejects_reusable_chunked_optimizer_buffers(create_args): + """Checkpoint entry must preserve the CLI guard when startup validation was bypassed.""" + + create_args.async_save = True + create_args.ckpt_format = 'torch_dist' + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + pytest.raises(RuntimeError, match="does not support --async-save"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_rejects_mutated_format_with_chunked_optimizer_offload(create_args): + """Checkpoint entry revalidates the format after resume-time argument restoration.""" + + create_args.ckpt_format = 'torch' + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + pytest.raises(RuntimeError, match="requires --ckpt-format torch_dist"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_treats_zero_offload_fraction_as_disabled(create_args): + """A zero fraction must not impose async-save or distributed-format restrictions.""" + + create_args.async_save = True + create_args.ckpt_format = 'torch' + optimizer = SimpleNamespace( + config=SimpleNamespace( + chunked_optimizer_state_offload=True, optimizer_state_offload_fraction=0.0 + ) + ) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + mock.patch( + "megatron.training.checkpointing.is_empty_async_queue", + side_effect=RuntimeError("continued past optimizer offload guard"), + ), + pytest.raises(RuntimeError, match="continued past optimizer offload guard"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_skips_offload_guards_without_optimizer_state(create_args): + """Model-only saves do not expose reusable optimizer buffers to the async writer.""" + + create_args.async_save = True + create_args.ckpt_format = 'torch_dist' + create_args.no_save_optim = True + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + mock.patch( + "megatron.training.checkpointing.is_empty_async_queue", + side_effect=RuntimeError("continued past optimizer offload guard"), + ), + pytest.raises(RuntimeError, match="continued past optimizer offload guard"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + +def test_checkpoint_save_allows_legacy_format_without_optimizer_state(create_args): + """A model-only save is not forced through optimizer-specific distributed hooks.""" + + create_args.ckpt_format = 'torch' + create_args.no_save_optim = True + optimizer = SimpleNamespace(config=SimpleNamespace(chunked_optimizer_state_offload=True)) + with ( + mock.patch("megatron.training.checkpointing.get_args", return_value=create_args), + mock.patch( + "megatron.training.checkpointing.on_save_checkpoint_start", + side_effect=RuntimeError("continued past optimizer offload guard"), + ), + pytest.raises(RuntimeError, match="continued past optimizer offload guard"), + ): + save_checkpoint(0, [], optimizer, None, 0) + + @pytest.fixture def create_ckpt_load_args(create_args): """Setup dummy args allowing checkpoint load.""" @@ -360,7 +449,6 @@ def test_save_checkpoint(init_model_parallel, create_args, tmp_path_dist_ckpt, c with TempNamedDir(tmp_path_dist_ckpt / "test_save_checkpoint", sync=True) as save_dir: args.save = save_dir - args.save_tokenizer_assets = False set_args(args) save_checkpoint( @@ -397,7 +485,6 @@ def test_load_checkpoint( with TempNamedDir(tmp_path_dist_ckpt / "test_load_checkpoint", sync=True) as ckpt_dir: args.load = ckpt_dir args.save = ckpt_dir - args.save_tokenizer_assets = False set_args(args) # Create and save a checkpoint first. @@ -433,79 +520,6 @@ def test_load_checkpoint( assert new_opt_param_scheduler.state_dict() == opt_param_scheduler.state_dict() -@pytest.mark.parametrize("ckpt_format", ["torch"]) -def test_load_checkpoint_override_opt_param_scheduler( - init_model_parallel, create_ckpt_load_args, tmp_path_dist_ckpt, ckpt_format -): - """Test override_opt_param_scheduler behavior during checkpoint load.""" - args = create_ckpt_load_args - args.ckpt_format = ckpt_format - args.use_distributed_optimizer = False - args.use_dist_ckpt = ckpt_format != "torch" - args.override_opt_param_scheduler = True - args.lr = 1.0 - args.min_lr = 0.1 - args.decoupled_lr = 0.5 - args.decoupled_min_lr = 0.05 - args.consumed_train_samples = 42 - - with TempNamedDir( - tmp_path_dist_ckpt / "test_load_checkpoint_override_opt_param_scheduler", sync=True - ) as ckpt_dir: - args.load = ckpt_dir - args.save = ckpt_dir - args.save_tokenizer_assets = False - set_args(args) - - # Create and save a checkpoint first. - iteration = 123 - config = TransformerConfig(num_layers=1, kv_channels=1) - model = MockModel(config) - - optimizer = MockOptimizer({"optimizer": "optimizer_state"}) - optimizer.param_groups = [ - {"is_decoupled_lr": False, "max_lr": -1.0, "min_lr": -1.0}, - {"is_decoupled_lr": True, "max_lr": -1.0, "min_lr": -1.0}, - ] - opt_param_scheduler = MockOptParamScheduler( - {"opt_param_scheduler": "scheduler_state", "num_steps": 3} - ) - num_floating_point_operations_so_far = 456 - - save_checkpoint( - iteration, [model], optimizer, opt_param_scheduler, num_floating_point_operations_so_far - ) - - # Create new model, optimizer, and scheduler instances to load into. - new_model = MockModel(config) - new_optimizer = MockOptimizer({"optimizer": "dummy1"}) - new_optimizer.param_groups = [ - {"is_decoupled_lr": False, "max_lr": -2.0, "min_lr": -2.0}, - {"is_decoupled_lr": True, "max_lr": -2.0, "min_lr": -2.0}, - ] - new_opt_param_scheduler = MockOptParamScheduler( - {"opt_param_scheduler": "dummy2", "num_steps": 0} - ) - - # Load checkpoint and verify runtime overrides are restored. - loaded_iter, loaded_flops = load_checkpoint( - [new_model], new_optimizer, new_opt_param_scheduler, strict=True - ) - assert loaded_iter == iteration - assert loaded_flops == num_floating_point_operations_so_far - assert new_optimizer.param_groups[0]["max_lr"] == args.lr - assert new_optimizer.param_groups[0]["min_lr"] == args.min_lr - assert new_optimizer.param_groups[1]["max_lr"] == args.decoupled_lr - assert new_optimizer.param_groups[1]["min_lr"] == args.decoupled_min_lr - assert new_opt_param_scheduler.num_steps == args.consumed_train_samples - assert new_opt_param_scheduler.step_calls[-1] == 0 - - # Ensure loading without optimizer/scheduler remains safe. - loaded_iter_none, loaded_flops_none = load_checkpoint([new_model], None, None, strict=True) - assert loaded_iter_none == iteration - assert loaded_flops_none == num_floating_point_operations_so_far - - def test_dist_checkpoint_versioning(init_model_parallel, tmp_path_dist_ckpt, create_ckpt_load_args): """Test distributed checkpoint versioning.""" args = create_ckpt_load_args @@ -518,7 +532,6 @@ def test_dist_checkpoint_versioning(init_model_parallel, tmp_path_dist_ckpt, cre ) as ckpt_dir: args.load = ckpt_dir args.save = ckpt_dir - args.save_tokenizer_assets = False set_args(args) # Create and save a checkpoint first. diff --git a/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py b/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py index 0655e865c86..b616e140829 100644 --- a/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py +++ b/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py @@ -16,26 +16,24 @@ 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. """ import re -from types import SimpleNamespace import pytest import torch -from megatron.core.dist_checkpointing.core import CheckpointingException -from megatron.core.dist_checkpointing.validation import StrictHandling from megatron.core.distributed.fsdp.src.megatron_fsdp.utils import ( get_mcore_tensor_parallel_partition_dim, is_mcore_tensor_model_parallel, @@ -45,80 +43,13 @@ copy_tensor_model_parallel_attributes, set_tensor_model_parallel_attributes, ) -from megatron.core.transformer import fsdp_dtensor_checkpoint from megatron.core.transformer.fsdp_dtensor_checkpoint import ( - MLA_UNFUSED_DOWN_PROJS, - _intersect_slice, - _shift_slice, - _strip_wrapper_prefixes, - absorbed_input_layernorm_key, flatten_state_dict, get_expert_index_from_key, - get_mla_fused_down_proj_splits, - get_mtp_inner_layer_paths, - get_unexpected_model_keys, - handle_mla_down_proj_in_state_dict, - handle_mtp_in_state_dict, - match_mla_fused_down_proj_key, - rename_mtp_inner_layer_keys, - validate_fsdp_dtensor_model_load, + handle_gdn_in_state_dict, ) -# ============================================================================ -# Test the slice/path helpers shared by the SwiGLU, GDN and MLA handlers -# ============================================================================ -class TestSharedSliceHelpers: - """These three are used by every fused-parameter handler, so pin their semantics.""" - - def test_intersect_overlapping(self): - assert _intersect_slice(slice(0, 10), slice(4, 20)) == slice(4, 10) - - def test_intersect_contained(self): - assert _intersect_slice(slice(0, 10), slice(2, 5)) == slice(2, 5) - - @pytest.mark.parametrize( - "s1,s2", - [ - (slice(0, 4), slice(4, 8)), # touching but not overlapping - (slice(0, 4), slice(6, 8)), # disjoint - (slice(6, 8), slice(0, 4)), # disjoint, reversed - ], - ) - def test_intersect_without_overlap_is_empty(self, s1, s2): - """An FSDP shard that holds none of a section must yield a zero-length slice, not a - negative-length one, otherwise the section's reshape silently takes the wrong data.""" - result = _intersect_slice(s1, s2) - assert result.stop - result.start == 0 - assert list(range(10))[result] == [] - - def test_shift_slice(self): - assert _shift_slice(slice(6, 10), -4) == slice(2, 6) - - def test_shift_rebases_onto_local_storage(self): - """The handlers shift by -fsdp_slice.start to index into this rank's local shard.""" - fsdp_slice = slice(8, 16) - section = _intersect_slice(fsdp_slice, slice(12, 20)) - assert _shift_slice(section, -fsdp_slice.start) == slice(4, 8) - - @pytest.mark.parametrize( - "path,expected", - [ - ("module.decoder.layers.0.weight", "decoder.layers.0.weight"), - ("model.module.decoder.layers.0.weight", "decoder.layers.0.weight"), - ("module.module.decoder.weight", "decoder.weight"), - ("decoder.layers.0.weight", "decoder.layers.0.weight"), - ("", ""), - ], - ) - def test_strip_wrapper_prefixes(self, path, expected): - assert _strip_wrapper_prefixes(path) == expected - - def test_strip_keeps_inner_module_segments(self): - """Only leading wrapper segments go; 'module' deeper in the path is a real submodule.""" - assert _strip_wrapper_prefixes("module.decoder.module.weight") == "decoder.module.weight" - - # ============================================================================ # Test get_expert_index_from_key # ============================================================================ @@ -483,6 +414,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. @@ -730,633 +744,6 @@ def test_vlm_scenario(self): assert all("vision_encoder" in k for k in should_skip) -# ============================================================================ -# Fixtures for the fused-MLA / MTP handlers -# -# These build real nn.Modules so that named_modules() behaves exactly as it does -# in production; only the leaf parameters are placeholders. -# ============================================================================ -Q_LORA_RANK = 24 -KV_LORA_RANK = 16 -QK_POS_EMB_HEAD_DIM = 8 -MLA_SPLIT_SIZES = [Q_LORA_RANK, KV_LORA_RANK + QK_POS_EMB_HEAD_DIM] - - -def _mla_config(): - return SimpleNamespace( - q_lora_rank=Q_LORA_RANK, kv_lora_rank=KV_LORA_RANK, qk_pos_emb_head_dim=QK_POS_EMB_HEAD_DIM - ) - - -class _FusedMLAAttention(torch.nn.Module): - """Stands in for FusedMLASelfAttention: one fused down-proj plus a config.""" - - def __init__(self, absorb_layernorm=False): - super().__init__() - self.config = _mla_config() - self.linear_qkv_down_proj = torch.nn.Module() - self.linear_qkv_down_proj.weight = torch.nn.Parameter(torch.zeros(sum(MLA_SPLIT_SIZES), 4)) - if absorb_layernorm: - self.linear_qkv_down_proj.layer_norm_weight = torch.nn.Parameter(torch.zeros(4)) - - -class _UnfusedMLAAttention(torch.nn.Module): - """Stands in for MLASelfAttention: separate q/kv down-projections.""" - - def __init__(self): - super().__init__() - self.config = _mla_config() - self.linear_q_down_proj = torch.nn.Module() - self.linear_kv_down_proj = torch.nn.Module() - - -class _Layer(torch.nn.Module): - def __init__(self, attention): - super().__init__() - self.self_attention = attention - self.input_layernorm = torch.nn.Module() - - -class _MTPLayer(torch.nn.Module): - """Stands in for MultiTokenPredictionLayer.""" - - def __init__(self, mtp_layer_pattern=None): - super().__init__() - self.mtp_layer_pattern = mtp_layer_pattern - self.mtp_model_layer = torch.nn.Module() - - -class _Decoder(torch.nn.Module): - def __init__(self, layers): - super().__init__() - self.layers = torch.nn.ModuleList(layers) - - -class _Model(torch.nn.Module): - """Wrapper mirroring the ``module.`` indirection the handlers expect.""" - - def __init__(self, decoder_layers=(), mtp_layers=()): - super().__init__() - inner = torch.nn.Module() - inner.decoder = _Decoder(list(decoder_layers)) - if mtp_layers: - inner.mtp = _Decoder(list(mtp_layers)) - self.module = inner - - def get_parameter(self, name): - # The handlers only need an object carrying Megatron-FSDP shard metadata; the - # split itself is exercised separately. - return SimpleNamespace(name=name) - - -def _fused_model(num_layers=1, absorb_layernorm=False, with_mtp=False): - layers = [_Layer(_FusedMLAAttention(absorb_layernorm)) for _ in range(num_layers)] - mtp = [_MTPLayer()] if with_mtp else [] - return _Model(decoder_layers=layers, mtp_layers=mtp) - - -# ============================================================================ -# Test get_mla_fused_down_proj_splits -# ============================================================================ -class TestGetMLAFusedDownProjSplits: - """The fused down-proj split sizes must match FusedMLASelfAttention.sharded_state_dict.""" - - def test_fused_model_reports_q_and_kv_sizes(self): - splits = get_mla_fused_down_proj_splits(_fused_model()) - assert splits == {"decoder.layers.0.self_attention": MLA_SPLIT_SIZES} - - def test_kv_section_includes_pos_emb_head_dim(self): - """kv is kv_lora_rank + qk_pos_emb_head_dim, not kv_lora_rank alone.""" - splits = get_mla_fused_down_proj_splits(_fused_model()) - _, kv = splits["decoder.layers.0.self_attention"] - assert kv == KV_LORA_RANK + QK_POS_EMB_HEAD_DIM - - def test_split_sizes_sum_to_fused_width(self): - model = _fused_model() - splits = get_mla_fused_down_proj_splits(model) - fused_weight = model.module.decoder.layers[0].self_attention.linear_qkv_down_proj.weight - assert sum(splits["decoder.layers.0.self_attention"]) == fused_weight.shape[0] - - def test_unfused_model_reports_nothing(self): - model = _Model(decoder_layers=[_Layer(_UnfusedMLAAttention())]) - assert get_mla_fused_down_proj_splits(model) == {} - - def test_multiple_layers(self): - splits = get_mla_fused_down_proj_splits(_fused_model(num_layers=3)) - assert set(splits) == { - "decoder.layers.0.self_attention", - "decoder.layers.1.self_attention", - "decoder.layers.2.self_attention", - } - - def test_mtp_attention_included(self): - """The MTP layer's attention is fused too and must be rewritten alongside the decoder.""" - model = _Model(decoder_layers=[_Layer(_FusedMLAAttention())], mtp_layers=[_MTPLayer()]) - model.module.mtp.layers[0].mtp_model_layer = _Layer(_FusedMLAAttention()) - splits = get_mla_fused_down_proj_splits(model) - assert "mtp.layers.0.mtp_model_layer.self_attention" in splits - - -# ============================================================================ -# Test match_mla_fused_down_proj_key -# ============================================================================ -class TestMatchMLAFusedDownProjKey: - - SPLITS = {"decoder.layers.0.self_attention": MLA_SPLIT_SIZES} - - def test_weight_key_matched(self): - key = "decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - wrapper, path, sizes, leaf = match_mla_fused_down_proj_key(key, self.SPLITS) - assert wrapper == "" - assert path == "decoder.layers.0.self_attention" - assert sizes == MLA_SPLIT_SIZES - assert leaf == "weight" - - def test_wrapper_prefix_preserved(self): - key = "module.decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - wrapper, path, _, leaf = match_mla_fused_down_proj_key(key, self.SPLITS) - assert wrapper == "module." - assert path == "decoder.layers.0.self_attention" - assert leaf == "weight" - - def test_nested_wrapper_prefix_preserved(self): - key = "model.module.decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - wrapper, _, _, _ = match_mla_fused_down_proj_key(key, self.SPLITS) - assert wrapper == "model.module." - - def test_layer_norm_leaf_matched(self): - key = "decoder.layers.0.self_attention.linear_qkv_down_proj.layer_norm_weight" - _, _, _, leaf = match_mla_fused_down_proj_key(key, self.SPLITS) - assert leaf == "layer_norm_weight" - - def test_extra_state_leaf_matched(self): - key = "decoder.layers.0.self_attention.linear_qkv_down_proj._extra_state" - _, _, _, leaf = match_mla_fused_down_proj_key(key, self.SPLITS) - assert leaf == "_extra_state" - - def test_already_unfused_key_not_matched(self): - """An unfused checkpoint-style key must not be rewritten a second time.""" - key = "decoder.layers.0.self_attention.linear_q_down_proj.weight" - assert match_mla_fused_down_proj_key(key, self.SPLITS) is None - - def test_unrelated_key_not_matched(self): - key = "decoder.layers.0.mlp.linear_fc1.weight" - assert match_mla_fused_down_proj_key(key, self.SPLITS) is None - - def test_other_layer_not_matched(self): - key = "decoder.layers.1.self_attention.linear_qkv_down_proj.weight" - assert match_mla_fused_down_proj_key(key, self.SPLITS) is None - - def test_empty_splits_match_nothing(self): - key = "decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - assert match_mla_fused_down_proj_key(key, {}) is None - - -# ============================================================================ -# Test absorbed_input_layernorm_key -# ============================================================================ -class TestAbsorbedInputLayernormKey: - """With fuse_input_layernorm the layernorm moves up one level, per the layer spec's - sharded_state_dict_keys_map: self_attention.linear_qkv_down_proj.layer_norm_ -> - input_layernorm.""" - - def test_weight(self): - assert ( - absorbed_input_layernorm_key("", "decoder.layers.0.self_attention", "layer_norm_weight") - == "decoder.layers.0.input_layernorm.weight" - ) - - def test_bias(self): - assert ( - absorbed_input_layernorm_key("", "decoder.layers.3.self_attention", "layer_norm_bias") - == "decoder.layers.3.input_layernorm.bias" - ) - - def test_wrapper_preserved(self): - assert ( - absorbed_input_layernorm_key( - "module.", "decoder.layers.0.self_attention", "layer_norm_weight" - ) - == "module.decoder.layers.0.input_layernorm.weight" - ) - - def test_top_level_attention_rejected(self): - with pytest.raises(AssertionError): - absorbed_input_layernorm_key("", "self_attention", "layer_norm_weight") - - -# ============================================================================ -# Test handle_mla_down_proj_in_state_dict -# -# split_fused_fsdp_param needs real Megatron-FSDP DTensors, so it is replaced with a -# recorder; everything around it (key rewriting, split sizes, optimizer state) is real. -# ============================================================================ -@pytest.fixture -def recorded_splits(monkeypatch): - calls = [] - - def fake_split(data, dist_param, split_sizes, is_expert_param=False, split_dim=0): - calls.append({"data": data, "split_sizes": list(split_sizes)}) - return [f"{data}::q", f"{data}::kv"] - - monkeypatch.setattr(fsdp_dtensor_checkpoint, "split_fused_fsdp_param", fake_split) - return calls - - -class TestHandleMLADownProjInStateDict: - - WEIGHT_KEY = "decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - Q_KEY = "decoder.layers.0.self_attention.linear_q_down_proj.weight" - KV_KEY = "decoder.layers.0.self_attention.linear_kv_down_proj.weight" - - def test_fused_weight_becomes_q_and_kv(self, recorded_splits): - model_sd, _ = handle_mla_down_proj_in_state_dict( - _fused_model(), {self.WEIGHT_KEY: "W"}, None - ) - assert set(model_sd) == {self.Q_KEY, self.KV_KEY} - assert model_sd[self.Q_KEY] == "W::q" - assert model_sd[self.KV_KEY] == "W::kv" - - def test_unfused_projection_names_match_mcore(self): - assert MLA_UNFUSED_DOWN_PROJS == ("linear_q_down_proj", "linear_kv_down_proj") - - def test_config_split_sizes_are_forwarded(self, recorded_splits): - handle_mla_down_proj_in_state_dict(_fused_model(), {self.WEIGHT_KEY: "W"}, None) - assert recorded_splits[0]["split_sizes"] == MLA_SPLIT_SIZES - - def test_absorbed_layernorm_is_moved_not_split(self, recorded_splits): - ln_key = "decoder.layers.0.self_attention.linear_qkv_down_proj.layer_norm_weight" - model_sd, _ = handle_mla_down_proj_in_state_dict( - _fused_model(absorb_layernorm=True), {ln_key: "LN"}, None - ) - assert model_sd == {"decoder.layers.0.input_layernorm.weight": "LN"} - assert recorded_splits == [] - - def test_extra_state_left_alone(self, recorded_splits): - key = "decoder.layers.0.self_attention.linear_qkv_down_proj._extra_state" - model_sd, _ = handle_mla_down_proj_in_state_dict(_fused_model(), {key: "X"}, None) - assert model_sd == {key: "X"} - - def test_unrelated_keys_untouched(self, recorded_splits): - other = "decoder.layers.0.mlp.linear_fc1.weight" - model_sd, _ = handle_mla_down_proj_in_state_dict( - _fused_model(), {self.WEIGHT_KEY: "W", other: "M"}, None - ) - assert model_sd[other] == "M" - - def test_unfused_model_is_noop(self, recorded_splits): - model = _Model(decoder_layers=[_Layer(_UnfusedMLAAttention())]) - original = {self.Q_KEY: "Q", self.KV_KEY: "KV"} - model_sd, opt_sd = handle_mla_down_proj_in_state_dict(model, original, None) - assert model_sd == original - assert opt_sd is None - assert recorded_splits == [] - - def test_input_state_dict_not_mutated(self, recorded_splits): - original = {self.WEIGHT_KEY: "W"} - handle_mla_down_proj_in_state_dict(_fused_model(), original, None) - assert original == {self.WEIGHT_KEY: "W"} - - def test_all_layers_rewritten(self, recorded_splits): - model_sd, _ = handle_mla_down_proj_in_state_dict( - _fused_model(num_layers=3), - { - f"decoder.layers.{i}.self_attention.linear_qkv_down_proj.weight": f"W{i}" - for i in range(3) - }, - None, - ) - assert len(model_sd) == 6 - assert all("linear_qkv_down_proj" not in k for k in model_sd) - - def test_optimizer_state_split(self, recorded_splits): - opt_key = f"module.{self.WEIGHT_KEY}" - optimizer_state_dict = {"state": {opt_key: {"exp_avg": "A", "exp_avg_sq": "B", "step": 7}}} - _, opt_sd = handle_mla_down_proj_in_state_dict(_fused_model(), {}, optimizer_state_dict) - state = opt_sd["state"] - assert set(state) == {f"module.{self.Q_KEY}", f"module.{self.KV_KEY}"} - assert state[f"module.{self.Q_KEY}"]["exp_avg"] == "A::q" - assert state[f"module.{self.KV_KEY}"]["exp_avg"] == "A::kv" - assert state[f"module.{self.Q_KEY}"]["exp_avg_sq"] == "B::q" - # Non-tensor entries are carried over to both halves. - assert state[f"module.{self.Q_KEY}"]["step"] == 7 - - def test_optimizer_layernorm_moved_not_split(self, recorded_splits): - """Regression: the absorbed layernorm has no q/kv halves to split into.""" - ln_key = "module.decoder.layers.0.self_attention.linear_qkv_down_proj.layer_norm_weight" - _, opt_sd = handle_mla_down_proj_in_state_dict( - _fused_model(absorb_layernorm=True), {}, {"state": {ln_key: {"exp_avg": "A"}}} - ) - assert set(opt_sd["state"]) == {"module.decoder.layers.0.input_layernorm.weight"} - assert recorded_splits == [] - - def test_optimizer_unrelated_state_preserved(self, recorded_splits): - other = "module.decoder.layers.0.mlp.linear_fc1.weight" - _, opt_sd = handle_mla_down_proj_in_state_dict( - _fused_model(), {}, {"state": {other: {"exp_avg": "A"}}} - ) - assert opt_sd["state"][other] == {"exp_avg": "A"} - - def test_empty_optimizer_state_tolerated(self, recorded_splits): - _, opt_sd = handle_mla_down_proj_in_state_dict(_fused_model(), {}, {"state": {}}) - assert opt_sd["state"] == {} - - -# ============================================================================ -# Test get_mtp_inner_layer_paths -# ============================================================================ -class TestGetMTPInnerLayerPaths: - - def test_gpt_mtp_layer_included(self): - model = _Model(mtp_layers=[_MTPLayer(mtp_layer_pattern=None)]) - assert get_mtp_inner_layer_paths(model) == ["mtp.layers.0"] - - def test_mamba_mtp_layer_excluded(self): - """Mamba MTP keeps mtp_model_layer as its native checkpoint name.""" - model = _Model(mtp_layers=[_MTPLayer(mtp_layer_pattern="M-M-")]) - assert get_mtp_inner_layer_paths(model) == [] - - def test_model_without_mtp(self): - model = _Model(decoder_layers=[_Layer(_UnfusedMLAAttention())]) - assert get_mtp_inner_layer_paths(model) == [] - - def test_multiple_mtp_layers(self): - model = _Model(mtp_layers=[_MTPLayer(), _MTPLayer()]) - assert get_mtp_inner_layer_paths(model) == ["mtp.layers.0", "mtp.layers.1"] - - -# ============================================================================ -# Test rename_mtp_inner_layer_keys -# ============================================================================ -class TestRenameMTPInnerLayerKeys: - """The model uses mtp_model_layer; checkpoints use transformer_layer.""" - - PATHS = ["mtp.layers.0"] - - def test_key_renamed(self): - sd = {"mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight": 1} - out = rename_mtp_inner_layer_keys(sd, self.PATHS) - assert out == {"mtp.layers.0.transformer_layer.mlp.linear_fc1.weight": 1} - - def test_wrapper_prefix_preserved(self): - sd = {"module.mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight": 1} - out = rename_mtp_inner_layer_keys(sd, self.PATHS) - assert list(out) == ["module.mtp.layers.0.transformer_layer.mlp.linear_fc1.weight"] - - def test_expert_key_renamed(self): - """The MTP layer's expert weights were the bulk of the silently dropped keys.""" - sd = {"mtp.layers.0.mtp_model_layer.mlp.experts.linear_fc1.weight3": 1} - out = rename_mtp_inner_layer_keys(sd, self.PATHS) - assert list(out) == ["mtp.layers.0.transformer_layer.mlp.experts.linear_fc1.weight3"] - - def test_decoder_key_untouched(self): - sd = {"decoder.layers.0.mlp.linear_fc1.weight": 1} - assert rename_mtp_inner_layer_keys(sd, self.PATHS) == sd - - def test_mtp_key_outside_inner_layer_untouched(self): - sd = {"mtp.layers.0.enorm.weight": 1} - assert rename_mtp_inner_layer_keys(sd, self.PATHS) == sd - - def test_only_inner_layer_occurrence_renamed(self): - sd = {"mtp.layers.0.mtp_model_layer.sub.mtp_model_layer.weight": 1} - out = rename_mtp_inner_layer_keys(sd, self.PATHS) - assert list(out) == ["mtp.layers.0.transformer_layer.sub.mtp_model_layer.weight"] - - def test_unlisted_mtp_layer_untouched(self): - sd = {"mtp.layers.1.mtp_model_layer.mlp.linear_fc1.weight": 1} - assert rename_mtp_inner_layer_keys(sd, self.PATHS) == sd - - def test_input_not_mutated(self): - sd = {"mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight": 1} - rename_mtp_inner_layer_keys(sd, self.PATHS) - assert list(sd) == ["mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight"] - - -# ============================================================================ -# Test handle_mtp_in_state_dict -# ============================================================================ -class TestHandleMTPInStateDict: - - MODEL_KEY = "mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight" - RENAMED_KEY = "mtp.layers.0.transformer_layer.mlp.linear_fc1.weight" - - def test_model_state_renamed(self): - model_sd, _ = handle_mtp_in_state_dict( - _Model(mtp_layers=[_MTPLayer()]), {self.MODEL_KEY: 1}, None - ) - assert model_sd == {self.RENAMED_KEY: 1} - - def test_optimizer_state_renamed(self): - opt_key = f"module.{self.MODEL_KEY}" - _, opt_sd = handle_mtp_in_state_dict( - _Model(mtp_layers=[_MTPLayer()]), {}, {"state": {opt_key: {"exp_avg": "A"}}} - ) - assert list(opt_sd["state"]) == [f"module.{self.RENAMED_KEY}"] - - def test_noop_without_mtp(self): - model = _Model(decoder_layers=[_Layer(_UnfusedMLAAttention())]) - sd = {"decoder.layers.0.mlp.linear_fc1.weight": 1} - model_sd, opt_sd = handle_mtp_in_state_dict(model, sd, None) - assert model_sd == sd - assert opt_sd is None - - def test_noop_for_mamba_mtp(self): - model = _Model(mtp_layers=[_MTPLayer(mtp_layer_pattern="M-M-")]) - sd = {self.MODEL_KEY: 1} - model_sd, _ = handle_mtp_in_state_dict(model, sd, None) - assert model_sd == sd - - def test_empty_optimizer_state_tolerated(self): - _, opt_sd = handle_mtp_in_state_dict(_Model(mtp_layers=[_MTPLayer()]), {}, {"state": {}}) - assert opt_sd["state"] == {} - - -# ============================================================================ -# Test get_unexpected_model_keys -# ============================================================================ -class TestGetUnexpectedModelKeys: - """ "Unexpected" follows StrictHandling: requested by this rank, absent from the - checkpoint. This is the direction that silently zeroed weights.""" - - def test_missing_model_key_reported(self): - metadata = {"model": {"decoder.layers.0.mlp.linear_fc1.weight": object()}} - load = { - "model": { - "decoder.layers.0.mlp.linear_fc1.weight": 1, - "decoder.layers.0.self_attention.linear_qkv_down_proj.weight": 1, - } - } - assert get_unexpected_model_keys(metadata, load) == { - "model.decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - } - - def test_nothing_missing(self): - metadata = {"model": {"a.weight": object(), "b.weight": object()}} - load = {"model": {"a.weight": 1, "b.weight": 1}} - assert get_unexpected_model_keys(metadata, load) == set() - - def test_expert_parallel_extras_not_reported(self): - """Regression: under EP the checkpoint holds every expert but each rank owns a - slice. Reporting checkpoint-side extras would fail every EP run.""" - metadata = { - "model": { - f"decoder.layers.0.mlp.experts.linear_fc1.weight{i}": object() for i in range(8) - } - } - load = { - "model": {f"decoder.layers.0.mlp.experts.linear_fc1.weight{i}": 1 for i in range(2)} - } - assert get_unexpected_model_keys(metadata, load) == set() - - def test_optimizer_section_ignored(self): - """Optimizer state is legitimately absent when it is not being restored.""" - metadata = {"model": {"a.weight": object()}} - load = {"model": {"a.weight": 1}, "optimizer": {"state": {"a.weight": {"exp_avg": 1}}}} - assert get_unexpected_model_keys(metadata, load) == set() - - def test_rng_section_ignored(self): - metadata = {"model": {"a.weight": object()}} - load = {"model": {"a.weight": 1}, "rng_state": [1, 2, 3]} - assert get_unexpected_model_keys(metadata, load) == set() - - def test_extra_state_ignored(self): - """TE _extra_state is dropped from both sides before save and load.""" - metadata = {"model": {"a.weight": object()}} - load = {"model": {"a.weight": 1, "a._extra_state": 1}} - assert get_unexpected_model_keys(metadata, load) == set() - - def test_virtual_pipeline_model_sections_covered(self): - metadata = {"model0": {"a.weight": object()}, "model1": {}} - load = {"model0": {"a.weight": 1}, "model1": {"b.weight": 1}} - assert get_unexpected_model_keys(metadata, load) == {"model1.b.weight"} - - def test_non_model_top_level_key_ignored(self): - metadata = {"model": {}} - load = {"model": {}, "checkpoint_version": 3.0} - assert get_unexpected_model_keys(metadata, load) == set() - - -# ============================================================================ -# Test validate_fsdp_dtensor_model_load -# ============================================================================ -class TestValidateFsdpDtensorModelLoad: - - MISSING_KEY = "decoder.layers.0.self_attention.linear_qkv_down_proj.weight" - - def _mismatched(self): - return {"model": {}}, {"model": {self.MISSING_KEY: 1}} - - def _matched(self): - return {"model": {"a.weight": object()}}, {"model": {"a.weight": 1}} - - def test_raise_unexpected_raises(self): - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException) as excinfo: - validate_fsdp_dtensor_model_load( - metadata, load, "/ckpt", strict=StrictHandling.RAISE_UNEXPECTED - ) - assert self.MISSING_KEY in str(excinfo.value) - assert "/ckpt" in str(excinfo.value) - - def test_error_explains_the_cause(self): - """The message should point at sharded_state_dict(), which is the actual trap.""" - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException, match="sharded_state_dict"): - validate_fsdp_dtensor_model_load(metadata, load, "/ckpt") - - def test_raises_by_default(self): - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException): - validate_fsdp_dtensor_model_load(metadata, load, "/ckpt") - - def test_no_raise_when_complete(self): - metadata, load = self._matched() - assert validate_fsdp_dtensor_model_load(metadata, load, "/ckpt") == set() - - def test_raise_all_behaves_like_raise_unexpected(self): - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException): - validate_fsdp_dtensor_model_load( - metadata, load, "/ckpt", strict=StrictHandling.RAISE_ALL - ) - - @pytest.mark.parametrize("strict", [StrictHandling.LOG_UNEXPECTED, StrictHandling.LOG_ALL]) - def test_log_variants_do_not_raise(self, strict, caplog): - metadata, load = self._mismatched() - with caplog.at_level("WARNING"): - result = validate_fsdp_dtensor_model_load(metadata, load, "/ckpt", strict=strict) - assert result == {f"model.{self.MISSING_KEY}"} - assert self.MISSING_KEY in caplog.text - - def test_ignore_all_disables_the_check(self): - metadata, load = self._mismatched() - assert ( - validate_fsdp_dtensor_model_load( - metadata, load, "/ckpt", strict=StrictHandling.IGNORE_ALL - ) - == set() - ) - - def test_assume_ok_unexpected_raises(self): - """ASSUME_OK_UNEXPECTED defers to the underlying strategy, which never raises on a - partial DCP load, so for fsdp_dtensor it has to raise here instead. This is the - default value of --dist-ckpt-strictness, so the check must be on out of the box.""" - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException): - validate_fsdp_dtensor_model_load( - metadata, load, "/ckpt", strict=StrictHandling.ASSUME_OK_UNEXPECTED - ) - - def test_assume_ok_unexpected_quiet_when_complete(self): - metadata, load = self._matched() - assert ( - validate_fsdp_dtensor_model_load( - metadata, load, "/ckpt", strict=StrictHandling.ASSUME_OK_UNEXPECTED - ) - == set() - ) - - @pytest.mark.parametrize( - "strict", [StrictHandling.RETURN_UNEXPECTED, StrictHandling.RETURN_ALL] - ) - def test_return_variants_report_without_raising(self, strict): - metadata, load = self._mismatched() - assert validate_fsdp_dtensor_model_load(metadata, load, "/ckpt", strict=strict) == { - f"model.{self.MISSING_KEY}" - } - - def test_string_strict_flag_accepted(self): - metadata, load = self._mismatched() - with pytest.raises(CheckpointingException): - validate_fsdp_dtensor_model_load(metadata, load, "/ckpt", strict="raise_unexpected") - - def test_invalid_strict_flag_rejected(self): - metadata, load = self._matched() - with pytest.raises(ValueError): - validate_fsdp_dtensor_model_load(metadata, load, "/ckpt", strict="nonsense") - - def test_expert_parallel_load_does_not_raise(self): - """Regression: the validator must stay quiet on a healthy expert-parallel load.""" - metadata = { - "model": { - f"decoder.layers.0.mlp.experts.linear_fc1.weight{i}": object() for i in range(8) - } - } - load = { - "model": {f"decoder.layers.0.mlp.experts.linear_fc1.weight{i}": 1 for i in range(2)} - } - assert validate_fsdp_dtensor_model_load(metadata, load, "/ckpt") == set() - - def test_message_truncated(self): - metadata = {"model": {}} - load = {"model": {f"decoder.layers.{i}.weight": 1 for i in range(50)}} - with pytest.raises(CheckpointingException) as excinfo: - validate_fsdp_dtensor_model_load(metadata, load, "/ckpt", max_reported=5) - message = str(excinfo.value) - assert "and 45 more" in message - assert message.count("\n model.decoder.layers.") == 5 - - # ============================================================================ # Test checkpoint_inspector.py pretrained-only checkpoint handling # (hasattr mcore_data guard — part of PR #3912 but related) From 1f34327e6dc79e5b5f2e785f83d5a49e4d6b5c0b Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 16 Aug 2026 01:31:14 +0800 Subject: [PATCH 04/34] Add head-wise output gating to MLA Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit ed853d3b82a0448c5a09c2002d5ce6fef3449c32) --- megatron/core/models/gpt/gpt_layer_specs.py | 3 + .../transformer/multi_latent_attention.py | 158 ++++++++++------ .../core/transformer/transformer_config.py | 21 ++- .../test_multi_latent_attention.py | 172 +++++++++++++++++- 4 files changed, 289 insertions(+), 65 deletions(-) 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/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index d821c14effe..9880fd781ea 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -15,6 +15,7 @@ except ImportError: HAVE_EINOPS = False + from megatron.core import tensor_parallel from megatron.core.dist_checkpointing.mapping import ShardedObject from megatron.core.extensions.transformer_engine import HAVE_TE @@ -36,7 +37,6 @@ ) from megatron.core.transformer.attention import Attention, LinearProjBuilder from megatron.core.transformer.enums import AttnMaskType -from megatron.core.transformer.mla_qk_norm_config import QKNormConfigResolver from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.torch_norm import LayerNormBuilder from megatron.core.transformer.transformer_config import MLATransformerConfig @@ -125,6 +125,7 @@ class MLASelfAttentionSubmodules: linear_kv_up_proj: Union[ModuleSpec, type] = None linear_qkv_down_proj: Union[ModuleSpec, type] = None core_attention: Union[ModuleSpec, type] = None + linear_gate: Union[ModuleSpec, type] = None class MultiLatentAttention(Attention): @@ -144,6 +145,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ) -> None: # TODO(nschank): Restructure so that the Attention initializer knows which specific @@ -156,6 +158,7 @@ def __init__( attn_mask_type=attn_mask_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) self.config: MLATransformerConfig @@ -204,6 +207,11 @@ def __init__( "'rope' and 'yarn'" ) + if self.config.experimental_attention_variant == "dsa": + core_attn_extra_kwargs = {"is_mtp_layer": is_mtp_layer} + else: + core_attn_extra_kwargs = {} + self.core_attention = build_module( submodules.core_attention, config=self.config, @@ -215,8 +223,29 @@ def __init__( v_channels=self.config.v_head_dim, cp_comm_type=cp_comm_type, pg_collection=self.pg_collection, + **core_attn_extra_kwargs, ) + if self.config.attention_output_gate: + if submodules.linear_gate is None: + raise ValueError("MLA output gating requires a linear_gate module spec.") + self.linear_gate = build_module( + submodules.linear_gate, + self.config.hidden_size, + self.config.num_attention_heads, + 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 +354,8 @@ def forward( # hidden_states: [sq, b, h] + gate_input = hidden_states + inference_context = deprecate_inference_params(inference_context, inference_params) if inference_context and not inference_context.is_static_batching(): assert ( @@ -334,6 +365,26 @@ def forward( if self.config.cache_mla_latents: self.prepare_for_absorption() + # Set the right cp group for dynamic-cp. Mirrors Attention.forward: + # downstream RoPE uses self.pg_collection.cp, which must point at this + # microbatch's dynamic CP group. Restored before every return. + _orig_cp_group = self.pg_collection.cp + if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: + assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" + self.pg_collection.cp = packed_seq_params.cp_group + if ( + packed_seq_params is not None + and packed_seq_params.qkv_format == "thd" + and self.pg_collection.cp is not None + and get_pg_size(self.pg_collection.cp) > 1 + and packed_seq_params.cp_partition_mode != "zigzag" + ): + raise ValueError( + "MultiLatentAttention requires cp_partition_mode='zigzag', but " + f"packed_seq_params has {packed_seq_params.cp_partition_mode!r}. CP partition " + "conversion must be handled before entering MLA." + ) + # ===================== # Query, Key, and Value # ===================== @@ -377,10 +428,10 @@ def forward( # core attention computation # ================================== # Need corresponding TE change + needs_output_trim = False core_attn_manager = off_interface( self.offload_core_attention and self.training, query, "core_attn" ) - needs_output_trim = False if self.checkpoint_core_attention and self.training: core_attn_out = self._checkpointed_attention_forward( query, @@ -421,7 +472,6 @@ def forward( cu_kv_lengths, kv_lengths, block_table, - inference_context.is_decode_only(), ) # Only rearrange if not in absorption mode (Flash MLA handles format correctly) if not inference_context.is_decode_only(): @@ -457,6 +507,9 @@ def forward( self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out) self.qkv_up_checkpoint = None + if self.linear_gate is not None: + core_attn_out = self._project_and_apply_mla_output_gate(core_attn_out, gate_input) + # ================= # Output. [sq, b, h] # ================= @@ -465,8 +518,28 @@ def forward( output, bias = apply_module(self.linear_proj)(core_attn_out) output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out]) + self.pg_collection.cp = _orig_cp_group return output, bias + def _project_and_apply_mla_output_gate( + self, core_attn_out: torch.Tensor, gate_input: torch.Tensor + ) -> torch.Tensor: + """Project the gate after core attention, then apply it to the MLA output.""" + gate, _ = apply_module(self.linear_gate)(gate_input) + return self._apply_mla_output_gate(core_attn_out, gate) + + @staticmethod + def _apply_mla_output_gate(core_attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Apply one native-dtype sigmoid gate per local MLA attention head.""" + output_shape = core_attn_out.shape + core_attn_out = core_attn_out.view(*output_shape[:2], gate.size(-1), -1) + # Preserve the Tiny/native BF16 autograd contract. Casting to FP32 before + # sigmoid can produce the same BF16 forward values while changing the + # gate-input and gate-weight VJP. + gate = torch.sigmoid(gate) + core_attn_out = core_attn_out * gate.unsqueeze(-1) + return core_attn_out.reshape(output_shape) + class MLASelfAttention(MultiLatentAttention): """MLA Self-attention layer class @@ -484,6 +557,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): if pg_collection is None: @@ -498,17 +572,14 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) - # Resolve which classes to use for Q and KV linear up projections and norms, based on - # QK-norm selection. - layer_classes = self._resolve_qk_norm_config(submodules) - if self.config.q_lora_rank is None: # Not projecting query self.linear_q_proj = build_module( - layer_classes["linear_q_proj"], + submodules.linear_q_proj, self.config.hidden_size, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -555,7 +626,7 @@ def __init__( ) self.linear_q_up_proj = build_module( - layer_classes["linear_q_up_proj"], + submodules.linear_q_up_proj, self.config.q_lora_rank, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -602,7 +673,7 @@ def __init__( ) self.linear_kv_up_proj = build_module( - layer_classes["linear_kv_up_proj"], + submodules.linear_kv_up_proj, self.config.kv_lora_rank, self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim), config=self.config, @@ -617,24 +688,18 @@ def __init__( ) if self.config.q_lora_rank is not None: - self.q_layernorm = layer_classes["q_layernorm"]( + self.q_layernorm = submodules.q_layernorm( hidden_size=self.config.q_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, ) - self.kv_layernorm = layer_classes["kv_layernorm"]( + self.kv_layernorm = submodules.kv_layernorm( hidden_size=self.config.kv_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, ) - def _resolve_qk_norm_config( - self, submodules - ) -> dict[str, ModuleSpec | type | LayerNormBuilder]: - """Resolve which Q/KV norm and up-projection implementations to build.""" - return QKNormConfigResolver(self.config, submodules).resolve() - def _qkv_down_projection(self, hidden_states): """Unfused q/kv down projection path.""" if self.config.q_lora_rank is not None: @@ -681,11 +746,6 @@ def get_query_key_value_tensors( assert ( hidden_states.ndim == 3 ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" - if packed_seq_params is not None: - assert ( - packed_seq_params.local_cp_size is None - ), "hybrid_context_parallel is not supported with MLA yet and is planned for future. \ - Please disable hybrid_context_parallel." inference_context = deprecate_inference_params(inference_context, inference_params) @@ -730,14 +790,9 @@ def get_query_key_value_tensors( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv rope_max_seqlen_q = packed_seq_params.max_seqlen_q rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv - rope_freqs_max_seqlen = ( - max(rope_max_seqlen_q, rope_max_seqlen_kv) - if rope_max_seqlen_q is not None and rope_max_seqlen_kv is not None - else None - ) else: cu_seqlens_q = cu_seqlens_kv = None - rope_freqs_max_seqlen = None + rope_max_seqlen_q = rope_max_seqlen_kv = None # ========================================= # QKV down projection and layernorm @@ -762,7 +817,7 @@ def get_query_key_value_tensors( # k_pos_emb: [s, b, qk_pos_emb_head_dim] k_pos_emb = gather_from_sequence_parallel_region(k_pos_emb, group=self.tp_group) - if packed_seq_params is not None: + if thd_packed_seq: # If sequence packing, TE expect [t, h, d] shaped qkv input. # In Megatron-Core, the qkv shape is [t, 1, h, d]. # So we need to reshape qkv from [t, 1, h, d] to [t, h, d]. @@ -917,7 +972,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po sequence_start = inference_context.sequence_len_offset sequence_end = sequence_start + q_len rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end] - elif packed_seq_params is None or self.config.context_parallel_size == 1: + elif not thd_packed_seq or self.config.context_parallel_size == 1: # Shorten rotary_pos_emb to the sequence length when inference_params # is not provided. This makes sure we can run forward directly with # any sequence length. During training, the sequence length is always @@ -949,7 +1004,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, - max_seqlen=rope_freqs_max_seqlen, + max_seqlen=rope_max_seqlen_q, ) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( @@ -960,7 +1015,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, - max_seqlen=rope_freqs_max_seqlen, + max_seqlen=rope_max_seqlen_kv, ) # query: [num_tokens, n, (qk_head_dim + v_head_dim)] @@ -1092,6 +1147,15 @@ def backward_dw(self) -> NoReturn: """Execute weight gradient computation""" self._backward_kv_proj() self._backward_q_proj() + if self.linear_gate is not None: + self.linear_gate.backward_dw() + # For the 'dsa' experimental variant core_attention is DSAttention, whose + # indexer linears defer their wgrads under delay_wgrad_compute and need an + # explicit flush. For standard MLA the core is TEDotProductAttention, which + # owns no linears and defines no backward_dw hence the guard. + core_attention_backward_dw = getattr(self.core_attention, "backward_dw", None) + if core_attention_backward_dw is not None: + core_attention_backward_dw() self._backward_output_proj() def _backward_kv_proj(self): @@ -1249,6 +1313,7 @@ def __init__( attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, pp_layer_offset: Optional[int] = None, name: str | None = None, ): @@ -1264,6 +1329,7 @@ def __init__( attention_type="self", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, pp_layer_offset=pp_layer_offset, name=name, ) @@ -1272,9 +1338,6 @@ def __init__( "FusedMLASelfAttention requires q_lora_rank to be set; " "fallback to MLASelfAttention for q_lora_rank=None." ) - # Resolve which linear class to use for Q and KV up projections, - # based on QK-norm selection. - layer_classes = self._resolve_qk_norm_config(submodules) qkv_down_proj_kwargs = {} if submodules.linear_qkv_down_proj in [TELinear]: @@ -1310,7 +1373,7 @@ def __init__( ) self.linear_q_up_proj = build_module( - layer_classes["linear_q_up_proj"], + submodules.linear_q_up_proj, self.config.q_lora_rank, self.config.num_attention_heads * self.q_head_dim, config=self.config, @@ -1325,7 +1388,7 @@ def __init__( ) self.linear_kv_up_proj = build_module( - layer_classes["linear_kv_up_proj"], + submodules.linear_kv_up_proj, self.config.kv_lora_rank, self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim), config=self.config, @@ -1339,12 +1402,12 @@ def __init__( name=(name + ".linear_kv_up_proj") if name is not None else None, ) - self.q_layernorm = layer_classes["q_layernorm"]( + self.q_layernorm = submodules.q_layernorm( hidden_size=self.config.q_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, ) - self.kv_layernorm = layer_classes["kv_layernorm"]( + self.kv_layernorm = submodules.kv_layernorm( hidden_size=self.config.kv_lora_rank, config=self.config, eps=self.config.layernorm_epsilon, @@ -1428,8 +1491,7 @@ def _clone_sharded_object_with_key(obj: ShardedObject, new_key: str) -> ShardedO sharded_state_dict[kv_extra_key] = fused_obj for key in list(sharded_state_dict.keys()): - suffix = key[len(fused_prefix) :] if key.startswith(fused_prefix) else "" - if key.startswith(fused_prefix) and not suffix.startswith("layer_norm_"): + if key.startswith(fused_prefix): del sharded_state_dict[key] fused_weight = self.linear_qkv_down_proj.weight @@ -1473,12 +1535,8 @@ def _clone_sharded_object_with_key(obj: ShardedObject, new_key: str) -> ShardedO return sharded_state_dict - def _synthetic_state_dict_key_suffixes(self): - """Return source checkpoint keys used to locate this module in a state dict.""" - return ("linear_q_down_proj.weight",) - - def _synthesize_fused_qkv_down_weight(self, state_dict, prefix): - """Materialize fused qkv-down weight from old separate q/kv checkpoint keys.""" + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Load state dict with automatic unfused->fused conversion.""" q_key = f"{prefix}linear_q_down_proj.weight" kv_key = f"{prefix}linear_kv_down_proj.weight" fused_key = f"{prefix}linear_qkv_down_proj.weight" @@ -1495,8 +1553,4 @@ def _as_tensor(x): state_dict.pop(f"{prefix}linear_q_down_proj.bias", None) state_dict.pop(f"{prefix}linear_kv_down_proj.bias", None) - def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): - """Load state dict with automatic unfused->fused conversion.""" - self._synthesize_fused_qkv_down_weight(state_dict, prefix) - return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 2d30c38e6da..33d7408dfb7 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -263,15 +263,16 @@ class TransformerConfig(ModelParallelConfig): defualts to False. Setting qk_clip will automatically log the max logit""" attention_output_gate: bool = False - """Apply a full head_dim output gate (num_attention_heads * head_dim rows - fused inline per-group [q, gate, k, v] into linear_qkv). Mutually - exclusive with `head_wise_attn_gate` (per-head scalar gate).""" + """Whether to apply output gating to attention layers. + For MLA, this uses one sigmoid gate per attention head. For standard attention, + this retains the full head-dimension output gate behavior. It is mutually exclusive + with ``head_wise_attn_gate``. + """ 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 @@ -3785,8 +3786,16 @@ def __post_init__(self): ): raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.") - if self.attention_output_gate: - raise NotImplementedError("Output gate is not supported for MLA yet.") + if self.attention_output_gate and self.mla_down_proj_fusion: + # Fused MLA hides the post-input-LayerNorm activation inside the fused + # LayerNorm+linear module. Gated MLA must consume that activation as the + # gate input; using raw hidden_states would silently change the model. + # Keep this combination fail-fast until the fused API exposes the + # normalized activation. + raise ValueError( + "MLA output gating does not support fused down projections; " + "disable mla_down_proj_fusion to use the unfused path." + ) # DSv4 hybrid: derive qk_head_dim and kv_lora_rank from v_head_dim and qk_pos_emb_head_dim if self.experimental_attention_variant == "dsv4_hybrid": diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index 3131222f800..4c2327731aa 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/tests/unit_tests/transformer/test_multi_latent_attention.py @@ -741,6 +741,159 @@ def test_gpu_forward_thd(self): assert bias.shape[0] == config.hidden_size +class TestMLAOutputGate: + + @pytest.fixture(scope='function', autouse=True) + def setup_and_teardown(self): + 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, + ) + 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_head_wise_gate(self): + gate = self.parallel_attention.linear_gate + assert gate is not None + assert gate.weight.shape == ( + self.transformer_config.num_attention_heads, + self.transformer_config.hidden_size, + ) + assert gate.bias is None or gate.bias.numel() == 0 + + def test_gate_is_applied_per_head(self): + core_attn_out = torch.arange(24, dtype=torch.bfloat16).view(2, 1, 12) + gate = torch.tensor([[[-2.0, -1.0, 0.0, 1.0]], [[2.0, 3.0, 4.0, 5.0]]]).bfloat16() + + output = self.parallel_attention._apply_mla_output_gate(core_attn_out, gate) + expected = core_attn_out.view(2, 1, 4, 3) * torch.sigmoid(gate).unsqueeze(-1) + + torch.testing.assert_close(output, expected.reshape_as(output), atol=0, rtol=0) + + def test_gate_preserves_native_dtype_vjp(self): + if not torch.cuda.is_available(): + pytest.skip("native BF16 VJP check requires CUDA") + + generator = torch.Generator(device="cuda").manual_seed(20260814) + gate_input = torch.randn( + (8, 1, 4), generator=generator, device="cuda", dtype=torch.bfloat16 + ) + core_output = torch.randn( + (8, 1, 32), generator=generator, device="cuda", dtype=torch.bfloat16 + ) + output_gradient = torch.randn( + core_output.shape, generator=generator, device="cuda", 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 + ) + + reference_gate = gate_input.detach().clone().requires_grad_(True) + reference_core = core_output.detach().clone().requires_grad_(True) + reference_output = ( + reference_core.view(8, 1, 4, 8) * torch.sigmoid(reference_gate).unsqueeze(-1) + ).reshape_as(reference_core) + reference_gradients = torch.autograd.grad( + reference_output, (reference_gate, reference_core), output_gradient + ) + + torch.testing.assert_close(actual_output, reference_output, atol=0, rtol=0) + for actual, reference in zip(actual_gradients, reference_gradients): + torch.testing.assert_close(actual, reference, atol=0, rtol=0) + + 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): + 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, + 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 +1591,23 @@ def test_clip_qk_with_absorption_raises_error(self): ], ) @pytest.mark.parametrize( - ("tp", "sp", "cp"), + ("tp", "sp", "cp", "output_gate"), [ - (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), # TP w/o SP + (4, True, 1, False), # TP w/ SP + (1, False, 4, False), # CP + (2, False, 2, False), # CP + TP w/o SP + (2, True, 2, False), # CP + TP w/ SP + (4, True, 1, True), # Gated MLA with TP + SP + (1, False, 2, True), # Gated MLA 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 ): + 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 +1681,7 @@ def initialize_gpt_model( bf16=True, rope_type=rope_type, apply_rope_fusion=apply_rope_fusion, + attention_output_gate=output_gate, hidden_dropout=0.0, attention_dropout=0.0, ) From dc1ca15e8999be3a0d8762bd818d68de00485cb4 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 16 Aug 2026 01:31:14 +0800 Subject: [PATCH 05/34] Include nested MTP modules in QK clipping Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit fe92b2dc6eea23f5350dd08053bcad19060c0ad8) --- megatron/core/optimizer/qk_clip.py | 71 +++++++++------- tests/unit_tests/optimizer/test_qk_clip.py | 98 ++++++++++++++++++++++ 2 files changed, 141 insertions(+), 28 deletions(-) create mode 100644 tests/unit_tests/optimizer/test_qk_clip.py 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/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 From 98d52cf5d1d9f9518986bcd691308eff9f86bd9d Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 16 Aug 2026 01:31:15 +0800 Subject: [PATCH 06/34] Add a Ling-V3 Tiny hybrid model spec Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 2ec22972aa54d5c3ff2fdbf620ae4d7fcaf30d8f) --- examples/ling_v3_tiny/model_spec.py | 25 + megatron/core/models/hybrid/hybrid_block.py | 750 +++++++++++++++--- .../models/hybrid/hybrid_layer_allocation.py | 9 +- .../core/models/hybrid/hybrid_layer_specs.py | 22 + megatron/training/arguments.py | 2 +- .../models/test_ling_v3_tiny_model_spec.py | 205 +++++ tests/unit_tests/ssm/test_hybrid_block.py | 2 +- .../ssm/test_hybrid_layer_allocation.py | 26 +- 8 files changed, 906 insertions(+), 135 deletions(-) create mode 100644 examples/ling_v3_tiny/model_spec.py create mode 100644 tests/unit_tests/models/test_ling_v3_tiny_model_spec.py diff --git a/examples/ling_v3_tiny/model_spec.py b/examples/ling_v3_tiny/model_spec.py new file mode 100644 index 00000000000..158714e98e8 --- /dev/null +++ b/examples/ling_v3_tiny/model_spec.py @@ -0,0 +1,25 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from copy import deepcopy + +from megatron.core.extensions.transformer_engine import TEColumnParallelLinear +from megatron.core.models.hybrid.hybrid_layer_specs import ( + hybrid_stack_spec as default_hybrid_stack_spec, +) +from megatron.core.tensor_parallel import ColumnParallelLinear + +LING_V3_TINY_HYBRID_LAYER_PATTERN = "K-KEKE+EKEKEKE+EKEKEKE+EKEKEKE+EKEKEKE+EKEKEKE+E/+E" + + +# Ling-V3 Tiny uses tensor-parallel projections for the KDA beta path and the +# MLA latent down projections and output gate. Keep the default hybrid spec +# unchanged and expose this model recipe through the standard --spec hook. +hybrid_stack_spec = deepcopy(default_hybrid_stack_spec) + +_kda_submodules = hybrid_stack_spec.submodules.kda_layer.submodules.self_attention.submodules +_kda_submodules.beta_proj = ColumnParallelLinear + +_mla_submodules = hybrid_stack_spec.submodules.mla_layer.submodules.self_attention.submodules +_mla_submodules.linear_q_down_proj = TEColumnParallelLinear +_mla_submodules.linear_kv_down_proj = TEColumnParallelLinear +_mla_submodules.linear_gate = ColumnParallelLinear diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index fb455ad9ede..087ca8da675 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -1,11 +1,10 @@ -# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2024, Tri Dao, Albert Gu. # Some of this code was adopted from https://github.com/state-spaces/mamba/ # This source code is licensed under the Apache license found in the # LICENSE file in the root directory of this source tree. -import copy from contextlib import nullcontext from dataclasses import dataclass from typing import List, Optional, Tuple, Union @@ -13,30 +12,33 @@ import torch from torch import Tensor, nn -from megatron.core.context_parallel import ContextParallelLayoutManager from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding from megatron.core.enums import Fp8Recipe -from megatron.core.extensions.transformer_engine import TELayerNormColumnParallelLinear, TENorm +from megatron.core.extensions.transformer_engine import TENorm from megatron.core.fp4_utils import get_fp4_context from megatron.core.fp8_utils import get_fp8_context from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.inference.utils import InferenceMode from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols as LayerSymbols -from megatron.core.models.hybrid.layers.hybrid_hyper_connection import HyperConnectionHybridLayer from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward -from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager +from megatron.core.tensor_parallel.random import MHCCheckpointManager from megatron.core.transformer import TransformerConfig from megatron.core.transformer.cuda_graphs import annotate_first_last_layer +from megatron.core.transformer.enums import CudaGraphModule from megatron.core.transformer.hyper_connection import ( HyperConnectionModule, learned_output_contract, ) from megatron.core.transformer.identity_op import IdentityOp -from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 -from megatron.core.transformer.multi_latent_attention import FusedMLASelfAttention +from megatron.core.transformer.module import ( + GraphableMegatronModule, + MegatronModule, + convert_module_to_dtype_except_fp32_marked, + mark_keep_in_fp32, +) from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_layer import TransformerLayer from megatron.core.transformer.utils import ( @@ -59,11 +61,487 @@ class HybridStackSubmodules: attention_layer: Union[ModuleSpec, type] = IdentityOp dsa_layer: Union[ModuleSpec, type] = IdentityOp mla_layer: Union[ModuleSpec, type] = IdentityOp + csa_layer: Union[ModuleSpec, type] = IdentityOp + hca_layer: Union[ModuleSpec, type] = IdentityOp + window_layer: Union[ModuleSpec, type] = IdentityOp mlp_layer: Union[ModuleSpec, type] = IdentityOp moe_layer: Union[ModuleSpec, type] = IdentityOp mtp_block_spec: Optional[ModuleSpec] = None +class HyperConnectionHybridLayer(GraphableMegatronModule): + """Layer-boundary mHC wrapper for HybridStack layers. + + Hybrid layers already own their local residual paths. For this initial + integration we treat each hybrid layer as a single function by aggregating + n streams to the layer input, running the existing layer, and feeding only + the layer delta back through mHC expansion. The expansion path intentionally + uses zero additional dropout because the wrapped hybrid layer has already + applied its local dropout/residual update before the delta is computed. + + Checkpoint compatibility: this is a *wrapper* (the inner layer is held as + `self.inner_layer`), so wrapped-layer state_dict keys are nested under + `inner_layer.` (e.g. `layers.0.inner_layer.input_layernorm.weight` instead + of `layers.0.input_layernorm.weight`). HybridStack checkpoints saved with + `enable_hyper_connections=False` cannot be loaded into a model with + `enable_hyper_connections=True` (and vice versa) without a key-mapping + migration. Note: this differs from `HyperConnectionTransformerLayer`, + which subclasses `TransformerLayer` and only adds new sibling fields, + keeping all base keys stable. + + CUDA graphs: this wrapper subclasses ``GraphableMegatronModule`` so that, with + ``cuda_graph_impl="transformer_engine"``, wrapped layers are captured per-layer — + mirroring ``HyperConnectionTransformerLayer`` on the GPT path. Without this, the TE + graph discovery (``_layer_is_graphable``) only inspects the top-level layer type and + silently skips every wrapped layer, so an mHC-enabled HybridStack would run entirely + eager. Two capture modes: + + * Non-MoE inner layers (attention variants, Mamba): the whole wrapper forward + (mHC aggregate + inner layer + n-stream BDA) is captured as one graph. The inner + layer's own ``__call__`` graph routing is bypassed during capture (see + ``_call_inner_layer``) to avoid nested capture. + * MoE inner layers, when ``moe_router`` is in ``cuda_graph_modules``: the expert + all-to-all is not graph-safe, so only the deterministic prefix is graphed (mHC + ``compute_mappings``/``aggregate`` + the inner layer's router/preprocess). The graph + outputs the router intermediates, the mHC state (``h_post``, ``h_res``) and the + n-stream residual; on replay the experts run eagerly and the n-stream BDA (eager) + consumes the inner's raw ``mlp_output_with_bias`` as the layer delta. Routing the + residual *through the graph* (not reusing the layer input directly in the eager BDA) + keeps the backward gradient flowing into the captured graph, which is required for + bit-identical training — again mirroring ``HyperConnectionTransformerLayer``. + + ``_get_submodules_under_cudagraphs`` returns the submodules whose params the wrapper + graph's manual hooks must drive: ``[self]`` for whole-wrapper capture, or the mHC module + + the inner router/preprocess submodules for partial MoE capture (experts stay eager). + """ + + def __init__(self, config: TransformerConfig, layer: MegatronModule) -> None: + super().__init__(config=config) + if ( + config.cuda_graph_impl in ("transformer_engine", "full_iteration") + and config.recompute_granularity == "selective" + and "mhc" in (config.recompute_modules or []) + ): + raise ValueError( + "mHC selective recompute with CUDA Graphs (cuda_graph_impl=" + f"{config.cuda_graph_impl!r}) is not supported for HybridStack mHC " + "layers. For per-layer Transformer Engine capture, the hybrid wrapper " + "captures the mHC producer inside the graph, so per-microbatch " + "checkpoint registration cannot run; the guarded attention-only split " + "exists only on the GPT HyperConnectionTransformerLayer path. " + "Full-iteration capture is rejected here because it has only been " + "validated on that same GPT path -- the config-level gate that admits " + "it is not model-family aware. Disable CUDA graphs or remove 'mhc' " + "from recompute_modules." + ) + self.inner_layer = layer + self.layer_number = layer.layer_number + self.hyper_connection = HyperConnectionModule(config=config, layer_number=self.layer_number) + if config.params_dtype is not None: + convert_module_to_dtype_except_fp32_marked(self.hyper_connection, config.params_dtype) + if hasattr(layer, 'tp_group'): + self.tp_group = layer.tp_group + + def get_layer_static_inputs(self, seq_length, micro_batch_size): + """Override to produce n-stream hidden_states of shape [s, b, n*C]. + + CUDA graph capture allocates static buffers sized by this method. The base + returns [s, b, C], but mHC layers carry n-stream hidden states [s, b, n*C]. + Mirrors ``HyperConnectionTransformerLayer.get_layer_static_inputs``. + """ + if hasattr(self.inner_layer, "get_layer_static_inputs"): + static_inputs = self.inner_layer.get_layer_static_inputs(seq_length, micro_batch_size) + else: + static_inputs = super().get_layer_static_inputs(seq_length, micro_batch_size) + hs = static_inputs["hidden_states"] + n = self.config.num_residual_streams + static_inputs["hidden_states"] = torch.ones( + (hs.shape[0], hs.shape[1], n * self.config.hidden_size), + dtype=hs.dtype, + requires_grad=hs.requires_grad, + device=hs.device, + ) + return static_inputs + + @staticmethod + def _decompose_packed_seq_params_to_kwargs(kwargs): + """Decompose PackedSeqParams into tensor kwargs for TE CUDA graphs.""" + packed_seq_params = kwargs.pop('packed_seq_params', None) + if packed_seq_params is None: + return + kwargs['cu_seqlens_q'] = packed_seq_params.cu_seqlens_q + kwargs['cu_seqlens_kv'] = packed_seq_params.cu_seqlens_kv + kwargs['cu_seqlens_q_padded'] = packed_seq_params.cu_seqlens_q_padded + kwargs['cu_seqlens_kv_padded'] = packed_seq_params.cu_seqlens_kv_padded + + def _reconstruct_packed_seq_params_from_kwargs(self, kwargs): + """Reconstruct THD PackedSeqParams from tensor kwargs in the graph capture path.""" + if 'cu_seqlens_q' not in kwargs: + return + max_seqlen = self.config.max_seqlen_per_dp_cp_rank * self.config.context_parallel_size + packed_seq_params = PackedSeqParams( + qkv_format='thd', + cp_partition_mode=self.config.cp_partition_mode, + cu_seqlens_q=kwargs.pop('cu_seqlens_q'), + cu_seqlens_kv=kwargs.pop('cu_seqlens_kv'), + cu_seqlens_q_padded=kwargs.pop('cu_seqlens_q_padded'), + cu_seqlens_kv_padded=kwargs.pop('cu_seqlens_kv_padded'), + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + # This Python flag is baked into the captured graph and cannot vary + # between replay batches. Use the conservative THD-safe branch. + pad_between_seqs=True, + ) + kwargs['packed_seq_params'] = packed_seq_params + + def __call__(self, *args, **kwargs): + # Keep non-tensor recompute state out of TE CUDA graph inputs; the GPT + # hyper-connection path follows the same pattern. + self._mhc_recompute_manager = kwargs.pop("mhc_recompute_manager", None) + return super().__call__(*args, **kwargs) + + def _inner_is_moe(self) -> bool: + """True when the inner layer is an MoE ``TransformerLayer``. Such layers use the + GPT-style raw-delta path (feed the inner's ``mlp_output_with_bias`` straight to the + n-stream BDA) in both the eager forward and the CUDA-graph replay.""" + from megatron.core.transformer.moe.moe_layer import MoELayer + + return isinstance(self.inner_layer, TransformerLayer) and isinstance( + getattr(self.inner_layer, 'mlp', None), MoELayer + ) + + def _inner_is_partial_moe_capture(self) -> bool: + """True when the inner layer is MoE and the configured ``cuda_graph_modules`` request + partial MoE capture (``moe_router``). + + In that case the wrapper does NOT capture the whole forward as one graph (the expert + all-to-all is not graph-safe). Instead it graphs the deterministic prefix (mHC aggregate + + the inner layer's router/preprocess) and runs the experts + mHC BDA eagerly — mirroring + how ``HyperConnectionTransformerLayer`` graphs MoE layers on the GPT path. Whole-wrapper + capture is still used for non-MoE inner layers (attention variants, Mamba). + """ + return ( + self._inner_is_moe() + and bool(self.config.cuda_graph_modules) + and CudaGraphModule.moe_router in self.config.cuda_graph_modules + ) + + def _te_cuda_graph_capture(self, *args, **kwargs): + """Capture the graph-safe portion of the wrapper forward. + + For non-MoE inner layers the whole wrapper forward (mHC aggregate + inner layer + + n-stream BDA) is captured as one graph. For MoE inner layers under ``moe_router`` + partial capture, only the deterministic prefix is graphed: the mHC + ``compute_mappings``/``aggregate`` followed by the inner layer's router/preprocess. + The captured outputs are the inner router/preprocess intermediates plus the mHC + state (``h_post``, ``h_res``) and the aggregated single-stream input needed to + reconstruct the layer delta on replay. ``context`` is ``None`` for the graphed + hybrid layer types, so it is dropped (a tuple containing ``None`` cannot be a + CUDA-graph output). + """ + self._reconstruct_packed_seq_params_from_kwargs(kwargs) + + if self._inner_is_partial_moe_capture(): + hidden_states = args[0] if args else kwargs["hidden_states"] + aggregated, h_res, h_post, residual = self.hyper_connection(hidden_states) + inner_kwargs = dict(kwargs) + inner_kwargs.pop("hidden_states", None) + inner_out = list(self.inner_layer._te_cuda_graph_capture(aggregated, **inner_kwargs)) + # inner_out = router/preprocess intermediates ending in the inner residual; + # append the mHC state AND the n-stream `residual` returned by the (graphed) + # hyper_connection. Routing `residual` through the graph as an output keeps its + # backward grad flowing into the graph's backward (mirrors + # HyperConnectionTransformerLayer), instead of a second autograd path the captured + # backward does not account for. The experts' raw mlp_output_with_bias (produced + # on replay) is the layer delta, so `aggregated` need not be captured. + return tuple(inner_out) + (h_post, h_res, residual) + + hidden_states, context = self.forward(*args, **kwargs) + cuda_graph_outputs = [hidden_states] + if context is not None: + cuda_graph_outputs.append(context) + return tuple(cuda_graph_outputs) + + def _te_cuda_graph_replay(self, *args, **kwargs): + """Replay the captured graph and restore the (hidden_states, context) contract. + + Non-MoE inner layers: the whole wrapper forward was captured, so the only graph + output is the layer's n-stream hidden_states; re-append ``None`` for context. + + MoE inner layers (partial capture): replay the graphed prefix, then run the + experts eagerly and apply the mHC n-stream BDA — reproducing exactly the eager + wrapper tail (``layer_delta = layer_output - aggregated`` then + ``fused_h_res_h_post_bda``), just with the deterministic prefix graphed. + """ + self._decompose_packed_seq_params_to_kwargs(kwargs) + + if self._inner_is_partial_moe_capture(): + out = list(super()._te_cuda_graph_replay(*args, **kwargs)) + residual = out.pop() # n-stream [s, b, n*C] — graph output (see capture) + h_res = out.pop() + h_post = out.pop() + # Resume the inner MoE experts eagerly to the raw delta (mlp_output_with_bias), + # then let the n-stream BDA own the residual — identical to the eager forward + # (`_call_inner_transformer_layer_without_local_bda` fast path → + # fused_h_res_h_post_bda), just with the router/preprocess prefix graphed. + # Mirror the eager fast path's BDA args (it feeds the inner layer's + # `hidden_dropout` / `bias_dropout_fusion`) so replay == eager bit-for-bit. + mlp_output_with_bias = self.inner_layer.resume_moe_experts_after_partial_cudagraph(out) + hidden_states = self.hyper_connection.fused_h_res_h_post_bda( + h_res, + residual, + h_post, + mlp_output_with_bias, + dropout_prob=self.inner_layer.hidden_dropout, + training=self.training, + fused=self.inner_layer.config.bias_dropout_fusion, + manager=None, + ) + if ( + self.config.fp32_residual_connection + and self.config.params_dtype is not None + and hidden_states.dtype != self.config.params_dtype + ): + hidden_states = hidden_states.to(self.config.params_dtype) + return hidden_states, None + + cuda_graph_output = list(super()._te_cuda_graph_replay(*args, **kwargs)) + return cuda_graph_output[0], None + + def _get_submodules_under_cudagraphs(self): + """Submodules whose params are driven by the wrapper graph's manual hooks. + + Whole-wrapper capture covers the entire wrapper (``[self]``, the base default). + For partial MoE capture only the graphed prefix is covered — the mHC module plus + the inner layer's router/preprocess submodules — so the experts (run eagerly) + keep their normal forward hooks. + """ + if self._inner_is_partial_moe_capture(): + return [self.hyper_connection] + self.inner_layer._get_submodules_under_cudagraphs() + return super()._get_submodules_under_cudagraphs() + + def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]: + """Delegate Mamba inference state shape requests to the wrapped layer.""" + if not hasattr(self.inner_layer, 'mamba_state_shapes_per_request'): + return None + return self.inner_layer.mamba_state_shapes_per_request() + + def _call_inner_layer( + self, + hidden_states: Tensor, + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext], + rotary_pos_emb: Optional[Tensor], + sequence_len_offset: Optional[Tensor], + packed_seq_params: Optional[PackedSeqParams], + padding_mask: Optional[Tensor], + input_ids: Optional[Tensor] = None, + ) -> Tuple[Tensor, Optional[Tensor]]: + # When this wrapper is itself being CUDA-graph captured, the inner layer + # must run as a plain forward: routing through its ``__call__`` would + # trigger nested TE graph capture (the inner layer is also a + # GraphableMegatronModule). During eager steps we keep ``__call__`` so the + # inner layer's forward pre-hooks (e.g. param all-gather) fire normally; + # under graph replay these are driven by the wrapper's manual hooks. + from megatron.core.transformer.cuda_graphs import is_graph_capturing + + inner = self.inner_layer.forward if is_graph_capturing() else self.inner_layer + + if isinstance(self.inner_layer, TransformerLayer): + output = inner( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + sequence_len_offset=sequence_len_offset, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + input_ids=input_ids, + _called_from_hybrid_mhc_wrapper=True, + ) + else: + # Non-transformer layers (e.g. MambaLayer; GatedDeltaNet which does + # accept `sequence_len_offset` is currently always wrapped inside a + # TransformerLayer spec, so it takes the branch above) do not accept + # rotary_pos_emb / sequence_len_offset / padding_mask — pass only + # the common arguments. New layer types that consume any of these + # must add explicit handling here. + output = inner( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + ) + + if isinstance(output, tuple): + context = output[1] if len(output) > 1 else None + return output[0], context + return output, None + + def _call_inner_transformer_layer_without_local_bda( + self, + hidden_states: Tensor, + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext], + rotary_pos_emb: Optional[Tensor], + sequence_len_offset: Optional[Tensor], + packed_seq_params: Optional[PackedSeqParams], + padding_mask: Optional[Tensor], + input_ids: Optional[Tensor] = None, + ) -> Optional[Tuple[Tuple[Tensor, Optional[Tensor]], Optional[Tensor], float, bool]]: + """Return a raw TransformerLayer branch output when the wrapped layer is split. + + Hybrid DSv4 layers are usually attention-only (`W/C/H/D`) or MLP/MoE-only (`-/E`) + TransformerLayer instances. For those layers, skip the inner layer's local + residual+BDA and feed the raw branch output directly into the mHC BDA, matching the + GPT mHC path and avoiding a residual add followed by `layer_output - aggregated`. + """ + if not isinstance(self.inner_layer, TransformerLayer): + return None + + layer = self.inner_layer + if (not layer.training) and layer.config.inference_fuse_tp_communication: + return None + + has_attention = not isinstance(layer.self_attention, IdentityOp) + has_cross_attention = not isinstance(layer.cross_attention, IdentityOp) + has_mlp = not isinstance(layer.mlp, IdentityOp) + + if has_cross_attention or has_attention == has_mlp: + return None + + if has_attention: + output_with_bias, attn_norm_manager, residual = ( + layer._forward_self_attention_output_with_bias( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + ) + ) + output_with_bias = layer._group_offload_output_with_bias( + output_with_bias, attn_norm_manager, forced_released_tensors=[residual] + ) + return output_with_bias, None, layer.hidden_dropout, layer.config.bias_dropout_fusion + + output_with_bias, residual = layer._forward_mlp_output_with_bias( + hidden_states, + inference_context=inference_context, + padding_mask=padding_mask, + input_ids=input_ids, + packed_seq_params=packed_seq_params, + ) + if layer.mlp_norm_manager is not None: + output_with_bias = layer._group_offload_output_with_bias( + output_with_bias, layer.mlp_norm_manager, forced_released_tensors=[residual] + ) + layer.mlp_norm_manager = None + return output_with_bias, None, layer.hidden_dropout, layer.config.bias_dropout_fusion + + def forward( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor] = None, + inference_context: Optional[BaseInferenceContext] = None, + rotary_pos_emb: Optional[Tensor] = None, + sequence_len_offset: Optional[Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + padding_mask: Optional[Tensor] = None, + input_ids: Optional[Tensor] = None, + mhc_recompute_manager=None, + ) -> Tuple[Tensor, Optional[Tensor]]: + """Run the wrapped hybrid layer through one layer-boundary mHC update. + + ``attention_mask`` defaults to ``None`` so that CUDA-graph capture, which + calls this forward with only the static ``hidden_states`` input, does not + fail on a missing positional argument (causal masking is inferred by the + attention backend when the mask is ``None``). + """ + + if mhc_recompute_manager is None: + mhc_recompute_manager = getattr(self, '_mhc_recompute_manager', None) + + aggregated, h_res, h_post, residual = self.hyper_connection( + hidden_states, mhc_recompute_manager=mhc_recompute_manager + ) + fast_path_result = self._call_inner_transformer_layer_without_local_bda( + aggregated, + attention_mask, + inference_context, + rotary_pos_emb, + sequence_len_offset, + packed_seq_params, + padding_mask, + input_ids, + ) + + if fast_path_result is None: + layer_output, context = self._call_inner_layer( + aggregated, + attention_mask, + inference_context, + rotary_pos_emb, + sequence_len_offset, + packed_seq_params, + padding_mask, + input_ids, + ) + # The inner hybrid layer already applied its own local residual/dropout, so + # it returns `aggregated + f(aggregated)`. We feed only the function + # delta `f(aggregated)` into the n-stream BDA so it does not double-count + # the residual that mHC owns. + if self.config.fp32_residual_connection and aggregated.dtype != layer_output.dtype: + aggregated = aggregated.to(layer_output.dtype) + layer_output_with_bias = (layer_output - aggregated, None) + dropout_prob = 0.0 + bias_dropout_fusion = False + else: + layer_output_with_bias, context, dropout_prob, bias_dropout_fusion = fast_path_result + + layer_output = layer_output_with_bias[0] + # Sanity check: this contract requires the branch output to preserve shape; + # any mismatch indicates a future layer type is breaking the residual assumption + # and would silently corrupt the n-stream state. + if layer_output.shape != aggregated.shape: + raise RuntimeError( + "HyperConnectionHybridLayer requires wrapped branches to preserve " + f"hidden-state shape. Got {tuple(layer_output.shape)} from wrapped branch " + f"vs {tuple(aggregated.shape)} input." + ) + is_last_in_recompute_block = bool( + mhc_recompute_manager is not None + and getattr(mhc_recompute_manager, "is_last_layer_in_recompute_block", False) + ) + mhc_bda_manager = None if is_last_in_recompute_block else mhc_recompute_manager + + hidden_states = self.hyper_connection.fused_h_res_h_post_bda( + h_res, + residual, + h_post, + layer_output_with_bias, + dropout_prob=dropout_prob, + training=self.training, + fused=bias_dropout_fusion, + manager=mhc_bda_manager, + ) + # In `HyperConnectionTransformerLayer` the n-stream output stays in compute + # dtype because the post-attention `x` is in compute dtype. In the hybrid + # wrapper, `layer_delta` may be fp32 (when `fp32_residual_connection=True` + # or an inner layer upcasts), so `fused_h_res_h_post_bda`'s `output.to(x.dtype)` + # would leave the result in fp32 and silently propagate fp32 n-stream + # hidden states to every subsequent layer (~2x activation memory). Restore + # the compute-dtype contract here. + if ( + self.config.fp32_residual_connection + and self.config.params_dtype is not None + and hidden_states.dtype != self.config.params_dtype + ): + hidden_states = hidden_states.to(self.config.params_dtype) + return hidden_states, context + + class HybridStack(MegatronModule): """ Constructor for the HybridStack class. @@ -87,6 +565,7 @@ class HybridStack(MegatronModule): pg_collection (ProcessGroupCollection): the required model communication process groups to use. is_mtp_layer (bool, optional): whether this is an MTP layer. Defaults to False. + mtp_layer_number (int, optional): enclosing MTP depth for logging nested MTP metrics. """ def __init__( @@ -102,6 +581,7 @@ def __init__( dtype=None, pg_collection: ProcessGroupCollection = None, is_mtp_layer: bool = False, + mtp_layer_number: Optional[int] = None, name: str | None = None, ) -> None: """ @@ -113,18 +593,19 @@ def __init__( self.post_layer_norm = post_layer_norm self.post_process = post_process self.is_mtp_layer = is_mtp_layer + self.mtp_layer_number = mtp_layer_number assert pg_collection is not None, "pg_collection must be provided for HybridStack" self.pp_group = pg_collection.pp self.tp_group = pg_collection.tp - self.cp_group = pg_collection.cp - self.tp_cp_group = pg_collection.tp_cp # Required for pipeline parallel schedules self.input_tensor = None self.pg_collection = pg_collection + # Lazily populated mHC recompute layout cache (deterministic from config + # and num_layers); see `_build_mhc_recompute_layer_plan`. self._mhc_block_end_plan: Optional[List[bool]] = None assert layer_type_list is not None, ( @@ -132,26 +613,6 @@ def __init__( "--hybrid-layer-pattern by HybridModel." ) self.layer_type_list = layer_type_list - self._cp_layout_manager = None - if self.cp_group.size() > 1: - layer_layouts = tuple( - ( - self.config.attention_cp_layout - if layer_type in LayerSymbols.ATTENTION_LAYERS - else self.config.linear_cp_layout - ) - for layer_type in self.layer_type_list - ) - self._cp_layout_manager = ContextParallelLayoutManager( - layer_layouts=layer_layouts, - boundary_layout=self.config.linear_cp_layout, - sequence_parallel=self.config.sequence_parallel, - cp_group=self.cp_group, - tp_group=self.tp_group, - tp_cp_group=self.tp_cp_group, - ) - if getattr(self.config, "mla_down_proj_fusion", False): - submodules = self._fuse_mla_down_proj(submodules) # Build layers from the pre-selected segment self.layers = nn.ModuleList() @@ -204,6 +665,41 @@ def __init__( is_mtp_layer=is_mtp_layer, add_layer_offset=False, pp_layer_offset=pp_layer_offset, + name=(name + f".layers.{i}") if name is not None else None, + ) + elif layer_type == LayerSymbols.CSA: + # DSv4 Compressed Sparse Attention (compress_ratio fixed by the spec). + layer = build_module( + submodules.csa_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + ) + elif layer_type == LayerSymbols.HCA: + # DSv4 Heavily Compressed Attention (compress_ratio fixed by the spec). + layer = build_module( + submodules.hca_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + ) + elif layer_type == LayerSymbols.WINDOW: + # DSv4 sliding-window-only attention (compress_ratio=0 fixed by the spec; + # no compressor / no top-k indexer — attends only within csa_window_size). + layer = build_module( + submodules.window_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, ) elif layer_type == LayerSymbols.MLP: layer = build_module( @@ -225,22 +721,13 @@ def __init__( name=(name + f".layers.{i}") if name is not None else None, ) elif layer_type == LayerSymbols.GDN: - gdn_layer_spec = submodules.gdn_layer - if self.config.experimental_attention_variant == "gdn2": - # 'G' layers build the GDN2 variant when the gdn2 experimental - # attention variant is selected. - from megatron.core.ssm.gated_delta_net import GatedDeltaNet2 - - gdn_layer_spec = copy.deepcopy(gdn_layer_spec) - gdn_layer_spec.submodules.self_attention.module = GatedDeltaNet2 layer = build_module( - gdn_layer_spec, + submodules.gdn_layer, config=self.config, layer_number=layer_number, pg_collection=pg_collection, # Set to False as we do not want to change offset. add_layer_offset=False, - pp_layer_offset=pp_layer_offset, name=(name + f".layers.{i}") if name is not None else None, ) elif layer_type == LayerSymbols.KDA: @@ -254,7 +741,9 @@ def __init__( ) else: raise ValueError("unexpected layer_type") - if self.config.enable_mhc_connections: + if self.is_mtp_layer and self.mtp_layer_number is not None: + self._set_mtp_layer_number_for_moe_metrics(layer, self.mtp_layer_number) + if self.config.enable_hyper_connections: layer = HyperConnectionHybridLayer(config=self.config, layer=layer) self.layers.append(layer) @@ -272,8 +761,12 @@ def __init__( eps=self.config.layernorm_epsilon, ) - if self.config.enable_mhc_connections and self.post_process and not self.is_mtp_layer: - hc_mult = self.config.mhc_num_residual_streams + # Skip hc_head_* params inside the nested MTP HybridStack — `forward()` + # no longer calls `learned_output_contract` there (MTP owns that), so these + # params would be orphaned and break DDP's per-param grad-ready accounting + # with a `len(per_param_grad_ready_counts) != len(params)` AssertionError. + if self.config.enable_hyper_connections and self.post_process and not self.is_mtp_layer: + hc_mult = self.config.num_residual_streams hc_dim = self.config.hidden_size * hc_mult self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim))) self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult))) @@ -284,25 +777,15 @@ def __init__( setattr(self.hc_head_base, 'sequence_parallel', True) setattr(self.hc_head_scale, 'sequence_parallel', True) - def _fuse_mla_down_proj(self, submodules: HybridStackSubmodules) -> HybridStackSubmodules: - # Avoid modifying the original object so users don't get surprised about their `submodules` - # being modified underneath them. - submodules = copy.deepcopy(submodules) - mla_spec = submodules.mla_layer - # We always fuse the input layernorm because Hybrid always uses TransformerEngine. - mla_spec.submodules.input_layernorm = IdentityOp - mla_spec.submodules.self_attention.module = FusedMLASelfAttention - mla_spec.submodules.self_attention.submodules.linear_qkv_down_proj = ( - TELayerNormColumnParallelLinear - ) - mla_spec.submodules.self_attention.submodules.linear_q_down_proj = None - mla_spec.submodules.self_attention.submodules.linear_kv_down_proj = None - mla_spec.submodules.sharded_state_dict_keys_map = { - "self_attention.linear_q_down_proj.layer_norm_": "input_layernorm.", - "self_attention.linear_kv_down_proj.layer_norm_": "input_layernorm.", - "self_attention.linear_qkv_down_proj.layer_norm_": "input_layernorm.", - } - return submodules + @staticmethod + def _set_mtp_layer_number_for_moe_metrics( + layer: torch.nn.Module, mtp_layer_number: int + ) -> None: + """Tell nested MTP MoE routers which MTP depth they belong to for logging.""" + for module in layer.modules(): + router = getattr(module, "router", None) + if router is not None and getattr(router, "is_mtp_layer", False): + router.mtp_layer_number = mtp_layer_number def set_input_tensor(self, input_tensor: Tensor): """Set input tensor to be used instead of forward()'s input. @@ -316,62 +799,66 @@ def set_input_tensor(self, input_tensor: Tensor): def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]: """ - Returns the recurrent mixer's conv and SSM state shapes per input sequence - if this block contains Mamba or GDN layers (this may not be the case with PP > 1). + Returns the Mamba conv and ssm states shapes per input sequence + if this block contains Mamba layers (this may not be the case with PP > 1). """ for layer_type, layer in zip(self.layer_type_list, self.layers): if layer_type == LayerSymbols.MAMBA: return layer.mamba_state_shapes_per_request() - if layer_type == LayerSymbols.GDN: - if hasattr(layer, 'mamba_state_shapes_per_request'): - state_shapes = layer.mamba_state_shapes_per_request() - if state_shapes is not None: - return state_shapes - return layer.self_attention.mamba_state_shapes_per_request() return None def _compute_mhc_block_end_plan(self) -> List[bool]: - """Compute deterministic per-layer mHC recompute block boundaries.""" + """Compute per-layer block-end markers (deterministic from config).""" num_layers = len(self.layers) - block_ends: List[bool] = [False] * num_layers + is_recompute_block_end: List[bool] = [False] * num_layers if num_layers == 0: - return block_ends - - layers_per_block = self.config.mhc_recompute_layer_num - for layer_idx in range(num_layers): - is_last_in_stack = layer_idx == num_layers - 1 - block_ends[layer_idx] = is_last_in_stack or ( - layers_per_block is not None and (layer_idx + 1) % layers_per_block == 0 - ) - return block_ends + return is_recompute_block_end + mhc_recompute_layer_num = self.config.mhc_recompute_layer_num + for l_no in range(num_layers): + is_last_in_stack = l_no == num_layers - 1 + is_last_in_recompute_block = is_last_in_stack + if mhc_recompute_layer_num is not None: + is_last_in_recompute_block = is_last_in_stack or ( + (l_no + 1) % mhc_recompute_layer_num == 0 + ) + is_recompute_block_end[l_no] = is_last_in_recompute_block + return is_recompute_block_end def _build_mhc_recompute_layer_plan( self, use_mhc_recompute: bool - ) -> Tuple[List[Optional[CheckpointWithoutOutputManager]], List[bool]]: - """Build single-use recompute managers for this forward pass.""" + ) -> Tuple[List[Optional[MHCCheckpointManager]], List[bool]]: + """Pre-build per-layer MHC recompute managers and block-end markers. + + The block-end plan is deterministic from config and cached on the + instance; only the per-block ``MHCCheckpointManager`` instances are + allocated fresh per forward pass (managers are single-use). Mirrors + the caching scheme used by ``TransformerBlock``. + """ num_layers = len(self.layers) if not use_mhc_recompute or num_layers == 0: return [None] * num_layers, [False] * num_layers if self._mhc_block_end_plan is None: self._mhc_block_end_plan = self._compute_mhc_block_end_plan() - block_ends = self._mhc_block_end_plan + is_recompute_block_end = self._mhc_block_end_plan - layer_managers: List[Optional[CheckpointWithoutOutputManager]] = [None] * num_layers - manager = CheckpointWithoutOutputManager() - for layer_idx in range(num_layers): - layer_managers[layer_idx] = manager - if block_ends[layer_idx] and layer_idx != num_layers - 1: - manager = CheckpointWithoutOutputManager() - return layer_managers, block_ends + layer_managers: List[Optional[MHCCheckpointManager]] = [None] * num_layers + mhc_manager = MHCCheckpointManager() + for l_no in range(num_layers): + layer_managers[l_no] = mhc_manager + if is_recompute_block_end[l_no] and l_no != num_layers - 1: + mhc_manager = MHCCheckpointManager() + return layer_managers, is_recompute_block_end @staticmethod def _finalize_mhc_recompute_layer( - manager: Optional[CheckpointWithoutOutputManager], hidden_states: Tensor, is_block_end: bool + mhc_manager: Optional[MHCCheckpointManager], + hidden_states: Tensor, + is_last_in_recompute_block: bool, ) -> None: - """Finalize the current mHC recompute block when its last layer finishes.""" - if manager is not None and is_block_end: - manager.discard_all_outputs_and_register_unified_recompute(hidden_states) + """Finalize MHC recompute state for the current layer when a block ends.""" + if mhc_manager is not None and is_last_in_recompute_block: + mhc_manager.discard_all_outputs_and_register_unified_recompute(hidden_states) def forward( self, @@ -383,7 +870,8 @@ def forward( inference_params: Optional[BaseInferenceContext] = None, packed_seq_params: Optional[PackedSeqParams] = None, padding_mask=None, - ): + input_ids: Optional[Tensor] = None, + ) -> Union[Tensor, Tuple[Tensor, Tensor]]: """ Forward function of the HybridStack class. @@ -399,15 +887,15 @@ def forward( rotary_pos_emb (Tensor, optional): the rotary positional embeddings. Defaults to None. Returns: - Tensor: the output tensor. + Tensor in the common case. A 2-tuple ``(hidden_states, mhc_multistream)`` ONLY when + ``enable_hyper_connections and post_process and mtp_num_layers > 0 and not + is_mtp_layer`` — the extra element is the pre-contraction multi-stream tensor that + MTP's ``_concat_embeddings`` consumes. Callers (e.g. ``HybridModel.forward``) must + handle both; pipeline send/recv only ever transfers the contracted ``hidden_states``. """ inference_context = deprecate_inference_params(inference_context, inference_params) - cp_layout_state = None - if self._cp_layout_manager is not None: - cp_layout_state = self._cp_layout_manager.build_forward_state(packed_seq_params) - if not self.pre_process: # See set_input_tensor() hidden_states = self.input_tensor @@ -416,9 +904,13 @@ def forward( if isinstance(hidden_states, WrappedTensor): hidden_states = hidden_states.unwrap() - if self.config.enable_mhc_connections and self.pre_process and not self.is_mtp_layer: + # Skip input_expand inside MTP nested HybridStack: when mHC + MTP, the outer + # decoder hands in already-multi-stream hidden_states via mhc_multistream + # (see multi_token_prediction.py _concat_embeddings), so expanding again would + # produce [s, b, n*(n*h)] instead of [s, b, n*h] and break HC mapping_proj. + if self.config.enable_hyper_connections and self.pre_process and not self.is_mtp_layer: hidden_states = HyperConnectionModule.input_expand( - hidden_states, self.config.mhc_num_residual_streams + hidden_states, self.config.num_residual_streams ) if inference_context and inference_context.is_static_batching(): @@ -470,11 +962,13 @@ def get_inner_quant_context(config, layer_number): use_mhc_recompute = ( self.training - and self.config.enable_mhc_connections + and self.config.enable_hyper_connections and self.config.recompute_granularity == 'selective' and "mhc" in self.config.recompute_modules ) - mhc_layer_managers, mhc_block_ends = self._build_mhc_recompute_layer_plan(use_mhc_recompute) + mhc_layer_managers, mhc_is_last_in_recompute_block = self._build_mhc_recompute_layer_plan( + use_mhc_recompute + ) with outer_fp8_context: if self.config.recompute_granularity == 'full' and self.training: @@ -488,23 +982,21 @@ def get_inner_quant_context(config, layer_number): attention_bias=None, packed_seq_params=packed_seq_params, padding_mask=padding_mask, + input_ids=input_ids, use_inner_quantization_context=(use_inner_fp8_context or use_fp4_context), - cp_layout_state=cp_layout_state, ) else: - for layer_idx, layer in enumerate(self.layers): - layer_packed_seq_params = packed_seq_params - if cp_layout_state is not None: - hidden_states, layer_packed_seq_params = cp_layout_state.prepare_layer( - layer_idx, hidden_states - ) + for l_no, layer in enumerate(self.layers): # Layers have 1-indexed layer numbers attribute. inner_quant_context = get_inner_quant_context( self.config, layer.layer_number - 1 ) - mhc_manager = mhc_layer_managers[layer_idx] + + mhc_manager = mhc_layer_managers[l_no] if mhc_manager is not None: - mhc_manager.is_last_layer_in_recompute_block = mhc_block_ends[layer_idx] + mhc_manager.is_last_layer_in_recompute_block = ( + mhc_is_last_in_recompute_block[l_no] + ) with inner_quant_context: if isinstance(layer, (TransformerLayer, HyperConnectionHybridLayer)): @@ -514,9 +1006,11 @@ def get_inner_quant_context(config, layer_number): inference_context=inference_context, rotary_pos_emb=rotary_pos_emb, sequence_len_offset=sequence_len_offset, - packed_seq_params=layer_packed_seq_params, + packed_seq_params=packed_seq_params, padding_mask=padding_mask, ) + if input_ids is not None: + layer_kwargs["input_ids"] = input_ids if mhc_manager is not None and isinstance( layer, HyperConnectionHybridLayer ): @@ -527,7 +1021,7 @@ def get_inner_quant_context(config, layer_number): hidden_states=hidden_states, attention_mask=attention_mask, inference_context=inference_context, - packed_seq_params=layer_packed_seq_params, + packed_seq_params=packed_seq_params, ) # The attention layer (currently a simplified transformer layer) @@ -535,17 +1029,27 @@ def get_inner_quant_context(config, layer_number): # for cross-attention, and is not needed in our model. if isinstance(hidden_states, tuple): hidden_states = hidden_states[0] - if cp_layout_state is not None: - hidden_states = cp_layout_state.finalize_layer(layer_idx, hidden_states) self._finalize_mhc_recompute_layer( - manager=mhc_manager, + mhc_manager=mhc_manager, hidden_states=hidden_states, - is_block_end=mhc_block_ends[layer_idx], + is_last_in_recompute_block=mhc_is_last_in_recompute_block[l_no], ) + # When mHC + MTP, save the pre-contraction multi-stream tensor for MTP input. + # MTP's _concat_embeddings mHC branch expects [s, b, n*h] (multi-stream), while + # the contracted hidden_states is [s, b, h]. Mirrors transformer_block.py:948-988. + # Only the OUTER decoder stack does this; nested MTP stacks (is_mtp_layer=True) + # must keep returning a single Tensor so MTP's _postprocess receives the right + # type for learned_output_contract. + # On the final stage of a (non-MTP) stack with mHC active, capture the pre-contraction + # multi-stream tensor for MTP's `_concat_embeddings` (only meaningful when MTP layers + # exist, i.e. mtp_num_layers > 0), THEN contract the streams. Combining capture and + # contraction avoids repeating the condition. Nested MTP HybridStacks (is_mtp_layer=True) + # must NOT contract here — MTP's own `_postprocess` calls learned_output_contract + + # final_layernorm itself, so doing it here would double-collapse the multi-stream tensor. mhc_multistream = None - if self.config.enable_mhc_connections and self.post_process and not self.is_mtp_layer: + if self.config.enable_hyper_connections and self.post_process and not self.is_mtp_layer: if (self.config.mtp_num_layers or 0) > 0: mhc_multistream = hidden_states hidden_states = learned_output_contract( @@ -553,7 +1057,7 @@ def get_inner_quant_context(config, layer_number): self.hc_head_fn, self.hc_head_base, self.hc_head_scale, - self.config.mhc_num_residual_streams, + self.config.num_residual_streams, self.config.layernorm_epsilon, ) @@ -636,7 +1140,7 @@ def sharded_state_dict( make_sharded_tensors_for_checkpoint( local_state_dict, prefix, - sharded_offsets=sharded_offsets, + sharded_offsets=sharded_offsets or (), tp_group=self.tp_group, dp_cp_group=metadata['dp_cp_group'], ) diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index f0750064936..65ff8cdf8c5 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -19,6 +19,7 @@ class Symbols: 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) @@ -26,9 +27,9 @@ class Symbols: MOE = 'E' PIPE = '|' MTP_SEPARATOR = "/" - VALID_LAYERS = {MAMBA, GDN, KDA, ATTENTION, DS_ATTENTION, CSA, HCA, WINDOW, MLP, MOE} + 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 = {DS_ATTENTION, CSA, HCA, WINDOW} + MLA_ATTENTION = {MLA, DS_ATTENTION, CSA, HCA, WINDOW} @classmethod def name_sorted_valid_layer_symbols(cls) -> list[str]: @@ -298,7 +299,7 @@ def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False) f"Valid symbols are: {valid_chars}" ) - # Disallow standard Attention ('*') mixed with any MLA-based attention (D/C/H/W). + # 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( @@ -329,7 +330,7 @@ def validate_segment_layers(segment: str) -> List[str]: f"one of {Symbols.VALID_LAYERS}" ) - # Disallow standard Attention ('*') mixed with any MLA-based attention (D/C/H/W). + # 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" diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index 1d422d32dc6..2aa40c1cbed 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -195,6 +195,28 @@ self_attn_bda=get_bias_dropout_add, ), ), + mla_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, + self_attention=ModuleSpec( + module=MLASelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=MLASelfAttentionSubmodules( + linear_q_proj=TEColumnParallelLinear, + linear_q_down_proj=TELinear, + linear_q_up_proj=TEColumnParallelLinear, + linear_kv_down_proj=TELinear, + linear_kv_up_proj=TEColumnParallelLinear, + core_attention=TEDotProductAttention, + linear_proj=TERowParallelLinear, + q_layernorm=IdentityOp, + kv_layernorm=IdentityOp, + ), + ), + self_attn_bda=get_bias_dropout_add, + ), + ), # Started with spec from gpt_layer_specs.py # Using the TE spec because we had problems getting the non-TE spec # working diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 095dce88cb6..c6374a6c178 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -5082,7 +5082,7 @@ def _add_experimental_args(parser): type=str, default=None, help='Specify a hybrid layer pattern using M (mamba), G (gdn), K (kda), ' - '* (attention), D (dsa), - (mlp), E (moe). Use | to define pipeline ' + '* (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". ' diff --git a/tests/unit_tests/models/test_ling_v3_tiny_model_spec.py b/tests/unit_tests/models/test_ling_v3_tiny_model_spec.py new file mode 100644 index 00000000000..c55511054a0 --- /dev/null +++ b/tests/unit_tests/models/test_ling_v3_tiny_model_spec.py @@ -0,0 +1,205 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import pytest +import torch +import torch.nn.functional as F + +from examples.ling_v3_tiny.model_spec import LING_V3_TINY_HYBRID_LAYER_PATTERN +from examples.ling_v3_tiny.model_spec import hybrid_stack_spec as ling_v3_tiny_hybrid_stack_spec +from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TELinear +from megatron.core.models.hybrid.hybrid_block import HybridStack +from megatron.core.models.hybrid.hybrid_layer_allocation import ( + Symbols, + get_hybrid_layer_counts, + get_hybrid_total_layer_count, + parse_hybrid_pattern, +) +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.ssm.gated_delta_net import HAVE_FLA_KDA, KimiDeltaAttention +from megatron.core.ssm.mlp_layer import MLPLayer +from megatron.core.tensor_parallel import ColumnParallelLinear +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.multi_latent_attention import MLASelfAttention +from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper +from megatron.core.transformer.transformer_config import MLATransformerConfig +from megatron.core.transformer.transformer_layer import MoETransformerLayer +from tests.unit_tests.test_utilities import Utils + + +def _kda_submodules(spec): + return spec.submodules.kda_layer.submodules.self_attention.submodules + + +def _mla_submodules(spec): + return spec.submodules.mla_layer.submodules.self_attention.submodules + + +def test_ling_v3_tiny_hybrid_pattern_matches_model_topology(): + parsed = parse_hybrid_pattern(LING_V3_TINY_HYBRID_LAYER_PATTERN) + counts = get_hybrid_layer_counts(LING_V3_TINY_HYBRID_LAYER_PATTERN) + + assert get_hybrid_total_layer_count(LING_V3_TINY_HYBRID_LAYER_PATTERN) == 48 + assert parsed.mtp_pattern == Symbols.MLA + Symbols.MOE + assert parsed.mtp_num_depths == 1 + + main_counts = {symbol: parsed.main_pattern.count(symbol) for symbol in Symbols.VALID_LAYERS} + expected_nonzero_counts = {Symbols.KDA: 18, Symbols.MLA: 6, Symbols.MLP: 1, Symbols.MOE: 23} + assert { + symbol: count for symbol, count in main_counts.items() if count + } == expected_nonzero_counts + assert counts[Symbols.MLA] == 7 + assert counts[Symbols.MOE] == 24 + + +def test_ling_v3_tiny_spec_selects_required_projection_backends(): + default_kda = _kda_submodules(hybrid_stack_spec) + default_mla = _mla_submodules(hybrid_stack_spec) + tiny_kda = _kda_submodules(ling_v3_tiny_hybrid_stack_spec) + tiny_mla = _mla_submodules(ling_v3_tiny_hybrid_stack_spec) + + assert tiny_kda.beta_proj is ColumnParallelLinear + assert tiny_mla.linear_q_down_proj is TEColumnParallelLinear + assert tiny_mla.linear_kv_down_proj is TEColumnParallelLinear + assert tiny_mla.linear_gate is ColumnParallelLinear + + assert default_kda.beta_proj is TEColumnParallelLinear + assert default_mla.linear_q_down_proj is TELinear + assert default_mla.linear_kv_down_proj is TELinear + assert default_mla.linear_gate is None + + +@pytest.mark.internal +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available.") +@pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_ling_v3_tiny_reduced_hybrid_model_forward_backward(tp_size): + """Exercise every Tiny layer type and the nested MTP recipe.""" + + Utils.initialize_distributed() + if torch.distributed.get_world_size() % tp_size != 0: + Utils.destroy_model_parallel() + pytest.skip("world size must be divisible by tensor parallel size") + + Utils.initialize_model_parallel(tp_size, 1) + try: + model_parallel_cuda_manual_seed(123) + config = MLATransformerConfig( + num_layers=6, + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=4, + num_query_groups=4, + kv_channels=32, + q_lora_rank=32, + kv_lora_rank=32, + qk_head_dim=32, + qk_pos_emb_head_dim=16, + v_head_dim=32, + rope_type="rope", + rotary_base=10000, + rotary_percent=0.5, + qk_layernorm=True, + attention_output_gate=True, + 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, + num_moe_experts=2, + moe_ffn_hidden_size=64, + moe_router_topk=2, + moe_router_load_balancing_type="none", + moe_router_score_function="sigmoid", + moe_router_dtype="fp32", + moe_router_topk_scaling_factor=1.0, + moe_shared_expert_intermediate_size=64, + moe_grouped_gemm=True, + moe_token_dispatcher_type="alltoall", + mtp_num_layers=1, + mtp_loss_scaling_factor=0.1, + normalization="RMSNorm", + layernorm_epsilon=1.0e-6, + activation_func=F.silu, + gated_linear_unit=True, + add_bias_linear=False, + attention_dropout=0.0, + hidden_dropout=0.0, + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + use_cpu_initialization=True, + is_hybrid_model=True, + calculate_per_token_loss=True, + tensor_model_parallel_size=tp_size, + sequence_parallel=tp_size > 1, + ) + model = HybridModel( + config=config, + hybrid_stack_spec=ling_v3_tiny_hybrid_stack_spec, + vocab_size=128, + max_sequence_length=16, + hybrid_layer_pattern="K-KE+E/+E", + position_embedding_type="rope", + rotary_percent=0.5, + rotary_base=10000, + ).cuda() + + main_layers = model.decoder.layers + assert isinstance(main_layers[0].self_attention, KimiDeltaAttention) + assert isinstance(main_layers[1], MLPLayer) + assert isinstance(main_layers[2].self_attention, KimiDeltaAttention) + assert isinstance(main_layers[3], MoETransformerLayer) + assert isinstance(main_layers[4].self_attention, MLASelfAttention) + assert isinstance(main_layers[5], MoETransformerLayer) + + assert isinstance(model.mtp.layers[0].mtp_model_layer, HybridStack) + mtp_layers = model.mtp.layers[0].mtp_model_layer.layers + assert isinstance(mtp_layers[0].self_attention, MLASelfAttention) + assert isinstance(mtp_layers[1], MoETransformerLayer) + + sharded_state_dict = model.sharded_state_dict() + required_sharded_keys = [ + "decoder.layers.0.self_attention.beta_proj.weight", + "decoder.layers.4.self_attention.linear_q_down_proj.weight", + "decoder.layers.4.self_attention.linear_kv_down_proj.weight", + "decoder.layers.4.self_attention.linear_gate.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_gate.weight", + ] + for key in required_sharded_keys: + assert key in sharded_state_dict, key + + input_ids = torch.arange(16, device="cuda", dtype=torch.long).unsqueeze(0) + position_ids = torch.arange(16, device="cuda", dtype=torch.long).unsqueeze(0) + labels = (input_ids + 1) % 128 + loss_mask = torch.ones_like(input_ids, dtype=torch.float32) + loss = model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + labels=labels, + loss_mask=loss_mask, + ) + + assert loss.shape == input_ids.shape + assert torch.isfinite(loss).all() + assert "loss_sums" in MTPLossLoggingHelper.tracker + assert torch.isfinite(MTPLossLoggingHelper.tracker["loss_sums"]).all() + assert "num_tokens" in MTPLossLoggingHelper.tracker + assert (MTPLossLoggingHelper.tracker["num_tokens"] > 0).all() + + loss.float().mean().backward() + required_gradients = [ + "decoder.layers.0.self_attention.beta_proj.weight", + "decoder.layers.4.self_attention.linear_gate.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_gate.weight", + ] + parameters = dict(model.named_parameters()) + for name in required_gradients: + assert parameters[name].grad is not None, name + assert torch.isfinite(parameters[name].grad).all(), name + finally: + MTPLossLoggingHelper.tracker = {} + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index 42ec59b6cfa..7db5eb2170e 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -399,7 +399,7 @@ def test_hyper_connection_pipeline_boundary_shapes(self): assert output.shape == (sequence_length, micro_batch_size, transformer_config.hidden_size) def test_invalid_layer_types_cause_failure(self): - invalid_symbol = '+' + invalid_symbol = 'X' assert invalid_symbol not in Symbols.VALID_LAYERS # sanity check. layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP + invalid_symbol # validate_segment_layers() in hybrid_layer_allocation.py throws a ValueError. diff --git a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py index e33fe34082e..e8383752e12 100644 --- a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py +++ b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py @@ -78,8 +78,8 @@ 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', '+']), + ("K+K", ['K', '+', 'K']), + ("M+M+", ['M', '+', 'M', '+']), ] for pattern, expected in test_cases: result = validate_segment_layers(pattern) @@ -329,6 +329,7 @@ def test_simple_pattern(self): 'D': 0, 'G': 0, 'K': 0, + '+': 0, 'M': 2, '-': 0, 'E': 0, @@ -344,6 +345,7 @@ def test_all_layer_types(self): 'D': 0, 'G': 1, 'K': 0, + '+': 0, 'M': 1, '-': 1, 'E': 1, @@ -372,6 +374,7 @@ def test_with_pipes(self): 'D': 0, 'G': 0, 'K': 0, + '+': 0, 'M': 2, '-': 0, 'E': 0, @@ -384,6 +387,7 @@ def test_with_pipes(self): 'D': 0, 'G': 0, 'K': 0, + '+': 0, 'M': 4, '-': 4, 'E': 0, @@ -399,6 +403,7 @@ def test_with_mtp(self): 'D': 0, 'G': 0, 'K': 0, + '+': 0, 'M': 6, '-': 0, 'E': 0, @@ -415,6 +420,7 @@ def test_with_pipes_and_mtp(self): 'D': 0, 'G': 0, 'K': 0, + '+': 0, 'M': 8, '-': 4, 'E': 0, @@ -429,6 +435,7 @@ def test_moe_pattern(self): 'D': 0, 'G': 0, 'K': 0, + '+': 0, 'M': 2, '-': 0, 'E': 2, @@ -444,6 +451,7 @@ def test_mtp_with_attention(self): 'D': 0, 'G': 0, 'K': 0, + '+': 0, 'M': 7, '-': 0, 'E': 0, @@ -458,6 +466,7 @@ def test_gdn_pattern(self): 'D': 0, 'G': 2, 'K': 0, + '+': 0, 'M': 2, '-': 0, 'E': 0, @@ -473,6 +482,7 @@ def test_gdn_hybrid_pattern(self): 'D': 0, 'G': 2, 'K': 0, + '+': 0, 'M': 1, '-': 0, 'E': 0, @@ -488,10 +498,10 @@ def test_dsa_pattern(self): 'G': 0, 'K': 0, 'M': 2, - '+': 0, - '-': 0, - 'E': 0, - } + '+': 0, + '-': 0, + 'E': 0, + } def test_mla_pattern(self): assert get_hybrid_layer_counts("+M+M") == { @@ -510,6 +520,9 @@ def test_mla_pattern(self): def test_kda_pattern(self): assert get_hybrid_layer_counts("K+K/K") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 0, 'D': 0, 'G': 0, @@ -529,6 +542,7 @@ def test_empty_pattern(self): 'D': 0, 'G': 0, 'K': 0, + '+': 0, 'M': 0, '-': 0, 'E': 0, From f4cafaf5e661f1613d16aee78e3cdece20b416da Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 16 Aug 2026 01:31:15 +0800 Subject: [PATCH 07/34] Add a Ling-V3 Tiny pretraining example Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 5fe61a8935009caed61739a6407da7ec883dc9d3) --- examples/ling_v3_tiny/README.md | 14 +++ .../ling_v3_tiny/train_ling_v3_tiny_8gpu.sh | 103 ++++++++++++++++++ 2 files changed, 117 insertions(+) create mode 100644 examples/ling_v3_tiny/README.md create mode 100755 examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh diff --git a/examples/ling_v3_tiny/README.md b/examples/ling_v3_tiny/README.md new file mode 100644 index 00000000000..fdcd150bc3d --- /dev/null +++ b/examples/ling_v3_tiny/README.md @@ -0,0 +1,14 @@ +# Ling-V3 Tiny pretraining example + +This example runs the full Ling-V3 Tiny hybrid pattern from fresh initialization with +`MockGPTDataset`. + +Run it on one 8-GPU node: + +```bash +bash examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh +``` + +The default is a two-step BF16 training smoke run without validation or test iterations. +`TRAIN_ITERS`, `SEQ_LENGTH`, and `MASTER_PORT` may be overridden. No checkpoint is written unless +`SAVE_DIR` is explicitly set. diff --git a/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh b/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh new file mode 100755 index 00000000000..12fda0c7296 --- /dev/null +++ b/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh @@ -0,0 +1,103 @@ +#!/usr/bin/env bash + +# Full Ling-V3 Tiny fresh-initialization smoke run. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)" + +PYTHON="${PYTHON:-python}" +MASTER_PORT="${MASTER_PORT:-29501}" +SEQ_LENGTH="${SEQ_LENGTH:-128}" +TRAIN_ITERS="${TRAIN_ITERS:-2}" +SAVE_ARGS=() + +if [[ -n "${SAVE_DIR:-}" ]]; then + mkdir -p "${SAVE_DIR}" + SAVE_ARGS+=(--save "${SAVE_DIR}" --save-interval "${SAVE_INTERVAL:-1000000000}") +fi + +cd "${REPO_ROOT}" +export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" + +"${PYTHON}" -m torch.distributed.run \ + --nproc_per_node=8 \ + --master_port="${MASTER_PORT}" \ + pretrain_hybrid.py \ + --spec examples.ling_v3_tiny.model_spec hybrid_stack_spec \ + --hybrid-layer-pattern "K-KEKE+EKEKEKE+EKEKEKE+EKEKEKE+EKEKEKE+EKEKEKE+E/+E" \ + --mock-data \ + --tokenizer-type NullTokenizer \ + --vocab-size 157184 \ + --seq-length "${SEQ_LENGTH}" \ + --max-position-embeddings "${SEQ_LENGTH}" \ + --micro-batch-size 1 \ + --global-batch-size 8 \ + --train-iters "${TRAIN_ITERS}" \ + --eval-interval 1000000000 \ + --eval-iters 0 \ + --log-interval 1 \ + --seed 1234 \ + --bf16 \ + --transformer-impl transformer_engine \ + --normalization RMSNorm \ + --norm-epsilon 1e-6 \ + --swiglu \ + --disable-bias-linear \ + --hidden-size 1536 \ + --num-attention-heads 16 \ + --kv-channels 128 \ + --ffn-hidden-size 4608 \ + --untie-embeddings-and-output-weights \ + --position-embedding-type rope \ + --rotary-percent 0.5 \ + --rotary-base 6000000 \ + --tensor-model-parallel-size 1 \ + --pipeline-model-parallel-size 1 \ + --expert-model-parallel-size 8 \ + --context-parallel-size 2 \ + --cp-comm-type p2p \ + --num-experts 128 \ + --moe-ffn-hidden-size 512 \ + --moe-shared-expert-intermediate-size 512 \ + --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 \ + --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 \ + --linear-conv-kernel-dim 4 \ + --linear-key-head-dim 128 \ + --linear-value-head-dim 128 \ + --linear-num-key-heads 16 \ + --linear-num-value-heads 16 \ + --kda-safe-gate \ + --kda-lower-bound -5.0 \ + --multi-latent-attention \ + --q-lora-rank 256 \ + --kv-lora-rank 512 \ + --qk-head-dim 128 \ + --qk-pos-emb-head-dim 64 \ + --v-head-dim 128 \ + --qk-layernorm \ + --attention-output-gate \ + --mtp-num-layers 1 \ + --mtp-loss-scaling-factor 0.1 \ + --recompute-granularity full \ + --recompute-method uniform \ + --recompute-num-layers 1 \ + --lr 1e-4 \ + --min-lr 1e-4 \ + --lr-decay-style constant \ + --weight-decay 0.1 \ + --clip-grad 1.0 \ + --distributed-backend nccl \ + "${SAVE_ARGS[@]}" From 94b5ca0cadf13a1f180a1f861a0080da59d81c16 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 16 Aug 2026 01:31:15 +0800 Subject: [PATCH 08/34] Account for KDA and MLA in training metrics Count KDA and MLA operations in the hybrid training FLOPs estimate so that Ling-V3 Tiny TFLOPs and MFU include every attention layer used by the model. Keep the architecture-specific formulas separate and cover direct KDA projections and state updates, MLA latent projections, output gating, and ragged dense attention. Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 8296e8b79b774720212ba88016b331ba9fa1e960) --- megatron/training/training.py | 3318 ++++++++--------- .../test_num_floating_point_operations.py | 67 + 2 files changed, 1598 insertions(+), 1787 deletions(-) diff --git a/megatron/training/training.py b/megatron/training/training.py index f92cf6baa96..32e0587fd8d 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -41,9 +41,8 @@ _LEGACY_TRAIN_START_TIME = time.time() # NOTE(asolergi-nv): Legacy timestamp # First-party. -from megatron.core._rank_utils import safe_get_rank from megatron.core import mpu, nccl_allocator, tensor_parallel -from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper +from megatron.core.datasets.data_schedule import wrap_data_iterator from megatron.core.distributed import DistributedDataParallel as DDP from megatron.core.distributed import ( DistributedDataParallelConfig, @@ -51,20 +50,17 @@ finalize_model_grads, ) from megatron.core.distributed.fsdp.mcore_fsdp_adapter import ( - FullyShardedDataParallel, - FullyShardedDataParallelV1, - FullyShardedDataParallelV2, + FullyShardedDataParallel as megatron_FSDP, ) from megatron.core.enums import ModelType from megatron.core.fp8_utils import correct_amax_history_if_needed -from megatron.core.full_cuda_graph import FullCudaGraphWrapper, get_shared_capture_stream +from megatron.core.full_cuda_graph import FullCudaGraphWrapper from megatron.core.inference.symmetric_memory import SymmetricMemoryManager from megatron.core.inference.unified_memory import create_unified_mempool from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( - is_gated_delta_net_variant, is_linear_attention_variant, ) -from megatron.core.msc_utils import maybe_msc +from megatron.core.msc_utils import MultiStorageClientFeature, open_file from megatron.core.num_microbatches_calculator import ( destroy_num_microbatches_calculator, get_current_global_batch_size, @@ -96,7 +92,7 @@ destroy_global_memory_buffer, destroy_model_parallel, get_context_parallel_group, - get_hybrid_data_context_parallel_groups, + get_dynamic_data_context_parallel_groups, update_pg_timeout, ) from megatron.core.pipeline_parallel import get_forward_backward_func @@ -153,7 +149,7 @@ set_jit_fusion_options, write_args_to_tensorboard, ) -from megatron.training.utils import is_gtp_remat_active, is_hybrid_model +from megatron.training.utils import is_hybrid_model # Local. from . import ft_integration, one_logger_utils @@ -173,9 +169,9 @@ get_energy_monitor, get_one_logger, get_signal_handler, - get_telemetry, get_tensorboard_writer, get_timers, + get_tokenizer, get_wandb_writer, ) from .theoretical_memory_usage import report_theoretical_memory @@ -211,8 +207,6 @@ try: from modelopt.torch.distill.plugins.megatron import get_tensor_shapes_adjust_fn_for_distillation - from megatron.post_training.utils import maybe_enable_modelopt - has_nvidia_modelopt = True except ImportError: has_nvidia_modelopt = False @@ -235,9 +229,6 @@ # Module-level globals. # Startup timestamps for tracking program initialization phases. _STARTUP_TIMESTAMPS = { - 'slurm_job_start_time': None, # Set by Slurm itself (SLURM_JOB_START_TIME env var) - 'launch_script_start': None, # Set by the sbatch script, its own first line - 'launch_script_presrun': None, # Set by the sbatch script, right before srun 'program_start': None, # Set by entry script before imports 'main_entry': None, # Set by entry script at start of __main__ 'pretrain_entry': None, # Set at top of pretrain() @@ -267,13 +258,7 @@ MAX_NUM_CHECKPOINTS_MEMORY_REPORTED = 3 -def set_startup_timestamps( - program_start=None, - main_entry=None, - launch_script_start=None, - launch_script_presrun=None, - slurm_job_start_time=None, -): +def set_startup_timestamps(program_start=None, main_entry=None): """Set startup timestamps from the entry script. Call this after imports but before calling pretrain() to register @@ -282,18 +267,6 @@ def set_startup_timestamps( Args: program_start: Timestamp captured at very start of program, before any imports. main_entry: Timestamp captured right after entering __main__ block. - launch_script_start: Timestamp captured at the very top of the sbatch launch - script, before any of its own setup work (mkdir, checkpoint symlinking, - etc). Optional -- only meaningful if the launch script actually passes - it through an env var; entry points that don't have no way to know this. - launch_script_presrun: Timestamp captured by the launch script right before - it invokes srun. The gap to program_start is container/scheduler - overhead the launch script itself can't observe from inside Python. - slurm_job_start_time: Slurm's own record of when the job was granted its - allocation and started (SLURM_JOB_START_TIME env var, set by Slurm for - every process in the job -- no cooperation from the launch script - needed). Can be earlier than launch_script_start if there's - prolog/scheduling overhead before the launch script's first line runs. """ global _TRAIN_START_TIME, _STARTUP_TIMESTAMPS if program_start is not None: @@ -301,394 +274,6 @@ def set_startup_timestamps( _STARTUP_TIMESTAMPS['program_start'] = program_start if main_entry is not None: _STARTUP_TIMESTAMPS['main_entry'] = main_entry - if launch_script_start is not None: - _STARTUP_TIMESTAMPS['launch_script_start'] = launch_script_start - if launch_script_presrun is not None: - _STARTUP_TIMESTAMPS['launch_script_presrun'] = launch_script_presrun - if slurm_job_start_time is not None: - _STARTUP_TIMESTAMPS['slurm_job_start_time'] = slurm_job_start_time - - -# OTel: module-level helpers imported once at startup. -try: - from nemo.lens.state import is_span_group_enabled as _otel_sg_enabled - from nemo.lens.helpers import managed_span as _otel_managed_span - from nemo.lens.helpers import safe_set_span_attributes as _otel_safe_set_attrs - from nemo.lens.helpers import trace_fn as _otel_trace_fn -except ImportError: - from megatron.core.telemetry.fallbacks import is_span_group_enabled as _otel_sg_enabled - from megatron.core.telemetry.fallbacks import managed_span as _otel_managed_span - from megatron.core.telemetry.fallbacks import safe_set_span_attributes as _otel_safe_set_attrs - from megatron.core.telemetry.fallbacks import trace_fn as _otel_trace_fn - - -# OTel: module-level state for the job-level spans -- module level rather than -# closures local to pretrain(), because train()'s own early-exit path (a few -# thousand lines away, a different function) needs to reach the same -# end-spans-then-shutdown sequence. Closures defined inside pretrain() aren't -# callable from there; these are. -# -# Structure (audit section L: NO run-long umbrellas). An incarnation emits a set of -# SEGMENTED traces rather than one tree, because a run-long root span is dead weight -# (its children live in their own traces anyway) and is LOST outright when a fault kill -# arrives before it can be ended and exported. There is deliberately no in-process -# 'workload' / 'megatron.pretrain' span: the authoritative run envelope is the -# reckoner's out-of-band slurm.job (from sacct), and every span here carries run -# identity (run_uuid / slurm.sluid / nvrx.cycle) on its RESOURCE attributes, which is -# what actually stitches the traces back together. -# -# trace A: pre_startup (top-level, fully backdated, CYCLE 0 / non-NVRx only: -# SLURM_JOB_START_TIME -> launch_script_start, i.e. the -# queue tail / prolog before our first line. Ends exactly -# where the agent's nvrx.cold_start begins, so the two -# tile instead of overlapping -- audit section K.) -# trace B: megatron.startup (own trace, fresh trace_id via an empty Context()) -# |- megatron.startup.launch_script (backdated; plain-srun path only) -# |- megatron.startup.container_load (backdated; plain-srun path only) -# |- megatron.startup.imports / .arg_parse / .inprocess_setup / -# | .in_job_setup / .initialize_megatron / .jit_fusion_options (backdated) -# `- whatever runs live before train()'s loop (weight-hash check, sniff test, -# cuda-graph setup) -- they attach to the current context -# trace C_n: megatron.train (one per checkpoint interval, each its own trace -- -# see _reroot_otel_interval; C_0 Links back to trace B, C_n Links to C_n-1) -# -# The launch_script/container_load children are only emitted when this process actually -# owns that window (the launch script passed presrun through AND we're not under NVRx, -# where nvrx.cold_start covers it -- see pretrain_hybrid.py). When they are emitted the -# startup trace is anchored at launch_script_start so they tile inside it; otherwise it -# is anchored at program_start. -_otel_startup_span = None -_otel_startup_span_ctx = None # kept after the span ends, for the first interval's Link -_otel_ctx_module = None -_otel_startup_ctx_token = None -_otel_shutdown_done = False -# Trace segmentation: the steady-state loop is emitted as a series of megatron.train -# spans -- one per checkpoint interval -- instead of one run-long umbrella. Each is a -# compact [N iterations + checkpoint (+eval/sniff)] unit in its own trace, so it stays -# short and survives a fault-kill (it .end()s at the next boundary; a single umbrella -# only ends at loop exit, which a SIGKILL'd cycle never reaches). _otel_interval_span -# is the current block's root (a fresh trace_id, Linked back to the prior block -- and, -# for the first block, to this incarnation's megatron.startup trace); -# _otel_trace_interval_step is a dedicated per-loop counter (not -# `iteration`, which carries the resume offset) driving the save_interval-boundary -# re-root. Iterations/checkpoint/eval/sniff nest under it via _otel_managed_span. -_otel_interval_span = None -_otel_interval_ctx_token = None -_otel_trace_interval_step = 0 -# Set once the atexit hook + SIGTERM handler have been installed, so a second -# pretrain() call in the same process can't stack duplicate registrations. -_otel_exit_hooks_installed = False - - -def _otel_telemetry_active(): - """True only when telemetry is initialized AND actually exporting. - - Gates the process-global side effects (the atexit hook, the SIGTERM handler) that - exist purely to flush telemetry. With telemetry off there is nothing to flush, and - installing them anyway would change the process's signal semantics -- a Python-level - SIGTERM handler makes syscalls EINTR-restartable-through-Python and inserts us into - the torchelastic/ft_launcher handler chain -- for a run that asked for none of it. - - Returns False both when nemo-lens is absent (get_telemetry() is None) and when it is - present but disabled (the handle exists but is not exporting). - """ - handle = get_telemetry() - return handle is not None and bool(getattr(handle, 'is_exporting', False)) - - -def _start_otel_job_spans(model_type, program_start): - """Emit pre_startup and open this incarnation's megatron.startup trace. - - See the module comment above for the full picture. There are no run-long - umbrella spans any more (audit section L): 'workload' and 'megatron.pretrain' - used to wrap everything, but their children live in their own segmented traces, - the goodput denominator is the reckoner's out-of-band slurm.job, and a run-long - span is exactly the thing a fault SIGKILL destroys before it can be exported. - So this function now does just two things: - - 1. pre_startup -- fully backdated, top-level (emitted before any context is - attached, so it roots its own trace). Only when the SLURM->launch-script gap - is real; pretrain_hybrid.py suppresses it on NVRx restart cohorts. - 2. megatron.startup -- opened in an EMPTY Context() so it becomes the ROOT of a - fresh trace instead of inheriting whatever is ambient, the same technique - _reroot_otel_interval uses for the per-interval blocks. It uses the explicit - span API (start_span + context.attach) rather than managed_span because it - must stay open past the end of this function -- it closes in train(), just - before the loop (see _end_otel_startup_span). - - Its backdated children (launch_script/container_load/imports/arg_parse/...) are - already fully elapsed by the time we get here -- no tracer existed during imports - and arg parsing, telemetry isn't initialized until parse_and_validate_args() - returns -- so they are created-and-closed with explicit timestamps rather than - being live spans; see _backdated_otel_span. - """ - global _otel_startup_span, _otel_startup_span_ctx, _otel_ctx_module - global _otel_startup_ctx_token - - if not _otel_sg_enabled('job'): - return - - from opentelemetry import context as _otel_ctx, trace as _otel_trace - from opentelemetry.context import Context as _OtelContext - from nemo.lens.helpers import safe_set_span_attributes as _otel_set_attrs - - _otel_ctx_module = _otel_ctx - _otel_tracer = get_telemetry().tracer - - launch_script_start = _STARTUP_TIMESTAMPS.get('launch_script_start') - launch_script_presrun = _STARTUP_TIMESTAMPS.get('launch_script_presrun') - slurm_job_start_time = _STARTUP_TIMESTAMPS.get('slurm_job_start_time') - _program_start_ns = int(program_start * 1e9) if program_start is not None else None - - # pre_startup: the gap between Slurm marking the job started and our own launch - # script's first line actually running -- queue tail / prolog scripts / node setup. - # Kept as a COARSE fallback for when the out-of-band SLURM reckoner isn't available - # (that reconstructs this authoritatively from sacct). Emitted here, before anything - # is attached to the context, so it is a top-level span in its own trace. - # - # It ends at launch_script_start, which on cycle 0 is the sbatch script's first line - # -- the exact instant the NVRx agent starts nvrx.cold_start -- so pre_startup and - # cold_start TILE. (They used to overlap by the whole cold-start window because - # pretrain_hybrid.py overrode launch_script_start with the worker-spawn stamp on - # every cohort including cycle 0; that override is now restart-only. Audit sect. K.) - # On restart cohorts pretrain_hybrid.py drops slurm_job_start_time entirely, so the - # guard below skips pre_startup: the agent's restart tree owns that window instead. - if ( - launch_script_start is not None - and slurm_job_start_time is not None - and slurm_job_start_time < launch_script_start - ): - _backdated_otel_span('pre_startup', slurm_job_start_time, launch_script_start) - - # Do we own the pre-Python launch window (launch script start -> interpreter up)? - # Only on the plain-srun path: the launch script has to have passed its pre-srun - # stamp through, AND we must not be under NVRx (where the agent's nvrx.cold_start / - # restart-cycle spans already describe that window -- pretrain_hybrid.py signals - # this by dropping presrun on every NVRx cohort). When we do own it, fold it into - # the startup trace: anchor the trace at launch_script_start so launch_script and - # container_load tile inside it ahead of imports. Otherwise the startup trace starts - # at program_start, as before, and neither child is emitted. - _fold_launch_phases = ( - launch_script_start is not None - and launch_script_presrun is not None - and program_start is not None - and launch_script_start <= launch_script_presrun <= program_start - ) - _startup_start_ns = ( - int(launch_script_start * 1e9) if _fold_launch_phases else _program_start_ns - ) - - # Empty Context() -> no parent -> megatron.startup roots its OWN trace. - _otel_startup_span = _otel_tracer.start_span( - "megatron.startup", context=_OtelContext(), start_time=_startup_start_ns - ) - _otel_mark_goodput(_otel_startup_span) - # model_type used to ride on megatron.pretrain; keep it on the startup root now that - # the umbrella is gone, so it is still queryable per incarnation. - _otel_set_attrs(_otel_startup_span, {'megatron.model_type': str(model_type)}) - try: - _otel_startup_span_ctx = _otel_startup_span.get_span_context() - except Exception: # noqa: BLE001 -- telemetry must never break training - _otel_startup_span_ctx = None - _otel_startup_ctx_token = _otel_ctx.attach(_otel_trace.set_span_in_context(_otel_startup_span)) - - if _fold_launch_phases: - _backdated_otel_span( - 'megatron.startup.launch_script', launch_script_start, launch_script_presrun - ) - _backdated_otel_span( - 'megatron.startup.container_load', launch_script_presrun, program_start - ) - - -def _otel_mark_goodput(span): - """Mark a manually created span as a goodput boundary. - - The job/startup structural spans are created via the explicit start_span API - (they must outlive the function that opens them), so they cannot pass - attributes through a managed_span call and are stamped by hand instead. - Absent this attribute a span is profiling-only, which is the default for - everything except the resiliency-accounting boundaries. - """ - try: - span.set_attribute('is_goodput_span', True) - except Exception: # noqa: BLE001 -- telemetry must never break training - pass - - -def _backdated_otel_span(name, start, end): - """Record an already-elapsed phase as a closed span with explicit timestamps.""" - if not _otel_sg_enabled('job') or start is None or end is None: - return - _otel_tracer = get_telemetry().tracer - _s = _otel_tracer.start_span(name, start_time=int(start * 1e9)) - _otel_mark_goodput(_s) - _s.end(end_time=int(end * 1e9)) - - -def _end_otel_startup_span(): - """End the megatron.startup trace root and detach its token. - - This closes (and therefore exports) the whole startup trace. Safe to call more - than once -- the None check makes it idempotent, since it may run once explicitly - inside train() (after the preamble, before the loop) and again later via - _end_otel_job_spans() if that first call never happened (early crash). - - _otel_startup_span_ctx is deliberately NOT cleared: the first interval block Links - back to the startup trace, and by then the span itself is gone. - """ - global _otel_startup_span - if _otel_startup_span is not None: - try: - if _otel_ctx_module is not None and _otel_startup_ctx_token is not None: - _otel_ctx_module.detach(_otel_startup_ctx_token) - except Exception: # noqa: BLE001 -- telemetry must never break training - pass - try: - _otel_startup_span.end() - except Exception: # noqa: BLE001 - pass - _otel_startup_span = None - - -def _start_otel_train_span(): - """Prepare the steady-state training loop's tracing. - - The loop is emitted as one megatron.train span per checkpoint interval rather - than a single umbrella span; see the trace-segmentation note at the top of this - module for why. This just resets the per-interval counter; the first loop pass - opens the first megatron.train interval. - - Called after _end_otel_startup_span() (so the startup trace is closed and detached; - the interval blocks root their own traces regardless, and the first one Links back - to the startup trace via the span context kept in _otel_startup_span_ctx). - """ - global _otel_trace_interval_step - if not _otel_sg_enabled('job'): - return - _otel_trace_interval_step = 0 # restart the per-interval trace counter for this loop - - -def _reroot_otel_interval(): - """Close the current interval span and open the next as a new megatron.train - span in a new trace (empty context -> fresh trace_id). See the trace-segmentation - note at the top of this module for why the loop is segmented this way. - - Each interval is Linked to the prior interval (a chain), and the FIRST interval of - this incarnation is Linked to its megatron.startup trace instead -- so a backend can - walk startup -> block -> block -> ... even though they are all separate traces. There - is no run-long span to link to any more (audit section L); run identity - (run_uuid/sluid/nvrx.cycle) rides the resource attrs on every span, which is the - primary correlation key. Iterations/checkpoint/eval/sniff nest under the block - because they use _otel_managed_span, which attaches to the current context.""" - global _otel_interval_span, _otel_interval_ctx_token - if get_telemetry() is None or not _otel_sg_enabled('job'): - return - from opentelemetry import context as _octx, trace as _otr - from opentelemetry.context import Context - from opentelemetry.trace import Link - prev = _otel_interval_span - links = [] - try: - if prev is not None: - links.append(Link(prev.get_span_context())) - elif _otel_startup_span_ctx is not None: - # First block of this incarnation -> chain it to the startup trace. - links.append(Link(_otel_startup_span_ctx)) - except Exception: # noqa: BLE001 - links = [] - # Close the prior interval (detach its context, end the span) before opening the next. - if prev is not None: - try: - if _otel_interval_ctx_token is not None: - _octx.detach(_otel_interval_ctx_token) - except Exception: # noqa: BLE001 - pass - try: - prev.end() - except Exception: # noqa: BLE001 - pass - _otel_interval_span = None - _otel_interval_ctx_token = None - try: - sp = get_telemetry().tracer.start_span( - 'megatron.train', context=Context(), links=links - ) - _otel_mark_goodput(sp) - _otel_interval_span = sp - _otel_interval_ctx_token = _octx.attach(_otr.set_span_in_context(sp)) - except Exception: # noqa: BLE001 - _otel_interval_span = None - _otel_interval_ctx_token = None - - -def _maybe_reroot_otel_interval(): - """Called once at the top of each training-loop pass. Re-roots the trace at each - checkpoint-frequency boundary using a dedicated counter (so groups are clean - save_interval-sized chunks regardless of the resume offset in `iteration`). - save_interval<=0 (no checkpointing) -> never re-roots (single trace, as before).""" - global _otel_trace_interval_step - freq = getattr(get_args(), 'save_interval', None) or 0 - if freq > 0 and (_otel_trace_interval_step % freq == 0): - _reroot_otel_interval() - _otel_trace_interval_step += 1 - - -def _end_otel_interval_span(): - """Close the final interval root (loop exit / teardown). Idempotent.""" - global _otel_interval_span, _otel_interval_ctx_token - if _otel_interval_span is not None: - try: - from opentelemetry import context as _octx - if _otel_interval_ctx_token is not None: - _octx.detach(_otel_interval_ctx_token) - except Exception: # noqa: BLE001 - pass - try: - _otel_interval_span.end() - except Exception: # noqa: BLE001 - pass - _otel_interval_span = None - _otel_interval_ctx_token = None - - -def _end_otel_train_span(): - """Close the final megatron.train interval span. Idempotent -- called on - train()'s normal return and via _end_otel_job_spans() on the sys.exit() path - (before the telemetry shutdown, so the last block survives).""" - _end_otel_interval_span() - - -def _end_otel_job_spans(): - """Close whatever is still open (the current interval block, and startup if a - crash/early-exit skipped its normal end point), then flush+shutdown telemetry. - - Since the run-long umbrellas are gone (audit section L) this is a short list -- - which is the point: there is no longer a long-lived span that a fault SIGKILL can - destroy in the window between "training died" and "exporter flushed". - - This is the single place that does all of this, and every exit path -- - pretrain()'s own normal completion, train()'s early-exit - (--exit-interval/duration/signal, which calls sys.exit() directly from inside - train(), never returning control to the rest of pretrain()), the uncaught-exception - path, the hard-terminate SIGTERM handler, and the atexit fallback -- must call this - same function rather than reimplementing the sequence. Ending the spans before - shutdown() matters: shutdown() flushes and tears down the provider, so anything - ended afterward has nowhere left to export into. Safe to call more than once (e.g. - once from train()'s early exit, again via atexit) -- span-ending is idempotent via - the None checks in the two helpers, and the shutdown() call is separately guarded - so it only actually fires once. - """ - global _otel_shutdown_done - - _end_otel_train_span() # end innermost first (train() sys.exit() path) - _end_otel_startup_span() - - if not _otel_shutdown_done: - _otel_handle = get_telemetry() - if _otel_handle is not None: - _otel_handle.shutdown() - _otel_shutdown_done = True def destroy_global_state(): @@ -702,7 +287,7 @@ def destroy_global_state(): def print_datetime(string, override_timestamp=None): """Note that this call will sync across all ranks. Use override_timestamp if provided; - otherwise use current timestamp.""" + otherwise use current timestamp.""" torch.distributed.barrier() if override_timestamp is None: time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f') @@ -799,11 +384,118 @@ def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float return total_real_tokens / dedup, seqlen_squared_sum / dedup +def _dsv4_hybrid_self_attention_flops( + *, + hidden_size, + num_attention_heads, + v_head_dim, + q_lora_rank, + o_groups, + o_lora_rank, + csa_window_size, + seq_length, + n_layers_r0, + n_layers_r4, + n_layers_r128, + dsa_indexer_n_heads, + dsa_indexer_head_dim, + dsa_indexer_topk, +): + """Per-iteration DSv4-hybrid self-attention FLOPs coefficients. + + DSv4 attention layers are MLA-based but replace full ``O(L^2)`` core + attention with sparse attention (sliding window + compressed-KV), plus a + main compressor and, on ``ratio==4`` layers, a learned indexer (DSA). This + is the SINGLE SOURCE OF TRUTH shared by the standard-model path + (``transformer_flops``) and the hybrid-model path (``hybrid_flops``): the + two differ only in how they obtain the per-ratio attention-layer counts + (a ``csa_compress_ratios`` list vs. Window/CSA/HCA pattern symbols). + + Layer-type ratios: + * ``r==0`` (Window): window-only attention; no compressor / indexer. + * ``r==4`` (CSA): window + learned-topk over compressed KV + (compressor + indexer). + * ``r==128`` (HCA): window + all compressed KV (compressor only). + + Returns ``(token_linear, core)`` WITHOUT the fwd+bwd (x3) or FMA (x2) + expansion factors -- the caller applies those. Multiply ``token_linear`` by + the real (unpadded) token count and ``core`` by ``sum_i(L_i ** 2)``. + """ + n_attn_layers = n_layers_r0 + n_layers_r4 + n_layers_r128 + + # ---- MLA projections (token-linear, per attention layer) ---- + # DSv4 hybrid MLA: ``qk_head_dim + qk_pos_emb_head_dim == v_head_dim`` and the + # joint KV is a single ``hidden -> v_head_dim`` projection; the output uses a + # grouped low-rank ``(o_groups x o_lora_rank)`` projection (NOT the dense + # ``num_heads * v_head_dim -> hidden`` projection of plain MLA). + q_term = q_lora_rank * (hidden_size + num_attention_heads * v_head_dim + 1) + kv_term = hidden_size * v_head_dim + v_head_dim + o_term = num_attention_heads * v_head_dim * o_lora_rank + o_groups * o_lora_rank * hidden_size + mla_proj_term = (q_term + kv_term + o_term) * n_attn_layers + + # ---- Sparse attention (replaces full core attention) ---- + # Split into token-linear parts (window attention, constant per token) and + # L^2 parts (compressed-KV attention, scales with sequence length) so THD + # packed sequences get correct ``seqlen_squared_sum`` scaling. + # r=0: window-only, fixed per-token cost. + sparse_attn_r0 = n_layers_r0 * num_attention_heads * csa_window_size * v_head_dim * 2 + # r=128: window (token-linear) + all compressed KV (L^2). Compressed + # positions per token ~= L/(128*2) (causal /2), x2 for QK^T + softmax@V -> + # L^2 coefficient ``num_heads * v_head_dim / 128``. + sparse_attn_r128_window = n_layers_r128 * num_attention_heads * csa_window_size * v_head_dim * 2 + sparse_attn_r128_core = n_layers_r128 * num_attention_heads * v_head_dim / 128 + + # ---- Main compressor (ratio > 0 layers) ---- + # Two projections per layer (wkv + wgate): ``hidden -> coff * v_head_dim``. + # ratio == 4: coff = 2 (overlapping windows); ratio == 128: coff = 1. + main_compressor_term = ( + n_layers_r4 * hidden_size * (2 * v_head_dim) * 2 + + n_layers_r128 * hidden_size * (1 * v_head_dim) * 2 + ) + + # ---- r=4 layers: sparse attention + indexer ---- + if n_layers_r4 > 0: + assert ( + dsa_indexer_n_heads is not None + ), "dsa_indexer_n_heads must be set for dsv4_hybrid with ratio==4 layers." + assert ( + dsa_indexer_head_dim is not None + ), "dsa_indexer_head_dim must be set for dsv4_hybrid with ratio==4 layers." + assert ( + dsa_indexer_topk is not None + ), "dsa_indexer_topk must be set for dsv4_hybrid with ratio==4 layers." + effective_topk_4 = min(dsa_indexer_topk, seq_length // 4) + avg_comp_4 = effective_topk_4 * (1 - effective_topk_4 * 4 / (2 * seq_length)) + sparse_attn_r4 = ( + n_layers_r4 * num_attention_heads * (csa_window_size + avg_comp_4) * v_head_dim * 2 + ) + # Indexer token-linear: compressor (coff=2, wkv + wgate), Q proj, weights proj. + indexer_token_term = ( + n_layers_r4 * hidden_size * (2 * dsa_indexer_head_dim) * 2 + + n_layers_r4 * q_lora_rank * dsa_indexer_n_heads * dsa_indexer_head_dim + + n_layers_r4 * hidden_size * dsa_indexer_n_heads + ) + # Indexer L^2: scoring each query against ~L/4 compressed positions. + indexer_scoring_core = n_layers_r4 * dsa_indexer_n_heads * dsa_indexer_head_dim / 4 + else: + sparse_attn_r4 = 0 + indexer_token_term = 0 + indexer_scoring_core = 0 + + token_linear = ( + mla_proj_term + + sparse_attn_r0 + + sparse_attn_r4 + + sparse_attn_r128_window + + main_compressor_term + + indexer_token_term + ) + core = sparse_attn_r128_core + indexer_scoring_core + return token_linear, core + + def num_floating_point_operations( - args, - batch_size, - seqlen_squared_sum_in_batch=None, - total_real_tokens_in_batch=None, + args, batch_size, seqlen_squared_sum_in_batch=None, total_real_tokens_in_batch=None ): """Compute the number of floating-point operations for one global batch. @@ -842,26 +534,49 @@ def mlp_layer_flops(total_tokens, hidden_size, expansion=4.0, swiglu=False): scale_factor = 3.0 / 2.0 if swiglu else 1.0 return 4 * expansion * scale_factor * total_tokens * hidden_size**2 - def moe_layer_flops(total_tokens, hidden_size, moe_ffn_hidden_size, - shared_expert_ffn_hidden_size, num_experts_routed_to, - moe_latent_size=None, swiglu=False): + def moe_layer_flops( + total_tokens, + hidden_size, + moe_ffn_hidden_size, + shared_expert_ffn_hidden_size, + num_experts_routed_to, + moe_latent_size=None, + swiglu=False, + ): """Calculate FLOPs for an MoE layer.""" scale_factor = 3.0 / 2.0 if swiglu else 1.0 if moe_latent_size is None: - routed_flops = (4 * total_tokens * hidden_size * - moe_ffn_hidden_size * num_experts_routed_to * scale_factor) + routed_flops = ( + 4 + * total_tokens + * hidden_size + * moe_ffn_hidden_size + * num_experts_routed_to + * scale_factor + ) else: # Routed experts run on moe_latent_size. - routed_flops = (4 * total_tokens * moe_latent_size * - moe_ffn_hidden_size * num_experts_routed_to * scale_factor) + routed_flops = ( + 4 + * total_tokens + * moe_latent_size + * moe_ffn_hidden_size + * num_experts_routed_to + * scale_factor + ) # Up proj and down proj. - routed_flops += (4 * total_tokens * hidden_size * moe_latent_size) + routed_flops += 4 * total_tokens * hidden_size * moe_latent_size shared_flops = 4 * total_tokens * hidden_size * shared_expert_ffn_hidden_size * scale_factor 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, ): """Calculate FLOPs for an attention layer. @@ -877,13 +592,57 @@ def attn_layer_flops( # 4 * total_tokens * h * p * (h + h*(g/n)): QKV + output projections (fwd*3, with FMA*2). # 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 * 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, + ): + """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 + ) + # 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 * num_heads if attention_output_gate else 0) + ) + # 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 @@ -894,115 +653,238 @@ def mamba_layer_flops(total_tokens, hidden_size, state_dim=16, nheads = d_in // head_dim return ( ( - 2 - * total_tokens - * hidden_size - * (2 * d_in + 2 * num_groups * state_dim + nheads) + 2 * total_tokens * hidden_size * (2 * d_in + 2 * num_groups * state_dim + nheads) ) # in_proj + (7 * total_tokens * d_in * state_dim) # scan + (2 * total_tokens * d_in * hidden_size) # out_proj ) - def gated_delta_product_layer_flops( + def gdn_layer_flops( total_tokens, hidden_size, - num_householder, - state_dim=128, - head_dim=64, - num_groups=8, - num_heads=None, + qk_head_dim=128, + v_head_dim=128, + num_qk_heads=16, + num_v_heads=32, conv_kernel_dim=4, ): - """Calculate FLOPs for a Gated Delta Product (GDP) layer.""" - if num_heads is None: - d_inner = 2 * hidden_size - num_heads = d_inner // head_dim - else: - d_inner = num_heads * head_dim - in_proj_dim = ( - d_inner * (1 + num_householder) - + num_groups * state_dim * (1 + num_householder) - + num_heads * (1 + num_householder) - ) - conv_dim = d_inner * num_householder + num_groups * state_dim * (1 + num_householder) - non_core_flops = 2 * total_tokens * ( - hidden_size * in_proj_dim - + conv_kernel_dim * conv_dim - + d_inner * hidden_size - ) - # Best-case recurrent GDP core estimate. The FLA chunk kernel may do additional - # score/solve/WY work, but this keeps the implementation-agnostic lower bound explicit. - core_flops = (4 * num_householder + 3) * total_tokens * d_inner * state_dim - return non_core_flops + core_flops - - def gdn_layer_flops(total_tokens, hidden_size, - qk_head_dim=128, v_head_dim=128, - num_qk_heads=16, num_v_heads=32, - conv_kernel_dim=4, use_gdn2=False): """Calculate FLOPs for a Gated Delta Net (GDN) layer.""" qk_dim = qk_head_dim * num_qk_heads v_dim = v_head_dim * num_v_heads - if use_gdn2: - # GDN2 in_proj: hidden_size -> (4*qk_dim + 3*v_dim) for q, k, v, z, f, b, w - in_proj_dim = 4 * qk_dim + 3 * v_dim - else: - # GDN in_proj: hidden_size -> (2*qk_dim + 2*v_dim + 2*num_v_heads) - in_proj_dim = 2 * qk_dim + 2 * v_dim + 2 * num_v_heads return ( - 2 * total_tokens * ( - # in_proj - hidden_size * in_proj_dim + 2 + * total_tokens + * ( + # in_proj: hidden_size -> (2*qk_dim + 2*v_dim + 2*num_v_heads) + hidden_size * (2 * qk_dim + 2 * v_dim + 2 * num_v_heads) # conv1d + conv_kernel_dim * (2 * qk_dim + v_dim) # gated delta rule: KK^T, VK^T, S(a(I-bKK^T)), and SQ - + num_v_heads * (v_head_dim ** 2) * 4 + + num_v_heads * (v_head_dim**2) * 4 # out_proj: v_dim -> hidden_size + hidden_size * v_dim ) ) - 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): - """Calculate total FLOPs for the hybrid model.""" - 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) - 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) + 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, + 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, + multi_latent_attention=False, + 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, + 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: + if multi_latent_attention: + plain_attn_layer_flops = 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, + ) + else: + plain_attn_layer_flops = attn_layer_flops( + total_tokens, + seqlen_squared_sum, + hidden_size, + num_attn_heads, + gqa, + gqa_groups, + kv_channels, + ) + 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, + ) + + 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, + ) + 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_layer_flops( + total_tokens, + hidden_size, + mamba_state_dim, + mamba_head_dim, + mamba_num_groups, + mamba_num_heads, + ) + + 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 @@ -1091,50 +973,57 @@ def transformer_flops(): https://arxiv.org/abs/2305.10403 https://arxiv.org/abs/2205.05198 ''' - ## MLA - if args.q_lora_rank is None: - q_term = ( - args.hidden_size - * args.num_attention_heads - * (args.qk_head_dim + args.qk_pos_emb_head_dim) - ) + if args.experimental_attention_variant == "dsv4_hybrid": + ## DSv4 hybrid: the MLA projections AND the sparse attention that + ## replaces full core attention are both computed together by + ## ``_dsv4_hybrid_self_attention_flops`` in the dsv4_hybrid branch + ## below (the single source of truth shared with ``hybrid_flops``). + ## Zero the standard MLA terms here so they are not double-counted. + standard_self_attn_term = 0 + standard_self_attn_core_term = 0 else: - q_term = args.q_lora_rank * ( - args.hidden_size - + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim) - + 1 - ) - # Token-linear part of MLA self-attention (lora projs, kv proj, RoPE, output proj). - standard_self_attn_term = ( - forward_backward_expansion_factor - * fma_expansion_factor - * ( - ## q lora + rope + q norm - q_term - ## kv lora + rope + kv norm - + args.kv_lora_rank - * ( + ## MLA + if args.q_lora_rank is None: + q_term = ( args.hidden_size - + args.num_attention_heads * (args.qk_head_dim + args.v_head_dim) + * args.num_attention_heads + * (args.qk_head_dim + args.qk_pos_emb_head_dim) + ) + else: + q_term = args.q_lora_rank * ( + args.hidden_size + + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim) + 1 ) - + args.hidden_size * args.qk_pos_emb_head_dim - ## o proj - + (args.num_attention_heads * args.v_head_dim) * args.hidden_size + # Token-linear part of MLA self-attention (lora projs, kv proj, RoPE, output proj). + standard_self_attn_term = ( + forward_backward_expansion_factor + * fma_expansion_factor + * ( + ## q lora + rope + q norm + q_term + ## kv lora + rope + kv norm + + args.kv_lora_rank + * ( + args.hidden_size + + args.num_attention_heads * (args.qk_head_dim + args.v_head_dim) + + 1 + ) + + args.hidden_size * args.qk_pos_emb_head_dim + ## o proj + + (args.num_attention_heads * args.v_head_dim) * args.hidden_size + ) ) - ) - # Core-attention (L^2) part: ``QK^T`` and ``(softmax(QK^T)) V``. The - # ``/2`` accounts for the causal mask and the ``*2`` cancels it via FMA. - standard_self_attn_core_term = ( - forward_backward_expansion_factor - * fma_expansion_factor - * ( - args.num_attention_heads - * (args.qk_head_dim + args.qk_pos_emb_head_dim) - / 2 - + args.num_attention_heads * args.v_head_dim / 2 + # Core-attention (L^2) part: ``QK^T`` and ``(softmax(QK^T)) V``. The + # ``/2`` accounts for the causal mask and the ``*2`` cancels it via FMA. + standard_self_attn_core_term = ( + forward_backward_expansion_factor + * fma_expansion_factor + * ( + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim) / 2 + + args.num_attention_heads * args.v_head_dim / 2 + ) ) - ) else: ## MHA or GQA @@ -1156,8 +1045,7 @@ def transformer_flops(): + gate_projection_size ) ## out proj - + query_projection_size - * args.hidden_size + + query_projection_size * args.hidden_size ) ) # Core-attention (L^2) part: ``QK^T`` and ``(softmax(QK^T)) V``. @@ -1169,13 +1057,15 @@ def transformer_flops(): * 2 # QK^T and (QK^T)V ) + dsv4_hybrid_extra_term = 0 + dsv4_hybrid_extra_core_term = 0 if is_linear_attention_variant(args.experimental_attention_variant): # Calculate number of dense and MoE Transformer MLPs. if isinstance(args.linear_attention_freq, int): linear_attention_pattern = [ # [1,1,...,1,0,1,1,...,1,0,...] - 0 if ((i + 1) % args.linear_attention_freq == 0) - else 1 for i in range(num_layers) + 0 if ((i + 1) % args.linear_attention_freq == 0) else 1 + for i in range(num_layers) ] elif isinstance(args.linear_attention_freq, list): linear_attention_pattern = args.linear_attention_freq @@ -1199,7 +1089,7 @@ 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 == "gated_delta_net": # 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 @@ -1207,28 +1097,18 @@ def transformer_flops(): num_v_heads = args.linear_num_value_heads qk_dim = qk_head_dim * num_qk_heads v_dim = v_head_dim * num_v_heads - if args.experimental_attention_variant == "gdn2": - # GDN2 in_proj: q, k, v, z, f, b, w - in_proj_dim = 4 * qk_dim + 3 * v_dim - else: - in_proj_dim = 2 * qk_dim + 2 * v_dim + 2 * num_v_heads linear_self_attn_term = ( forward_backward_expansion_factor * fma_expansion_factor * ( ## in proj - args.hidden_size - * in_proj_dim + args.hidden_size * (2 * qk_dim + 2 * v_dim + 2 * num_v_heads) ## conv1d - + args.linear_conv_kernel_dim - * (2 * qk_dim + v_dim) + + args.linear_conv_kernel_dim * (2 * qk_dim + v_dim) ## gated delta rule - + num_v_heads - * (v_head_dim ** 2) - * 4 # KK^T, VK^T, S(a(I-bKK^T)), and SQ + + num_v_heads * (v_head_dim**2) * 4 # KK^T, VK^T, S(a(I-bKK^T)), and SQ ## out proj - + args.hidden_size - * v_dim + + args.hidden_size * v_dim ) ) else: @@ -1236,6 +1116,45 @@ def transformer_flops(): "Invalid experimental_attention_variant: " f"{args.experimental_attention_variant}" ) + elif args.experimental_attention_variant == "dsv4_hybrid": + # DSv4 hybrid: MLA projections + sparse attention (window / + # compressed-KV), main compressor, and learned indexer (DSA) are all + # computed by the shared ``_dsv4_hybrid_self_attention_flops`` helper. + # The standard MLA terms were zeroed above to avoid double-counting. + num_linear_attention_layers = 0 + linear_self_attn_term = 0 + num_standard_attention_layers = num_layers + + compress_ratios = args.csa_compress_ratios + assert compress_ratios is not None, "csa_compress_ratios must be set for dsv4_hybrid" + assert len(compress_ratios) == num_layers, ( + f"Invalid length of csa_compress_ratios: {len(compress_ratios)}, " + f"expected num_layers + mtp_num_layers ({num_layers})." + ) + # ratio == 0: window-only; ratio == 4: window + topk compressed KV + # (compressor + indexer); ratio == 128: window + all compressed KV. + dsv4_token_term, dsv4_core_term = _dsv4_hybrid_self_attention_flops( + hidden_size=args.hidden_size, + num_attention_heads=args.num_attention_heads, + v_head_dim=args.v_head_dim, + q_lora_rank=args.q_lora_rank, + o_groups=args.o_groups, + o_lora_rank=args.o_lora_rank, + csa_window_size=args.csa_window_size, + seq_length=args.seq_length, + n_layers_r0=sum(1 for r in compress_ratios if r == 0), + n_layers_r4=sum(1 for r in compress_ratios if r == 4), + n_layers_r128=sum(1 for r in compress_ratios if r == 128), + dsa_indexer_n_heads=args.dsa_indexer_n_heads, + dsa_indexer_head_dim=args.dsa_indexer_head_dim, + dsa_indexer_topk=args.dsa_indexer_topk, + ) + dsv4_hybrid_extra_term = ( + forward_backward_expansion_factor * fma_expansion_factor * dsv4_token_term + ) + dsv4_hybrid_extra_core_term = ( + forward_backward_expansion_factor * fma_expansion_factor * dsv4_core_term + ) else: num_linear_attention_layers = 0 linear_self_attn_term = 0 @@ -1246,9 +1165,14 @@ def transformer_flops(): self_attn_term = ( linear_self_attn_term * num_linear_attention_layers + standard_self_attn_term * num_standard_attention_layers + + dsv4_hybrid_extra_term + ) + # Core attention (L^2) FLOPs. Standard attention has a uniform per-layer + # coefficient; DSv4 sparse attention varies by layer type and is pre-summed. + self_attn_core_term = ( + standard_self_attn_core_term * num_standard_attention_layers + + dsv4_hybrid_extra_core_term ) - # Core attention (L^2) FLOPs per standard-attention layer. - self_attn_core_term = standard_self_attn_core_term * num_standard_attention_layers # Token-linear FLOPs scale with the real (unpadded) token count. # For BSHD this falls back to ``batch_size * seq_length`` (no padding). @@ -1261,8 +1185,7 @@ def transformer_flops(): * args.hidden_size * ( # dense layer (deepseek v2, v3 style) - (args.ffn_hidden_size * ffn_expansion_factor) - * num_dense_layers + (args.ffn_hidden_size * ffn_expansion_factor) * num_dense_layers # routed experts + ( (moe_ffn_hidden_size * num_experts_routed_to * ffn_expansion_factor) @@ -1280,8 +1203,7 @@ def transformer_flops(): ) * num_moe_layers # Shared Experts. - + (shared_expert_ffn_hidden_size * ffn_expansion_factor) - * num_moe_layers + + (shared_expert_ffn_hidden_size * ffn_expansion_factor) * num_moe_layers ) # Self Attention (token-linear part). + self_attn_term @@ -1310,48 +1232,53 @@ def transformer_flops(): ) return total_floating_point_operations - def _uses_gated_delta_product_spec(args): - """Return True when the selected hybrid stack spec swaps Mamba layers to GDP.""" - def _split_spec_part(part): - return str(part).replace('[', ' ').replace(']', ' ').replace(',', ' ').split() - - spec = getattr(args, 'spec', None) - if spec is None: - return False - if isinstance(spec, str): - spec_parts = _split_spec_part(spec) - else: - spec_parts = [] - for part in spec: - spec_parts.extend(_split_spec_part(part)) - if not spec_parts: - return False - return spec_parts[-1] in {'gdp_stack_spec', 'gated_delta_product_stack_spec'} - # 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] + + # Ling uses + for dense MLA. DSv4 D/C/H/W layers are MLA variants too, + # 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,28 +1287,54 @@ 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, num_attn_heads=args.num_attention_heads, gqa=args.group_query_attention, gqa_groups=args.num_query_groups, kv_channels=args.kv_channels, mlp_expansion=args.ffn_hidden_size / args.hidden_size, swiglu=args.swiglu, - use_gated_delta_product=_uses_gated_delta_product_spec(args), moe_latent_size=args.moe_latent_size, - moe_ffn_hidden_size=(args.moe_ffn_hidden_size if args.moe_ffn_hidden_size is not None - else args.ffn_hidden_size), - shared_expert_ffn_hidden_size=(0 if args.moe_shared_expert_intermediate_size is None - else args.moe_shared_expert_intermediate_size), + moe_ffn_hidden_size=( + args.moe_ffn_hidden_size + if args.moe_ffn_hidden_size is not None + else args.ffn_hidden_size + ), + shared_expert_ffn_hidden_size=( + 0 + if args.moe_shared_expert_intermediate_size is None + else args.moe_shared_expert_intermediate_size + ), num_experts_routed_to=args.moe_router_topk, gdn_qk_head_dim=args.linear_key_head_dim or 128, gdn_v_head_dim=args.linear_value_head_dim or 128, 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, + multi_latent_attention=args.multi_latent_attention, + 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), + 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. @@ -1410,7 +1363,7 @@ def get_start_time_from_progress_log(): def _get_field(string, type): return type(string.split(': ')[1]) - with maybe_msc.open(progress_log_filename, 'r') as f: + with open_file(progress_log_filename, 'r') as f: for line in f: line = line.strip() line_tokens = line.split('\t') @@ -1427,7 +1380,9 @@ def _get_field(string, type): # save_checkpoint was called directly (without save_checkpoint_and_time), # which writes "Saved async checkpoint" but not "Saving async checkpoint". if latest_num_floating_point_operations_uncommitted is not None: - latest_num_floating_point_operations = latest_num_floating_point_operations_uncommitted + latest_num_floating_point_operations = ( + latest_num_floating_point_operations_uncommitted + ) latest_num_floating_point_operations_uncommitted = None if world_size_in_line != args.world_size: # Re-start search if we see a different world size. @@ -1441,8 +1396,10 @@ def _get_field(string, type): assert ( start_time is not None and start_num_floating_point_operations is not None ), "Should have seen at least one 'Starting job' entry with same world_size" - print_rank_0(f"megatron.training.get_start_time_from_progress_log: " - f"{start_time=}, {start_num_floating_point_operations=}") + print_rank_0( + f"megatron.training.get_start_time_from_progress_log: " + f"{start_time=}, {start_num_floating_point_operations=}" + ) return datetime.strptime(start_time, '%Y-%m-%d %H:%M:%S'), start_num_floating_point_operations @@ -1459,6 +1416,7 @@ def preprocess_common_state_dict(common_state_dict): preprocessed_common_state_dict['args']['use_distributed_optimizer'] and "optimizer" in preprocessed_common_state_dict ): + def reorder_inner_param_groups(optimizer_state_dict): # When distributed optimizer loading, source param groups will be reordered, # so we reorder the param groups here to prevent warning. @@ -1474,13 +1432,7 @@ def reorder_inner_param_groups(optimizer_state_dict): if "param_groups" not in inner_optimizer: return param_groups = inner_optimizer["param_groups"] - # Treat missing and explicit None identifier values as equivalent. - # Wrap each component so None never compares directly with floats or strings. - def key_fn(pg): - return [ - (value is not None, value) - for value in (pg.get(key) for key in param_group_identifier_keys) - ] + key_fn = lambda pg: [pg[key] for key in param_group_identifier_keys] param_groups.sort(key=key_fn) inner_optimizer["param_groups"] = param_groups @@ -1592,15 +1544,10 @@ def pretrain( args = get_args() timers = get_timers() - # OTel span setup (_start_otel_job_spans) is deferred until after - # set_jit_fusion_options() below, where program_start/main_entry/pretrain_entry - # and the other startup timestamps are all available -- see the block right - # after those are extracted from _STARTUP_TIMESTAMPS. - if args.fine_grained_activation_offloading: from megatron.core.pipeline_parallel.utils import set_ideal_affinity_for_current_gpu - set_ideal_affinity_for_current_gpu() + set_ideal_affinity_for_current_gpu() if cfg_container.logger.log_progress: append_to_progress_log(args.save, "Starting job") @@ -1616,110 +1563,12 @@ def pretrain( main_entry = _STARTUP_TIMESTAMPS.get('main_entry') pretrain_entry = _STARTUP_TIMESTAMPS.get('pretrain_entry') - # OTel: emit the backdated pre_startup span and open this incarnation's - # megatron.startup trace (its own trace_id -- there are no run-long umbrella spans; - # see _start_otel_job_spans()'s docstring and the module comment above it for the - # full picture and why these use the explicit span API and backdating instead of - # managed_span). The megatron.startup.* spans below then nest inside that trace. - # The per-interval megatron.train blocks are separate traces again, chained by - # Links -- see _end_otel_startup_span(), called right before train()'s loop. - _start_otel_job_spans(model_type, program_start) - - _backdated_otel_span('megatron.startup.imports', program_start, main_entry) - _backdated_otel_span('megatron.startup.arg_parse', main_entry, pretrain_entry) - _backdated_otel_span( - 'megatron.startup.inprocess_setup', pretrain_entry, timestamp_after_inprocess_setup - ) - _backdated_otel_span( - 'megatron.startup.in_job_setup', - timestamp_after_inprocess_setup, - timestamp_after_in_job_setup, - ) - _backdated_otel_span( - 'megatron.startup.initialize_megatron', - timestamp_after_in_job_setup, - timestamp_after_initialize_megatron, - ) - _backdated_otel_span( - 'megatron.startup.jit_fusion_options', - timestamp_after_initialize_megatron, - timestamp_after_set_jit_fusion_options, - ) - - # Both hooks below exist ONLY to flush telemetry, and both are process-global side - # effects, so they are installed only when telemetry is actually active. With it off - # there is nothing to flush and Megatron must behave exactly as it does upstream -- - # in particular it must NOT grow a Python-level SIGTERM handler it never had, which - # would change EINTR behaviour and splice us into the ft_launcher/torchelastic - # handler chain. _otel_exit_hooks_installed keeps the installation itself idempotent. - global _otel_exit_hooks_installed - if _otel_telemetry_active() and not _otel_exit_hooks_installed: - _otel_exit_hooks_installed = True - - # An atexit fallback in case pretrain() exits via unhandled exception -- - # ensures both spans still get ended and telemetry flushed. Idempotent, so - # this doesn't double-flush if _end_otel_job_spans() already ran normally. - import atexit - atexit.register(_end_otel_job_spans) - - # atexit does NOT run on SIGTERM (CPython terminates without unwinding), and the - # ft_launcher tears down ranks with SIGTERM-then-SIGKILL on a fault/restart -- so - # without a SIGTERM handler, megatron.train (and the open span tree) leaks as a - # never-exported span in exactly the faulted runs we care about. - import signal as _signal - - _otel_prev_sigterm = _signal.getsignal(_signal.SIGTERM) - # with --exit-signal-handler, SIGTERM means "finish the current iteration, save a final - # checkpoint, then exit" (DistributedSignalHandler drain -> should_exit). We must not end - # spans / shut down the provider in that case -- that truncates megatron.train and drops the - # final-iteration + final-checkpoint spans; the normal should_exit path runs - # _end_otel_job_spans. There we only bounded-flush so nothing already-ended is lost if - # SIGKILL beats the drain. - # For an immediate/hard terminate (SIG_DFL / ft hard kill) we end + export the tree now. - _otel_graceful_drain = False - try: - _otel_graceful_drain = bool(getattr(get_args(), 'exit_signal_handler', False)) - except Exception: - pass - _otel_sigterm_fired = [False] # re-entry guard so a 2nd SIGTERM can't re-enter shutdown() - - def _otel_force_flush(): - try: - from opentelemetry import trace as _ot - _prov = _ot.get_tracer_provider() - if hasattr(_prov, 'force_flush'): - _prov.force_flush() - except Exception: - pass - - def _otel_sigterm_handler(signum, frame): - if not _otel_sigterm_fired[0]: - _otel_sigterm_fired[0] = True - try: - if _otel_graceful_drain: - _otel_force_flush() # graceful drain: flush only, keep provider alive - else: - _end_otel_job_spans() # hard terminate: end + export the tree now - except Exception: - pass - # chain to whatever handler was already installed (ft/torchelastic drain, or default) - if callable(_otel_prev_sigterm): - _otel_prev_sigterm(signum, frame) - elif _otel_prev_sigterm == _signal.SIG_DFL: - _signal.signal(_signal.SIGTERM, _signal.SIG_DFL) - os.kill(os.getpid(), signum) - # SIG_IGN: previous handler ignored SIGTERM -> nothing more to do. - - try: - _signal.signal(_signal.SIGTERM, _otel_sigterm_handler) - except (ValueError, OSError): - # Not the main thread / platform unsupported -> atexit remains the fallback. - pass - # Initialize program_start_global with a fallback value in case set_startup_timestamps() wasn't called program_start_global = _TRAIN_START_TIME if _STARTUP_TIMESTAMPS['program_start'] is not None: - program_start_global = torch.tensor([_STARTUP_TIMESTAMPS['program_start']], dtype=torch.double, device='cuda') + program_start_global = torch.tensor( + [_STARTUP_TIMESTAMPS['program_start']], dtype=torch.double, device='cuda' + ) torch.distributed.all_reduce(program_start_global, op=torch.distributed.ReduceOp.MIN) program_start_global = program_start_global.item() set_startup_timestamps(program_start=program_start_global) @@ -1739,26 +1588,41 @@ def _otel_sigterm_handler(signum, frame): # Print basic megatron init time (using global min start) # NOTE(asolergi-nv): This is not entirely accurate, but we keep it for backwards compatibility. print_rank_0( - 'time to initialize megatron (seconds): {:.3f}'.format(megatron_init_end - _LEGACY_TRAIN_START_TIME) + 'time to initialize megatron (seconds): {:.3f}'.format( + megatron_init_end - _LEGACY_TRAIN_START_TIME + ) ) # Note, not entirely accurate as rank 0 might not be the first or last to hit these timestamps - print_datetime('after in-process setup and before initialize_megatron', timestamp_after_inprocess_setup) - print_datetime('after in-job setup and before initialize_megatron', timestamp_after_in_job_setup) + print_datetime( + 'after in-process setup and before initialize_megatron', timestamp_after_inprocess_setup + ) + print_datetime( + 'after in-job setup and before initialize_megatron', timestamp_after_in_job_setup + ) if program_start is not None and main_entry is not None and pretrain_entry is not None: # Inject startup deltas into timers startup_timers = { - 'startup-program-entry-spread': program_start - program_start_global, # Local program start timestamp vs the global earliest program start timestamp - 'startup-library-setup': main_entry - program_start, # Local library imports - 'startup-program-setup': pretrain_entry - main_entry, # Local __main__ entry to pretrain entry - 'startup-in-process-setup': timestamp_after_inprocess_setup - pretrain_entry, # Local in-process setup - 'startup-in-job-setup': timestamp_after_in_job_setup - timestamp_after_inprocess_setup, # Local in-job setup - 'startup-initialize-megatron': timestamp_after_initialize_megatron - timestamp_after_in_job_setup, # Local initialize megatron - 'startup-set-jit-fusion-options': timestamp_after_set_jit_fusion_options - timestamp_after_initialize_megatron, # Local set JIT fusion options - 'all-reduce-start-timestamps-tensor': megatron_init_end - timestamp_after_set_jit_fusion_options, # 2x All-reduce, first collective call - 'startup-megatron-init-local': megatron_init_end - pretrain_entry, # Local megatron init - 'startup-megatron-init-global': megatron_init_end - program_start_global, # Local megatron init vs the global earliest program start timestamp + 'startup-program-entry-spread': program_start + - program_start_global, # Local program start timestamp vs the global earliest program start timestamp + 'startup-library-setup': main_entry - program_start, # Local library imports + 'startup-program-setup': pretrain_entry + - main_entry, # Local __main__ entry to pretrain entry + 'startup-in-process-setup': timestamp_after_inprocess_setup + - pretrain_entry, # Local in-process setup + 'startup-in-job-setup': timestamp_after_in_job_setup + - timestamp_after_inprocess_setup, # Local in-job setup + 'startup-initialize-megatron': timestamp_after_initialize_megatron + - timestamp_after_in_job_setup, # Local initialize megatron + 'startup-set-jit-fusion-options': timestamp_after_set_jit_fusion_options + - timestamp_after_initialize_megatron, # Local set JIT fusion options + 'all-reduce-start-timestamps-tensor': megatron_init_end + - timestamp_after_set_jit_fusion_options, # 2x All-reduce, first collective call + 'startup-megatron-init-local': megatron_init_end + - pretrain_entry, # Local megatron init + 'startup-megatron-init-global': megatron_init_end + - program_start_global, # Local megatron init vs the global earliest program start timestamp } for name, delta in startup_timers.items(): timers(name, log_level=0).set_elapsed(delta) @@ -1801,7 +1665,8 @@ def _otel_sigterm_handler(signum, frame): if cfg_container.checkpoint.replication: repl_strategy = CliqueReplicationStrategy.from_replication_params( - cfg_container.checkpoint.replication_jump, cfg_container.checkpoint.replication_factor + cfg_container.checkpoint.replication_jump, + cfg_container.checkpoint.replication_factor, ) else: repl_strategy = None @@ -1816,19 +1681,19 @@ def _otel_sigterm_handler(signum, frame): # Model, optimizer, and learning rate. timers('model-and-optimizer-setup', log_level=0).start(barrier=True) - with _otel_managed_span('model_init', 'megatron.startup.model_init', is_goodput_span=True): - model, optimizer, opt_param_scheduler = setup_model_and_optimizer( - model_type, - model_provider_func=model_provider, - checkpointing_context=checkpointing_context, - cfg_container=cfg_container, - # TODO (@maanug): temporary until initialize.py builds a pgcollection as bridge does. - pg_collection=( - pg_collection - if pg_collection is not None - else ProcessGroupCollection.use_mpu_process_groups() - ), - ) + model, optimizer, opt_param_scheduler = setup_model_and_optimizer( + model_type, + model_provider_func=model_provider, + checkpointing_context=checkpointing_context, + cfg_container=cfg_container, + # TODO (@maanug): temporary until initialize.py builds a pgcollection as bridge does. + pg_collection=( + pg_collection + if pg_collection is not None + else ProcessGroupCollection.use_mpu_process_groups() + ), + ) + timers('model-and-optimizer-setup').stop() print_datetime('after model, optimizer, and learning rate ' 'scheduler are built') model_cfg = get_model_config(model[0]) @@ -1869,7 +1734,9 @@ def _otel_sigterm_handler(signum, frame): # Build an isolated inference config so training config remains unchanged inference_config = copy.deepcopy(model_cfg) if args.rl_inference_tensor_model_parallel_size is not None: - inference_config.tensor_model_parallel_size = args.rl_inference_tensor_model_parallel_size + inference_config.tensor_model_parallel_size = ( + args.rl_inference_tensor_model_parallel_size + ) if args.rl_inference_pipeline_model_parallel_size is not None: inference_config.pipeline_model_parallel_size = ( args.rl_inference_pipeline_model_parallel_size @@ -1927,36 +1794,35 @@ def _otel_sigterm_handler(signum, frame): "This flag is only useful when doing refit since the weights are shared with the training model." ) - # Data stuff. Dataset index / dataloader construction (GPTDataset/BlendedDataset - # index building or loading from the cache) can be a multi-second chunk of - # startup on its own -- the 'data_loading' span group exists exactly for this. - # This shares scope with the existing perfetto-native trace_region; the two - # tracing systems are independent. + # Data stuff. app_metrics['app_build_dataiters_start_time'] = one_logger_utils.get_timestamp_in_ms() timers('train/valid/test-data-iterators-setup', log_level=0).start(barrier=True) - with _otel_managed_span('data_loading', 'megatron.startup.dataloader', is_goodput_span=True): - if args.virtual_pipeline_model_parallel_size is not None: - train_data_iterator = [] - valid_data_iterator = [] - test_data_iterator = [] - for vp_stage in range(len(model)): - dataset_provider_parameters = inspect.signature(train_valid_test_dataset_provider).parameters - assert "vp_stage" in dataset_provider_parameters, \ - "vp_stage must be a kwarg in train_valid_test_dataset_provider when using virtual pipeline parallelism" - vp_stage_train_valid_test_dataset_provider = \ - functools.partial(train_valid_test_dataset_provider, vp_stage=vp_stage) - if getattr(train_valid_test_dataset_provider, 'is_distributed', False): - vp_stage_train_valid_test_dataset_provider.is_distributed = True - iterators = build_train_valid_test_data_iterators( - vp_stage_train_valid_test_dataset_provider - ) - train_data_iterator.append(iterators[0]) - valid_data_iterator.append(iterators[1]) - test_data_iterator.append(iterators[2]) - else: - train_data_iterator, valid_data_iterator, test_data_iterator = ( - build_train_valid_test_data_iterators(train_valid_test_dataset_provider) + if args.virtual_pipeline_model_parallel_size is not None: + train_data_iterator = [] + valid_data_iterator = [] + test_data_iterator = [] + for vp_stage in range(len(model)): + dataset_provider_parameters = inspect.signature( + train_valid_test_dataset_provider + ).parameters + assert ( + "vp_stage" in dataset_provider_parameters + ), "vp_stage must be a kwarg in train_valid_test_dataset_provider when using virtual pipeline parallelism" + vp_stage_train_valid_test_dataset_provider = functools.partial( + train_valid_test_dataset_provider, vp_stage=vp_stage + ) + if getattr(train_valid_test_dataset_provider, 'is_distributed', False): + vp_stage_train_valid_test_dataset_provider.is_distributed = True + iterators = build_train_valid_test_data_iterators( + vp_stage_train_valid_test_dataset_provider ) + train_data_iterator.append(iterators[0]) + valid_data_iterator.append(iterators[1]) + test_data_iterator.append(iterators[2]) + else: + train_data_iterator, valid_data_iterator, test_data_iterator = ( + build_train_valid_test_data_iterators(train_valid_test_dataset_provider) + ) timers('train/valid/test-data-iterators-setup').stop() print_datetime('after dataloaders are built') app_metrics['app_build_dataiters_finish_time'] = one_logger_utils.get_timestamp_in_ms() @@ -2007,14 +1873,7 @@ def _otel_sigterm_handler(signum, frame): ) print_rank_0(f'[RLProfiler] Profiling enabled, output: {profile_dir}') - # OTel: megatron.startup deliberately stays open into train() -- its - # preamble (weight-hash check, sniff test, cuda-graph setup) is init, not - # the loop, so it should parent to startup. train() ends startup and starts - # megatron.train itself, right before its loop. If train() isn't called - # (do_train False), _end_otel_job_spans() ends startup as an idempotent - # fallback. - - if not cfg_container.validation.skip_train or (args.perform_rl_step and args.do_train): + if not cfg_container.validation.skip_train or args.perform_rl_step: if cfg_container.validation.skip_train: print_rank_0('RL inference-only mode (--skip-train --perform-rl-step) ...') else: @@ -2023,37 +1882,30 @@ def _otel_sigterm_handler(signum, frame): iteration = 0 args.curr_iteration = iteration if args.do_train and (args.train_iters or 0) > 0: - try: - iteration, num_floating_point_operations_so_far = train( - forward_step_func, - model, - optimizer, - opt_param_scheduler, - train_data_iterator, - valid_data_iterator, - process_non_loss_data_func, - model_cfg, - checkpointing_context, - non_loss_data_func, - inference_model, - p2p_communicator=p2p_communicator, - pg_collection=pg_collection, - ) - except Exception: - # OTel: an uncaught training exception (a real hardware/CUDA/NCCL - # fault, or an injected fault) would skip the _end_otel_job_spans() - # at the bottom of pretrain(), leaving the in-flight megatron.train block - # unended -> unexported (only the graceful-SIGTERM and normal-exit paths - # flush it today). End + flush the open spans here, then re-raise so the - # process still dies exactly as before. Idempotent -> no double-flush. - # (SystemExit from the exit-interval path is NOT caught here; it already - # flushed before sys.exit(). SIGKILL is unrecoverable regardless.) - _end_otel_job_spans() - raise + iteration, num_floating_point_operations_so_far = train( + forward_step_func, + model, + optimizer, + opt_param_scheduler, + train_data_iterator, + valid_data_iterator, + process_non_loss_data_func, + model_cfg, + checkpointing_context, + non_loss_data_func, + inference_model, + p2p_communicator=p2p_communicator, + pg_collection=pg_collection, + ) print_datetime('after training is done') - if not cfg_container.validation.skip_train and cfg_container.checkpoint.save and iteration != 0 and iteration % cfg_container.checkpoint.save_interval != 0: + if ( + not cfg_container.validation.skip_train + and cfg_container.checkpoint.save + and iteration != 0 + and iteration % cfg_container.checkpoint.save_interval != 0 + ): save_checkpoint_and_time( iteration, model, @@ -2061,7 +1913,7 @@ def _otel_sigterm_handler(signum, frame): opt_param_scheduler, num_floating_point_operations_so_far, checkpointing_context, - train_data_iterator=train_data_iterator + train_data_iterator=train_data_iterator, ) one_logger and one_logger.log_metrics( @@ -2072,7 +1924,6 @@ def _otel_sigterm_handler(signum, frame): print_rank_0('skipping training (--skip-train is on) ...') iteration = args.iteration - args.curr_iteration = iteration if args.do_valid: prefix = f'iteration {iteration} on validation set' @@ -2097,12 +1948,18 @@ def _otel_sigterm_handler(signum, frame): ) else: evaluate_and_print_results( - prefix, forward_step_func, - valid_data_iterator, model, - iteration, process_non_loss_data_func, model_cfg, - verbose=True, write_to_tensorboard=not cfg_container.validation.skip_train, + prefix, + forward_step_func, + valid_data_iterator, + model, + iteration, + process_non_loss_data_func, + model_cfg, + verbose=True, + write_to_tensorboard=not cfg_container.validation.skip_train, non_loss_data_func=non_loss_data_func, - pg_collection=pg_collection, p2p_communicator=p2p_communicator + pg_collection=pg_collection, + p2p_communicator=p2p_communicator, ) if args.do_test: @@ -2137,20 +1994,9 @@ def _otel_sigterm_handler(signum, frame): if args.perform_rl_step: rl_utils.rl_inference_interface_shutdown() - if getattr(args, 'gtp_remat_nccl_ub', False) or getattr(args, 'gtp_expert_remat_nccl_ub', False): - from megatron.core.tensor_parallel.gtp_api import deregister_and_clear_gtp_symm_pools - - # Deregister the GTP symmetric-memory pools: windows left registered when the - # process groups are destroyed make NCCL abort. - deregister_and_clear_gtp_symm_pools() - ft_integration.shutdown() one_logger_utils.finish() - # OTel: close any still-open span (interval block / startup) and flush all pending - # telemetry data. - _end_otel_job_spans() - def update_train_iters(args): @@ -2200,7 +2046,6 @@ def wrap_model_chunks_with_ddp( ddp_config, *, use_layer_wise_distributed_optimizer=False, - use_layer_wise_param_layout=True, DP=DDP, pg_collection=None, bucket_sizes=None, @@ -2211,12 +2056,13 @@ def wrap_model_chunks_with_ddp( Centralises the DDP-wrapping wiring shared between :func:`get_model` and unit tests. - For ``use_layer_wise_distributed_optimizer=True`` and ``use_layer_wise_param_layout=True``: - forces ``ddp_config.use_distributed_optimizer=True`` (mutated in place; needed - for reduce-scatter), and computes per-chunk shard-aligned layouts via - :meth:`LayerWiseDistributedOptimizer.compute_full_param_layout`. With - ``use_layer_wise_param_layout=False``, no layout is supplied and LayerWise falls back - to its legacy ``allgather_params`` sync path. + For ``use_layer_wise_distributed_optimizer=True``: forces + ``ddp_config.use_distributed_optimizer=True`` (mutated in place; needed for reduce-scatter) + and computes per-chunk layouts via + :meth:`LayerWiseDistributedOptimizer.compute_full_param_layout`. That method picks the + padded shard-aligned LayerWise layout or the compact decoupled layout per + ``ddp_config.use_layer_wise_param_layout`` (True → padded, False → compact, with LayerWise + buffers treated as non-DistOpt and synced via legacy ``allgather_params``). For non-layerwise with ``ddp_config.use_distributed_optimizer=True``: computes per-chunk byte-level layouts via @@ -2232,11 +2078,9 @@ def wrap_model_chunks_with_ddp( model_chunks: List of model chunks to wrap (un-DDP-wrapped). config: :class:`TransformerConfig`. ddp_config: :class:`DistributedDataParallelConfig`. Mutated in place when - ``use_layer_wise_distributed_optimizer=True`` and ``use_layer_wise_param_layout=True``. + ``use_layer_wise_distributed_optimizer=True``. Its ``use_layer_wise_param_layout`` + field selects the padded vs compact decoupled LayerWise layout (default). use_layer_wise_distributed_optimizer: Whether the layerwise wiring runs. - use_layer_wise_param_layout: When ``use_layer_wise_distributed_optimizer=True``, - controls whether to compute and supply a shard-aligned param layout - to DDP. ``False`` keeps LayerWise on its legacy sync path. DP: The DDP class to construct (``DistributedDataParallel`` or an FSDP variant). pg_collection: Optional :class:`ProcessGroupCollection`. When provided, @@ -2259,13 +2103,18 @@ def wrap_model_chunks_with_ddp( # Compute per-chunk layouts (DDP only). per_chunk_layouts = [None] * n if DP is DDP: - if use_layer_wise_distributed_optimizer and use_layer_wise_param_layout: + if use_layer_wise_distributed_optimizer: + # LayerWise (Muon) optimizer. Force use_distributed_optimizer=True so sibling + # non-LayerWise buffers (embeddings, biases, layernorm) shard with the byte-level + # DistributedOptimizer layout, and tag params so DDP buffer grouping routes + # LayerWise-managed matrices (Muon's Newton-Schulz domain) to a separate buffer. + # The padded-vs-compact LayerWise layout decision is made inside + # compute_full_param_layout / _ParamAndGradBuffer from use_layer_wise_param_layout: + # by default (compact) LayerWise buffers get the no-padding layout and the per-buffer + # override flips use_distributed_optimizer off for them; with + # --use-layer-wise-param-layout they stay on the padded DistOpt layout. ddp_config.use_distributed_optimizer = True compute_layout = LayerWiseDistributedOptimizer.compute_full_param_layout - # Tag params so DDP buffer grouping routes LayerWise-managed matrices - # (Muon's Newton-Schulz domain) to a shard-aligned buffer and routes - # everything else (embeddings, biases, layernorm) to a separate - # DistOpt-style buffer. tag_params_for_buffer_routing(model_chunks) elif not use_layer_wise_distributed_optimizer and ddp_config.use_distributed_optimizer: compute_layout = DistributedOptimizer.compute_full_param_layout @@ -2281,22 +2130,8 @@ def wrap_model_chunks_with_ddp( "wrap_model_chunks_with_ddp requires a dp_cp process group to size " "the distributed-optimizer parameter layout" ) - # The distributed optimizer shards each bucket over the intra-instance group, which - # is what DDP hands to the buffer as its data_parallel_group. Size the layout by that - # same group, otherwise the layout reports more shards than the reduce-scatter uses - # and the trailing shards of every bucket end up owned by no rank. intra_dp_cp is - # the full dp_cp when num_distributed_optimizer_instances is 1, so this only differs - # when there are several instances. - intra_dp_cp_group = getattr(layout_pgs, "intra_dp_cp", None) - intra_expt_dp_group = getattr(layout_pgs, "intra_expt_dp", None) - data_parallel_world_size = get_pg_size( - intra_dp_cp_group if intra_dp_cp_group is not None else layout_pgs.dp_cp - ) - expert_data_parallel_world_size = get_pg_size( - intra_expt_dp_group - if intra_expt_dp_group is not None - else getattr(layout_pgs, "expt_dp", None) - ) + data_parallel_world_size = get_pg_size(layout_pgs.dp_cp) + expert_data_parallel_world_size = get_pg_size(getattr(layout_pgs, "expt_dp", None)) for i, (chunk, bucket_size) in enumerate(zip(model_chunks, bucket_sizes)): all_params = [p for p in chunk.parameters() if p.requires_grad] per_chunk_layouts[i] = compute_layout( @@ -2330,32 +2165,13 @@ def wrap_model_chunks_with_ddp( return wrapped -def _freeze_all_model_chunks(model_list): - """Freeze all parameters in a list of model chunks (for logits-saving runs).""" - for model_module in model_list: - model_module.requires_grad_(False) - # Additionally freeze expert biases of routers - for module in model_module.modules(): - if hasattr(module, "frozen_expert_bias"): - module.frozen_expert_bias = True - return model_list - - -def _forward_backward_grad_context(args): - """Grad context for a train step's forward/backward pass. - - Returns a tuple of (grad_context, forward_only). - grad_context is ``torch.no_grad()`` when all layers are frozen (e.g. teacher logits - dumps), no parameter needs gradients, so there is no reason to build the - autograd graph. Otherwise returns a no-op context. - forward_only is True when all layers are frozen, False otherwise. - """ - grad_context = torch.no_grad() if getattr(args, "freeze_all_layers", False) else nullcontext() - forward_only = getattr(args, "freeze_all_layers", False) - return grad_context, forward_only - - -def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True, config=None, pg_collection=None): +def get_model( + model_provider_func, + model_type=ModelType.encoder_or_decoder, + wrap_with_ddp=True, + config=None, + pg_collection=None, +): """Build the model.""" args = get_args() args.model_type = model_type @@ -2363,10 +2179,13 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap pg_collection = ProcessGroupCollection.use_mpu_process_groups() if args.create_all_gather_group: - timeout = timedelta(minutes=args.distributed_timeout_minutes) if args.distributed_timeout_minutes else None + timeout = ( + timedelta(minutes=args.distributed_timeout_minutes) + if args.distributed_timeout_minutes + else None + ) dp_cp_ag, expt_dp_ag = create_all_gather_groups( - for_expert_parallelism=(args.expert_model_parallel_size > 1), - timeout=timeout, + for_expert_parallelism=(args.expert_model_parallel_size > 1), timeout=timeout ) pg_collection.dp_cp_ag = dp_cp_ag pg_collection.expt_dp_ag = expt_dp_ag @@ -2376,7 +2195,16 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap print_rank_0("> including expert parallelism AG group") if has_nvidia_modelopt: - maybe_enable_modelopt(args) + from megatron.post_training.checkpointing import has_modelopt_state + + # [ModelOpt]: Check if the checkpoint is a ModelOpt checkpoint and + # set a flag to use our model provider if so. + if args.load is not None and has_modelopt_state(args.load): + print_rank_0(f'ModelOpt checkpoint detected') + args.modelopt_enabled = True + elif getattr(args, "export_kd_teacher_load", None): + # For distillation ckpts without ModelOpt state + args.modelopt_enabled = True # Build model. def build_model(): @@ -2416,7 +2244,6 @@ def build_model(): model.model_type = model_type return model - if args.init_model_with_meta_device: with torch.device('meta'): model = build_model() @@ -2428,7 +2255,12 @@ def build_model(): # For rare operations like post-training logits saving if args.freeze_all_layers: - _freeze_all_model_chunks(model) + for model_module in model: + model_module.requires_grad_(False) + # Additionally freeze expert biases of routers + for module in model_module.modules(): + if hasattr(module, "frozen_expert_bias"): + module.frozen_expert_bias = True # Set tensor model parallel attributes if not set. # Only parameters that are already tensor model parallel have these @@ -2444,12 +2276,9 @@ def build_model(): ) if get_pg_rank(pg_collection.dp) == 0 and get_pg_rank(pg_collection.cp) == 0: print( - ' > number of parameters on (tensor, gtp_weight_remat, pipeline) ' - 'model parallel rank ({}, {}, {}): {}'.format( - get_pg_rank(pg_collection.tp), - get_pg_rank(pg_collection.gtp_remat), - get_pg_rank(pg_collection.pp), - num_parameters, + ' > number of parameters on (tensor, pipeline) ' + 'model parallel rank ({}, {}): {}'.format( + get_pg_rank(pg_collection.tp), get_pg_rank(pg_collection.pp), num_parameters ), flush=True, ) @@ -2471,7 +2300,10 @@ def build_model(): # Materialize tensors on meta device (GPU allocation) if not using FSDP2 and not using Megatron FSDP. if args.init_model_with_meta_device and not args.use_torch_fsdp2 and not args.use_megatron_fsdp: - model = [to_empty_if_meta_device(model_module, device=torch.device("cuda")) for model_module in model] + model = [ + to_empty_if_meta_device(model_module, device=torch.device("cuda")) + for model_module in model + ] # Before TE2.x: The model_module.bfloat16()/model_module.half() above will call the inplace # copy of TE's Float8Tensor, which will write an unwanted value (amax calculated @@ -2485,7 +2317,7 @@ def build_model(): assert HAVE_FSDP2, "Torch FSDP2 requires torch>=2.4.0" DP = torch_FSDP elif args.use_megatron_fsdp: - DP = FullyShardedDataParallel + DP = megatron_FSDP else: DP = DDP @@ -2495,10 +2327,7 @@ def build_model(): if not getattr(args, "use_torch_fsdp2", False): # In the Megatron FSDP and DDP use path, we need to initialize the bucket size. ddp_config.bucket_size = resolve_ddp_bucket_size( - ddp_config, - pg_collection.dp_cp, - ddp_config.overlap_grad_reduce, - num_parameters, + ddp_config, pg_collection.dp_cp, ddp_config.overlap_grad_reduce, num_parameters ) # Compute per-chunk bucket sizes / disable_bucketing flags. Bucketing is @@ -2514,15 +2343,12 @@ def build_model(): for disable in per_chunk_disable_bucketing ] - if config.cuda_graph_impl == "full_iteration": - # DDP initialization must use the full-iteration capture stream so its retained - # AccumulateGrad nodes do not reference a different, non-capturing stream. - ddp_stream = get_shared_capture_stream() - else: - # Preserve a dedicated initialization stream for all other implementations. - ddp_stream = torch.cuda.Stream() + # Setup stream for ddp initialization. The side-stream may be necessary for cuda graph + # capture support with DDP, but we sync it with the current stream to avoid races. + ddp_stream = torch.cuda.Stream() + # Wait for the default stream to complete before starting ddp_stream ddp_stream.wait_stream(torch.cuda.current_stream()) - + # Make ddp_stream start after whatever the default stream already queued with torch.cuda.stream(ddp_stream): model = wrap_model_chunks_with_ddp( model, @@ -2531,15 +2357,13 @@ def build_model(): use_layer_wise_distributed_optimizer=getattr( args, 'use_layer_wise_distributed_optimizer', False ), - use_layer_wise_param_layout=getattr( - args, 'use_layer_wise_param_layout', True - ), DP=DP, pg_collection=pg_collection if args.use_megatron_fsdp else None, bucket_sizes=per_chunk_bucket_sizes, disable_bucketing_per_chunk=per_chunk_disable_bucketing, ) - # Ensure initialization-stream work completes before touching params on the default stream. + # End of setup_stream + # Critical: ensure side-stream work completes before touching params on default stream torch.cuda.current_stream().wait_stream(ddp_stream) # Broadcast params from data parallel src rank to other data parallel ranks. @@ -2621,6 +2445,7 @@ def get_megatron_optimizer_config(args: Any) -> OptimizerConfig: return config, config_overrides + def get_megatron_ddp_config(args: argparse.Namespace) -> DistributedDataParallelConfig: """Return an MCore DDPConfig from the argparse arguments.""" @@ -2638,16 +2463,18 @@ def get_megatron_ddp_config(args: argparse.Namespace) -> DistributedDataParallel kwargs["num_buckets"] = args.ddp_num_buckets kwargs["bucket_size"] = args.ddp_bucket_size kwargs["pad_buckets_for_high_nccl_busbw"] = args.ddp_pad_buckets_for_high_nccl_busbw - kwargs["reduce_scatter_with_fp32_accumulation"] = args.ddp_reduce_scatter_with_fp32_accumulation - kwargs["param_name_patterns_for_fp32_local_accumulation"] = \ - tuple(args.ddp_param_name_patterns_for_fp32_local_accumulation) + kwargs["reduce_scatter_with_fp32_accumulation"] = ( + args.ddp_reduce_scatter_with_fp32_accumulation + ) + kwargs["param_name_patterns_for_fp32_local_accumulation"] = tuple( + args.ddp_param_name_patterns_for_fp32_local_accumulation + ) kwargs["average_in_collective"] = args.ddp_average_in_collective # Megatron-FSDP arguments. kwargs["megatron_fsdp_main_params_dtype"] = args.megatron_fsdp_main_params_dtype kwargs["megatron_fsdp_main_grads_dtype"] = args.megatron_fsdp_main_grads_dtype kwargs["megatron_fsdp_grad_comm_dtype"] = args.megatron_fsdp_grad_comm_dtype - if args.use_megatron_fsdp and args.megatron_fsdp_version == 1: - kwargs["megatron_fsdp_use_decoupled_grad"] = args.use_precision_aware_optimizer + kwargs["megatron_fsdp_use_decoupled_grad"] = args.use_precision_aware_optimizer if args.use_megatron_fsdp and args.cuda_graph_impl != "none": # Run Megatron-FSDP in CUDA graph-safe mode. Avoids some graph-unsafe host-side # operations (such as pointer dereferencing) that can break CUDA graph replay. @@ -2682,9 +2509,6 @@ def setup_model_and_optimizer( skip_optimizer = not (has_normal_optimizer or has_rl_optimizer) wrap_with_ddp = not skip_optimizer - if has_nvidia_modelopt: - maybe_enable_modelopt(args) - def _build_model_wrapper(wrap_with_ddp: bool): if cfg_container is not None and getattr(cfg_container, "model", None) is not None: from megatron.training.utils import start_memory_history_recording @@ -2695,12 +2519,6 @@ def _build_model_wrapper(wrap_with_ddp: bool): model_config = cfg.model builder_cls = model_config.get_builder_cls() builder = builder_cls(model_config) - - # Inject freeze_all_layers as a pre-wrap hook so DDP sees requires_grad=False - # and skips grad-buffer allocation for all params (matching get_model behavior). - if args.freeze_all_layers: - model_config.pre_wrap_hooks.append(_freeze_all_model_chunks) - return builder.build_distributed_models( pg_collection=pg_collection, ddp_config=cfg.ddp, @@ -2710,54 +2528,22 @@ def _build_model_wrapper(wrap_with_ddp: bool): wrap_with_ddp=wrap_with_ddp, data_parallel_random_init=cfg.rng.data_parallel_random_init, use_layer_wise_distributed_optimizer=cfg.optimizer.use_layer_wise_distributed_optimizer, - use_layer_wise_param_layout=getattr(args, 'use_layer_wise_param_layout', True), ) else: - assert model_provider_func is not None, "Must provide a model config via config_container or a model_provider_func." - return get_model(model_provider_func, model_type, wrap_with_ddp=wrap_with_ddp, pg_collection=pg_collection) - - # Configure GTP weight-remat padding/loss reduction before model construction (pad - # alignment governs how dim-0 shards are built). Placed here (not in get_model) so it - # also covers the config-container builder path, which does not call get_model. - if is_gtp_remat_active(args): - from megatron.core.process_groups_config import resolve_gtp_remat_group - from megatron.core.tensor_parallel.gtp_api import ( - configure_gtp_remat_from_recipe, - register_gtp_symm_pool, - ) - - configure_gtp_remat_from_recipe( - fp4=getattr(args, 'fp4', None) is not None, - fp8_recipe=getattr(args, 'fp8_recipe', None), - fp8=getattr(args, 'fp8', None) is not None, - calculate_per_token_loss=getattr(args, 'calculate_per_token_loss', False), - reduce_scatter_with_fp32_accumulation=getattr( - args, 'gtp_remat_reduce_scatter_with_fp32_accumulation', False - ), - ) - - if getattr(args, 'gtp_remat_nccl_ub', False): - register_gtp_symm_pool(resolve_gtp_remat_group(pg_collection, is_expert=False)) - if getattr(args, 'gtp_expert_remat_nccl_ub', False): - register_gtp_symm_pool(resolve_gtp_remat_group(pg_collection, is_expert=True)) + assert ( + model_provider_func is not None + ), "Must provide a model config via config_container or a model_provider_func." + return get_model( + model_provider_func, + model_type, + wrap_with_ddp=wrap_with_ddp, + pg_collection=pg_collection, + ) model = _build_model_wrapper(wrap_with_ddp) unwrapped_model = unwrap_model(model) - # Classify each GTP param's prefetch chain after model build + DDP wrap, before the - # first forward. Placed here (not in get_model) so it also covers the config-container - # builder path. - if is_gtp_remat_active(args): - from megatron.core.tensor_parallel.gtp_api import classify_gtp_remat_chains - - classify_gtp_remat_chains( - model, - cuda_graph_modules=getattr(args, 'cuda_graph_modules', None), - moe_shared_expert_overlap=getattr(args, 'moe_shared_expert_overlap', False), - cuda_graph_impl=getattr(args, 'cuda_graph_impl', 'none'), - ) - - if args.logits_save_dir is not None and mpu.is_pipeline_last_stage(): + if args.logits_save_dir is not None: from megatron.training.distillation import LogitsSaverHooks logits_saver = LogitsSaverHooks( @@ -2769,13 +2555,15 @@ def _build_model_wrapper(wrap_with_ddp: bool): ) logits_saver.attach_hooks(unwrapped_model[-1]) - if args.logits_load_dir is not None and mpu.is_pipeline_last_stage(): + if args.logits_load_dir is not None: from megatron.training.distillation import StudentLogitsCapture student_logits_capture = StudentLogitsCapture() student_logits_capture.attach_hooks(unwrapped_model[-1]) - one_logger and one_logger.log_metrics({"app_build_optimzer_start_time": one_logger_utils.get_timestamp_in_ms()}) + one_logger and one_logger.log_metrics( + {"app_build_optimzer_start_time": one_logger_utils.get_timestamp_in_ms()} + ) if skip_optimizer: optimizer, opt_param_scheduler = None, None # In RL inference-only mode, train_iters must still be set despite having no optimizer. @@ -2806,7 +2594,9 @@ def _build_model_wrapper(wrap_with_ddp: bool): ) opt_param_scheduler = get_optimizer_param_scheduler(optimizer) - one_logger and one_logger.log_metrics({"app_build_optimzer_finish_time": one_logger_utils.get_timestamp_in_ms()}) + one_logger and one_logger.log_metrics( + {"app_build_optimzer_finish_time": one_logger_utils.get_timestamp_in_ms()} + ) if args.moe_use_upcycling: torch.distributed.barrier() @@ -2862,31 +2652,23 @@ def _build_model_wrapper(wrap_with_ddp: bool): {'load_checkpoint_start_time': one_logger_utils.get_timestamp_in_ms()} ) timers('load-checkpoint', log_level=0).start(barrier=True) - with _otel_managed_span('load_checkpoint', 'megatron.checkpoint.load', is_goodput_span=True): - ckpt_pgc = getattr(unwrapped_model[0], "pg_collection", None) - args.iteration, args.num_floating_point_operations_so_far = load_checkpoint( - model, - optimizer, - opt_param_scheduler, - checkpointing_context=checkpointing_context, - skip_load_to_model_and_opt=HAVE_FSDP2 - and getattr(args, "use_torch_fsdp2", False) - and args.ckpt_format == "torch_dist", - tp_group=ckpt_pgc.tp if ckpt_pgc is not None else None, - pp_group=ckpt_pgc.pp if ckpt_pgc is not None else None, - # Replica_id must match the save path (see save_checkpoint_and_time): use the - # gtp_remat-inclusive group, not replicate dp_cp, or gtp_remat peers collide. - dp_cp_group=getattr(ckpt_pgc, "dp_cp_gtp_remat", None), - dp_group=ckpt_pgc.dp if ckpt_pgc is not None else None, - expt_dp_group=ckpt_pgc.expt_dp if ckpt_pgc is not None else None, - rng_state_key_prefix=getattr(unwrapped_model[0], "rng_state_key_prefix", ""), - ) - # Barrier + min/max all-reduce right after the load. Unlike the checkpoint - # SAVE (ragged writers -> cross-rank skew at timers.log), the fully-parallel - # LOAD is uniform across ranks (~ms spread), so no meaningful skew - # serializes here -- left unspanned; megatron.checkpoint.load covers the - # (uniform) restart cost. + ckpt_pgc = getattr(unwrapped_model[0], "pg_collection", None) + args.iteration, args.num_floating_point_operations_so_far = load_checkpoint( + model, + optimizer, + opt_param_scheduler, + checkpointing_context=checkpointing_context, + skip_load_to_model_and_opt=HAVE_FSDP2 + and getattr(args, "use_torch_fsdp2", False) + and args.ckpt_format == "torch_dist", + tp_group=ckpt_pgc.tp if ckpt_pgc is not None else None, + pp_group=ckpt_pgc.pp if ckpt_pgc is not None else None, + dp_cp_group=ckpt_pgc.dp_cp if ckpt_pgc is not None else None, + dp_group=ckpt_pgc.dp if ckpt_pgc is not None else None, + expt_dp_group=ckpt_pgc.expt_dp if ckpt_pgc is not None else None, + rng_state_key_prefix=getattr(unwrapped_model[0], "rng_state_key_prefix", ""), + ) timers('load-checkpoint').stop(barrier=True) timers.log(['load-checkpoint']) one_logger and one_logger.log_metrics( @@ -2899,14 +2681,6 @@ def _build_model_wrapper(wrap_with_ddp: bool): args.iteration = 0 args.num_floating_point_operations_so_far = 0 - # [ModelOpt]: Load the teacher checkpoint for ModelOpt distillation if applicable. - # Import locally to prevent circular import: megatron.post_training.checkpointing - # imports `get_args` from megatron.training at module scope. - if has_nvidia_modelopt: - from megatron.post_training.checkpointing import load_kd_teacher_checkpoint - - load_kd_teacher_checkpoint(model) - # Validate that the world size can accommodate the current batch size. # This catches the case where GPUs were scaled up mid-training but the # current position in the batch size schedule yields a batch size that @@ -2964,12 +2738,21 @@ def dummy_train_step(data_iterator): """Single dummy training step.""" args = get_args() tp_rank = mpu.get_tensor_model_parallel_rank() - has_cu_seqlens = getattr(args, 'sft', False) or getattr(args, 'dataloader_inter_document_masking', False) - is_hybrid_cp = args.hybrid_context_parallel + has_cu_seqlens = getattr(args, 'sft', False) or getattr( + args, 'dataloader_inter_document_masking', False + ) + is_hybrid_cp = args.dynamic_context_parallel BATCH_KEYS = [ - "tokens", "labels", "loss_mask", "position_ids", "attention_mask", - "cu_seqlens", "cu_seqlens_padded", "max_seqlen", "local_cp_size", + "tokens", + "labels", + "loss_mask", + "position_ids", + "attention_mask", + "cu_seqlens", + "cu_seqlens_padded", + "max_seqlen", + "local_cp_size", "hybrid_cp_group", ] @@ -2982,7 +2765,11 @@ def dummy_train_step(data_iterator): if tp_rank == 0: batch = next(data_iterator) for key in BATCH_KEYS: - batch[key] = batch[key].cuda(non_blocking=True) if key in batch and batch[key] is not None else None + batch[key] = ( + batch[key].cuda(non_blocking=True) + if key in batch and batch[key] is not None + else None + ) batch = get_batch_on_this_tp_rank( batch, broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), @@ -3003,11 +2790,22 @@ def dummy_train_step(data_iterator): batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), - hybrid_cp_group_func=get_hybrid_data_context_parallel_groups, + hybrid_cp_group_func=get_dynamic_data_context_parallel_groups, ) -def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=None, pg_collection: Optional[ProcessGroupCollection | MultiModuleProcessGroupCollection] = None, p2p_communicator: Optional[P2PCommunicator] = None): +def train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=None, + pg_collection: Optional[ProcessGroupCollection | MultiModuleProcessGroupCollection] = None, + p2p_communicator: Optional[P2PCommunicator] = None, +): """Single training step. pg_collection: optional carrier forwarded to the schedule for the cross-grid case; None @@ -3017,25 +2815,88 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch """ args = get_args() timers = get_timers() + num_microbatches = get_num_microbatches() + + optimizer_config = getattr(optimizer, 'config', None) + chunked_optimizer_state_offload = bool( + optimizer_config is not None + and getattr(optimizer_config, 'chunked_optimizer_state_offload', False) + and getattr(optimizer_config, 'optimizer_state_offload_fraction', 1.0) > 0.0 + ) + pre_forward_param_sync_before_master_offload = ( + chunked_optimizer_state_offload + and optimizer.optimizer_state_offload_requires_pre_forward_param_sync() + ) + delay_master_offload_for_param_buffer = ( + chunked_optimizer_state_offload + and args.reuse_grad_buf_for_mxfp8_param_ag + and args.overlap_param_gather + ) + if chunked_optimizer_state_offload: + # Prefetch masters and the first state chunk as late as possible so H2D overlaps + # gradient finalization. Preserve custom hooks and avoid wrapping more than once. + finalize_model_grads_func = getattr(config, 'finalize_model_grads_func', None) + if ( + getattr( + finalize_model_grads_func, + '_chunked_optimizer_state_offload_wrapped_optimizer', + None, + ) + is not optimizer + ): + base_finalize_model_grads_func = getattr( + finalize_model_grads_func, + '_chunked_optimizer_state_offload_base_finalize_model_grads_func', + None, + ) + if base_finalize_model_grads_func is None: + base_finalize_model_grads_func = finalize_model_grads_func or finalize_model_grads + + def finalize_model_grads_with_state_reload(*fmg_args, **fmg_kwargs): + optimizer.prefetch_optimizer_state_for_gradient_finalization() + return base_finalize_model_grads_func(*fmg_args, **fmg_kwargs) - # OTel: set up per-step sub-span support. - _otel_step_tracer = None - if _otel_sg_enabled('forward_backward') or _otel_sg_enabled('optimizer'): - from nemo.lens.helpers import span_cm, safe_set_span_attributes as _otel_set_attrs - _otel_step_tracer = get_telemetry().tracer + setattr( + finalize_model_grads_with_state_reload, + '_chunked_optimizer_state_offload_wrapped_optimizer', + optimizer, + ) + setattr( + finalize_model_grads_with_state_reload, + '_chunked_optimizer_state_offload_base_finalize_model_grads_func', + base_finalize_model_grads_func, + ) + config.finalize_model_grads_func = finalize_model_grads_with_state_reload rerun_state_machine = get_rerun_state_machine() - save_params_in_this_iteration = (args.save_params_interval is not None and - (iteration + 1) % args.save_params_interval == 0) - save_activations_in_this_iteration = (args.save_activations_interval is not None and - (iteration + 1) % args.save_activations_interval == 0) - save_tpe_in_this_iteration = (args.save_tokens_per_expert_interval is not None and - (iteration + 1) % args.save_tokens_per_expert_interval == 0) - save_wgrads_in_this_iteration = (args.save_wgrads_interval is not None and - (iteration + 1) % args.save_wgrads_interval == 0) - save_dgrads_in_this_iteration = (args.save_dgrads_interval is not None and - (iteration + 1) % args.save_dgrads_interval == 0) - while rerun_state_machine.should_run_forward_backward(data_iterator): + packed_data_iterator = None + has_wrapped_data_iterator = False + rerun_data_iterator = data_iterator + save_params_in_this_iteration = ( + args.save_params_interval is not None and (iteration + 1) % args.save_params_interval == 0 + ) + save_activations_in_this_iteration = ( + args.save_activations_interval is not None + and (iteration + 1) % args.save_activations_interval == 0 + ) + save_tpe_in_this_iteration = ( + args.save_tokens_per_expert_interval is not None + and (iteration + 1) % args.save_tokens_per_expert_interval == 0 + ) + save_wgrads_in_this_iteration = ( + args.save_wgrads_interval is not None and (iteration + 1) % args.save_wgrads_interval == 0 + ) + save_dgrads_in_this_iteration = ( + args.save_dgrads_interval is not None and (iteration + 1) % args.save_dgrads_interval == 0 + ) + while rerun_state_machine.should_run_forward_backward(rerun_data_iterator): + # Start D2H before zeroing gradients. In the MXFP8 staging path, updated masters + # remain readable until their delayed D2H after the main-param copy below. + if chunked_optimizer_state_offload and not pre_forward_param_sync_before_master_offload: + optimizer.offload_optimizer_state_for_forward( + offload_master=not delay_master_offload_for_param_buffer + ) + # Set grad to zero. for model_chunk in model: model_chunk.zero_grad_buffer() @@ -3043,6 +2904,16 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch model_chunk.force_all_reduce = save_wgrads_in_this_iteration optimizer.zero_grad() + if pre_forward_param_sync_before_master_offload: + # Compact LayerWise FP8 gather consumes fp32 masters. Finish the normally + # overlapped gather after grad-buffer reset. If a sibling DistOpt still needs + # MXFP8 staging, keep masters until that copy below completes. + optimizer.ensure_master_weights_for_pre_forward_param_sync() + optimizer.start_param_sync_for_bucket_group_subset(force_sync=True) + optimizer.offload_optimizer_state_for_forward( + offload_master=not delay_master_offload_for_param_buffer + ) + if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False): # Distillation shape-adjust reads parallel_state; only for modelopt-enabled runs. adjust_tensor_shapes_fn = get_tensor_shapes_adjust_fn_for_distillation( @@ -3068,54 +2939,71 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch # and subsequent param_data zero are baked into the graph and replay # unconditionally. We must populate param_data so the replayed AG gathers # correct weights, even when forward pre-hooks are disabled. - # A model component's DDP config may differ from the global args. Check the - # DistributedOptimizer config that owns the buffers and hooks; non-overlapped - # optimizers stage during optimizer.step(). - optimizer_instances = getattr(optimizer, 'chained_optimizers', [optimizer]) - mxfp8_overlap_optimizers = [ - optim_instance - for optim_instance in optimizer_instances - if ( - isinstance(optim_instance, DistributedOptimizer) - and optim_instance.ddp_config.reuse_grad_buf_for_mxfp8_param_ag - and optim_instance.ddp_config.overlap_param_gather - ) - ] - if mxfp8_overlap_optimizers: + if args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather: # Check if forward_pre_hook is enabled by checking if hooks are registered. forward_pre_hook_enabled = len(model[0].remove_forward_pre_hook_handles) > 0 full_cg_captured = FullCudaGraphWrapper.cuda_graph.get("training") is not None if forward_pre_hook_enabled or full_cg_captured: - for optim_instance in mxfp8_overlap_optimizers: - optim_instance._copy_main_params_to_param_buffer() + for optim_instance in optimizer.chained_optimizers: + if isinstance(optim_instance, DistributedOptimizer): + # Only this DistOpt sibling consumes masters in the MXFP8 param-buffer + # staging pass. The staging entry restores its own offloaded masters; + # do not perturb the LayerWise/Muon master lifecycle here. + optim_instance._copy_main_params_to_param_buffer() + + # In the delayed MXFP8 path, optimizer state was offloaded above while masters + # remained resident for main-param staging. Start master D2H after that copy. + if delay_master_offload_for_param_buffer: + optimizer.offload_optimizer_state_for_forward() # Forward pass. if save_activations_in_this_iteration: - enable_activation_logging(model, args.save) - if save_tpe_in_this_iteration: - enable_tokens_per_expert_logging(model, args.save) - if save_dgrads_in_this_iteration: - enable_dgrad_logging(model, args.save) - grad_context, forward_only = _forward_backward_grad_context(args) - _fb_cm = ( - span_cm("megatron.train.iteration.forward_backward", tracer=_otel_step_tracer, num_microbatches=get_num_microbatches()) - if _otel_sg_enabled('forward_backward') and _otel_step_tracer is not None else nullcontext() - ) - with grad_context, _fb_cm: - losses_reduced = forward_backward_func( - forward_step_func=forward_step_func, - data_iterator=data_iterator, - model=model, - num_microbatches=get_num_microbatches(), - seq_length=args.seq_length, - micro_batch_size=args.micro_batch_size, - decoder_seq_length=args.decoder_seq_length, - forward_only=forward_only, - adjust_tensor_shapes_fn=adjust_tensor_shapes_fn, - force_all_reduce=save_wgrads_in_this_iteration, - p2p_communicator=p2p_communicator, - pg_collection=pg_collection, + enable_activation_logging(model, args.save) + if save_tpe_in_this_iteration: + enable_tokens_per_expert_logging(model, args.save) + if save_dgrads_in_this_iteration: + enable_dgrad_logging(model, args.save) + if getattr(config, 'sequence_packing_scheduler', None) is not None: + # Dynamic-CP / sequence packing must happen after the rerun state machine has + # observed the original RerunDataIterator. The scheduler returns another + # RerunDataIterator containing the packed microbatches for this step. + if not has_wrapped_data_iterator: + ( + packed_data_iterator, + num_microbatches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) = wrap_data_iterator(data_iterator, config, get_num_microbatches()) + has_wrapped_data_iterator = True + rerun_data_iterator = packed_data_iterator + forward_backward_data_iterator = packed_data_iterator + else: + num_microbatches = get_num_microbatches() + seqlen_sum_this_global_batch = args.seq_length * args.global_batch_size + seqlen_squared_sum_this_global_batch = args.seq_length**2 * args.global_batch_size + forward_backward_data_iterator = data_iterator + + if getattr(config, "mtp_num_layers", None): + # Writer objects only exist on the last rank, so use globally consistent + # configuration flags to keep the acceptance collective branch rank-aligned. + has_acceptance_consumer = bool( + getattr(args, "tensorboard_dir", None) or getattr(args, "wandb_project", "") ) + MTPLossLoggingHelper.configure_acceptance_collection(enabled=has_acceptance_consumer) + losses_reduced = forward_backward_func( + forward_step_func=forward_step_func, + data_iterator=forward_backward_data_iterator, + model=model, + num_microbatches=num_microbatches, + seq_length=args.seq_length, + micro_batch_size=args.micro_batch_size, + decoder_seq_length=args.decoder_seq_length, + forward_only=False, + adjust_tensor_shapes_fn=adjust_tensor_shapes_fn, + force_all_reduce=save_wgrads_in_this_iteration, + p2p_communicator=p2p_communicator, + pg_collection=pg_collection, + ) if save_activations_in_this_iteration: save_activations(iteration + 1) disable_activation_logging() @@ -3155,7 +3043,19 @@ def _save_state_dict(attr_name, label): should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit() if should_exit: - return {}, True, should_checkpoint, should_exit, exit_code, None, None, 0 + return ( + {}, + True, + should_checkpoint, + should_exit, + exit_code, + None, + None, + 0, + num_microbatches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) # Empty unused memory. if args.empty_unused_memory_level >= 1: @@ -3169,16 +3069,7 @@ def _save_state_dict(attr_name, label): # Update parameters. timers('optimizer', log_level=1).start(barrier=args.barrier_with_L1_time) - _opt_cm = ( - span_cm("megatron.train.iteration.optimizer", tracer=_otel_step_tracer) - if _otel_sg_enabled('optimizer') and _otel_step_tracer is not None else nullcontext() - ) - with _opt_cm as _opt_span: - update_successful, grad_norm, num_zeros_in_grad = optimizer.step() - if _opt_span is not None: - _otel_set_attrs(_opt_span, { - "megatron.update_successful": bool(update_successful), - }) + update_successful, grad_norm, num_zeros_in_grad = optimizer.step() # get max attention logit for logging and run clip_qk() # Part of MuonClip Optimizer step @@ -3197,13 +3088,11 @@ def _save_state_dict(attr_name, label): if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() for _required in ("mp", "pp", "dp_cp"): - assert getattr(pg_collection, _required, None) is not None, ( - f"model pg_collection used by train_step must define {_required}" - ) + assert ( + getattr(pg_collection, _required, None) is not None + ), f"model pg_collection used by train_step must define {_required}" mp_group = pg_collection.mp - # gtp_remat-inclusive: the reported global per-token loss must cover gtp_remat peers' distinct - # tokens (replicate dp_cp would report a 1/gtp_remat subsample -> per-step noisy). Display-only. - dp_cp_group = getattr(pg_collection, 'dp_cp_gtp_remat', None) or pg_collection.dp_cp + dp_cp_group = pg_collection.dp_cp is_last_stage = is_pp_last_stage(pg_collection.pp) # when freezing sub-models we may have a mixture of successful and unsucessful ranks, # so we must gather across mp ranks @@ -3223,15 +3112,7 @@ def _save_state_dict(attr_name, label): # Update learning rate. if update_successful: - # data_parallel_size excludes the GTP-remat axis (it's folded into total_model_size at - # arguments.py:446); each gtp-remat peer consumes a distinct microbatch, so multiply it - # back in for the sample count. - increment = ( - get_num_microbatches() - * args.micro_batch_size - * args.data_parallel_size - * args.gtp_weight_remat_size - ) + increment = get_num_microbatches() * args.micro_batch_size * args.data_parallel_size opt_param_scheduler.step(increment=increment) skipped_iter = 0 else: @@ -3268,8 +3149,23 @@ def _save_state_dict(attr_name, label): grad_norm, num_zeros_in_grad, log_max_attention_logit, + num_microbatches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, ) - return {}, skipped_iter, should_checkpoint, should_exit, exit_code, grad_norm, num_zeros_in_grad, log_max_attention_logit + return ( + {}, + skipped_iter, + should_checkpoint, + should_exit, + exit_code, + grad_norm, + num_zeros_in_grad, + log_max_attention_logit, + num_microbatches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) def training_log( @@ -3288,6 +3184,7 @@ def training_log( is_first_iteration=False, seqlen_squared_sum_in_batch: float | None = None, total_real_tokens_in_batch: float | None = None, + num_microbatches: int | None = None, ): """Log training information such as losses, timing, ....""" args = get_args() @@ -3329,48 +3226,45 @@ def training_log( # Logging. timers_to_log = [] if args.timing_log_level >= 1: - timers_to_log.extend([ - 'forward-backward', - 'layernorm-grads-all-reduce', - 'embedding-grads-all-reduce', - 'all-grads-sync', - 'params-all-gather', - 'optimizer-copy-to-main-grad', - 'optimizer-unscale-and-check-inf', - 'optimizer-clip-main-grad', - 'optimizer-count-zeros', - 'optimizer-inner-step', - 'optimizer-copy-main-to-model-params', - 'optimizer', - ]) + timers_to_log.extend( + [ + 'forward-backward', + 'layernorm-grads-all-reduce', + 'embedding-grads-all-reduce', + 'all-grads-sync', + 'params-all-gather', + 'optimizer-copy-to-main-grad', + 'optimizer-unscale-and-check-inf', + 'optimizer-clip-main-grad', + 'optimizer-count-zeros', + 'optimizer-inner-step', + 'optimizer-copy-main-to-model-params', + 'optimizer', + ] + ) if args.timing_log_level >= 2: - timers_to_log.extend([ - 'batch-generator', - 'forward-compute', - 'backward-compute', - 'forward-recv', - 'forward-send', - 'backward-recv', - 'backward-send', - 'forward-send-forward-recv', - 'forward-send-backward-recv', - 'backward-send-forward-recv', - 'backward-send-backward-recv', - 'forward-backward-send-forward-backward-recv', - ]) + timers_to_log.extend( + [ + 'batch-generator', + 'forward-compute', + 'backward-compute', + 'forward-recv', + 'forward-send', + 'backward-recv', + 'backward-send', + 'forward-send-forward-recv', + 'forward-send-backward-recv', + 'backward-send-forward-recv', + 'backward-send-backward-recv', + 'forward-backward-send-forward-backward-recv', + ] + ) # Add timers from RL loop if needed. if args.perform_rl_step: timers_to_log.extend(RL_LOGGABLE_TIMER_NAMES) - # Calculate batch size. data_parallel_size excludes the GTP-remat axis (it's folded into - # total_model_size at arguments.py:446); each gtp-remat peer consumes a distinct microbatch, - # so multiply it back in for the global sample count. - batch_size = ( - args.micro_batch_size - * args.data_parallel_size - * args.gtp_weight_remat_size - * get_num_microbatches() - ) + # Calculate batch size. + batch_size = args.micro_batch_size * args.data_parallel_size * get_num_microbatches() # Track app tag & app tag ID one_logger_utils.track_app_tag(batch_size, args.world_size, args.seq_length) @@ -3390,7 +3284,9 @@ def training_log( wandb_writer.log({'samples vs steps': args.consumed_train_samples}, iteration) if learning_rate is not None: writer.add_scalar('learning-rate', learning_rate, iteration) - writer.add_scalar('learning-rate vs samples', learning_rate, args.consumed_train_samples) + writer.add_scalar( + 'learning-rate vs samples', learning_rate, args.consumed_train_samples + ) if wandb_writer: wandb_writer.log({'learning-rate': learning_rate}, iteration) if args.skipped_train_samples > 0: @@ -3441,10 +3337,14 @@ def training_log( if wandb_writer: wandb_writer.log({'params-norm': params_norm}, iteration) if args.perform_rl_step: - grpo_collection_iteration = iteration // (args.grpo_iterations * ( ( args.grpo_samples_per_iteration )// args.global_batch_size )) + grpo_collection_iteration = iteration // ( + args.grpo_iterations * ((args.grpo_samples_per_iteration) // args.global_batch_size) + ) writer.add_scalar('grpo_collection_iteration', grpo_collection_iteration, iteration) if wandb_writer: - wandb_writer.log({'grpo_collection_iteration': grpo_collection_iteration}, iteration) + wandb_writer.log( + {'grpo_collection_iteration': grpo_collection_iteration}, iteration + ) if args.log_memory_to_tensorboard: mem_stats = torch.cuda.memory_stats() writer.add_scalar( @@ -3464,7 +3364,7 @@ def training_log( # Log MoE metrics. moe_log_string = "" if args.num_experts is not None: - moe_loss_scale = 1 / get_num_microbatches() + moe_loss_scale = 1 / (num_microbatches or get_num_microbatches()) track_names = [] if "aux_loss" in args.moe_router_load_balancing_type: track_names.append("load_balancing_loss") @@ -3475,44 +3375,25 @@ def training_log( if args.moe_z_loss_coeff is not None: track_names.append("z_loss") + moe_layer_freq = args.moe_layer_freq + mtp_num_layers = args.mtp_num_layers if is_hybrid_model(args): - from megatron.core.models.hybrid.hybrid_layer_allocation import ( + from megatron.core.ssm.mamba_hybrid_layer_allocation import ( Symbols, parse_hybrid_pattern, ) - parsed_pattern = parse_hybrid_pattern(args.hybrid_layer_pattern) - main_pattern = parsed_pattern.main_pattern or "" - mtp_pattern = parsed_pattern.mtp_pattern or "" - main_moe_layers = main_pattern.count(Symbols.MOE) - mtp_moe_layers_per_depth = mtp_pattern.count(Symbols.MOE) - if parsed_pattern.mtp_num_depths > 0 and mtp_moe_layers_per_depth > 0: - mtp_moe_layers = ( - mtp_moe_layers_per_depth - if args.mtp_use_repeated_layer - else mtp_moe_layers_per_depth * parsed_pattern.mtp_num_depths - ) - else: - mtp_moe_layers = 0 - num_moe_layers = main_moe_layers + mtp_moe_layers + parsed_hybrid_pattern = parse_hybrid_pattern(args.hybrid_layer_pattern) + main_pattern = (parsed_hybrid_pattern.main_pattern or "").replace(Symbols.PIPE, "") + layers = len(main_pattern) + parsed_hybrid_pattern.mtp_num_depths + moe_layer_freq = [int(layer_type == Symbols.MOE) for layer_type in main_pattern] + moe_layer_freq.extend( + parsed_hybrid_pattern.mtp_pattern.count(Symbols.MOE) + for _ in range(parsed_hybrid_pattern.mtp_num_depths) + ) + mtp_num_layers = None else: - if args.moe_layer_freq is None: - moe_layer_pattern = [1] * args.num_layers - elif isinstance(args.moe_layer_freq, int): - moe_layer_pattern = [ - 1 if (i % args.moe_layer_freq == 0) else 0 for i in range(args.num_layers) - ] - elif isinstance(args.moe_layer_freq, list): - moe_layer_pattern = args.moe_layer_freq - else: - raise ValueError(f"Invalid moe_layer_freq: {args.moe_layer_freq}") - main_moe_layers = sum(moe_layer_pattern) - mtp_moe_layers = 0 - if args.mtp_num_layers and moe_layer_pattern[-1]: - mtp_moe_layers = 1 if args.mtp_use_repeated_layer else args.mtp_num_layers - num_moe_layers = main_moe_layers + mtp_moe_layers - - layers = args.num_layers + (args.mtp_num_layers or 0) + layers = args.num_layers moe_log_string = get_moe_metrics_tracker().report( loss_scale=moe_loss_scale, @@ -3523,15 +3404,22 @@ def training_log( force_initialize=True, track_names=track_names, num_layers=layers, - num_moe_layers=num_moe_layers, - moe_layer_freq=args.moe_layer_freq, + moe_layer_freq=moe_layer_freq, + mtp_num_layers=mtp_num_layers, pg_collection=pg_collection, total_loss_dict=total_loss_dict, ) # Log MTP metrics. if args.mtp_num_layers is not None: - mtp_loss_scale = 1 / get_num_microbatches() + if args.calculate_per_token_loss: + # The tracker already reduces raw loss sums and token counts into a + # per-token loss, matching the main loss normalization path. + mtp_loss_scale = 1.0 + else: + # Legacy mode accumulates microbatch-normalized losses, so average + # by the scheduled microbatch count for this step. + mtp_loss_scale = 1 / (num_microbatches or get_num_microbatches()) MTPLossLoggingHelper.track_mtp_metrics( mtp_loss_scale, iteration, writer, wandb_writer, total_loss_dict ) @@ -3545,27 +3433,35 @@ def training_log( writer=writer, wandb_writer=wandb_writer, total_loss_dict=total_loss_dict, + num_layers=args.num_layers + (args.mtp_num_layers or 0), + num_indexer_layers=( + sum(ratio == 4 for ratio in args.csa_compress_ratios) + if args.csa_compress_ratios is not None + else None + ), + preserve_groups=args.cuda_graph_impl != "none", ) # Dump memory snapshot and print metrics to stdout. if iteration % args.log_interval == 0 or is_first_iteration: - should_prof_rank = (args.profile_ranks == [] or safe_get_rank() in args.profile_ranks) # [] is all ranks - if args.record_memory_history and (should_prof_rank or torch.distributed.get_backend() == 'fake'): - rank = safe_get_rank() - base, ext = os.path.splitext(args.memory_snapshot_path) - snapshot_filename = f"{base}_{rank}{ext}" - torch.cuda.memory._dump_snapshot(snapshot_filename) + if args.record_memory_history and ( + is_last_rank() or torch.distributed.get_backend() == 'fake' + ): + snapshot = torch.cuda.memory._snapshot() + from pickle import dump + + with open(args.memory_snapshot_path, 'wb') as f: + dump(snapshot, f) elapsed_time = timers('interval-time').elapsed(barrier=True, reset=should_reset) elapsed_time_per_iteration = elapsed_time / total_iterations - llm_world_size = getattr(args, 'mimo_llm_world_size', args.world_size) throughput = num_floating_point_operations( args, batch_size, seqlen_squared_sum_in_batch=seqlen_squared_sum_in_batch, total_real_tokens_in_batch=total_real_tokens_in_batch, - ) / (elapsed_time_per_iteration * 10**12 * llm_world_size) + ) / (elapsed_time_per_iteration * 10**12 * args.world_size) one_logger_utils.track_e2e_metrics(args.log_throughput, throughput) @@ -3609,27 +3505,6 @@ def training_log( if learning_rate is not None: log_string += f' learning rate: {learning_rate:.6E} |' log_string += f' global batch size: {batch_size:5d} |' - - # OTel: snapshot the accumulator state BEFORE it is torn down. The loop directly - # below zeroes each loss tensor when should_reset, and the should_reset block - # further down zeroes advanced/skipped/nan_iters -- so anything read after that - # point sees 0.0 loss and 0 skipped iterations on EVERY export, which is exactly - # what the metrics emission (further below, where the other emission inputs like - # `throughput` are in scope) used to do. Take the values here and emit them there. - # The .item() is a device sync, so it stays inside the telemetry guard and is not - # paid at all when telemetry is off -- the same guard the emission itself uses. - _otel_telemetry_log = get_telemetry() - _otel_loss_snapshot = None - _otel_skipped_iters_snapshot = 0 - if _otel_telemetry_log is not None and _otel_telemetry_log.is_exporting: - _meta_keys = (advanced_iters_key, skipped_iters_key, nan_iters_key) - _loss_keys = [k for k in total_loss_dict if k not in _meta_keys] - if _loss_keys: - _otel_loss_snapshot = total_loss_dict[_loss_keys[0]].item() / float( - max(1, total_loss_dict.get(advanced_iters_key, 1)) - ) - _otel_skipped_iters_snapshot = int(total_loss_dict.get(skipped_iters_key, 0)) - for key in total_loss_dict: if key not in [advanced_iters_key, skipped_iters_key, nan_iters_key]: avg = total_loss_dict[key].item() / float( @@ -3656,7 +3531,7 @@ def training_log( # RL token throughput metrics. if args.perform_rl_step: log_string += rl_utils.log_rl_throughput_metrics( - args, batch_size, elapsed_time_per_iteration, iteration, wandb_writer, + args, batch_size, elapsed_time_per_iteration, iteration, wandb_writer ) if should_reset: @@ -3664,30 +3539,6 @@ def training_log( total_loss_dict[skipped_iters_key] = 0 total_loss_dict[nan_iters_key] = 0 print_rank_last(log_string) - - # OTel: emit training metrics at log interval (export rank only). Loss and - # skipped-iteration counts come from the snapshot taken above, before the - # accumulators were reset; everything else is still live at this point. - if _otel_telemetry_log is not None and _otel_telemetry_log.is_exporting: - from megatron.core.telemetry.training_metrics import record_training_metrics - _avg_loss = _otel_loss_snapshot - _tokens_per_sec = ( - batch_size * args.seq_length / elapsed_time_per_iteration - if elapsed_time_per_iteration > 0 else None - ) - _mem_gb = torch.cuda.max_memory_allocated() / (1024 ** 3) if torch.cuda.is_available() else None - record_training_metrics( - meter=_otel_telemetry_log.meter, - step_duration_ms=elapsed_time_per_iteration * 1000.0, - loss=_avg_loss, - throughput_tflops=throughput if args.log_throughput else None, - grad_norm=grad_norm, - learning_rate=learning_rate, - skipped_iters=_otel_skipped_iters_snapshot, - tokens_per_sec=_tokens_per_sec, - memory_allocated_gb=_mem_gb, - ) - reported_memory_in_this_iteration = False if report_memory_flag: # Report memory after optimizer state has been initialized. @@ -3703,8 +3554,11 @@ def training_log( if iteration > (loaded_iteration + 1): # Make sure the memory after the second iteration is reported to include optimizer state memory. report_memory_flag = False - if args.log_memory_interval is not None and iteration % args.log_memory_interval == 0 and \ - not reported_memory_in_this_iteration: + if ( + args.log_memory_interval is not None + and iteration % args.log_memory_interval == 0 + and not reported_memory_in_this_iteration + ): report_memory( f'(after {iteration} iterations)', process_group=pg_collection.dp if pg_collection is not None else None, @@ -3724,8 +3578,12 @@ def training_log( # Write timers to wandb, don't reset the counts. if args.log_timers_to_tensorboard: - timers.write(timers_to_log, writer, iteration, normalizer=args.log_interval, reset=False) - timers.write(timers_to_log, wandb_writer, iteration, normalizer=args.log_interval, reset=False) + timers.write( + timers_to_log, writer, iteration, normalizer=args.log_interval, reset=False + ) + timers.write( + timers_to_log, wandb_writer, iteration, normalizer=args.log_interval, reset=False + ) # Log timers to stdout timers.log(timers_to_log, normalizer=args.log_interval, reset=should_reset) @@ -3736,7 +3594,6 @@ def compute_throughputs_and_append_to_progress_log(iteration, num_floating_point args = get_args() if args.save is None: return - llm_world_size = getattr(args, 'mimo_llm_world_size', args.world_size) # Compute job throughput. # args.num_floating_point_operations_so_far keeps track of floating-point operations @@ -3744,7 +3601,7 @@ def compute_throughputs_and_append_to_progress_log(iteration, num_floating_point global _TRAIN_START_TIME job_throughput = ( num_floating_point_operations_so_far - args.num_floating_point_operations_so_far - ) / ((time.time() - _TRAIN_START_TIME) * 10**12 * llm_world_size) + ) / ((time.time() - _TRAIN_START_TIME) * 10**12 * args.world_size) # Compute cumulative throughput since jobs of this world size were launched. # `get_start_time_from_progress_log` returns start time and number of floating-point @@ -3753,7 +3610,7 @@ def compute_throughputs_and_append_to_progress_log(iteration, num_floating_point elapsed_time = (datetime.now() - start_time).total_seconds() cumulative_throughput = ( num_floating_point_operations_so_far - start_num_floating_point_operations - ) / (elapsed_time * 10**12 * llm_world_size) + ) / (elapsed_time * 10**12 * args.world_size) tokens_so_far = args.consumed_train_samples * args.seq_length saved_ckpt_prefix = 'Saving async checkpoint' if args.async_save else 'Saved checkpoint' @@ -3763,33 +3620,13 @@ def compute_throughputs_and_append_to_progress_log(iteration, num_floating_point f"Job throughput: {job_throughput:.1f} TFLOP/s/GPU\t" f"Cumulative throughput: {cumulative_throughput:.1f} TFLOP/s/GPU\t" f"Floating-point operations: {num_floating_point_operations_so_far:.2e}\t" - f"Tokens (in billions): {tokens_so_far / 10**9:.2f}" - ) - - -def _assert_param_gather_overlap_model(model_chunk): - """Assert that a model chunk implements the parameter-gather overlap lifecycle.""" - # MimoModel is a composite wrapper rather than a DDP instance, but delegates this - # interface to its active inner DDP modules. - required_methods = ( - 'enable_forward_pre_hook', - 'disable_forward_pre_hook', - 'start_param_sync', - ) - missing_methods = [ - method_name - for method_name in required_methods - if not callable(getattr(model_chunk, method_name, None)) - ] - assert not missing_methods, ( - f'{type(model_chunk).__name__} does not support parameter-gather overlap; ' - f'missing callable methods: {", ".join(missing_methods)}' + f"Tokens (in billions): {tokens_so_far / 10**9:.2f}", ) def enable_forward_pre_hook(model_chunks): for model_chunk in model_chunks: - _assert_param_gather_overlap_model(model_chunk) + assert isinstance(model_chunk, DDP) model_chunk.enable_forward_pre_hook() @@ -3797,15 +3634,15 @@ def disable_forward_pre_hook(model_chunks, optimizer=None, param_sync=True): if param_sync and optimizer is not None: optimizer.prepare_model_params_for_param_sync() for model_chunk in model_chunks: - _assert_param_gather_overlap_model(model_chunk) + assert isinstance(model_chunk, DDP) model_chunk.disable_forward_pre_hook(param_sync=param_sync) -def force_param_sync(model_chunks, optimizer=None) -> None: +def force_param_sync(model_chunks: list[DDP], optimizer=None) -> None: if optimizer is not None: optimizer.prepare_model_params_for_param_sync() for model_chunk in model_chunks: - _assert_param_gather_overlap_model(model_chunk) + assert isinstance(model_chunk, DDP) model_chunk.start_param_sync(force_sync=True) @@ -3823,149 +3660,108 @@ def save_checkpoint_and_time( timers = get_timers() energy_monitor = get_energy_monitor() - # OTel: single goodput span parenting ALL exposed (main-thread) checkpoint - # cost -- the total wall-clock the training loop is blocked on the checkpoint: - # pre-save buffer free, state-dict generation, the async dispatch, AND the - # post-save cross-rank skew wait at timers.log. Everything below nests under - # it (megatron.func.save_checkpoint_and_time, megatron.checkpoint.save, - # timers_log, post_save_barrier). start_span + context attach (not a `with`) - # so the whole body isn't re-indented; ended explicitly at the function end. - _exposed_save_span = None - _exposed_save_token = None - if _otel_sg_enabled('checkpoint'): - from opentelemetry import context as _octx, trace as _otr - _exposed_save_span = get_telemetry().tracer.start_span('megatron.checkpoint.exposed_save') - _otel_mark_goodput(_exposed_save_span) - _exposed_save_span.set_attribute('megatron.iteration', iteration) - _exposed_save_token = _octx.attach(_otr.set_span_in_context(_exposed_save_span)) - try: - - # Synchronize forward pre-hook state before checkpoint save to avoid race conditions - if should_disable_forward_pre_hook(args): - force_param_sync(model, optimizer=optimizer) + # Synchronize forward pre-hook state before checkpoint save to avoid race conditions + if should_disable_forward_pre_hook(args): + force_param_sync(model, optimizer=optimizer) - # Stop timer to get accurate train interval time and exclude checkpointing duration - timers('interval-time').stop() - energy_monitor.pause() + # Stop timer to get accurate train interval time and exclude checkpointing duration + timers('interval-time').stop() + energy_monitor.pause() - # Log E2E metrics before save-checkpoint - one_logger_utils.track_e2e_metrics() - # Free overlap param-gather buffers so that the async checkpoint worker - # process has enough GPU headroom for D2H tensor transfers. - with _otel_managed_span('checkpoint', 'megatron.checkpoint.reclaim_memory', is_goodput_span=True): - for model_chunk in model: - if hasattr(model_chunk, 'free_overlap_buffers'): - model_chunk.free_overlap_buffers() - - # timer.log() reports the min & max time. We do not need a barrier here. - timer_key = 'save-checkpoint-non-persistent' if non_persistent_ckpt else 'save-checkpoint' - timers(timer_key, log_level=0).start(barrier=False) - - # Resolve checkpoint groups from this rank's module PGC; None for stock runs - # falls back to the mpu groups inside save_checkpoint (byte-identical). - ckpt_pgc = getattr(unwrap_model(model)[0], "pg_collection", None) - tp_group = getattr(ckpt_pgc, "tp", None) if ckpt_pgc is not None else None - pp_group = getattr(ckpt_pgc, "pp", None) if ckpt_pgc is not None else None - dp_group = getattr(ckpt_pgc, "dp", None) if ckpt_pgc is not None else None - # Replica_id needs the gtp_remat-inclusive group (dp_cp_gtp_remat), not replicate dp_cp. - dp_cp_group = getattr(ckpt_pgc, "dp_cp_gtp_remat", None) if ckpt_pgc is not None else None - expt_dp_group = getattr(ckpt_pgc, "expt_dp", None) if ckpt_pgc is not None else None - # Per-grid rng key namespace set by a multi-grid model; '' for stock single-grid. - rng_state_key_prefix = getattr(unwrap_model(model)[0], "rng_state_key_prefix", "") - - global num_checkpoints_memory_reported, MAX_NUM_CHECKPOINTS_MEMORY_REPORTED - should_report_memory = num_checkpoints_memory_reported < MAX_NUM_CHECKPOINTS_MEMORY_REPORTED - - if should_report_memory: - # Track memory before checkpoint save. - report_memory(f"(before save_checkpoint for iteration {iteration})", process_group=dp_group) - - # Save checkpoint. - with _otel_managed_span('checkpoint', 'megatron.checkpoint.save', is_goodput_span=True, **{'megatron.iteration': iteration}): - save_checkpoint( - iteration, - model, - optimizer, - opt_param_scheduler, - num_floating_point_operations_so_far, - checkpointing_context, - non_persistent_ckpt=non_persistent_ckpt, - train_data_iterator=train_data_iterator, - preprocess_common_state_dict_fn=preprocess_common_state_dict, - tp_group=tp_group, - pp_group=pp_group, - dp_cp_group=dp_cp_group, - dp_group=dp_group, - expt_dp_group=expt_dp_group, - rng_state_key_prefix=rng_state_key_prefix, - ) + # Log E2E metrics before save-checkpoint + one_logger_utils.track_e2e_metrics() + # Free overlap param-gather buffers and release cached GPU memory so + # that the async checkpoint worker process has enough GPU headroom for + # D2H tensor transfers. + for model_chunk in model: + if hasattr(model_chunk, 'free_overlap_buffers'): + model_chunk.free_overlap_buffers() + torch.cuda.empty_cache() + + # timer.log() reports the min & max time. We do not need a barrier here. + timer_key = 'save-checkpoint-non-persistent' if non_persistent_ckpt else 'save-checkpoint' + timers(timer_key, log_level=0).start(barrier=False) + + # Resolve checkpoint groups from this rank's module PGC; None for stock runs + # falls back to the mpu groups inside save_checkpoint (byte-identical). + ckpt_pgc = getattr(unwrap_model(model)[0], "pg_collection", None) + tp_group = getattr(ckpt_pgc, "tp", None) if ckpt_pgc is not None else None + pp_group = getattr(ckpt_pgc, "pp", None) if ckpt_pgc is not None else None + dp_group = getattr(ckpt_pgc, "dp", None) if ckpt_pgc is not None else None + dp_cp_group = getattr(ckpt_pgc, "dp_cp", None) if ckpt_pgc is not None else None + expt_dp_group = getattr(ckpt_pgc, "expt_dp", None) if ckpt_pgc is not None else None + # Per-grid rng key namespace set by a multi-grid model; '' for stock single-grid. + rng_state_key_prefix = getattr(unwrap_model(model)[0], "rng_state_key_prefix", "") + + global num_checkpoints_memory_reported, MAX_NUM_CHECKPOINTS_MEMORY_REPORTED + should_report_memory = num_checkpoints_memory_reported < MAX_NUM_CHECKPOINTS_MEMORY_REPORTED + + if should_report_memory: + # Track memory before checkpoint save. + report_memory(f"(before save_checkpoint for iteration {iteration})", process_group=dp_group) + + # Save checkpoint. + save_checkpoint( + iteration, + model, + optimizer, + opt_param_scheduler, + num_floating_point_operations_so_far, + checkpointing_context, + non_persistent_ckpt=non_persistent_ckpt, + train_data_iterator=train_data_iterator, + preprocess_common_state_dict_fn=preprocess_common_state_dict, + tp_group=tp_group, + pp_group=pp_group, + dp_cp_group=dp_cp_group, + dp_group=dp_group, + expt_dp_group=expt_dp_group, + rng_state_key_prefix=rng_state_key_prefix, + ) - # Stop timer and compute time elapsed to save checkpoint. Stop timer before timers.log() call as it resets the timer. - # Since timer.log() reports the min & max time, we do not need a barrier here. - timers(timer_key).stop(barrier=False) - save_checkpoint_duration = timers(timer_key).elapsed(reset=False) - if should_report_memory: - # Track memory after checkpoint save. - with _otel_managed_span('checkpoint', 'megatron.checkpoint.report_memory', is_goodput_span=True): - report_memory(f"(after save_checkpoint for iteration {iteration})", process_group=dp_group) - num_checkpoints_memory_reported += 1 - - if args.fp8: - # Run garbage collection after checkpoint saving to free memory from - # dequantized bf16 tensors that were temporarily created during fp8 - # model checkpoint saving. - with _otel_managed_span('checkpoint', 'megatron.checkpoint.gc_collect', is_goodput_span=True): - gc.collect() + # Stop timer and compute time elapsed to save checkpoint. Stop timer before timers.log() call as it resets the timer. + # Since timer.log() reports the min & max time, we do not need a barrier here. + timers(timer_key).stop(barrier=False) + save_checkpoint_duration = timers(timer_key).elapsed(reset=False) + + if should_report_memory: + # Track memory after checkpoint save. + report_memory(f"(after save_checkpoint for iteration {iteration})", process_group=dp_group) + num_checkpoints_memory_reported += 1 - # timers.log reports min & max across ranks -> a collective. Sitting right - # after the per-rank-imbalanced save (and outside the save spans), it's a - # prime spot for the fast ranks to serialize waiting for the slowest saver. - with _otel_managed_span('step', 'megatron.checkpoint.timers_log', is_goodput_span=True): - timers.log([timer_key]) + if args.fp8: + # Run garbage collection after checkpoint saving to free memory from + # dequantized bf16 tensors that were temporarily created during fp8 + # model checkpoint saving. + gc.collect() - # Log E2E metrics after save-checkpoint - one_logger_utils.track_e2e_metrics() + timers.log([timer_key]) - one_logger_utils.on_save_checkpoint_end(save_checkpoint_duration, iteration, args.async_save) + # Log E2E metrics after save-checkpoint + one_logger_utils.track_e2e_metrics() - if args.log_progress and not non_persistent_ckpt: - compute_throughputs_and_append_to_progress_log( - iteration, num_floating_point_operations_so_far - ) + one_logger_utils.on_save_checkpoint_end(save_checkpoint_duration, iteration, args.async_save) - # Recover timing - energy_monitor.resume() - # Interval-time timer restarts with an explicit GLOBAL BARRIER -- the first - # hard sync after the checkpoint. If the skew didn't already serialize at - # timers.log above, it serializes HERE: fast ranks block for the slowest - # saver. Outside every save span, so it read as the dark post-checkpoint gap. - timers('interval-time', log_level=0).start(barrier=True) + if args.log_progress and not non_persistent_ckpt: + compute_throughputs_and_append_to_progress_log( + iteration, num_floating_point_operations_so_far + ) - # OTel: close the exposed-save goodput span that parents all of the above. - finally: - if _exposed_save_span is not None: - from opentelemetry import context as _octx - _octx.detach(_exposed_save_token) - _exposed_save_span.end() + # Recover timing + energy_monitor.resume() + timers('interval-time', log_level=0).start(barrier=True) -def _run_gpu_sniff_test(tag, span_name='megatron.train.sniff_test'): +def _run_gpu_sniff_test(tag): from megatron.core.process_groups_config import ProcessGroupCollection from megatron.training.gpu_sniff_test import run_gpu_sniff_test - # Two call sites with distinct span names (span_name): the once-at-start - # 'megatron.startup.sniff_test' run in train()'s preamble, and the periodic - # --gpu-sniff-test-interval 'megatron.train.sniff_test' runs in the step - # loop. The tag attribute additionally records which invocation this is. - with _otel_managed_span('job', span_name, is_goodput_span=True, **{'megatron.sniff_test.tag': tag}): - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=['ep', 'dp', 'tp'], - ) - print_datetime(f'running GPU sniff test ({tag})') - timers = get_timers() - timers('gpu-sniff-test', log_level=0).start(barrier=True) - run_gpu_sniff_test(tag, pg_collection=pg_collection) - timers('gpu-sniff-test').stop(barrier=True) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['ep', 'dp', 'tp']) + print_datetime(f'running GPU sniff test ({tag})') + timers = get_timers() + timers('gpu-sniff-test', log_level=0).start(barrier=True) + run_gpu_sniff_test(tag, pg_collection=pg_collection) + timers('gpu-sniff-test').stop(barrier=True) timers.log(['gpu-sniff-test']) print_datetime(f'finished GPU sniff test ({tag})') @@ -3977,7 +3773,7 @@ def post_training_step_callbacks( iteration, prof, num_floating_point_operations_since_last_log_event, - nsys_nvtx_context = None, + nsys_nvtx_context=None, ): """Run all post-training-step functions (e.g., FT heartbeats, GC).""" args = get_args() @@ -4015,8 +3811,7 @@ def post_training_step_callbacks( if ( args.profile and iteration == args.profile_step_end - and (len(args.profile_ranks) == 0 or - torch.distributed.get_rank() in args.profile_ranks) + and (len(args.profile_ranks) == 0 or torch.distributed.get_rank() in args.profile_ranks) ): # Disable NVTX range when profiling ends. if args.nvtx_ranges: @@ -4032,19 +3827,13 @@ def post_training_step_callbacks( nsys_nvtx_context.__exit__(None, None, None) # GPU sniff test. - if ( - args.gpu_sniff_test_interval is not None - and iteration % args.gpu_sniff_test_interval == 0 - ): + if args.gpu_sniff_test_interval is not None and iteration % args.gpu_sniff_test_interval == 0: _run_gpu_sniff_test(f'iteration {iteration:7d}') - # Manual garbage collection. With --manual-gc the interpreter's automatic - # collector is off; this synchronous full collection is the only GC. With - # manual_gc_interval=0 (the common case) it never fires. + # Manual garbage collection. if args.manual_gc: if args.manual_gc_interval != 0 and iteration % args.manual_gc_interval == 0: - with _otel_managed_span('step', 'megatron.train.gc_collect', is_goodput_span=True): - gc.collect() + gc.collect() # Return updated FLOPs accumulator so caller can persist the reset return num_floating_point_operations_since_last_log_event @@ -4120,10 +3909,6 @@ def checkpoint_and_decide_exit( done_cuda = torch.tensor( [train_time > args.exit_duration_in_mins], dtype=torch.int, device='cuda' ) - # MAX all-reduce so every rank makes the SAME exit decision despite each - # reading its own wall-clock (a determinism guarantee, not a sync point). - # ~0.1ms and left unspanned: the GPU is already drained by here, and the - # real post-checkpoint wait is the cross-rank save skew at timers.log. torch.distributed.all_reduce(done_cuda, op=torch.distributed.ReduceOp.MAX) done = done_cuda.item() if done: @@ -4142,12 +3927,8 @@ def checkpoint_and_decide_exit( return True # Exit based on iterations. - if ( - args.exit_interval - and iteration % args.exit_interval == 0 - ) or ( - args.phase_transition_iterations - and iteration in args.phase_transition_iterations + if (args.exit_interval and iteration % args.exit_interval == 0) or ( + args.phase_transition_iterations and iteration in args.phase_transition_iterations ): if args.save and not saved_checkpoint: save_checkpoint_and_time( @@ -4190,7 +3971,6 @@ def train( """ args = get_args() timers = get_timers() - fault_injector_kwargs = {} for f in dataclasses.fields(FaultInjectorConfig): if hasattr(args, f.name): @@ -4234,29 +4014,31 @@ def train( args.load = None args.finetune = True load_checkpoint( - model, - None, # Don't load optimizer state - None, # Don't load scheduler state - checkpointing_context=checkpointing_context, - skip_load_to_model_and_opt=HAVE_FSDP2 - and getattr(args, "use_torch_fsdp2", False) - and args.ckpt_format == "torch_dist", - ) - ref_state_dict = {k: (v.cpu() if v is not None else v) for k, v in model[0].state_dict().items()} + model, + None, # Don't load optimizer state + None, # Don't load scheduler state + checkpointing_context=checkpointing_context, + skip_load_to_model_and_opt=HAVE_FSDP2 + and getattr(args, "use_torch_fsdp2", False) + and args.ckpt_format == "torch_dist", + ) + ref_state_dict = { + k: (v.cpu() if v is not None else v) for k, v in model[0].state_dict().items() + } # Reload RL training checkpoint weights args.load = load args.finetune = finetune print_rank_0("> Reloading RL training checkpoint...") load_checkpoint( - model, - None, - None, - checkpointing_context=checkpointing_context, - skip_load_to_model_and_opt=HAVE_FSDP2 - and getattr(args, "use_torch_fsdp2", False) - and args.ckpt_format == "torch_dist", - ) + model, + None, + None, + checkpointing_context=checkpointing_context, + skip_load_to_model_and_opt=HAVE_FSDP2 + and getattr(args, "use_torch_fsdp2", False) + and args.ckpt_format == "torch_dist", + ) args.no_load_optim = no_load_optim @@ -4268,15 +4050,12 @@ def train( ) def _dp_world_size(): - # Full DP x gtp_remat degree (num_microbatches spans the full data-distribution axis). - gtp_remat = args.gtp_weight_remat_size if lang_pgc is not None: - return lang_pgc.dp.size() * gtp_remat + return lang_pgc.dp.size() if mpu.model_parallel_is_initialized(): return mpu.get_data_parallel_world_size() - # args.data_parallel_size is the language (llm) dp on all ranks (set in validate_args) and - # excludes gtp_remat, so scale by gtp_remat to span the full data-distribution axis. - return args.data_parallel_size * gtp_remat + # args.data_parallel_size equals the language (llm) dp on all ranks (entry validate_args). + return args.data_parallel_size # IMPORTANT FIX: For RL training, reinitialize the microbatch calculator with the correct configuration if args.perform_rl_step: @@ -4303,9 +4082,6 @@ def _dp_world_size(): energy_monitor = get_energy_monitor() one_logger = get_one_logger() - if args.hybrid_context_parallel: - train_data_iterator = iter(HybridCPDataLoaderWrapper(train_data_iterator, config)) - if args.run_workload_inspector_server: try: import threading @@ -4335,8 +4111,10 @@ def _dp_world_size(): # Make sure rerun_state_machine has the right iteration loaded from checkpoint. rerun_state_machine = get_rerun_state_machine() if rerun_state_machine.current_iteration != iteration: - print_rank_0(f"Overwriting rerun_state_machine.current_iteration from " - f"{rerun_state_machine.current_iteration} to {iteration}...") + print_rank_0( + f"Overwriting rerun_state_machine.current_iteration from " + f"{rerun_state_machine.current_iteration} to {iteration}..." + ) rerun_state_machine.current_iteration = iteration # Track E2E metrics at the start of training. @@ -4357,9 +4135,7 @@ def _dp_world_size(): # Setup some training config params. config.grad_scale_func = optimizer.scale_loss if optimizer is not None else None config.timers = timers - if isinstance( - model[0], (FullyShardedDataParallelV1, FullyShardedDataParallelV2, DDP) - ) and args.overlap_grad_reduce: + if isinstance(model[0], (megatron_FSDP, DDP)) and args.overlap_grad_reduce: assert config.no_sync_func is None, ( 'When overlap_grad_reduce is True, config.no_sync_func must be None; ' 'a custom no_sync_func is not supported when overlapping grad-reduce' @@ -4388,7 +4164,7 @@ def _dp_world_size(): # GPU sniff test at start of training. if args.gpu_sniff_test_interval is not None: - _run_gpu_sniff_test('before training', span_name='megatron.startup.sniff_test') + _run_gpu_sniff_test('before training') # Initialize router trace if requested. The tracer attaches forward hooks # to all TopKRouter modules and writes one JSONL record per (iteration, @@ -4441,9 +4217,7 @@ def _dp_world_size(): eval_duration = 0.0 eval_iterations = 0 # Wrap forward_backward_func for Full iteration CUDA graph - forward_backward_func = get_forward_backward_func( - schedule_pg_collection=pg_collection - ) + forward_backward_func = get_forward_backward_func(schedule_pg_collection=pg_collection) if args.cuda_graph_impl == "full_iteration": forward_backward_func = FullCudaGraphWrapper( forward_backward_func, @@ -4454,11 +4228,7 @@ def _dp_world_size(): if args.moe_expert_rank_capacity_factor is not None: copy_main_params = args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather forward_backward_func = PagedStashRunner( - config, - copy_main_params, - model, - optimizer, - forward_backward_func, + config, copy_main_params, model, optimizer, forward_backward_func ) if args.optimizer_cuda_graph: optimizer.step = OptimizerCudaGraphWrapper( @@ -4480,7 +4250,7 @@ def get_e2e_base_metrics(): 'total_flops_since_current_train_start': num_floating_point_operations_since_current_train_start, 'num_floating_point_operations_so_far': num_floating_point_operations_so_far, 'consumed_train_samples': args.consumed_train_samples, - 'world_size': getattr(args, 'mimo_llm_world_size', args.world_size), + 'world_size': args.world_size, 'seq_length': args.seq_length, } @@ -4490,23 +4260,26 @@ def get_e2e_base_metrics(): one_logger.store_set('get_e2e_base_metrics', get_e2e_base_metrics) prof = None - nsys_nvtx_context = None # reference to context for nsys profiling, so it can be cleaned up + nsys_nvtx_context = None # reference to context for nsys profiling, so it can be cleaned up if ( args.profile - and (len(args.profile_ranks) == 0 or - torch.distributed.get_rank() in args.profile_ranks) + and (len(args.profile_ranks) == 0 or torch.distributed.get_rank() in args.profile_ranks) and args.use_pytorch_profiler ): if args.pytorch_profiler_collect_chakra: et_dir = Path(f"{args.tensorboard_dir}/../chakra") et_dir.mkdir(parents=True, exist_ok=True) - et = torch.profiler.ExecutionTraceObserver().register_callback(f"{et_dir}/rank-{torch.distributed.get_rank()}.json.gz") + et = torch.profiler.ExecutionTraceObserver().register_callback( + f"{et_dir}/rank-{torch.distributed.get_rank()}.json.gz" + ) else: et = None + def trace_handler(p): profile_dir = Path(f"{args.tensorboard_dir}/../torch_profile") profile_dir.mkdir(parents=True, exist_ok=True) p.export_chrome_trace(f"{profile_dir}/rank-{torch.distributed.get_rank()}.json.gz") + prof = torch.profiler.profile( schedule=torch.profiler.schedule( wait=max(args.profile_step_start - 1, 0), @@ -4532,16 +4305,12 @@ def trace_handler(p): param_sync_func = config.param_sync_func config.param_sync_func = None pre_hook_enabled = False - # Also, check weight hash across DP replicas to be very pedantic. This hashes - # every parameter and all-reduces across DP replicas + a barrier -- it's - # typically the dominant cost of the train() preamble, so it gets its own - # span (parented to megatron.startup, still open here). + # Also, check weight hash across DP replicas to be very pedantic. if args.check_weight_hash_across_dp_replicas_interval is not None: - with _otel_managed_span('job', 'megatron.startup.weight_hash_check', is_goodput_span=True): - assert check_param_hashes_across_dp_replicas( - model, cross_check=True - ), "Parameter hashes not matching across DP replicas" - torch.distributed.barrier() + assert check_param_hashes_across_dp_replicas( + model, cross_check=True + ), "Parameter hashes not matching across DP replicas" + torch.distributed.barrier() print_rank_0(f">>> Weight hashes match after {iteration} iterations...") # Initialize CUDA Graphs helper. @@ -4552,28 +4321,15 @@ def trace_handler(p): seq_length=args.seq_length, micro_batch_size=args.micro_batch_size, optimizers=[optimizer], + thd_sequence_length_upper_bound=_get_thd_sequence_length_upper_bound(args), ) - # OTel: everything above in train() (the preamble: weight-hash check, sniff - # test, cuda-graph helper init) was init and parented to megatron.startup, which - # stayed open through it. Close startup now; the steady-state loop below is - # emitted as one megatron.train span per interval (see - # _reroot_otel_interval), not one umbrella. _start_otel_train_span just arms the - # per-interval counter; the first loop pass opens the first block. - _end_otel_startup_span() - _start_otel_train_span() - # Run training iterations till done. buffered_rollouts = None while iteration < args.train_iters: - # At each checkpoint-interval boundary, re-root into a new trace so this - # pass's iteration + (this interval's) checkpoint/eval/sniff form one compact - # trace instead of accreting into a run-long one. Must be the first thing in - # the pass so everything below nests under the current interval root. - _maybe_reroot_otel_interval() - if (args.profile - and (len(args.profile_ranks) == 0 or - torch.distributed.get_rank() in args.profile_ranks)): + if args.profile and ( + len(args.profile_ranks) == 0 or torch.distributed.get_rank() in args.profile_ranks + ): # Enable NVTX range when profiling starts and nvtx_ranges is set. if iteration == args.profile_step_start and args.nvtx_ranges: configure_nvtx_profiling(True) @@ -4581,27 +4337,21 @@ def trace_handler(p): prof.step() elif iteration == args.profile_step_start: torch.cuda.check_error(torch.cuda.cudart().cudaProfilerStart()) - if args.record_shapes: - nsys_nvtx_context = torch.autograd.profiler.emit_nvtx(record_shapes=args.record_shapes) - nsys_nvtx_context.__enter__() - - # Fault-tolerance heartbeat at the top of the loop -- uninstrumented - # main-thread work that sits in the post-checkpoint gap alongside the - # exit-duration barrier. - with _otel_managed_span('checkpoint', 'megatron.checkpoint.ft_heartbeat', is_goodput_span=True): - ft_integration.on_checkpointing_start() - # Non-blocking finalize of the *previous* async checkpoint, on the - # training critical path -- this is the "exposed" cost the async save - # imposes back on the loop (the flip side of the background write span - # on the worker), and shows up as a per-iteration gap near checkpoint - # boundaries. Cheap most iterations, blocks when finalizing. - with _otel_managed_span('checkpoint', 'megatron.checkpoint.save.finalize', is_goodput_span=True): - maybe_finalize_async_save(blocking=False) + nsys_nvtx_context = torch.autograd.profiler.emit_nvtx( + record_shapes=args.record_shapes + ) + nsys_nvtx_context.__enter__() + + ft_integration.on_checkpointing_start() + maybe_finalize_async_save(blocking=False) ft_integration.on_checkpointing_end(is_async_finalization=True) # Update the timeout for all process groups after initialization # We update the timeout after the first successful iteration, # which takes longer than others usually - if args.distributed_timeout_seconds_after_init is not None and iteration == start_iteration+1: + if ( + args.distributed_timeout_seconds_after_init is not None + and iteration == start_iteration + 1 + ): # TODO: some dynamic timeout setting is required # based on the iteration time considering interval-based steps (e.g. eval, checkpoint) # e.g. timeout for normal iterations vs timeout for iterations with checkpoint @@ -4615,9 +4365,9 @@ def trace_handler(p): # Standard microbatch update (sequence packing overrides this in rl_utils.py) update_num_microbatches(args.consumed_train_samples, consistency_check=False, verbose=True) # Skip automatic checkpoint on microbatch changes when sequence packing is active - # as it intentionally reconfigures microbatches + # as it intentionally reconfigures microbatches. if get_num_microbatches() != num_microbatches and iteration != 0: - if args.rl_use_sequence_packing: + if args.rl_use_sequence_packing or args.sequence_packing_scheduler is not None: print_rank_0( f"[Sequence Packing] Skipping automatic checkpoint at iteration {iteration} " f"(microbatch change: {num_microbatches} -> {get_num_microbatches()})" @@ -4640,23 +4390,18 @@ def trace_handler(p): num_microbatches = get_num_microbatches() update_num_microbatches(args.consumed_train_samples, consistency_check=True, verbose=True) - # Capture CUDA Graphs. One-off, at the warmup-step boundary -- the actual - # graph capture (create_cudagraphs) is a notable one-time cost worth its - # own span, distinct from the megatron.train.iteration spans around it. + # Capture CUDA Graphs. if ( args.cuda_graph_impl == "transformer_engine" and not cuda_graph_helper.capture_finished() and iteration - start_iteration == args.cuda_graph_warmup_steps ): - with _otel_managed_span( - 'job', 'megatron.train.cuda_graph_capture', is_goodput_span=True, **{'megatron.iteration': iteration} - ): - if args.cuda_graph_warmup_steps > 0 and should_disable_forward_pre_hook(args): - disable_forward_pre_hook(model, param_sync=False) - cuda_graph_helper.create_cudagraphs() - if args.cuda_graph_warmup_steps > 0 and should_disable_forward_pre_hook(args): - enable_forward_pre_hook(model) - cuda_graph_helper.cuda_graph_set_manual_hooks() + if args.cuda_graph_warmup_steps > 0 and should_disable_forward_pre_hook(args): + disable_forward_pre_hook(model, param_sync=False) + cuda_graph_helper.create_cudagraphs() + if args.cuda_graph_warmup_steps > 0 and should_disable_forward_pre_hook(args): + enable_forward_pre_hook(model) + cuda_graph_helper.cuda_graph_set_manual_hooks() # Completely skip iteration if needed. if (iteration + 1) in args.iterations_to_skip: @@ -4665,9 +4410,7 @@ def trace_handler(p): if iteration == start_iteration: start_iteration = iteration + 1 iteration += 1 - batch_size = ( - _dp_world_size() * args.micro_batch_size * get_num_microbatches() - ) + batch_size = _dp_world_size() * args.micro_batch_size * get_num_microbatches() args.consumed_train_samples += batch_size args.skipped_train_samples += batch_size continue @@ -4682,7 +4425,11 @@ def trace_handler(p): torch.cuda.empty_cache() with torch.no_grad(): train_data_iterator = rl_utils.get_grpo_data_iterator( - model, inference_model, optimizer, iteration, ref_state_dict, + model, + inference_model, + optimizer, + iteration, + ref_state_dict, grpo_iterations=args.grpo_iterations, grpo_prompts_per_step=args.grpo_prompts_per_step, grpo_group_size=args.grpo_group_size, @@ -4706,52 +4453,46 @@ def trace_handler(p): grad_norm = 0.0 num_zeros_in_grad = 0 max_attention_logit = None - _step_span = None + num_microbatches = get_num_microbatches() + seqlen_sum_this_global_batch = None + seqlen_squared_sum_this_global_batch = None else: - # OTel: dedicated span for the first iteration actually executed in this - # process (post checkpoint-resume, post iteration-skip) — not iteration 1, - # just the first one that runs. Kept separate from the per-step span below - # since it captures one-off warmup costs (compilation, CUDA graph capture, - # prefetch) that steady-state iterations don't pay. - _first_iter_span_cm = ( - _otel_managed_span( - 'first_iteration', 'megatron.train.first_iteration', is_goodput_span=True, - **{'megatron.iteration': iteration}, - ) - if is_first_iteration - else nullcontext() + ft_integration.on_training_step_start() + ( + loss_dict, + skipped_iter, + should_checkpoint, + should_exit, + exit_code, + grad_norm, + num_zeros_in_grad, + max_attention_logit, + num_microbatches, + seqlen_sum_this_global_batch, + seqlen_squared_sum_this_global_batch, + ) = train_step( + forward_step_func, + train_data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, + pg_collection=pg_collection, + p2p_communicator=p2p_communicator, ) - # OTel: optional per-step span wrapping the real train_step. - with _first_iter_span_cm, _otel_managed_span('step', 'megatron.train.iteration', is_goodput_span=True, **{'megatron.iteration': iteration}) as _step_span: - ft_integration.on_training_step_start() - ( - loss_dict, - skipped_iter, - should_checkpoint, - should_exit, - exit_code, - grad_norm, - num_zeros_in_grad, - max_attention_logit, - ) = train_step( - forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=iteration, - pg_collection=pg_collection, - p2p_communicator=p2p_communicator, - ) - ft_integration.on_training_step_end() - if _maybe_raise_workload_exception is not None and iteration != start_iteration: - _maybe_raise_workload_exception() - # Fault delay timing can start at the end of iteration N. Self-firing faults - # (signals, GIL, GPU) may then manifest in iteration N or N+1 depending on the - # configured delay; workload-exception faults manifest on a later poll. - if _maybe_raise_workload_exception is not None and should_setup_fault_injection_at_iteration( - fault_injector_config, iteration - ): - setup_fault_injection(fault_injector_config) - if _step_span is not None: - _otel_safe_set_attrs( - _step_span, {'megatron.skipped': bool(skipped_iter)} - ) + ft_integration.on_training_step_end() + if _maybe_raise_workload_exception is not None and iteration != start_iteration: + _maybe_raise_workload_exception() + # Fault delay timing can start at the end of iteration N. Self-firing faults + # (signals, GIL, GPU) may then manifest in iteration N or N+1 depending on the + # configured delay; workload-exception faults manifest on a later poll. + if ( + _maybe_raise_workload_exception is not None + and should_setup_fault_injection_at_iteration(fault_injector_config, iteration) + ): + setup_fault_injection(fault_injector_config) if should_checkpoint: save_checkpoint_and_time( iteration, @@ -4777,13 +4518,9 @@ def trace_handler(p): else: # Enable forward pre-hook after training step has successfully run. All subsequent # forward passes will use the forward pre-hook / `param_sync_func` in - # `forward_backward_func`. One-time (iteration == start_iteration) - # param-gather pre-hook re-enable -- runs between the train_step - # and train_log spans on the first iteration only, which is - # exactly the ~1.8s post-first-iteration gap. + # `forward_backward_func`. if should_disable_forward_pre_hook(args): - with _otel_managed_span('first_iteration', 'megatron.train.forward_pre_hook', is_goodput_span=True): - enable_forward_pre_hook(model) + enable_forward_pre_hook(model) config.param_sync_func = param_sync_func pre_hook_enabled = True # Set the manual hooks here since it's not set right after the capturing. @@ -4803,12 +4540,10 @@ def trace_handler(p): getattr(args, "fsdp_manual_registration", False) and getattr(args, "nccl_ub", False) and getattr(args, "use_megatron_fsdp", False) - and iteration == start_iteration + 1 + and iteration == start_iteration + 1 ): for model_chunk in model: - if isinstance( - model_chunk, (FullyShardedDataParallelV1, FullyShardedDataParallelV2) - ) and getattr( + if isinstance(model_chunk, megatron_FSDP) and getattr( model_chunk.ddp_config, "fsdp_manual_registration", False ): param_and_grad_buffer = getattr(model_chunk, "param_and_grad_buffer", None) @@ -4818,14 +4553,10 @@ def trace_handler(p): if args.perform_rl_step and args.rl_use_sequence_packing: iteration_sequences = rl_utils.get_iteration_sequence_count(args) # Track bins separately for packed mode - bin_count = ( - _dp_world_size() * args.micro_batch_size * get_num_microbatches() - ) + bin_count = _dp_world_size() * args.micro_batch_size * get_num_microbatches() args.consumed_train_bins += bin_count else: - batch_size = ( - _dp_world_size() * args.micro_batch_size * get_num_microbatches() - ) + batch_size = _dp_world_size() * args.micro_batch_size * get_num_microbatches() iteration_sequences = batch_size # Update consumed samples (always means sequences now) @@ -4842,14 +4573,22 @@ def trace_handler(p): else: assert num_skipped_samples_in_batch == 0 args.skipped_train_samples += num_skipped_samples_in_batch - # Drain the per-iteration packed-sequence stats so the FLOPs computation - # reflects THD per-chunk causal attention AND excludes padding tokens - # from token-linear work. Returns ``(None, None)`` for unpacked BSHD - # runs (no collective issued), letting ``num_floating_point_operations`` - # fall back to its closed-form defaults. - total_real_tokens_in_batch, seqlen_squared_sum_in_batch = ( - consume_seqlen_stats_in_iteration() - ) + if getattr(config, 'sequence_packing_scheduler', None) is not None and not args.skip_train: + # The scheduler computed these from the real sequence lengths before + # CP padding and rerouting, so use them directly for FLOPs accounting. + assert seqlen_sum_this_global_batch is not None + assert seqlen_squared_sum_this_global_batch is not None + total_real_tokens_in_batch = seqlen_sum_this_global_batch + seqlen_squared_sum_in_batch = seqlen_squared_sum_this_global_batch + else: + # Drain the per-iteration packed-sequence stats so the FLOPs computation + # reflects THD per-chunk causal attention AND excludes padding tokens + # from token-linear work. Returns ``(None, None)`` for unpacked BSHD + # runs (no collective issued), letting ``num_floating_point_operations`` + # fall back to its closed-form defaults. + total_real_tokens_in_batch, seqlen_squared_sum_in_batch = ( + consume_seqlen_stats_in_iteration() + ) num_floating_point_operations_in_batch = num_floating_point_operations( args, batch_size, @@ -4859,75 +4598,46 @@ def trace_handler(p): num_floating_point_operations_so_far += num_floating_point_operations_in_batch num_floating_point_operations_since_last_log_event += num_floating_point_operations_in_batch - # OTel: super-span over the whole post-step REPORTING block (loss-scale - # sync, param-norm reduction, throughput/tensorboard/wandb logging). One - # goodput span that captures ALL the exposed reporting overhead -- - # including the loss_scale.item() device sync and misc bookkeeping that - # would otherwise be dark -- WITHOUT instrumenting each line. params_norm - # and train.log nest under it; its own uninstrumented time is the blank - # you'd otherwise see. start_span + attach (no re-indent); closed after - # training_log below. - _report_span = None - _report_token = None - if _otel_sg_enabled('step'): - from opentelemetry import context as _octx, trace as _otr - _report_span = get_telemetry().tracer.start_span('megatron.train.iteration_report') - _otel_mark_goodput(_report_span) - _report_token = _octx.attach(_otr.set_span_in_context(_report_span)) - try: + # Logging. + if optimizer is not None and not optimizer.is_stub_optimizer: + loss_scale = optimizer.get_loss_scale().item() + else: + loss_scale = 1.0 + params_norm = None - # Logging. - if optimizer is not None and not optimizer.is_stub_optimizer: - # First .item() after the train_step: a device sync draining the - # iteration's pending GPU queue (captured under iteration_report). - loss_scale = optimizer.get_loss_scale().item() - else: - loss_scale = 1.0 - params_norm = None - - if args.log_params_norm: - # Cross-rank param L2 norm (--log-params-norm): a full-model reduction - # + all-reduce that BLOCKS the training loop -- exposed goodput cost - # (~1.5s cold on the first iteration, ~10ms steady). Kept as a real - # cost span (it stalls the critical path), unlike passive monitors. - with _otel_managed_span('step', 'megatron.train.params_norm', is_goodput_span=True): - params_norm = calc_params_l2_norm(model) - if optimizer is not None: - learning_rate = get_canonical_lr_for_logging(optimizer.param_groups) - else: - learning_rate = None - # Per-iteration logging (throughput calc, tensorboard/wandb writes) -- - # uninstrumented per-iteration overhead outside the train_step span. - with _otel_managed_span('step', 'megatron.train.log', is_goodput_span=True): - report_memory_flag = training_log( - loss_dict, - total_loss_dict, - learning_rate, - iteration, - loss_scale, - report_memory_flag, - skipped_iter, - grad_norm, - params_norm, - num_zeros_in_grad, - max_attention_logit, - pg_collection=model_pg_collection, - is_first_iteration=is_first_iteration, - seqlen_squared_sum_in_batch=seqlen_squared_sum_in_batch, - total_real_tokens_in_batch=total_real_tokens_in_batch, - ) - # OTel: close the iteration-report super-span (parents params_norm + log; - # its own uninstrumented time is the loss_scale sync + FLOPs bookkeeping). - finally: - if _report_span is not None: - from opentelemetry import context as _octx - _octx.detach(_report_token) - _report_span.end() + if args.log_params_norm: + params_norm = calc_params_l2_norm(model) + if optimizer is not None: + learning_rate = get_canonical_lr_for_logging(optimizer.param_groups) + else: + learning_rate = None + report_memory_flag = training_log( + loss_dict, + total_loss_dict, + learning_rate, + iteration, + loss_scale, + report_memory_flag, + skipped_iter, + grad_norm, + params_norm, + num_zeros_in_grad, + max_attention_logit, + pg_collection=model_pg_collection, + is_first_iteration=is_first_iteration, + seqlen_squared_sum_in_batch=seqlen_squared_sum_in_batch, + total_real_tokens_in_batch=total_real_tokens_in_batch, + num_microbatches=num_microbatches, + ) is_first_iteration = False # Evaluation. - if args.eval_interval and iteration % args.eval_interval == 0 and args.do_valid \ - and (args.start_eval_at_iter is None or iteration >= args.start_eval_at_iter): + if ( + args.eval_interval + and iteration % args.eval_interval == 0 + and args.do_valid + and (args.start_eval_at_iter is None or iteration >= args.start_eval_at_iter) + ): if args.log_energy: energy_monitor.pause() timers('interval-time').stop() @@ -4961,16 +4671,25 @@ def trace_handler(p): training_model=rl_training_model, ) else: - evaluate_and_print_results(prefix, forward_step_func, - valid_data_iterator, model, - iteration, process_non_loss_data_func, - config, verbose=False, write_to_tensorboard=True, - non_loss_data_func=non_loss_data_func, - pg_collection=pg_collection, - p2p_communicator=p2p_communicator) + evaluate_and_print_results( + prefix, + forward_step_func, + valid_data_iterator, + model, + iteration, + process_non_loss_data_func, + config, + verbose=False, + write_to_tensorboard=True, + non_loss_data_func=non_loss_data_func, + pg_collection=pg_collection, + p2p_communicator=p2p_communicator, + ) eval_duration += timers('eval-time').elapsed() - eval_iterations += sum(args.eval_iters) if isinstance(args.eval_iters, list) else args.eval_iters + eval_iterations += ( + sum(args.eval_iters) if isinstance(args.eval_iters, list) else args.eval_iters + ) timers('eval-time').stop() one_logger_utils.track_e2e_metrics() @@ -5012,13 +4731,6 @@ def trace_handler(p): if should_exit: break - # Early-exit paths (exit-duration / exit-interval / signal handler) sys.exit() - # below before the normal-path logging, so record the train-loop finish time here. - if should_exit: - one_logger and one_logger.log_metrics( - {'app_train_loop_finish_time': one_logger_utils.get_timestamp_in_ms()} - ) - # Destroy CUDA Graphs. if args.cuda_graph_impl == "transformer_engine" and cuda_graph_helper.graphs_created(): cuda_graph_helper.delete_cuda_graphs() @@ -5039,11 +4751,9 @@ def trace_handler(p): disable_forward_pre_hook(model, optimizer=optimizer) ft_integration.on_checkpointing_start() - # Finalize all unfinished async requests and terminate the persistent - # async worker (if enabled) if the code is meant to exit and not return from this - # function. - with _otel_managed_span('checkpoint', 'megatron.checkpoint.exit_finalize', is_goodput_span=True): - maybe_finalize_async_save(blocking=True, terminate=should_exit) + # This will finalize all unfinalized async request and terminate + # a persistent async worker if persistent ckpt worker is enabled + maybe_finalize_async_save(blocking=True, terminate=True) ft_integration.on_checkpointing_end(is_async_finalization=True) if args.log_energy: @@ -5063,21 +4773,13 @@ def trace_handler(p): # ncclCommDeregister on handles created by ncclCommWindowRegister, # causing "NCCL WARN Deregister: Could not find handle" and a crash. torch.distributed.barrier() - if getattr(args, 'gtp_remat_nccl_ub', False) or getattr(args, 'gtp_expert_remat_nccl_ub', False): - from megatron.core.tensor_parallel.gtp_api import deregister_and_clear_gtp_symm_pools - - # Deregister the GTP symmetric-memory pools: windows left registered when the - # process groups are destroyed make NCCL abort. - deregister_and_clear_gtp_symm_pools() - for model_module in model: if isinstance(model_module, DDP): for buf in model_module.buffers + model_module.expert_parallel_buffers: if getattr(buf, 'nccl_mem_pool', None) is not None: - nccl_allocator.deregister_mem_pool(buf.nccl_mem_pool, buf.data_parallel_group) - one_logger and one_logger.log_metrics( - {'app_finish_time': one_logger_utils.get_timestamp_in_ms()} - ) + nccl_allocator.deregister_mem_pool( + buf.nccl_mem_pool, buf.data_parallel_group + ) wandb_writer = get_wandb_writer() if wandb_writer: wandb_writer.finish() @@ -5085,26 +4787,11 @@ def trace_handler(p): one_logger_utils.finish() if args.perform_rl_step: rl_utils.rl_inference_interface_shutdown() - # This sys.exit() bypasses the rest of pretrain() -- including its normal - # _end_otel_job_spans() call -- so the open megatron.train interval block (and - # megatron.startup, if it somehow never closed) would otherwise only get ended - # later via the atexit fallback, which runs AFTER a bare shutdown() here already - # tore down the provider (nothing left to export into). _end_otel_job_spans() - # ends the spans first, then shuts down, so it must be this call, not a - # standalone get_telemetry().shutdown(). - _end_otel_job_spans() sys.exit(exit_code) - # OTel: normal loop completion -- end the final megatron.train interval block here - # so it doesn't bleed into pretrain()'s post-train eval/test. Those now root their - # own traces (there is no umbrella left to parent them). The sys.exit() path above - # ends the block via _end_otel_job_spans() instead. - _end_otel_train_span() - return iteration, num_floating_point_operations_so_far -@_otel_trace_fn('evaluate', 'megatron.evaluate') def evaluate( forward_step_func, data_iterator, @@ -5137,22 +4824,12 @@ def evaluate( # make validation batch size independent from training batch size eval_batch_size = args.eval_global_batch_size eval_micro_batch_size = args.eval_micro_batch_size - # data_parallel_size excludes the GTP-remat axis (it's folded into total_model_size at - # arguments.py:446); each gtp-remat peer consumes a distinct microbatch, so include it in the - # global sample breadth we divide out to recover the microbatch count. - eval_num_microbatches = eval_batch_size // ( - eval_micro_batch_size * args.data_parallel_size * args.gtp_weight_remat_size - ) + eval_num_microbatches = eval_batch_size // (eval_micro_batch_size * args.data_parallel_size) forward_backward_func = get_forward_backward_func(schedule_pg_collection=pg_collection) # Reductions source per-rank groups from the model (encoder rank -> encoder groups). eval_pgc = get_attr_wrapped_model(model[0], "pg_collection") if eval_pgc is None: eval_pgc = ProcessGroupCollection.use_mpu_process_groups() - # gtp_remat-inclusive, mirroring train_step: gtp_remat peers hold distinct micro-batches, so - # the reduction must cover every distinct-data rank. Reducing over the replicate dp_cp would - # average a 1/gtp_remat subsample of the eval batch (768 sequences read, 12 averaged at - # dp=6/gtp_remat=64) while consumed_valid_samples still books the full eval_batch_size. - eval_dp_cp_group = getattr(eval_pgc, 'dp_cp_gtp_remat', None) or eval_pgc.dp_cp if args.cuda_graph_impl == "full_iteration": forward_backward_func = FullCudaGraphWrapper( forward_backward_func, @@ -5163,11 +4840,7 @@ def evaluate( if args.moe_expert_rank_capacity_factor is not None: copy_main_params = args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather forward_backward_func = PagedStashRunner( - config, - copy_main_params, - model, - None, - forward_backward_func, + config, copy_main_params, model, None, forward_backward_func ) if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False): @@ -5196,21 +4869,32 @@ def evaluate( # Don't care about timing during evaluation config.timers = None ft_integration.on_eval_step_start() - with _otel_managed_span('evaluate', 'megatron.evaluate.step', - **{'megatron.eval_iteration': iteration}): - loss_dicts = forward_backward_func( - forward_step_func=forward_step_func, - data_iterator=data_iterator, - model=model, - num_microbatches=eval_num_microbatches, - seq_length=args.seq_length, - micro_batch_size=eval_micro_batch_size, - decoder_seq_length=args.decoder_seq_length, - forward_only=True, - adjust_tensor_shapes_fn=adjust_tensor_shapes_fn, - pg_collection=pg_collection, - p2p_communicator=p2p_communicator, - ) + if getattr(config, 'sequence_packing_scheduler', None) is not None: + try: + (packed_data_iterator, scheduled_eval_num_microbatches, _, _) = ( + wrap_data_iterator(data_iterator, config, eval_num_microbatches) + ) + except StopIteration: + # Validation data iterator exhausted, stop evaluation early. + ft_integration.on_eval_step_end() + config.timers = get_timers() + break + else: + packed_data_iterator = data_iterator + scheduled_eval_num_microbatches = eval_num_microbatches + loss_dicts = forward_backward_func( + forward_step_func=forward_step_func, + data_iterator=packed_data_iterator, + model=model, + num_microbatches=scheduled_eval_num_microbatches, + seq_length=args.seq_length, + micro_batch_size=eval_micro_batch_size, + decoder_seq_length=args.decoder_seq_length, + forward_only=True, + adjust_tensor_shapes_fn=adjust_tensor_shapes_fn, + pg_collection=pg_collection, + p2p_communicator=p2p_communicator, + ) ft_integration.on_eval_step_end() config.timers = get_timers() @@ -5222,7 +4906,9 @@ def evaluate( # Reduce across processes. for key in loss_dicts[0].keys(): if key not in total_loss_dict: - total_loss_dict[key] = torch.tensor([0.0, 0.0], dtype=torch.float, device='cuda') + total_loss_dict[key] = torch.tensor( + [0.0, 0.0], dtype=torch.float, device='cuda' + ) val = [x[key].view(-1) for x in loss_dicts] if val[0].numel() == 2: @@ -5231,13 +4917,13 @@ def evaluate( val = torch.vstack(val) val = val[:, 0] / val[:, 1].clamp(min=1) val = val.mean() - torch.distributed.all_reduce(val, group=eval_dp_cp_group) - val /= torch.distributed.get_world_size(group=eval_dp_cp_group) + torch.distributed.all_reduce(val, group=eval_pgc.dp_cp) + val /= torch.distributed.get_world_size(group=eval_pgc.dp_cp) total_loss_dict[key][0] += val total_loss_dict[key][1] += 1 - else : + else: val = torch.vstack(val).sum(dim=0) - torch.distributed.all_reduce(val, group=eval_dp_cp_group) + torch.distributed.all_reduce(val, group=eval_pgc.dp_cp) total_loss_dict[key] += val elif val[0].numel() == 1: val = torch.cat(val).sum() @@ -5291,18 +4977,6 @@ def evaluate( rerun_state_machine.set_mode(rerun_mode) - - # OTel: set eval_iters on the active span started by the @_otel_trace_fn decorator. - # get_current_span() always returns a NonRecordingSpan (no-op) when no span is active, - # so this is safe whether or not the group is enabled or the decorator fell back to no-op. - try: - from opentelemetry import trace as _ot - _ot.get_current_span().set_attribute( - 'megatron.eval_iters', eval_iters if eval_iters is not None else 0 - ) - except Exception: # noqa: BLE001 — OTel must never crash evaluation - pass - return total_loss_dict, collected_non_loss_data, False @@ -5353,11 +5027,13 @@ def evaluate_and_print_results( eval_iters = args.eval_iters if args.validation_set_names: - assert args.multiple_validation_sets, \ - "--validation-set-names requires --multiple-validation-sets" - assert len(args.validation_set_names) == len(data_iterators), \ - f"Number of --validation-set-names ({len(args.validation_set_names)}) must match " \ + assert ( + args.multiple_validation_sets + ), "--validation-set-names requires --multiple-validation-sets" + assert len(args.validation_set_names) == len(data_iterators), ( + f"Number of --validation-set-names ({len(args.validation_set_names)}) must match " f"the number of validation datasets ({len(data_iterators)})" + ) for index, (iterator, iterations) in enumerate(zip(data_iterators, eval_iters)): suffix = "" @@ -5387,7 +5063,9 @@ def evaluate_and_print_results( ppl = math.exp(min(20, total_loss_dict[key].item())) string += '{} PPL: {:.6E} | '.format(key, ppl) if writer: - writer.add_scalar('{} validation{}'.format(key, suffix), total_loss_dict[key].item(), iteration) + writer.add_scalar( + '{} validation{}'.format(key, suffix), total_loss_dict[key].item(), iteration + ) writer.add_scalar( '{} validation{} vs samples'.format(key, suffix), total_loss_dict[key].item(), @@ -5396,11 +5074,14 @@ def evaluate_and_print_results( if args.log_validation_ppl_to_tensorboard: writer.add_scalar('{} validation{} ppl'.format(key, suffix), ppl, iteration) writer.add_scalar( - '{} validation{} ppl vs samples'.format(key, suffix), ppl, args.consumed_train_samples + '{} validation{} ppl vs samples'.format(key, suffix), + ppl, + args.consumed_train_samples, ) if wandb_writer and is_last_rank(): wandb_writer.log( - {'{} validation{}'.format(key, suffix): total_loss_dict[key].item()}, iteration + {'{} validation{}'.format(key, suffix): total_loss_dict[key].item()}, + iteration, ) if process_non_loss_data_func is not None and writer and is_last_rank(): @@ -5455,7 +5136,11 @@ def get_train_valid_test_num_samples(): # Get train_samples in current phase. if args.phase_transition_iterations: - phase_transition_samples = [0] + [t * args.global_batch_size for t in args.phase_transition_iterations] + [args.train_samples] + phase_transition_samples = ( + [0] + + [t * args.global_batch_size for t in args.phase_transition_iterations] + + [args.train_samples] + ) current_sample = args.iteration * args.global_batch_size last_transition_sample = max(s for s in phase_transition_samples if s <= current_sample) next_transition_sample = min(s for s in phase_transition_samples if s > current_sample) @@ -5466,7 +5151,9 @@ def get_train_valid_test_num_samples(): return (train_samples_in_current_phase, eval_samples, test_samples) -def build_train_valid_test_datasets(build_train_valid_test_datasets_provider, train_valid_test_num_samples=None): +def build_train_valid_test_datasets( + build_train_valid_test_datasets_provider, train_valid_test_num_samples=None +): """Build pretraining datasets.""" if train_valid_test_num_samples is None: train_valid_test_num_samples = get_train_valid_test_num_samples() @@ -5505,8 +5192,14 @@ def build_train_valid_test_data_loaders(build_train_valid_test_datasets_provider # Get consumed train samples in this phase. if args.phase_transition_iterations: - last_transition = max(iteration for iteration in (0, *args.phase_transition_iterations) if iteration <= args.iteration) - consumed_train_samples_in_current_phase = (args.iteration - last_transition) * args.global_batch_size + last_transition = max( + iteration + for iteration in (0, *args.phase_transition_iterations) + if iteration <= args.iteration + ) + consumed_train_samples_in_current_phase = ( + args.iteration - last_transition + ) * args.global_batch_size else: consumed_train_samples_in_current_phase = args.consumed_train_samples @@ -5523,17 +5216,21 @@ def build_train_valid_test_data_loaders(build_train_valid_test_datasets_provider valid_dataloaders = None test_dataloader = None do_train = (args.train_iters or 0) > 0 - do_valid = (args.full_validation or args.eval_iters > 0) - do_test = False + do_valid = args.full_validation or args.eval_iters > 0 + do_test = args.full_validation or args.eval_iters > 0 else: # Build datasets. - train_ds, valid_ds, test_ds = build_train_valid_test_datasets(build_train_valid_test_datasets_provider) + train_ds, valid_ds, test_ds = build_train_valid_test_datasets( + build_train_valid_test_datasets_provider + ) valid_ds = [valid_ds] if not isinstance(valid_ds, list) else valid_ds if args.skip_train: train_dataloader = None else: - train_dataloader = build_pretraining_data_loader(train_ds, consumed_train_samples_in_current_phase) + train_dataloader = build_pretraining_data_loader( + train_ds, consumed_train_samples_in_current_phase + ) valid_dataloaders = [] for valid_d in valid_ds: if args.skip_train or args.full_validation: @@ -5542,13 +5239,19 @@ def build_train_valid_test_data_loaders(build_train_valid_test_datasets_provider if args.multiple_validation_sets: # TODO(bnorick): for multiple validation sets without full validation, args.consumed_valid_samples is not # correct and needs to be calculated/set per validation set - raise NotImplementedError("--multiple-validation-sets currently requires --full-validation") - valid_dataloaders.append(build_pretraining_data_loader(valid_d, args.consumed_valid_samples)) + raise NotImplementedError( + "--multiple-validation-sets currently requires --full-validation" + ) + valid_dataloaders.append( + build_pretraining_data_loader(valid_d, args.consumed_valid_samples) + ) if not args.multiple_validation_sets: assert len(valid_dataloaders) == 1 test_dataloader = build_pretraining_data_loader(test_ds, 0) do_train = train_dataloader is not None and (args.skip_train or args.train_iters > 0) - do_valid = valid_dataloaders is not None and (args.full_validation or args.eval_iters > 0) + do_valid = valid_dataloaders is not None and ( + args.full_validation or args.eval_iters > 0 + ) do_test = test_dataloader is not None and (args.full_validation or args.eval_iters > 0) flags = torch.tensor( @@ -5608,7 +5311,9 @@ def _get_iterator(dataloader_type, dataloader): args.eval_iters = [None] * len(valid_dataloaders) else: local_eval_iters = [len(dl) for dl in valid_dataloaders] - eval_iters_tensor = torch.tensor(local_eval_iters, dtype=torch.long, device='cuda') + eval_iters_tensor = torch.tensor( + local_eval_iters, dtype=torch.long, device='cuda' + ) torch.distributed.all_reduce( eval_iters_tensor, op=torch.distributed.ReduceOp.MAX, @@ -5617,7 +5322,9 @@ def _get_iterator(dataloader_type, dataloader): args.eval_iters = eval_iters_tensor.tolist() else: local_eval_iters = len(valid_dataloaders[0]) - eval_iters_tensor = torch.tensor([local_eval_iters], dtype=torch.long, device='cuda') + eval_iters_tensor = torch.tensor( + [local_eval_iters], dtype=torch.long, device='cuda' + ) torch.distributed.all_reduce( eval_iters_tensor, op=torch.distributed.ReduceOp.MAX, @@ -5667,3 +5374,40 @@ def should_disable_forward_pre_hook(args): ) and args.overlap_param_gather ) + + +def _get_thd_sequence_length_upper_bound(args): + """Return the padded per-sample THD length upper bound used for graph sizing.""" + max_sequence_length = getattr(args, "seq_length", None) + mock_config_spec = None + if getattr(args, "use_varlen_dataset", False): + mock_config_spec = getattr(args, "varlen_mock_dataset_config_json", None) + elif getattr(args, "sft", False): + mock_config_spec = getattr(args, "sft_mock_dataset_config_json", None) + + if mock_config_spec is not None: + from megatron.training.datasets.utils import load_json_arg + + mock_config = load_json_arg(mock_config_spec) + if isinstance(mock_config, dict) and mock_config.get("max_seq_len") is not None: + max_sequence_length = int(mock_config["max_seq_len"]) + + if max_sequence_length is None: + return None + + if getattr(args, "seq_length", None) is not None: + max_sequence_length = min(int(max_sequence_length), int(args.seq_length)) + + cp_size = int(getattr(args, "context_parallel_size", 1) or 1) + if getattr(args, "dynamic_context_parallel", False): + cp_pad = int(getattr(args, "data_parallel_size", 1) or 1) * cp_size * 2 + else: + cp_pad = cp_size * 2 if cp_size > 1 else 1 + + sp_pad = ( + int(getattr(args, "tensor_model_parallel_size", 1) or 1) + if getattr(args, "sequence_parallel", False) + else 1 + ) + pad_granularity = cp_pad * sp_pad + return int(math.ceil(max_sequence_length / pad_granularity) * pad_granularity) diff --git a/tests/unit_tests/test_num_floating_point_operations.py b/tests/unit_tests/test_num_floating_point_operations.py index df5e4191843..0f08008ec1b 100644 --- a/tests/unit_tests/test_num_floating_point_operations.py +++ b/tests/unit_tests/test_num_floating_point_operations.py @@ -100,6 +100,24 @@ def _make_hybrid_args(*, num_layers=4, hidden_size=512, num_attention_heads=8, s return args +def _make_ling_hybrid_args(pattern): + """Minimal direct-KDA/MLA dimensions for hybrid FLOPs tests.""" + args = _make_hybrid_args() + args.hybrid_layer_pattern = pattern + 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 + 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 + return args + + class TestBSHDBackwardCompat: """For unpacked BSHD, the new optional arg must not change the result.""" @@ -286,6 +304,55 @@ def test_householder_count_changes_flops(self): assert flops_m4 - flops_m3 == expected_delta +class TestLingHybridAttentionFlops: + """KDA and MLA layers must contribute to Ling-V3 Tiny MFU accounting.""" + + def test_kda_direct_projection_formula(self): + args = _make_ling_hybrid_args("K") + 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) + + def test_mla_uses_ragged_attention_work(self): + args = _make_ling_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 + + class TestPaddingRemoval: """``total_real_tokens_in_batch`` removes padding from token-linear FLOPs. From 824ca51fb7df311139487bd5af87c5ce4ac6e6f5 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Mon, 17 Aug 2026 22:27:36 +0800 Subject: [PATCH 09/34] Fix Hybrid MTP z-loss metric indexing Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 7f0ae9f300d5e07528ceb4d4c793a8b6b798f4f0) --- megatron/core/transformer/moe/router.py | 565 ++++++++++-------- .../transformer/moe/test_routers.py | 22 + 2 files changed, 334 insertions(+), 253 deletions(-) diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 6ce93c76837..b26bf1a0d3e 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -1,15 +1,14 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from abc import ABC, abstractmethod -from typing import Optional, Union +from dataclasses import dataclass +from typing import Optional, Sequence, Union import torch from megatron.core.inference.utils import InferenceMode from megatron.core.jit import jit_fuser -from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( - is_batch_invariant_mode_enabled, -) +from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.moe.moe_logging import get_moe_metrics_tracker from megatron.core.transformer.moe.moe_utils import ( @@ -20,17 +19,30 @@ apply_router_token_dropping, compute_routing_scores_for_aux_loss, get_tokens_per_expert_and_token_count, - qb_dual_update, router_gating_linear, sinkhorn, switch_load_balancing_loss_func, topk_routing_with_score_function, - z_loss_func, ) from megatron.core.transformer.moe.router_replay import RouterReplay from megatron.core.transformer.transformer_config import TransformerConfig +@dataclass(frozen=True) +class _AuxLossGroupConfig: + """Process groups for local aux/seq-aux loss and its metric logging.""" + + loss_reduce_groups: Sequence[torch.distributed.ProcessGroup] + metric_reduce_group: Optional[torch.distributed.ProcessGroup] + metric_avg_group: Optional[torch.distributed.ProcessGroup] + metric_needs_dp_avg: bool + + @property + def metric_pre_reduce_groups(self) -> Optional[Sequence[torch.distributed.ProcessGroup]]: + """Groups to reduce eagerly before recording metrics, if tracker reduction is unsafe.""" + return self.loss_reduce_groups if self.metric_avg_group is not None else None + + class Router(ABC, MegatronModule): """Base Router class""" @@ -39,6 +51,7 @@ def __init__( config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None, is_mtp_layer: bool = False, + layer_number: Optional[int] = None, ) -> None: """ Initialize the Router module. @@ -52,7 +65,7 @@ def __init__( self.config = config self.num_experts = self.config.num_moe_experts self.moe_aux_loss_func = None - self.layer_number = None + self.layer_number = layer_number self.is_mtp_layer = is_mtp_layer self.tp_group = pg_collection.tp self.cp_group = pg_collection.cp @@ -70,10 +83,6 @@ def __init__( ) else: self.bias = None - if self.config.moe_router_skip_muon: - setattr(self.weight, 'use_muon', False) - if self.bias is not None: - setattr(self.bias, 'use_muon', False) # If calculate per token loss, we need to scale up moe aux loss by the number of tokens. # So we need to know if the model is configured to calculate per token loss. self.calculate_per_token_loss = self.config.calculate_per_token_loss @@ -166,6 +175,7 @@ def __init__( config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None, is_mtp_layer: bool = False, + layer_number: Optional[int] = None, ) -> None: """Initialize the zero token dropping router. @@ -174,14 +184,43 @@ def __init__( pg_collection (ProcessGroupCollection, optional): Process groups for MoE operations. is_mtp_layer (bool): Flag indicating if this router is part of an MTP layer. """ - super().__init__(config=config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer) + super().__init__( + config=config, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + layer_number=layer_number, + ) self.topk = self.config.moe_router_topk self.routing_type = self.config.moe_router_load_balancing_type self.score_function = self.config.moe_router_score_function self.input_jitter = None + self.mtp_layer_number: Optional[int] = None self.frozen_expert_bias = False - self.enable_expert_bias = self.config.moe_router_enable_expert_bias + if self.config.moe_n_hash_layers > 0: + assert layer_number is not None, "layer_number is required for the hash-based router." + self.is_hash_layer = ( + not self.is_mtp_layer + and self.config.moe_n_hash_layers > 0 + and layer_number <= self.config.moe_n_hash_layers + ) + if self.is_hash_layer: + # DSv4-Pro ships a pre-trained tid2eid table in its inference checkpoint; + # no public initialization recipe is documented. Round-robin is used here + # only as a placeholder so the layer is runnable from scratch. + vocab_size = self.config.actual_vocab_size + num_experts = self.config.num_moe_experts + ids = torch.arange(vocab_size, device=torch.cuda.current_device()) + tid2eid = torch.stack([(ids + k) % num_experts for k in range(self.topk)], dim=1).to( + torch.int32 + ) + self.register_buffer('tid2eid', tid2eid) + else: + self.tid2eid = None + + self.enable_expert_bias = ( + self.config.moe_router_enable_expert_bias and not self.is_hash_layer + ) if self.enable_expert_bias: self.register_buffer( 'local_tokens_per_expert', @@ -224,41 +263,6 @@ def __init__( self.global_tokens_per_expert = None self.ga_steps = None - # Quantile balancing replaces the aux loss with a per-expert bias `qb_beta`. - # `qb_beta_accum`/`qb_beta_count` collect the per-microbatch quantile, reduced - # and reset each global batch. - if self.routing_type == "quantile_balancing": - assert not self.is_aux_loss_enabled(), ( - "Quantile balancing handles load balance via the bias update; " - "aux losses must be disabled (set moe_aux_loss_coeff to 0)." - ) - self.register_buffer( - 'qb_beta', - torch.zeros( - self.config.num_moe_experts, - dtype=torch.float32, - device=torch.cuda.current_device(), - ), - ) - self.register_buffer( - 'qb_beta_accum', - torch.zeros( - self.config.num_moe_experts, - dtype=torch.float32, - device=torch.cuda.current_device(), - ), - persistent=False, - ) - self.register_buffer( - 'qb_beta_count', - torch.zeros((), dtype=torch.long, device=torch.cuda.current_device()), - persistent=False, - ) - else: - self.qb_beta = None - self.qb_beta_accum = None - self.qb_beta_count = None - self.router_replay = None if self.config.moe_enable_routing_replay: self.router_replay = RouterReplay() @@ -273,13 +277,6 @@ def _maintain_float32_expert_bias(self): if hasattr(self, 'expert_bias') and self.expert_bias is not None: if self.expert_bias.dtype != torch.float32: self.expert_bias.data = self.expert_bias.data.to(torch.float32) - # Keep the QB bias in fp32 for the same reason. - if hasattr(self, 'qb_beta') and self.qb_beta is not None: - if self.qb_beta.dtype != torch.float32: - self.qb_beta.data = self.qb_beta.data.to(torch.float32) - if hasattr(self, 'qb_beta_accum') and self.qb_beta_accum is not None: - if self.qb_beta_accum.dtype != torch.float32: - self.qb_beta_accum.data = self.qb_beta_accum.data.to(torch.float32) def sinkhorn_load_balancing(self, logits: torch.Tensor): """Apply sinkhorn routing to the logits tensor. @@ -314,81 +311,6 @@ def _sinkhorn_activation(logits): scores = logits * map return scores, map - def quantile_balancing(self, logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Apply quantile-balancing (QB) routing to the logits tensor. - - Selects top-k experts per token using a dual coordinate-descent update on - a per-expert bias ``qb_beta``. Load balance is handled entirely by the bias - update; auxiliary losses must be disabled when QB is active. - - Args: - logits (torch.Tensor): The logits tensor, shape ``[num_tokens, num_experts]``. - - Returns: - Tuple[torch.Tensor, torch.Tensor]: Sparse routing probs and boolean - routing map, each shaped ``[num_tokens, num_experts]``. - """ - assert ( - not self.config.moe_router_fusion - ), "Quantile balancing routing does not support moe_router_fusion." - assert ( - self.config.moe_router_num_groups is None and self.config.moe_router_group_topk is None - ), "Quantile balancing routing does not support group-limited routing." - - local_num_tokens = logits.shape[0] - # Gather logits across TP/CP so the quantile sees a whole sequence's tokens. - # The DP reduction and qb_beta update run at the global-batch boundary in - # finalize_model_grads._update_router_qb_beta. - gather_group = self.tp_cp_group - gather_size = gather_group.size() if gather_group is not None else 1 - - should_update_beta = self.training and torch.is_grad_enabled() - - with torch.no_grad(): - logits_fp32 = logits.detach().to(dtype=torch.float32) - - if gather_size > 1: - full_logits = torch.empty( - (local_num_tokens * gather_size, self.config.num_moe_experts), - dtype=logits_fp32.dtype, - device=logits_fp32.device, - ) - torch.distributed.all_gather_into_tensor( - full_logits, logits_fp32.contiguous(), group=gather_group - ) - gather_rank = torch.distributed.get_rank(group=gather_group) - else: - full_logits = logits_fp32 - gather_rank = 0 - - # Route with the previous batch's qb_beta; in training, accumulate this - # microbatch's quantile for the next update. - full_indices, beta_local = qb_dual_update( - full_logits, self.topk, self.qb_beta, update_beta=should_update_beta - ) - if should_update_beta: - self.qb_beta_accum.add_(beta_local) - self.qb_beta_count.add_(1) - - # Take this rank's rows (all_gather orders rows by rank). - if gather_size > 1: - indices = full_indices[ - gather_rank * local_num_tokens : (gather_rank + 1) * local_num_tokens - ].contiguous() - else: - indices = full_indices - - # QB only picks the experts; reuse the shared score function for the probs. - return topk_routing_with_score_function( - logits, - self.topk, - use_pre_softmax=self.config.moe_router_pre_softmax, - scaling_factor=self.config.moe_router_topk_scaling_factor, - score_function=self.score_function, - fused=self.config.moe_router_fusion, - precomputed_indices=indices, - ) - def get_aux_loss_coeff(self, aux_loss_type: str) -> float: """Return the aux loss coeff for the given auxiliary loss type. If the auxiliary loss type is not found, return 0.0. @@ -417,16 +339,19 @@ def _apply_aux_loss( scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor, with_padding_mask: bool = False, + packed_seq_params: Optional[PackedSeqParams] = None, ): """Apply the auxiliary loss for the given scores and routing map.""" aux_loss_coeff = self.get_aux_loss_coeff("aux_loss") if aux_loss_coeff == 0: return probs + aux_loss_groups = self._get_aux_loss_groups(packed_seq_params) global_tokens_per_expert, local_num_tokens, total_num_tokens = ( get_tokens_per_expert_and_token_count( routing_map=routing_map, - reduce_group=self.tp_cp_group, + reduce_group=aux_loss_groups.loss_reduce_groups[0], + reduce_groups=aux_loss_groups.loss_reduce_groups, topk=self.topk, with_padding_mask=with_padding_mask, ) @@ -441,13 +366,19 @@ def _apply_aux_loss( moe_aux_loss_coeff=aux_loss_coeff, fused=self.config.moe_router_fusion, ) + probs = self.attach_and_log_load_balancing_loss( probs, aux_loss_coeff, aux_loss, "load_balancing_loss", - self.tp_cp_group, + aux_loss_groups.metric_reduce_group, + avg_group=aux_loss_groups.metric_avg_group, + needs_dp_avg=aux_loss_groups.metric_needs_dp_avg, valid_token_count=local_num_tokens, + aux_loss_logging_reduce_groups=aux_loss_groups.metric_pre_reduce_groups, + aux_loss_scale_reduce_groups=aux_loss_groups.loss_reduce_groups, + aux_loss_scale_num_tokens=total_num_tokens, ) return probs @@ -459,6 +390,7 @@ def _apply_seq_aux_loss( seq_length: int, bsz: int, with_padding_mask: bool = False, + packed_seq_params: Optional[PackedSeqParams] = None, ): """Apply the sequence-level auxiliary loss for the given scores and routing map. @@ -474,10 +406,12 @@ def _apply_seq_aux_loss( scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1) routing_map = routing_map.reshape(seq_length, -1) + aux_loss_groups = self._get_aux_loss_groups(packed_seq_params) global_tokens_per_expert, local_num_tokens, total_num_tokens = ( get_tokens_per_expert_and_token_count( routing_map=routing_map, - reduce_group=self.tp_cp_group, + reduce_group=aux_loss_groups.loss_reduce_groups[0], + reduce_groups=aux_loss_groups.loss_reduce_groups, with_padding_mask=with_padding_mask, topk=self.topk * bsz, ) @@ -501,10 +435,13 @@ def _apply_seq_aux_loss( seq_aux_loss_coeff, aux_loss, "seq_load_balancing_loss", - self.tp_cp_group, - # local_num_tokens is per-sequence (bsz folded into the expert dim above); - # * bsz recovers the micro-batch total, else per-token-loss scaling keeps a 1/MBS. - valid_token_count=local_num_tokens * bsz, + aux_loss_groups.metric_reduce_group, + avg_group=aux_loss_groups.metric_avg_group, + needs_dp_avg=aux_loss_groups.metric_needs_dp_avg, + valid_token_count=local_num_tokens, + aux_loss_logging_reduce_groups=aux_loss_groups.metric_pre_reduce_groups, + aux_loss_scale_reduce_groups=aux_loss_groups.loss_reduce_groups, + aux_loss_scale_num_tokens=total_num_tokens, ) return probs @@ -520,7 +457,8 @@ def _apply_global_aux_loss( if global_aux_loss_coeff == 0: return probs - # Use unified function to compute tokens_per_expert and num_tokens + # Global aux loss intentionally uses the full static TP x DP x CP domain. + # Dynamic CP subgroups only affect local aux/seq-aux domains. global_tokens_per_expert, local_num_tokens, total_num_tokens = ( get_tokens_per_expert_and_token_count( routing_map=routing_map, @@ -529,7 +467,6 @@ def _apply_global_aux_loss( topk=self.topk, ) ) - self.global_tokens_per_expert += global_tokens_per_expert self.ga_steps += 1 averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps @@ -543,6 +480,7 @@ def _apply_global_aux_loss( moe_aux_loss_coeff=global_aux_loss_coeff, fused=self.config.moe_router_fusion, ) + probs = self.attach_and_log_load_balancing_loss( probs, global_aux_loss_coeff, @@ -551,18 +489,64 @@ def _apply_global_aux_loss( self.tp_dp_cp_group, needs_dp_avg=False, valid_token_count=local_num_tokens, + # The global aux-loss statistics/logging domain is TP x DP x CP, but + # per-token-loss gradient normalization already reduces the denominator + # across DP x CP in finalize_model_grads. Scale the aux-loss numerator + # over TP x CP only, matching the original static behavior while still + # using an exact valid-token count when padding is present. + aux_loss_scale_reduce_groups=(self.tp_cp_group,), ) return probs + def _get_aux_loss_groups( + self, packed_seq_params: Optional[PackedSeqParams] = None + ) -> _AuxLossGroupConfig: + """Return process groups for MoE aux-loss statistics and logging.""" + if ( + packed_seq_params is not None + and packed_seq_params.local_cp_size is not None + and packed_seq_params.cp_group is not None + ): + return _AuxLossGroupConfig( + loss_reduce_groups=(packed_seq_params.cp_group, self.tp_group), + metric_reduce_group=None, + metric_avg_group=self.tp_dp_cp_group, + metric_needs_dp_avg=False, + ) + + return _AuxLossGroupConfig( + loss_reduce_groups=(self.tp_cp_group,), + metric_reduce_group=self.tp_cp_group, + metric_avg_group=None, + metric_needs_dp_avg=True, + ) + + def _get_metric_layer_number(self) -> tuple[int, int]: + """Return the 1-based metric index and the total number of metric slots.""" + num_layers = self.config.num_layers + (self.config.mtp_num_layers or 0) + if not self.is_mtp_layer: + return self.layer_number, num_layers + + # Hybrid MTP depths can contain multiple internal sublayers (for example `/WE`). + # Metrics are allocated per MTP depth, not per internal hybrid sublayer. + mtp_layer_number = self.mtp_layer_number or self.layer_number + if self.config.mtp_num_layers is not None: + mtp_layer_number = min(mtp_layer_number, self.config.mtp_num_layers) + return self.config.num_layers + mtp_layer_number, num_layers + def attach_and_log_load_balancing_loss( self, activation: torch.Tensor, aux_loss_coeff: float, aux_loss: torch.Tensor, aux_loss_name: str, - reduce_group: torch.distributed.ProcessGroup, + reduce_group: Optional[torch.distributed.ProcessGroup], + avg_group: Optional[torch.distributed.ProcessGroup] = None, needs_dp_avg: bool = True, valid_token_count: Optional[Union[int, torch.Tensor]] = None, + aux_loss_logging_reduce_groups: Optional[Sequence[torch.distributed.ProcessGroup]] = None, + aux_loss_scale_reduce_groups: Optional[Sequence[torch.distributed.ProcessGroup]] = None, + aux_loss_scale_num_tokens: Optional[Union[int, torch.Tensor]] = None, ): """Attach aux loss function to activation and add to logging. @@ -571,7 +555,10 @@ def attach_and_log_load_balancing_loss( aux_loss_coeff (float): Coefficient for the aux loss. aux_loss (torch.Tensor): Computed aux loss. aux_loss_name (str): Name of the aux loss for logging. - reduce_group (torch.distributed.ProcessGroup): Process group for reduction. + reduce_group (torch.distributed.ProcessGroup, optional): Process group for deferred + logging reduction. + avg_group (torch.distributed.ProcessGroup, optional): Process group for deferred + logging average. needs_dp_avg (bool): Whether to average this metric across DP ranks after reduce_group. valid_token_count (int or torch.Tensor, optional): Number of valid tokens excluding padding tokens. Can be a Python int or a torch.Tensor (typically 0-d tensor). @@ -591,53 +578,49 @@ def attach_and_log_load_balancing_loss( # add the aux loss logging value to other layer's since it is difficult to get the # correct layer_number for MTP. It does not affect the correctness of the calculation # results and the reduced load_balancing_loss logging value. - num_layers = self.config.num_layers - if self.config.mtp_num_layers is not None: - num_layers += self.config.mtp_num_layers + layer_number, num_layers = self._get_metric_layer_number() - if self.is_mtp_layer: - layer_number = self.layer_number + self.config.num_layers - else: - layer_number = self.layer_number + metric_value = aux_loss / aux_loss_coeff + if aux_loss_logging_reduce_groups is not None: + metric_value = metric_value.detach().clone() + for group in aux_loss_logging_reduce_groups: + torch.distributed.all_reduce(metric_value, group=group) get_moe_metrics_tracker().record( aux_loss_name, - aux_loss / aux_loss_coeff, + metric_value, layer_number, num_layers, reduce_group=reduce_group, + avg_group=avg_group, needs_dp_avg=needs_dp_avg, ) if self.calculate_per_token_loss: - # Target final scaling on aux_loss gradients: 1 / (num_micro_batches * dp_size), - # matching the !calculate_per_token_loss path. - # - # --calculate-per-token-loss already divides every parameter gradient by - # total_global_tokens (the global non-padded token count summed in - # finalize_model_grads). The router's `num_local_tokens` (= activation.shape[0]) - # is sequence-parallel sharded — the router weight is marked - # `sequence_parallel=True` in Router.reset_parameters (see - # `setattr(self.weight, 'sequence_parallel', ...)` above), so each TP rank - # computes a partial gradient on the router weight from its local sequence - # shard, and `_allreduce_non_tensor_model_parallel_grads` SUMS those partial - # gradients across the TP group. Re-expressing total_global_tokens in terms of the - # router's `num_local_tokens`: - # total_global_tokens - # = num_micro_batches * dp_cp_size * loss_func_local_tokens - # = num_micro_batches * dp_cp_size * tp_size * num_local_tokens - # = num_micro_batches * dp_size * (num_local_tokens * tp_cp_group.size()) - # (using loss_func_local_tokens = tp_size * num_local_tokens, then regrouping - # dp_cp_size * tp_size as dp_size * tp_cp_group.size()). - # - # So pre-multiplying aux_loss by num_local_tokens * tp_cp_group.size() cancels - # that same factor in total_global_tokens above, leaving 1 / (num_micro_batches * - # dp_size) as the effective scaling on the aux_loss gradient — the target. - # Use valid_token_count (excluding padding) if provided, otherwise use total tokens. - num_local_tokens = ( - valid_token_count if valid_token_count is not None else activation.shape[0] - ) + # --calculate-per-token-loss divides all parameter gradients by the global + # non-padded token count in finalize_model_grads. Pre-multiplying by the + # valid-token count from this aux-loss domain makes the final objective a + # token-weighted average of per-domain aux losses. Use the reduced count + # directly: with THD padding or dynamic CP, valid token counts can differ + # by rank/group, so local_num_tokens * group_size is not generally correct. + if aux_loss_scale_num_tokens is None: + num_local_tokens = ( + valid_token_count if valid_token_count is not None else activation.shape[0] + ) + if torch.is_tensor(num_local_tokens): + aux_loss_scale_num_tokens = num_local_tokens.clone().to( + device=activation.device + ) + else: + aux_loss_scale_num_tokens = torch.tensor( + num_local_tokens, device=activation.device + ) + if aux_loss_scale_reduce_groups is None: + assert reduce_group is not None, "reduce_group is required for aux-loss scaling" + aux_loss_scale_reduce_groups = (reduce_group,) + for group in aux_loss_scale_reduce_groups: + torch.distributed.all_reduce(aux_loss_scale_num_tokens, group=group) activation = MoEAuxLossAutoScaler.apply( - activation, aux_loss * num_local_tokens * self.tp_cp_group.size() + activation, aux_loss * aux_loss_scale_num_tokens ) else: activation = MoEAuxLossAutoScaler.apply(activation, aux_loss) @@ -649,66 +632,52 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None): Args: logits (torch.Tensor): The logits of the router. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape in [num_tokens]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [num_tokens]. True = padding, + False = valid. Defaults to None. Returns: torch.Tensor: The logits after applying the z-loss. """ if self.config.moe_z_loss_coeff is not None and self.training and torch.is_grad_enabled(): # Skip Z loss calculations when using torch.no_grad() or checkpointing. - moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size() - z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask) - if self.calculate_per_token_loss: - # Same derivation as in attach_and_log_load_balancing_loss: - # - Target final scaling on z_loss gradients: 1 / (num_micro_batches * dp_size). - # - In terms of the router's `num_local_tokens`, the total_global_tokens - # divisor that finalize_model_grads applies factors as - # num_micro_batches * dp_size * (num_local_tokens * tp_cp_group.size()). - # - Pre-multiplying z_loss by num_local_tokens * tp_cp_group.size() cancels - # that same factor in total_global_tokens, leaving - # 1 / (num_micro_batches * dp_size) as the effective scaling — the target. - # The /tp_cp_group.size() on moe_z_loss_coeff above is a separate forward-side - # correction: z_loss is computed independently on each TP+CP rank's local - # logits and must be averaged across TP+CP rather than summed. - # Count valid tokens: sum of inverted mask (False -> True = valid) - num_local_tokens = ( - (~padding_mask).sum() if padding_mask is not None else logits.shape[0] - ) - logits = MoEAuxLossAutoScaler.apply( - logits, z_loss * num_local_tokens * self.tp_cp_group.size() - ) + logsum = torch.logsumexp(logits, dim=-1) + z_loss_values = torch.square(logsum) + if padding_mask is not None: + valid_mask = ~padding_mask + z_loss_values = z_loss_values * valid_mask + num_local_tokens = valid_mask.sum() + z_loss_sum = z_loss_values.sum() + z_loss_mean = z_loss_sum / torch.clamp(num_local_tokens, min=1) else: - logits = MoEAuxLossAutoScaler.apply(logits, z_loss) + z_loss_sum = z_loss_values.sum() + # Keep the token count as a Python scalar so CUDA graph capture does not + # record a CPU-to-CUDA tensor creation. + z_loss_mean = z_loss_sum / max(logits.shape[0], 1) - # TODO: repeated-MTP z_loss is scaled after MoEAuxLossAutoScaler.apply(), so this - # adjusts logging only; move the scaling above the attach point if z_loss is used. - # When using repeated MTP layers, the same MTP layer is called mtp_num_layers times. - # To avoid accumulating the z_loss multiple times, we scale it by 1/mtp_num_layers - # so the total loss is correct. + mtp_loss_scale = 1 if ( self.is_mtp_layer and self.config.mtp_use_repeated_layer and self.config.mtp_num_layers is not None ): - z_loss = z_loss / self.config.mtp_num_layers - - num_layers = self.config.num_layers - if self.config.mtp_num_layers is not None: - num_layers += self.config.mtp_num_layers + mtp_loss_scale = self.config.mtp_num_layers - if self.is_mtp_layer: - layer_number = self.layer_number + self.config.num_layers + if self.calculate_per_token_loss: + # --calculate-per-token-loss divides gradients by the global non-padded + # token count. Attach the local z-loss numerator directly so the final + # objective is a token-weighted z-loss over valid tokens. + z_loss = z_loss_sum * self.config.moe_z_loss_coeff / mtp_loss_scale + logits = MoEAuxLossAutoScaler.apply(logits, z_loss) else: - layer_number = self.layer_number + moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size() + z_loss = z_loss_mean * moe_z_loss_coeff / mtp_loss_scale + logits = MoEAuxLossAutoScaler.apply(logits, z_loss) + + layer_number, num_layers = self._get_metric_layer_number() get_moe_metrics_tracker().record( - "z_loss", - z_loss / moe_z_loss_coeff, - layer_number, - num_layers, - avg_group=self.tp_dp_cp_group, + "z_loss", z_loss_mean / mtp_loss_scale, layer_number, num_layers ) return logits @@ -744,17 +713,77 @@ def _apply_expert_bias( if self.enable_expert_bias and torch.is_grad_enabled(): with torch.no_grad(): if padding_mask is not None: - routing_map = routing_map & (~padding_mask).unsqueeze(-1) + flat_mask = padding_mask.reshape(-1) + assert ( + flat_mask.shape[0] == routing_map.shape[0] + ), f"padding_mask flat {flat_mask.shape} vs routing_map {routing_map.shape}" + routing_map = routing_map & (~flat_mask).unsqueeze(-1) self.local_tokens_per_expert += routing_map.sum(dim=0) - def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def _hash_routing(self, logits: torch.Tensor, input_ids: torch.Tensor): + """Hash-based routing: expert indices come from the tid2eid lookup table. + + Scores are still computed from the gating logits for weight computation, + but expert selection is determined by the pre-computed hash table. + + Args: + logits (torch.Tensor): Gating logits, shape [num_tokens, num_experts]. + input_ids (torch.Tensor): Token IDs, shape [seq_length, bsz]. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: routing_probs and routing_map. + """ + num_tokens, num_experts = logits.shape + + if self.score_function == "softmax": + scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) + elif self.score_function == "sigmoid": + scores = torch.sigmoid(logits.float()).type_as(logits) + elif self.score_function == "sqrtsoftplus": + scores = torch.nn.functional.softplus(logits.float()).sqrt().type_as(logits) + else: + raise ValueError(f"Invalid score_function: {self.score_function}") + + # input_ids is [b, s] from the model, but hidden_states are [s, b, h] + # and get flattened to [s*b, h]. Transpose to match. + flat_ids = input_ids.T.reshape(-1) + top_indices = self.tid2eid[flat_ids].long() # [num_tokens, topk] + if ( + self.config.moe_router_force_load_balancing + or self.config.moe_router_force_biased is not None + ): + # override top_indices with random topk indices + # logits in processed by apply_random_logits or apply_biased_logits + _, top_indices = torch.topk(logits, k=self.topk, dim=1) + + probs = scores.gather(1, top_indices) + if self.score_function != "softmax": + probs = probs / (probs.sum(dim=-1, keepdim=True) + 1e-20) + + if self.config.moe_router_topk_scaling_factor: + probs = probs * self.config.moe_router_topk_scaling_factor + + routing_probs = torch.zeros_like(logits).scatter(1, top_indices, probs) + routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool() + + return routing_probs, routing_map + + def routing( + self, + logits: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): """Top-k routing function Args: logits (torch.Tensor): Logits tensor after gating. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape [seq_length, bsz]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [seq_length, bsz]. True = padding, + False = valid. Defaults to None. + input_ids (torch.Tensor, optional): The input IDs tensor. Shape [seq_length, bsz]. + Defaults to None. Returns: probs (torch.Tensor): The probabilities of token to experts assignment. @@ -772,13 +801,14 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N logits = self.apply_z_loss(logits, padding_mask=padding_mask) # Calculate probs and routing_map for token dispatching - if self.routing_type == "sinkhorn": + if self.is_hash_layer: + assert input_ids is not None, ( + "input_ids is required for hash-based routing but was None. " + "Ensure --moe-n-hash-layers is set correctly and input_ids are passed." + ) + probs, routing_map = self._hash_routing(logits, input_ids) + elif self.routing_type == "sinkhorn": probs, routing_map = self.sinkhorn_load_balancing(logits) - elif self.routing_type == "quantile_balancing": - assert ( - padding_mask is None - ), "Quantile balancing routing does not support padding masks yet." - probs, routing_map = self.quantile_balancing(logits) else: probs, routing_map = topk_routing_with_score_function( logits, @@ -793,6 +823,20 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N router_replay=self.router_replay, ) + # Dropless HybridEP consumes the sparse routing map directly, so exclude padding + # rows before dispatch. Other dispatchers retain their existing fixed-route + # assumptions until they support sparse routing maps end to end. + use_dropless_hybridep = ( + self.config.moe_token_dispatcher_type == "flex" + and self.config.moe_flex_dispatcher_backend == "hybridep" + and self.config.moe_expert_capacity_factor is None + and self.config.moe_expert_rank_capacity_factor is None + ) + if padding_mask is not None and use_dropless_hybridep: + valid_tokens = (~padding_mask).unsqueeze(-1) + probs = probs * valid_tokens + routing_map = routing_map & valid_tokens + # Apply token dropping to probs and routing_map. if self.config.moe_expert_capacity_factor is not None: probs, routing_map = apply_router_token_dropping( @@ -819,6 +863,7 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N scores_for_aux_loss, routing_map_for_aux_loss, with_padding_mask=padding_mask is not None, + packed_seq_params=packed_seq_params, ) probs = self._apply_seq_aux_loss( probs, @@ -827,6 +872,7 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N seq_length, bsz, with_padding_mask=padding_mask is not None, + packed_seq_params=packed_seq_params, ) probs = self._apply_global_aux_loss( probs, @@ -846,15 +892,23 @@ def reset_global_aux_loss_tracker(self): self.global_tokens_per_expert.zero_() self.ga_steps.zero_() - def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def forward( + self, + input: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): """ Forward pass of the router. Args: input (torch.Tensor): Input tensor. - padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. - Shape [seq_length, bsz]. True for valid tokens, - False for padding tokens. Defaults to None. + padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions. + Shape [seq_length, bsz]. True = padding, + False = valid. Defaults to None. + input_ids (torch.Tensor, optional): The input IDs tensor. Shape [seq_length, bsz]. + Defaults to None. """ self._maintain_float32_expert_bias() @@ -872,7 +926,12 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No logits, self.config.moe_router_force_biased, self.layer_number ) - probs, routing_map = self.routing(logits, padding_mask=padding_mask) + probs, routing_map = self.routing( + logits, + padding_mask=padding_mask, + input_ids=input_ids, + packed_seq_params=packed_seq_params, + ) return probs, routing_map @@ -904,6 +963,7 @@ def __init__( config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None, is_mtp_layer: bool = False, + layer_number: Optional[int] = None, ) -> None: """Initialize the specialized inference top-k router. @@ -921,7 +981,12 @@ def __init__( f"['sigmoid', 'softmax'], got '{config.moe_router_score_function}'" ) - super().__init__(config=config, pg_collection=pg_collection) + super().__init__( + config=config, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + layer_number=layer_number, + ) @staticmethod @torch.compile @@ -937,7 +1002,6 @@ def _compiled_topk_routing( fused, router_replay, dense_output, - precomputed_indices, ): return topk_routing_with_score_function( logits, @@ -951,23 +1015,12 @@ def _compiled_topk_routing( fused=fused, router_replay=router_replay, dense_output=dense_output, - precomputed_indices=precomputed_indices, ) def _forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): logits = self.gating(input).squeeze(1) # [num_tokens, num_experts] - # QB selects on (logits - qb_beta); at inference qb_beta is fixed, so it's per-token. - precomputed_indices = None - if self.qb_beta is not None: - precomputed_indices = (logits - self.qb_beta).topk(self.topk, dim=1).indices - - routing = ( - topk_routing_with_score_function - if is_batch_invariant_mode_enabled() - else self._compiled_topk_routing - ) - probs, top_indices = routing( + probs, top_indices = self._compiled_topk_routing( logits, self.topk, use_pre_softmax=self.config.moe_router_pre_softmax, @@ -979,16 +1032,22 @@ def _forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = N fused=self.config.moe_router_fusion, router_replay=self.router_replay, dense_output=True, - precomputed_indices=precomputed_indices, ) return probs.squeeze(1), top_indices.squeeze(1) - def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): + def forward( + self, + input: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): """Simplified forward pass for inference - returns dense tensors only. Args: input (torch.Tensor): Input tensor of shape [seq_length, bsz, hidden_size]. padding_mask (torch.Tensor, optional): Not used in inference. + input_ids (torch.Tensor, optional): Not used in inference. Returns: Tuple[torch.Tensor, torch.Tensor]: @@ -997,6 +1056,6 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No """ if not InferenceMode.is_active(): - return super().forward(input, padding_mask) + return super().forward(input, padding_mask, input_ids, packed_seq_params) return self._forward(input, padding_mask) diff --git a/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): From e069d17db817d85ba2c4ac108edd02991114fea0 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 16 Aug 2026 01:31:15 +0800 Subject: [PATCH 10/34] Add Ling-V3 Tiny functional regression Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 7ff158e789757d4c963dfc1287071af327a2fdda) --- .../get_test_results_from_tensorboard_logs.py | 21 ++- .../test_pretraining_regular_pipeline.py | 4 +- .../model_config.yaml | 122 ++++++++++++++++++ tests/test_utils/recipes/h100/mamba.yaml | 18 +-- 4 files changed, 143 insertions(+), 22 deletions(-) create mode 100644 tests/functional_tests/test_cases/hybrid/hybrid_ling_v3_tiny_tp1_pp1_ep8_cp2_dgx_h100_1N8G/model_config.yaml diff --git a/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py b/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py index 2dcf06c6b57..9eb573e63de 100644 --- a/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py +++ b/tests/functional_tests/python_test_utils/get_test_results_from_tensorboard_logs.py @@ -45,27 +45,26 @@ def collect_train_test_metrics( if golden_value_key in [ "iteration-time", + "grad-norm", "mem-allocated-bytes", "mem-max-allocated-bytes", "lm loss", "num-zeros", "mtp_1 loss", - "mtp_2 loss", - "total loss", + "z_loss", ] } if output_path is not None: - serialized_summaries = json.dumps( - { - golden_value_key: golden_values.model_dump() - for golden_value_key, golden_values in summaries.items() - }, - indent=4, - allow_nan=False, - ) with open(output_path, "w") as fh: - fh.write(serialized_summaries) + json.dump( + { + golden_value_key: golden_values.model_dump() + for golden_value_key, golden_values in summaries.items() + }, + fh, + indent=4, + ) if __name__ == "__main__": diff --git a/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py b/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py index 832558adaad..a39f7ebb92e 100644 --- a/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py +++ b/tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py @@ -13,15 +13,15 @@ CHECK_THRESHOLDS = { "iteration-time": [common.ApproximateTest(atol=0, rtol=0.05)], + "grad-norm": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], "mem-allocated-bytes": [common.ApproximateTest(atol=0, rtol=0.05)], "mem-max-allocated-bytes": [common.ApproximateTest(atol=0, rtol=0.05)], "lm loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], "mtp_1 loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], - "mtp_2 loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], + "z_loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], "num-zeros": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.20)], "generated_tokens": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], "logprobs": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], - "total loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)], } diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_ling_v3_tiny_tp1_pp1_ep8_cp2_dgx_h100_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_ling_v3_tiny_tp1_pp1_ep8_cp2_dgx_h100_1N8G/model_config.yaml new file mode 100644 index 00000000000..963a5742943 --- /dev/null +++ b/tests/functional_tests/test_cases/hybrid/hybrid_ling_v3_tiny_tp1_pp1_ep8_cp2_dgx_h100_1N8G/model_config.yaml @@ -0,0 +1,122 @@ +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 + +MODEL_ARGS: + # Reduced Ling-V3 Tiny stack: 3 KDA, 1 MLA, 1 dense MLP, 3 MoE, and 1 MTP depth. + --spec: "[examples.ling_v3_tiny.model_spec 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 preserve Tiny's attention and FFN ratios. + --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 Tiny's 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 + --cp-comm-type: p2p + --distributed-backend: nccl + + # Ling-V3 Tiny MoE routing, reduced from 128 to 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 + + # 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} + --timing-log-level: 0 + +TEST_TYPE: regular +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..4a98172c499 100644 --- a/tests/test_utils/recipes/h100/mamba.yaml +++ b/tests/test_utils/recipes/h100/mamba.yaml @@ -13,9 +13,7 @@ spec: script_setup: | set -euo pipefail unset https_proxy - umask 077 - printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" - trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT + echo "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" | tee -a /root/.netrc # Checkout latest cd /opt @@ -36,11 +34,7 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ - rm -f "$HOME/.netrc" - unset RO_API_TOKEN script: |- - rm -f "${{HOME}}/.netrc" - unset RO_API_TOKEN set -euo pipefail ls cd /opt/megatron-lm @@ -87,6 +81,12 @@ products: # - environment: [lts] # disabled until triton is bumped # scope: [nightly] + - test_case: [hybrid_mr_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp] + products: + - environment: [dev] + scope: [mr, mr-github, mr-github-slim] + platforms: [dgx_h100] + - test_case: [hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G] products: - environment: [dev] @@ -103,9 +103,9 @@ products: # - environment: [lts] # disabled until triton is bumped # scope: [nightly] - - test_case: [hybrid_nemotron_v3_pico_7b_a1b_tp1_ep8_QAD_dgx_h100_1N8G] + - test_case: [hybrid_ling_v3_tiny_tp1_pp1_ep8_cp2_dgx_h100_1N8G] products: - # Disabled while the deterministic total-loss mismatch is investigated. + # Enable after the official H100 functional run produces golden values. - environment: [dev] scope: [mr-broken, mr-github-broken] platforms: [dgx_h100] From e239b61a13fef8166fbce1c7259055492d03b2e3 Mon Sep 17 00:00:00 2001 From: Fan Kun <784819644@qq.com> Date: Thu, 20 Aug 2026 15:46:39 +0800 Subject: [PATCH 11/34] Apply suggestions from code review Co-authored-by: Yuzhong Wang (cherry picked from commit 529e5e44a54a3a4b7c15313c7a8183a92f4d3ebd) Signed-off-by: svcnemo-autobot --- megatron/core/transformer/multi_latent_attention.py | 3 ++- megatron/core/transformer/transformer_config.py | 9 +++++++++ tests/test_utils/recipes/h100/mamba.yaml | 2 +- 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 9880fd781ea..23d6272024b 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -529,6 +529,7 @@ def _project_and_apply_mla_output_gate( return self._apply_mla_output_gate(core_attn_out, gate) @staticmethod + @jit_fuser def _apply_mla_output_gate(core_attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: """Apply one native-dtype sigmoid gate per local MLA attention head.""" output_shape = core_attn_out.shape @@ -536,7 +537,7 @@ def _apply_mla_output_gate(core_attn_out: torch.Tensor, gate: torch.Tensor) -> t # Preserve the Tiny/native BF16 autograd contract. Casting to FP32 before # sigmoid can produce the same BF16 forward values while changing the # gate-input and gate-weight VJP. - gate = torch.sigmoid(gate) + 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) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 33d7408dfb7..c1c9caaca12 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -2230,6 +2230,15 @@ def __post_init__(self): "multi_latent_attention." ) + if ( + "gdn_norm_out" in self.recompute_modules + and self.experimental_attention_variant not in ["gdn", "kda"] + ): + raise ValueError( + "gdn_norm_out in recompute_modules is only supported with " + f"GDN-family layers, but got {self.experimental_attention_variant=}." + ) + if ( "gdn" in self.recompute_modules and self.experimental_attention_variant != "gated_delta_net" diff --git a/tests/test_utils/recipes/h100/mamba.yaml b/tests/test_utils/recipes/h100/mamba.yaml index 4a98172c499..e3f85dcbe5d 100644 --- a/tests/test_utils/recipes/h100/mamba.yaml +++ b/tests/test_utils/recipes/h100/mamba.yaml @@ -107,5 +107,5 @@ products: products: # Enable after the official H100 functional run produces golden values. - environment: [dev] - scope: [mr-broken, mr-github-broken] + scope: [mr, mr-github, mr-github-slim] platforms: [dgx_h100] From 8c734278622d912eb74a5d476774a7a5c4780416 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Thu, 20 Aug 2026 19:26:59 +0800 Subject: [PATCH 12/34] Rename Ling-V3 Tiny functional test Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit c3bf41458267d36e308d6f16d66c1dfc142f0ea8) --- .../model_config.yaml | 0 tests/test_utils/recipes/h100/mamba.yaml | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename tests/functional_tests/test_cases/hybrid/{hybrid_ling_v3_tiny_tp1_pp1_ep8_cp2_dgx_h100_1N8G => hybrid_mr_mcore_te_tp1_pp1_ep8_cp2_kda_gated_mla_1N8G}/model_config.yaml (100%) diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_ling_v3_tiny_tp1_pp1_ep8_cp2_dgx_h100_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 similarity index 100% rename from tests/functional_tests/test_cases/hybrid/hybrid_ling_v3_tiny_tp1_pp1_ep8_cp2_dgx_h100_1N8G/model_config.yaml rename to tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp1_ep8_cp2_kda_gated_mla_1N8G/model_config.yaml diff --git a/tests/test_utils/recipes/h100/mamba.yaml b/tests/test_utils/recipes/h100/mamba.yaml index e3f85dcbe5d..fcfe1c937c8 100644 --- a/tests/test_utils/recipes/h100/mamba.yaml +++ b/tests/test_utils/recipes/h100/mamba.yaml @@ -103,7 +103,7 @@ products: # - environment: [lts] # disabled until triton is bumped # scope: [nightly] - - test_case: [hybrid_ling_v3_tiny_tp1_pp1_ep8_cp2_dgx_h100_1N8G] + - 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] From 344b2702f3fd0f4dc3fb2760eb032227248ffc07 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Thu, 20 Aug 2026 19:27:56 +0800 Subject: [PATCH 13/34] Fix MLA output gate compilation and tests Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit f08a9895ddd989f98c673785c2f5cc1eb814462e) --- .../core/transformer/multi_latent_attention.py | 8 ++++---- .../transformer/test_multi_latent_attention.py | 16 ++++++++++------ 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 23d6272024b..8711cfe6af2 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -19,6 +19,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, @@ -531,12 +532,11 @@ def _project_and_apply_mla_output_gate( @staticmethod @jit_fuser def _apply_mla_output_gate(core_attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - """Apply one native-dtype sigmoid gate per local MLA attention head.""" + """Apply an FP32 sigmoid gate per local MLA attention head.""" output_shape = core_attn_out.shape core_attn_out = core_attn_out.view(*output_shape[:2], gate.size(-1), -1) - # Preserve the Tiny/native BF16 autograd contract. Casting to FP32 before - # sigmoid can produce the same BF16 forward values while changing the - # gate-input and gate-weight VJP. + # 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) diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index 4c2327731aa..c9690a5e08d 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/tests/unit_tests/transformer/test_multi_latent_attention.py @@ -786,13 +786,15 @@ def test_gate_is_applied_per_head(self): gate = torch.tensor([[[-2.0, -1.0, 0.0, 1.0]], [[2.0, 3.0, 4.0, 5.0]]]).bfloat16() output = self.parallel_attention._apply_mla_output_gate(core_attn_out, gate) - expected = core_attn_out.view(2, 1, 4, 3) * torch.sigmoid(gate).unsqueeze(-1) + expected = core_attn_out.view(2, 1, 4, 3) * torch.sigmoid(gate.float()).to( + core_attn_out.dtype + ).unsqueeze(-1) torch.testing.assert_close(output, expected.reshape_as(output), atol=0, rtol=0) - def test_gate_preserves_native_dtype_vjp(self): + def test_gate_uses_fp32_sigmoid_vjp(self): if not torch.cuda.is_available(): - pytest.skip("native BF16 VJP check requires CUDA") + pytest.skip("FP32 sigmoid VJP check requires CUDA") generator = torch.Generator(device="cuda").manual_seed(20260814) gate_input = torch.randn( @@ -815,15 +817,17 @@ def test_gate_preserves_native_dtype_vjp(self): reference_gate = gate_input.detach().clone().requires_grad_(True) reference_core = core_output.detach().clone().requires_grad_(True) reference_output = ( - reference_core.view(8, 1, 4, 8) * torch.sigmoid(reference_gate).unsqueeze(-1) + reference_core.view(8, 1, 4, 8) + * torch.sigmoid(reference_gate.float()).to(reference_core.dtype).unsqueeze(-1) ).reshape_as(reference_core) reference_gradients = torch.autograd.grad( reference_output, (reference_gate, reference_core), output_gradient ) - torch.testing.assert_close(actual_output, reference_output, atol=0, rtol=0) + # Allow one BF16 quantization step from the compiled sigmoid/cast path. + torch.testing.assert_close(actual_output, reference_output, atol=1e-2, rtol=1e-2) for actual, reference in zip(actual_gradients, reference_gradients): - torch.testing.assert_close(actual, reference, atol=0, rtol=0) + torch.testing.assert_close(actual, reference, atol=1e-2, rtol=1e-2) def test_missing_gate_spec_raises(self): submodules = get_mla_self_attn_submodules() From 40f04231dc2837b722aba8edb09aa945778a8c7a Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Thu, 20 Aug 2026 20:09:02 +0800 Subject: [PATCH 14/34] Validate GDN norm recompute for hybrid models Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 545c76184fd7e1452e31f2ec785094a443f0a433) --- megatron/core/transformer/transformer_config.py | 3 ++- tests/unit_tests/ssm/gated_delta_net/test_gdn.py | 11 +++++++++++ tests/unit_tests/ssm/test_kimi_delta_attention.py | 8 ++++++-- 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index c1c9caaca12..2d94fc8e478 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -2232,7 +2232,8 @@ def __post_init__(self): if ( "gdn_norm_out" in self.recompute_modules - and self.experimental_attention_variant not in ["gdn", "kda"] + and not self.is_hybrid_model + and self.experimental_attention_variant != "gated_delta_net" ): raise ValueError( "gdn_norm_out in recompute_modules is only supported with " 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 ec53ca490d6..6ded0dfae8f 100644 --- a/tests/unit_tests/ssm/gated_delta_net/test_gdn.py +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py @@ -102,6 +102,16 @@ def test_gdn_norm_out_recompute_accepts_gdn_variant(): 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( @@ -114,6 +124,7 @@ def test_gdn_norm_out_recompute_accepts_non_experimental_hybrid_config(): # 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"], diff --git a/tests/unit_tests/ssm/test_kimi_delta_attention.py b/tests/unit_tests/ssm/test_kimi_delta_attention.py index 5e1bc44c9b9..9cc9de3dc0b 100644 --- a/tests/unit_tests/ssm/test_kimi_delta_attention.py +++ b/tests/unit_tests/ssm/test_kimi_delta_attention.py @@ -70,9 +70,13 @@ def _make_config( ) -def test_kda_norm_out_recompute_config_accepts_gdn_family_selector(): +def test_kda_norm_out_recompute_config_accepts_hybrid_model(): config = replace( - _make_config(), recompute_granularity="selective", recompute_modules=["gdn_norm_out"] + _make_config(), + experimental_attention_variant=None, + is_hybrid_model=True, + recompute_granularity="selective", + recompute_modules=["gdn_norm_out"], ) assert config.recompute_modules == ["gdn_norm_out"] From f6eba65b6035169f35ed081abca88eb91fd5acd8 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Thu, 20 Aug 2026 20:37:51 +0800 Subject: [PATCH 15/34] Tighten MLA output gate VJP tolerance Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 27f29bf2f362c68e62761c6ccf432e090e23d6db) --- .../unit_tests/transformer/test_multi_latent_attention.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index c9690a5e08d..d5e9f2fd7b8 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/tests/unit_tests/transformer/test_multi_latent_attention.py @@ -824,10 +824,11 @@ def test_gate_uses_fp32_sigmoid_vjp(self): reference_output, (reference_gate, reference_core), output_gradient ) - # Allow one BF16 quantization step from the compiled sigmoid/cast path. - torch.testing.assert_close(actual_output, reference_output, atol=1e-2, rtol=1e-2) + # Allow exactly one BF16 quantization step from the compiled sigmoid/cast path. + bf16_eps = torch.finfo(torch.bfloat16).eps + torch.testing.assert_close(actual_output, reference_output, atol=bf16_eps, rtol=0.0) for actual, reference in zip(actual_gradients, reference_gradients): - torch.testing.assert_close(actual, reference, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(actual, reference, atol=bf16_eps, rtol=0.0) def test_missing_gate_spec_raises(self): submodules = get_mla_self_attn_submodules() From 8ebbc564288b9e5d73b4f6ac3213ff230ff85b07 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Fri, 21 Aug 2026 11:17:56 +0800 Subject: [PATCH 16/34] Validate KDA safe-gate bounds Reject missing or out-of-range lower bounds during KDA construction when safe gate is enabled, matching the FLA kernel contract. Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 22458cd25babee85ed279809641ffde5b7078e76) --- megatron/core/ssm/gated_delta_net/kda.py | 11 ++++++++ .../ssm/test_kimi_delta_attention.py | 27 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/megatron/core/ssm/gated_delta_net/kda.py b/megatron/core/ssm/gated_delta_net/kda.py index d278ce003c3..ada0fcbcf15 100644 --- a/megatron/core/ssm/gated_delta_net/kda.py +++ b/megatron/core/ssm/gated_delta_net/kda.py @@ -121,6 +121,17 @@ def __init__( def _validate_config(config: TransformerConfig) -> None: """Validate the direct-projection KDA layout supported by this implementation.""" + if config.kda_safe_gate: + if config.kda_lower_bound is None: + raise ValueError( + "KimiDeltaAttention requires kda_lower_bound when kda_safe_gate=True." + ) + if not (-5.0 <= config.kda_lower_bound < 0.0): + raise ValueError( + "KimiDeltaAttention requires kda_lower_bound to be in [-5, 0) " + "when kda_safe_gate=True." + ) + if config.linear_num_key_heads != config.linear_num_value_heads: raise ValueError( "KimiDeltaAttention currently requires equal key and value head counts." diff --git a/tests/unit_tests/ssm/test_kimi_delta_attention.py b/tests/unit_tests/ssm/test_kimi_delta_attention.py index 9cc9de3dc0b..82c5244e9db 100644 --- a/tests/unit_tests/ssm/test_kimi_delta_attention.py +++ b/tests/unit_tests/ssm/test_kimi_delta_attention.py @@ -113,6 +113,33 @@ def test_kda_rejects_unsupported_head_layouts(): KimiDeltaAttention._validate_config(config) +@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): + config = _make_config() + config.kda_safe_gate = True + config.kda_lower_bound = lower_bound + + with pytest.raises(ValueError, match="kda_lower_bound"): + KimiDeltaAttention._validate_config(config) + + +@pytest.mark.parametrize("lower_bound", [-5.0, -1.0, -1e-6]) +def test_kda_accepts_valid_safe_gate_lower_bound(lower_bound): + config = _make_config() + config.kda_safe_gate = True + config.kda_lower_bound = lower_bound + + KimiDeltaAttention._validate_config(config) + + +def test_kda_ignores_lower_bound_when_safe_gate_is_disabled(): + config = _make_config() + config.kda_safe_gate = False + config.kda_lower_bound = None + + KimiDeltaAttention._validate_config(config) + + def test_kda_rejects_invalid_packed_boundaries(): q = torch.tensor([0, 8, 16], dtype=torch.int32) KimiDeltaAttention._validate_packed_cu_seqlens(q, q.clone()) From f709ee9897e030f5b76d3f25266d2df170d14b1c Mon Sep 17 00:00:00 2001 From: Yuzhong Wang Date: Fri, 21 Aug 2026 01:50:50 -0700 Subject: [PATCH 17/34] refactor: Align KDA structure with GatedDeltaNet Signed-off-by: Yuzhong Wang (cherry picked from commit 4b03bfb69173c787418cd0bbcf25c7459b3649a1) --- .../ling_v3_tiny/train_ling_v3_tiny_8gpu.sh | 1 + ...rimental_attention_variant_module_specs.py | 151 ++++- megatron/core/ssm/gated_delta_net/common.py | 4 +- megatron/core/ssm/gated_delta_net/kda.py | 639 ++++++++++++------ .../core/transformer/transformer_config.py | 161 +++-- megatron/training/argument_utils.py | 237 +++++-- megatron/training/arguments.py | 5 +- megatron/training/training.py | 13 +- .../model_config.yaml | 1 + ...rimental_attention_variant_module_specs.py | 63 +- .../ssm/gated_delta_net/test_gdn.py | 4 +- .../test_kda.py} | 91 ++- 12 files changed, 950 insertions(+), 420 deletions(-) rename tests/unit_tests/ssm/{test_kimi_delta_attention.py => gated_delta_net/test_kda.py} (85%) diff --git a/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh b/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh index 12fda0c7296..1712f908170 100755 --- a/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh +++ b/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh @@ -58,6 +58,7 @@ export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" --pipeline-model-parallel-size 1 \ --expert-model-parallel-size 8 \ --context-parallel-size 2 \ + --linear-cp-mode headwise \ --cp-comm-type p2p \ --num-experts 128 \ --moe-ffn-hidden-size 512 \ 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..289c1128627 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -1,16 +1,33 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import warnings from typing import List, Optional from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.models.backends import BackendSpecProvider -from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNet2, GatedDeltaNetSubmodules +from megatron.core.ssm.gated_delta_net import ( + GatedDeltaNet, + GatedDeltaNetSubmodules, + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, +) from megatron.core.transformer.enums import AttnMaskType, LayerType from megatron.core.transformer.experimental_attention_variant.absorbed_mla import ( AbsorbedMLASelfAttention, AbsorbedMLASelfAttentionSubmodules, ) +from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, +) +from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + DSv4HybridSelfAttentionSubmodules, +) from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexer, DSAIndexerSubmodules, @@ -19,6 +36,7 @@ is_dsa_skip_topk_layer, source_dsa_compute_layer, ) +from megatron.core.transformer.hyper_connection import HyperConnectionModule from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_block import ( @@ -27,6 +45,7 @@ ) from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, MlpBuilder, TransformerLayer, TransformerLayerSubmodules, @@ -58,7 +77,7 @@ ########## # Canonical ``experimental_attention_variant`` names served by the gated delta net family. -GDN_ATTENTION_VARIANTS = ("gdn", "gdn2") +GDN_ATTENTION_VARIANTS = ("gdn", "kda") # Deprecated ``experimental_attention_variant`` spellings mapped to their canonical name. _DEPRECATED_ATTENTION_VARIANT_ALIASES = {"gated_delta_net": "gdn"} @@ -72,25 +91,33 @@ def get_gated_delta_net_module_spec( config: TransformerConfig, backend: BackendSpecProvider = None ) -> ModuleSpec: - """Build module spec for GatedDeltaNet attention.""" + """Build a module spec for a GDN-family attention variant.""" if backend is None: backend = _get_backend_spec_provider(config=config) rms_norm = config.normalization == "RMSNorm" - # gdn2 reuses the GDN submodules and spec structure with the GatedDeltaNet2 module. - gdn_module = ( - GatedDeltaNet2 if config.experimental_attention_variant == "gdn2" else GatedDeltaNet - ) - attention = ModuleSpec( - module=gdn_module, - submodules=GatedDeltaNetSubmodules( - in_proj=backend.column_parallel_layer_norm_linear(), - out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False), - out_proj=backend.row_parallel_linear(), - ), - metainfo={"fuse_input_layernorm": True}, - ) + if config.experimental_attention_variant == "kda": + attention = ModuleSpec( + module=KimiDeltaAttention, + submodules=KimiDeltaAttentionSubmodules( + in_proj=backend.column_parallel_linear(), + beta_proj=backend.column_parallel_linear(), + out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False), + out_proj=backend.row_parallel_linear(), + ), + metainfo={"fuse_input_layernorm": False}, + ) + else: + attention = ModuleSpec( + module=GatedDeltaNet, + submodules=GatedDeltaNetSubmodules( + in_proj=backend.column_parallel_layer_norm_linear(), + out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False), + out_proj=backend.row_parallel_linear(), + ), + metainfo={"fuse_input_layernorm": True}, + ) return attention @@ -146,6 +173,63 @@ def get_dsa_module_spec_for_backend( return attention +def get_dsv4_hybrid_module_spec_for_backend( + config: TransformerConfig, backend: BackendSpecProvider = None +) -> ModuleSpec: + """Helper function to get module spec for DSv4 Hybrid Sparse Attention.""" + assert config.multi_latent_attention, "Currently only MLA supports sparse attention." + assert config.qk_l2_norm is False, "qk_l2_norm is not supported with MLA." + + # Adjust for RMS norm. + rms_norm = config.normalization == "RMSNorm" + # DSA indexer requires normalized q as input, so here we cannot fuse qk layernorm + # with linear projection and have to use unfused qk layernorm. + qk_norm = ( + backend.layer_norm(rms_norm=rms_norm, for_qk=True) if config.qk_layernorm else IdentityOp + ) + + compressor_spec = ModuleSpec( + module=Compressor, + submodules=CompressorSubmodules( + linear_wkv=backend.linear(), + linear_wgate=backend.linear(), + norm=backend.layer_norm(rms_norm=True, for_qk=False), + ), + ) + + indexer_spec = ModuleSpec( + module=CSAIndexer, + submodules=CSAIndexerSubmodules( + linear_wq_b=backend.linear(), + linear_weights_proj=backend.linear(), + compressor=compressor_spec, + ), + ) + + core_attention = ModuleSpec( + module=CompressedSparseAttention, + submodules=CompressedSparseAttentionSubmodules( + compressor=compressor_spec, indexer=indexer_spec + ), + ) + + attention = ModuleSpec( + module=DSv4HybridSelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=DSv4HybridSelfAttentionSubmodules( + linear_q_down_proj=backend.linear(), + linear_q_up_proj=backend.column_parallel_linear(), + linear_kv_proj=backend.column_parallel_linear(), + core_attention=core_attention, + linear_proj=backend.row_parallel_linear(), + q_layernorm=qk_norm, + kv_layernorm=qk_norm, + ), + metainfo={"fuse_input_layernorm": False}, + ) + return attention + + def get_experimental_attention_variant_module_spec( config: TransformerConfig, backend: BackendSpecProvider = None ) -> ModuleSpec: @@ -158,6 +242,8 @@ def get_experimental_attention_variant_module_spec( return get_gated_delta_net_module_spec(config=config, backend=backend) elif config.experimental_attention_variant == "dsa": return get_dsa_module_spec_for_backend(config=config, backend=backend) + elif config.experimental_attention_variant == "dsv4_hybrid": + return get_dsv4_hybrid_module_spec_for_backend(config=config, backend=backend) else: raise ValueError( f"Invalid experimental attention variant: {config.experimental_attention_variant}" @@ -245,6 +331,10 @@ def get_transformer_layer_with_experimental_attention_variant_spec( # Get GPT decoder block layer specs rms_norm = config.normalization == "RMSNorm" + enable_hc = config.enable_hyper_connections + hc_module = HyperConnectionModule if enable_hc else IdentityOp + layer_module = HyperConnectionTransformerLayer if enable_hc else TransformerLayer + layer_specs = [] for layer_number in range(config.num_layers): attention = ( @@ -271,14 +361,16 @@ def get_transformer_layer_with_experimental_attention_variant_spec( layer_specs.append( ModuleSpec( - module=TransformerLayer, + module=layer_module, submodules=TransformerLayerSubmodules( input_layernorm=input_layernorm, self_attention=attention, self_attn_bda=get_bias_dropout_add, + self_attention_hyper_connection=hc_module, pre_mlp_layernorm=pre_mlp_layernorm, mlp=not_none(mlp), mlp_bda=get_bias_dropout_add, + mlp_hyper_connection=hc_module, ), ) ) @@ -349,19 +441,7 @@ def get_transformer_block_with_experimental_attention_variant_spec( def normalize_experimental_attention_variant( experimental_attention_variant: Optional[str], ) -> Optional[str]: - """Resolve a deprecated ``experimental_attention_variant`` spelling to its canonical name. - - ``gated_delta_net`` is the deprecated spelling of ``gdn``. Passing it emits a - ``DeprecationWarning`` and returns the canonical name so that every downstream - consumer only has to handle ``gdn``. - - Args: - experimental_attention_variant: The configured variant name, possibly a - deprecated alias. - - Returns: - The canonical variant name, or the argument unchanged when it is not an alias. - """ + """Resolve a deprecated attention-variant spelling to its canonical name.""" canonical = _DEPRECATED_ATTENTION_VARIANT_ALIASES.get(experimental_attention_variant) if canonical is None: return experimental_attention_variant @@ -376,11 +456,7 @@ def normalize_experimental_attention_variant( def is_gated_delta_net_variant(experimental_attention_variant: Optional[str]) -> bool: - """Check if the experimental attention variant is served by a gated delta net layer. - - Accepts the deprecated ``gated_delta_net`` spelling without warning; use - :func:`normalize_experimental_attention_variant` to emit the deprecation notice. - """ + """Return whether a name selects a GDN-family attention implementation.""" canonical = _DEPRECATED_ATTENTION_VARIANT_ALIASES.get( experimental_attention_variant, experimental_attention_variant ) @@ -501,7 +577,7 @@ def _get_backend_spec_provider(config: TransformerConfig) -> BackendSpecProvider ) backend: BackendSpecProvider = ( KitchenSpecProvider( - fallback=TESpecProvider(), + fallback=TESpecProvider(fallback_to_eager_attn=config.fallback_to_eager_attn), use_kitchen_attention=config.use_kitchen_attention, kitchen_attention_backend=config.kitchen_attention_backend, ) @@ -537,6 +613,7 @@ def _get_self_attention_module_spec( qk_l2_norm=config.qk_l2_norm, use_kitchen=config.use_kitchen, use_te_activation_func=config.use_te_activation_func, + fallback_to_eager_attn=config.fallback_to_eager_attn, use_kitchen_attention=config.use_kitchen_attention, kitchen_attention_backend=config.kitchen_attention_backend, mla_down_proj_fusion=getattr(config, "mla_down_proj_fusion", False), diff --git a/megatron/core/ssm/gated_delta_net/common.py b/megatron/core/ssm/gated_delta_net/common.py index 30e0cfa5d9b..0483fec8fe4 100644 --- a/megatron/core/ssm/gated_delta_net/common.py +++ b/megatron/core/ssm/gated_delta_net/common.py @@ -362,9 +362,7 @@ def reset_parameters(self): nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init) self._reset_dt_bias() A = torch.empty( - self.A_log.shape[0], - dtype=self.config.params_dtype, - device=torch.cuda.current_device(), + self.A_log.shape[0], dtype=self.A_log.dtype, device=self.A_log.device ).uniform_(*self.A_init_range) self.A_log.data.copy_(torch.log(A)) diff --git a/megatron/core/ssm/gated_delta_net/kda.py b/megatron/core/ssm/gated_delta_net/kda.py index ada0fcbcf15..0019c406e5c 100644 --- a/megatron/core/ssm/gated_delta_net/kda.py +++ b/megatron/core/ssm/gated_delta_net/kda.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh. """Kimi Delta Attention, a channel-wise Gated DeltaNet variant.""" @@ -6,14 +6,16 @@ import math from dataclasses import dataclass from functools import partial +from typing import Optional import torch import torch.nn.functional as F from megatron.core import tensor_parallel +from megatron.core.context_parallel_layout import convert_module_input_tensors_cp_partition_mode from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.jit import jit_fuser -from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gated_delta_net.common import ( HAVE_FLA, @@ -21,9 +23,9 @@ _GDNBase, a2a_cp_to_hp, a2a_hp_to_cp, + build_cp_context, causal_conv1d, get_parameter_local_cp, - l2norm, ) from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import mark_keep_in_fp32 @@ -70,8 +72,9 @@ def __init__( *, name: str | None = None, cp_comm_type: str | None = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, ) -> None: - self._validate_config(config) 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." @@ -89,8 +92,12 @@ def __init__( 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, @@ -117,40 +124,16 @@ def __init__( # Pass raw g to FLA and let the KDA kernel apply A_log and dt_bias. self.use_gate_in_kernel = True - @staticmethod - def _validate_config(config: TransformerConfig) -> None: - """Validate the direct-projection KDA layout supported by this implementation.""" - - if config.kda_safe_gate: - if config.kda_lower_bound is None: - raise ValueError( - "KimiDeltaAttention requires kda_lower_bound when kda_safe_gate=True." - ) - if not (-5.0 <= config.kda_lower_bound < 0.0): - raise ValueError( - "KimiDeltaAttention requires kda_lower_bound to be in [-5, 0) " - "when kda_safe_gate=True." - ) - - if config.linear_num_key_heads != config.linear_num_value_heads: - raise ValueError( - "KimiDeltaAttention currently requires equal key and value head counts." - ) - if config.linear_key_head_dim != config.linear_value_head_dim: - raise ValueError( - "KimiDeltaAttention currently requires equal key and value head dimensions." - ) - tp_cp_size = config.tensor_model_parallel_size * config.context_parallel_size - if config.linear_num_key_heads % tp_cp_size != 0: - raise ValueError( - "KimiDeltaAttention key heads must be divisible by tensor parallel size times " - "context parallel size." - ) - def _setup_variant_attrs(self) -> None: - """Set KDA dimensions and projection checkpoint metadata.""" + """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, @@ -159,38 +142,48 @@ def _setup_variant_attrs(self) -> None: self.qk_dim_local_tp, self.v_dim_local_tp, ) - self.feat_dim_split = ( - (2 * self.qk_dim_local_tp + self.v_dim_local_tp) // self.cp_size, - self.qk_dim_local_tp // self.cp_size, - self.v_dim_local_tp // 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 reset_parameters(self) -> None: - """Initialize convolution and KDA gate parameters.""" + 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.""" - if not self.config.perform_initialization: - return + 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, + ) - with tensor_parallel.get_cuda_rng_tracker().fork(): - if self.conv_init is not None: - torch.nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init) + def _reset_dt_bias(self) -> None: + """Initialize the KDA channel-wise step-size bias in inverse-softplus space.""" - dt_min, dt_max, dt_init_floor = 0.001, 0.1, 1e-4 - dt = torch.exp( - torch.rand(self.dt_bias_dim, device=self.dt_bias.device, dtype=torch.float32) - * (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))) + dt_min, dt_max, dt_init_floor = 0.001, 0.1, 1e-4 + dt = torch.exp( + torch.rand(self.dt_bias_dim, device=self.dt_bias.device, dtype=self.dt_bias.dtype) + * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min) + ).clamp(min=dt_init_floor) + self.dt_bias.data.copy_(dt + torch.log(-torch.expm1(-dt))) - A = torch.empty(self.a_log_dim, device=self.A_log.device, dtype=torch.float32).uniform_( - *self.A_init_range - ) - self.A_log.data.copy_(torch.log(A)) + @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: @@ -202,155 +195,352 @@ def _apply_gated_norm(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor gate = gate.reshape(-1, self.value_head_dim) return (x * torch.sigmoid(gate.float())).to(x_dtype) - def _prepare_kda_inputs( - self, - qkv: torch.Tensor, - raw_g: torch.Tensor, - beta: torch.Tensor, - output_gate: torch.Tensor, - batch: int, - seq_len: int, - ) -> tuple[torch.Tensor, ...]: - """Split, normalize, and reshape tensors for the FLA KDA kernel.""" - - 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, - ) - query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) - value = value.reshape(batch, seq_len, -1, self.value_head_dim) - if self.use_qk_l2norm: - query_key = l2norm(query_key.contiguous()) - - num_key_heads = self.num_k_heads_local_tp // self.cp_size - query, key = torch.split(query_key, [num_key_heads, num_key_heads], dim=2) - 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() - output_gate = output_gate.reshape( - batch, seq_len, self.num_v_heads_local_tp // self.cp_size, self.value_head_dim - ) - return ( - query.contiguous(), - key.contiguous(), - value.contiguous(), - raw_g.contiguous(), - beta.contiguous(), - output_gate.contiguous(), - ) - - @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 _gated_norm_and_a2a( - self, - core_attn_out: torch.Tensor, - gate: torch.Tensor, - thd_cp_a2a_inv: torch.Tensor | None, - batch: int, - seq_len: int, - packed_seq_params: PackedSeqParams | None = None, - ) -> torch.Tensor: - """Apply KDA's gated output norm and restore its context-parallel layout.""" - - nvtx_range_push(suffix="gated_norm") - norm_out_hp = self._apply_gated_norm(core_attn_out, gate) - nvtx_range_pop(suffix="gated_norm") - - norm_out_hp = norm_out_hp.reshape(batch, seq_len, -1) - norm_out_hp = norm_out_hp.transpose(0, 1).contiguous() - - return a2a_hp_to_cp( - norm_out_hp, self.cp_size, self.pg_collection.cp, packed_seq_params, thd_cp_a2a_inv - ) - def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, - inference_context: BaseInferenceContext | None = None, - packed_seq_params: PackedSeqParams | None = None, - sequence_len_offset: int | None = None, + inference_context: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[int] = None, *, - inference_params: BaseInferenceContext | None = 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) - if inference_context is not None: # pragma: no cover - raise NotImplementedError("KDA recurrent inference is not implemented.") - seq_len, batch, _ = hidden_states.shape - total_seq_len = seq_len * self.sp_size * self.cp_size + active_pg_collection = pg_collection if pg_collection is not None else self.pg_collection + base_cp_group = active_pg_collection.cp + cp_group = resolve_cp_group(base_cp_group, packed_seq_params) + if self.config.linear_cp_mode == "chunkwise": + cp_group_chunkwise = cp_group + cp_group_headwise = None + elif self.config.linear_cp_mode == "headwise": + cp_group_chunkwise = None + cp_group_headwise = cp_group + elif cp_group.size() == 1: + cp_group_chunkwise = None + cp_group_headwise = None + else: + raise ValueError( + f"Unsupported linear_cp_mode {self.config.linear_cp_mode!r}; " + "expected 'headwise' or 'chunkwise'." + ) + + cp_size_chunkwise = cp_group_chunkwise.size() if cp_group_chunkwise is not None else 1 + cp_size_headwise = cp_group_headwise.size() if cp_group_headwise is not None else 1 + cp_size_runtime = cp_group.size() + back_to_input_converter = None + if self.config.linear_cp_mode == "chunkwise": + hidden_states, back_to_input_converter = convert_module_input_tensors_cp_partition_mode( + hidden_states=hidden_states, + packed_seq_params=packed_seq_params, + cp_group=cp_group_chunkwise, + tp_group=self.tp_group, + tp_cp_group=getattr(active_pg_collection, "tp_cp", None), + target_partition_mode="contiguous", + sequence_parallel=self.config.sequence_parallel, + config=self.config, + ) + + 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 packed_seq_params.cp_partition_mode != "zigzag" + ) + or ( + (packed_seq_params is None or packed_seq_params.qkv_format != "thd") + and self.config.cp_partition_mode != "zigzag" + ) + ): + raise ValueError( + "KimiDeltaAttention with headwise CP requires zigzag layout. CP partition " + "conversion must be handled before calling KimiDeltaAttention." + ) - cu_seqlens_q = None 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, - total_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, - total_seq_len, + seq_len_global, "cu_seqlens_kv", - cp_size=self.cp_size, + 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 back_to_input_converter is not None: + out = back_to_input_converter.convert( + out, seq_dim=0, sequence_parallel=self.config.sequence_parallel + ) - nvtx_range_push(suffix="kda_in_proj") + 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="kda_in_proj") + nvtx_range_pop(suffix="in_proj") qkvfg, thd_cp_a2a_inv = a2a_cp_to_hp( qkvfg, self.in_proj_split_sections, - self.cp_size, - self.pg_collection.cp, + cp_size_headwise, + cp_group_headwise, cu_seqlens_q, - total_seq_len, + seq_len_post_headwise, packed_seq_params, ) beta, _ = a2a_cp_to_hp( beta, (self.num_k_heads_local_tp,), - self.cp_size, - self.pg_collection.cp, + cp_size_headwise, + cp_group_headwise, cu_seqlens_q, - total_seq_len, + 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) - seq_len = qkvfg.shape[1] - qkv, raw_g, output_gate = torch.split(qkvfg, self.feat_dim_split, dim=-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="kda_conv1d") + 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, @@ -359,14 +549,14 @@ def forward( conv1d_weight = get_parameter_local_cp( self.conv1d.weight, dim=0, - cp_group=self.pg_collection.cp, + cp_group=cp_group_headwise, split_sections=qkv_channels_split_sections, ) conv1d_bias = ( get_parameter_local_cp( self.conv1d.bias, dim=0, - cp_group=self.pg_collection.cp, + cp_group=cp_group_headwise, split_sections=qkv_channels_split_sections, ) if self.conv_bias @@ -375,79 +565,102 @@ def forward( if self.config.deterministic_mode or causal_conv1d is None: qkv = qkv.transpose(1, 2).contiguous() conv_out = F.conv1d( - qkv, - conv1d_weight, + input=qkv, + weight=conv1d_weight, bias=conv1d_bias, stride=self.conv1d.stride, padding=self.conv1d.padding, dilation=self.conv1d.dilation, - groups=self.conv_dim_local_tp // self.cp_size, + groups=self.conv_dim_local_tp // cp_size_headwise, ) - qkv = self.act_fn(conv_out[..., :seq_len]).transpose(1, 2).contiguous() + 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=qkv, + x=conv_input, weight=conv1d_weight.squeeze(1), bias=conv1d_bias, activation=self.activation, initial_state=None, output_final_state=False, - cu_seqlens=cu_seqlens_q, + cu_seqlens=conv_cu_seqlens, + cp_context=conv_cp_context, ) - nvtx_range_pop(suffix="kda_conv1d") + if pad_n > 0: + qkv = qkv[:, :orig_seq, :] + nvtx_range_pop(suffix="conv1d") - query, key, value, raw_g, beta, output_gate = self._prepare_kda_inputs( - qkv, raw_g, beta, output_gate, batch, seq_len - ) - A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp) - dt_bias_local_cp = get_parameter_local_cp( - self.dt_bias, dim=0, cp_group=self.pg_collection.cp - ) + A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group_headwise) + dt_bias_local_cp = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group_headwise) - nvtx_range_push(suffix="kda_chunk") - core_attn_out, _ = self.gated_delta_rule( - q=query, - k=key, - v=value, - g=raw_g, - beta=beta, - A_log=A_log_local_cp, - dt_bias=dt_bias_local_cp, - 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, + 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="kda_chunk") + nvtx_range_pop(suffix="prepare_input_for_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, output_gate) - else: - norm_out = self._gated_norm_and_a2a( - core_attn_out, output_gate, thd_cp_a2a_inv, batch, seq_len, packed_seq_params - ) + gate = kernel_inputs.pop("gate") - nvtx_range_push(suffix="kda_out_proj") - out, out_bias = self.out_proj(norm_out) - nvtx_range_pop(suffix="kda_out_proj") + 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, + ) - if self.recompute_norm_out: - self.norm_out_checkpoint.discard_output_and_register_recompute(out) + @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.""" - return out, out_bias + 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.""" diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 2d94fc8e478..0283df062e4 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -305,10 +305,15 @@ class TransformerConfig(ModelParallelConfig): #################### # attention variant #################### - experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = ( - None - ) - """Type of attention variant to use. Currently support gated_delta_net, dsa, and dsv4_hybrid.""" + experimental_attention_variant: Optional[ + Literal['gdn', 'kda', 'gated_delta_net', 'dsa', 'dsv4_hybrid'] + ] = None + """Type of experimental attention variant to use. + + ``gdn`` and ``kda`` select Gated DeltaNet and Kimi Delta Attention, respectively. + ``gated_delta_net`` is a deprecated compatibility alias for ``gdn`` and is normalized in + ``__post_init__`` with a ``DeprecationWarning``. + """ cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag" """How THD sequence rows are partitioned across context-parallel ranks. @@ -617,9 +622,9 @@ class TransformerConfig(ModelParallelConfig): "mhc": recompute HyperConnection intermediate activations via CheckpointWithoutOutput + MHCCheckpointManager. Requires enable_hyper_connections=True. Cannot be used with "mlp". - "gdn": recompute the entire GatedDeltaNet module (in_proj, conv1d, gated delta rule, - gated norm, CP all-to-all and out_proj). Requires - experimental_attention_variant="gated_delta_net". + "gdn": recompute the entire GDN-family layer, including GatedDeltaNet and KDA + (input projections, conv1d, gated delta rule, gated norm, CP all-to-all, and + out_proj). Requires a GDN-family experimental attention variant or a hybrid model. "gdn_norm_out": recompute gated output normalization and layout restoration for Gated DeltaNet-family layers, including GatedDeltaNet and KDA. "moe_act", "layernorm", "mla_up_proj", "mhc", and "gdn_norm_out" use @@ -1472,6 +1477,12 @@ def __post_init__(self): """ super().__post_init__() + # Imported lazily because the module-spec module imports TransformerConfig. + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + is_gated_delta_net_variant, + normalize_experimental_attention_variant, + ) + # When fp32 residual connections are enabled, pipeline parallel communication must # use fp32 to match the dtype of the residual stream between pipeline stages. if self.fp32_residual_connection and self.pipeline_dtype is not None: @@ -1567,6 +1578,11 @@ def __post_init__(self): self.experimental_attention_variant = self.linear_attention_type self.linear_attention_type = None + if self.experimental_attention_variant is not None: + self.experimental_attention_variant = normalize_experimental_attention_variant( + self.experimental_attention_variant + ) + if self.cp_partition_mode not in ("zigzag", "contiguous"): raise ValueError(f"Unsupported cp_partition_mode: {self.cp_partition_mode}") @@ -1601,19 +1617,21 @@ def __post_init__(self): "cp_partition_mode='contiguous' is not supported with " "multi_latent_attention outside dsv4_hybrid." ) - if self.experimental_attention_variant not in ("dsv4_hybrid", "gated_delta_net"): + 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 either 'dsv4_hybrid' " - "or 'gated_delta_net'." + "requires experimental_attention_variant to be 'dsv4_hybrid', 'gdn', " + "or 'kda'." ) if ( - self.experimental_attention_variant == "gated_delta_net" + is_gated_delta_net_variant(self.experimental_attention_variant) and self.linear_cp_mode == "headwise" ): raise ValueError( "cp_partition_mode='contiguous' is incompatible with " - "gated_delta_net linear_cp_mode='headwise'." + "GDN-family linear_cp_mode='headwise'." ) elif self.cp_partition_mode == "zigzag": if self.experimental_attention_variant == "dsv4_hybrid": @@ -1652,50 +1670,67 @@ def __post_init__(self): ) self.dsa_kernel_backend = legacy_backend - if self.experimental_attention_variant in ["gated_delta_net"]: - assert ( - self.linear_attention_freq is not None - ), f"linear_attention_freq must be set for linear attention." - - if self.experimental_attention_variant == "gated_delta_net": + if is_gated_delta_net_variant(self.experimental_attention_variant): + if not self.is_hybrid_model: + assert ( + self.linear_attention_freq is not None + ), "linear_attention_freq must be set for linear attention." if self.pad_packed_seq_alignment is not None: tail_policy = self.thd_tail_padding_policy or 'append_dummy_seq' assert tail_policy == 'append_dummy_seq', ( - "gated_delta_net with pad_packed_seq_alignment requires " + "GDN-family attention with pad_packed_seq_alignment requires " "thd_tail_padding_policy='append_dummy_seq'." ) - # Check required parameters - assert ( - self.linear_conv_kernel_dim is not None - ), "linear_conv_kernel_dim must be set for gated delta net." - assert ( - self.linear_key_head_dim is not None - ), "linear_key_head_dim must be set for gated delta net." - assert ( - self.linear_value_head_dim is not None - ), "linear_value_head_dim must be set for gated delta net." - assert ( - self.linear_num_key_heads is not None - ), "linear_num_key_heads must be set for gated delta net." - assert ( - self.linear_num_value_heads is not None - ), "linear_num_value_heads must be set for gated delta net." - assert self.linear_num_value_heads % self.linear_num_key_heads == 0, ( - f"linear_num_value_heads ({self.linear_num_value_heads}) must be a multiple of " - f"linear_num_key_heads ({self.linear_num_key_heads})." - ) - 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}." - ) + # 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 a GDN-family layer." + assert ( + self.linear_key_head_dim is not None + ), "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 a GDN-family layer." + assert ( + self.linear_num_key_heads is not None + ), "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 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." + ) - if self.context_parallel_size > 1: - assert self.linear_cp_mode in ("headwise", "chunkwise"), ( - f"linear_cp_mode must be one of 'headwise' or 'chunkwise', " + 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}." + ) + + 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 " @@ -1716,7 +1751,7 @@ def __post_init__(self): f"{self.linear_num_value_heads=} must be a multiple of " 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( @@ -1824,13 +1859,15 @@ def __post_init__(self): "build or set dsa_kernel_backend='none'." ) - if ( - self.gdn_pre_gated_delta_rule_fusion - and self.experimental_attention_variant != "gated_delta_net" - ): + 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='gated_delta_net'." + "experimental_attention_variant='gdn'." ) if self.fp8: @@ -2233,7 +2270,7 @@ def __post_init__(self): if ( "gdn_norm_out" in self.recompute_modules and not self.is_hybrid_model - and self.experimental_attention_variant != "gated_delta_net" + and not is_gated_delta_net_variant(self.experimental_attention_variant) ): raise ValueError( "gdn_norm_out in recompute_modules is only supported with " @@ -2242,17 +2279,18 @@ def __post_init__(self): if ( "gdn" in self.recompute_modules - and self.experimental_attention_variant != "gated_delta_net" + 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 " - "experimental_attention_variant='gated_delta_net'." + "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 GatedDeltaNet module, including gated norm." + "'gdn' recomputes the full GDN-family layer, including gated norm." ) if "core_attn" in self.recompute_modules: warnings.warn( @@ -3179,10 +3217,9 @@ def _scope_to_str(s): and (not self.cuda_graph_modules or CudaGraphModule.attn in self.cuda_graph_modules) ) - cp_layout_conversion_required = self.experimental_attention_variant == "gated_delta_net" - # TODO: Extend this predicate as GDN2/KDA are introduced, and for DSv4 when - # dsa_cp_balance_indexer is introduced; those paths will also require module-local THD CP - # layout conversion. + cp_layout_conversion_required = is_gated_delta_net_variant( + self.experimental_attention_variant + ) if ( (self.context_parallel_size > 1 or self.dynamic_context_parallel) and self.sequence_packing_scheduler is not None diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py index 124083ead83..db927573166 100644 --- a/megatron/training/argument_utils.py +++ b/megatron/training/argument_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import ast import builtins @@ -17,7 +17,7 @@ import torch.nn.functional as F from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.spec_utils import import_module +from megatron.core.transformer.spec_utils import ModuleSpec, import_module from megatron.training.config import ( CheckpointConfig, DistributedInitConfig, @@ -38,17 +38,20 @@ # TODO: support arg renames + class TypeInferenceError(Exception): """Custom exception type to be conditionally handled by ArgumentGroupFactory.""" + pass + class ArgumentGroupFactory: """Utility that adds an argument group to an ArgumentParser based on the attributes of a dataclass. This utility uses dataclass metadata including type annotations and docstrings to automatically infer the type, default, and other argparse keyword arguments. - You can override or supplement the automatically inferred argparse kwargs for any + You can override or supplement the automatically inferred argparse kwargs for any dataclass field by providing an "argparse_meta" key in the field's metadata dict. The value should be a dict of kwargs that will be passed to ArgumentParser.add_argument(). These metadata kwargs take precedence over the automatically inferred values. @@ -77,13 +80,13 @@ class YourConfig: that require some customized or additional handling. Args: - src_cfg_class: The source dataclass type (not instance) whose fields will be - converted into command-line arguments. Each field's type annotation determines - the argument type, default values become argument defaults, and field-level + src_cfg_class: The source dataclass type (not instance) whose fields will be + converted into command-line arguments. Each field's type annotation determines + the argument type, default values become argument defaults, and field-level docstrings are extracted to populate argument help text. - exclude: Optional list of attribute names from `src_cfg_class` to exclude from + exclude: Optional list of attribute names from `src_cfg_class` to exclude from argument generation. Useful for omitting internal fields, computed properties, - or attributes that should be configured through other means. If None, all + or attributes that should be configured through other means. If None, all dataclass fields will be converted to command-line arguments. Default: None. """ @@ -97,7 +100,7 @@ def _format_arg_name(self, config_attr_name: str, prefix: Optional[str] = None) Args: config_attr_name: dataclass attribute name - prefix: prefix string to add to the dataclass attribute name. e.g. 'no' for bool + prefix: prefix string to add to the dataclass attribute name. e.g. 'no' for bool settings that are default True. A hyphen is added after the prefix. Default: None """ arg_name = config_attr_name @@ -112,6 +115,7 @@ def _get_enum_kwargs(self, config_type: enum.EnumMeta) -> dict[str, Any]: With these settings, the user must provide a valid enum value, e.g. 'flash', for `AttnBackend.flash`. """ + def enum_type_handler(cli_arg): return config_type[cli_arg] @@ -135,7 +139,9 @@ def _extract_type(self, config_type: type) -> dict[str, Any]: if origin in [types.UnionType, typing.Union]: # Handle Optional and Union - if type_tuple[1] == type(None): # Optional type. First element is value inside Optional[] + if type_tuple[1] == type( + None + ): # Optional type. First element is value inside Optional[] return self._extract_type(type_tuple[0]) else: raise TypeInferenceError(f"Unions not supported by argparse: {config_type}") @@ -146,17 +152,20 @@ def _extract_type(self, config_type: type) -> dict[str, Any]: kwargs["nargs"] = "+" return kwargs else: - raise TypeInferenceError(f"Multi-type lists not supported by argparse: {config_type}") + raise TypeInferenceError( + f"Multi-type lists not supported by argparse: {config_type}" + ) elif origin is typing.Literal: choices_types = [type(choice) for choice in type_tuple] - assert all([t == choices_types[0] for t in choices_types]), "Type of each choice in a Literal type should all be the same." + assert all( + [t == choices_types[0] for t in choices_types] + ), "Type of each choice in a Literal type should all be the same." kwargs = {"type": choices_types[0], "choices": type_tuple} return kwargs else: raise TypeInferenceError(f"Unsupported type: {config_type}") - def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: """Assemble kwargs for add_argument(). @@ -166,7 +175,9 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: argparse_kwargs = {} argparse_kwargs["arg_names"] = [self._format_arg_name(attribute.name)] argparse_kwargs["dest"] = attribute.name - argparse_kwargs["help"] = self.field_docstrings[attribute.name] if attribute.name in self.field_docstrings else "" + argparse_kwargs["help"] = ( + self.field_docstrings[attribute.name] if attribute.name in self.field_docstrings else "" + ) # dataclasses specifies that both should not be set if isinstance(attribute.default, type(dataclasses.MISSING)): @@ -180,7 +191,6 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: # save metadata here, but update at the end so the metadata has highest precedence attr_argparse_meta = attribute.metadata["argparse_meta"] - # if we cannot infer the argparse type, all of this logic may fail. we try to defer # to the developer-specified metadata if present try: @@ -188,12 +198,17 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: # use store_true or store_false action for enable/disable flags, which doesn't accept a 'type' if argparse_kwargs["type"] == bool: - argparse_kwargs["action"] = "store_true" if attribute.default == False else "store_false" + argparse_kwargs["action"] = ( + "store_true" if attribute.default == False else "store_false" + ) argparse_kwargs.pop("type") # add '--no-*' and '--disable-*' prefix if this is a store_false argument if argparse_kwargs["action"] == "store_false": - argparse_kwargs["arg_names"] = [self._format_arg_name(attribute.name, prefix="no"), self._format_arg_name(attribute.name, prefix="disable")] + argparse_kwargs["arg_names"] = [ + self._format_arg_name(attribute.name, prefix="no"), + self._format_arg_name(attribute.name, prefix="disable"), + ] except TypeInferenceError as e: if attr_argparse_meta is not None: print( @@ -205,7 +220,7 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: else: raise e - # metadata provided by field takes precedence + # metadata provided by field takes precedence if attr_argparse_meta is not None: argparse_kwargs.update(attr_argparse_meta) @@ -255,8 +270,12 @@ def _get_field_docstrings(self, src_cfg_class: type) -> dict[str, str]: if a_cond and b_cond: # These should be guaranteed by typechecks above, but assert just in case - assert isinstance(a.target.id, str), "Dataclass attribute not in the expected format. Name is not a string." - assert isinstance(b.value.value, str), "Dataclass attribute docstring is not a string." + assert isinstance( + a.target.id, str + ), "Dataclass attribute not in the expected format. Name is not a string." + assert isinstance( + b.value.value, str + ), "Dataclass attribute docstring is not a string." # Formatting docstring = inspect.cleandoc(b.value.value) @@ -274,6 +293,100 @@ def _get_field_docstrings(self, src_cfg_class: type) -> dict[str, str]: return field_docstrings +def _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, pattern): + """Normalize DSv4 hybrid compression ratios for args and TransformerConfig.""" + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + + variant = kw_args.get( + 'experimental_attention_variant', getattr(args, 'experimental_attention_variant', None) + ) + if variant != 'dsv4_hybrid': + return + + fixed_ratio_map = {Symbols.WINDOW: 0, Symbols.CSA: 4, Symbols.HCA: 128} + ratio_symbols = set(fixed_ratio_map) + sections = pattern.split(Symbols.MTP_SEPARATOR) + layers = ''.join(section.replace(Symbols.PIPE, '') for section in sections) + attn_symbols = [symbol for symbol in layers if symbol in ratio_symbols] + compact_len = len(attn_symbols) + full_len = len(layers) + + def padded_to_compact_and_full(provided): + compact = [] + full = [] + compact_iter = iter(provided) + for symbol in layers: + if symbol in ratio_symbols: + ratio = next(compact_iter) + expected = fixed_ratio_map[symbol] + assert ratio == expected, ( + f"csa_compress_ratios has ratio {ratio} for hybrid symbol " + f"'{symbol}', expected {expected}." + ) + compact.append(ratio) + full.append(ratio) + else: + full.append(0) + return compact, full + + # args keeps the compact CLI form; TransformerConfig gets the full form for per-layer indexing. + if getattr(args, 'csa_compress_ratios', None) is None: + compact_ratios = [fixed_ratio_map[symbol] for symbol in attn_symbols] + args.csa_compress_ratios, kw_args['csa_compress_ratios'] = padded_to_compact_and_full( + compact_ratios + ) + else: + provided = list(args.csa_compress_ratios) + if len(provided) == compact_len: + args.csa_compress_ratios, kw_args['csa_compress_ratios'] = padded_to_compact_and_full( + provided + ) + elif len(provided) == full_len: + compact = [] + for ratio, symbol in zip(provided, layers): + if symbol in ratio_symbols: + expected = fixed_ratio_map[symbol] + assert ratio == expected, ( + f"csa_compress_ratios has ratio {ratio} for hybrid symbol " + f"'{symbol}', expected {expected}." + ) + compact.append(ratio) + else: + assert ratio == 0, ( + f"csa_compress_ratios should not pad non-attention hybrid symbol " + f"'{symbol}' with non-zero ratio {ratio}." + ) + args.csa_compress_ratios = compact + kw_args['csa_compress_ratios'] = provided + else: + raise AssertionError( + f"csa_compress_ratios length ({len(provided)}) must equal either the " + f"number of W/C/H attention symbols ({compact_len}) or the legacy " + f"number of all layers in the hybrid pattern ({full_len}) for pattern " + f"'{pattern}'." + ) + + +def _resolve_dsa_kernel_backend_cli_default(args, kw_args): + """Resolve an omitted --dsa-kernel-backend flag to its effective backend. + + The CLI flag carries no static default so an omitted flag is distinguishable + from an explicit "none". DSv4 hybrid launches historically ran fused kernels + by default (the deprecated --no-dsa-kernel-fusion flag defaulted to True), so + an omitted flag resolves to "cudnn" for that variant unless the deprecated + switch was passed; every other configuration resolves to "none". Runs before + TransformerConfig.__post_init__ folds the deprecated switch into + dsa_kernel_backend, so explicit legacy values still win. + """ + if 'dsa_kernel_backend' not in kw_args or kw_args['dsa_kernel_backend'] is not None: + return + variant = kw_args.get('experimental_attention_variant') or kw_args.get('linear_attention_type') + if variant == 'dsv4_hybrid' and getattr(args, 'apply_dsa_kernel_fusion', None) is None: + kw_args['dsa_kernel_backend'] = 'cudnn' + else: + kw_args['dsa_kernel_backend'] = 'none' + + def core_transformer_config_from_args(args, config_class=None): from megatron.core.activations import squared_relu from megatron.core.fusions.fused_bias_geglu import quick_gelu @@ -293,7 +406,9 @@ def core_transformer_config_from_args(args, config_class=None): config_class = MLATransformerConfig if args.heterogeneous_layers_config_path is not None: - assert not args.multi_latent_attention, "Multi latent attention with heterogeneous layers is not supported." + assert ( + not args.multi_latent_attention + ), "Multi latent attention with heterogeneous layers is not supported." config_class = HeterogeneousTransformerConfig # Translate args to core transformer configuration @@ -306,9 +421,10 @@ def core_transformer_config_from_args(args, config_class=None): kw_args['pipeline_dtype'] = args.params_dtype kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm kw_args['num_moe_experts'] = args.num_experts + kw_args['actual_vocab_size'] = args.padded_vocab_size kw_args['rotary_interleaved'] = args.rotary_interleaved - kw_args['num_layers_in_first_pipeline_stage']= args.decoder_first_pipeline_num_layers - kw_args['num_layers_in_last_pipeline_stage']= args.decoder_last_pipeline_num_layers + kw_args['num_layers_in_first_pipeline_stage'] = args.decoder_first_pipeline_num_layers + kw_args['num_layers_in_last_pipeline_stage'] = args.decoder_last_pipeline_num_layers kw_args['fp8_param'] = args.fp8_param_gather kw_args['fp4_param'] = args.fp4_param_gather if args.swiglu: @@ -336,17 +452,31 @@ def core_transformer_config_from_args(args, config_class=None): # Pop 'rope_type' to let the config class use the default value. kw_args.pop('rope_type', None) else: - assert (args.multi_latent_attention or args.rope_type == 'rope'), ( - f'Common attention only support rope_type="rope", but got {args.rope_type}.' - ) + assert ( + args.multi_latent_attention or args.rope_type == 'rope' + ), f'Common attention only support rope_type="rope", but got {args.rope_type}.' if len(args.cp_comm_type) == 1: kw_args['cp_comm_type'] = args.cp_comm_type[0] if args.hybrid_layer_pattern is not None: kw_args['is_hybrid_model'] = True from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols - if Symbols.DS_ATTENTION in args.hybrid_layer_pattern: - kw_args['experimental_attention_variant'] = 'dsa' + + pattern = args.hybrid_layer_pattern + has_kda = Symbols.KDA in pattern + has_dsv4_csa = any(s in pattern for s in (Symbols.WINDOW, Symbols.CSA, Symbols.HCA)) + has_dsa = Symbols.DS_ATTENTION in pattern + if getattr(args, 'experimental_attention_variant', None) is None: + if has_dsv4_csa: + kw_args['experimental_attention_variant'] = 'dsv4_hybrid' + elif has_dsa: + kw_args['experimental_attention_variant'] = 'dsa' + elif has_kda: + kw_args['experimental_attention_variant'] = 'kda' + + _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, pattern) + + _resolve_dsa_kernel_backend_cli_default(args, kw_args) kw_args['inference_sampling_seed'] = args.seed @@ -430,16 +560,8 @@ def _default_config_from_args(cls: type, args: Namespace, return_instance: bool return kwargs -def gpt_config_from_args( - args: Namespace, config: TransformerConfig | None = None, model_config_cls: type = GPTModelConfig -) -> Any: - """Create a GPTModelConfig (or a compatible subclass) from the `args` Namespace. - - `model_config_cls` lets callers reuse this same arg-derivation logic for - subclasses that only override metadata (e.g. `builder`) and add no new fields, - such as `ModelOptModelConfig`. - """ - assert issubclass(model_config_cls, GPTModelConfig) +def gpt_config_from_args(args: Namespace, config: TransformerConfig | None = None) -> Any: + """Create a GPTModelConfig from the appropriate values in the `args` Namespace.""" kwargs = {} if config is None: @@ -457,7 +579,6 @@ def gpt_config_from_args( kwargs["transformer_layer_spec"] = import_module(args.spec) kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy - kwargs["logit_dtype"] = getattr(args, "logit_dtype", None) kwargs["position_embedding_type"] = args.position_embedding_type kwargs["rotary_percent"] = args.rotary_percent kwargs["rotary_base"] = args.rotary_base @@ -475,23 +596,17 @@ def gpt_config_from_args( kwargs["vocab_size"] = args.padded_vocab_size kwargs["should_pad_vocab"] = False else: - assert args.vocab_size is not None, "Either --padded-vocab-size or --vocab-size must be specified." + assert ( + args.vocab_size is not None + ), "Either --padded-vocab-size or --vocab-size must be specified." kwargs["vocab_size"] = args.vocab_size kwargs["should_pad_vocab"] = True - return model_config_cls(**kwargs) + return GPTModelConfig(**kwargs) -def hybrid_config_from_args( - args: Namespace, config: TransformerConfig | None = None, model_config_cls: type = HybridModelConfig -) -> Any: - """Create a HybridModelConfig (or a compatible subclass) from the `args` Namespace. - - `model_config_cls` lets callers reuse this same arg-derivation logic for - subclasses that only override metadata (e.g. `builder`) and add no new fields, - such as `ModelOptHybridModelConfig`. - """ - assert issubclass(model_config_cls, HybridModelConfig) +def hybrid_config_from_args(args: Namespace, config: TransformerConfig | None = None) -> Any: + """Create a HybridModelConfig from the appropriate values in the `args` Namespace.""" kwargs = {} if config is None: @@ -505,10 +620,15 @@ def hybrid_config_from_args( not transformer_cfg.inference_fuse_tp_communication ), "inference_fuse_tp_communication is not supported for HybridModel" elif args.spec is not None: - kwargs["hybrid_stack_spec"] = import_module(args.spec) + hybrid_stack_spec = import_module(args.spec) + # ModuleSpec implements __call__ to build the described module, so a + # generic callable check would instantiate static specs here before a + # ProcessGroupCollection exists. Only resolve callable spec factories. + if callable(hybrid_stack_spec) and not isinstance(hybrid_stack_spec, ModuleSpec): + hybrid_stack_spec = hybrid_stack_spec(transformer_cfg) + kwargs["hybrid_stack_spec"] = hybrid_stack_spec kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy - kwargs["logit_dtype"] = getattr(args, "logit_dtype", None) kwargs["hybrid_layer_pattern"] = args.hybrid_layer_pattern kwargs["position_embedding_type"] = args.position_embedding_type kwargs["rotary_percent"] = args.rotary_percent @@ -526,11 +646,13 @@ def hybrid_config_from_args( kwargs["vocab_size"] = args.padded_vocab_size kwargs["should_pad_vocab"] = False else: - assert args.vocab_size is not None, "Either --padded-vocab-size or --vocab-size must be specified." + assert ( + args.vocab_size is not None + ), "Either --padded-vocab-size or --vocab-size must be specified." kwargs["vocab_size"] = args.vocab_size kwargs["should_pad_vocab"] = True - return model_config_cls(**kwargs) + return HybridModelConfig(**kwargs) def pretrain_cfg_container_from_args(args: Namespace, model_cfg=None) -> PretrainConfigContainer: @@ -574,7 +696,6 @@ def pretrain_cfg_container_from_args(args: Namespace, model_cfg=None) -> Pretrai checkpoint=CheckpointConfig(**ckpt_kwargs), profiling=ProfilingConfig(**prof_kwargs), tokenizer=_default_config_from_args(TokenizerConfig, args), - rerun_state_machine=RerunStateMachineConfig(**rerunsm_kwargs), straggler=_default_config_from_args(StragglerDetectionConfig, args), ) @@ -595,9 +716,7 @@ def inference_cfg_from_args(args: Namespace) -> InferenceSetupConfig: return _default_config_from_args(InferenceSetupConfig, args) -def inference_cfg_container_from_args( - args: Namespace, model_cfg=None -) -> InferenceConfigContainer: +def inference_cfg_container_from_args(args: Namespace, model_cfg=None) -> InferenceConfigContainer: """Build an InferenceConfigContainer from the argparse arguments. This mirrors ``pretrain_cfg_container_from_args`` but assembles only the configs that diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index c6374a6c178..873c904675a 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.""" @@ -2273,6 +2273,7 @@ def core_transformer_config_from_args(args, config_class=None): from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols _pat = args.hybrid_layer_pattern + _has_kda = Symbols.KDA in _pat _has_dsv4_csa = (Symbols.CSA in _pat) or (Symbols.HCA in _pat) or (Symbols.WINDOW in _pat) _has_dsa = Symbols.DS_ATTENTION in _pat if getattr(args, 'experimental_attention_variant', None) is None: @@ -2285,6 +2286,8 @@ def core_transformer_config_from_args(args, config_class=None): kw_args['experimental_attention_variant'] = 'dsv4_hybrid' elif _has_dsa: kw_args['experimental_attention_variant'] = 'dsa' + elif _has_kda: + kw_args['experimental_attention_variant'] = 'kda' # Normalize compact and legacy-padded ratios through the shared migration helper. _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, _pat) diff --git a/megatron/training/training.py b/megatron/training/training.py index 32e0587fd8d..66ffac27395 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -58,6 +58,7 @@ from megatron.core.inference.symmetric_memory import SymmetricMemoryManager from megatron.core.inference.unified_memory import create_unified_mempool from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + is_gated_delta_net_variant, is_linear_attention_variant, ) from megatron.core.msc_utils import MultiStorageClientFeature, open_file @@ -1089,7 +1090,17 @@ def transformer_flops(): num_linear_attention_layers = sum(linear_attention_pattern) num_standard_attention_layers = num_layers - num_linear_attention_layers - if args.experimental_attention_variant == "gated_delta_net": + 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 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 index 963a5742943..6bc5042b78e 100644 --- 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 @@ -45,6 +45,7 @@ MODEL_ARGS: --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 diff --git a/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py b/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py index d9c4c14b1f7..3b277d4f5b8 100644 --- a/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py +++ b/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from unittest.mock import MagicMock, patch @@ -8,7 +8,10 @@ from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_block import TransformerBlockSubmodules -from megatron.core.transformer.transformer_layer import TransformerLayer +from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, + TransformerLayer, +) # --------------------------------------------------------------------------- # Helpers: fake backend and config builders @@ -82,6 +85,7 @@ def _make_config(**overrides): use_kitchen_attention=False, kitchen_attention_backend="sdpa", fallback_to_eager_attn=False, + enable_hyper_connections=False, ) defaults.update(overrides) cfg = MagicMock() @@ -108,9 +112,10 @@ def _fn(variant): "variant, expected", [ ("gdn", True), - ("gdn2", True), + ("kda", True), ("gated_delta_net", True), ("dsa", False), + ("dsv4_hybrid", False), (None, False), ("some_unknown_variant", False), ], @@ -120,6 +125,30 @@ def test_variants(self, variant, expected): assert self._fn(variant) is expected +class TestNormalizeExperimentalAttentionVariant: + """Validate canonical variant names and deprecated aliases.""" + + @staticmethod + def _fn(variant): + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + normalize_experimental_attention_variant, + ) + + return normalize_experimental_attention_variant(variant) + + @pytest.mark.parametrize("variant", ["gdn", "kda", "dsa", "dsv4_hybrid", None]) + def test_canonical_names_are_unchanged(self, variant): + """Canonical and unrelated variant names pass through unchanged.""" + assert self._fn(variant) == variant + + def test_legacy_gated_delta_net_alias_warns_and_normalizes(self): + """The legacy spelling warns and resolves to canonical GDN.""" + with pytest.warns(DeprecationWarning, match="Use 'gdn' instead"): + normalized = self._fn("gated_delta_net") + + assert normalized == "gdn" + + # =================================================================== # Tests for get_moe_layer_pattern # =================================================================== @@ -247,6 +276,26 @@ def test_returns_correct_module_spec(self): assert spec.module is GatedDeltaNet assert spec.metainfo == {"fuse_input_layernorm": True} + def test_kda_uses_direct_projection_submodules(self): + """KDA uses separate input and beta projections without fused input norm.""" + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_gated_delta_net_module_spec, + ) + from megatron.core.ssm.gated_delta_net import ( + KimiDeltaAttention, + KimiDeltaAttentionSubmodules, + ) + + backend = _make_backend() + cfg = _make_config(experimental_attention_variant="kda", normalization="RMSNorm") + spec = get_gated_delta_net_module_spec(cfg, backend=backend) + + assert spec.module is KimiDeltaAttention + assert isinstance(spec.submodules, KimiDeltaAttentionSubmodules) + assert spec.submodules.in_proj == _FakeColumnParallelLinear + assert spec.submodules.beta_proj == _FakeColumnParallelLinear + assert spec.metainfo == {"fuse_input_layernorm": False} + def test_submodules_use_backend_modules(self): """Verify backend-provided projection/norm modules are wired into submodules.""" from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( @@ -420,9 +469,10 @@ class TestGetExperimentalAttentionVariantModuleSpec: "variant, target_fn", [ ("gdn", "get_gated_delta_net_module_spec"), - ("gdn2", "get_gated_delta_net_module_spec"), + ("kda", "get_gated_delta_net_module_spec"), ("gated_delta_net", "get_gated_delta_net_module_spec"), ("dsa", "get_dsa_module_spec_for_backend"), + ("dsv4_hybrid", "get_dsv4_hybrid_module_spec_for_backend"), ], ) def test_dispatches_to_variant_handler(self, variant, target_fn): @@ -539,7 +589,7 @@ def test_hybrid_attention_pattern(self): assert specs[2].submodules.self_attention is exp_attn_spec assert specs[3].submodules.self_attention is std_attn_spec - def test_hybrid_moe_pattern(self): + def test_hybrid_moe_pattern_with_mhc(self): """Verify MLP alternates between MoE and dense specs per moe_layer_freq pattern.""" from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( get_transformer_layer_with_experimental_attention_variant_spec, @@ -551,6 +601,7 @@ def test_hybrid_moe_pattern(self): num_moe_experts=8, moe_layer_freq=2, normalization="RMSNorm", + enable_hyper_connections=True, ) backend = _make_backend() attn_spec = self._make_attention_spec(fuse_input_layernorm=False) @@ -574,6 +625,8 @@ def test_hybrid_moe_pattern(self): assert specs[1].submodules.mlp is dense_spec assert specs[2].submodules.mlp is moe_spec assert specs[3].submodules.mlp is dense_spec + for s in specs: + assert s.module is HyperConnectionTransformerLayer # =================================================================== diff --git a/tests/unit_tests/ssm/gated_delta_net/test_gdn.py b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py index 6ded0dfae8f..325b1f5c2d1 100644 --- a/tests/unit_tests/ssm/gated_delta_net/test_gdn.py +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py @@ -58,7 +58,7 @@ def _make_gdn_config(**overrides): "num_attention_heads": 8, "activation_func": F.silu, "bf16": True, - "experimental_attention_variant": "gated_delta_net", + "experimental_attention_variant": "gdn", "linear_attention_freq": [1], "transformer_impl": "transformer_engine", } @@ -89,7 +89,7 @@ def test_gdn_pre_gated_delta_rule_fusion_accepts_gdn_variant(): def test_gdn_pre_gated_delta_rule_fusion_requires_gdn_variant(): - with pytest.raises(ValueError, match="experimental_attention_variant='gated_delta_net'"): + with pytest.raises(ValueError, match="experimental_attention_variant='gdn'"): _make_gdn_config( experimental_attention_variant=None, linear_attention_freq=None, diff --git a/tests/unit_tests/ssm/test_kimi_delta_attention.py b/tests/unit_tests/ssm/gated_delta_net/test_kda.py similarity index 85% rename from tests/unit_tests/ssm/test_kimi_delta_attention.py rename to tests/unit_tests/ssm/gated_delta_net/test_kda.py index 82c5244e9db..1b0b5ecfc78 100644 --- a/tests/unit_tests/ssm/test_kimi_delta_attention.py +++ b/tests/unit_tests/ssm/gated_delta_net/test_kda.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import copy import os @@ -43,6 +43,7 @@ def _make_config( 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, @@ -66,6 +67,9 @@ def _make_config( 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", ) @@ -73,7 +77,6 @@ def _make_config( def test_kda_norm_out_recompute_config_accepts_hybrid_model(): config = replace( _make_config(), - experimental_attention_variant=None, is_hybrid_model=True, recompute_granularity="selective", recompute_modules=["gdn_norm_out"], @@ -82,6 +85,17 @@ def test_kda_norm_out_recompute_config_accepts_hybrid_model(): 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( @@ -96,48 +110,40 @@ def _build_kda(config: TransformerConfig) -> KimiDeltaAttention: ).cuda() -def test_kda_rejects_unsupported_head_layouts(): - config = _make_config() - config.linear_num_value_heads = 8 +def test_kda_config_rejects_unsupported_head_layouts(): with pytest.raises(ValueError, match="equal key and value head counts"): - KimiDeltaAttention._validate_config(config) + replace(_make_config(), linear_num_value_heads=8) - config = _make_config() - config.linear_value_head_dim = 64 with pytest.raises(ValueError, match="equal key and value head dimensions"): - KimiDeltaAttention._validate_config(config) + 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") - config = _make_config() - config.context_parallel_size = 3 - with pytest.raises(ValueError, match="tensor parallel size times context parallel size"): - KimiDeltaAttention._validate_config(config) + 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): - config = _make_config() - config.kda_safe_gate = True - config.kda_lower_bound = lower_bound - with pytest.raises(ValueError, match="kda_lower_bound"): - KimiDeltaAttention._validate_config(config) + 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 = _make_config() - config.kda_safe_gate = True - config.kda_lower_bound = lower_bound + config = replace(_make_config(), kda_safe_gate=True, kda_lower_bound=lower_bound) - KimiDeltaAttention._validate_config(config) + assert config.kda_lower_bound == lower_bound def test_kda_ignores_lower_bound_when_safe_gate_is_disabled(): - config = _make_config() - config.kda_safe_gate = False - config.kda_lower_bound = None + config = replace(_make_config(), kda_safe_gate=False, kda_lower_bound=None) - KimiDeltaAttention._validate_config(config) + assert config.kda_lower_bound is None def test_kda_rejects_invalid_packed_boundaries(): @@ -237,7 +243,8 @@ def test_kda_packed_matches_unpacked_cp1(): @pytest.mark.internal @pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") -def test_kda_selective_recompute_norm_out(): +@pytest.mark.parametrize("recompute_module", ["gdn_norm_out", "gdn"]) +def test_kda_selective_recompute(recompute_module): Utils.initialize_model_parallel(1, 1) try: @@ -254,7 +261,7 @@ def run(module, hidden_states): base_config = _make_config() recompute_config = copy.deepcopy(base_config) recompute_config.recompute_granularity = "selective" - recompute_config.recompute_modules = ["gdn_norm_out"] + recompute_config.recompute_modules = [recompute_module] model_parallel_cuda_manual_seed(42) base_kda = _build_kda(base_config) @@ -277,16 +284,20 @@ def run(module, hidden_states): ) 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 True - assert recompute_kda.norm_out_checkpoint is not 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: - # The checkpointed gate/norm backward can choose a different BF16 reduction order. - # Keep output and input-gradient checks bitwise, while allowing only BF16-scale - # rounding in parameter gradients. + # 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], @@ -388,12 +399,18 @@ def get_packed_tensor_on_this_cp_rank(tensor): @pytest.mark.parametrize( - ("sequence_packing", "tp", "sp", "cp"), - [(False, 2, True, 1), (True, 1, False, 2), (True, 2, True, 2)], + ("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): - config = _make_config() +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, From 1a0706097d6133faf484f2a357aee2076137c6d1 Mon Sep 17 00:00:00 2001 From: Yuzhong Wang Date: Fri, 21 Aug 2026 01:51:20 -0700 Subject: [PATCH 18/34] feat: Add attention output gate granularity Signed-off-by: Yuzhong Wang (cherry picked from commit a7855d9f029272b8beec2bf12d125991ae6d15a7) --- .../ling_v3_tiny/train_ling_v3_tiny_8gpu.sh | 1 + megatron/core/transformer/attention.py | 5 + .../transformer/multi_latent_attention.py | 50 +- .../core/transformer/transformer_config.py | 24 +- megatron/training/arguments.py | 16 +- megatron/training/training.py | 90 +++- .../model_config.yaml | 1 + .../models/test_ling_v3_tiny_model_spec.py | 1 + .../test_num_floating_point_operations.py | 467 +++++++++++++++++- .../test_multi_latent_attention.py | 135 +++-- 10 files changed, 678 insertions(+), 112 deletions(-) diff --git a/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh b/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh index 1712f908170..5a1bed5dafa 100755 --- a/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh +++ b/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh @@ -90,6 +90,7 @@ export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" --v-head-dim 128 \ --qk-layernorm \ --attention-output-gate \ + --gated-attention-proj-granularity headwise \ --mtp-num-layers 1 \ --mtp-loss-scaling-factor 0.1 \ --recompute-granularity full \ diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index ab7c9614cf5..62c16ff8b4c 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -1701,6 +1701,11 @@ def __init__( Args: name (str | None): module instance name passed top-down from its paranet module """ + if config.gated_attention_proj_granularity != 'elementwise': + raise ValueError( + "Regular attention does not support headwise " + "gated_attention_proj_granularity; use 'elementwise'." + ) super().__init__( config=config, submodules=submodules, diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 8711cfe6af2..ff3ae5991de 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -165,6 +165,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 @@ -233,7 +237,7 @@ def __init__( self.linear_gate = build_module( submodules.linear_gate, self.config.hidden_size, - self.config.num_attention_heads, + self.gate_projection_size, config=self.config, init_method=self.config.init_method, gather_output=False, @@ -529,10 +533,50 @@ def _project_and_apply_mla_output_gate( 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_output_gate(core_attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - """Apply an FP32 sigmoid gate per local MLA attention head.""" + 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 diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 0283df062e4..e8cd33aedb8 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -265,9 +265,16 @@ class TransformerConfig(ModelParallelConfig): attention_output_gate: bool = False """Whether to apply output gating to attention layers. - For MLA, this uses one sigmoid gate per attention head. For standard attention, - this retains the full head-dimension output gate behavior. It is mutually exclusive - with ``head_wise_attn_gate``. + 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 @@ -1520,6 +1527,17 @@ def __post_init__(self): if self.num_query_groups is None: self.num_query_groups = self.num_attention_heads + if self.gated_attention_proj_granularity not in ('elementwise', 'headwise'): + raise ValueError( + "gated_attention_proj_granularity must be either 'elementwise' or 'headwise', " + f"got {self.gated_attention_proj_granularity!r}." + ) + if self.gated_attention_proj_granularity == 'headwise' and not self.multi_latent_attention: + raise ValueError( + "Regular attention does not support headwise " + "gated_attention_proj_granularity; use 'elementwise'." + ) + if ( self.num_query_groups % self.tensor_model_parallel_size != 0 and self.tensor_model_parallel_size % self.num_query_groups != 0 diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 873c904675a..2d08bcf91dc 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -820,6 +820,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). @@ -899,9 +901,8 @@ def validate_args(args, defaults={}): # Infer mtp_num_layers from unified pattern if args.hybrid_layer_pattern and sep in args.hybrid_layer_pattern: - parsed = parse_hybrid_pattern(args.hybrid_layer_pattern) - if parsed.mtp_pattern and parsed.mtp_num_depths > 0: - inferred_mtp_num_layers = parsed.mtp_num_depths + if parsed_hybrid_pattern.mtp_pattern and parsed_hybrid_pattern.mtp_num_depths > 0: + inferred_mtp_num_layers = parsed_hybrid_pattern.mtp_num_depths if args.mtp_num_layers is None: args.mtp_num_layers = inferred_mtp_num_layers elif args.mtp_num_layers != inferred_mtp_num_layers: @@ -945,8 +946,13 @@ def validate_args(args, defaults={}): args.rank, ) - # Infer use of MLA from unified pattern - if args.hybrid_layer_pattern and 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 # === End of hybrid layer pattern: deprecation handling and validation === diff --git a/megatron/training/training.py b/megatron/training/training.py index 66ffac27395..3db5fff8913 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.""" @@ -578,6 +578,8 @@ def attn_layer_flops( gqa=True, gqa_groups=8, kv_channels=None, + attention_output_gate=False, + gated_attention_proj_granularity="elementwise", ): """Calculate FLOPs for an attention layer. @@ -590,10 +592,23 @@ 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)) + + 2 * total_tokens * hidden_size * gate_projection_size + 2 * seqlen_squared_sum * hidden_size * p ) @@ -608,6 +623,7 @@ def mla_attn_layer_flops( 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. @@ -627,13 +643,24 @@ def mla_attn_layer_flops( 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 * num_heads if attention_output_gate else 0) + + hidden_size * gate_projection_size ) # Core attention (L^2) part: QK^T and (softmax(QK^T))V. /2 (causal) cancels *2 (FMA). core = fma * ( @@ -754,13 +781,13 @@ def hybrid_flops( kda_conv_kernel_dim=4, vocab_size=256000, mtp_num_layers=0, - multi_latent_attention=False, 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, @@ -799,29 +826,19 @@ def hybrid_flops( dsv4_token_term * total_tokens + dsv4_core_term * seqlen_squared_sum ) else: - if multi_latent_attention: - plain_attn_layer_flops = 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, - ) - else: - plain_attn_layer_flops = attn_layer_flops( - total_tokens, - seqlen_squared_sum, - hidden_size, - num_attn_heads, - gqa, - gqa_groups, - kv_channels, - ) + # 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( @@ -835,6 +852,7 @@ def hybrid_flops( qk_pos_emb_head_dim, v_head_dim, attention_output_gate, + gated_attention_proj_granularity, ) kda_flops_total = 0 @@ -996,6 +1014,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 @@ -1013,6 +1045,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 @@ -1328,13 +1362,15 @@ def transformer_flops(): kda_conv_kernel_dim=args.linear_conv_kernel_dim or 4, vocab_size=args.padded_vocab_size, mtp_num_layers=mtp_num_layers, - multi_latent_attention=args.multi_latent_attention, 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, 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 index 6bc5042b78e..0eed047d204 100644 --- 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 @@ -84,6 +84,7 @@ MODEL_ARGS: --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 diff --git a/tests/unit_tests/models/test_ling_v3_tiny_model_spec.py b/tests/unit_tests/models/test_ling_v3_tiny_model_spec.py index c55511054a0..bb31ed768c7 100644 --- a/tests/unit_tests/models/test_ling_v3_tiny_model_spec.py +++ b/tests/unit_tests/models/test_ling_v3_tiny_model_spec.py @@ -101,6 +101,7 @@ def test_ling_v3_tiny_reduced_hybrid_model_forward_backward(tp_size): rotary_percent=0.5, qk_layernorm=True, attention_output_gate=True, + gated_attention_proj_granularity="headwise", linear_conv_kernel_dim=4, linear_key_head_dim=32, linear_value_head_dim=32, diff --git a/tests/unit_tests/test_num_floating_point_operations.py b/tests/unit_tests/test_num_floating_point_operations.py index 0f08008ec1b..7ebf6bc4e7d 100644 --- a/tests/unit_tests/test_num_floating_point_operations.py +++ b/tests/unit_tests/test_num_floating_point_operations.py @@ -54,6 +54,7 @@ def _make_gpt_args( args.group_query_attention = False args.num_query_groups = num_attention_heads args.attention_output_gate = False + args.gated_attention_proj_granularity = "elementwise" args.multi_latent_attention = False # MoE / MTP disabled. args.num_experts = None @@ -96,7 +97,6 @@ def _make_hybrid_args(*, num_layers=4, hidden_size=512, num_attention_heads=8, s args.mamba_head_dim = 64 args.mamba_num_groups = 8 args.mamba_num_heads = 128 - args.gdp_num_householder = 3 return args @@ -115,6 +115,7 @@ def _make_ling_hybrid_args(pattern): args.kv_lora_rank = 256 args.v_head_dim = 64 args.attention_output_gate = True + args.gated_attention_proj_granularity = "headwise" return args @@ -273,35 +274,145 @@ def test_hybrid_attention_layers_count(self): expected_delta = num_attn_layers * expected_delta_per_layer_per_unit_sum * bshd_sum assert flops_doubled - flops_bshd == expected_delta + def test_mla_attention_variants_are_counted(self): + """Regression: DSv4 MLA attention variants (CSA 'C', HCA 'H', Window 'W', + DS_ATTENTION 'D') must contribute attention FLOPs. Previously these symbols + were omitted, so a pattern made entirely of C/H/W layers reported zero + attention FLOPs -- roughly halving the throughput estimate.""" + args = _make_hybrid_args() + args.multi_latent_attention = True + # MLA dims (DSv4-like; values only need to be self-consistent and positive). + args.q_lora_rank = 256 + args.kv_lora_rank = 32 + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.v_head_dim = 64 + batch_size = 4 -class TestGatedDeltaProductFlops: - """GDP FLOPs must use the Householder count from the model configuration.""" + # MLA attention variants (C, H) must be counted exactly like dense MLA '+'. + # Both patterns are 2 MLA + 2 Mamba layers, so they yield identical FLOPs. + # Regular attention '*' is intentionally not the reference: hybrid grammar + # assigns different projection work to '*' and '+'. + args.hybrid_layer_pattern = "CMHM" + flops_mla_variants = num_floating_point_operations(args, batch_size) + args.hybrid_layer_pattern = "+M+M" + flops_dense_mla = num_floating_point_operations(args, batch_size) + + assert flops_mla_variants == flops_dense_mla + # And attention must be a non-trivial contributor: dropping the two + # attention layers entirely (all Mamba) changes the estimate. + args.hybrid_layer_pattern = "MMMM" + flops_no_attn = num_floating_point_operations(args, batch_size) + assert flops_mla_variants != flops_no_attn + + +class TestHybridMatchesStandard: + """The hybrid-model path (``hybrid_flops`` / ``mla_attn_layer_flops``) must + agree with the standard-model path (``transformer_flops``) for an equivalent + architecture. + + A dense Transformer with ``N`` layers has ``N`` attention sub-layers and + ``N`` MLP sub-layers. The hybrid path models those as separate layers, so the + matching hybrid pattern is ``N`` attention layers interleaved with ``N`` MLP + layers (``"*-" * N``). With identical hidden / head / ffn / vocab dims the + two estimates must be bit-for-bit equal -- this pins ``mla_attn_layer_flops`` + (and the dense ``attn_layer_flops``) against the reference MLA / MHA / GQA + terms in ``transformer_flops`` (see GitHub PR #5358), both for BSHD and for + the THD ``sum(L_i^2)`` / real-token scaling. + """ - def test_householder_count_changes_flops(self): - args = _make_hybrid_args() - args.spec = ["megatron.core.models.hybrid.hybrid_layer_specs", "gdp_stack_spec"] + @staticmethod + def _equivalent_pair(configure, num_layers=3): + """Build (standard_args, hybrid_args) describing the same architecture.""" + standard = _make_gpt_args(num_layers=num_layers) + configure(standard) + + hybrid = _make_gpt_args(num_layers=num_layers) + configure(hybrid) + # The hybrid path reads Mamba dims even when no Mamba layers are present. + hybrid.mamba_state_dim = 128 + hybrid.mamba_head_dim = 64 + hybrid.mamba_num_groups = 8 + hybrid.mamba_num_heads = 128 + # Match the standard model's attention implementation explicitly: + # '*' is regular attention and '+' is dense MLA in the hybrid grammar. + attention_symbol = "+" if hybrid.multi_latent_attention else "*" + hybrid.hybrid_layer_pattern = (attention_symbol + "-") * num_layers + return standard, hybrid + + def _assert_match(self, configure): + standard, hybrid = self._equivalent_pair(configure) batch_size = 4 - flops_m3 = num_floating_point_operations(args, batch_size) - args.gdp_num_householder = 4 - flops_m4 = num_floating_point_operations(args, batch_size) + # BSHD: closed-form defaults. + assert num_floating_point_operations(standard, batch_size) == num_floating_point_operations( + hybrid, batch_size + ) - total_tokens = batch_size * args.seq_length - d_inner = args.mamba_num_heads * args.mamba_head_dim - group_state_dim = args.mamba_num_groups * args.mamba_state_dim - forward_delta_per_layer = ( - 2 - * total_tokens - * ( - args.hidden_size * (d_inner + group_state_dim + args.mamba_num_heads) - + 4 * (d_inner + group_state_dim) - ) - + 4 * total_tokens * d_inner * args.mamba_state_dim + # THD: both paths must consume the ragged stats identically. + s = standard.seq_length + thd_sum_sq = batch_size * 4 * (s // 4) ** 2 + thd_tokens = batch_size * s * 3 // 4 + assert num_floating_point_operations( + standard, + batch_size, + seqlen_squared_sum_in_batch=thd_sum_sq, + total_real_tokens_in_batch=thd_tokens, + ) == num_floating_point_operations( + hybrid, + batch_size, + seqlen_squared_sum_in_batch=thd_sum_sq, + total_real_tokens_in_batch=thd_tokens, ) - num_gdp_layers = 2 - expected_delta = 3 * num_gdp_layers * forward_delta_per_layer - assert flops_m4 - flops_m3 == expected_delta + def test_dense_mha(self): + self._assert_match(lambda args: None) + + def test_gqa(self): + def configure(args): + args.group_query_attention = True + args.num_query_groups = 2 + + self._assert_match(configure) + + def test_mla_with_q_lora(self): + def configure(args): + args.multi_latent_attention = True + args.group_query_attention = False + args.q_lora_rank = 256 + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.kv_lora_rank = 256 + args.v_head_dim = 64 + + self._assert_match(configure) + + def test_mla_without_q_lora(self): + def configure(args): + args.multi_latent_attention = True + args.group_query_attention = False + args.q_lora_rank = None + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.kv_lora_rank = 256 + args.v_head_dim = 64 + + self._assert_match(configure) + + @pytest.mark.parametrize("gate_granularity", ("elementwise", "headwise")) + def test_mla_output_gate(self, gate_granularity): + def configure(args): + args.multi_latent_attention = True + args.group_query_attention = False + args.q_lora_rank = 256 + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.kv_lora_rank = 256 + args.v_head_dim = 64 + args.attention_output_gate = True + args.gated_attention_proj_granularity = gate_granularity + + self._assert_match(configure) class TestLingHybridAttentionFlops: @@ -352,6 +463,47 @@ def test_mla_uses_ragged_attention_work(self): ) assert flops_full - flops_half == expected_delta + def test_mla_gate_granularity_changes_only_projection_work(self): + args = _make_ling_hybrid_args("+") + batch_size = 2 + total_tokens = batch_size * args.seq_length + + args.gated_attention_proj_granularity = "headwise" + headwise_flops = num_floating_point_operations(args, batch_size) + args.gated_attention_proj_granularity = "elementwise" + elementwise_flops = num_floating_point_operations(args, batch_size) + + expected_delta = ( + 3 + * 2 + * total_tokens + * args.hidden_size + * args.num_attention_heads + * (args.v_head_dim - 1) + ) + assert elementwise_flops - headwise_flops == expected_delta + + @pytest.mark.parametrize( + ("gate_granularity", "gate_projection_size"), (("elementwise", 512), ("headwise", 8)) + ) + def test_hybrid_regular_attention_output_gate(self, gate_granularity, gate_projection_size): + args = _make_ling_hybrid_args("*") + # 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. @@ -742,3 +894,272 @@ def test_dedup_across_topology(self, tp, cp, pp): f"topology tp={tp} cp={cp} pp={pp} dp={dp_size}: " f"got seqlen_squared_sum={seqlen_squared_sum}, expected {expected_sum_sq}" ) + + +def _make_dsv4_args(): + """Minimal args for a DSv4-hybrid MLA model with sparse attention. + + 4 layers with compress_ratios [0, 4, 128, 128] (1 r0, 1 r4, 2 r128). + No MoE / MTP to keep the golden reference simple. + """ + args = _make_gpt_args( + num_layers=4, + hidden_size=512, + num_attention_heads=8, + seq_length=256, + ffn_hidden_size=2048, + padded_vocab_size=1024, + ) + args.multi_latent_attention = True + args.group_query_attention = False + args.q_lora_rank = 128 + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.kv_lora_rank = 64 + args.v_head_dim = 64 + args.o_lora_rank = 64 + args.o_groups = 2 + args.experimental_attention_variant = "dsv4_hybrid" + args.csa_window_size = 64 + args.csa_compress_ratios = [0, 4, 128, 128] + args.dsa_indexer_n_heads = 4 + args.dsa_indexer_head_dim = 32 + args.dsa_indexer_topk = 16 + return args + + +def _dsv4_golden_flops(args, total_tokens, seqlen_squared_sum): + """Independent golden calculator for DSv4-hybrid FLOPs. + + Reimplements the formula from ``num_floating_point_operations`` so that the + test does not just call the same code twice. Assumes no MoE / MTP. + """ + fwd_bwd = 3 + fma = 2 + ffn_exp = 3 if args.swiglu else 2 + + # ---- MLA projections (token-linear, per layer) ---- + q_term = args.q_lora_rank * (args.hidden_size + args.num_attention_heads * args.v_head_dim + 1) + kv_term = args.hidden_size * args.v_head_dim + args.v_head_dim + o_term = ( + args.num_attention_heads * args.v_head_dim * args.o_lora_rank + + args.o_groups * args.o_lora_rank * args.hidden_size + ) + mla_proj_per_layer = fwd_bwd * fma * (q_term + kv_term + o_term) + + # ---- DSv4 sparse attention extra (token-linear + L^2) ---- + ratios = args.csa_compress_ratios + n_r0 = sum(1 for r in ratios if r == 0) + n_r4 = sum(1 for r in ratios if r == 4) + n_r128 = sum(1 for r in ratios if r == 128) + nh = args.num_attention_heads + vhd = args.v_head_dim + w = args.csa_window_size + + # Token-linear sparse attention + sparse_r0 = n_r0 * nh * w * vhd * 2 + sparse_r128_win = n_r128 * nh * w * vhd * 2 + if n_r4 > 0: + eff_topk = min(args.dsa_indexer_topk, args.seq_length // 4) + avg_comp_4 = eff_topk * (1 - eff_topk * 4 / (2 * args.seq_length)) + sparse_r4 = n_r4 * nh * (w + avg_comp_4) * vhd * 2 + idx_tok = ( + n_r4 * args.hidden_size * (2 * args.dsa_indexer_head_dim) * 2 + + n_r4 * args.q_lora_rank * args.dsa_indexer_n_heads * args.dsa_indexer_head_dim + + n_r4 * args.hidden_size * args.dsa_indexer_n_heads + ) + idx_core = n_r4 * args.dsa_indexer_n_heads * args.dsa_indexer_head_dim / 4 + else: + sparse_r4, idx_tok, idx_core = 0, 0, 0 + + compressor = n_r4 * args.hidden_size * (2 * vhd) * 2 + n_r128 * args.hidden_size * (1 * vhd) * 2 + dsv4_token = fwd_bwd * fma * (sparse_r0 + sparse_r4 + sparse_r128_win + compressor + idx_tok) + # L^2 core: r=128 compressed-KV + r=4 indexer scoring + r128_core = n_r128 * nh * vhd / 128 + dsv4_core = fwd_bwd * fma * (r128_core + idx_core) + + # ---- Aggregation ---- + num_layers = args.num_layers + self_attn_term = mla_proj_per_layer * num_layers + dsv4_token + self_attn_core_term = dsv4_core # standard core is 0 for DSv4 + + mlp = fwd_bwd * fma * args.hidden_size * (args.ffn_hidden_size * ffn_exp * num_layers) + logit = fwd_bwd * fma * args.hidden_size * args.padded_vocab_size + + return total_tokens * (mlp + self_attn_term + logit) + seqlen_squared_sum * self_attn_core_term + + +class TestDSv4Hybrid: + """DSv4 hybrid sparse-attention FLOPs against an independent golden calculator.""" + + def test_bshd(self): + """BSHD (uniform sequences) must match the golden calculator.""" + args = _make_dsv4_args() + batch_size = 2 + total_tokens = batch_size * args.seq_length + sum_sq = batch_size * args.seq_length**2 + + flops = num_floating_point_operations(args, batch_size) + expected = _dsv4_golden_flops(args, total_tokens, sum_sq) + assert flops == expected + + def test_thd(self): + """THD (packed variable-length subsequences) must match the golden + calculator and be strictly less than BSHD due to the L^2 sparse-attention + components (r=128 compressed-KV, r=4 indexer scoring).""" + args = _make_dsv4_args() + batch_size = 2 + packed_lengths = [64, 64, 128, 256] + total_tokens = sum(packed_lengths) + thd_sum_sq = sum(L**2 for L in packed_lengths) + + flops = num_floating_point_operations( + args, + batch_size, + seqlen_squared_sum_in_batch=thd_sum_sq, + total_real_tokens_in_batch=total_tokens, + ) + expected = _dsv4_golden_flops(args, total_tokens, thd_sum_sq) + assert flops == expected + # THD must be strictly less than BSHD. + bshd_flops = num_floating_point_operations(args, batch_size) + assert flops < bshd_flops + + +# ``compress_ratio -> hybrid attention symbol``: Window (r0) / CSA (r4) / HCA (r128). +_RATIO_TO_SYMBOL = {0: "W", 4: "C", 128: "H"} + + +def _make_dsv4_pair(*, ratios, mtp, moe): + """Build (standard_args, hybrid_args) for the SAME DSv4-hybrid model. + + The standard-model path expresses it as ``len(ratios)`` fused layers + (MLA attention + FFN per layer, ``experimental_attention_variant`` = + ``dsv4_hybrid``); the hybrid-model path expresses the identical model as a + pattern with one attention layer (W/C/H per ratio) and one FFN layer (MoE + ``E`` or dense ``-``) each. With one MTP depth, the standard path appends a + fused MTP layer (ratio 0 / Window) and the hybrid path appends ``/W``. + """ + ffn_symbol = "E" if moe else "-" + + def _configure(args): + args.multi_latent_attention = True + args.group_query_attention = False + args.q_lora_rank = 256 + args.qk_head_dim = 64 + args.qk_pos_emb_head_dim = 32 + args.kv_lora_rank = 64 + args.v_head_dim = 64 + args.o_lora_rank = 64 + args.o_groups = 2 + args.experimental_attention_variant = "dsv4_hybrid" + args.csa_window_size = 64 + args.dsa_indexer_n_heads = 4 + args.dsa_indexer_head_dim = 32 + args.dsa_indexer_topk = 16 + if moe: + args.num_experts = 256 + args.moe_layer_freq = 1 + args.moe_router_topk = 6 + args.moe_ffn_hidden_size = 2048 + args.moe_shared_expert_intermediate_size = 2048 + + # ---- standard-model args (fused layers) ---- + standard = _make_gpt_args( + num_layers=len(ratios), + hidden_size=512, + num_attention_heads=8, + seq_length=256, + ffn_hidden_size=2048, + padded_vocab_size=1024, + ) + _configure(standard) + # ``csa_compress_ratios`` must cover main + MTP layers; the MTP layer is a + # Window (ratio 0) attention layer. + standard.csa_compress_ratios = list(ratios) + [0] * mtp + standard.mtp_num_layers = mtp if mtp else None + + # ---- hybrid-model args (split layers, same dims) ---- + hybrid = _make_gpt_args( + num_layers=len(ratios), + hidden_size=512, + num_attention_heads=8, + seq_length=256, + ffn_hidden_size=2048, + padded_vocab_size=1024, + ) + _configure(hybrid) + hybrid.mamba_state_dim = 128 + hybrid.mamba_head_dim = 64 + hybrid.mamba_num_groups = 8 + hybrid.mamba_num_heads = 128 + main_pattern = "".join(_RATIO_TO_SYMBOL[r] + ffn_symbol for r in ratios) + mtp_pattern = ("/" + "W" + ffn_symbol) * mtp # MTP layer = Window attention + hybrid.hybrid_layer_pattern = main_pattern + mtp_pattern + hybrid.mtp_num_layers = mtp if mtp else None + return standard, hybrid + + +class TestDSv4HybridMatchesStandard: + """The hybrid-model path and the standard-model path must agree for a DSv4 + model with sparse attention. + + Regression for GitHub PR #5485: ``hybrid_flops`` previously routed DSv4 + attention through the plain full-MLA helper, which (a) computes the full + ``O(L^2)`` core attention that DSv4 actually replaces with sparse attention, + and (b) uses the dense output projection instead of DSv4's grouped low-rank + one. That inflated the hybrid-model FLOPs estimate by ~30% vs. the + (correct) standard-model ``dsv4_hybrid`` estimate, even though both describe + the same model. Both paths now share ``_dsv4_hybrid_self_attention_flops``, + so they must match bit-for-bit -- for BSHD and THD, with and without MTP, + and for dense and MoE FFN layers. + """ + + @staticmethod + def _assert_match(*, ratios, mtp, moe): + standard, hybrid = _make_dsv4_pair(ratios=ratios, mtp=mtp, moe=moe) + batch_size = 2 + + # BSHD (closed-form defaults). + assert num_floating_point_operations(standard, batch_size) == num_floating_point_operations( + hybrid, batch_size + ) + + # THD: both paths must consume the ragged stats identically. + s = standard.seq_length + thd_sum_sq = batch_size * 4 * (s // 4) ** 2 + thd_tokens = batch_size * s * 3 // 4 + assert num_floating_point_operations( + standard, + batch_size, + seqlen_squared_sum_in_batch=thd_sum_sq, + total_real_tokens_in_batch=thd_tokens, + ) == num_floating_point_operations( + hybrid, + batch_size, + seqlen_squared_sum_in_batch=thd_sum_sq, + total_real_tokens_in_batch=thd_tokens, + ) + + def test_dense_no_mtp(self): + self._assert_match(ratios=[0, 4, 128], mtp=0, moe=False) + + def test_dense_with_mtp(self): + self._assert_match(ratios=[0, 4, 128], mtp=1, moe=False) + + def test_moe_no_mtp(self): + self._assert_match(ratios=[0, 4, 128, 128], mtp=0, moe=True) + + def test_moe_with_mtp(self): + # Production-like: MoE FFN, window/CSA/HCA mix, one MTP depth. + self._assert_match(ratios=[0, 0, 4, 128, 4, 128], mtp=1, moe=True) + + def test_hybrid_overcount_is_gone(self): + """The OLD bug made the hybrid estimate strictly larger than the standard + one (phantom full-MLA O(L^2) core attention). Pin that they are equal, + not merely close.""" + standard, hybrid = _make_dsv4_pair(ratios=[0, 4, 128, 128], mtp=1, moe=True) + batch_size = 2 + std_flops = num_floating_point_operations(standard, batch_size) + hyb_flops = num_floating_point_operations(hybrid, batch_size) + assert hyb_flops == std_flops diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index d5e9f2fd7b8..e6dbcfc1caf 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,10 +741,11 @@ 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): + def setup_and_teardown(self, gate_granularity): Utils.initialize_model_parallel(1, 1) model_parallel_cuda_manual_seed(123) self.transformer_config = MLATransformerConfig( @@ -761,6 +762,7 @@ def setup_and_teardown(self): 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, @@ -772,40 +774,51 @@ def setup_and_teardown(self): def teardown_method(self, method): Utils.destroy_model_parallel() - def test_constructor_builds_head_wise_gate(self): + def test_constructor_builds_configured_gate_projection(self, gate_granularity): gate = self.parallel_attention.linear_gate assert gate is not None - assert gate.weight.shape == ( - self.transformer_config.num_attention_heads, - self.transformer_config.hidden_size, + 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_is_applied_per_head(self): - core_attn_out = torch.arange(24, dtype=torch.bfloat16).view(2, 1, 12) - gate = torch.tensor([[[-2.0, -1.0, 0.0, 1.0]], [[2.0, 3.0, 4.0, 5.0]]]).bfloat16() - - output = self.parallel_attention._apply_mla_output_gate(core_attn_out, gate) - expected = core_attn_out.view(2, 1, 4, 3) * torch.sigmoid(gate.float()).to( - core_attn_out.dtype - ).unsqueeze(-1) - - torch.testing.assert_close(output, expected.reshape_as(output), atol=0, rtol=0) - - def test_gate_uses_fp32_sigmoid_vjp(self): - if not torch.cuda.is_available(): - pytest.skip("FP32 sigmoid VJP check requires CUDA") - - generator = torch.Generator(device="cuda").manual_seed(20260814) - gate_input = torch.randn( - (8, 1, 4), generator=generator, device="cuda", dtype=torch.bfloat16 + 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 ) - core_output = torch.randn( - (8, 1, 32), generator=generator, device="cuda", dtype=torch.bfloat16 + gate_size = ( + output_size + if gate_granularity == "elementwise" + else self.transformer_config.num_attention_heads ) - output_gradient = torch.randn( - core_output.shape, generator=generator, device="cuda", dtype=torch.bfloat16 + 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) @@ -814,21 +827,37 @@ def test_gate_uses_fp32_sigmoid_vjp(self): actual_output, (actual_gate, actual_core), output_gradient ) - reference_gate = gate_input.detach().clone().requires_grad_(True) - reference_core = core_output.detach().clone().requires_grad_(True) - reference_output = ( - reference_core.view(8, 1, 4, 8) - * torch.sigmoid(reference_gate.float()).to(reference_core.dtype).unsqueeze(-1) - ).reshape_as(reference_core) - reference_gradients = torch.autograd.grad( - reference_output, (reference_gate, reference_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 + ) + expanded_fp32_scale = torch.repeat_interleave( + sigmoid_fp32, self.transformer_config.v_head_dim, dim=-1 + ) + expected_core_gradient = (output_gradient.float() * expanded_fp32_scale).to( + core_output.dtype + ) + 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 ) - # Allow exactly one BF16 quantization step from the compiled sigmoid/cast path. - bf16_eps = torch.finfo(torch.bfloat16).eps - torch.testing.assert_close(actual_output, reference_output, atol=bf16_eps, rtol=0.0) - for actual, reference in zip(actual_gradients, reference_gradients): - torch.testing.assert_close(actual, reference, atol=bf16_eps, rtol=0.0) + torch.testing.assert_close(actual_output, expected_output, atol=0, rtol=0) + torch.testing.assert_close(actual_gradients[0], expected_gate_gradient, atol=0, rtol=0) + torch.testing.assert_close(actual_gradients[1], expected_core_gradient, atol=0, rtol=0) def test_missing_gate_spec_raises(self): submodules = get_mla_self_attn_submodules() @@ -841,7 +870,7 @@ def test_missing_gate_spec_raises(self): attn_mask_type=AttnMaskType.causal, ) - def test_fused_down_projection_is_rejected(self): + 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, @@ -856,6 +885,7 @@ def test_fused_down_projection_is_rejected(self): rotary_base=10000, original_max_position_embeddings=32, attention_output_gate=True, + gated_attention_proj_granularity=gate_granularity, mla_down_proj_fusion=True, ) @@ -1596,20 +1626,22 @@ def test_clip_qk_with_absorption_raises_error(self): ], ) @pytest.mark.parametrize( - ("tp", "sp", "cp", "output_gate"), + ("tp", "sp", "cp", "output_gate", "gate_granularity"), [ - (4, False, 1, False), # TP w/o SP - (4, True, 1, False), # TP w/ SP - (1, False, 4, False), # CP - (2, False, 2, False), # CP + TP w/o SP - (2, True, 2, False), # CP + TP w/ SP - (4, True, 1, True), # Gated MLA with TP + SP - (1, False, 2, True), # Gated MLA with CP + (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, output_gate + 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.") @@ -1687,6 +1719,7 @@ def initialize_gpt_model( 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 de28a8214cee5fb9355199f8d1fd7263a00835e1 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 23 Aug 2026 01:44:51 +0800 Subject: [PATCH 19/34] Remove Ling-specific onboarding artifacts Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit f9e20e58f30eed25501dc48d5a083a8b4d57d8bb) --- examples/ling_v3_tiny/README.md | 14 -- examples/ling_v3_tiny/model_spec.py | 25 --- .../ling_v3_tiny/train_ling_v3_tiny_8gpu.sh | 105 --------- .../models/test_ling_v3_tiny_model_spec.py | 206 ------------------ 4 files changed, 350 deletions(-) delete mode 100644 examples/ling_v3_tiny/README.md delete mode 100644 examples/ling_v3_tiny/model_spec.py delete mode 100755 examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh delete mode 100644 tests/unit_tests/models/test_ling_v3_tiny_model_spec.py diff --git a/examples/ling_v3_tiny/README.md b/examples/ling_v3_tiny/README.md deleted file mode 100644 index fdcd150bc3d..00000000000 --- a/examples/ling_v3_tiny/README.md +++ /dev/null @@ -1,14 +0,0 @@ -# Ling-V3 Tiny pretraining example - -This example runs the full Ling-V3 Tiny hybrid pattern from fresh initialization with -`MockGPTDataset`. - -Run it on one 8-GPU node: - -```bash -bash examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh -``` - -The default is a two-step BF16 training smoke run without validation or test iterations. -`TRAIN_ITERS`, `SEQ_LENGTH`, and `MASTER_PORT` may be overridden. No checkpoint is written unless -`SAVE_DIR` is explicitly set. diff --git a/examples/ling_v3_tiny/model_spec.py b/examples/ling_v3_tiny/model_spec.py deleted file mode 100644 index 158714e98e8..00000000000 --- a/examples/ling_v3_tiny/model_spec.py +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -from copy import deepcopy - -from megatron.core.extensions.transformer_engine import TEColumnParallelLinear -from megatron.core.models.hybrid.hybrid_layer_specs import ( - hybrid_stack_spec as default_hybrid_stack_spec, -) -from megatron.core.tensor_parallel import ColumnParallelLinear - -LING_V3_TINY_HYBRID_LAYER_PATTERN = "K-KEKE+EKEKEKE+EKEKEKE+EKEKEKE+EKEKEKE+EKEKEKE+E/+E" - - -# Ling-V3 Tiny uses tensor-parallel projections for the KDA beta path and the -# MLA latent down projections and output gate. Keep the default hybrid spec -# unchanged and expose this model recipe through the standard --spec hook. -hybrid_stack_spec = deepcopy(default_hybrid_stack_spec) - -_kda_submodules = hybrid_stack_spec.submodules.kda_layer.submodules.self_attention.submodules -_kda_submodules.beta_proj = ColumnParallelLinear - -_mla_submodules = hybrid_stack_spec.submodules.mla_layer.submodules.self_attention.submodules -_mla_submodules.linear_q_down_proj = TEColumnParallelLinear -_mla_submodules.linear_kv_down_proj = TEColumnParallelLinear -_mla_submodules.linear_gate = ColumnParallelLinear diff --git a/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh b/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh deleted file mode 100755 index 5a1bed5dafa..00000000000 --- a/examples/ling_v3_tiny/train_ling_v3_tiny_8gpu.sh +++ /dev/null @@ -1,105 +0,0 @@ -#!/usr/bin/env bash - -# Full Ling-V3 Tiny fresh-initialization smoke run. - -set -euo pipefail - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)" - -PYTHON="${PYTHON:-python}" -MASTER_PORT="${MASTER_PORT:-29501}" -SEQ_LENGTH="${SEQ_LENGTH:-128}" -TRAIN_ITERS="${TRAIN_ITERS:-2}" -SAVE_ARGS=() - -if [[ -n "${SAVE_DIR:-}" ]]; then - mkdir -p "${SAVE_DIR}" - SAVE_ARGS+=(--save "${SAVE_DIR}" --save-interval "${SAVE_INTERVAL:-1000000000}") -fi - -cd "${REPO_ROOT}" -export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" -export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" - -"${PYTHON}" -m torch.distributed.run \ - --nproc_per_node=8 \ - --master_port="${MASTER_PORT}" \ - pretrain_hybrid.py \ - --spec examples.ling_v3_tiny.model_spec hybrid_stack_spec \ - --hybrid-layer-pattern "K-KEKE+EKEKEKE+EKEKEKE+EKEKEKE+EKEKEKE+EKEKEKE+E/+E" \ - --mock-data \ - --tokenizer-type NullTokenizer \ - --vocab-size 157184 \ - --seq-length "${SEQ_LENGTH}" \ - --max-position-embeddings "${SEQ_LENGTH}" \ - --micro-batch-size 1 \ - --global-batch-size 8 \ - --train-iters "${TRAIN_ITERS}" \ - --eval-interval 1000000000 \ - --eval-iters 0 \ - --log-interval 1 \ - --seed 1234 \ - --bf16 \ - --transformer-impl transformer_engine \ - --normalization RMSNorm \ - --norm-epsilon 1e-6 \ - --swiglu \ - --disable-bias-linear \ - --hidden-size 1536 \ - --num-attention-heads 16 \ - --kv-channels 128 \ - --ffn-hidden-size 4608 \ - --untie-embeddings-and-output-weights \ - --position-embedding-type rope \ - --rotary-percent 0.5 \ - --rotary-base 6000000 \ - --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 \ - --num-experts 128 \ - --moe-ffn-hidden-size 512 \ - --moe-shared-expert-intermediate-size 512 \ - --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 \ - --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 \ - --linear-conv-kernel-dim 4 \ - --linear-key-head-dim 128 \ - --linear-value-head-dim 128 \ - --linear-num-key-heads 16 \ - --linear-num-value-heads 16 \ - --kda-safe-gate \ - --kda-lower-bound -5.0 \ - --multi-latent-attention \ - --q-lora-rank 256 \ - --kv-lora-rank 512 \ - --qk-head-dim 128 \ - --qk-pos-emb-head-dim 64 \ - --v-head-dim 128 \ - --qk-layernorm \ - --attention-output-gate \ - --gated-attention-proj-granularity headwise \ - --mtp-num-layers 1 \ - --mtp-loss-scaling-factor 0.1 \ - --recompute-granularity full \ - --recompute-method uniform \ - --recompute-num-layers 1 \ - --lr 1e-4 \ - --min-lr 1e-4 \ - --lr-decay-style constant \ - --weight-decay 0.1 \ - --clip-grad 1.0 \ - --distributed-backend nccl \ - "${SAVE_ARGS[@]}" diff --git a/tests/unit_tests/models/test_ling_v3_tiny_model_spec.py b/tests/unit_tests/models/test_ling_v3_tiny_model_spec.py deleted file mode 100644 index bb31ed768c7..00000000000 --- a/tests/unit_tests/models/test_ling_v3_tiny_model_spec.py +++ /dev/null @@ -1,206 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -import pytest -import torch -import torch.nn.functional as F - -from examples.ling_v3_tiny.model_spec import LING_V3_TINY_HYBRID_LAYER_PATTERN -from examples.ling_v3_tiny.model_spec import hybrid_stack_spec as ling_v3_tiny_hybrid_stack_spec -from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TELinear -from megatron.core.models.hybrid.hybrid_block import HybridStack -from megatron.core.models.hybrid.hybrid_layer_allocation import ( - Symbols, - get_hybrid_layer_counts, - get_hybrid_total_layer_count, - parse_hybrid_pattern, -) -from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec -from megatron.core.models.hybrid.hybrid_model import HybridModel -from megatron.core.ssm.gated_delta_net import HAVE_FLA_KDA, KimiDeltaAttention -from megatron.core.ssm.mlp_layer import MLPLayer -from megatron.core.tensor_parallel import ColumnParallelLinear -from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer.multi_latent_attention import MLASelfAttention -from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper -from megatron.core.transformer.transformer_config import MLATransformerConfig -from megatron.core.transformer.transformer_layer import MoETransformerLayer -from tests.unit_tests.test_utilities import Utils - - -def _kda_submodules(spec): - return spec.submodules.kda_layer.submodules.self_attention.submodules - - -def _mla_submodules(spec): - return spec.submodules.mla_layer.submodules.self_attention.submodules - - -def test_ling_v3_tiny_hybrid_pattern_matches_model_topology(): - parsed = parse_hybrid_pattern(LING_V3_TINY_HYBRID_LAYER_PATTERN) - counts = get_hybrid_layer_counts(LING_V3_TINY_HYBRID_LAYER_PATTERN) - - assert get_hybrid_total_layer_count(LING_V3_TINY_HYBRID_LAYER_PATTERN) == 48 - assert parsed.mtp_pattern == Symbols.MLA + Symbols.MOE - assert parsed.mtp_num_depths == 1 - - main_counts = {symbol: parsed.main_pattern.count(symbol) for symbol in Symbols.VALID_LAYERS} - expected_nonzero_counts = {Symbols.KDA: 18, Symbols.MLA: 6, Symbols.MLP: 1, Symbols.MOE: 23} - assert { - symbol: count for symbol, count in main_counts.items() if count - } == expected_nonzero_counts - assert counts[Symbols.MLA] == 7 - assert counts[Symbols.MOE] == 24 - - -def test_ling_v3_tiny_spec_selects_required_projection_backends(): - default_kda = _kda_submodules(hybrid_stack_spec) - default_mla = _mla_submodules(hybrid_stack_spec) - tiny_kda = _kda_submodules(ling_v3_tiny_hybrid_stack_spec) - tiny_mla = _mla_submodules(ling_v3_tiny_hybrid_stack_spec) - - assert tiny_kda.beta_proj is ColumnParallelLinear - assert tiny_mla.linear_q_down_proj is TEColumnParallelLinear - assert tiny_mla.linear_kv_down_proj is TEColumnParallelLinear - assert tiny_mla.linear_gate is ColumnParallelLinear - - assert default_kda.beta_proj is TEColumnParallelLinear - assert default_mla.linear_q_down_proj is TELinear - assert default_mla.linear_kv_down_proj is TELinear - assert default_mla.linear_gate is None - - -@pytest.mark.internal -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available.") -@pytest.mark.skipif(not HAVE_FLA_KDA, reason="FLA with KDA support is not installed.") -@pytest.mark.parametrize("tp_size", [1, 2]) -def test_ling_v3_tiny_reduced_hybrid_model_forward_backward(tp_size): - """Exercise every Tiny layer type and the nested MTP recipe.""" - - Utils.initialize_distributed() - if torch.distributed.get_world_size() % tp_size != 0: - Utils.destroy_model_parallel() - pytest.skip("world size must be divisible by tensor parallel size") - - Utils.initialize_model_parallel(tp_size, 1) - try: - model_parallel_cuda_manual_seed(123) - config = MLATransformerConfig( - num_layers=6, - hidden_size=128, - ffn_hidden_size=256, - num_attention_heads=4, - num_query_groups=4, - kv_channels=32, - q_lora_rank=32, - kv_lora_rank=32, - qk_head_dim=32, - qk_pos_emb_head_dim=16, - v_head_dim=32, - rope_type="rope", - rotary_base=10000, - rotary_percent=0.5, - qk_layernorm=True, - attention_output_gate=True, - gated_attention_proj_granularity="headwise", - 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, - num_moe_experts=2, - moe_ffn_hidden_size=64, - moe_router_topk=2, - moe_router_load_balancing_type="none", - moe_router_score_function="sigmoid", - moe_router_dtype="fp32", - moe_router_topk_scaling_factor=1.0, - moe_shared_expert_intermediate_size=64, - moe_grouped_gemm=True, - moe_token_dispatcher_type="alltoall", - mtp_num_layers=1, - mtp_loss_scaling_factor=0.1, - normalization="RMSNorm", - layernorm_epsilon=1.0e-6, - activation_func=F.silu, - gated_linear_unit=True, - add_bias_linear=False, - attention_dropout=0.0, - hidden_dropout=0.0, - bf16=True, - params_dtype=torch.bfloat16, - pipeline_dtype=torch.bfloat16, - use_cpu_initialization=True, - is_hybrid_model=True, - calculate_per_token_loss=True, - tensor_model_parallel_size=tp_size, - sequence_parallel=tp_size > 1, - ) - model = HybridModel( - config=config, - hybrid_stack_spec=ling_v3_tiny_hybrid_stack_spec, - vocab_size=128, - max_sequence_length=16, - hybrid_layer_pattern="K-KE+E/+E", - position_embedding_type="rope", - rotary_percent=0.5, - rotary_base=10000, - ).cuda() - - main_layers = model.decoder.layers - assert isinstance(main_layers[0].self_attention, KimiDeltaAttention) - assert isinstance(main_layers[1], MLPLayer) - assert isinstance(main_layers[2].self_attention, KimiDeltaAttention) - assert isinstance(main_layers[3], MoETransformerLayer) - assert isinstance(main_layers[4].self_attention, MLASelfAttention) - assert isinstance(main_layers[5], MoETransformerLayer) - - assert isinstance(model.mtp.layers[0].mtp_model_layer, HybridStack) - mtp_layers = model.mtp.layers[0].mtp_model_layer.layers - assert isinstance(mtp_layers[0].self_attention, MLASelfAttention) - assert isinstance(mtp_layers[1], MoETransformerLayer) - - sharded_state_dict = model.sharded_state_dict() - required_sharded_keys = [ - "decoder.layers.0.self_attention.beta_proj.weight", - "decoder.layers.4.self_attention.linear_q_down_proj.weight", - "decoder.layers.4.self_attention.linear_kv_down_proj.weight", - "decoder.layers.4.self_attention.linear_gate.weight", - "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_gate.weight", - ] - for key in required_sharded_keys: - assert key in sharded_state_dict, key - - input_ids = torch.arange(16, device="cuda", dtype=torch.long).unsqueeze(0) - position_ids = torch.arange(16, device="cuda", dtype=torch.long).unsqueeze(0) - labels = (input_ids + 1) % 128 - loss_mask = torch.ones_like(input_ids, dtype=torch.float32) - loss = model( - input_ids=input_ids, - position_ids=position_ids, - attention_mask=None, - labels=labels, - loss_mask=loss_mask, - ) - - assert loss.shape == input_ids.shape - assert torch.isfinite(loss).all() - assert "loss_sums" in MTPLossLoggingHelper.tracker - assert torch.isfinite(MTPLossLoggingHelper.tracker["loss_sums"]).all() - assert "num_tokens" in MTPLossLoggingHelper.tracker - assert (MTPLossLoggingHelper.tracker["num_tokens"] > 0).all() - - loss.float().mean().backward() - required_gradients = [ - "decoder.layers.0.self_attention.beta_proj.weight", - "decoder.layers.4.self_attention.linear_gate.weight", - "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_gate.weight", - ] - parameters = dict(model.named_parameters()) - for name in required_gradients: - assert parameters[name].grad is not None, name - assert torch.isfinite(parameters[name].grad).all(), name - finally: - MTPLossLoggingHelper.tracker = {} - Utils.destroy_model_parallel() From ef8e9afff89773388d955be37c575ee53c494fbc Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 23 Aug 2026 01:45:03 +0800 Subject: [PATCH 20/34] Remove FSDP DTensor support from Ling onboarding Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit fe1cacaf413194405295bf65396fe65d8d27ce05) --- .../transformer/fsdp_dtensor_checkpoint.py | 158 ++++------------ megatron/training/checkpointing.py | 13 +- .../mfsdp_v1/test_gdn_fsdp_checkpoint.py | 176 ------------------ tests/unit_tests/test_checkpointing.py | 135 -------------- .../test_fsdp_dtensor_checkpoint.py | 89 +-------- 5 files changed, 50 insertions(+), 521 deletions(-) delete mode 100644 tests/unit_tests/distributed/mfsdp_v1/test_gdn_fsdp_checkpoint.py diff --git a/megatron/core/transformer/fsdp_dtensor_checkpoint.py b/megatron/core/transformer/fsdp_dtensor_checkpoint.py index 3bf2e4c1c10..2dde56c9c73 100644 --- a/megatron/core/transformer/fsdp_dtensor_checkpoint.py +++ b/megatron/core/transformer/fsdp_dtensor_checkpoint.py @@ -414,29 +414,27 @@ 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-family fused projections in model and optimizer state dicts. + """Handle GDN (Gated DeltaNet) 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, 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. + ``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. This is analogous to :func:`handle_swiglu_in_state_dict` which splits ``linear_fc1`` into ``weight_w`` / ``weight_v``. - Sub-key naming follows each module's ``in_proj_split_names`` and - ``in_proj_split_sections`` metadata:: + Sub-key naming follows ``GatedDeltaNet.sharded_state_dict()``:: - GDN in_proj.weight → .query / .key / .value / .z / .beta / .alpha (6-way) - KDA in_proj.weight → .query / .key / .value / .g / .gate (5-way) + in_proj.weight → .query / .key / .value / .z / .beta / .alpha (6-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"] def _strip_wrappers(path): @@ -447,117 +445,36 @@ def _strip_wrappers(path): return '.'.join(parts) # ------------------------------------------------------------------ - # 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. + # 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__). # ------------------------------------------------------------------ - 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: + _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')): continue - 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_wrappers(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), + 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_wrappers(name)] = { + 'in_proj_sizes': [qk, qk, v, v, nvh, nvh], + 'conv1d_sizes': [qk, qk, v], } - if not projection_info: + if not _gdn_info: return model_state_dict, optimizer_state_dict - parameter_map = {} - for name, parameter in model.named_parameters(): - normalized_name = _strip_wrappers(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_wrappers(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 metadata when ``key`` names a fused GDN-family parameter.""" + """Return (split_sizes, sub_names, split_dim) if *key* is a GDN fused + parameter that needs splitting, else ``None``.""" norm = _strip_wrappers(key) - 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 + for gdn_path, info in _gdn_info.items(): + if not norm.startswith(gdn_path + '.'): + continue + rel = norm[len(gdn_path) + 1 :] if rel == 'in_proj.weight': - return info['in_proj_sizes'], info['in_proj_names'], 0 + return info['in_proj_sizes'], GDN_IN_PROJ_NAMES, 0 if rel in ('conv1d.weight', 'conv1d.bias'): return info['conv1d_sizes'], GDN_CONV1D_NAMES, 0 return None @@ -571,10 +488,15 @@ def offset_slice(s, offset): return slice(s.start + offset, s.stop + offset) def split_gdn_fused(data, dist_param, split_sizes, split_dim): - """Split a fused GDN-family projection DTensor into component DTensors. + """Split a fused GDN projection DTensor into per-component DTensors. - The implementation handles both model-state DTensors and optimizer-state - tensors while preserving FSDP and tensor-parallel metadata. + Args: + data: The tensor to split. May be a DTensor (model state dict) or a + plain Tensor (optimizer states from FusedAdam). + dist_param: The corresponding model parameter (always a DTensor). + Used for global shape, numel, FSDP slice, and dist index metadata. + split_sizes: List of sizes for each component along split_dim. + split_dim: Dimension along which to split. """ total_split = sum(split_sizes) if isinstance(data, DTensor) and data.shape[split_dim] == total_split: @@ -658,7 +580,7 @@ def split_gdn_fused(data, dist_param, split_sizes, split_dim): if match is None: continue sizes, names, dim = match - dist_param = _get_parameter(key) + dist_param = model.get_parameter(f"module.{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 @@ -686,7 +608,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 = _get_parameter(key) + dist_param = model.get_parameter(key[len("module.") :]) 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/training/checkpointing.py b/megatron/training/checkpointing.py index d092e82e54e..395d5a57eca 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1379,12 +1379,15 @@ def preprocess_fsdp_dtensor_state_dict(args, raw_state_dict, model): else: model_state_dict, _ = handle_swiglu_in_state_dict(model, state_dict["model"], None) state_dict["model"] = model_state_dict - model_state_dict, optimizer_state_dict = handle_gdn_in_state_dict( - model, state_dict["model"], state_dict.get("optimizer") - ) - state_dict["model"] = model_state_dict - if optimizer_state_dict is not None: + if "optimizer" in state_dict: + model_state_dict, optimizer_state_dict = handle_gdn_in_state_dict( + model, state_dict["model"], state_dict["optimizer"] + ) + state_dict["model"] = model_state_dict state_dict["optimizer"] = optimizer_state_dict + else: + model_state_dict, _ = handle_gdn_in_state_dict(model, state_dict["model"], None) + state_dict["model"] = model_state_dict if args.num_experts: state_dict["model"] = handle_experts_in_state_dict(state_dict["model"], args.num_experts) preprocess_state_dict_for_uneven_dtensor(state_dict) 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 deleted file mode 100644 index e2c90f760c4..00000000000 --- a/tests/unit_tests/distributed/mfsdp_v1/test_gdn_fsdp_checkpoint.py +++ /dev/null @@ -1,176 +0,0 @@ -# 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/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index d5edfa7d779..5e35d3bda9f 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -8,7 +8,6 @@ import pytest import torch import torch.distributed.checkpoint -from torch.distributed.checkpoint import FileSystemReader from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel @@ -26,7 +25,6 @@ _load_base_checkpoint, get_checkpoint_tracker_filename, load_checkpoint, - preprocess_fsdp_dtensor_state_dict, read_metadata, save_checkpoint, ) @@ -50,32 +48,6 @@ def sharded_state_dict(self, *args, metadata: Optional[dict] = None, **kwargs): return self.state_dict() -class MockDeltaRule(torch.nn.Module): - """Minimal delta-rule module exposing production checkpoint metadata.""" - - def __init__(self, split_names, split_sections): - super().__init__() - self.qk_dim = 4 - self.v_dim = 6 - 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) - - -class MockDeltaRuleModel(MegatronModule): - """Minimal Megatron model used to exercise the real FSDP checkpoint path.""" - - def __init__(self, config, split_names, split_sections): - super().__init__(config=config) - self.delta_rule = MockDeltaRule(split_names, split_sections) - - def sharded_state_dict(self, *args, metadata: Optional[dict] = None, **kwargs): - return self.state_dict() - - class MockState: def __init__(self, state_dict): self._state_dict = state_dict @@ -102,113 +74,6 @@ def sharded_state_dict(self, *args, metadata: Optional[dict] = None, **kwargs): return self.state_dict() -def test_fsdp_dtensor_preprocessing_calls_delta_rule_handler(): - args = SimpleNamespace(swiglu=False, num_experts=None) - model = mock.sentinel.model - model_state = {"weight": mock.sentinel.model_weight} - optimizer_state = {"state": {}} - split_model_state = {"weight.query": mock.sentinel.query_weight} - split_optimizer_state = {"state": {"weight.query": {}}} - - with ( - mock.patch("megatron.training.checkpointing.handle_fp8_extra_state_case"), - mock.patch( - "megatron.training.checkpointing.handle_gdn_in_state_dict", - return_value=(split_model_state, split_optimizer_state), - ) as handle_gdn, - mock.patch("megatron.training.checkpointing.preprocess_state_dict_for_uneven_dtensor"), - ): - state_dict = preprocess_fsdp_dtensor_state_dict( - args, {"model": model_state, "optimizer": optimizer_state}, model - ) - - handle_gdn.assert_called_once_with(model, model_state, optimizer_state) - assert state_dict["model"] is split_model_state - assert state_dict["optimizer"] is split_optimizer_state - - -def _local_checkpoint_tensor(value): - if hasattr(value, "to_local"): - value = value.to_local() - return value.detach().clone() - - -@pytest.mark.parametrize( - ("split_names", "split_sections"), - [ - (("query", "key", "value", "z", "f", "b", "w"), (4, 4, 6, 6, 4, 4, 6)), - (("query", "key", "value", "g", "gate"), (4, 4, 6, 4, 6)), - ], -) -def test_fsdp_dtensor_delta_rule_save_load_roundtrip( - init_model_parallel, create_ckpt_load_args, tmp_path_dist_ckpt, split_names, split_sections -): - """Real Megatron-FSDP save/load preserves variant-specific fused projections.""" - - args = create_ckpt_load_args - args.ckpt_format = "fsdp_dtensor" - args.use_megatron_fsdp = True - args.use_distributed_optimizer = True - args.use_dist_ckpt = True - args.swiglu = False - args.num_experts = None - args.save_tokenizer_assets = False - - config = TransformerConfig(num_layers=1, kv_channels=1) - ddp_config = DistributedDataParallelConfig( - use_distributed_optimizer=True, use_megatron_fsdp=True - ) - - def build_model(seed): - torch.manual_seed(seed) - return FullyShardedDataParallel( - config=config, - ddp_config=ddp_config, - module=MockDeltaRuleModel(config, split_names, split_sections), - ) - - source_model = build_model(123) - source_optimizer = MockState({"state": {}}) - source_scheduler = MockState({"scheduler": "source"}) - - with TempNamedDir(tmp_path_dist_ckpt / "test_delta_rule_fsdp_roundtrip", sync=True) as ckpt_dir: - args.save = ckpt_dir - args.load = ckpt_dir - set_args(args) - save_checkpoint(7, [source_model], source_optimizer, source_scheduler, 0) - - checkpoint_dir = ckpt_dir / "iter_0000007" - metadata_keys = FileSystemReader(checkpoint_dir).read_metadata().state_dict_metadata - metadata_prefix = "model.module.delta_rule" - for name in split_names: - assert f"{metadata_prefix}.in_proj.weight.{name}" in metadata_keys - for parameter in ("weight", "bias"): - for name in ("query", "key", "value"): - assert f"{metadata_prefix}.conv1d.{parameter}.{name}" in metadata_keys - assert f"{metadata_prefix}.in_proj.weight" not in metadata_keys - assert f"{metadata_prefix}.conv1d.weight" not in metadata_keys - assert f"{metadata_prefix}.conv1d.bias" not in metadata_keys - - source_state = { - key: _local_checkpoint_tensor(value) for key, value in source_model.state_dict().items() - } - target_model = build_model(456) - target_optimizer = MockState({"state": {}}) - target_scheduler = MockState({"scheduler": "target"}) - - loaded_iteration, _ = load_checkpoint( - [target_model], target_optimizer, target_scheduler, strict=True - ) - - assert loaded_iteration == 7 - target_state = { - key: _local_checkpoint_tensor(value) for key, value in target_model.state_dict().items() - } - assert source_state.keys() == target_state.keys() - for key in source_state: - torch.testing.assert_close(target_state[key], source_state[key], rtol=0, atol=0) - - def create_checkpoint(load_path, ckpt_format): """Setup a dummy checkpoint directory.""" iteration = 123 diff --git a/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py b/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py index b616e140829..e591d7cd2d8 100644 --- a/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py +++ b/tests/unit_tests/transformer/test_fsdp_dtensor_checkpoint.py @@ -16,14 +16,13 @@ Tests cover: 1. Per-TransformerLayer gated_linear_unit check in handle_swiglu_in_state_dict - 2. GDN-family fused projection splitting in handle_gdn_in_state_dict + 2. GDN 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, -projection metadata validation) directly and cover real DTensor splitting in the -distributed M-FSDP test suite. +_match_gdn_key) by extracting the logic into standalone testable units. 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. @@ -46,7 +45,6 @@ from megatron.core.transformer.fsdp_dtensor_checkpoint import ( flatten_state_dict, get_expert_index_from_key, - handle_gdn_in_state_dict, ) @@ -414,89 +412,6 @@ 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. From a7c70aa499bf8a768b643ad75f8f4d04e78a80d6 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 23 Aug 2026 01:45:15 +0800 Subject: [PATCH 21/34] Use the generic Hybrid spec in functional coverage Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 611c8ac088270d5b1a64728358a9b53985cdac88) --- .../core/models/hybrid/hybrid_layer_specs.py | 1 + .../model_config.yaml | 16 ++++++++++------ 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index 2aa40c1cbed..c09d6887b33 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -209,6 +209,7 @@ 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/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 index 0eed047d204..5b6bfd80e9d 100644 --- 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 @@ -5,10 +5,11 @@ ENV_VARS: NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8 PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True + ENABLE_LIGHTWEIGHT_MODE: true MODEL_ARGS: - # Reduced Ling-V3 Tiny stack: 3 KDA, 1 MLA, 1 dense MLP, 3 MoE, and 1 MTP depth. - --spec: "[examples.ling_v3_tiny.model_spec hybrid_stack_spec]" + # 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. @@ -24,7 +25,7 @@ MODEL_ARGS: --eval-interval: 1000000000 --eval-iters: 0 - # Reduced model dimensions preserve Tiny's attention and FFN ratios. + # Reduced model dimensions exercise the attention and FFN paths. --hidden-size: 512 --num-attention-heads: 8 --kv-channels: 64 @@ -40,7 +41,7 @@ MODEL_ARGS: --transformer-impl: transformer_engine --attention-backend: fused - # One-node topology exercises Tiny's EP=8 and CP=2 paths. + # 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 @@ -49,7 +50,7 @@ MODEL_ARGS: --cp-comm-type: p2p --distributed-backend: nccl - # Ling-V3 Tiny MoE routing, reduced from 128 to 32 experts. + # Grouped sigmoid MoE routing with 32 experts. --num-experts: 32 --moe-ffn-hidden-size: 256 --moe-shared-expert-intermediate-size: 256 @@ -110,9 +111,12 @@ MODEL_ARGS: --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: regular +TEST_TYPE: ckpt-resume METRICS: - lm loss - mtp_1 loss From 0e06b450267943f359d5069a1574537006f145ef Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 23 Aug 2026 01:45:27 +0800 Subject: [PATCH 22/34] Refine KDA and hybrid FLOPs test contracts Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 1f6ffd19d0e47f7e1629eb15046f37d28081c272) --- .../ssm/gated_delta_net/test_kda.py | 9 +-- .../test_num_floating_point_operations.py | 61 +++++++++++-------- 2 files changed, 36 insertions(+), 34 deletions(-) diff --git a/tests/unit_tests/ssm/gated_delta_net/test_kda.py b/tests/unit_tests/ssm/gated_delta_net/test_kda.py index 1b0b5ecfc78..5427b19092b 100644 --- a/tests/unit_tests/ssm/gated_delta_net/test_kda.py +++ b/tests/unit_tests/ssm/gated_delta_net/test_kda.py @@ -14,8 +14,7 @@ ) 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, GatedDeltaNet, KimiDeltaAttention -from megatron.core.ssm.gated_delta_net.common import _GDNBase +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 @@ -32,12 +31,6 @@ _KDA_INPUT_GRAD_RELATIVE_L2 = 5e-3 -def test_kda_and_gdn_are_sibling_gdn_family_variants(): - assert issubclass(GatedDeltaNet, _GDNBase) - assert issubclass(KimiDeltaAttention, _GDNBase) - assert not issubclass(KimiDeltaAttention, GatedDeltaNet) - - def _make_config( tp_size: int = 1, cp_size: int = 1, diff --git a/tests/unit_tests/test_num_floating_point_operations.py b/tests/unit_tests/test_num_floating_point_operations.py index 7ebf6bc4e7d..f1872e3f744 100644 --- a/tests/unit_tests/test_num_floating_point_operations.py +++ b/tests/unit_tests/test_num_floating_point_operations.py @@ -100,15 +100,22 @@ def _make_hybrid_args(*, num_layers=4, hidden_size=512, num_attention_heads=8, s return args -def _make_ling_hybrid_args(pattern): - """Minimal direct-KDA/MLA dimensions for hybrid FLOPs tests.""" +def _make_kda_hybrid_args(): + """Minimal KDA dimensions for hybrid FLOPs tests.""" args = _make_hybrid_args() - args.hybrid_layer_pattern = pattern + 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 @@ -415,11 +422,11 @@ def configure(args): self._assert_match(configure) -class TestLingHybridAttentionFlops: - """KDA and MLA layers must contribute to Ling-V3 Tiny MFU accounting.""" +class TestKimiDeltaAttentionFlops: + """KDA layers must contribute their projection and kernel work.""" - def test_kda_direct_projection_formula(self): - args = _make_ling_hybrid_args("K") + 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 @@ -445,8 +452,12 @@ def test_kda_direct_projection_formula(self): assert num_floating_point_operations(args, batch_size) == 3 * (kda_forward + logits_forward) - def test_mla_uses_ragged_attention_work(self): - args = _make_ling_hybrid_args("+") + +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( @@ -463,31 +474,29 @@ def test_mla_uses_ragged_attention_work(self): ) assert flops_full - flops_half == expected_delta - def test_mla_gate_granularity_changes_only_projection_work(self): - args = _make_ling_hybrid_args("+") + def test_headwise_gate_projection_work(self): + args = _make_mla_hybrid_args() batch_size = 2 total_tokens = batch_size * args.seq_length - args.gated_attention_proj_granularity = "headwise" - headwise_flops = num_floating_point_operations(args, batch_size) - args.gated_attention_proj_granularity = "elementwise" - elementwise_flops = num_floating_point_operations(args, batch_size) + 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 - expected_delta = ( - 3 - * 2 - * total_tokens - * args.hidden_size - * args.num_attention_heads - * (args.v_head_dim - 1) - ) - assert elementwise_flops - headwise_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_hybrid_regular_attention_output_gate(self, gate_granularity, gate_projection_size): - args = _make_ling_hybrid_args("*") + 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 From c723eba8157b4c4480cd31d4accf9cab661782e5 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 23 Aug 2026 01:45:40 +0800 Subject: [PATCH 23/34] Correct MLA headwise gate VJP reference Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 4a68b8a8e10252ebbef1506230c6c36d90e89aca) --- .../unit_tests/transformer/test_multi_latent_attention.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index e6dbcfc1caf..88562a78fae 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/tests/unit_tests/transformer/test_multi_latent_attention.py @@ -833,12 +833,7 @@ def test_gate_matches_fp32_sigmoid_reference_vjp(self, gate_granularity): native_scale = torch.repeat_interleave( native_scale, self.transformer_config.v_head_dim, dim=-1 ) - expanded_fp32_scale = torch.repeat_interleave( - sigmoid_fp32, self.transformer_config.v_head_dim, dim=-1 - ) - expected_core_gradient = (output_gradient.float() * expanded_fp32_scale).to( - core_output.dtype - ) + 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) From 7a31ac0b8005a76ba2750e12a61ffe1e7024ae7b Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Sun, 23 Aug 2026 12:10:52 +0800 Subject: [PATCH 24/34] Correct KDA source attribution Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 87c96f8d372276a92a242d8a92f93669249118c5) --- megatron/core/ssm/gated_delta_net/kda.py | 1 - 1 file changed, 1 deletion(-) diff --git a/megatron/core/ssm/gated_delta_net/kda.py b/megatron/core/ssm/gated_delta_net/kda.py index 0019c406e5c..39e64539f47 100644 --- a/megatron/core/ssm/gated_delta_net/kda.py +++ b/megatron/core/ssm/gated_delta_net/kda.py @@ -1,5 +1,4 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh. """Kimi Delta Attention, a channel-wise Gated DeltaNet variant.""" From 8ff53072341ec725bf753bab36a1551d282cd2c7 Mon Sep 17 00:00:00 2001 From: Fan Kun <784819644@qq.com> Date: Mon, 24 Aug 2026 11:33:54 +0800 Subject: [PATCH 25/34] Update megatron/training/training.py Co-authored-by: Yuzhong Wang Signed-off-by: Fan Kun <784819644@qq.com> (cherry picked from commit be824cc2b3e16eb5dd7cdf1c86dec4f43d1c5a06) --- megatron/training/training.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/training/training.py b/megatron/training/training.py index 3db5fff8913..85fc717d901 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1291,8 +1291,8 @@ def transformer_flops(): num_mlp_layers = layer_counts[Symbols.MLP] num_moe_layers = layer_counts[Symbols.MOE] - # Ling uses + for dense MLA. DSv4 D/C/H/W layers are MLA variants too, - # but their sparse attention FLOPs are handled by the DSv4 branch. + # 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] From c7f07f75c84eec5867eb0ec2db7633d71bfdcef4 Mon Sep 17 00:00:00 2001 From: Fan Kun <784819644@qq.com> Date: Mon, 24 Aug 2026 11:34:10 +0800 Subject: [PATCH 26/34] Update megatron/core/transformer/multi_latent_attention.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Fan Kun <784819644@qq.com> (cherry picked from commit aa516dc6751e2691d9bf386b8088af5ffbc2cf90) --- megatron/core/transformer/multi_latent_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index ff3ae5991de..97eee00aac2 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -1197,7 +1197,7 @@ def backward_dw(self) -> NoReturn: # 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. + # 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() From de4b5387b9b6891eca9699b02ea3575c5e616b48 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Mon, 24 Aug 2026 11:47:20 +0800 Subject: [PATCH 27/34] Fix hybrid layer count doctest examples Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit bd98edc192eac975b4882035e87926f69f64f06e) --- megatron/core/models/hybrid/hybrid_layer_allocation.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index 65ff8cdf8c5..6aa866ba158 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -180,10 +180,10 @@ def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]: Examples: >>> get_hybrid_layer_counts("M*M*") - {'*': 2, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'K': 0, 'M': 2, '-': 0, 'E': 0, 'W': 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, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'K': 0, 'M': 8, '-': 4, 'E': 0, 'W': 0} + >>> 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} """ parsed = parse_hybrid_pattern(pattern) counts = {symbol: 0 for symbol in Symbols.name_sorted_valid_layer_symbols()} From e51cb513206d8bd6b48ebe2273497de39665cba8 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Mon, 24 Aug 2026 15:21:24 +0800 Subject: [PATCH 28/34] Fix Ling-V3 Tiny CI assertions Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit f4e113e3e2c841a4d4ef01271bb0a853e1c13b21) --- tests/unit_tests/models/test_hybrid_moe_model.py | 3 +++ tests/unit_tests/transformer/test_multi_latent_attention.py | 4 ++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 2eac69295e1..1113c296e34 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -141,6 +141,7 @@ "fp8_wgrad": True, "fused_residual_rmsnorm": False, "fused_single_qkv_rope": False, + "gated_attention_proj_granularity": "elementwise", "gated_linear_unit": False, "gdp_cutedsl_kernel": False, "gdp_num_chunk_states_to_recompute": 2, @@ -165,6 +166,8 @@ "init_method_std": 0.014, "init_model_with_meta_device": False, "is_hybrid_model": True, + "kda_lower_bound": None, + "kda_safe_gate": False, "kitchen_attention_backend": "sdpa", "kv_channels": 128, "layernorm_epsilon": 1e-05, diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index 88562a78fae..f4c07430400 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/tests/unit_tests/transformer/test_multi_latent_attention.py @@ -851,8 +851,8 @@ def test_gate_matches_fp32_sigmoid_reference_vjp(self, gate_granularity): ) torch.testing.assert_close(actual_output, expected_output, atol=0, rtol=0) - torch.testing.assert_close(actual_gradients[0], expected_gate_gradient, atol=0, rtol=0) - torch.testing.assert_close(actual_gradients[1], expected_core_gradient, 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() From 2e6d0a651355a206ccd65bfe768d53fdfd2e2235 Mon Sep 17 00:00:00 2001 From: fkuner <784819644@qq.com> Date: Mon, 24 Aug 2026 17:43:19 +0800 Subject: [PATCH 29/34] Remove unused KDA parameter markers Signed-off-by: fkuner <784819644@qq.com> (cherry picked from commit 14e28c6ed58f096937fa904e1c0d29736d75ccb8) --- megatron/core/ssm/gated_delta_net/kda.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net/kda.py b/megatron/core/ssm/gated_delta_net/kda.py index 39e64539f47..04cb3ae439d 100644 --- a/megatron/core/ssm/gated_delta_net/kda.py +++ b/megatron/core/ssm/gated_delta_net/kda.py @@ -116,9 +116,6 @@ def __init__( # FP32 through model casting, optimizer construction, and checkpointing. mark_keep_in_fp32(self.dt_bias) mark_keep_in_fp32(self.A_log) - self.dt_bias._no_reinit = True - self.dt_bias._no_weight_decay = True - self.A_log._no_weight_decay = True # Pass raw g to FLA and let the KDA kernel apply A_log and dt_bias. self.use_gate_in_kernel = True From daddbfcdebd65be95741eed9fdbee586bb73c91e Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Thu, 27 Aug 2026 23:00:29 +0000 Subject: [PATCH 30/34] chore: preserve current inference context while applying KDA Signed-off-by: svcnemo-autobot --- .../inference/contexts/dynamic_context.py | 1716 +++++++++++++---- 1 file changed, 1378 insertions(+), 338 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index bc699392830..7f516bb584a 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1,17 +1,18 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging import math import operator import warnings from contextlib import nullcontext -from typing import Dict, List, Optional, Sequence, Tuple +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple import torch # type: ignore import torch.nn.functional as F # type: ignore from torch import Tensor # type: ignore from megatron.core import parallel_state +from megatron.core.inference.batch_dimensions_utils import TOKEN_ROUNDER as _TOKEN_ROUNDER from megatron.core.inference.batch_dimensions_utils import ( CUDAGraphBatchDimensionBuilder, InferenceBatchDimensions, @@ -22,6 +23,8 @@ PrefixCachingEvictionPolicy, ) from megatron.core.inference.inference_request import DynamicInferenceRequest +from megatron.core.inference.moe import InferenceGroupedGemmBackend +from megatron.core.inference.moe.vllm_fused_moe import VllmFusedMoeBuffers from megatron.core.inference.sampling.base import Sampling from megatron.core.inference.sampling_params import SamplingParams from megatron.core.inference.unified_memory import ( @@ -35,8 +38,11 @@ get_layer_maps_from_layer_type_list, ) from megatron.core.package_info import __version__ as mcore_version +from megatron.core.ssm.ops.gdp.metadata import max_gdp_chunk_counts from megatron.core.transformer import MLATransformerConfig, TransformerConfig +from megatron.core.transformer.enums import InferenceCudaGraphScope from megatron.core.transformer.moe.token_dispatcher_inference import ( + InferenceAllGatherDispatcherBase, NCCLAllGatherDispatcher, NVLSAllGatherVDispatcher, ) @@ -49,9 +55,12 @@ from .base_context import BaseInferenceContext from .gpu_view import ContextGPUView from .kv_block_allocator import KVBlockAllocator -from .mamba_slot_allocator import MambaSlotAllocator +from .mamba_slot_allocator import MAX_INTERMEDIATE_OFFSETS_PER_REQUEST, MambaSlotAllocator from .routing_metadata import RoutingMetadata +# These callbacks are currently consumed only by the Dynamo frontend. +KVEventListener = Callable[[str, dict[str, Any]], None] + try: from .fused_kv_append_kernel import triton_append_key_value_cache except ImportError: @@ -224,6 +233,69 @@ def get_mem_size_str(n_bytes: int) -> str: raise Exception(f"something went wrong, n_bytes={n_bytes}.") +class DynamoHelper: + """Manage KV-cache lifecycle events consumed by the Dynamo frontend.""" + + def __init__(self) -> None: + self._kv_event_listeners: list[KVEventListener] = [] + self._pending_kv_stored_events: list[dict[str, Any]] = [] + + @property + def has_kv_event_listeners(self) -> bool: + """Return whether any KV-event listeners are registered.""" + return bool(self._kv_event_listeners) + + def add_kv_event_listener(self, listener: KVEventListener) -> None: + """Register a KV-cache lifecycle listener. + + Args: + listener: Callback invoked with the event kind and payload. + """ + self._kv_event_listeners.append(listener) + + def queue_kv_stored_event(self, payload: dict[str, Any]) -> None: + """Queue a stored event for publication after a successful forward pass. + + Args: + payload: Stored-event payload. + """ + self._pending_kv_stored_events.append(payload) + + def publish_pending_kv_stored_events(self) -> None: + """Publish blocks whose KV contents were produced by a successful forward pass.""" + pending, self._pending_kv_stored_events = self._pending_kv_stored_events, [] + for payload in pending: + self._emit_kv_event("stored", payload) + + def discard_pending_kv_stored_events(self) -> None: + """Discard registrations left by an interrupted or failed forward pass.""" + self._pending_kv_stored_events.clear() + + def on_kv_blocks_deregistered(self, _block_ids: list[int], hashes: set[int]) -> None: + """Publish removal events for deregistered KV blocks. + + Args: + _block_ids: Deregistered block IDs, unused by Dynamo. + hashes: Hashes of the deregistered blocks. + """ + if hashes: + self._emit_kv_event("removed", {"block_hashes": list(hashes)}) + + def notify_kv_cache_cleared(self) -> None: + """Notify listeners that no previously advertised block is routable.""" + self.discard_pending_kv_stored_events() + if self._kv_event_listeners: + self._emit_kv_event("cleared", {}) + + def _emit_kv_event(self, kind: str, payload: dict[str, Any]) -> None: + """Notify Dynamo listeners without allowing frontend failures to stop inference.""" + for listener in tuple(self._kv_event_listeners): + try: + listener(kind, payload) + except Exception: # pylint: disable=broad-exception-caught + logging.exception("KV-event listener failed while handling %r", kind) + + @internal_api # pylint: disable=line-too-long class DynamicInferenceContext(BaseInferenceContext): @@ -246,7 +318,7 @@ class DynamicInferenceContext(BaseInferenceContext): """ DEFAULT_MAX_TOKENS = 16384 - TOKEN_ROUNDER = 64 + TOKEN_ROUNDER = _TOKEN_ROUNDER REQUEST_ROUNDER = 4 TMS_TAG = "inference_context" @@ -267,6 +339,8 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Hyperparameter for choosing to prioritize prefix hit matches vs minimizing idle load self.prefix_caching_routing_alpha = inference_config.prefix_caching_routing_alpha + self.media_cache_coordinator_policy = inference_config.media_cache_coordinator_policy + self.media_cache_routing_weight = inference_config.media_cache_routing_weight # Monotonic clock for prefix caching LRU eviction ordering. # Incremented each engine step but kept independent so the engine step @@ -276,12 +350,24 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Prefix caching hit tracking (accumulated, reset by engine after logging). self.prefix_cache_hits = 0 # requests that matched at least one cached block self.prefix_cache_blocks_matched = 0 # total matched blocks across all requests + # Prefill compute accounting (drained into engine accumulators each step). + # computed = prompt tokens actually run through the model this step; + # skipped = prompt tokens whose prefill was skipped via a prefix-cache hit. + # A high skipped fraction confirms prefix caching is saving prefill compute + # (so any per-step latency growth is attention-over-context, not re-prefill). + self.prefix_cache_prefill_computed_tokens = 0 + self.prefix_cache_prefill_skipped_tokens = 0 # Engine step counter (used for logging, metrics, and event tracking) self.step_count = 0 self.async_sched_step_count = 0 self.async_sched_compaction_step_count = 0 + # Per-request VLM data (empty when not using multimodal). + self._request_to_image_embeddings: Dict[int, Optional[Tensor]] = {} + self._request_to_image_token_mask: Dict[int, Optional[Tensor]] = {} + self._request_to_image_token_count: Dict[int, int] = {} + self.cache_mla_latent = ( isinstance(model_config, MLATransformerConfig) and model_config.cache_mla_latents ) @@ -306,11 +392,13 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC else: self.num_attention_heads_per_partition = 1 + self.batch_invariant_mode = model_config.batch_invariant_mode self.num_speculative_tokens = inference_config.num_speculative_tokens assert self.num_speculative_tokens < inference_config.block_size_tokens, ( f"num_speculative_tokens ({self.num_speculative_tokens}) must be < " f"block_size_tokens ({inference_config.block_size_tokens})" ) + self._async_sched_token_offsets = None # Cache the PP group we should use for PP collectives inside the context. # If the model provides a pg_collection with a pp group, prefer it. @@ -340,30 +428,53 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Mamba states. mamba_inference_state_config = inference_config.mamba_inference_state_config self.is_hybrid_model = mamba_inference_state_config is not None + self.gdp_num_householder = 0 if self.is_hybrid_model: self.mamba_conv_states_shape = mamba_inference_state_config.conv_states_shape self.mamba_ssm_states_shape = mamba_inference_state_config.ssm_states_shape self.mamba_conv_states_dtype = mamba_inference_state_config.conv_states_dtype self.mamba_ssm_states_dtype = mamba_inference_state_config.ssm_states_dtype self.mamba_chunk_size = mamba_inference_state_config.mamba_chunk_size + self.ssm_chunk_alignment = mamba_inference_state_config.ssm_chunk_alignment + self.gdp_num_householder = mamba_inference_state_config.gdp_num_householder + + if self.batch_invariant_mode: + # Gated Delta Product does not implement batch-invariant mode yet. + assert self.gdp_num_householder == 0, ( + "batch_invariant_mode does not support Gated Delta Product layers; " + "set batch_invariant_mode=False." + ) + assert not self.enable_prefix_caching, ( + "batch_invariant_mode does not support Mamba prefix caching; " + "set enable_prefix_caching=False." + ) + assert self.num_speculative_tokens == 0, ( + "batch_invariant_mode for Mamba dynamic inference only supports " + "one-token decode; set num_speculative_tokens=0." + ) + assert self.mamba_ssm_states_dtype == torch.float32, ( + "batch_invariant_mode requires FP32 Mamba SSM states so state-passing " + "boundaries are not rounded between decode chunks." + ) - # For hybrid models, the layer map converts the global layer index to the - # corresponding attention layer index or Mamba layer index depending on the - # layer type. - attention_layer_map, dsa_layer_map, gdn_layer_map, kda_layer_map, mamba_layer_map = ( - operator.itemgetter( - Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.GDN, Symbols.KDA, Symbols.MAMBA - )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)) - ) - - if len(gdn_layer_map) > 0: - raise NotImplementedError("GDN layers are not supported for inference.") - if len(kda_layer_map) > 0: + # 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. + attention_layer_map, dsa_layer_map = operator.itemgetter( + Symbols.ATTENTION, Symbols.DS_ATTENTION + )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)) + if Symbols.KDA in mamba_inference_state_config.layer_type_list: raise NotImplementedError("KDA layers are not supported for inference.") + recurrent_layer_map = {} + for global_layer_idx, layer_type in enumerate( + mamba_inference_state_config.layer_type_list + ): + if layer_type in (Symbols.MAMBA, Symbols.GDN): + recurrent_layer_map[global_layer_idx] = len(recurrent_layer_map) self.num_attention_layers = len(attention_layer_map) + len(dsa_layer_map) - self.num_mamba_layers = len(mamba_layer_map) - self.layer_map = attention_layer_map | dsa_layer_map | mamba_layer_map + self.num_mamba_layers = len(recurrent_layer_map) + self.layer_map = attention_layer_map | dsa_layer_map | recurrent_layer_map else: # The layer map is the identity function for pure Transformer models. # Use the same per-PP-rank layer count as TransformerBlock (handles @@ -492,7 +603,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC paused_buffer_size_bytes = int(paused_buffer_size_bytes * (1.0 - mamba_memory_ratio)) block_count = buffer_size_bytes // self.block_size_bytes - block_count = max(2, block_count) # need >= 1 active block + 1 dummy block + block_count = max(2, block_count) # need >= 1 usable block + 1 dummy block paused_block_count = paused_buffer_size_bytes // self.block_size_bytes elif self.is_hybrid_model and inference_config.max_requests is not None: # Auto-derive mamba/KV split from max_requests. Allocate exactly enough @@ -506,19 +617,19 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC ) mamba_max_requests = inference_config.max_requests - # Subtract mamba memory proportionally from active and paused buffers. + # Subtract Mamba memory proportionally from both configured buffer inputs. mamba_memory_ratio = mamba_memory_needed / total_memory buffer_size_bytes = int(buffer_size_bytes * (1.0 - mamba_memory_ratio)) paused_buffer_size_bytes = int(paused_buffer_size_bytes * (1.0 - mamba_memory_ratio)) block_count = buffer_size_bytes // self.block_size_bytes - block_count = max(2, block_count) # need >= 1 active block + 1 dummy block + block_count = max(2, block_count) # need >= 1 usable block + 1 dummy block paused_block_count = paused_buffer_size_bytes // self.block_size_bytes else: block_count = buffer_size_bytes // ( self.block_size_bytes + mamba_states_memory_per_request ) - block_count = max(2, block_count) # need >= 1 active block + 1 dummy block + block_count = max(2, block_count) # need >= 1 usable block + 1 dummy block paused_block_count = paused_buffer_size_bytes // ( self.block_size_bytes + mamba_states_memory_per_request ) @@ -543,13 +654,17 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.kv_block_allocator = KVBlockAllocator( context=self, - total_count=( + pool_size=( block_count if self.unified_memory_level == 0 else block_count + paused_block_count ), - paused_count=paused_block_count, + paused_limit=paused_block_count, enable_prefix_caching=self.enable_prefix_caching, prefix_caching_eviction_policy=self.prefix_caching_eviction_policy, ) + self.dynamo_helper = DynamoHelper() + self.kv_block_allocator.add_blocks_deregistered_observer( + self.dynamo_helper.on_kv_blocks_deregistered + ) # Track request metadata. request_metadata_types = inference_config.request_metadata_types @@ -559,6 +674,8 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Initialize context state. self.params_dtype = model_config.params_dtype + self.hidden_size = model_config.hidden_size + self.inference_cuda_graph_scope = model_config.inference_cuda_graph_scope self.max_sequence_length = inference_config.max_sequence_length # Block ids. With speculative decoding, blocks are pre-allocated when the @@ -571,7 +688,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Set max_requests, max_tokens. if inference_config.max_requests is None: # Maximize compute utilization by defaulting to 1 block per request. - self.max_requests = self.kv_block_allocator.total_count - 1 # -1 for dummy block + self.max_requests = self.kv_block_allocator.pool_size - 1 # -1 for dummy block # Adjust max_requests for Mamba memory constraints if necessary if self.is_hybrid_model and mamba_max_requests < self.max_requests: @@ -589,6 +706,15 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.max_tokens = inference_config.max_tokens or self.DEFAULT_MAX_TOKENS + # Per-step upper bound on Mamba intermediate-state extractions, shared with + # MambaMetadata and MambaSlotAllocator so scratch/metadata buffers and the + # budget accounting agree. Bounded both by the token budget (one block + # boundary per block_size_tokens) and by the request budget + # (MAX_INTERMEDIATE_OFFSETS_PER_REQUEST per request); + token_based_count = math.ceil(self.max_tokens / self.block_size_tokens) + request_based_count = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * self.max_requests + self.max_mamba_intermediate_states_per_step = min(token_based_count, request_based_count) + assert self.max_tokens >= self.max_requests, ( f"max_tokens ({self.max_tokens}) must be >= " f"max_requests ({self.max_requests}), " @@ -601,7 +727,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.non_graph_attn_metadata = {} self.graph_attn_metadata["mha_metadata"] = GraphedMHAMetadata( - block_count_total=self.kv_block_allocator.total_count, + block_count_total=self.kv_block_allocator.pool_size, max_kv_block_count=self.max_kv_block_count, max_requests=self.max_requests, block_size_tokens=self.block_size_tokens, @@ -609,7 +735,7 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC ) self.non_graph_attn_metadata["mha_metadata"] = NonGraphedMHAMetadata( - block_count_total=self.kv_block_allocator.total_count, + block_count_total=self.kv_block_allocator.pool_size, max_kv_block_count=self.max_kv_block_count, max_requests=self.max_requests, block_size_tokens=self.block_size_tokens, @@ -629,6 +755,12 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC and model_config.inference_moe_token_dispatcher_type == 'nccl' ) + # are we using the inference_optimized nvls ep dispatcher for MoEs? + self._nvls_dispatcher = ( + get_pg_size(self.expert_model_parallel_group) > 1 + and model_config.inference_moe_token_dispatcher_type == 'nvls' + ) + # are we using the training a2a dispatcher for MoEs? # Note that this is not optimal for speed. self._training_ep_dispatcher = ( @@ -647,12 +779,15 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC ) # CUDA graph token budget for prefill/mixed graphs. Decode graphs are always - # capped at max_requests * (num_speculative_tokens + 1) inside the helper; this - # only widens the prefill/mixed range when `cuda_graph_all_prefills` is set. + # capped at max_requests * (num_speculative_tokens + 1) inside the helper. By + # default the prefill/mixed range is bounded by `cuda_graph_max_tokens`, clamped + # to never fall below that decode bound nor exceed `max_tokens`; setting + # `cuda_graph_all_prefills` widens the range to the full `max_tokens`. + decode_bound = self.max_requests * (self.num_speculative_tokens + 1) cuda_graph_max_tokens = ( self.max_tokens if inference_config.cuda_graph_all_prefills - else self.max_requests * (self.num_speculative_tokens + 1) + else min(max(inference_config.cuda_graph_max_tokens, decode_bound), self.max_tokens) ) # CUDA graph config list. @@ -673,22 +808,58 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Allocate per-step dispatcher buffers upfront so update_metadata never # triggers an allocation inside a captured CUDA graph. - if get_pg_size(self.expert_model_parallel_group) > 1: - if self._nccl_ep_dispatcher: - NCCLAllGatherDispatcher.allocate_buffers() - else: - # Use moe_latent_size if set (latent MoE: SuperV3, UltraV3), else hidden_size. - moe_hidden_size = model_config.moe_latent_size or model_config.hidden_size - NVLSAllGatherVDispatcher.allocate_buffers( - per_rank_worst_case_token_count=self.round_up_tokens(self.max_tokens) - // tp_size, - topk=model_config.moe_router_topk, - hidden_size=moe_hidden_size, - ep_group=self.expert_model_parallel_group, - ) + # + # The shared _valid_tokens_tensor scalar is read as a pointer by both fused + # MoE backends (mcore_fused_moe and vllm_fused_moe) regardless of EP size, so + # allocate it unconditionally (covers EP=1, where no dispatcher comm buffers + # exist). The EP>1 dispatchers below reallocate it as part of their own buffer + # setup, which is harmless. + InferenceAllGatherDispatcherBase.allocate_valid_tokens_tensor() + if self._nccl_ep_dispatcher: + NCCLAllGatherDispatcher.allocate_buffers() + elif self._nvls_dispatcher: + # Use moe_latent_size if set (latent MoE: SuperV3, UltraV3), else hidden_size. + moe_hidden_size = model_config.moe_latent_size or model_config.hidden_size + NVLSAllGatherVDispatcher.allocate_buffers( + per_rank_worst_case_token_count=self.round_up_tokens(self.max_tokens) // tp_size, + topk=model_config.moe_router_topk, + hidden_size=moe_hidden_size, + ep_group=self.expert_model_parallel_group, + ) + + # Pre-allocate the vLLM fused-MoE intermediates so no allocation happens + # inside CUDA graph capture; one buffer set is shared by all MoE layers + # and graphs. Like the dispatcher buffers above, these persist across + # engine suspend/resume. + if ( + model_config.num_moe_experts + and model_config.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.VLLM + ): + ep_size = get_pg_size(self.expert_model_parallel_group) + moe_hidden_size = model_config.moe_latent_size or model_config.hidden_size + # Worst-case rows entering the MoE: the fixed NVLS AGV buffer height + # (per-rank worst case * ep_size); max_tokens covers the EP=1 / NCCL paths. + moe_max_rows = max( + self.max_tokens, self.round_up_tokens(self.max_tokens) // tp_size * ep_size + ) + VllmFusedMoeBuffers.allocate_buffers( + max_tokens=moe_max_rows, + topk=model_config.moe_router_topk, + fc1_output_size=model_config.moe_ffn_hidden_size + * (2 if model_config.gated_linear_unit else 1), + hidden_size=moe_hidden_size, + num_local_experts=model_config.num_moe_experts // ep_size, + ) # Deal with chunked prefill self.enable_chunked_prefill = inference_config.enable_chunked_prefill + if self.batch_invariant_mode and self.is_hybrid_model and self.enable_chunked_prefill: + # A chunk plus its final token must fit in one step; otherwise a prompt + # of that length can never advance without an invalid one-token tail. + assert self.max_tokens > self.ssm_chunk_alignment, ( + "batch-invariant SSM chunked prefill requires max_tokens > " + f"ssm_chunk_alignment ({self.ssm_chunk_alignment})." + ) # FlashInfer. if inference_config.use_flashinfer_fused_rope is True: @@ -698,17 +869,30 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC self.use_flashinfer_fused_rope = inference_config.use_flashinfer_fused_rope self.inference_grouped_gemm_backend = model_config.inference_grouped_gemm_backend + # Placeholder for the MTP decoder hidden-states buffer; allocated inside + # initialize_all_tensors() when num_speculative_tokens > 0. + self.mtp_decoder_hidden_states = None + # Allocate GPU state. self.is_tensor_state_allocated = False + self._bookkeeping_no_real_work = False self.initialize_all_tensors() + # Bind the GPU real-token-count tensor onto the NVLS dispatcher class + # so it can mask out CUDA-graph padding tokens during routing. The + # tensor lives inside gpu_view._buf (fixed address) and is refreshed + # each step by transfer_bookkeeping_to_gpu(). NVLS-only — the NCCL + # dispatcher requires equal token counts across ranks already. + if self._nvls_dispatcher: + NVLSAllGatherVDispatcher.set_real_token_count_tensor(self.gpu_view.real_token_count) + # Print info. - active_blocks = self.kv_block_allocator.active_count - total_blocks = self.kv_block_allocator.total_count - paused_blocks = self.kv_block_allocator.paused_count - active_kv_bytes = active_blocks * self.block_size_bytes - total_kv_bytes = total_blocks * self.block_size_bytes - paused_kv_bytes = paused_blocks * self.block_size_bytes + pool_size = self.kv_block_allocator.pool_size + usable_blocks = pool_size - 1 + paused_limit = self.kv_block_allocator.paused_limit + pool_size_bytes = pool_size * self.block_size_bytes + usable_bytes = usable_blocks * self.block_size_bytes + paused_limit_bytes = paused_limit * self.block_size_bytes log_lines = [ "DynamicInferenceContext: configuration summary", @@ -719,9 +903,10 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC f" max_kv_blocks_per_req: {self.max_kv_block_count}", f" KV cache:", f" block_size_bytes: {get_mem_size_str(self.block_size_bytes)}", - f" active_blocks: {active_blocks} ({get_mem_size_str(active_kv_bytes)})", - f" paused_blocks: {paused_blocks} ({get_mem_size_str(paused_kv_bytes)})", - f" total_blocks: {total_blocks} ({get_mem_size_str(total_kv_bytes)})", + f" pool_size: {pool_size} ({get_mem_size_str(pool_size_bytes)})", + f" usable_blocks: {usable_blocks} ({get_mem_size_str(usable_bytes)})", + f" paused_limit: {paused_limit} " + f"({get_mem_size_str(paused_limit_bytes)})", ] if self.is_hybrid_model: @@ -762,12 +947,21 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC and prefix_caching_mamba_gb is not None and prefix_caching_mamba_gb > 0 ): + assert self.mamba_slot_allocator is not None prefix_cache_bytes = int(prefix_caching_mamba_gb * 1024**3) - prefix_cache_slots = prefix_cache_bytes // mamba_bytes_per_req + # The allocator contains the PP-synchronized durable capacity. + # Scratch remains stage-local because it is temporary forward + # storage and does not participate in prefix-cache matching. + scratch_slots = self.max_mamba_intermediate_states_per_step + scratch_bytes = scratch_slots * mamba_bytes_per_req + durable_slots = self.mamba_slot_allocator.max_slots log_lines += [ f" Mamba prefix cache:", f" budget: {get_mem_size_str(prefix_cache_bytes)}", - f" slots: {prefix_cache_slots}", + f" extraction_scratch: {scratch_slots} slots " + f"({get_mem_size_str(scratch_bytes)})", + f" durable_slots: {durable_slots} " + f"({get_mem_size_str(durable_slots * mamba_bytes_per_req)})", f" per_slot: {get_mem_size_str(mamba_bytes_per_req)}", ] @@ -780,7 +974,7 @@ def _allocate_memory_buffer(self): self.memory_buffer = torch.empty( ( self.num_attention_layers, - self.kv_block_allocator.total_count, + self.kv_block_allocator.pool_size, self.block_size_tokens, self.kv_reduced_dim, ), @@ -792,7 +986,7 @@ def _allocate_memory_buffer(self): ( 2, # key and value self.num_attention_layers, - self.kv_block_allocator.total_count, + self.kv_block_allocator.pool_size, self.block_size_tokens, self.num_attention_heads_per_partition, self.hidden_size_per_attention_head, @@ -816,25 +1010,35 @@ def _allocate_mamba_states(self): self.mamba_metadata = MambaMetadata( max_requests=self.max_requests, max_tokens=self.max_tokens, + max_intermediate_count=self.max_mamba_intermediate_states_per_step, mamba_chunk_size=self.mamba_chunk_size, d_conv=self.mamba_conv_states_shape[-1], + decode_indices_dtype=self._mamba_decode_indices_dtype, + gdp_num_householder=self.gdp_num_householder, ) # Bind the unified CPU/GPU buffers so the per-step Mamba metadata # fields ride along with the single coalesced H2D in # transfer_bookkeeping_to_gpu(). - self.mamba_metadata.bind_cpu_buffers( - { - "batch_indices_decode": self._cpu_mamba_batch_indices_decode, - "batch_indices_prefill": self._cpu_mamba_batch_indices_prefill, - "seq_idx": self._cpu_mamba_seq_idx, - "cu_seqlens": self._cpu_mamba_cu_seqlens, - "cu_chunk_seqlens": self._cpu_mamba_cu_chunk_seqlens, - "last_chunk_indices": self._cpu_mamba_last_chunk_indices, - "seq_idx_for_varlen": self._cpu_mamba_seq_idx_for_varlen, - "conv_seq_idx": self._cpu_mamba_conv_seq_idx, - "conv_seq_start": self._cpu_mamba_conv_seq_start, - } - ) + _mamba_cpu_bufs = { + "batch_indices_decode": self._cpu_mamba_batch_indices_decode, + "batch_indices_prefill": self._cpu_mamba_batch_indices_prefill, + "seq_idx": self._cpu_mamba_seq_idx, + "cu_seqlens": self._cpu_mamba_cu_seqlens, + "cu_chunk_seqlens": self._cpu_mamba_cu_chunk_seqlens, + "last_chunk_indices": self._cpu_mamba_last_chunk_indices, + "seq_idx_for_varlen": self._cpu_mamba_seq_idx_for_varlen, + "conv_seq_idx": self._cpu_mamba_conv_seq_idx, + "conv_seq_start": self._cpu_mamba_conv_seq_start, + } + if self.gdp_num_householder > 0: + _mamba_cpu_bufs.update( + { + "gdp_chunk_indices": self._cpu_gdp_chunk_indices, + "gdp_chunk_indices_dp": self._cpu_gdp_chunk_indices_dp, + "gdp_chunk_offsets": self._cpu_gdp_chunk_offsets, + } + ) + self.mamba_metadata.bind_cpu_buffers(_mamba_cpu_bufs) self.mamba_metadata.bind_gpu_buffers(self.gpu_view) self.mamba_conv_states = torch.empty( (self.num_mamba_layers, self.max_requests) + self.mamba_conv_states_shape, @@ -948,6 +1152,9 @@ def initialize_all_tensors(self) -> None: device='cpu', pin_memory=True, ) + self._async_sched_token_offsets = torch.arange( + self.num_speculative_tokens + 1, device='cpu' + ) # Track request metadata. Backed by pinned CPU memory: bookkeeping is # CPU-resident; GPU consumers read from the active-slice mirror in @@ -995,6 +1202,10 @@ def initialize_all_tensors(self) -> None: _tok_int32_bytes = self.max_tokens * 4 # Request-level fields are all 4 bytes wide (5 int32 + 2 float32 = 7 fields). _req_4byte_bytes = self.max_requests * 4 + # Scalar: real (unpadded) token count for the current step. Refreshed + # in transfer_bookkeeping_to_gpu(); read on GPU via + # `gpu_view.real_token_count` (MoE routing masks padding tokens). + _real_token_count_bytes = 4 # MHA section: 5 fields (int32) shared between GraphedMHAMetadata and # NonGraphedMHAMetadata. max_bs == max_requests. _mha_query_lengths_bytes = self.max_requests * 4 @@ -1006,6 +1217,7 @@ def initialize_all_tensors(self) -> None: 3 * _tok_int64_bytes + 3 * _tok_int32_bytes + 7 * _req_4byte_bytes + + _real_token_count_bytes + _mha_query_lengths_bytes + _mha_cu_query_seq_lengths_bytes + _mha_kv_seq_lengths_bytes @@ -1014,12 +1226,18 @@ def initialize_all_tensors(self) -> None: ) # Mamba section (hybrid models only). Must match the MambaMetadata # shapes (mirrors the layout documented in ContextGPUView). - # batch_indices_decode is int64; all other fields are int32. + # batch_indices_decode is int32 in batch-invariant mode and int64 otherwise; + # all other fields are int32. if self.is_hybrid_model: - # mamba_batch_indices_decode is int64; pad to 8-byte alignment. - _mamba_align_pad = (8 - _pre_mamba_bytes % 8) % 8 + self._mamba_decode_indices_dtype = ( + torch.int32 if self.batch_invariant_mode else torch.int64 + ) + _decode_index_bytes = 4 if self.batch_invariant_mode else 8 + _mamba_align_pad = ( + _decode_index_bytes - _pre_mamba_bytes % _decode_index_bytes + ) % _decode_index_bytes self._max_mamba_chunks = self.max_tokens // self.mamba_chunk_size + self.max_requests - _mamba_batch_indices_decode_bytes = self.max_requests * 8 + _mamba_batch_indices_decode_bytes = self.max_requests * _decode_index_bytes _mamba_batch_indices_prefill_bytes = self.max_requests * 4 _mamba_seq_idx_bytes = self.max_tokens * 4 _mamba_cu_seqlens_bytes = (self.max_requests + 1) * 4 @@ -1053,6 +1271,27 @@ def initialize_all_tensors(self) -> None: + _mamba_conv_seq_idx_bytes + _mamba_conv_seq_start_bytes ) + # GDP chunk descriptor section, only present for GDP models. All int32, + # so no alignment padding is needed after the Mamba section. + # gdp_chunk_indices int32 (max_gdp_chunks, 2) + # gdp_chunk_indices_dp int32 (max_gdp_chunks_dp, 2) + # gdp_chunk_offsets int32 (max_requests + 1,) + if self.gdp_num_householder > 0: + self._max_gdp_chunks, self._max_gdp_chunks_dp = max_gdp_chunk_counts( + self.max_tokens, self.max_requests, self.gdp_num_householder + ) + _gdp_chunk_indices_bytes = self._max_gdp_chunks * 2 * 4 + _gdp_chunk_indices_dp_bytes = self._max_gdp_chunks_dp * 2 * 4 + _gdp_chunk_offsets_bytes = (self.max_requests + 1) * 4 + else: + self._max_gdp_chunks = 0 + self._max_gdp_chunks_dp = 0 + _gdp_chunk_indices_bytes = 0 + _gdp_chunk_indices_dp_bytes = 0 + _gdp_chunk_offsets_bytes = 0 + _total_bytes += ( + _gdp_chunk_indices_bytes + _gdp_chunk_indices_dp_bytes + _gdp_chunk_offsets_bytes + ) self._cpu_bookkeeping_buf = torch.empty( _total_bytes, dtype=torch.uint8, device='cpu', pin_memory=True ) @@ -1132,6 +1371,14 @@ def initialize_all_tensors(self) -> None: ].view(torch.int32) _off += _req_4byte_bytes + # Scalar staging slot for the real (unpadded) token count. Refreshed + # from `self.batch_dimensions.token_count` in transfer_bookkeeping_to_gpu() + # and read on GPU via `gpu_view.real_token_count`. + self._staging_real_token_count = self._cpu_bookkeeping_buf[ + _off : _off + _real_token_count_bytes + ].view(torch.int32) + _off += _real_token_count_bytes + # Static tensor addresses to make `last_token_logits` graphable with speculative decoding. max_logit_idxs = self.max_requests * (self.num_speculative_tokens + 1) self.active_logit_idxs = torch.zeros( @@ -1175,7 +1422,7 @@ def initialize_all_tensors(self) -> None: _off += _mamba_align_pad self._cpu_mamba_batch_indices_decode = self._cpu_bookkeeping_buf[ _off : _off + _mamba_batch_indices_decode_bytes - ].view(torch.int64) + ].view(self._mamba_decode_indices_dtype) _off += _mamba_batch_indices_decode_bytes self._cpu_mamba_batch_indices_prefill = self._cpu_bookkeeping_buf[ _off : _off + _mamba_batch_indices_prefill_bytes @@ -1212,6 +1459,24 @@ def initialize_all_tensors(self) -> None: ].view(torch.int32) _off += _mamba_conv_seq_start_bytes + if self.gdp_num_householder > 0: + self._cpu_gdp_chunk_indices = ( + self._cpu_bookkeeping_buf[_off : _off + _gdp_chunk_indices_bytes] + .view(torch.int32) + .view(self._max_gdp_chunks, 2) + ) + _off += _gdp_chunk_indices_bytes + self._cpu_gdp_chunk_indices_dp = ( + self._cpu_bookkeeping_buf[_off : _off + _gdp_chunk_indices_dp_bytes] + .view(torch.int32) + .view(self._max_gdp_chunks_dp, 2) + ) + _off += _gdp_chunk_indices_dp_bytes + self._cpu_gdp_chunk_offsets = self._cpu_bookkeeping_buf[ + _off : _off + _gdp_chunk_offsets_bytes + ].view(torch.int32) + _off += _gdp_chunk_offsets_bytes + assert _off == _total_bytes, f"layout bug: wrote {_off} of {_total_bytes} bytes" # GPU view: the single interface for GPU code to read context state. @@ -1222,7 +1487,16 @@ def initialize_all_tensors(self) -> None: max_kv_blocks=self.max_kv_block_count, device=torch.cuda.current_device(), max_mamba_chunks=self._max_mamba_chunks, + max_gdp_chunks=self._max_gdp_chunks, + max_gdp_chunks_dp=self._max_gdp_chunks_dp, + # Same presence predicate as the CPU-side GDP section above; the two + # layouts must byte-match, so they gate on the same flag. + has_gdp=self.gdp_num_householder > 0, + mamba_decode_indices_dtype=( + self._mamba_decode_indices_dtype if self.is_hybrid_model else torch.int64 + ), ) + self._bookkeeping_h2d_done_event = torch.cuda.Event() # Cache of (input_ids_view, pos_ids_view) keyed by num_tokens. Instead of slicing and # unsqueezing on every new inference step (constructing new TensorImpls at 30-60 us), @@ -1269,6 +1543,40 @@ def initialize_all_tensors(self) -> None: and self.config.enable_prefix_caching ): self._allocate_mamba_cache(self.config.prefix_caching_mamba_gb) + elif self.is_hybrid_model and self.config.enable_prefix_caching: + # Memory-only mode: prefix caching on a hybrid model without a Mamba + # cache budget deduplicates identical KV prefixes for memory savings, + # but does NOT cache Mamba recurrent state. Prefill skipping is + # therefore disabled (prefix_skip_tokens is forced to 0) and every + # token is recomputed, so results stay correct -- but the main latency + # benefit of prefix caching is forgone. Warn so a user who expected + # full caching knows to set prefix_caching_mamba_gb. + logging.warning( + "enable_prefix_caching is set on a hybrid (Mamba) model but " + "prefix_caching_mamba_gb is not configured (got %r). Running in " + "memory-only mode: identical KV prefixes are deduplicated for " + "memory savings, but Mamba state caching and prefill skipping are " + "disabled (every token is recomputed). Set prefix_caching_mamba_gb " + "> 0 to enable full prefix caching.", + self.config.prefix_caching_mamba_gb, + ) + + # MTP speculative decoding: persistent buffer for decoder hidden states. + # Only needed for block-scope CUDA graphs, where the Python assignment in + # forward() runs only during graph capture. Using copy_() into a fixed + # buffer ensures every batch-size graph replay writes to the same GPU + # address. Sized to max_tokens; only [:actual_tokens] is valid each step. + if ( + self.num_speculative_tokens > 0 + and self.inference_cuda_graph_scope == InferenceCudaGraphScope.block + ): + self.mtp_decoder_hidden_states = torch.empty( + self.max_tokens, + 1, + self.hidden_size, + device=torch.cuda.current_device(), + dtype=self.params_dtype, + ) # Reset tensor-related metadata. self.reset_metadata() @@ -1489,6 +1797,10 @@ def pad_active_slices(self): # Padded gather indices fan in to row 0 harmlessly when used by FlashInfer. self.active_request_last_token_idxs[padding_request_slice].fill_(0) + def _clear_input_id_padding(self) -> None: + """Restore zero-filled padding without touching active VLM ``-1`` placeholders.""" + self.token_to_input_ids[self.padding_slice].zero_() + def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None: """Append to KV cache. @@ -1509,6 +1821,7 @@ def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) padded_active_token_count=self.padded_active_token_count, token_to_block_idx=self.gpu_view.token_to_block_idx, token_to_local_position_within_kv_block=self.gpu_view.token_to_local_position_within_kv_block, + dummy_block_idx=self.kv_block_allocator.dummy_block_idx, ) block_idx = self.gpu_view.token_to_block_idx[: self.padded_active_token_count] @@ -1597,15 +1910,58 @@ def _allocate_mamba_cache(self, mamba_gb: float) -> None: ssm_size = _math.prod(self.mamba_ssm_states_shape) * self.mamba_ssm_states_dtype.itemsize per_slot_bytes = self.num_mamba_layers * (conv_size + ssm_size) total_bytes = int(mamba_gb * 1024**3) - max_slots = total_bytes // per_slot_bytes + + # MambaSlotAllocator allocates two GPU buffer families with the same + # per-slot footprint, both of which must fit in this budget: + # - "durable" cache: self.ssm_states / self.conv_states, sized to + # `max_slots` slots (computed below). + # - "scratch" buffers: self.intermediate_ssm_out / self.intermediate_conv_out, + # fixed CUDA-graph-safe staging for intermediate-state + # extraction, sized to the per-step token-budget cap + # `scratch_slots` = max_mamba_intermediate_states_per_step. + # The scratch is not part of the durable cache but consumes the same + # per-slot bytes, so reserve it from the budget up front before sizing the + # durable cache; otherwise total usage silently exceeds mamba_gb (and can + # OOM) when scratch_slots > max_slots. + scratch_slots = self.max_mamba_intermediate_states_per_step + scratch_bytes = scratch_slots * per_slot_bytes + max_slots = (total_bytes - scratch_bytes) // per_slot_bytes # durable slots + local_max_slots = max_slots + + # Prefix-cache state is replicated across pipeline stages. Use the + # smallest stage-local capacity so identical cache operations produce + # identical block-to-slot mappings on every stage. A recurrent prefix + # is executable only when all stages retain its state, so additional + # slots on an individual stage would not increase usable cache capacity. + if get_pg_size(self.pipeline_parallel_group) > 1: + max_slots_tensor = torch.tensor( + max_slots, dtype=torch.int64, device=torch.cuda.current_device() + ) + torch.distributed.all_reduce( + max_slots_tensor, + op=torch.distributed.ReduceOp.MIN, + group=self.pipeline_parallel_group, + ) + max_slots = int(max_slots_tensor.item()) + if max_slots < 1: - logging.warning( - "Mamba cache budget (%.3f GB) too small for even 1 slot " - "(need %.3f GB per slot). Mamba caching disabled.", - mamba_gb, - per_slot_bytes / 1024**3, + if local_max_slots >= 1: + raise ValueError( + f"Mamba prefix cache budget (prefix_caching_mamba_gb={mamba_gb:.4g} GB) " + f"has room for {local_max_slots} durable slots on this pipeline stage, but " + f"another stage has room for fewer than one. Cache capacity is mirrored to " + f"the minimum across stages; increase prefix_caching_mamba_gb, reduce " + f"max_tokens, or rebalance the Mamba layers." + ) + raise ValueError( + f"Mamba prefix cache budget (prefix_caching_mamba_gb={mamba_gb:.4g} GB) " + f"is too small. The CUDA-graph extraction scratch reserves " + f"{scratch_bytes / 1024**3:.4g} GB ({scratch_slots} slots x " + f"{per_slot_bytes / 1024:.1f} KB/slot), leaving room for " + f"fewer than one durable cache slot. Increase prefix_caching_mamba_gb " + f"to at least {(scratch_bytes + per_slot_bytes) / 1024**3:.4g} GB, or " + f"reduce max_tokens." ) - return self.mamba_slot_allocator = MambaSlotAllocator( context=self, @@ -1621,9 +1977,14 @@ def _allocate_mamba_cache(self, mamba_gb: float) -> None: ) logging.info( - "Mamba prefix cache: %d slots (%.3f GB), per-slot %.1f KB", + "Mamba prefix cache: %d durable slots (%.3f GB) + %d scratch slots " + "(%.3f GB) = %.3f GB total within %.3f GB budget, per-slot %.1f KB", max_slots, max_slots * per_slot_bytes / 1024**3, + scratch_slots, + scratch_bytes / 1024**3, + (max_slots + scratch_slots) * per_slot_bytes / 1024**3, + mamba_gb, per_slot_bytes / 1024, ) @@ -1691,7 +2052,6 @@ def apply_rotary_emb_query( cp_group=cp_group, mscale=mscale, mla_rotary_interleaved=config.multi_latent_attention, - max_seqlen=query_emb.size(0), ) return query @@ -2018,6 +2378,8 @@ def add_dummy_requests_for_expert_parallel_step( self.request_to_kv_block_ids[0:N, 0] = dummy_block_idx # 3. Token-level state consumed by the triton KV append kernel. + # Dummy slots are active, so initialize their embedding IDs explicitly. + self.token_to_input_ids[0:T].zero_() self.token_to_block_idx[0:T] = dummy_block_idx # Compute per-request token positions: e.g. query_lengths [3,2] -> [0,1,2,0,1] query_lengths = self.request_query_lengths[0:N] @@ -2051,7 +2413,9 @@ def initialize_attention_state( *, construct_graph_dimensions: Optional[InferenceBatchDimensions] = None, is_expert_parallel_dummy_cuda_graph_step: bool = False, - ) -> None: + transfer_bookkeeping_to_gpu: bool = True, + record_bookkeeping_done_event: bool = False, + ) -> Optional[torch.cuda.Event]: """Initialize attention state so that every layer can use it. Args: @@ -2059,8 +2423,17 @@ def initialize_attention_state( The graph config to use for constructing the cuda graphs. is_expert_parallel_dummy_cuda_graph_step (bool): Whether this is a dummy expert model parallel step. - Return: - None. + transfer_bookkeeping_to_gpu (bool): Whether to publish the prepared + CPU bookkeeping snapshot to GPU before returning. Legacy + callers publish immediately; async scheduling binds the GPU + views here and publishes their values later. + record_bookkeeping_done_event (bool): Whether to record an event + after the bookkeeping H2D transfer. + + Returns: + Optional[torch.cuda.Event]: Event marking bookkeeping H2D + completion, or `None` when no event was requested or no + transfer was performed. """ # Launch deferred Mamba GPU ops first (state zeroing/restore) so they # overlap with the CPU work below. These are non-blocking GPU kernels. @@ -2138,6 +2511,7 @@ def initialize_attention_state( self.padded_active_token_count = self.padded_batch_dimensions.token_count self.padded_active_request_count = self.padded_batch_dimensions.req_count self.padding_slice = slice(self.active_token_count, self.padded_active_token_count) + self._clear_input_id_padding() self.build_active_slices( min(self.padded_active_request_count, self.max_requests - self.paused_request_count) @@ -2225,10 +2599,19 @@ def initialize_attention_state( self._cpu_mha_cu_kv_seq_lengths[real_bs] ) - # Block table: [0:real_bs] real, [real_bs:padded_bs] = -1 sentinel. + # Block table: [0:real_bs] real, [real_bs:padded_bs] = dummy block. self._cpu_mha_block_table[:real_bs] = request_to_kv_block_ids_view[:real_bs] if real_bs < padded_bs: - self._cpu_mha_block_table[real_bs:padded_bs] = -1 + self._cpu_mha_block_table[real_bs:padded_bs] = self.kv_block_allocator.dummy_block_idx + # Real rows must avoid having a -1 sentinel in their trailing columns, + # because the kernel treats the -1 sentinel as a real value. + # We cannot avoid writing -1 into `request_to_kv_block_ids`; other logic needs it. + # The only option is to overwrite the -1 with a dummy block index via `masked_fill`. + if real_bs > 0: + _real_rows = self._cpu_mha_block_table[:real_bs] + # masked_fill_ over the pinned int32 view: no index_put/nonzero + # temporaries on the per-step CPU path. + _real_rows.masked_fill_(_real_rows < 0, self.kv_block_allocator.dummy_block_idx) # Max sequence lengths (Python scalars; consumed as kernel launch args). if not self.using_cuda_graph_this_step() and real_bs > 0: @@ -2298,12 +2681,17 @@ def initialize_attention_state( # No-op when the queue is already empty (regular non-warmup steps). self._execute_pending_mamba_ops() - # Run the H2D transfer here so callers that bypass the controller - # (e.g. unit tests that call `model.forward()` directly after - # `initialize_attention_state()`) see populated GPU bookkeeping. The - # text-generation controller still calls `transfer_bookkeeping_to_gpu` - # explicitly; that second call is a cheap idempotent re-copy. - self.transfer_bookkeeping_to_gpu() + # Record whether this step produces real output — false on CUDA-graph + # capture (warmup) or dummy EP steps. Used by transfer_bookkeeping_to_gpu + # to publish real_token_count=0 so MoE routing masks all padding tokens. + self._bookkeeping_no_real_work = ( + construct_graph_dimensions is not None or is_expert_parallel_dummy_cuda_graph_step + ) + + # Preserve the existing behavior for callers that do not publish explicitly. + if transfer_bookkeeping_to_gpu: + return self.transfer_bookkeeping_to_gpu(record_done_event=record_bookkeeping_done_event) + return None def _execute_pending_mamba_ops(self) -> None: """Execute Mamba GPU operations deferred from add_request() / update_requests(). @@ -2329,18 +2717,33 @@ def _execute_pending_mamba_ops(self) -> None: self.mamba_ssm_states[:, indices] = 0.0 self._pending_mamba_zeros.clear() - def transfer_bookkeeping_to_gpu(self) -> None: + def transfer_bookkeeping_to_gpu( + self, skip_token_input_ids: bool = False, record_done_event: bool = False + ) -> Optional[torch.cuda.Event]: """Batch transfer CPU bookkeeping state to GPU staging buffers. - Called after initialize_attention_state() and before the forward pass. - All copies use non_blocking=True with pinned CPU memory. CUDA stream - ordering guarantees the forward pass sees completed transfers. + Legacy steps call this from initialize_attention_state(). Async + scheduling instead delays publication until after preparation and the + GPU sample-to-input copy. Legacy transfers block because the pinned CPU + source is re-staged in place. Async scheduling requests an event-tracked + non-blocking copy and synchronizes that event before reusing the source. The bookkeeping fields are backed by one contiguous pinned CPU buffer - and one contiguous GPU buffer; a single cudaMemcpyAsync suffices. - Request-level staging slots are refreshed from the persistent CPU - tensors immediately before the H2D (GPU reads them at `[:n_active]` + and one contiguous GPU buffer; a single memcpy covers the whole + transfer. Request-level staging slots are refreshed from the persistent + CPU tensors immediately before the H2D (GPU reads them at `[:n_active]` while CPU bookkeeping keeps them at `[paused_count:total_count)`). + + Args: + skip_token_input_ids (bool): If true, leave + `gpu_view.token_to_input_ids` unchanged while copying the rest + of the bookkeeping buffer. + record_done_event (bool): Whether to record and return an event after + an asynchronous bookkeeping transfer. + + Returns: + Optional[torch.cuda.Event]: Event marking H2D completion, or `None` + when no event was requested. """ n_active = self.total_request_count - self.paused_request_count active_slice = slice(self.paused_request_count, self.total_request_count) @@ -2374,11 +2777,32 @@ def transfer_bookkeeping_to_gpu(self) -> None: self._staging_request_query_lengths[n_active:padded_active] = 0 self._staging_request_kv_length_offsets[n_active:padded_active] = 0 - # Coalesced H2D: one cudaMemcpyAsync for the entire bookkeeping buffer. + # Real (unpadded) token count for this step. CUDA-graph replay pads + # the token dim to a captured size; MoE routing reads this on GPU and + # rewrites padding rows' routing entries to -1 so they don't go to + # any expert. Set to 0 on CUDA-graph capture / dummy EP steps so + # every row gets masked out. + self._staging_real_token_count[0] = ( + 0 if self._bookkeeping_no_real_work else self.batch_dimensions.token_count + ) + + # Coalesced H2D: one copy for the entire bookkeeping buffer. # Copying the whole (max_tokens + max_requests)-sized buffer including # unused slots is cheap (~71 KB total, ~3-5 us on PCIe Gen4) and saves - # 8 redundant launch overheads vs. the prior per-field copies. - self.gpu_view._buf.copy_(self._cpu_bookkeeping_buf, non_blocking=True) + # redundant launch overheads vs. per-field copies. Async scheduling + # decode steps skip token_to_input_ids here because sampled tokens are + # already GPU-resident and copied directly into the GPU input buffer. + if skip_token_input_ids: + token_to_input_ids_offset = ( + self.token_to_input_ids.numel() * self.token_to_input_ids.element_size() + ) + else: + token_to_input_ids_offset = 0 + # Only event-tracked callers may leave the copy in flight; legacy callers + # block before the pinned CPU source can be re-staged. + self.gpu_view._buf[token_to_input_ids_offset:].copy_( + self._cpu_bookkeeping_buf[token_to_input_ids_offset:], non_blocking=record_done_event + ) # MHA metadata GPU views were already bound to state_data in # initialize_attention_state(); the H2D above populates the underlying @@ -2389,6 +2813,54 @@ def transfer_bookkeeping_to_gpu(self) -> None: self.mamba_metadata.load_from_cpu(self._pending_mamba_transfer) self._pending_mamba_transfer = None + done_event = None + if record_done_event: + done_event = self._bookkeeping_h2d_done_event + done_event.record(torch.cuda.current_stream()) + + return done_event + + def copy_async_sched_sample_to_forward( + self, sampled_tokens_cuda: Tensor, sampled_mtp_tokens_cuda: Optional[Tensor] = None + ) -> None: + """Populate GPU input token IDs from sampled CUDA tokens for async scheduling. + + Async scheduling keeps sampled tokens GPU-resident for the next decode + forward. CPU bookkeeping is prepared independently and published later; + this direct GPU copy populates the live input-ID view without waiting + for the sample's CPU copy. + + Args: + sampled_tokens_cuda (Tensor): 1D CUDA tensor containing one sampled + token per active decode request. + sampled_mtp_tokens_cuda (Optional[Tensor]): MTP draft tokens with shape + ``[num_speculative_tokens, active_request_count]``. + """ + active_request_count = self.total_request_count - self.paused_request_count + + if self.num_speculative_tokens > 0: + expected_shape = (self.num_speculative_tokens, active_request_count) + if sampled_mtp_tokens_cuda is None or tuple(sampled_mtp_tokens_cuda.shape) != ( + expected_shape + ): + actual_shape = ( + None if sampled_mtp_tokens_cuda is None else sampled_mtp_tokens_cuda.shape + ) + raise RuntimeError( + f"Expected MTP draft token shape {expected_shape}, got {actual_shape}." + ) + + tokens_per_request = self.num_speculative_tokens + 1 + token_count = active_request_count * tokens_per_request + grouped_tokens = self.gpu_view.token_to_input_ids[:token_count].view( + active_request_count, tokens_per_request + ) + grouped_tokens[:, 0].copy_(sampled_tokens_cuda, non_blocking=True) + if sampled_mtp_tokens_cuda is not None: + grouped_tokens[:, 1:].copy_(sampled_mtp_tokens_cuda.transpose(0, 1), non_blocking=True) + if token_count < self.padded_active_token_count: + self.gpu_view.token_to_input_ids[token_count : self.padded_active_token_count].zero_() + def reset_tensors(self) -> None: """Fill all bookkeeping tensors with sentinel values.""" @@ -2415,19 +2887,34 @@ def reset_tensors(self) -> None: self.token_to_block_idx.fill_(-1) self.token_to_local_position_within_kv_block.fill_(0) - def reset_metadata(self) -> None: + def reset_metadata( + self, preserve_prefix_cache: bool = False, *, preserve_counters: bool = False + ) -> None: """Reset all bookkeeping state: counters, block allocator, attention/mamba state. This must be called after ``initialize_all_tensors()`` and after any suspend/resume cycle to bring the context back to a clean state. + + Args: + preserve_prefix_cache: When True, keep the KV block allocator's prefix-cache + state (hash index, ref counts, cached blocks) intact. Used by the idle + ``dummy_forward`` path, which only needs to clear the transient one-token + step state -- wiping the allocator there would destroy cross-request prefix + reuse for any subsequent request (the engine idles between requests at low + concurrency, especially with EP > 1). + preserve_counters: When True, keep engine-step, prefix-cache clock, + prefill-token, and async-scheduling counters intact. """ + # There is no prefix-cache state to preserve when caching is disabled. + preserve_prefix_cache = preserve_prefix_cache and self.enable_prefix_caching # Reset request/token counts. self.total_request_count = 0 self.active_token_count = 0 - self.lifetime_prefill_token_count = 0 - self.async_sched_step_count = 0 - self.async_sched_compaction_step_count = 0 + if not preserve_counters: + self.lifetime_prefill_token_count = 0 + self.async_sched_step_count = 0 + self.async_sched_compaction_step_count = 0 self.paused_request_count = 0 self.batch_dimensions = InferenceBatchDimensions( token_count=0, prefill_req_count=0, decode_req_count=0 @@ -2443,7 +2930,8 @@ def reset_metadata(self) -> None: # Reset attention, mamba, and block allocator state. self.reset_attention_state() self.reset_mamba_state() - self.kv_block_allocator.reset() + if not preserve_prefix_cache: + self.kv_block_allocator.reset() self.request_to_kv_block_ids.fill_(-1) # Reset chunked prefill state @@ -2455,7 +2943,15 @@ def reset_metadata(self) -> None: token_count=0, prefill_req_count=0, decode_req_count=0 ) - def reset(self) -> None: + # Reset VLM data. + # Guarded because tests can construct via __new__ and skip __init__. + getattr(self, "_request_to_image_embeddings", {}).clear() + getattr(self, "_request_to_image_token_mask", {}).clear() + getattr(self, "_request_to_image_token_count", {}).clear() + + def reset( + self, preserve_prefix_cache: bool = False, *, preserve_counters: bool = False + ) -> None: """Reset entire context. This method does: @@ -2466,18 +2962,32 @@ def reset(self) -> None: This method is useful after cuda graph warmup iterations, where the context's memory buffer is referenced by the cuda graph system and cannot be deallocated. + + Args: + preserve_prefix_cache: When True, keep the KV and Mamba prefix-cache + state (hash indices and cached blocks/slots) intact. Used by + the idle ``dummy_forward`` path so an idle step between requests does + not destroy cross-request prefix reuse. + preserve_counters: When True, keep engine-step, prefix-cache clock, + prefill-token, and async-scheduling counters intact. """ + # There is no prefix-cache state to preserve when caching is disabled. + preserve_prefix_cache = preserve_prefix_cache and self.enable_prefix_caching + self.dynamo_helper.discard_pending_kv_stored_events() self.reset_tensors() - self.reset_metadata() + self.reset_metadata( + preserve_prefix_cache=preserve_prefix_cache, preserve_counters=preserve_counters + ) - # Reset lifetime counters (not reset in reset_metadata, which is also - # called during suspend/resume where these must persist). - self.step_count = 0 - self.prefix_cache_lru_clock = 0 + if not preserve_counters: + self.step_count = 0 + self.prefix_cache_lru_clock = 0 - # Reset Mamba cache state - if self.mamba_slot_allocator is not None: + # Reset Mamba cache state. + if not preserve_prefix_cache and self.mamba_slot_allocator is not None: self.mamba_slot_allocator.reset() + if not preserve_prefix_cache: + self.dynamo_helper.notify_kv_cache_cleared() def current_input_and_position_ids( self, *, num_warmup_tokens: Optional[int] = None @@ -2566,13 +3076,46 @@ def last_token_logits(self, logits: Tensor) -> Tensor: ) return logits.squeeze(0)[self.active_logit_idxs[: self.num_last_token_logits], :] + def _find_mamba_match_count( + self, req: DynamicInferenceRequest, start_block: int, end_block: int + ) -> int: + """Find the farthest cached Mamba state within a chunk-local block range. + + Mamba state restore is only valid for blocks that the current chunk also + assigns from the KV cache. Chunked prefill can schedule a prompt prefix + that is shorter than the farthest cached full-prompt Mamba boundary, so + this helper intentionally uses the same block domain as KV matching. + """ + if self.mamba_slot_allocator is None or not req.precomputed_block_hashes: + return 0 + + end_block = min(end_block, len(req.precomputed_block_hashes)) + if start_block >= end_block: + return 0 + + mamba_map = self.mamba_slot_allocator.hash_to_block_id + hashes = req.precomputed_block_hashes[start_block:end_block] + for i in range(len(hashes) - 1, -1, -1): + if hashes[i] in mamba_map: + return i + 1 + return 0 + def _compute_prefix_match( - self, req: DynamicInferenceRequest, prefill_chunk_length: int + self, + req: DynamicInferenceRequest, + prefill_chunk_length: int, + record_mamba_match: bool = False, ) -> Tuple[list, int, int, int, int, int]: """Compute prefix match results and skip counts for a request chunk. Shared by check_availability (budget checks) and add_request (execution). + Args: + req: Request being scheduled. + prefill_chunk_length: Number of prompt tokens considered in this chunk. + record_mamba_match: If True, store the chunk-local executable Mamba + match count on the request for diagnostics/tests. + Returns: Tuple of (matched_block_ids, num_blocks_from_pool, already_allocated_blocks, overall_required_blocks, @@ -2611,7 +3154,11 @@ def _compute_prefix_match( # Only applies to the first chunk (finished == 0); continuation chunks # already had Mamba state restored during the first chunk. if self.is_hybrid_model and self.mamba_slot_allocator is not None and finished == 0: - num_mamba_matched = getattr(req, '_mamba_num_matched_blocks', 0) + num_mamba_matched = self._find_mamba_match_count( + req, already_allocated_blocks, already_allocated_blocks + num_matched + ) + if record_mamba_match: + req._mamba_num_matched_blocks = num_mamba_matched assert ( num_mamba_matched <= num_matched ), f"Mamba match ({num_mamba_matched}) > KV match ({num_matched})" @@ -2631,6 +3178,8 @@ def _compute_prefix_match( else: prefix_skip_tokens = 0 elif self.is_hybrid_model and finished == 0: + if record_mamba_match: + req._mamba_num_matched_blocks = 0 prefix_skip_tokens = 0 # Clamp so that effective_prefill_chunk_length >= 2 when possible. @@ -2660,19 +3209,36 @@ def check_availability(self, req: DynamicInferenceRequest) -> Tuple[bool, bool, Check if the request can be added to the context. """ # Note that for hybrid models checking the total request count is sufficient - # because we allocate a single set of Mamba state tensors for each request + # because we allocate a single set of Mamba state tensors for each request. request_can_be_added = ( self.total_request_count < self.max_requests and self.paused_request_count == 0 ) + # A handoff is the exception: it keeps its live recurrent-state slot after + # its request row is removed, so both capacities must be checked independently. + if self.is_hybrid_model and self.kv_block_allocator.enable_handoff_pinning: + request_can_be_added &= self.mamba_metadata.mamba_state_free_slot_count > 0 - _, num_blocks_from_pool, _, _, _, effective_prefill_chunk_length = ( + matched_block_ids, num_blocks_from_pool, _, _, _, effective_prefill_chunk_length = ( self._compute_prefix_match(req, req.remaining_prompt_length) ) request_tokens_can_be_added = ( self.active_token_count + effective_prefill_chunk_length <= self.max_tokens ) - kv_cache_available = self.kv_block_allocator.is_memory_available(num_blocks_from_pool) + # add_request pins the matched blocks before allocating. Only matches that + # are currently evictable (ref_count == 0) count against the evictable + # pool; matches already pinned by another in-flight request are not in + # get_evictable_block_count() and pinning them frees nothing. Reserve only + # the ref_count == 0 matches so availability is not under-reported. + potential_matched_count = 0 + if matched_block_ids: + matched_tensor = torch.tensor(matched_block_ids, dtype=torch.int32, device='cpu') + potential_matched_count = int( + (self.kv_block_allocator.block_ref_counts[matched_tensor] == 0).sum() + ) + kv_cache_available = self.kv_block_allocator.is_memory_available( + num_blocks_from_pool, potential_matched_count=potential_matched_count + ) return request_can_be_added, request_tokens_can_be_added, kv_cache_available def _find_kv_match_count( @@ -2757,31 +3323,47 @@ def add_request( overall_required_blocks, prefix_skip_tokens, effective_prefill_chunk_length, - ) = self._compute_prefix_match(req, prefill_chunk_length) + ) = self._compute_prefix_match(req, prefill_chunk_length, record_mamba_match=True) num_matched_blocks = len(matched_block_ids) effective_kv_offset = req.finished_chunk_token_count + prefix_skip_tokens - # Track prefix cache hits. + # Track prefix cache hits. num_cached_tokens accumulates across prefill + # chunks: each chunk matches a disjoint block range (start advances with + # finished_chunk_token_count), so a long cached prefix is discovered + # incrementally and must be summed, not overwritten. if num_matched_blocks > 0: self.prefix_cache_hits += 1 self.prefix_cache_blocks_matched += num_matched_blocks + req.num_cached_tokens += num_matched_blocks * self.block_size_tokens # Slice tokens to skip matched prefix this_round_tokens = req.remaining_prompt_tokens[prefix_skip_tokens:prefill_chunk_length] - new_block_ids = None - if num_blocks_from_pool > 0: - new_block_ids = self.kv_block_allocator.allocate_memory_blocks(num_blocks_from_pool) - if new_block_ids is None or len(new_block_ids) != num_blocks_from_pool: - raise BlockOverflowError(req.request_id) - - # Increment ref counts and update timestamps for matched (shared) blocks + # Pin matched (shared) blocks BEFORE allocation. allocate_memory_blocks() + # may trigger LRU eviction, and a matched block still at ref_count == 0 + # would be an eviction candidate — descendant-first LRU could evict a + # matched leaf and immediately reuse its ID for a new block, leaving the + # block table with a duplicate ID and a dangling parent. Incrementing ref + # counts first removes matched blocks from the evictable set (see + # evict_lru_blocks / get_evictable_block_count), so eviction falls back to + # genuinely unused cached blocks. + matched_tensor = None if num_matched_blocks > 0: matched_tensor = torch.tensor(matched_block_ids, dtype=torch.int32, device='cpu') self.kv_block_allocator.block_ref_counts[matched_tensor] += 1 if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU: self.kv_block_allocator.update_timestamps(matched_tensor) + new_block_ids = None + if num_blocks_from_pool > 0: + new_block_ids = self.kv_block_allocator.allocate_memory_blocks(num_blocks_from_pool) + if new_block_ids is None or len(new_block_ids) != num_blocks_from_pool: + # Roll back the pin so a failed add does not leak ref counts on + # the matched blocks (which would make them permanently unevictable). + if matched_tensor is not None: + self.kv_block_allocator.block_ref_counts[matched_tensor] -= 1 + raise BlockOverflowError(req.request_id) + # Note that we decremented the total_request_count for the chunked prefill request # in update_requests, so setting current_id to the total_request_count will again # make the last request the continuing chunked prefill request if one exists. @@ -2863,6 +3445,63 @@ def add_request( self.token_to_block_idx[ self.active_token_count : self.active_token_count + effective_prefill_chunk_length ] = self.request_to_kv_block_ids[current_id][token_offset_range // self.block_size_tokens] + if num_matched_blocks > 0: + # Some tokens we are about to compute may land inside a block we matched + # by hash. That block already holds the correct KV for exactly these + # tokens and is shared with whoever cached it, so send those writes to the + # dummy block rather than perturb a concurrent reader's values. + # + # Only the write mapping moves. `request_to_kv_block_ids` still points at + # the real block, so attention reads the cached KV through the block table, + # and `token_to_block_idx` is rebuilt every step so nothing needs restoring. + # + # All positions below are absolute token positions within the request. + # A chunk resuming mid-block, with one fresh block past the match: + # + # token 0 B 2B 3B 4B + # blocks |--0--|--1--|--2--|--3--| + # matched [=================] blocks 0-2 were hash-matched + # chunk [==============] tokens computed this step + # redirect [========] the overlap -> dummy block + # write [=====] the rest -> real block 3 + + # 1. The blocks matched here sit right after the blocks this request + # already owns. The partial block carried over from the previous + # chunk joins them when it was itself matched earlier -- the tokens + # completing it are cached too, and it is still shared with whoever + # cached it. Everything before that is a prefix of matched blocks, so + # starting the span at 0 is safe; the intersection below clips it. + matched_prefix_start_token = ( + 0 + if req.num_matched_prefix_blocks >= already_allocated_blocks + else already_allocated_blocks * self.block_size_tokens + ) + matched_prefix_end_token = ( + already_allocated_blocks + num_matched_blocks + ) * self.block_size_tokens + + # 2. This chunk computes the tokens starting at `effective_kv_offset` + # (the prompt offset left over after skipping the cached prefix). + chunk_start_token = effective_kv_offset + chunk_end_token = effective_kv_offset + effective_prefill_chunk_length + + # 3. The redundant tokens are where those two spans overlap. Non-empty + # whenever we matched more blocks than we skipped tokens for: the + # ">= 2 computed tokens" clamp, the Mamba back-off, or memory-only + # hybrid mode where nothing is skipped but blocks are still shared. + overlap_start_token = max(matched_prefix_start_token, chunk_start_token) + overlap_end_token = min(matched_prefix_end_token, chunk_end_token) + + # 4. Rebase the overlap onto this chunk's slots and redirect it. The + # overlap starts past the chunk start only when this chunk's own + # matched blocks begin later than the tokens it computes, i.e. the + # chunk opens inside a freshly allocated block that it must write. + if overlap_end_token > overlap_start_token: + redirect_start = self.active_token_count + overlap_start_token - chunk_start_token + redirect_end = self.active_token_count + overlap_end_token - chunk_start_token + self.token_to_block_idx[redirect_start:redirect_end] = ( + self.kv_block_allocator.dummy_block_idx + ) self.token_to_local_position_within_kv_block[ self.active_token_count : self.active_token_count + effective_prefill_chunk_length ] = (token_offset_range % self.block_size_tokens) @@ -2883,9 +3522,29 @@ def _register_range(start: int, end: int): return block_ids_to_hash = self.request_to_kv_block_ids[current_id][start:end].tolist() block_hashes_slice = req.precomputed_block_hashes[start:end] + # Parent hash of block k is the hash of block k-1 in the chain; + # block 0 is a root (parent hash 0). Enables LRU eviction to keep + # parents cached until their children are gone. + parent_hashes_slice = [ + req.precomputed_block_hashes[k - 1] if k > 0 else 0 for k in range(start, end) + ] self.kv_block_allocator.register_kv_block_hashes( - block_ids_to_hash, block_hashes_slice + block_ids_to_hash, block_hashes_slice, parent_hashes_slice ) + if self.dynamo_helper.has_kv_event_listeners: + token_start = start * self.block_size_tokens + token_end = end * self.block_size_tokens + token_ids = req.prompt_tokens[token_start:token_end].tolist() + self.dynamo_helper.queue_kv_stored_event( + { + "block_hashes": list(block_hashes_slice), + "token_ids": token_ids, + "num_block_tokens": [self.block_size_tokens] * (end - start), + "parent_hash": ( + int(req.precomputed_block_hashes[start - 1]) if start > 0 else None + ), + } + ) # Range 1: prior-chunk partial block that this chunk just completed _register_range(previously_complete, min(already_allocated_blocks, num_complete_blocks)) @@ -2909,23 +3568,35 @@ def _register_range(start: int, end: int): else: self._pending_mamba_zeros.append(mamba_idx) - # compute_and_store_offsets sets both CPU state (hash_to_block_id, - # _eos_cache_block_id_gpu) and GPU staging buffers. Runs immediately - # because commit_intermediate_states() reads the CPU state after the - # forward pass. - if self.mamba_slot_allocator is not None: - self.mamba_slot_allocator.compute_and_store_offsets( - req, - current_id, - prefix_skip_tokens, - prefill_chunk_length, - num_matched_blocks, - matched_block_ids, - overall_required_blocks, - ) + # compute_and_store_offsets sets CPU state + GPU staging buffers that + # commit_intermediate_states() consumes after the forward pass. Run it for + # EVERY prefill chunk (not just the first): the last complete block of a + # multi-chunk prompt falls in a continuation chunk, and caching its Mamba + # state is precisely what lets a later turn skip prefill on a hybrid model. + # Mamba slot allocation / state restore above stays first-chunk-only. + if self.is_hybrid_model and self.mamba_slot_allocator is not None: + self.mamba_slot_allocator.compute_and_store_offsets( + req, + current_id, + prefix_skip_tokens, + prefill_chunk_length, + num_matched_blocks, + matched_block_ids, + overall_required_blocks, + ) + + # Extend the request's run of matched blocks, so a later chunk knows not to + # rewrite the partial block it inherits from this one. Only extend when this + # chunk's matches continue the run: a gap means the blocks in between were + # computed by this request and the tokens completing them must be written. + if num_matched_blocks > 0 and req.num_matched_prefix_blocks >= already_allocated_blocks: + req.num_matched_prefix_blocks = already_allocated_blocks + num_matched_blocks self.active_token_count += effective_prefill_chunk_length self.lifetime_prefill_token_count += effective_prefill_chunk_length + if self.enable_prefix_caching: + self.prefix_cache_prefill_computed_tokens += effective_prefill_chunk_length + self.prefix_cache_prefill_skipped_tokens += prefix_skip_tokens self.total_request_count += 1 self.num_prefill_requests += 1 @@ -3044,53 +3715,117 @@ def release_memory_blocks_from_request_indexes(self, request_indexes) -> None: sa._intermediate_block_ids_cpu[request_indexes] = -1 sa._eos_cache_block_id_cpu[request_indexes] = -1 + def _get_paused_request_count_within_block_budget(self) -> int: + """Count the left-most paused requests whose blocks fit the paused budget.""" + block_budget = self.kv_block_allocator.paused_limit + if self.paused_request_count == 0 or block_budget == 0: + return 0 + + if not self.enable_prefix_caching: + paused_block_counts = self.request_kv_block_counts[: self.paused_request_count] + cumulative_block_counts = paused_block_counts.cumsum(dim=0) + return int(torch.count_nonzero(cumulative_block_counts <= block_budget).item()) + + seen_block_ids: set[int] = set() + retained_request_count = 0 + for request_idx in range(self.paused_request_count): + block_ids = self.request_to_kv_block_ids[request_idx] + seen_block_ids.update(block_ids[block_ids >= 0].tolist()) + if len(seen_block_ids) > block_budget: + break + retained_request_count += 1 + return retained_request_count + + def _get_releasable_block_counts( + self, request_start_idx: int, request_end_idx: int + ) -> list[int]: + """Count blocks made allocatable by each right-most request suffix. + + The returned list has one entry for every possible suffix length, including + zero. For prefix caching, a shared block is credited only once all selected + request references account for its current allocator reference count. + """ + suffix_request_count = request_end_idx - request_start_idx + releasable_counts = [0] * (suffix_request_count + 1) + if suffix_request_count == 0: + return releasable_counts + + if not self.enable_prefix_caching: + suffix_block_counts = ( + self.request_kv_block_counts[request_start_idx:request_end_idx] + .flip(dims=[0]) + .cumsum(dim=0) + ) + return [0, *suffix_block_counts.tolist()] + + block_rows = self.request_to_kv_block_ids[request_start_idx:request_end_idx] + request_offsets, block_offsets = torch.where(block_rows >= 0) + if request_offsets.numel() == 0: + return releasable_counts + + block_ids = block_rows[request_offsets, block_offsets].long() + unique_block_ids, inverse, selected_reference_counts = torch.unique( + block_ids, return_inverse=True, return_counts=True + ) + current_reference_counts = self.kv_block_allocator.block_ref_counts[unique_block_ids] + assert torch.all( + selected_reference_counts <= current_reference_counts + ), "selected more KV block references than the allocator owns" + + # A fully selected block becomes releasable when the suffix reaches its + # left-most reference. + suffix_depths = suffix_request_count - request_offsets + release_suffix_depths = torch.zeros_like(unique_block_ids) + release_suffix_depths.scatter_reduce_( + 0, inverse, suffix_depths, reduce="amax", include_self=True + ) + fully_selected = selected_reference_counts == current_reference_counts + releases_by_suffix = torch.bincount( + release_suffix_depths[fully_selected], minlength=suffix_request_count + 1 + ) + return releases_by_suffix.cumsum(dim=0).tolist() + def resume_paused_requests( - self, active_request_count: int, newly_paused_request_ids: torch.Tensor - ) -> tuple[int, torch.Tensor]: - """Resume as many paused requests as we have space for in the active buffer. + self, active_request_count: int, newly_paused_request_ids: Optional[torch.Tensor] + ) -> tuple[int, Optional[torch.Tensor]]: + """Resume as many paused requests as compute and KV capacity permit. Args: active_request_count (int): Number of active requests. - newly_paused_request_ids (torch.Tensor): List of newly paused request ids. - next_tokens (torch.Tensor): Sampled tokens. + newly_paused_request_ids (Optional[torch.Tensor]): Newly paused request ids. Returns: - (tuple[int, torch.Tensor]) active_request_count, newly_paused_request_ids. + (tuple[int, Optional[torch.Tensor]]) Updated active count and newly paused ids. """ # Assign released blocks to paused requests. # todo: @shanmugamr, un-pause requests using FIFO, rather than LIFO. resume_request_count = 0 if self.paused_request_count > 0: - active_block_count_avail = self.kv_block_allocator.get_active_avail() - # Clone not needed: flip() makes a copy. - paused_block_counts = self.request_kv_block_counts[: self.paused_request_count] - # Flip counts before cumsum, since paused requests are resumed from - # the right-most index, so we must count resumed blocks starting from - # the right side. - paused_block_counts = paused_block_counts.flip(dims=[0]) - - # Check which paused requests will actually need a new block upon resuming + # Check which paused requests will actually need a new block upon + # resuming. Flip before cumsum because requests resume from the right. offsets = self.request_last_kv_block_offset[: self.paused_request_count] needs_new_block = ( offsets >= self.block_size_tokens - 1 - self.num_speculative_tokens - ).to(paused_block_counts.dtype) + ).to(self.request_kv_block_counts.dtype) needs_new_block = needs_new_block.flip(dims=[0]) - # Add +1 ONLY to the block counts of requests that finished their previous memory block - paused_block_counts += needs_new_block - paused_block_counts_cumsum = paused_block_counts.cumsum(dim=0) - resume_request_count = min( - torch.nonzero(paused_block_counts_cumsum <= active_block_count_avail).numel(), - self.kv_block_allocator.total_avail, - ) - # Constrain resumptions by the maximum allowed active requests and tokens max_allowed_active = min( self.max_requests, self.max_tokens // (self.num_speculative_tokens + 1) ) allowed_to_resume = max(0, max_allowed_active - active_request_count) - resume_request_count = min(resume_request_count, allowed_to_resume) + candidate_count = min(self.paused_request_count, allowed_to_resume) + + if candidate_count > 0: + new_block_counts_cumsum = needs_new_block[:candidate_count].cumsum(dim=0) + max_new_block_count = int(new_block_counts_cumsum[-1].item()) + block_count_avail = self.kv_block_allocator.pool_avail + if max_new_block_count > block_count_avail: + block_count_avail = self.kv_block_allocator.get_allocatable_count() + resume_request_count = int( + torch.count_nonzero(new_block_counts_cumsum <= block_count_avail).item() + ) self.paused_request_count -= resume_request_count active_request_count += resume_request_count @@ -3106,8 +3841,10 @@ def resume_paused_requests( num_new_blocks = needs_new_block.sum().item() if num_new_blocks > 0: - assert num_new_blocks <= self.kv_block_allocator.total_avail block_ids = self.kv_block_allocator.allocate_memory_blocks(num_new_blocks) + assert ( + block_ids is not None and block_ids.numel() == num_new_blocks + ), f"failed to allocate {num_new_blocks} blocks for resumed requests" # Apply updates only to the requests that required a new block relative_row_idx = torch.nonzero(needs_new_block).squeeze(1) @@ -3133,53 +3870,65 @@ def evict_overflow_paused_requests( active_request_count: int, next_tokens: torch.Tensor, new_speculative_tokens: Optional[torch.Tensor] = None, - ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: - """Evict requests that overflow the paused buffer. + ) -> Optional[torch.Tensor]: + """Evict requests that overflow the paused-block retention budget. Args: active_request_count (int): Number of active requests. next_tokens (torch.Tensor): Sampled tokens. + new_speculative_tokens (Optional[torch.Tensor]): Speculative tokens. Returns: - (torch.Tensor) Evicted request ids. + (Optional[torch.Tensor]) Evicted request ids. """ - # Overflow paused block count. - overflow_paused_block_count = ( - self.kv_block_allocator.get_paused_used() - self.kv_block_allocator.paused_count - ) - - # Nothing to evict? - if overflow_paused_block_count <= 0: - return None - - # Overflow paused block count. - paused_block_counts = self.request_kv_block_counts[: self.paused_request_count] - paused_block_counts_cumsum = paused_block_counts.cumsum(dim=0) - valid_paused_request_count = torch.nonzero( - paused_block_counts_cumsum <= self.kv_block_allocator.paused_count - ).numel() - overflow_paused_request_count = self.paused_request_count - valid_paused_request_count - - # Nothing to evict? (Similar to checking overflow_paused_block_count - # above, but here we allow up to one paused request to overflow into the - # active buffer. + # Keep the largest left-most paused prefix that fits the configured + # paused-block retention budget. Requests to its right must either resume + # or be evicted. + retained_paused_request_count = self._get_paused_request_count_within_block_budget() + overflow_paused_request_count = self.paused_request_count - retained_paused_request_count if overflow_paused_request_count == 0: return None - # Evict request count. (Flip paused_block_counts because evictions are - # counted from the right-most paused requests. - paused_block_counts = paused_block_counts[-overflow_paused_request_count:].flip(dims=[0]) - paused_block_counts_cumsum = paused_block_counts.cumsum(dim=0) - remaining_paused_request_counts = torch.arange( - overflow_paused_request_count - 1, - -1, - -1, - dtype=paused_block_counts_cumsum.dtype, - device='cpu', + max_allowed_active = min( + self.max_requests, self.max_tokens // (self.num_speculative_tokens + 1) + ) + allowed_to_resume = max(0, max_allowed_active - active_request_count) + + needs_new_block = ( + self.request_last_kv_block_offset[: self.paused_request_count] + >= self.block_size_tokens - 1 - self.num_speculative_tokens + ) + overflow_needs_new_block = needs_new_block[ + retained_paused_request_count : self.paused_request_count + ] + releasable_block_counts = self._get_releasable_block_counts( + retained_paused_request_count, self.paused_request_count ) - net_block_counts = paused_block_counts_cumsum - remaining_paused_request_counts - evict_request_count = torch.nonzero(net_block_counts >= 0)[0].item() + 1 + new_block_counts_by_survivor_count = overflow_needs_new_block.cumsum(dim=0) + new_block_counts_by_survivor_count = torch.cat( + (new_block_counts_by_survivor_count.new_zeros(1), new_block_counts_by_survivor_count) + ) + block_count_avail = self.kv_block_allocator.get_allocatable_count() + + # Preserve current ordering: evict the smallest right-most suffix whose + # physical releases let every remaining overflow request resume. + evict_request_count = None + for candidate_evict_count in range(overflow_paused_request_count + 1): + survivor_count = overflow_paused_request_count - candidate_evict_count + if survivor_count > allowed_to_resume: + continue + survivor_new_block_count = int( + new_block_counts_by_survivor_count[survivor_count].item() + ) + candidate_avail = block_count_avail + releasable_block_counts[candidate_evict_count] + if survivor_new_block_count <= candidate_avail: + evict_request_count = candidate_evict_count + break + + assert evict_request_count is not None + if evict_request_count == 0: + return None # Eviction index range. evict_start_idx = self.paused_request_count - evict_request_count @@ -3241,97 +3990,184 @@ def evict_overflow_paused_requests( return evict_request_ids - def prepare_requests(self, new_tokens: Tensor) -> None: - """Speculatively prepare active decode requests for the next forward pass. + def _get_async_sched_rows_requiring_new_block(self) -> Tensor: + """Return active request rows that need a block during the next prepare. - Async scheduling only supports decode-only steps with no pause, - evict, or resume lifecycle changes. If preparing the next token would - require one of those lifecycle changes, this method raises and the caller - should treat async scheduling as unsupported for that workload. + Returns: + Tensor: Boolean mask over active request rows. + """ + active_slice = slice(self.paused_request_count, self.total_request_count) + tokens_per_request = self.num_speculative_tokens + 1 + return ( + self.request_last_kv_block_offset[active_slice] + tokens_per_request + >= self.block_size_tokens + ) - Args: - new_tokens (Tensor): Newly sampled token for each active request. + def can_prepare_requests(self) -> bool: + """Return whether requests can be prepared without lifecycle changes. + + Returns: + bool: Whether all requests are active decode requests and the shared + KV-block pool can satisfy the exact next-step allocation demand. """ - if new_tokens.is_cuda: - new_tokens = new_tokens.cpu() + if self.num_prefill_requests != 0 or self.paused_request_count != 0: + return False + + rows_requiring_new_block = self._get_async_sched_rows_requiring_new_block() + num_new_blocks = int(rows_requiring_new_block.sum().item()) + return num_new_blocks <= self.kv_block_allocator.get_allocatable_count() + + def prepare_requests(self) -> None: + """Speculatively prepare active decode requests for the next forward pass. + Async scheduling only supports decode-only steps with no pause, + evict, or resume lifecycle changes. If preparation cannot allocate the + required KV blocks without a lifecycle change, this method raises. The + prepared decode layout establishes the active token count. + """ active_request_count = self.total_request_count - self.paused_request_count - if self.num_speculative_tokens != 0: - raise RuntimeError("Async scheduling does not support speculative tokens.") if self.num_prefill_requests != 0: raise RuntimeError("Async scheduling only supports decode-only steps.") if self.paused_request_count != 0: raise RuntimeError("Async scheduling does not support paused requests.") - if new_tokens.numel() != active_request_count: - raise RuntimeError( - f"Expected {active_request_count} new tokens, got {new_tokens.numel()}." - ) if active_request_count == 0: self.active_token_count = 0 return active_slice = slice(0, active_request_count) - rows_requiring_new_block = ( - self.request_last_kv_block_offset[active_slice] >= self.block_size_tokens - 1 - ) - num_new_blocks = rows_requiring_new_block.sum().item() + tokens_per_request = self.num_speculative_tokens + 1 + last_block_offsets = self.request_last_kv_block_offset[active_slice] + token_offsets = self._async_sched_token_offsets + rows_requiring_new_block = self._get_async_sched_rows_requiring_new_block() + num_new_blocks = int(rows_requiring_new_block.sum().item()) + + block_ids = None if num_new_blocks > 0: - active_block_count_avail = self.kv_block_allocator.get_active_avail() - if num_new_blocks > active_block_count_avail: + if num_new_blocks > self.kv_block_allocator.get_allocatable_count(): raise RuntimeError("Async scheduling cannot pause requests to allocate new blocks.") block_ids = self.kv_block_allocator.allocate_memory_blocks(num_new_blocks) if block_ids is None: raise RuntimeError("Async scheduling cannot evict requests to allocate new blocks.") + self.active_token_count = active_request_count * tokens_per_request + active_token_slice = slice(0, self.active_token_count) + grouped_token_block_ids = self.token_to_block_idx[active_token_slice].view( + active_request_count, tokens_per_request + ) + grouped_token_block_ids.copy_(self.request_last_kv_block_id[active_slice, None]) + + if block_ids is not None: row_idx = torch.nonzero(rows_requiring_new_block, as_tuple=True)[0] col_idx = self.request_kv_block_counts[row_idx] self.request_to_kv_block_ids[row_idx, col_idx] = block_ids self.request_kv_block_counts[row_idx] += 1 self.request_last_kv_block_id[row_idx] = block_ids + grouped_token_block_ids[row_idx] = torch.where( + last_block_offsets[row_idx, None] + 1 + token_offsets[None, :] + >= self.block_size_tokens, + block_ids[:, None], + grouped_token_block_ids[row_idx], + ) self.request_kv_length_offsets[active_slice].add_(self.request_query_lengths[active_slice]) - self.request_query_lengths[active_slice].fill_(1) - + self.request_query_lengths[active_slice].fill_(tokens_per_request) self.request_last_kv_block_offset[active_slice] = ( - self.request_last_kv_block_offset[active_slice] + 1 + last_block_offsets + tokens_per_request ) % self.block_size_tokens - self.active_token_count = active_request_count - self.token_to_input_ids[:active_request_count] = new_tokens - self.token_to_pos_ids[:active_request_count] = self.request_kv_length_offsets[active_slice] - self.token_to_request_idx[:active_request_count] = torch.arange( - active_request_count, device='cpu' + token_positions = ( + self.request_kv_length_offsets[active_slice, None] + token_offsets[None, :] + ) + token_request_idxs = torch.arange(active_request_count, device='cpu').repeat_interleave( + tokens_per_request ) - self.token_to_position_in_request[:active_request_count] = self.token_to_pos_ids[ - :active_request_count + self.token_to_pos_ids[active_token_slice] = token_positions.flatten() + self.token_to_request_idx[active_token_slice] = token_request_idxs + self.token_to_position_in_request[active_token_slice] = self.token_to_pos_ids[ + active_token_slice ] - self.token_to_local_position_within_kv_block[:active_request_count] = ( - self.token_to_pos_ids[:active_request_count] % self.block_size_tokens + self.token_to_local_position_within_kv_block[active_token_slice] = ( + self.token_to_pos_ids[active_token_slice] % self.block_size_tokens + ) + + def commit_sampled_tokens( + self, sampled_tokens_cpu: Tensor, sampled_mtp_tokens_cpu: Optional[Tensor] = None + ) -> None: + """Commit sampled CPU token IDs to the prepared request state. + + This establishes the post-resolution active token count and populates + the CPU input-ID staging rows in survivor order. Overlapped async + scheduling has already copied the same samples into the live GPU input + view for the speculative forward. + + Args: + sampled_tokens_cpu (Tensor): Sampled CPU token for each active request. + sampled_mtp_tokens_cpu (Optional[Tensor]): MTP draft tokens with shape + ``[num_speculative_tokens, active_request_count]``. + """ + assert sampled_tokens_cpu.device == torch.device( + 'cpu' + ), "Sampled tokens must be on the CPU before they are committed." + + if sampled_mtp_tokens_cpu is not None: + assert sampled_mtp_tokens_cpu.device == torch.device( + 'cpu' + ), "MTP draft tokens must be on the CPU before they are committed." + + active_request_count = self.total_request_count - self.paused_request_count + if sampled_tokens_cpu.numel() != active_request_count: + raise RuntimeError( + f"Expected {active_request_count} new tokens, got {sampled_tokens_cpu.numel()}." + ) + + expected_mtp_shape = (self.num_speculative_tokens, active_request_count) + if self.num_speculative_tokens == 0: + if sampled_mtp_tokens_cpu is not None and sampled_mtp_tokens_cpu.numel() != 0: + raise RuntimeError( + "Received MTP draft tokens when speculative decoding is disabled." + ) + else: + if sampled_mtp_tokens_cpu is None or tuple(sampled_mtp_tokens_cpu.shape) != ( + expected_mtp_shape + ): + actual_shape = ( + None if sampled_mtp_tokens_cpu is None else sampled_mtp_tokens_cpu.shape + ) + raise RuntimeError( + f"Expected MTP draft token shape {expected_mtp_shape}, got {actual_shape}." + ) + + tokens_per_request = self.num_speculative_tokens + 1 + active_token_count = active_request_count * tokens_per_request + self.active_token_count = active_token_count + grouped_tokens = self.token_to_input_ids[:active_token_count].view( + active_request_count, tokens_per_request ) - self.token_to_block_idx[:active_request_count] = self.request_last_kv_block_id[active_slice] + grouped_tokens[:, 0] = sampled_tokens_cpu + if sampled_mtp_tokens_cpu is not None: + grouped_tokens[:, 1:] = sampled_mtp_tokens_cpu.transpose(0, 1) - def resolve_requests(self, active_requests_mask: Tensor) -> Tensor: + def resolve_requests(self, active_requests_mask: Tensor) -> Tuple[Tensor, Tensor]: """Resolve finished requests after an async scheduling forward pass. - Async scheduling supports only request completion. The active request rows - and current decode-token rows are compacted in survivor order so any - following legacy or async scheduling step sees a consistent context. + Prefill requests transition to decode during resolution. Request rows use + the same hole-filling order as ``update_requests`` so seeded sampling stays + consistent with legacy scheduling. Token tensors and the active token + count are left untouched; prepare rebuilds derived token metadata and the + controller commits sampled input IDs after resolution. Args: active_requests_mask (Tensor): 1D mask marking requests that remain active. Returns: - Tensor: Request IDs for requests that finished during resolution. + Tuple[Tensor, Tensor]: Request IDs that finished and source row indices + for surviving requests in their resolved destination order. """ if active_requests_mask.is_cuda: active_requests_mask = active_requests_mask.cpu() - if self.num_speculative_tokens != 0: - raise RuntimeError("Async scheduling does not support speculative tokens.") - if self.num_prefill_requests != 0: - raise RuntimeError("Async scheduling only supports decode-only steps.") if self.paused_request_count != 0: raise RuntimeError("Async scheduling does not support paused requests.") @@ -3342,29 +4178,38 @@ def resolve_requests(self, active_requests_mask: Tensor) -> Tensor: f"got {active_requests_mask.numel()}." ) - survivor_idxs = torch.nonzero(active_requests_mask == 1, as_tuple=True)[0] + self.num_prefill_requests = 0 + self.request_in_prefill_status_tensor[self.request_in_prefill_status_tensor == 1] = 0 + finished_idxs = torch.nonzero(active_requests_mask == 0, as_tuple=True)[0] finished_request_ids = self.request_ids[finished_idxs].clone() + active_request_count = int(active_requests_mask.sum().item()) + survivor_idxs = torch.arange(active_request_count, device='cpu') + finished_idxs_on_left = torch.nonzero( + active_requests_mask[:active_request_count] == 0, as_tuple=True + )[0] + active_idxs_on_right = ( + torch.nonzero(active_requests_mask[active_request_count:] == 1, as_tuple=True)[0] + + active_request_count + ) + assert finished_idxs_on_left.numel() == active_idxs_on_right.numel() + survivor_idxs[finished_idxs_on_left] = active_idxs_on_right + self.reset_attention_state() if finished_idxs.numel() > 0: self.release_memory_blocks_from_request_indexes(finished_idxs) - active_request_count = survivor_idxs.numel() if active_request_count == 0: self.request_to_kv_block_ids.fill_(-1) self.total_request_count = 0 - self.active_token_count = 0 self.reset_mamba_state() - return finished_request_ids + return finished_request_ids, survivor_idxs dst_idxs = torch.arange(active_request_count, device='cpu') if not torch.equal(survivor_idxs, dst_idxs): self.request_kv_length_offsets[dst_idxs] = self.request_kv_length_offsets[survivor_idxs] - self.request_in_prefill_status_tensor[dst_idxs] = self.request_in_prefill_status_tensor[ - survivor_idxs - ] self.request_query_lengths[dst_idxs] = self.request_query_lengths[survivor_idxs] self.request_output_lengths[dst_idxs] = self.request_output_lengths[survivor_idxs] self.request_ids[dst_idxs] = self.request_ids[survivor_idxs] @@ -3374,25 +4219,18 @@ def resolve_requests(self, active_requests_mask: Tensor) -> Tensor: self.request_last_kv_block_offset[dst_idxs] = self.request_last_kv_block_offset[ survivor_idxs ] + if self.is_hybrid_model: + self.mamba_metadata.request_to_mamba_state_idx[dst_idxs] = ( + self.mamba_metadata.request_to_mamba_state_idx[survivor_idxs] + ) for metadata_tensor in self.request_metadata.values(): metadata_tensor[dst_idxs] = metadata_tensor[survivor_idxs] - - self.token_to_input_ids[dst_idxs] = self.token_to_input_ids[survivor_idxs] - self.token_to_pos_ids[dst_idxs] = self.token_to_pos_ids[survivor_idxs] - self.token_to_block_idx[dst_idxs] = self.token_to_block_idx[survivor_idxs] - self.token_to_local_position_within_kv_block[dst_idxs] = ( - self.token_to_local_position_within_kv_block[survivor_idxs] - ) - self.token_to_position_in_request[dst_idxs] = self.token_to_position_in_request[ - survivor_idxs - ] - - self.token_to_request_idx[:active_request_count] = dst_idxs stale_slice = slice(active_request_count, old_active_request_count) self.request_to_kv_block_ids[stale_slice] = -1 + if self.is_hybrid_model: + self.mamba_metadata.request_to_mamba_state_idx[stale_slice] = -1 self.total_request_count = active_request_count - self.active_token_count = active_request_count - return finished_request_ids + return finished_request_ids, survivor_idxs def update_requests( self, @@ -3642,7 +4480,7 @@ def update_requests( active_request_count, newly_paused_request_ids ) - # 6.b. Evict requests that overflow the paused buffer. + # 6.b. Evict requests that overflow the paused-block retention budget. evict_request_ids = self.evict_overflow_paused_requests( active_request_count, next_tokens, new_speculative_tokens ) @@ -3652,8 +4490,13 @@ def update_requests( active_request_count, newly_paused_request_ids ) - assert active_request_count > 0 or self.chunked_prefill_request_id != -1, ( - "active_request_count == %d with no hidden chunked prefill." % active_request_count + assert ( + active_request_count > 0 + or self.total_request_count == 0 + or self.chunked_prefill_request_id != -1 + ), "active_request_count == %d with %d paused requests and no hidden chunked prefill." % ( + active_request_count, + self.paused_request_count, ) # 6.d. Swap the chunked prefill request to the end of the active requests @@ -3857,25 +4700,83 @@ def _processed_log_probs( n_active: int, active_query_lengths: Optional[Tensor], sampling: Optional[Sampling], + row_to_request: Optional[Tensor] = None, ) -> Tensor: - """Sample the logprobs if desired.""" + """Calculate raw or sampling-processed per-row log probabilities. + + Args: + logits (Tensor): Raw logits with shape `[num_rows, vocab_size]`. + n_active (int): Number of active requests represented by the rows. + active_query_lengths (Optional[Tensor]): CPU token counts used to + map prefill rows to active requests, or `None` for decode. + sampling (Optional[Sampling]): Backend providing processed logprobs. + row_to_request (Optional[Tensor]): Explicit CPU mapping from each + logit row to an active request. + + Returns: + Tensor: Per-row log probabilities over the vocabulary. + """ if self.config.logprobs_mode == "raw_logprobs": return F.log_softmax(logits, dim=-1) assert sampling is not None, "processed_logprobs requires a sampling backend" # Map each logits row to its active request. - request_idx = torch.arange(n_active, device=logits.device) - row_to_request = ( - request_idx - if active_query_lengths is None - else request_idx.repeat_interleave(active_query_lengths) - ) - md = self.active_request_metadata - temperature = md["temperature"][:n_active].to(logits.device, torch.float32)[row_to_request] - top_k = md["top_k"][:n_active].to(logits.device, torch.long)[row_to_request] - top_p = md["top_p"][:n_active].to(logits.device, torch.float32)[row_to_request] - return sampling.log_probs_kernel(logits, temperature, top_k, top_p) + if row_to_request is None and active_query_lengths is not None: + row_to_request = torch.arange(n_active).repeat_interleave(active_query_lengths) + return sampling.log_probs_kernel(logits, self, token_to_request_index=row_to_request) + + def calculate_log_probs_tensors( + self, + logits: Tensor, + new_tokens: Tensor, + only_last_token_logits: Optional[bool] = False, + sampling: Optional[Sampling] = None, + row_to_request: Optional[Tensor] = None, + ) -> Tuple[Tensor, Tensor]: + """Calculate selected-token and full-distribution log probabilities. + + Args: + logits (Tensor): Raw model output logits with shape + `[1, sequence_length, vocab_size]`. + new_tokens (Tensor): Newly sampled tokens for active requests. + only_last_token_logits (Optional[bool]): Whether logits contain only + each request's final token row. + sampling (Optional[Sampling]): Sampling backend used for processed + log probabilities. + row_to_request (Optional[Tensor]): Explicit CPU mapping from each + logit row to an active request. + + Returns: + Tuple[Tensor, Tensor]: Selected-token log probabilities flattened in + active-token order and the full per-row log-probability tensor. + """ + logits_squeezed = logits.squeeze(0) + n_active = self.total_request_count - self.paused_request_count + + if only_last_token_logits or self.is_decode_only(): + seq_idx = torch.arange(len(new_tokens), dtype=torch.int32, device=logits.device) + active_logits = logits_squeezed[: len(new_tokens)].float() + log_probs = self._processed_log_probs( + active_logits, n_active, None, sampling, row_to_request + ) + return log_probs[seq_idx, new_tokens], log_probs + + logits_squeezed = logits_squeezed.float() + active_slice = slice(self.paused_request_count, self.total_request_count) + active_query_lengths_cpu = self.request_query_lengths[active_slice] + active_query_lengths_gpu = self.gpu_view.request_query_lengths[:n_active] + + # Shift away each request's first prompt token, then insert its sampled token. + active_token_ids = self.gpu_view.token_to_input_ids[: self.active_token_count].roll(-1, 0) + new_token_idx = active_query_lengths_gpu.cumsum(0) - 1 + active_token_ids[new_token_idx] = new_tokens + + log_probs = self._processed_log_probs( + logits_squeezed, n_active, active_query_lengths_cpu, sampling + ) + seq_idx = torch.arange(self.active_token_count, device=log_probs.device) + return log_probs[seq_idx, active_token_ids], log_probs def calculate_log_probs( self, @@ -3900,58 +4801,15 @@ def calculate_log_probs( log_probs (Tensor): Used to compute top n logprobs later if required. """ - # Calculate log_probs (sequence_length x vocab_size) - logits_squeezed = logits.squeeze(0).float() - n_active = self.total_request_count - self.paused_request_count + selected_log_probs, log_probs = self.calculate_log_probs_tensors( + logits, new_tokens, only_last_token_logits=only_last_token_logits, sampling=sampling + ) if only_last_token_logits or self.is_decode_only(): - seq_idx = torch.arange(len(new_tokens), dtype=torch.int32, device=logits.device) - log_probs = self._processed_log_probs( - logits_squeezed[seq_idx], n_active, None, sampling - ) - selected_log_probs = log_probs[seq_idx, new_tokens] return [[lp] for lp in selected_log_probs.tolist()], log_probs - # Get the selected token ids for all tokens. - # We shift the active token window left by one to remove the first prompt token for - # prefill requests and then set the token ids explicitly for the newly generated tokens. - # This is necessary because we calculate the log probs *before* updating the request metadata. - # - # Example (decode & prefill mix): - # - # active_query_lengths: [ 1 | 1 | 2 | 5 ] - # - # new_tokens : [ 52 | 12 | 3 | 86 ] - # - # seq_idx : [ 0 | 1 | 2 3 | 4 5 6 7 8 ] - # - # new_token_idx : [ 0 | 1 | 3 | 8 ] - # - # active_token_ids before left shift: - # : [ 31 | 75 | 45 16 | 90 12 72 24 88 ] - # - # active_token_ids after shift: - # : [ XX | XX | 16 XX | 12 72 24 88 XX ] (XX = undefined) - # - # active_token_ids[new_token_idx] = new_tokens - # : [ 52 | 12 | 16 3 | 12 72 24 88 86 ] - active_token_ids = self.gpu_view.token_to_input_ids[: self.active_token_count].roll(-1, 0) - active_query_lengths = self.gpu_view.request_query_lengths[:n_active] - - new_token_idx = active_query_lengths.cumsum(0) - 1 - active_token_ids[new_token_idx] = new_tokens - - # Compute (possibly processed) log-probs over all active-token rows. - log_probs = self._processed_log_probs( - logits_squeezed, n_active, active_query_lengths, sampling - ) - - # Extract the log probs for only the selected tokens. - # (sequence_length x vocab_size) -> (sequence_length) - seq_idx = torch.arange(self.active_token_count, device=log_probs.device) - selected_log_probs = log_probs[seq_idx, active_token_ids] - - # Split the log probs across request boundaries + active_slice = slice(self.paused_request_count, self.total_request_count) + active_query_lengths = self.request_query_lengths[active_slice] selected_log_probs_list = selected_log_probs.cpu().split( active_query_lengths.tolist(), dim=0 ) @@ -3979,11 +4837,11 @@ def get_kvcache_utilization_stats(self) -> dict: } """ # Total usable blocks exclude the reserved dummy block. - total_blocks = max(self.kv_block_allocator.total_count - 1, 1) - block_count_avail = int(self.kv_block_allocator.total_avail) + total_blocks = max(self.kv_block_allocator.pool_size - 1, 1) + block_count_avail = int(self.kv_block_allocator.pool_avail) # Overall allocated blocks in the buffer right now. - allocated_blocks = (self.kv_block_allocator.total_count - 1) - block_count_avail + allocated_blocks = (self.kv_block_allocator.pool_size - 1) - block_count_avail allocated_blocks = int(max(0, allocated_blocks)) # Active unique blocks referenced by current active requests only. @@ -4019,3 +4877,185 @@ def get_kvcache_utilization_stats(self) -> dict: 'total_request_count': int(total_request_count), 'max_requests': int(self.max_requests), } + + # ----- VLM per-request data management ----- + + def add_vlm_request_data( + self, + request_id: int, + image_embeddings: Optional[Tensor] = None, + image_token_mask: Optional[Tensor] = None, + ) -> None: + """Attach per-request image data to the context. + + Called at add_request time when a multimodal request has images. + No-op overhead for text-only requests that never call this. + + Args: + request_id: The request identifier. + image_embeddings: Tensor of shape [seq_img, 1, hidden] or None. + image_token_mask: 1-D tensor with -1 for text positions and + non-negative indices for image embedding positions; or None. + """ + self._request_to_image_embeddings[request_id] = image_embeddings + self._request_to_image_token_mask[request_id] = image_token_mask + if image_embeddings is not None: + self._request_to_image_token_count[request_id] = int( + image_embeddings.shape[0] * image_embeddings.shape[1] + ) + else: + self._request_to_image_token_count[request_id] = 0 + + def remove_vlm_request_data(self, request_id: int) -> None: + """Remove image data for a finished request.""" + self._request_to_image_embeddings.pop(request_id, None) + self._request_to_image_token_mask.pop(request_id, None) + self._request_to_image_token_count.pop(request_id, None) + + @property + def has_vlm_data(self) -> bool: + """True iff any active request has attached image data. + + Used by callers on the hot path (``TextGenerationController._ + dynamic_step_forward_logits``) to skip ``current_image_token_mask`` + / ``current_image_embeddings`` entirely on text-only workloads. + """ + return bool(self._request_to_image_token_mask) + + def current_image_token_mask(self) -> Optional[Tensor]: + """Flattened image-token mask aligned with current_input_ids. + + Returns a [1, padded_active_token_count] tensor with -1 for non-image + positions and non-negative indices into the concatenated image + embeddings, or None when there are no active tokens. During + decode-only steps (no image tokens consumed) returns all -1. + """ + # Text-only workloads never populate the per-request VLM dicts; short- + # circuit so the decode hot path doesn't pay for .tolist() / torch.cat. + if not self._request_to_image_token_mask: + return None + + if self.padded_active_token_count is None or self.padded_active_token_count == 0: + return None + + if self.is_decode_only(): + # Return the same [1, padded_active_token_count] shape the prefill + # branch does. Callers advanced-index a batch-first [b, seq, h] + # embedding tensor with this mask; the 2D form is the one that + # broadcasts correctly across the batch axis, a 1D mask would + # attempt to select along dim 0 and either error or silently + # index wrong. + return torch.full( + (1, self.padded_active_token_count), + -1, + dtype=torch.long, + device=torch.cuda.current_device(), + ) + + # Batch the three device→host copies into a single sync to keep the + # prefill hot path from stalling three separate times per step. A full + # vectorization (build the flat mask on-device via cumsum + gather) + # would remove even this remaining sync; TODO before this path takes + # heavy multimodal traffic. + active_slice = slice(self.paused_request_count, self.total_request_count) + _sync_batch = torch.stack( + [ + self.request_ids[active_slice], + self.request_query_lengths[active_slice], + # KV offset — how far into the stored prompt this step begins. + # Non-zero under chunked prefill and after pause/resume, so + # slicing per_request_mask must start at kv_offset, not 0. + self.request_kv_length_offsets[active_slice], + ], + dim=0, + ).tolist() + active_request_ids, active_query_lengths, active_kv_offsets = ( + _sync_batch[0], + _sync_batch[1], + _sync_batch[2], + ) + + segments: List[Tensor] = [] + cumulative_offset = 0 + for request_id, query_len, kv_offset in zip( + active_request_ids, active_query_lengths, active_kv_offsets + ): + per_request_mask = self._request_to_image_token_mask.get(request_id, None) + if per_request_mask is None or kv_offset >= per_request_mask.numel(): + # Past the end of the stored prompt mask — decode step, or a + # text-only request. Return an all-text (-1) segment; the + # stored mask is prompt-length and doesn't cover generated + # positions. + seg = torch.full( + (query_len,), -1, dtype=torch.long, device=torch.cuda.current_device() + ) + else: + # kv_offset is normally 0 for a fresh prefill and grows with + # chunked prefill / pause-and-resume — slice into the stored + # mask from where this step actually resumes. + seg = per_request_mask[kv_offset : kv_offset + query_len].clone() + # Pad if the request is being prefilled beyond the mask (mixed + # prefill+decode chunk). + if seg.numel() < query_len: + pad = torch.full( + (query_len - seg.numel(),), -1, dtype=seg.dtype, device=seg.device + ) + seg = torch.cat([seg, pad]) + positive = seg >= 0 + if positive.any(): + seg[positive] += cumulative_offset + + segments.append(seg) + cumulative_offset += int(self._request_to_image_token_count.get(request_id, 0)) + + if len(segments) == 0: + return None + + mask = torch.cat(segments, dim=0) + if mask.numel() < int(self.padded_active_token_count): + pad = torch.full( + (int(self.padded_active_token_count) - mask.numel(),), + -1, + dtype=torch.long, + device=mask.device, + ) + mask = torch.cat([mask, pad], dim=0) + return mask.unsqueeze(0) + + def current_image_embeddings(self) -> Optional[Tensor]: + """Concatenate image embeddings for active requests. + + Each per-request embedding has shape [seq_img_i, 1, hidden] where + seq_img_i may vary across requests under dynamic resolution. Each is + flattened to [seq_img_i, hidden], concatenated along dim 0, and + unsqueezed to [sum(seq_img_i), 1, hidden]. The downstream consumer + does ``permute(1, 0, 2).reshape(-1, hidden)`` to yield the flat + [total_image_tokens, hidden] array indexed by + :meth:`current_image_token_mask`. + + Returns tensor of shape [total_image_tokens, 1, hidden] or None. + """ + # Text-only workloads never populate the per-request VLM dicts; short- + # circuit so the decode hot path doesn't pay for .tolist() / torch.cat. + if not self._request_to_image_embeddings: + return None + + # Decode steps consume no image tokens, so the concatenated embedding + # is unused. Short-circuit before the .tolist() sync + torch.cat to + # keep the decode critical path free of the D2H stall that would + # otherwise fire on every step, mirroring the mask helper above. + if self.is_decode_only(): + return None + + active_request_ids = self.request_ids[ + self.paused_request_count : self.total_request_count + ].tolist() + + parts: List[Tensor] = [] + for request_id in active_request_ids: + emb = self._request_to_image_embeddings.get(request_id, None) + if emb is not None: + parts.append(emb.reshape(-1, emb.shape[-1])) + if not parts: + return None + return torch.cat(parts, dim=0).unsqueeze(1) From 8b36b00dd573b4a9a1f64ce6e7c808b50313ebb9 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Fri, 28 Aug 2026 03:15:12 +0000 Subject: [PATCH 31/34] fix(ci): secure mamba recipe credentials Signed-off-by: svcnemo-autobot --- tests/test_utils/recipes/h100/mamba.yaml | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_utils/recipes/h100/mamba.yaml b/tests/test_utils/recipes/h100/mamba.yaml index fcfe1c937c8..8ca2fed4168 100644 --- a/tests/test_utils/recipes/h100/mamba.yaml +++ b/tests/test_utils/recipes/h100/mamba.yaml @@ -13,7 +13,9 @@ spec: script_setup: | set -euo pipefail unset https_proxy - echo "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" | tee -a /root/.netrc + umask 077 + printf '%s\n' "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" >> "$HOME/.netrc" + trap 'rm -f "$HOME/.netrc"; unset RO_API_TOKEN' EXIT # Checkout latest cd /opt @@ -34,7 +36,11 @@ spec: git checkout $MCORE_BACKWARDS_COMMIT git rev-parse HEAD rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ + rm -f "$HOME/.netrc" + unset RO_API_TOKEN script: |- + rm -f "${{HOME}}/.netrc" + unset RO_API_TOKEN set -euo pipefail ls cd /opt/megatron-lm From 88a24d50e702d76783694fdd80952a8eb4e699c4 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Fri, 28 Aug 2026 04:53:01 +0000 Subject: [PATCH 32/34] fix: adapt Ling-V3 imports to main Signed-off-by: svcnemo-autobot --- ...rimental_attention_variant_module_specs.py | 71 ------------------- megatron/core/models/hybrid/hybrid_block.py | 18 ++--- megatron/core/packed_seq_params.py | 10 +++ megatron/core/ssm/gated_delta_net/gdn.py | 40 ++++++----- megatron/core/ssm/gated_delta_net/kda.py | 41 ++++++----- .../core/transformer/transformer_config.py | 6 +- pretrain_hybrid.py | 5 ++ 7 files changed, 75 insertions(+), 116 deletions(-) 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 289c1128627..96592b42028 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -16,18 +16,6 @@ AbsorbedMLASelfAttention, AbsorbedMLASelfAttentionSubmodules, ) -from megatron.core.transformer.experimental_attention_variant.csa import ( - CompressedSparseAttention, - CompressedSparseAttentionSubmodules, - Compressor, - CompressorSubmodules, - CSAIndexer, - CSAIndexerSubmodules, -) -from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( - DSv4HybridSelfAttention, - DSv4HybridSelfAttentionSubmodules, -) from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexer, DSAIndexerSubmodules, @@ -173,63 +161,6 @@ def get_dsa_module_spec_for_backend( return attention -def get_dsv4_hybrid_module_spec_for_backend( - config: TransformerConfig, backend: BackendSpecProvider = None -) -> ModuleSpec: - """Helper function to get module spec for DSv4 Hybrid Sparse Attention.""" - assert config.multi_latent_attention, "Currently only MLA supports sparse attention." - assert config.qk_l2_norm is False, "qk_l2_norm is not supported with MLA." - - # Adjust for RMS norm. - rms_norm = config.normalization == "RMSNorm" - # DSA indexer requires normalized q as input, so here we cannot fuse qk layernorm - # with linear projection and have to use unfused qk layernorm. - qk_norm = ( - backend.layer_norm(rms_norm=rms_norm, for_qk=True) if config.qk_layernorm else IdentityOp - ) - - compressor_spec = ModuleSpec( - module=Compressor, - submodules=CompressorSubmodules( - linear_wkv=backend.linear(), - linear_wgate=backend.linear(), - norm=backend.layer_norm(rms_norm=True, for_qk=False), - ), - ) - - indexer_spec = ModuleSpec( - module=CSAIndexer, - submodules=CSAIndexerSubmodules( - linear_wq_b=backend.linear(), - linear_weights_proj=backend.linear(), - compressor=compressor_spec, - ), - ) - - core_attention = ModuleSpec( - module=CompressedSparseAttention, - submodules=CompressedSparseAttentionSubmodules( - compressor=compressor_spec, indexer=indexer_spec - ), - ) - - attention = ModuleSpec( - module=DSv4HybridSelfAttention, - params={"attn_mask_type": AttnMaskType.causal}, - submodules=DSv4HybridSelfAttentionSubmodules( - linear_q_down_proj=backend.linear(), - linear_q_up_proj=backend.column_parallel_linear(), - linear_kv_proj=backend.column_parallel_linear(), - core_attention=core_attention, - linear_proj=backend.row_parallel_linear(), - q_layernorm=qk_norm, - kv_layernorm=qk_norm, - ), - metainfo={"fuse_input_layernorm": False}, - ) - return attention - - def get_experimental_attention_variant_module_spec( config: TransformerConfig, backend: BackendSpecProvider = None ) -> ModuleSpec: @@ -242,8 +173,6 @@ def get_experimental_attention_variant_module_spec( return get_gated_delta_net_module_spec(config=config, backend=backend) elif config.experimental_attention_variant == "dsa": return get_dsa_module_spec_for_backend(config=config, backend=backend) - elif config.experimental_attention_variant == "dsv4_hybrid": - return get_dsv4_hybrid_module_spec_for_backend(config=config, backend=backend) else: raise ValueError( f"Invalid experimental attention variant: {config.experimental_attention_variant}" diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 087ca8da675..84637304d6f 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -24,7 +24,7 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward -from megatron.core.tensor_parallel.random import MHCCheckpointManager +from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager from megatron.core.transformer import TransformerConfig from megatron.core.transformer.cuda_graphs import annotate_first_last_layer from megatron.core.transformer.enums import CudaGraphModule @@ -826,13 +826,13 @@ def _compute_mhc_block_end_plan(self) -> List[bool]: def _build_mhc_recompute_layer_plan( self, use_mhc_recompute: bool - ) -> Tuple[List[Optional[MHCCheckpointManager]], List[bool]]: + ) -> Tuple[List[Optional[CheckpointWithoutOutputManager]], List[bool]]: """Pre-build per-layer MHC recompute managers and block-end markers. The block-end plan is deterministic from config and cached on the - instance; only the per-block ``MHCCheckpointManager`` instances are - allocated fresh per forward pass (managers are single-use). Mirrors - the caching scheme used by ``TransformerBlock``. + instance; only the per-block ``CheckpointWithoutOutputManager`` + instances are allocated fresh per forward pass (managers are single-use). + This mirrors the caching scheme used by ``TransformerBlock``. """ num_layers = len(self.layers) if not use_mhc_recompute or num_layers == 0: @@ -842,17 +842,17 @@ def _build_mhc_recompute_layer_plan( self._mhc_block_end_plan = self._compute_mhc_block_end_plan() is_recompute_block_end = self._mhc_block_end_plan - layer_managers: List[Optional[MHCCheckpointManager]] = [None] * num_layers - mhc_manager = MHCCheckpointManager() + layer_managers: List[Optional[CheckpointWithoutOutputManager]] = [None] * num_layers + mhc_manager = CheckpointWithoutOutputManager() for l_no in range(num_layers): layer_managers[l_no] = mhc_manager if is_recompute_block_end[l_no] and l_no != num_layers - 1: - mhc_manager = MHCCheckpointManager() + mhc_manager = CheckpointWithoutOutputManager() return layer_managers, is_recompute_block_end @staticmethod def _finalize_mhc_recompute_layer( - mhc_manager: Optional[MHCCheckpointManager], + mhc_manager: Optional[CheckpointWithoutOutputManager], hidden_states: Tensor, is_last_in_recompute_block: bool, ) -> None: diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 52e7dbaa516..63c68d9881f 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -26,6 +26,7 @@ class PackedSeqParams: seq_idx: Tensor = None tokens_per_sample: int = None pad_between_seqs: bool = None + cp_partition_mode: str = "zigzag" cp_scatter_cache: object = None def __post_init__(self): @@ -67,3 +68,12 @@ def __post_init__(self): .to(torch.int32) .unsqueeze(0) # Add a batch dimension ) + + +def resolve_cp_group( + static_cp_group: dist.ProcessGroup, packed_seq_params: PackedSeqParams = None +) -> dist.ProcessGroup: + """Return the packed-sequence CP group when present, otherwise the static group.""" + if packed_seq_params is not None and packed_seq_params.cp_group is not None: + return packed_seq_params.cp_group + return static_cp_group diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py index 15ecc4485c9..908174d4e4e 100644 --- a/megatron/core/ssm/gated_delta_net/gdn.py +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -12,7 +12,7 @@ import torch.nn.functional as F from megatron.core import tensor_parallel -from megatron.core.context_parallel_layout import convert_module_input_tensors_cp_partition_mode +from megatron.core.context_parallel import contiguous_to_zigzag, zigzag_to_contiguous from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.jit import jit_fuser from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group @@ -131,17 +131,18 @@ def forward( 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() - back_to_input_converter = None - if self.config.linear_cp_mode == "chunkwise": - hidden_states, back_to_input_converter = convert_module_input_tensors_cp_partition_mode( - hidden_states=hidden_states, - packed_seq_params=packed_seq_params, - cp_group=cp_group_chunkwise, - tp_group=self.tp_group, - tp_cp_group=getattr(active_pg_collection, "tp_cp", None), - target_partition_mode="contiguous", - sequence_parallel=self.config.sequence_parallel, - config=self.config, + convert_from_zigzag = ( + self.config.linear_cp_mode == "chunkwise" + and cp_size_chunkwise > 1 + and self.config.cp_partition_mode == "zigzag" + ) + if convert_from_zigzag: + hidden_states = zigzag_to_contiguous( + hidden_states, + cp_group_chunkwise, + self.config.sequence_parallel, + self.tp_group, + getattr(active_pg_collection, "tp_cp", None), ) seq_len_local, batch, _ = hidden_states.shape @@ -160,7 +161,10 @@ def forward( ( packed_seq_params is not None and packed_seq_params.qkv_format == "thd" - and packed_seq_params.cp_partition_mode != "zigzag" + and getattr( + packed_seq_params, "cp_partition_mode", self.config.cp_partition_mode + ) + != "zigzag" ) or ( (packed_seq_params is None or packed_seq_params.qkv_format != "thd") @@ -265,9 +269,13 @@ def _checkpointed_compute(hidden_states): chunkwise_cp_context, ) - if back_to_input_converter is not None: - out = back_to_input_converter.convert( - out, seq_dim=0, sequence_parallel=self.config.sequence_parallel + if convert_from_zigzag: + out = contiguous_to_zigzag( + out, + cp_group_chunkwise, + self.config.sequence_parallel, + self.tp_group, + getattr(active_pg_collection, "tp_cp", None), ) return out, out_bias diff --git a/megatron/core/ssm/gated_delta_net/kda.py b/megatron/core/ssm/gated_delta_net/kda.py index 04cb3ae439d..ce18221e29f 100644 --- a/megatron/core/ssm/gated_delta_net/kda.py +++ b/megatron/core/ssm/gated_delta_net/kda.py @@ -11,7 +11,7 @@ import torch.nn.functional as F from megatron.core import tensor_parallel -from megatron.core.context_parallel_layout import convert_module_input_tensors_cp_partition_mode +from megatron.core.context_parallel import contiguous_to_zigzag, zigzag_to_contiguous from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.jit import jit_fuser from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group @@ -225,21 +225,21 @@ def forward( 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() - back_to_input_converter = None - if self.config.linear_cp_mode == "chunkwise": - hidden_states, back_to_input_converter = convert_module_input_tensors_cp_partition_mode( - hidden_states=hidden_states, - packed_seq_params=packed_seq_params, - cp_group=cp_group_chunkwise, - tp_group=self.tp_group, - tp_cp_group=getattr(active_pg_collection, "tp_cp", None), - target_partition_mode="contiguous", - sequence_parallel=self.config.sequence_parallel, - config=self.config, + convert_from_zigzag = ( + self.config.linear_cp_mode == "chunkwise" + and cp_size_chunkwise > 1 + and self.config.cp_partition_mode == "zigzag" + ) + if convert_from_zigzag: + hidden_states = zigzag_to_contiguous( + hidden_states, + cp_group_chunkwise, + self.config.sequence_parallel, + self.tp_group, + getattr(active_pg_collection, "tp_cp", None), ) seq_len_local, batch, _ = hidden_states.shape @@ -257,7 +257,10 @@ def forward( ( packed_seq_params is not None and packed_seq_params.qkv_format == "thd" - and packed_seq_params.cp_partition_mode != "zigzag" + and getattr( + packed_seq_params, "cp_partition_mode", self.config.cp_partition_mode + ) + != "zigzag" ) or ( (packed_seq_params is None or packed_seq_params.qkv_format != "thd") @@ -349,9 +352,13 @@ def _checkpointed_compute(hidden_states): chunkwise_cp_context, ) - if back_to_input_converter is not None: - out = back_to_input_converter.convert( - out, seq_dim=0, sequence_parallel=self.config.sequence_parallel + if convert_from_zigzag: + out = contiguous_to_zigzag( + out, + cp_group_chunkwise, + self.config.sequence_parallel, + self.tp_group, + getattr(active_pg_collection, "tp_cp", None), ) return out, out_bias diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e8cd33aedb8..c2063b58242 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -627,7 +627,7 @@ class TransformerConfig(ModelParallelConfig): "moe": recompute the MoE layer. "shared_experts": recompute the shared experts in the MoE layer. "mhc": recompute HyperConnection intermediate activations via - CheckpointWithoutOutput + MHCCheckpointManager. Requires + CheckpointWithoutOutput + CheckpointWithoutOutputManager. 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 @@ -1258,8 +1258,8 @@ class TransformerConfig(ModelParallelConfig): """Number of layers in each mHC recompute group. Layers are grouped in their local transformer-block order. The last layer in each group leaves - its final MLP BDA output uncheckpointed and closes the current ``MHCCheckpointManager``; a new - manager is created for the next group. + its final MLP BDA output uncheckpointed and closes the current + ``CheckpointWithoutOutputManager``; a new manager is created for the next group. In the standard forward path, the group-ending layer registers a unified recompute hook on its MLP BDA output. In the fine-grained expert-parallel overlap schedule, the group outputs are diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 7f70205a51c..e9169d77144 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -334,6 +334,11 @@ def forward_step(data_iterator, model: HybridModel): cu_seqlens_for_params = cu_seqlens packed_seq_params = PackedSeqParams( qkv_format="thd", + cp_partition_mode=( + "contiguous" + if args.linear_cp_mode == "chunkwise" and args.context_parallel_size > 1 + else "zigzag" + ), cu_seqlens_q=cu_seqlens_for_params, cu_seqlens_kv=cu_seqlens_for_params, cu_seqlens_q_padded=cu_seqlens_padded, From bc17ee52219adf2d2d57775c64ff0def0a5e8f70 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Fri, 28 Aug 2026 06:17:13 +0000 Subject: [PATCH 33/34] fix: restrict legacy checkpoint deserialization Signed-off-by: svcnemo-autobot --- megatron/core/model_parallel_config.py | 15 ++++++++++ megatron/core/msc_utils.py | 9 +++++- megatron/training/checkpointing.py | 21 ++------------ tests/unit_tests/test_checkpointing.py | 39 ++++++++++++++++++++++++++ 4 files changed, 64 insertions(+), 20 deletions(-) diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index c22d8784495..c3e32cca772 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -7,6 +7,21 @@ import torch +def _parse_pad_packed_seq_alignment(value): + """Parse THD packed-sequence padding alignment.""" + if value == "max": + return value + try: + alignment = int(value) + except (TypeError, ValueError) as exc: + raise ValueError( + "pad_packed_seq_alignment must be 'max' or a positive integer alignment." + ) from exc + if alignment <= 0: + raise ValueError("pad_packed_seq_alignment must be 'max' or a positive integer alignment.") + return alignment + + def resolve_tensor_parallel_weight_shards( tensor_model_parallel_size: int, tensor_parallel_num_weight_shards: Optional[int], diff --git a/megatron/core/msc_utils.py b/megatron/core/msc_utils.py index c7dd7dc07d6..1764ed756d9 100644 --- a/megatron/core/msc_utils.py +++ b/megatron/core/msc_utils.py @@ -55,6 +55,13 @@ def __setstate__(self, state): MultiStorageClientFeature = _FeatureFlag(default=False) +def open_file(*args, **kwargs): + """Open a path through MSC when enabled, otherwise use the builtin open.""" + if MultiStorageClientFeature.is_enabled(): + return MultiStorageClientFeature.import_package().open(*args, **kwargs) + return open(*args, **kwargs) + + class MaybeMultiStorageClient: """ Helper class to use MultiStorageClient @@ -108,4 +115,4 @@ def __dir__(self): maybe_msc = MaybeMultiStorageClient() -__all__ = ['MultiStorageClientFeature', 'maybe_msc'] +__all__ = ['MultiStorageClientFeature', 'maybe_msc', 'open_file'] diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 395d5a57eca..a9f6fc97146 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1735,24 +1735,7 @@ def _load_base_checkpoint( load_dir, iteration, release, return_base_dir=False ) try: - state_dict = torch.load(checkpoint_name, map_location='cpu') - except ModuleNotFoundError: - from megatron.legacy.fp16_deprecated import loss_scaler - - # For backward compatibility. - if not rank0: - print_rank_0(' > deserializing using the old code structure ...') - sys.modules['fp16.loss_scaler'] = sys.modules[ - 'megatron.legacy.fp16_deprecated.loss_scaler' - ] - sys.modules['megatron.fp16.loss_scaler'] = sys.modules[ - 'megatron.legacy.fp16_deprecated.loss_scaler' - ] - sys.modules['megatron.model'] = sys.modules['megatron.legacy.model'] - state_dict = torch.load(checkpoint_name, map_location='cpu') - sys.modules.pop('fp16.loss_scaler', None) - sys.modules.pop('megatron.fp16.loss_scaler', None) - sys.modules.pop('megatron.model', None) + state_dict = torch.load(checkpoint_name, map_location='cpu', weights_only=True) except Exception as e: print('could not load the checkpoint') print(e) @@ -2619,7 +2602,7 @@ def load_biencoder_checkpoint( ) ) - state_dict = torch.load(checkpoint_name, map_location='cpu') + state_dict = torch.load(checkpoint_name, map_location='cpu', weights_only=True) ret_state_dict = state_dict['model'] if only_query_model: diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index 5e35d3bda9f..e4cec83056b 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -1,6 +1,7 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Note: --ckpt-format torch_dist has tests in tests/unit_tests/dist_checkpointing. import os +from pathlib import Path from types import SimpleNamespace from typing import Optional from unittest import mock @@ -244,6 +245,44 @@ def init_model_parallel(): unset_num_microbatches_calculator() +def test_legacy_checkpoint_load_is_restricted(create_ckpt_load_args, tmp_path): + """Legacy checkpoint loading must use the restricted torch deserializer.""" + + load_dir = tmp_path / "test_legacy_checkpoint_load" + load_dir.mkdir() + create_checkpoint(load_dir, "torch") + create_ckpt_load_args.ckpt_format = "torch" + + with mock.patch( + "megatron.training.checkpointing.torch.load", + return_value={"args": "dummy", "iteration": 123}, + ) as load: + _load_base_checkpoint(load_dir, create_ckpt_load_args, rank0=True) + + load.assert_called_once() + assert load.call_args.kwargs["weights_only"] is True + + +def test_legacy_checkpoint_load_rejects_unsafe_pickle(create_ckpt_load_args, tmp_path): + """Loading a crafted legacy checkpoint must not execute its pickle payload.""" + + class UnsafeCheckpointPayload: + def __reduce__(self): + return Path.write_text, (tmp_path / "executed", "unsafe") + + load_dir = tmp_path / "test_unsafe_legacy_checkpoint" + load_dir.mkdir() + create_checkpoint(load_dir, "torch") + create_ckpt_load_args.ckpt_format = "torch" + checkpoint_name = load_dir / "iter_0000123" / "mp_rank_00" / "model_optim_rng.pt" + torch.save({"payload": UnsafeCheckpointPayload()}, checkpoint_name) + + with pytest.raises(SystemExit): + _load_base_checkpoint(load_dir, create_ckpt_load_args, rank0=True) + + assert not (tmp_path / "executed").exists() + + @pytest.mark.parametrize("ckpt_format", ["torch_dcp"]) def test_load_base_checkpoint( init_model_parallel, create_ckpt_load_args, ckpt_format, tmp_path_dist_ckpt From f205386d1f5235f429f48a4030215f62d95d37f0 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Sat, 29 Aug 2026 02:23:29 +0000 Subject: [PATCH 34/34] style: format gated delta net CP checks Signed-off-by: svcnemo-autobot --- megatron/core/ssm/gated_delta_net/gdn.py | 4 +--- megatron/core/ssm/gated_delta_net/kda.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py index 908174d4e4e..e56d3a95d72 100644 --- a/megatron/core/ssm/gated_delta_net/gdn.py +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -161,9 +161,7 @@ def forward( ( packed_seq_params is not None and packed_seq_params.qkv_format == "thd" - and getattr( - packed_seq_params, "cp_partition_mode", self.config.cp_partition_mode - ) + and getattr(packed_seq_params, "cp_partition_mode", self.config.cp_partition_mode) != "zigzag" ) or ( diff --git a/megatron/core/ssm/gated_delta_net/kda.py b/megatron/core/ssm/gated_delta_net/kda.py index ce18221e29f..cad16568472 100644 --- a/megatron/core/ssm/gated_delta_net/kda.py +++ b/megatron/core/ssm/gated_delta_net/kda.py @@ -257,9 +257,7 @@ def forward( ( packed_seq_params is not None and packed_seq_params.qkv_format == "thd" - and getattr( - packed_seq_params, "cp_partition_mode", self.config.cp_partition_mode - ) + and getattr(packed_seq_params, "cp_partition_mode", self.config.cp_partition_mode) != "zigzag" ) or (