add R3 for dsv4

DeepSeek-V4-Flash/install.sh

@@ -64,4 +64,10 @@ git apply --whitespace=nowarn ../verl-ascend-recipe/DeepSeek-V4-Flash/patch/vllm
 cd verl
 git apply --whitespace=nowarn ../verl-ascend-recipe/DeepSeek-V4-Flash/patch/verl.patch && cd ..
 
+cd vllm
+git apply --whitespace=nowarn ../verl-ascend-recipe/DeepSeek-V4-Flash/patch/vllm.patch && cd ..
+
+cd MindSpeed-LLM
+git apply --whitespace=nowarn ../verl-ascend-recipe/DeepSeek-V4-Flash/patch/mindspeed-llm.patch && cd ..
+
 

DeepSeek-V4-Flash/patch/megatron.patch

@@ -30,8 +30,155 @@ index 26a96f457..f8cbc5fd4 100644
      if recv_next_shape_tensor is not None:
          recv_next_shape = recv_next_shape_tensor.tolist()
 
+diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
+index dafedfd17..8cc172886 100644
+--- a/megatron/core/transformer/moe/moe_utils.py
++++ b/megatron/core/transformer/moe/moe_utils.py
+@@ -1,7 +1,7 @@
+ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+ import math
+-from typing import List, Optional, Union
++from typing import List, Optional, Union, Tuple
+
+ import torch
+
+@@ -15,10 +15,12 @@ try:
+         fused_sort_chunks_by_index,
+         fused_sort_chunks_by_index_with_probs,
+         fused_unpermute,
++        fused_compute_score_for_moe_aux_loss,
+     )
+
+     HAVE_TE = True
+ except ImportError:
++    fused_compute_score_for_moe_aux_loss = None
+     HAVE_TE = False
+
+
+@@ -759,3 +761,119 @@ def maybe_move_tensor_to_cpu(tensor, as_numpy=False, record_stream=False):
+             tensor.record_stream(torch.cuda.current_stream())
+         tensor = cpu_tensor
+     return tensor
++
++
++def apply_router_token_dropping(
++    routing_probs: torch.Tensor,
++    routing_map: torch.Tensor,
++    router_topk: int,
++    capacity_factor: float,
++    drop_policy: str = "probs",
++    pad_to_capacity: bool = False,
++) -> Tuple[torch.Tensor, torch.Tensor]:
++    """Apply token dropping to top-k expert selection.
++
++    This function enforces expert capacity limits by dropping tokens that exceed
++    the capacity and optionally padding to capacity.
++
++    Args:
++        routing_probs (torch.Tensor): Tensor of shape [num_tokens, num_experts]
++            containing the routing probabilities for selected experts.
++        routing_map (torch.Tensor): Boolean tensor of shape [num_tokens, num_experts]
++            indicating which experts were selected for each token.
++        router_topk (int): Number of experts selected per token.
++        capacity_factor (float): The capacity factor of each expert.
++        drop_policy (str, optional): Policy to drop tokens - "probs" or "position".
++                                     Defaults to "probs".
++        pad_to_capacity (bool, optional): Whether to pad to capacity. Defaults to False.
++
++    Returns:
++        Tuple[torch.Tensor, torch.Tensor]:
++            - final_probs: Routing probabilities after applying capacity constraints
++            - final_map: Boolean mask after applying capacity constraints
++    """
++    assert routing_probs.ndim == 2 and routing_map.ndim == 2
++    num_tokens, num_experts = routing_probs.shape
++    # Calculate expert capacity
++    expert_capacity = get_capacity(
++        num_tokens=num_tokens * router_topk,
++        num_experts=num_experts,
++        capacity_factor=capacity_factor,
++    )
++
++    # Create capacity mask based on drop policy
++    if expert_capacity > num_tokens:
++        # No need to drop tokens if capacity exceeds the number of tokens
++        capacity_mask = torch.ones_like(routing_probs).bool()
++    else:
++        if drop_policy == "probs":
++            _, capacity_indices = torch.topk(routing_probs, k=expert_capacity, dim=0, sorted=False)
++            capacity_mask = torch.zeros_like(routing_probs).scatter(0, capacity_indices, 1).bool()
++        elif drop_policy == "position":
++            _, capacity_indices = torch.topk(
++                routing_map.int(), k=expert_capacity, dim=0, sorted=False
++            )
++            capacity_mask = torch.zeros_like(routing_probs).scatter(0, capacity_indices, 1).bool()
++        else:
++            raise ValueError(f"Invalid drop_policy: {drop_policy}")
++
++    # Apply capacity constraints
++    if pad_to_capacity:
++        final_map = capacity_mask
++        final_probs = routing_probs * final_map
++    else:
++        # Get exceed mask and maskout exceeded probs and indices
++        final_map = torch.logical_and(routing_map, capacity_mask)
++        final_probs = routing_probs * final_map
++
++    return final_probs, final_map
++
++
++def compute_routing_scores_for_aux_loss(
++    logits: torch.Tensor,
++    topk: int,
++    score_function: str,
++    fused: bool = False,
++    padding_mask: Optional[torch.Tensor] = None,
++) -> Tuple[torch.Tensor, torch.Tensor]:
++    """Compute routing scores based on the score function.
++
++    Args:
++        logits (torch.Tensor): The logits tensor after gating, shape: [num_tokens, num_experts].
++        topk (int): The number of top-k indices to compute.
++        score_function (str): The score function to use. Can be either "softmax" or "sigmoid".
++        fused (bool, optional): Whether to use the fused version. Defaults to False.
++        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
++                                               Shape in [num_tokens]. True for valid tokens,
++                                               False for padding tokens. Defaults to None.
++
++    Returns:
++        Tuple[torch.Tensor, torch.Tensor]: The routing map and the normalized routing scores.
++    """
++    if fused:
++        if not HAVE_TE or fused_compute_score_for_moe_aux_loss is None:
++            raise ValueError(
++                "fused_compute_score_for_moe_aux_loss is not available. Please install TE >= 2.6.0."
++            )
++        routing_map, scores = fused_compute_score_for_moe_aux_loss(
++            logits=logits, topk=topk, score_function=score_function
++        )
++    else:
++        if score_function == "softmax":
++            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
++        elif score_function == "sigmoid":
++            scores = torch.sigmoid(logits)
++            scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
++        else:
++            raise ValueError(f"Invalid score_function: {score_function}")
++
++        _, top_indices = torch.topk(scores, k=topk, dim=1)
++        routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
++
++    # Apply padding mask to scores if provided
++    if padding_mask is not None:
++        # Invert padding_mask and make True indicates valid tokens
++        valid_mask = (~padding_mask).unsqueeze(-1)
++        routing_map = routing_map * valid_mask
++        scores = scores * valid_mask
++    return routing_map, scores
+\ No newline at end of file
 diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
-index 2b68a5a5f..e6cf2469b 100644
+index 2b68a5a5f..20d624610 100644
 --- a/megatron/core/transformer/transformer_config.py
 +++ b/megatron/core/transformer/transformer_config.py
 @@ -568,8 +568,8 @@ class TransformerConfig(ModelParallelConfig):
@@ -45,3 +192,11 @@ index 2b68a5a5f..e6cf2469b 100644
 
          if self.num_query_groups % self.tensor_model_parallel_size != 0:
              raise ValueError(
+@@ -1075,6 +1075,7 @@ class MLATransformerConfig(TransformerConfig):
+     The initialization function has an argument for each parameter, including those in
+     ModelParallelConfig. Included YaRN RoPE parameters that is fused in MLA.
+     """
++    enable_routing_replay: bool = False
+
+     multi_latent_attention: bool = True
+     """Whether to use Multi-Latent Attention."""

DeepSeek-V4-Flash/patch/mindspeed-llm.patch

@@ -0,0 +1,95 @@
+diff --git a/mindspeed_llm/core/transformer/moe/moe_utils.py b/mindspeed_llm/core/transformer/moe/moe_utils.py
+index d20e40a0..90816b98 100644
+--- a/mindspeed_llm/core/transformer/moe/moe_utils.py
++++ b/mindspeed_llm/core/transformer/moe/moe_utils.py
+@@ -18,6 +18,7 @@ from functools import wraps
+ from typing import Optional
+
+ import torch
++from enum import Enum
+ import torch.nn.functional as F
+ from megatron.core import parallel_state
+ from megatron.core.transformer.moe.moe_utils import get_capacity
+@@ -108,6 +109,7 @@ def topk_softmax_with_capacity_and_hash(
+     token_hash: bool = False,
+     tid2eid: Optional[torch.Tensor] = None,
+     input_ids: Optional[torch.Tensor] = None,
++    router_replay = None,
+ ):
+     """
+     patch hash operator in megatron topk_softmax_with_capacity
+@@ -115,7 +117,7 @@ def topk_softmax_with_capacity_and_hash(
+     assert logits.dim() == 2, f"Expected 2D logits [num_tokens, num_experts], got {logits.dim()}."
+     num_tokens, num_experts = logits.shape
+
+-    def compute_topk(scores, topk, num_groups=None, group_topk=None):
++    def _compute_topk(scores, topk, num_groups=None, group_topk=None):
+         if group_topk:
+             return group_limited_topk(
+                 scores=scores,
+@@ -128,6 +130,41 @@ def topk_softmax_with_capacity_and_hash(
+         else:
+             return torch.topk(scores, k=topk, dim=1)
+
++    def compute_topk(scores, topk, num_groups=None, group_topk=None):
++        from verl.utils.megatron.router_replay_patch import RouterReplayAction
++        # Default behavior if no replay is active
++        routing_action = router_replay.router_replay_action if router_replay is not None else None
++
++        if routing_action is None:
++            return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)
++
++        if routing_action == RouterReplayAction.RECORD:
++            probs, top_indices = _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)
++            if router_replay is not None:
++                router_replay.record_indices(top_indices)
++            return probs, top_indices
++        elif routing_action == RouterReplayAction.REPLAY_FORWARD:
++            if router_replay is None or router_replay.target_topk_idx is None:
++                # Fallback if replay data is not available
++                return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)
++            # Use the provided indices for replay
++            top_indices = router_replay.target_topk_idx
++            top_indices = top_indices.to(scores.device)
++            probs = scores.gather(1, top_indices)
++            return probs, top_indices
++        elif routing_action == RouterReplayAction.REPLAY_BACKWARD:
++            if router_replay is None or not router_replay.replay_backward_list:
++                # Fallback if replay data is not available
++                return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)
++            # Use the last recorded indices for backward replay
++            top_indices = router_replay.replay_backward_list.pop(0)
++            # Ensure indices are on the correct device
++            top_indices = top_indices.to(scores.device)
++            probs = scores.gather(1, top_indices)
++            return probs, top_indices
++        else:  # Unknown action, fallback
++            return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)
++
+     if score_function == "softmax":
+         if use_pre_softmax:
+             scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)
+diff --git a/mindspeed_llm/core/transformer/moe/router.py b/mindspeed_llm/core/transformer/moe/router.py
+index 5555ca2b..e0948d3c 100644
+--- a/mindspeed_llm/core/transformer/moe/router.py
++++ b/mindspeed_llm/core/transformer/moe/router.py
+@@ -599,6 +599,7 @@ def topk_router_routing(self, logits: torch.Tensor, input_ids: torch.Tensor = No
+                 token_hash=self.hash if hasattr(self, "hash") else None,
+                 tid2eid=self.tid2eid if hasattr(self, "tid2eid") else None,
+                 input_ids=input_ids,
++                router_replay=getattr(self, "router_replay", None)
+             )
+         else:
+             # A naive top-k routing without load balancing
+diff --git a/pretrain_deepseek4.py b/pretrain_deepseek4.py
+index 8e7ca23a..64d88cb4 100644
+--- a/pretrain_deepseek4.py
++++ b/pretrain_deepseek4.py
+@@ -54,6 +54,7 @@ def model_provider(
+         Union[DeepSeek4Model, megatron.legacy.model.DeepSeek4Model]: The returned model
+     """
+     args = get_args()
++    args.enable_routing_replay = model_provider.enable_routing_replay
+     use_te = args.transformer_impl == "transformer_engine"
+
+     print_rank_0('building GPT model ...')