[dev] [DeepSeek-v4] Part 3: MTP support with mHC and new mHC contract

gpt_builders.py

@@ -63,6 +63,13 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
             # Get the decoder layer spec explicitly if no decoder layer in the last stage,
             # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
             transformer_layer_spec_for_mtp = _get_transformer_layer_spec(use_te, config)
+        elif args.experimental_attention_variant is not None:
+            # get_gpt_decoder_layer_specs rejects experimental variants;
+            # build per-layer specs via the experimental entry point.
+            experimental_layer_specs = (
+                get_transformer_layer_with_experimental_attention_variant_spec(config=config)
+            )
+            transformer_layer_spec_for_mtp = experimental_layer_specs[-1]
         else:
             # Define the decoder block spec
             decoder_layer_specs = get_gpt_decoder_layer_specs(

megatron/core/fusions/fused_bias_swiglu.py

@@ -58,6 +58,12 @@ def clamped_swiglu(y, clamp_value):
     return res.to(dtype)
 
 
+@jit_fuser
+def bias_clamped_swiglu(y, bias, clamp_value):
+    """SwiGLU with clamping after bias addition."""
+    return clamped_swiglu(y + bias, clamp_value)
+
+
 @jit_fuser
 def clamped_weighted_swiglu(y, weights, clamp_value):
     dtype = y.dtype
@@ -134,6 +140,12 @@ def clamped_swiglu_back(g, y, clamp_value):
     return res.to(dtype)
 
 
+@jit_fuser
+def bias_clamped_swiglu_back(g, y, bias, clamp_value):
+    """Backward of SwiGLU with clamping after bias addition."""
+    return clamped_swiglu_back(g, y + bias, clamp_value)
+
+
 @jit_fuser
 def clamped_weighted_swiglu_back(g, y, weights, clamp_value):
     input_dtype = y.dtype
@@ -149,14 +161,18 @@ class BiasSwiGLUFunction(torch.autograd.Function):
 
     @staticmethod
     @nvtx_decorator()
-    def forward(ctx, input, bias, fp8_input_store, cpu_offload_input):
+    def forward(ctx, input, bias, fp8_input_store, cpu_offload_input, clamp_value):
         """Forward pass of biased SwiGLU activation.
 
         Args:
             ctx: Autograd context object for saving tensors for backward pass.
             input (torch.Tensor): Input tensor to apply SwiGLU to.
             bias (torch.Tensor): Bias tensor to be added to input before SwiGLU.
             fp8_input_store (bool): If True, stores intermediate values in FP8 format.
+            cpu_offload_input (bool): If True, mark saved tensors for activation offloading.
+            clamp_value (float | None): If set and positive, clamp the gate input to
+                ``[-inf, clamp_value]`` and the linear input to ``[-clamp_value, clamp_value]``
+                before applying SwiGLU.
 
         Returns:
             torch.Tensor: Result of applying bias addition followed by SwiGLU activation.
@@ -168,6 +184,9 @@ def forward(ctx, input, bias, fp8_input_store, cpu_offload_input):
         ctx.save_for_backward(input_for_backward, bias)
         ctx.ori_input_dtype = input.dtype
         ctx.fp8_input_store = fp8_input_store
+        ctx.clamp_value = clamp_value
+        if clamp_value is not None and clamp_value > 0:
+            return bias_clamped_swiglu(input, bias, clamp_value)
         return bias_swiglu(input, bias)
 
     @staticmethod
@@ -184,25 +203,34 @@ def backward(ctx, grad_output):
                 - Gradient with respect to the input tensor
                 - Gradient with respect to the bias tensor
                 - None for fp8_input_store parameter
+                - None for cpu_offload_input parameter
+                - None for clamp_value parameter
         """
         input, bias = ctx.saved_tensors
         input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
-        tmp = bias_swiglu_back(grad_output, input, bias)
-        return tmp, tmp, None, None
+        if ctx.clamp_value is not None and ctx.clamp_value > 0:
+            tmp = bias_clamped_swiglu_back(grad_output, input, bias, ctx.clamp_value)
+        else:
+            tmp = bias_swiglu_back(grad_output, input, bias)
+        return tmp, tmp, None, None, None
 
 
 class SwiGLUFunction(torch.autograd.Function):
     """Custom autograd function for SwiGLU activation without bias."""
 
     @staticmethod
     @nvtx_decorator()
-    def forward(ctx, input, fp8_input_store, cpu_offload_input):
+    def forward(ctx, input, fp8_input_store, cpu_offload_input, clamp_value):
         """Forward pass of SwiGLU activation.
 
         Args:
             ctx: Autograd context object for saving tensors for backward pass.
             input (torch.Tensor): Input tensor to apply SwiGLU to.
             fp8_input_store (bool): If True, stores intermediate values in FP8 format.
+            cpu_offload_input (bool): If True, mark saved tensors for activation offloading.
+            clamp_value (float | None): If set and positive, clamp the gate input to
+                ``[-inf, clamp_value]`` and the linear input to ``[-clamp_value, clamp_value]``
+                before applying SwiGLU.
 
         Returns:
             torch.Tensor: Result of applying SwiGLU activation.
@@ -213,6 +241,9 @@ def forward(ctx, input, fp8_input_store, cpu_offload_input):
         ctx.save_for_backward(input_for_backward)
         ctx.ori_input_dtype = input.dtype
         ctx.fp8_input_store = fp8_input_store
+        ctx.clamp_value = clamp_value
+        if clamp_value is not None and clamp_value > 0:
+            return clamped_swiglu(input, clamp_value)
         return swiglu(input)
 
     @staticmethod
@@ -228,11 +259,16 @@ def backward(ctx, grad_output):
             tuple: Tuple containing:
                 - Gradient with respect to the input tensor
                 - None for fp8_input_store parameter
+                - None for cpu_offload_input parameter
+                - None for clamp_value parameter
         """
         input = ctx.saved_tensors[0]
         input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
-        tmp = swiglu_back(grad_output, input)
-        return tmp, None, None
+        if ctx.clamp_value is not None and ctx.clamp_value > 0:
+            tmp = clamped_swiglu_back(grad_output, input, ctx.clamp_value)
+        else:
+            tmp = swiglu_back(grad_output, input)
+        return tmp, None, None, None
 
 
 class WeightedSwiGLUFunction(torch.autograd.Function):
@@ -260,7 +296,7 @@ def backward(ctx, grad_output):
         return tmp, wgrad, None, None
 
 
-def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False):
+def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False, clamp_value=None):
     """Implementation of biased SwiGLU that handles different input shapes.
 
     This function reshapes the input if necessary, applies the SwiGLU activation
@@ -272,6 +308,11 @@ def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False
             uses the bias-free SwiGLU variant.
         fp8_input_store (bool, optional): Whether to store intermediate values in FP8 format.
             Defaults to False.
+        cpu_offload_input (bool, optional): If True, mark saved tensors for activation
+            offloading. Defaults to False.
+        clamp_value (float | None, optional): If set and positive, clamp the gate input to
+            ``[-inf, clamp_value]`` and the linear input to ``[-clamp_value, clamp_value]``
+            before applying SwiGLU. Defaults to None (no clamping).
 
     Returns:
         torch.Tensor: Result of biased SwiGLU activation.
@@ -283,9 +324,11 @@ def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False
     assert len(ori_shape) in [2, 3]
     input = input.view(-1, ori_shape[-1])
     if bias is not None:
-        output = BiasSwiGLUFunction.apply(input, bias, fp8_input_store, cpu_offload_input)
+        output = BiasSwiGLUFunction.apply(
+            input, bias, fp8_input_store, cpu_offload_input, clamp_value
+        )
     else:
-        output = SwiGLUFunction.apply(input, fp8_input_store, cpu_offload_input)
+        output = SwiGLUFunction.apply(input, fp8_input_store, cpu_offload_input, clamp_value)
 
     return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1)
 

megatron/core/models/gpt/gpt_layer_specs.py

@@ -786,16 +786,10 @@ def get_gpt_mtp_block_spec_for_backend(
 
     transformer_layer_spec.submodules = copy.copy(transformer_layer_spec.submodules)
 
-    # MTP does not support hyper connections yet; strip HC modules and
-    # downgrade the layer class to plain TransformerLayer.
-    transformer_layer_spec.submodules.self_attention_hyper_connection = IdentityOp
-    transformer_layer_spec.submodules.cross_attention_hyper_connection = IdentityOp
-    transformer_layer_spec.submodules.mlp_hyper_connection = IdentityOp
-    if transformer_layer_spec.module is HyperConnectionTransformerLayer:
-        transformer_layer_spec.module = TransformerLayer
-
     mtp_layer_spec = get_mtp_layer_spec_for_backend(
-        mtp_model_layer_spec=transformer_layer_spec, backend=backend
+        mtp_model_layer_spec=transformer_layer_spec,
+        backend=backend,
+        enable_hyper_connections=config.enable_hyper_connections,
     )
     mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
     if config.mtp_use_repeated_layer:

megatron/core/models/gpt/gpt_model.py

@@ -550,7 +550,7 @@ def forward(
             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,
@@ -563,6 +563,13 @@ def forward(
             padding_mask=padding_mask,
             **decoder_extra_block_kwargs,
         )
+        # When mHC + MTP, the decoder returns (contracted, multi-stream).
+        # MTP needs multi-stream; lm_head needs contracted.
+        if isinstance(decoder_output, tuple):
+            hidden_states, mhc_multistream = decoder_output
+        else:
+            hidden_states = decoder_output
+            mhc_multistream = None
 
         return self._postprocess(
             hidden_states=hidden_states,
@@ -582,6 +589,7 @@ def forward(
             runtime_gather_output=runtime_gather_output,
             extra_block_kwargs=extra_block_kwargs,
             inference_context=inference_context,
+            mhc_multistream=mhc_multistream,
         )
 
     def _postprocess(
@@ -603,6 +611,7 @@ def _postprocess(
         runtime_gather_output=None,
         extra_block_kwargs=None,
         inference_context=None,
+        mhc_multistream=None,
     ):
         """Postprocesses decoder hidden states to generate logits or compute loss.
 
@@ -631,6 +640,7 @@ def _postprocess(
                 input_ids=input_ids,
                 position_ids=position_ids,
                 hidden_states=hidden_states,
+                mhc_multistream=mhc_multistream,
                 attention_mask=attention_mask,
                 inference_params=None,  # MTP layers don't use KV cache
                 rotary_pos_emb=rotary_pos_emb,

megatron/core/transformer/experimental_attention_variant/csa.py

@@ -120,6 +120,10 @@ def _apply_rope(
             ), "Fused MLA RoPE apply is not imported successfully"
         else:
             rotary_pos_emb, mscale = rotary_pos_emb_module(total_seq_len, packed_seq=False)
+            # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's
+            # concentration factor (mscale) is NOT part of the DSv4 model contract --
+            # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0.
+            mscale = 1.0
     if rotary_pos_emb is not None and ratio > 1:
         rotary_pos_emb = rotary_pos_emb[:total_seq_len:ratio][:rotary_seq_len]
     if rotary_pos_cos is not None and ratio > 1:
@@ -571,6 +575,7 @@ def __init__(
         pg_collection: Optional[ProcessGroupCollection] = None,
         rotary_pos_emb: nn.Module = None,
         compress_ratio: int = 0,
+        is_mtp_layer: bool = False,
     ):
         super().__init__(config=config)
 
@@ -579,6 +584,8 @@ def __init__(
         self.pg_collection = pg_collection
 
         self.layer_number = layer_number
+        if is_mtp_layer:
+            self.layer_number = self.layer_number + self.config.num_layers
         self.compress_ratio = compress_ratio
         self.window_size = config.csa_window_size
         self.v_head_dim = config.v_head_dim
@@ -728,7 +735,8 @@ def forward(
                             DSAIndexerLossLoggingHelper.save_loss_to_tracker(
                                 loss=indexer_loss,
                                 layer_number=self.layer_number,
-                                num_layers=self.config.num_layers,
+                                num_layers=self.config.num_layers
+                                + (self.config.mtp_num_layers or 0),
                             )
                     else:
                         _, topk_indices_compressed = self.indexer(

megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py

@@ -145,6 +145,7 @@ def __init__(
         core_attn_extra_kwargs = {
             "rotary_pos_emb": self.rotary_pos_emb,
             "compress_ratio": compress_ratio,
+            "is_mtp_layer": is_mtp_layer,
         }
         self.core_attention = build_module(
             submodules.core_attention,
@@ -330,6 +331,10 @@ def forward(
                 ), "Fused MLA RoPE apply is not imported successfully"
             else:
                 rotary_pos_emb, mscale = self.rotary_pos_emb(rope_seqlen, packed_seq=packed_seq)
+                # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's
+                # concentration factor (mscale) is NOT part of the DSv4 model contract --
+                # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0.
+                mscale = 1.0
         if self.config.apply_rope_fusion:
             core_attn_out = fused_mla_rope_inplace(
                 core_attn_out,
@@ -527,6 +532,10 @@ def get_query_key_value_tensors(
                 ), "Fused MLA RoPE apply is not imported successfully"
             else:
                 rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq)
+                # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's
+                # concentration factor (mscale) is NOT part of the DSv4 model contract --
+                # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0.
+                mscale = 1.0
 
         if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
             if packed_seq_params.cu_seqlens_q_padded is not None: