Deepseek v4 Support (#38)

### PR Category
<!-- One of [ Train | Inference | Compress | Serve | RL | Core |
Hardware | CICD | Tools | Others ] -->
[Train] Most of codes are copied from Megatron-LM Dev branch. The dev
branch is different with main branch or release version.
Megatron LM PR:
DeepSeek-V4:
NVIDIA#4458
NVIDIA#4481
NVIDIA#4518
mHC:
NVIDIA#2943
### PR Types
<!-- One of [ User Experience | New Features | Bug Fixes | Improvements
| Performance | Breaking Change| Deprecations | Test Case | Docs |
Others ] -->
[New features]
### PR Description
<!-- Describe what you’ve done -->
Add DeepSeek V4 model into FlagScale and Megatron-FL
Supported:
1. CSA and HCA
2. Hash Router
3. mHC
4. Engram(optional)

Unsupported:
1. Sqrtsoftpuls router score function. ✅
2. mHC recompute. ✅
3. Overlap_grad_reduce and overlap_param_gather when Zero 1. ✅
4. Any infra optimizations.

### NOTE: This is only a draft pr, please reivew to give more
suggestions.
such as:
1. File structure.
    - All modules are moved into Megatron-FL

### Next plan:
1. Distributed training. ✅
3. Muon optimizer with Zero 1 adaptation. 🚧
4. Low precision is out of scope of this pr, limited by resource. 
5. Maybe context parallel for sparse attention.
6. Welcome to give more suggestions.

---------

Co-authored-by: Hongxiao Bai <hongxiaob@nvidia.com>
Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com>

Author: 3 people

megatron/core/distributed/finalize_model_grads.py

@@ -290,7 +290,11 @@ def reset_model_temporary_tensors(config: TransformerConfig, model: List[torch.n
     """
     for model_chunk in model:
         for module in get_attr_wrapped_model(model_chunk, 'modules')():
-            if config.moe_router_enable_expert_bias and hasattr(module, 'expert_bias'):
+            if (
+                config.moe_router_enable_expert_bias
+                and hasattr(module, 'expert_bias')
+                and module.expert_bias is not None
+            ):
                 module.local_tokens_per_expert.zero_()
             if (
                 config.moe_router_load_balancing_type == "global_aux_loss"
@@ -312,7 +316,11 @@ def _update_router_expert_bias(model: List[torch.nn.Module], config: Transformer
             # cases where only the student is in training mode but the teacher is in eval mode
             # when using online knoweldge-distillation with Model-Optimizer. In this case, we want
             # to avoid updating teacher's expert_bias.
-            if hasattr(module, 'expert_bias') and module.training:
+            if (
+                hasattr(module, 'expert_bias')
+                and module.training
+                and module.expert_bias is not None
+            ):
                 tokens_per_expert_list.append(module.local_tokens_per_expert)
                 expert_bias_list.append(module.expert_bias)
     # For hybrid models with both MoE and Dense layers, this list can be empty.

megatron/core/distributed/param_and_grad_buffer.py

@@ -345,7 +345,7 @@ def start_param_sync(self, force_sync: bool = False):
                 # fp32 when grad_reduce_in_fp32 is enabled).
                 param_dtype = bucket.params_list[0].dtype
 
-                if max(bucket.layerwise_param_flat_sizes) == 0:
+                if bucket.layerwise_param_flat_sizes is None or max(bucket.layerwise_param_flat_sizes) == 0:    ##### FalgScale add #####
                     # All ranks have empty params for this bucket — skip.
                     bucket.layerwise_gather_list = None
                     continue

megatron/core/extensions/transformer_engine.py

@@ -80,6 +80,7 @@
         HAVE_TE = False
 
 _TE_CONFIG_TYPE_KEY = "transformer_engine_config_type"
+from megatron.plugin.hetero.parallel_context import get_parallel_context    ##### FlagScale add #####
 
 
 class TransformerEngineConfigType(enum.Enum):
@@ -1033,7 +1034,7 @@ def __init__(
             sequence_parallel=self.config.sequence_parallel,
             fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
             tp_group=tp_group if torch.distributed.is_initialized() else None,
-            tp_size=get_tensor_model_parallel_world_size(),  # FlagScale Add
+            tp_size=self.config.tensor_model_parallel_size if get_parallel_context() is None else get_tensor_model_parallel_world_size(), ##### FlagScale Add #####
             get_rng_state_tracker=(
                 get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
             ),
@@ -1543,7 +1544,7 @@ def __init__(
             ),
             attn_mask_type=attn_mask_type.name,
             sequence_parallel=self.config.sequence_parallel,
-            tp_size=get_tensor_model_parallel_world_size(),  # FlagScale Add
+            tp_size=self.config.tensor_model_parallel_size if get_parallel_context() is None else get_tensor_model_parallel_world_size(), ##### FlagScale Add #####
             get_rng_state_tracker=(
                 get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
             ),

megatron/core/fusions/fused_bias_dropout.py

@@ -1,10 +1,13 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
-from typing import Optional, Tuple
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+from typing import TYPE_CHECKING, Optional, Tuple
 
 import torch
 
 from megatron.core.jit import jit_fuser
 
+if TYPE_CHECKING:
+    from megatron.core.tensor_parallel.random import CheckpointManager
+
 # pylint: disable=missing-function-docstring
 
 
@@ -80,7 +83,26 @@ def bias_dropout_add_fused_inference(
     return _bias_dropout_add_func(x_with_bias, residual, prob, False)
 
 
-def get_bias_dropout_add(training, fused):
+def get_bias_dropout_add(
+    training, fused, mhc_recompute_manager: Optional['CheckpointManager'] = None
+):
+    """
+    Get the bias-dropout-add function.
+
+    Args:
+        training: Whether in training mode.
+        fused: Whether to use fused implementation.
+        mhc_recompute_manager: Optional CheckpointManager for checkpoint management.
+            When provided, the returned function will wrap the BDA operation with
+            CheckpointWithoutOutput for memory-efficient recomputation.
+
+    Returns:
+        A callable that performs bias-dropout-add operation.
+    """
+    if mhc_recompute_manager is not None:
+        # Return a checkpointed version that handles tuple unpacking internally
+        return _get_checkpointed_bda(training, fused, mhc_recompute_manager)
+
     if fused:
         # jit scripting for a nn.module (with dropout) is not
         # triggering the fusion kernel. For now, we use two
@@ -92,3 +114,68 @@ def get_bias_dropout_add(training, fused):
             return bias_dropout_add_fused_inference
     else:
         return bias_dropout_add_unfused(training)
+
+
+def _get_checkpointed_bda(training, fused, mhc_recompute_manager: 'CheckpointManager'):
+    """
+    Create a checkpointed bias-dropout-add function.
+
+    This function handles:
+    1. Tuple unpacking for x_with_bias (required because save_for_backward can't save tuples)
+    2. Non-tensor arguments like dropout probability (handled by CheckpointWithoutOutput)
+    3. Auto-registration to the CheckpointManager
+
+    Args:
+        training: Whether in training mode.
+        fused: Whether to use fused implementation.
+        mhc_recompute_manager: CheckpointManager for checkpoint management.
+
+    Returns:
+        A callable that performs checkpointed bias-dropout-add operation.
+    """
+    from megatron.core.tensor_parallel.random import CheckpointWithoutOutput
+
+    # Get the underlying BDA function
+    if fused:
+        if training:
+            bda_func = bias_dropout_add_fused_train
+        else:
+            bda_func = bias_dropout_add_fused_inference
+    else:
+        bda_func = bias_dropout_add_unfused(training)
+
+    def _checkpointed_bda(x_with_bias, residual, prob):
+        """
+        Checkpointed BDA that handles tuple unpacking internally.
+
+        Args:
+            x_with_bias: Either a tuple (x, bias) or a single tensor x.
+            residual: Residual tensor.
+            prob: Dropout probability.
+
+        Returns:
+            Output tensor after bias-dropout-add.
+        """
+        # Create checkpoint with manager
+        ckpt = CheckpointWithoutOutput(ckpt_manager=mhc_recompute_manager)
+
+        # Handle case where x_with_bias might be a single tensor (e.g., from IdentityOp)
+        if isinstance(x_with_bias, tuple):
+            x, bias = x_with_bias
+        else:
+            x = x_with_bias
+            bias = None
+
+        # Wrapper function that re-packs the tuple for the actual BDA function
+        def _bda_wrapper(output, bias, res, dropout):
+            return bda_func((output, bias), res, dropout)
+
+        # Call checkpoint with unpacked arguments
+        result = ckpt.checkpoint(_bda_wrapper, x, bias, residual, prob)
+
+        # No-op when manager is set - manager handles all discarding uniformly
+        ckpt.discard_output_and_register_recompute(result)
+
+        return result
+
+    return _checkpointed_bda

megatron/core/fusions/fused_bias_swiglu.py

@@ -48,6 +48,23 @@ def weighted_swiglu(y, weights):
     return res.to(dtype)
 
 
+@jit_fuser
+def clamped_swiglu(y, clamp_value):
+    dtype = y.dtype
+    y_1, y_2 = torch.chunk(y.to(torch.float32), 2, -1)
+    y_1 = y_1.clamp(min=None, max=clamp_value)
+    y_2 = y_2.clamp(min=-clamp_value, max=clamp_value)
+    res = F.silu(y_1) * y_2
+    return res.to(dtype)
+
+
+@jit_fuser
+def clamped_weighted_swiglu(y, weights, clamp_value):
+    dtype = y.dtype
+    res = clamped_swiglu(y, clamp_value) * weights
+    return res.to(dtype)
+
+
 # gradient of tanh approximation of gelu
 # gradient of actual gelu is:
 # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
@@ -97,6 +114,36 @@ def weighted_swiglu_back(g, y, weights):
     return input_grad.to(input_dtype), weights_grad.to(w_dtype)
 
 
+@jit_fuser
+def clamped_swiglu_back(g, y, clamp_value):
+    dtype = y.dtype
+    y_1, y_2 = torch.chunk(y.to(torch.float32), 2, -1)
+    y_1c = y_1.clamp(min=None, max=clamp_value)
+    y_2c = y_2.clamp(min=-clamp_value, max=clamp_value)
+    res = torch.cat(
+        (
+            g
+            * torch.sigmoid(y_1c)
+            * (1 + y_1c * (1 - torch.sigmoid(y_1c)))
+            * y_2c
+            * (y_1 <= clamp_value).to(g.dtype),
+            g * F.silu(y_1c) * ((y_2 >= -clamp_value) & (y_2 <= clamp_value)).to(g.dtype),
+        ),
+        -1,
+    )
+    return res.to(dtype)
+
+
+@jit_fuser
+def clamped_weighted_swiglu_back(g, y, weights, clamp_value):
+    input_dtype = y.dtype
+    w_dtype = weights.dtype
+    input_grad = clamped_swiglu_back(g * weights, y, clamp_value)
+    weights_grad = clamped_swiglu(y, clamp_value) * g.to(w_dtype)
+    weights_grad = torch.sum(weights_grad, dim=-1, keepdim=True)
+    return input_grad.to(input_dtype), weights_grad.to(w_dtype)
+
+
 class BiasSwiGLUFunction(torch.autograd.Function):
     """Custom autograd function for SwiGLU activation with bias support."""
 
@@ -190,20 +237,27 @@ def backward(ctx, grad_output):
 
 class WeightedSwiGLUFunction(torch.autograd.Function):
     @staticmethod
-    # bias is an optional argument
-    def forward(ctx, input, weights, fp8_input_store):
+    def forward(ctx, input, weights, fp8_input_store, clamp_value):
         input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
         ctx.save_for_backward(input_for_backward, weights)
         ctx.ori_input_dtype = input.dtype
         ctx.fp8_input_store = fp8_input_store
-        return weighted_swiglu(input, weights)
+        ctx.clamp_value = clamp_value
+        if clamp_value is not None and clamp_value > 0:
+            res = clamped_weighted_swiglu(input, weights, clamp_value)
+        else:
+            res = weighted_swiglu(input, weights)
+        return res
 
     @staticmethod
     def backward(ctx, grad_output):
         input, weights = ctx.saved_tensors
         input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
-        tmp, wgrad = weighted_swiglu_back(grad_output, input, weights)
-        return tmp, wgrad, None
+        if ctx.clamp_value is not None and ctx.clamp_value > 0:
+            tmp, wgrad = clamped_weighted_swiglu_back(grad_output, input, weights, ctx.clamp_value)
+        else:
+            tmp, wgrad = weighted_swiglu_back(grad_output, input, weights)
+        return tmp, wgrad, None, None
 
 
 def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False):
@@ -236,7 +290,7 @@ def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False
     return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1)
 
 
-def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False):
+def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False, clamp_value=None):
     """
     Token-wise-weighted bias swiglu fusion.
     """
@@ -246,7 +300,7 @@ def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False):
     if bias is not None:
         raise NotImplementedError("Bias is not supported for weighted swiglu fusion")
     else:
-        output = WeightedSwiGLUFunction.apply(input, weights, fp8_input_store)
+        output = WeightedSwiGLUFunction.apply(input, weights, fp8_input_store, clamp_value)
 
     return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1)