diff --git a/hybrid_builders.py b/hybrid_builders.py index 05b219277ef..d95002e1a21 100644 --- a/hybrid_builders.py +++ b/hybrid_builders.py @@ -2,8 +2,8 @@ from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_inference_stack_spec from megatron.core.models.hybrid.hybrid_model import HybridModel -from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.spec_utils import import_module +from megatron.core.transformer import MLATransformerConfig, TransformerConfig +from megatron.core.transformer.spec_utils import ModuleSpec, import_module from megatron.training import print_rank_0 from megatron.training.arguments import core_transformer_config_from_args from model_provider import count_parameters_in_layer @@ -13,6 +13,52 @@ def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None, print_rank_0('building Hybrid model ...') if config is None: config = core_transformer_config_from_args(args, TransformerConfig) + # MLA (and DSv4 hybrid) require MLATransformerConfig so that its __post_init__ runs + # the dsv4_hybrid derivation. The hybrid pretrain path can hand us a plain + # TransformerConfig, which silently skips that derivation; rebuild as MLA to match GPT. + if args.multi_latent_attention and not isinstance(config, MLATransformerConfig): + config = core_transformer_config_from_args(args) + # DSv4-hybrid head-dim contract: qk_head_dim and kv_lora_rank are derived from + # v_head_dim and qk_pos_emb_head_dim (MLATransformerConfig.__post_init__ does this for the + # GPT path). The hybrid config can reach here without that derivation applied, which breaks + # the MLA up-proj / fused-rope contract (q head dim must equal qk_head_dim + qk_pos_emb_head + # _dim == v_head_dim). Apply it for any DSv4 MLA attention: experimental_attention_variant + # == dsv4_hybrid, OR the layer pattern uses a DSv4 attention symbol (D/C/H/W). Idempotent. + _pattern = getattr(args, "hybrid_layer_pattern", None) or "" + _uses_dsv4_attn = ( + getattr(args, "experimental_attention_variant", None) == "dsv4_hybrid" + or any(sym in _pattern for sym in ("C", "H", "W")) + ) + if _uses_dsv4_attn: + derived = config.v_head_dim - config.qk_pos_emb_head_dim + if config.qk_head_dim != derived or config.kv_lora_rank != derived: + print_rank_0( + f"[hybrid dsv4] deriving qk_head_dim/kv_lora_rank = {config.v_head_dim} - " + f"{config.qk_pos_emb_head_dim} = {derived} (was qk_head_dim={config.qk_head_dim}, " + f"kv_lora_rank={config.kv_lora_rank})" + ) + config.qk_head_dim = derived + config.kv_lora_rank = derived + # 'C'/'H'/'W' layers carry their compress ratio via the spec, but array-driven 'D' layers + # AND the indexer-loss logger (which counts ratio==4 layers) read + # config.csa_compress_ratios. When not given explicitly, derive it from the pattern + # symbols (C->4, H->128, W/D/other->0) so the array is consistent with the symbols and + # the indexer loss is normalized correctly; pad MTP depths with 0. An explicit + # --csa-compress-ratios is always respected. + if config.csa_compress_ratios is None: + ratio_map = {"C": 4, "H": 128} + # One entry per ACTUAL layer: main layers, then every MTP layer of every MTP depth + # (a depth can hold multiple hybrid layers, e.g. "/MD-E"), mirroring the arguments.py + # derivation. Padding by mtp_num_layers (depth count) would be too short and an MTP + # attention that isn't first would IndexError at num_layers + layer_number - 1. + sections = _pattern.split("/") + ratios = [ratio_map.get(c, 0) for c in sections[0].replace("|", "")] + for mtp_sec in sections[1:]: + ratios += [ratio_map.get(c, 0) for c in mtp_sec.replace("|", "")] + config.csa_compress_ratios = ratios + print_rank_0( + f"[hybrid dsv4] derived csa_compress_ratios from pattern symbols: {ratios}" + ) if config.transformer_impl == "inference_optimized": hybrid_stack_spec = hybrid_inference_stack_spec @@ -21,6 +67,11 @@ def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None, ), "inference_fuse_tp_communication is not supported for HybridModel" elif args.spec is not None: hybrid_stack_spec = import_module(args.spec) + # Allow config-aware specs: if --spec resolves to a callable (not a ModuleSpec), + # call it with config to build the stack spec (e.g. hybrid_dsv4_stack_spec, which + # wires the DSv4 CompressedSparseAttention into the 'D' layer per config). + if not isinstance(hybrid_stack_spec, ModuleSpec) and callable(hybrid_stack_spec): + hybrid_stack_spec = hybrid_stack_spec(config) else: raise ValueError("You must provide a valid hybrid layer spec via --spec") diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 920b334a9e4..d31a5c6d7ae 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -2011,12 +2011,28 @@ def __init__( tp_group_for_te = None if is_te_min_version("2.14.0"): - extra_kwargs["single_grouped_weight"] = getattr( - config, "moe_single_grouped_weight", False - ) - extra_kwargs["single_grouped_bias"] = getattr( - config, "moe_single_grouped_bias", False - ) + # nemo_26.04 ships TE 2.14.0+71bbefbf whose GroupedLinear.__init__ does NOT + # yet accept single_grouped_{weight,bias}, even though the version string + # passes is_te_min_version("2.14.0"). Introspect the signature instead of + # version-gating, mirroring the patch in dsv4_fused_attn / main_megatron. + # The GroupedLinear.__init__ signature is constant for a given TE install, so + # introspect once and cache at module scope rather than on every TEGroupedLinear + # instantiation (matters for large MoE models with many expert groups). + global _TE_GROUPED_LINEAR_INIT_PARAMS + try: + _gl_params = _TE_GROUPED_LINEAR_INIT_PARAMS + except NameError: + _gl_params = _TE_GROUPED_LINEAR_INIT_PARAMS = set( + inspect.signature(te.pytorch.GroupedLinear.__init__).parameters + ) + if "single_grouped_weight" in _gl_params: + extra_kwargs["single_grouped_weight"] = getattr( + config, "moe_single_grouped_weight", False + ) + if "single_grouped_bias" in _gl_params: + extra_kwargs["single_grouped_bias"] = getattr( + config, "moe_single_grouped_bias", False + ) self.te_quant_params: Optional[TEQuantizationParams] = None quant_config = get_quant_config_or_none(name, config.quant_recipe) diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 266df5046ea..d462122c1c2 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -26,12 +26,13 @@ from megatron.core.recompute import checkpointed_forward from megatron.core.tensor_parallel.random import CheckpointManager from megatron.core.transformer import TransformerConfig +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 +from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_layer import TransformerLayer from megatron.core.transformer.utils import ( @@ -52,12 +53,15 @@ class HybridStackSubmodules: gdn_layer: Union[ModuleSpec, type] = IdentityOp attention_layer: Union[ModuleSpec, type] = IdentityOp dsa_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(MegatronModule): +class HyperConnectionHybridLayer(GraphableMegatronModule): """Layer-boundary mHC wrapper for HybridStack layers. Hybrid layers already own their local residual paths. For this initial @@ -76,6 +80,31 @@ class HyperConnectionHybridLayer(MegatronModule): 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: @@ -88,6 +117,138 @@ def __init__(self, config: TransformerConfig, layer: MegatronModule) -> None: 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``. + """ + 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 + + 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). + """ + 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_out = list(self.inner_layer._te_cuda_graph_capture(aggregated)) + # 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. + """ + 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'): @@ -103,9 +264,20 @@ def _call_inner_layer( 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 = self.inner_layer( + output = inner( hidden_states=hidden_states, attention_mask=attention_mask, inference_context=inference_context, @@ -113,6 +285,7 @@ def _call_inner_layer( 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: @@ -122,7 +295,7 @@ def _call_inner_layer( # 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 = self.inner_layer( + output = inner( hidden_states=hidden_states, attention_mask=attention_mask, inference_context=inference_context, @@ -134,22 +307,92 @@ def _call_inner_layer( return output[0], context return output, None - def forward( + 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, + ) + 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``). + """ + """Run the wrapped hybrid layer through one layer-boundary mHC update.""" aggregated, h_res, h_post, residual = self.hyper_connection( hidden_states, mhc_recompute_manager=mhc_recompute_manager ) - layer_output, context = self._call_inner_layer( + fast_path_result = self._call_inner_transformer_layer_without_local_bda( aggregated, attention_mask, inference_context, @@ -157,32 +400,42 @@ def forward( 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. The temporary [s, b, C] tensor here is the - # simplest correct form; a future optimization could fuse the subtraction - # into `fused_h_res_h_post_bda` to avoid the allocation. - # Sanity check: this contract requires the inner layer to preserve shape; - # any mismatch indicates a future layer type is breaking the residual - # assumption and would silently corrupt the n-stream state. + + 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 inner layers to preserve " - f"hidden-state shape. Got {tuple(layer_output.shape)} from inner layer " - f"vs {tuple(aggregated.shape)} input; layer must add its own residual." + "HyperConnectionHybridLayer requires wrapped branches to preserve " + f"hidden-state shape. Got {tuple(layer_output.shape)} from wrapped branch " + f"vs {tuple(aggregated.shape)} input." ) - # `fp32_residual_connection=True` may cause some inner layers (e.g., - # MambaLayer) to return `layer_output` in fp32 while `aggregated` is in - # compute dtype; explicitly upcast `aggregated` so the subtraction stays - # in fp32 instead of relying on PyTorch's implicit promotion. - if self.config.fp32_residual_connection and aggregated.dtype != layer_output.dtype: - aggregated = aggregated.to(layer_output.dtype) - layer_delta = layer_output - aggregated - # `dropout_prob=0.0` already disables dropout regardless of training mode; - # `training=self.training` is more semantically accurate than hard-coding - # False during a training-mode forward. is_last_in_recompute_block = bool( mhc_recompute_manager is not None and getattr(mhc_recompute_manager, "is_last_layer_in_recompute_block", False) @@ -193,10 +446,10 @@ def forward( h_res, residual, h_post, - (layer_delta, None), - dropout_prob=0.0, + layer_output_with_bias, + dropout_prob=dropout_prob, training=self.training, - fused=False, + fused=bias_dropout_fusion, manager=mhc_bda_manager, ) # In `HyperConnectionTransformerLayer` the n-stream output stays in compute @@ -238,6 +491,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__( @@ -253,6 +507,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: """ @@ -264,6 +519,7 @@ 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" @@ -326,6 +582,40 @@ def __init__( 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( submodules.mlp_layer, @@ -341,6 +631,7 @@ def __init__( config=self.config, layer_number=layer_number, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, add_layer_offset=False, name=(name + f".layers.{i}") if name is not None else None, ) @@ -356,6 +647,8 @@ def __init__( ) else: raise ValueError("unexpected layer_type") + 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) @@ -371,7 +664,11 @@ def __init__( eps=self.config.layernorm_epsilon, ) - if self.config.enable_hyper_connections and self.post_process: + # 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 = nn.Parameter(torch.randn(hc_mult, hc_dim)) @@ -383,6 +680,16 @@ def __init__( setattr(self.hc_head_base, 'sequence_parallel', True) setattr(self.hc_head_scale, 'sequence_parallel', True) + @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. @@ -466,7 +773,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. @@ -482,7 +790,11 @@ 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) @@ -495,7 +807,11 @@ def forward( if isinstance(hidden_states, WrappedTensor): hidden_states = hidden_states.unwrap() - if self.config.enable_hyper_connections and self.pre_process: + # 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.num_residual_streams ) @@ -569,6 +885,7 @@ 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), ) else: @@ -577,6 +894,7 @@ def get_inner_quant_context(config, layer_number): inner_quant_context = get_inner_quant_context( self.config, layer.layer_number - 1 ) + mhc_manager = mhc_layer_managers[l_no] if mhc_manager is not None: mhc_manager.is_last_layer_in_recompute_block = ( @@ -594,6 +912,8 @@ def get_inner_quant_context(config, layer_number): 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 ): @@ -619,7 +939,22 @@ def get_inner_quant_context(config, layer_number): is_last_in_recompute_block=mhc_is_last_in_recompute_block[l_no], ) - if self.config.enable_hyper_connections and self.post_process: + # 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_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( hidden_states, self.hc_head_fn, @@ -639,6 +974,8 @@ def get_inner_quant_context(config, layer_number): inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True ) + if mhc_multistream is not None: + return hidden_states, mhc_multistream return hidden_states def sharded_state_dict( diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index 67103fe67f1..a8d2006c3b0 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -18,11 +18,16 @@ class Symbols: GDN = 'G' ATTENTION = "*" DS_ATTENTION = "D" + 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, MLP, MOE} + VALID_LAYERS = {MAMBA, GDN, 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]: @@ -173,10 +178,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, '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} + {'*': 1, 'C': 0, 'D': 0, 'G': 0, 'H': 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()} @@ -292,9 +297,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: - 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]: @@ -320,9 +328,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: - 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 5b968f720c0..a18da8b5452 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -73,7 +73,12 @@ submodules=MultiTokenPredictionLayerSubmodules( enorm=TENorm, hnorm=TENorm, + # 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, mtp_model_layer=None, # Built via pattern + hybrid_submodules layer_norm=TENorm, ), @@ -292,7 +297,11 @@ submodules=MultiTokenPredictionLayerSubmodules( enorm=TENorm, hnorm=TENorm, + # 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, mtp_model_layer=None, # Built via pattern + hybrid_submodules layer_norm=TENorm, ), @@ -307,3 +316,56 @@ # Backward-compatible aliases mamba_stack_spec = hybrid_stack_spec mamba_inference_stack_spec = hybrid_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 ``MLASelfAttention + + DSAttention`` (DSv3-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/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 80a86e0cfac..1a5b852c7bd 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -504,15 +504,32 @@ def forward( # be None, so this assert will succeed. # assert attention_mask is None, "The attention mask is ignored and should be set to None" + # Pass input_ids to decoder for hash-based MoE routing. + decoder_extra_block_kwargs = {} + if self.config.moe_n_hash_layers > 0 and input_ids is not None: + decoder_extra_block_kwargs['input_ids'] = input_ids + # Run decoder. - hidden_states = self.decoder( + decoder_output = self.decoder( hidden_states=decoder_input, attention_mask=attention_mask, inference_context=inference_context, rotary_pos_emb=rotary_pos_emb, packed_seq_params=packed_seq_params, padding_mask=padding_mask, + **decoder_extra_block_kwargs, ) + # HybridStack.forward returns a single Tensor in the common case, but a 2-tuple + # (hidden_states, mhc_multistream) in exactly one case: enable_hyper_connections and + # post_process and mtp_num_layers > 0 and not is_mtp_layer — where MTP's mHC branch + # needs the pre-contraction multi-stream tensor for `_concat_embeddings`. Any other + # tuple return would be misinterpreted here, so keep that contract in sync with + # HybridStack.forward (see hybrid_block.py). + if isinstance(decoder_output, tuple): + hidden_states, mhc_multistream = decoder_output + else: + hidden_states = decoder_output + mhc_multistream = None output_weight = None if self.share_embeddings_and_output_weights: @@ -534,6 +551,7 @@ def forward( input_ids=input_ids, position_ids=position_ids, hidden_states=hidden_states, + mhc_multistream=mhc_multistream, attention_mask=attention_mask, inference_params=inference_params, rotary_pos_emb=rotary_pos_emb, diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 884444557a0..3ec40c3b2d3 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1728,12 +1728,43 @@ def _layer_is_graphable(layer, config): return True # import modules here to avoid a circular import + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer from megatron.core.ssm.mamba_layer import MambaLayer from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.mlp import MLP from megatron.core.transformer.moe.moe_layer import MoELayer from megatron.core.transformer.transformer_layer import TransformerLayer + # mHC wrapper: graphability is decided by the inner layer's type/scope. Non-MoE + # inner layers are captured as one whole-wrapper graph (mHC aggregate + inner + BDA). + # MoE inner layers use partial capture: the wrapper graphs the deterministic prefix + # (mHC aggregate + router/preprocess) and runs the expert all-to-all + mHC BDA eagerly + # (the all-to-all is not graph-safe), mirroring the GPT HyperConnectionTransformerLayer. + if isinstance(layer, HyperConnectionHybridLayer): + inner = layer.inner_layer + if isinstance(inner, MambaLayer) and CudaGraphModule.mamba in config.cuda_graph_modules: + return True + if isinstance(inner, TransformerLayer): + if isinstance(inner.mlp, MoELayer): + # MoE inner: graphable via partial (router/preprocess) capture. + if ( + CudaGraphModule.moe in config.cuda_graph_modules + or CudaGraphModule.moe_router in config.cuda_graph_modules + or CudaGraphModule.moe_preprocess in config.cuda_graph_modules + ): + return True + # attn-only scope on an MoE inner: the (identity) attention prefix has + # nothing graph-worthy, so leave eager. + return False + if CudaGraphModule.attn in config.cuda_graph_modules and not ( + isinstance(inner.self_attention, IdentityOp) + and isinstance(inner.cross_attention, IdentityOp) + ): + return True + if CudaGraphModule.mlp in config.cuda_graph_modules and isinstance(inner.mlp, MLP): + return True + return False + if isinstance(layer, MambaLayer) and CudaGraphModule.mamba in config.cuda_graph_modules: # mamba layer. return True diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 2e450de4826..de7e3f83cd7 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -571,7 +571,7 @@ class CompressedSparseAttention(MegatronModule): provides compressor and indexer submodule specs; this ``__init__`` inspects ``config.csa_compress_ratios[layer_idx]`` and conditionally builds them: - * ``ratio == 0``: window-only (compressor and indexer NOT built) + * ``ratio == 0``: window-only (compressor and indexer NOT built) — the 'W' layer symbol * ``ratio == 4``: window + 4x compressed + learned Indexer (both built) * ``ratio == 128``: window + 128x compressed, attend to all (compressor built only) """ @@ -617,7 +617,7 @@ def __init__( # Learnable attention sink per head self.attn_sink = nn.Parameter(torch.zeros(self.n_local_heads, dtype=torch.float32)) - # Conditionally build Compressor (ratio > 1) + # Conditionally build Compressor (ratio > 1). ratio == 0 is window-only ('W'): not built. if self.compress_ratio > 1 and submodules.compressor is not None: self.compressor = build_module( submodules.compressor, @@ -631,7 +631,7 @@ def __init__( else: self.compressor = None - # Conditionally build Indexer (ratio == 4) + # Conditionally build Indexer (ratio == 4). ratio == 0 is window-only ('W'): not built. if ( self.compress_ratio == 4 and not config.csa_dense_mode diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index fb526f1ca36..a8f5a65a185 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -71,6 +71,7 @@ def __init__( pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, is_mtp_layer: bool = False, + compress_ratio: Optional[int] = None, name: str | None = None, ) -> None: @@ -113,11 +114,14 @@ def __init__( self.softmax_scale = None - if is_mtp_layer: - layer_idx = self.config.num_layers + layer_number - 1 - compress_ratio = self.config.csa_compress_ratios[layer_idx] - else: - compress_ratio = self.config.csa_compress_ratios[layer_number - 1] + # Per-layer compress ratio. When set explicitly (e.g. hybrid 'C'/'H' layer symbols + # pass compress_ratio=4/128 via the spec), use it directly; otherwise fall back to the + # per-(global)-layer csa_compress_ratios array (GPT-parity / array-driven path). + _ratio_idx = self.config.num_layers + layer_number - 1 if is_mtp_layer else layer_number - 1 + if compress_ratio is None: + compress_ratio = self.config.csa_compress_ratios[_ratio_idx] + # compress_ratio == 0 is a sliding-window-only layer (the 'W' symbol): no compressor / + # no top-k indexer (see CompressedSparseAttention) AND standard (non-YARN) rope. use_compressed_yarn = compress_ratio > 1 rope_base = ( self.config.csa_compress_rotary_base if use_compressed_yarn else self.config.rotary_base @@ -416,8 +420,9 @@ def __init__( attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, - pp_layer_offset: Optional[int] = None, is_mtp_layer: bool = False, + pp_layer_offset: Optional[int] = None, + compress_ratio: Optional[int] = None, name: str | None = None, ): if pg_collection is None: @@ -431,8 +436,9 @@ def __init__( attention_type="self", cp_comm_type=cp_comm_type, pg_collection=pg_collection, - pp_layer_offset=pp_layer_offset, is_mtp_layer=is_mtp_layer, + pp_layer_offset=pp_layer_offset, + compress_ratio=compress_ratio, name=name, ) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 5ee9d07b886..da871c54522 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -73,8 +73,19 @@ def save_loss_to_tracker( return tracker = DSAIndexerLossLoggingHelper.tracker + # Tracker must be at least max(num_layers, layer_number) so hybrid MTP layers + # (whose layer_number can exceed config.num_layers + config.mtp_num_layers when + # each MTP depth contains multiple hybrid layers) don't index out of bounds. + # Grow lazily; with PP=1 every rank takes the same path, so sizes stay consistent. + needed = max(num_layers, layer_number) if "values" not in tracker: - tracker["values"] = torch.zeros(num_layers, device=torch.cuda.current_device()) + tracker["values"] = torch.zeros(needed, device=torch.cuda.current_device()) + elif tracker["values"].shape[0] < needed: + grown = torch.zeros( + needed, device=tracker["values"].device, dtype=tracker["values"].dtype + ) + grown[: tracker["values"].shape[0]] = tracker["values"] + tracker["values"] = grown tracker["values"][layer_number - 1] += loss.detach() tracker["reduce_group"] = reduce_group tracker["avg_group"] = avg_group @@ -102,15 +113,42 @@ def reduce_loss_in_tracker(num_layers: Optional[int] = None): tracker on ranks where no indexer layer ran. """ tracker = DSAIndexerLossLoggingHelper.tracker + pp_group = parallel_state.get_pipeline_model_parallel_group() + + # Agree on a consistent tracker size across the PP group BEFORE the collective. + # Ranks owning indexer layers may have grown the tracker via save_loss_to_tracker + # (e.g. an MTP layer whose layer_number exceeds num_layers), while ranks without any + # indexer layer have only a num_layers-sized (or absent) tracker. all_reduce requires + # identical shapes on every rank, so reduce-MAX the local size first, then pad to it + # (otherwise PP>1 hangs / errors on mismatched sizes). + # The agreed size (max over the PP group) is constant across iterations (num_layers and + # the layer numbering don't change), so compute it once and cache it. This avoids a + # per-iteration CPU-GPU sync (.item()); the size-negotiation all_reduce + .item() runs + # only on the first call. Every PP rank caches on the same (first) call, so later steps + # all skip it consistently. + if tracker.get("agreed_size") is not None: + size = tracker["agreed_size"] + else: + local_size = tracker["values"].shape[0] if "values" in tracker else (num_layers or 0) + size_t = torch.tensor( + [local_size], device=torch.cuda.current_device(), dtype=torch.long + ) + torch.distributed.all_reduce(size_t, op=torch.distributed.ReduceOp.MAX, group=pp_group) + size = int(size_t.item()) + tracker["agreed_size"] = size + if size == 0: + return if "values" not in tracker: - if num_layers is None: - return - tracker["values"] = torch.zeros(num_layers, device=torch.cuda.current_device()) + tracker["values"] = torch.zeros(size, device=torch.cuda.current_device()) + elif tracker["values"].shape[0] < size: + grown = torch.zeros( + size, device=tracker["values"].device, dtype=tracker["values"].dtype + ) + grown[: tracker["values"].shape[0]] = tracker["values"] + tracker["values"] = grown values = tracker["values"] - torch.distributed.all_reduce( - values, group=parallel_state.get_pipeline_model_parallel_group() - ) + torch.distributed.all_reduce(values, group=pp_group) # Reduce indexer losses across ranks. if tracker.get('reduce_group') is not None: torch.distributed.all_reduce(values, group=tracker.get('reduce_group')) @@ -1235,10 +1273,18 @@ def forward( ) # Save indexer loss for logging if indexer_loss_coeff > 0: + # On HybridModel, each MTP depth can contain multiple hybrid layers + # (e.g. `/MD-E` is 4 layers per depth), so `num_layers + mtp_num_layers` + # is an undercount when mtp_num_layers is depth, not layer count. Take + # the max with self.layer_number so the tracker grows to cover the + # largest layer index seen on this rank. DSAIndexerLossLoggingHelper.save_loss_to_tracker( loss=indexer_loss, layer_number=self.layer_number, - num_layers=self.config.num_layers + (self.config.mtp_num_layers or 0), + num_layers=max( + self.layer_number, + self.config.num_layers + (self.config.mtp_num_layers or 0), + ), ) # =================================== diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index 3bb7813fcd7..89d92c88af5 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -19,6 +19,7 @@ _MHC_COMPUTE_H_EPS = 1e-6 +# dynamic=True handles the hybrid mHC variable-shape path (was blanket-disabled) @torch.compile def _sinkhorn_iterations(input_logits: Tensor, num_iterations: int, eps: float) -> Tensor: M = input_logits.softmax(dim=-1) + eps @@ -62,12 +63,14 @@ def native_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6 return SinkhornKnopp.apply(input_logits, num_iterations, eps) +# dynamic=True handles the hybrid mHC variable-shape path (was blanket-disabled) @torch.compile def native_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: """Native n-stream weighted aggregation: out = sum_j(h_pre_j * x_j).""" return (x * h_pre.unsqueeze(-1)).sum(dim=2) +# dynamic=True handles the hybrid mHC variable-shape path (was blanket-disabled) @torch.compile def native_h_post_bda( h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] @@ -84,6 +87,7 @@ def native_h_post_bda( return x_expanded + mixed +# dynamic=True handles the hybrid mHC variable-shape path (was blanket-disabled) @torch.compile def native_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: """Native fused projection + RMS normalization.""" @@ -95,6 +99,7 @@ def native_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tenso return proj, r +# dynamic=True handles the hybrid mHC variable-shape path (was blanket-disabled) @torch.compile def native_fused_add_3(a: Tensor, b: Tensor, c: Tensor) -> Tensor: """Native 3-way elementwise add (torch.compile fuses into single kernel).""" @@ -258,6 +263,7 @@ def _projection_and_get_norm(self, x: Tensor) -> Tuple[Tensor, Tensor]: proj, r = self._proj_rms_op(x_2d, self.mapping_proj.weight, self.norm_eps) return proj.view(s, b, -1), r.view(s, b, 1) + # dynamic=True handles the hybrid mHC variable-shape path (was blanket-disabled) @torch.compile def _compute_h(self, proj: Tensor, r: Tensor) -> Tuple[Tensor, Tensor, Tensor]: """ @@ -338,6 +344,7 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: return h_pre, h_post, h_res + # dynamic=True handles the hybrid mHC variable-shape path (was blanket-disabled) @torch.compile def _apply_h_post(self, x: Tensor, h_post: Tensor) -> Tensor: """ @@ -437,6 +444,7 @@ def aggregate(self, x: Tensor, h_pre: Tensor) -> Tensor: x_streams = x.view(s, b, self.n, C) return self._h_aggregate_op(x_streams, h_pre) + # dynamic=True handles the hybrid mHC variable-shape path (was blanket-disabled) @torch.compile def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor: """ diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 01b0adea88d..580c8a5a650 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -176,6 +176,7 @@ def __init__( 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 if self.config.moe_n_hash_layers > 0: assert layer_number is not None, "layer_number is required for the hash-based router." @@ -497,7 +498,12 @@ def attach_and_log_load_balancing_loss( num_layers += self.config.mtp_num_layers if self.is_mtp_layer: - layer_number = self.layer_number + self.config.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) + layer_number = mtp_layer_number + self.config.num_layers else: layer_number = self.layer_number diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index 7e31b9fbb5b..2b7c3b56b36 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -1206,6 +1206,7 @@ def __init__( post_process=True, # MTP layer is self-contained pg_collection=pg_collection, is_mtp_layer=True, + mtp_layer_number=self.layer_number, name=(name + ".mtp_model_layer") if name is not None else None, ) elif self.config.mtp_num_layers is not None: @@ -1369,6 +1370,7 @@ def _proj_and_transformer_layer( self, hidden_states: Tensor, decoder_input: Tensor, + input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, padding_mask: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, @@ -1415,6 +1417,7 @@ def _proj_and_transformer_layer( rotary_pos_emb=rotary_pos_emb, inference_context=inference_params, packed_seq_params=packed_seq_params, + input_ids=input_ids, ) else: # GPT path: single TransformerLayer @@ -1431,6 +1434,7 @@ def _proj_and_transformer_layer( packed_seq_params=packed_seq_params, sequence_len_offset=sequence_len_offset, padding_mask=padding_mask, + input_ids=input_ids, ) if not self.mhc_enabled: @@ -1498,6 +1502,7 @@ def forward_single_position( hidden_states = self._proj_and_transformer_layer( hidden_states=hidden_states, decoder_input=decoder_input, + input_ids=next_token_ids, attention_mask=attention_mask, rotary_pos_emb=rotary_pos_emb, rotary_pos_cos=rotary_pos_cos, @@ -1511,6 +1516,7 @@ def _checkpointed_forward( self, hidden_states: Tensor, decoder_input: Tensor, + input_ids: Optional[Tensor] = None, attention_mask: Optional[Tensor] = None, padding_mask: Optional[Tensor] = None, context: Optional[Tensor] = None, @@ -1547,6 +1553,7 @@ def _checkpointed_forward( def custom_forward( hidden_states, decoder_input, + input_ids, attention_mask, padding_mask, context, @@ -1559,6 +1566,7 @@ def custom_forward( return self._proj_and_transformer_layer( hidden_states=hidden_states, decoder_input=decoder_input, + input_ids=input_ids, attention_mask=attention_mask, padding_mask=padding_mask, context=context, @@ -1606,6 +1614,7 @@ def checkpoint_handler(): parallel_state.get_tensor_model_parallel_group(), hidden_states, decoder_input, + input_ids, attention_mask, padding_mask, context, @@ -1626,6 +1635,7 @@ def checkpoint_handler(): self.config.distribute_saved_activations, hidden_states, decoder_input, + input_ids, attention_mask, padding_mask, context, @@ -1653,6 +1663,7 @@ def checkpoint_handler(): outputs = self._proj_and_transformer_layer( hidden_states=hidden_states, decoder_input=decoder_input, + input_ids=input_ids, attention_mask=attention_mask, padding_mask=padding_mask, context=context, @@ -1726,6 +1737,7 @@ def forward( hidden_states = self._checkpointed_forward( hidden_states=hidden_states, decoder_input=decoder_input, + input_ids=input_ids, attention_mask=attention_mask, padding_mask=padding_mask, context=context, @@ -1742,6 +1754,7 @@ def forward( hidden_states = self._proj_and_transformer_layer( hidden_states=hidden_states, decoder_input=decoder_input, + input_ids=input_ids, attention_mask=attention_mask, padding_mask=padding_mask, context=context, diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 10c8603dcac..8caffd7d004 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -333,7 +333,8 @@ class TransformerConfig(ModelParallelConfig): """Sliding window size for compressed sparse attention.""" csa_compress_ratios: Optional[List[int]] = None - """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...].""" + """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.""" @@ -1468,9 +1469,14 @@ def __post_init__(self): 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 equal " + 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( @@ -2423,7 +2429,7 @@ def __post_init__(self): 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: + 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)." diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index e29b57d5f39..b92efdfb8e7 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -589,6 +589,92 @@ def _get_layer_offset(config: TransformerConfig): ) return get_transformer_layer_offset(config) + @staticmethod + def _group_offload_output_with_bias( + output_with_bias, offload_manager, forced_released_tensors: Optional[list[Tensor]] = None + ): + """Commit a fine-grained offload group for a raw branch output tuple.""" + if isinstance(output_with_bias, tuple): + output = offload_manager.group_offload( + output_with_bias[0], forced_released_tensors=forced_released_tensors + ) + return (output, *output_with_bias[1:]) + output = offload_manager.group_offload( + output_with_bias, forced_released_tensors=forced_released_tensors + ) + return output + + def _forward_self_attention_output_with_bias( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + rotary_pos_cos: Optional[Tensor] = None, + rotary_pos_sin: Optional[Tensor] = None, + rotary_pos_cos_sin: Optional[Tensor] = None, + attention_bias: Optional[Tensor] = None, + inference_context: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[Tensor] = None, + *, + inference_params: Optional[Any] = None, + ): + """Run input norm + self-attention and return the raw output before BDA.""" + inference_context = deprecate_inference_params(inference_context, inference_params) + + attn_norm_manager = self.off_interface(self.offload_attn_norm, hidden_states, "attn_norm") + if self.recompute_input_layernorm: + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + with attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( + apply_module(self.input_layernorm), hidden_states + ) + else: + with attn_norm_manager as hidden_states: + input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) + + if isinstance(input_layernorm_output, tuple): + if len(input_layernorm_output) != 2: + raise ValueError( + f"When the output of input_layernorm is a tuple, it is " + f"expected to have 2 elements (output, residual), but " + f"got {len(input_layernorm_output)}" + ) + input_layernorm_output, residual = input_layernorm_output + else: + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + + using_fused_tp_inference_kernel = (not self.training) and ( + self.config.inference_fuse_tp_communication + ) + if using_fused_tp_inference_kernel: + self._set_proj_residual(residual) + + nvtx_range_push(suffix="self_attention") + attention_output_with_bias = self.self_attention( + input_layernorm_output, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + ) + nvtx_range_pop(suffix="self_attention") + + if self.recompute_input_layernorm: + self.input_layernorm_checkpoint.discard_output_and_register_recompute( + attention_output_with_bias[0] + ) + + return attention_output_with_bias, attn_norm_manager, residual + def _forward_attention( self, hidden_states: Tensor, @@ -793,31 +879,14 @@ def _forward_pre_mlp_layernorm(self, hidden_states: Tensor): return pre_mlp_layernorm_output - def _forward_mlp( + def _forward_mlp_output_with_bias( self, hidden_states: Tensor, inference_context: BaseInferenceContext | None = None, padding_mask: Tensor | None = None, input_ids: Optional[Tensor] = None, - ) -> Tensor | list[Tensor | None]: - """ - Perform a forward pass through the feed-forward layer. - - Args: - hidden_states (Tensor): Transformed hidden states before the MLP layernorm. - Shape [seq_length, batch_size, hidden_size]. - inference_context: Inference context for optimizations. - padding_mask (Tensor, optional): Padding mask for MoE routing. - Shape [bsz, seq_length]. True = padding (exclude), False = valid (include). - Only used for MoE layers to exclude padding tokens from aux loss computations. - The MoELayer will internally transform this to [seq_length, bsz] format. - input_ids (Tensor, optional): The input IDs tensor. Shape [seq_length, bsz]. - Only used for hash-based MoE routing. Defaults to None. - Returns: - output (Tensor): Transformed hidden states of shape [s, b, h]. - """ - - # Optional Layer norm post the cross-attention. + ) -> tuple[tuple[Tensor, Tensor | None], Tensor]: + """Run pre-MLP norm + MLP/MoE and return the raw output before BDA.""" pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) if isinstance(pre_mlp_layernorm_output, tuple): @@ -913,6 +982,38 @@ def _forward_mlp( ) nvtx_range_pop(suffix="mlp") + return mlp_output_with_bias, residual + + def _forward_mlp( + self, + hidden_states: Tensor, + inference_context: BaseInferenceContext | None = None, + padding_mask: Tensor | None = None, + input_ids: Optional[Tensor] = None, + ) -> Tensor | list[Tensor | None]: + """ + Perform a forward pass through the feed-forward layer. + + Args: + hidden_states (Tensor): Transformed hidden states before the MLP layernorm. + Shape [seq_length, batch_size, hidden_size]. + inference_context: Inference context for optimizations. + padding_mask (Tensor, optional): Padding mask for MoE routing. + Shape [bsz, seq_length]. True = padding (exclude), False = valid (include). + Only used for MoE layers to exclude padding tokens from aux loss computations. + The MoELayer will internally transform this to [seq_length, bsz] format. + input_ids (Tensor, optional): The input IDs tensor. Shape [seq_length, bsz]. + Only used for hash-based MoE routing. Defaults to None. + Returns: + output (Tensor): Transformed hidden states of shape [s, b, h]. + """ + + mlp_output_with_bias, residual = self._forward_mlp_output_with_bias( + hidden_states, + inference_context=inference_context, + padding_mask=padding_mask, + input_ids=input_ids, + ) if ( self.is_moe_layer @@ -1245,6 +1346,58 @@ def _te_cuda_graph_replay(self, *args, **kwargs): if self.config.delay_offload_until_cuda_graph: self.off_interface.exit_replay() + def resume_moe_experts_after_partial_cudagraph(self, cuda_graph_output): + """Resume the eager MoE *expert* compute after a partial (moe_router[/moe_preprocess]) + CUDA graph and return the raw ``mlp_output_with_bias`` (NOT the post-residual layer + output). ``cuda_graph_output`` is the list of captured router/preprocess intermediates, + with any outer (e.g. mHC) state already stripped by the caller. + + Used by the mHC hybrid wrapper (``HyperConnectionHybridLayer``): the wrapper's n-stream + BDA owns the residual combine (mirroring GPT's ``HyperConnectionTransformerLayer``), so + the inner layer's own ``_forward_post_mlp`` residual-add is intentionally skipped here. + EP-overlap is not supported on this path. + """ + assert not self.config.overlap_moe_expert_parallel_comm, ( + "HyperConnectionHybridLayer MoE CUDA-graph capture requires " + "overlap_moe_expert_parallel_comm=False." + ) + shared_expert_output, routing_map = None, None + # The inner residual is the last captured element; the mHC wrapper does not use it + # (the n-stream BDA combines residual), so drop it. + cuda_graph_output.pop() + if ( + self.config.moe_shared_expert_intermediate_size is not None + and not self.config.moe_shared_expert_overlap + ): + shared_expert_output = cuda_graph_output.pop() + + if CudaGraphModule.moe_preprocess in self.config.cuda_graph_modules: + (hidden_states, probs), attr_outputs = cuda_graph_output[:2], cuda_graph_output[2:] + valid_cudagraph_attrs = self.mlp.token_dispatcher.valid_cudagraph_attrs + assert len(attr_outputs) == len( + valid_cudagraph_attrs + ), f"attr_outputs: {len(attr_outputs)} != {len(valid_cudagraph_attrs)}" + for i, attr_name in enumerate(valid_cudagraph_attrs): + self.mlp.token_dispatcher.set_cudagraph_attr(attr_name, attr_outputs[i]) + else: + assert len(cuda_graph_output) == 3, ( + "CUDA graph output should be [hidden_states, probs, routing_map], " + f"but got {len(cuda_graph_output)} elements" + ) + hidden_states, probs, routing_map = cuda_graph_output + + nvtx_range_push(suffix="mlp") + self.mlp.cudagraph_tensor_store.set( + hidden_states=hidden_states, + probs=probs, + routing_map=routing_map, + shared_expert_output=shared_expert_output, + ) + mlp_output_with_bias = self.mlp(hidden_states) + self.mlp.cudagraph_tensor_store.clear() + nvtx_range_pop(suffix="mlp") + return mlp_output_with_bias + def _te_cuda_graph_replay_impl(self, args, kwargs, context): """Implementation of _te_cuda_graph_replay, separated for replay mode cleanup.""" cuda_graph_output = list(super()._te_cuda_graph_replay(*args, **kwargs)) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 29e21544599..ea0fce46c58 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2159,8 +2159,51 @@ def core_transformer_config_from_args(args, config_class=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' + _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' + # When the dsv4_hybrid variant is active (set above or explicitly) and the user did not + # provide --csa-compress-ratios, derive it from the pattern symbols (C->4, H->128, + # W/D/others->0; MTP slots 0) so the length-checked dsv4_hybrid validation passes and the + # per-layer ratios match the symbols. C/H/W layers also take their ratio via the spec. + _variant = kw_args.get('experimental_attention_variant', + getattr(args, 'experimental_attention_variant', None)) + if _variant == 'dsv4_hybrid' and getattr(args, 'csa_compress_ratios', None) is None: + _ratio_map = {Symbols.CSA: 4, Symbols.HCA: 128} + # One ratio entry per ACTUAL layer: main layers, then every MTP layer of every MTP + # depth (a depth can contain multiple hybrid layers, e.g. "/MD-E"). This makes the + # array long enough for the deepseek attn index (num_layers + layer_number - 1) for + # any MTP attention position, not just depth-first. C->4, H->128, others->0. + _sections = _pat.split(Symbols.MTP_SEPARATOR) + _ratios = [_ratio_map.get(c, 0) for c in _sections[0].replace(Symbols.PIPE, '')] + for _mtp_sec in _sections[1:]: + _ratios += [_ratio_map.get(c, 0) for c in _mtp_sec.replace(Symbols.PIPE, '')] + kw_args['csa_compress_ratios'] = _ratios + args.csa_compress_ratios = _ratios + # Exact length check (the pattern is known here, so the precise per-layer count is too): + # one ratio per main layer + one per MTP layer of every MTP depth. This catches a + # mis-sized user-provided --csa-compress-ratios with a clear error. (transformer_config + # keeps a >= backstop because it does not have the pattern to recompute this exactly.) + if _variant == 'dsv4_hybrid' and getattr(args, 'csa_compress_ratios', None) is not None: + _secs = _pat.split(Symbols.MTP_SEPARATOR) + _exact_len = sum(len(s.replace(Symbols.PIPE, '')) for s in _secs) + assert len(args.csa_compress_ratios) == _exact_len, ( + f"csa_compress_ratios length ({len(args.csa_compress_ratios)}) must equal the " + f"number of layers in the hybrid pattern (main + every MTP-depth layer) " + f"= {_exact_len} for pattern '{_pat}'." + ) kw_args['inference_sampling_seed'] = args.seed @@ -4702,7 +4745,8 @@ def _add_experimental_attention_variant_args(parser): '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). ' + 'transformer layer (valid values: 0, 4, 128; 0 = sliding-window-only, the "W" ' + 'hybrid layer symbol). ' 'The list length must equal num_layers.', ) group.add_argument( @@ -4712,6 +4756,10 @@ def _add_experimental_attention_variant_args(parser): 'and fall back to unfused PyTorch implementations.', dest='apply_dsa_kernel_fusion', ) + # 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 diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index 5707b4732ae..86eda9bd3e1 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -418,7 +418,22 @@ def set_jit_fusion_options(): """Set PyTorch JIT layer fusion options.""" # flags required to enable jit fusion kernels if is_torch_min_version("2.2.0a0"): - pass # we're using torch.compile for jit fusion + # we're using torch.compile for jit fusion. + # DSv4 hybrid: MoE-routing-driven fused ops (clamped_weighted_swiglu / + # bias_dropout_add_fused_train) see many distinct per-expert token counts and + # otherwise hit torch._dynamo's default cache_size_limit (8) -> eager fallback. + # Env-gated raise of the recompile cache limit; no-op when unset. + import os + + _dynamo_lim = os.environ.get("DYNAMO_CACHE_SIZE_LIMIT") + if _dynamo_lim: + import torch._dynamo + + _lim = int(_dynamo_lim) + torch._dynamo.config.cache_size_limit = _lim + torch._dynamo.config.accumulated_cache_size_limit = max( + _lim, torch._dynamo.config.accumulated_cache_size_limit + ) elif is_torch_min_version("1.10.0a0"): # nvfuser torch._C._jit_set_profiling_executor(True) diff --git a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py index faa553216da..2618d9cde50 100644 --- a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py +++ b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py @@ -102,6 +102,20 @@ def test_invalid_symbols_cause_failure(self): # 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): + validate_segment_layers("MWM*-") + @pytest.mark.internal class TestGetHybridTotalLayerCount: @@ -305,17 +319,60 @@ def test_dataclass_equality(self): class TestGetHybridLayerCounts: def test_simple_pattern(self): - assert get_hybrid_layer_counts("M*M*") == {'*': 2, 'D': 0, 'G': 0, 'M': 2, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("M*M*") == { + 'C': 0, + 'H': 0, + 'W': 0, + '*': 2, + 'D': 0, + 'G': 0, + 'M': 2, + '-': 0, + 'E': 0, + } def test_all_layer_types(self): # Not allowed to have both standard Attention and MLA/DSA, so we do separate asserts. - assert get_hybrid_layer_counts("MG*-E") == {'*': 1, 'D': 0, 'G': 1, 'M': 1, '-': 1, 'E': 1} - assert get_hybrid_layer_counts("MGD-E") == {'*': 0, 'D': 1, 'G': 1, 'M': 1, '-': 1, 'E': 1} + assert get_hybrid_layer_counts("MG*-E") == { + 'C': 0, + 'H': 0, + 'W': 0, + '*': 1, + 'D': 0, + 'G': 1, + 'M': 1, + '-': 1, + 'E': 1, + } + assert get_hybrid_layer_counts("MGD-E") == { + 'C': 0, + 'H': 0, + 'W': 0, + '*': 0, + 'D': 1, + 'G': 1, + 'M': 1, + '-': 1, + 'E': 1, + } def test_with_pipes(self): # Pipes should be skipped in counting - assert get_hybrid_layer_counts("M*|M*") == {'*': 2, 'D': 0, 'G': 0, 'M': 2, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("M*|M*") == { + 'C': 0, + 'H': 0, + 'W': 0, + '*': 2, + 'D': 0, + 'G': 0, + 'M': 2, + '-': 0, + 'E': 0, + } assert get_hybrid_layer_counts("M-M-|M-M*-") == { + 'C': 0, + 'H': 0, + 'W': 0, '*': 1, 'D': 0, 'G': 0, @@ -327,6 +384,9 @@ 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, @@ -339,6 +399,9 @@ 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, @@ -348,11 +411,24 @@ def test_with_pipes_and_mtp(self): } def test_moe_pattern(self): - assert get_hybrid_layer_counts("MEME") == {'*': 0, 'D': 0, 'G': 0, 'M': 2, '-': 0, 'E': 2} + assert get_hybrid_layer_counts("MEME") == { + 'C': 0, + 'H': 0, + 'W': 0, + '*': 0, + 'D': 0, + 'G': 0, + 'M': 2, + '-': 0, + 'E': 2, + } 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, @@ -362,17 +438,57 @@ def test_mtp_with_attention(self): } def test_gdn_pattern(self): - assert get_hybrid_layer_counts("GMGM") == {'*': 0, 'D': 0, 'G': 2, 'M': 2, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("GMGM") == { + 'C': 0, + 'H': 0, + 'W': 0, + '*': 0, + 'D': 0, + 'G': 2, + 'M': 2, + '-': 0, + 'E': 0, + } def test_gdn_hybrid_pattern(self): # GDN + Mamba + Attention - assert get_hybrid_layer_counts("G*GM*") == {'*': 2, 'D': 0, 'G': 2, 'M': 1, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("G*GM*") == { + 'C': 0, + 'H': 0, + 'W': 0, + '*': 2, + 'D': 0, + 'G': 2, + 'M': 1, + '-': 0, + 'E': 0, + } def test_dsa_pattern(self): - assert get_hybrid_layer_counts("DMDM") == {'*': 0, 'D': 2, 'G': 0, 'M': 2, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("DMDM") == { + 'C': 0, + 'H': 0, + 'W': 0, + '*': 0, + 'D': 2, + 'G': 0, + 'M': 2, + '-': 0, + 'E': 0, + } def test_empty_pattern(self): - assert get_hybrid_layer_counts("") == {'*': 0, 'D': 0, 'G': 0, 'M': 0, '-': 0, 'E': 0} + assert get_hybrid_layer_counts("") == { + 'C': 0, + 'H': 0, + 'W': 0, + '*': 0, + 'D': 0, + 'G': 0, + 'M': 0, + '-': 0, + 'E': 0, + } @pytest.mark.internal @@ -655,7 +771,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) == 6 + assert len(maps) == len(Symbols.VALID_LAYERS) attention_map, mamba_map, mlp_map, moe_map = operator.itemgetter( Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE )(maps) diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index 4d42f47ce62..4c6844f9b53 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -16,7 +16,7 @@ get_gpt_mtp_block_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel -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 validate_segment_layers from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.num_microbatches_calculator import ( @@ -34,10 +34,11 @@ CudaGraphManager, TECudaGraphHelper, _CudagraphGlobalRecord, + _layer_is_graphable, ) from megatron.core.transformer.enums import CudaGraphModule, CudaGraphScope, InferenceCudaGraphScope from megatron.core.transformer.mlp import MLPSubmodules -from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.moe.fused_a2a import reset_hybrid_ep_buffer from megatron.core.transformer.spec_utils import ModuleSpec, get_submodules from megatron.core.transformer.transformer_block import TransformerBlock @@ -805,6 +806,50 @@ def test_gpu_cudagraph(self): del parallel_mamba_block.layers[_].cudagraph_manager.cudagraph_runners[0].fwd_graph + def test_mhc_hybrid_layers_are_te_cudagraph_capturable(self): + """Regression: a mHC-enabled HybridStack must expose graph-capturable layers. + + When ``enable_hyper_connections=True``, ``HybridStack`` wraps every layer in + ``HyperConnectionHybridLayer``. That wrapper must subclass + ``GraphableMegatronModule`` and be recognized by ``_layer_is_graphable`` so TE + cuda-graph discovery finds the wrapped layers. Before the fix the wrapper + subclassed plain ``MegatronModule``, so discovery rejected every layer (0 + graphable) and CUDA graph capture was silently skipped for the whole hybrid + model -- making the mHC hybrid run fully eager (several times slower than the + graphed GPT mHC path). This test fails on the pre-fix code via both assertions. + """ + # The wrapper must be graph-capturable by construction. + assert issubclass(HyperConnectionHybridLayer, GraphableMegatronModule) + + layer_type_list = validate_segment_layers("M-M*-") # mamba / mlp / attention mix + config = TransformerConfig( + hidden_size=256, + num_layers=len(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + cuda_graph_impl="transformer_engine", + enable_hyper_connections=True, + num_residual_streams=4, + cuda_graph_modules=[CudaGraphModule.attn, CudaGraphModule.mamba, CudaGraphModule.mlp], + ) + block = HybridStack( + config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + pg_collection=ProcessGroupCollection.use_mpu_process_groups( + required_pgs=["tp", "pp", "cp"] + ), + ) + + # Every layer is wrapped, and the wrappers are discoverable as graphable. + assert all(isinstance(layer, HyperConnectionHybridLayer) for layer in block.layers) + graphable = [layer for layer in block.layers if _layer_is_graphable(layer, config)] + assert len(graphable) > 0, ( + "mHC HybridStack produced 0 graphable layers -- TE cuda-graph capture would " + "be silently skipped for the entire model (the pre-fix bug)." + ) + # Global storage for comparing unique buffer counts across different num_microbatches, # keyed by (pp_size, vpp_size) diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index 17006d64767..2cb600b19f7 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -223,6 +223,7 @@ def fake_proj_and_transformer_layer( self, hidden_states, decoder_input, + input_ids=None, attention_mask=None, padding_mask=None, context=None,