分享一个DeepSeek V3和R1中 Shared Experts和普通Experts融合的技巧

0x0. 前言

上周六的时候发现 @DiegoD94 在 vLLM 中尝试对 DeepSeek V3/R1 应用一个 fuse shared experts 到普通 256 个 expert 中的工作 (https://github.com/vllm-proje...)。还有一个技术文档：https://docs.google.com/docum... ，读了一下感觉比较有意义并且看起来对整体的吞吐和 TTFT/ITL 都有比较好的收益。所以利用周末时间在 SGLang 中实现了一下这个工作，由于我们之前在 SGLang 的 sgl moe_align_kernel 中已经支持了 num_experts>256 的情况，所以本次实现起来比较方便，不需要修改 sgl-kernel 里面的 cuda 代码。在 SGLang 中的细节如下：https://github.com/sgl-projec... 。下面来讲一下这个 fuse shared expert 的技巧，另外再次致谢 @DiegoD94

0x1. 效果

下面展示一下在 SGLang 中的端到端效果：

GSM8K Acc Test

➜  sglang git:(support_r1_shared_expers_fusion) ✗ python3 benchmark/gsm8k/bench_sglang.py --num-questions 2000 --parallel 2000 --num-shots 8                                            
100%|████████████████████████████████████████████████████████████████████████| 1319/1319 [01:08<00:00, 19.14it/s]
Accuracy: 0.952
Invalid: 0.000
Latency: 69.547 s
Output throughput: 1998.856 token/s

Benchmark in H200

注意，吞吐越高越好，TTFT/ITL越低越好。

从测试数据中可以看到，随着 qps 增大到 4，吞吐提升了 4%，TTFT 和 ITL 都降低了 15%-20%左右。表格对应的详细数据和服务启动命令以及 bench_serving 脚本的使用方式均贴在 https://github.com/sgl-projec... 中，感兴趣可以在这里找到相关信息。

0x2. 原理

如图 1 所示，DeepSeek 的 MoE 结构会将所有输入的 token 都送入共享专家(shared experts)，同时也会将每个 token 路由到由路由器选择的 top-k 个专家(routed experts)中，最后基于权重将所有专家的输出进行聚合。具体来说，DeepSeek R1 使用了 256 个路由专家和 1 个共享专家，对于每个 token，它会选择 top 8 个路由专家。在 VLLM 和 SGLang 的实现中，token 的隐藏状态首先通过共享专家，然后再通过 FusedMoE kernel 中的路由专家。最后将这两个输出相加作为最终输出。

一个简单的优化方法是：我们可以将共享专家的计算合并到 FusedMoE kernel 中，因为共享专家和路由专家具有完全相同的架构和形状。这样，对于每个 token，我们不再是从 256 个专家中选择 top 8 个，再单独选择 1 个共享专家，而是直接从 257 个专家(256 个路由专家+1 个共享专家)中选择 top 9 个，并且始终将第 9 个专家设置为共享专家。通过进一步调整这 9 个选定专家的聚合权重，我们可以得到与优化前的 MoE 层完全相同的输出结果。

这里的意思是，在原始的 Deepseek V3/R1 MoE forward 代码中：

def forward_normal(self, hidden_states: torch.Tensor) -> torch.Tensor:
    if self.n_shared_experts isnotNone:
        shared_output = self.shared_experts(hidden_states)
    # router_logits: (num_tokens, n_experts)
    router_logits = self.gate(hidden_states)
    final_hidden_states = (
        self.experts(hidden_states=hidden_states, router_logits=router_logits)
        * self.routed_scaling_factor
    )
    if shared_output isnotNone:
        final_hidden_states = final_hidden_states + shared_output
    if self.tp_size > 1:
        final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
    return final_hidden_states

普通的 experts 的计算结果需要乘以 self.routed_scaling_factor ，为了把 shared expert 融合到普通的 expert 中一起计算，我们需要在 grouped_topk 模块中把 top9 也就是 shared experts 对应的 topk_weights 提前除一下这个 self.routed_scaling_factor 参数，这样就可以保证数值等价，请看下面的关于 grouped_topk 中的修改。

0x3. 实现细节

除了上面提到的为了保持数值等价做的 Topk weights 的系数修改之外，我们还可以观察到我们始终将所有 token 的第 9 个专家分配给共享专家，并且这里的共享专家可能还不止一个。对应这行代码：

topk_ids[:, -1] = torch.randint(
    low=num_experts,
    high=num_experts + n_share_experts_fusion,
    size=(topk_ids.size(0),),
    dtype=topk_ids.dtype,
    device=topk_ids.device,
)

之所以设定多个共享专家，原因可能是可以对 token 做一个类似负载均衡的策略，防止一个共享专家分配到所有的 token 导致 Expert 间的 token 数差距过大导致计算效率低的问题。

我 tuning 并 Benchmark 了一下不使用任何额外复制的 shared experts 的情况，发现这种情况的性能确实比使用 tp_size 个 shared experts 复制的情况要糟糕，TTFT 甚至会延长。具体可以看这里的 benchmark 数据：https://github.com/sgl-projec...

要支持复制多个 shared experts 要修改一下 DeepseekV2ForCausalLM 类的 load_weights 函数：

if self.n_share_experts_fusion != 0:
            weights_list = list(weights)
            weights_dict = dict(weights_list)
            suffix_list = [
                "down_proj.weight",
                "down_proj.weight_scale_inv",
                "gate_proj.weight",
                "gate_proj.weight_scale_inv",
                "up_proj.weight",
                "up_proj.weight_scale_inv",
            ]
            names_to_remove = []
            for moe_layer in tqdm(
                range(
                    self.config.first_k_dense_replace,
                    self.config.num_hidden_layers,
                    self.config.moe_layer_freq,
                ),
                desc=f"Cloning {self.n_share_experts_fusion} "
                "replicas of the shared expert into MoE",
            ):
                for num_repeat in range(self.n_share_experts_fusion):
                    for suffix in suffix_list:
                        weights_list.append(
                            (
                                f"model.layers.{moe_layer}."
                                f"mlp.experts."
                                f"{self.config.n_routed_experts + num_repeat}"
                                f".{suffix}",
                                weights_dict[
                                    f"model.layers.{moe_layer}.mlp.shared_experts.{suffix}"
                                ].clone(),
                            )
                        )
                    names_to_remove += [
                        f"model.layers.{moe_layer}.mlp.shared_experts.{suffix}"
                        for suffix in suffix_list
                    ]
            weights = [w for w in weights_list if w[0] notin names_to_remove]

这里还把原始权重中每一层多余的那个 shared expert 的权重给移除了。

其它还有需要的注意的细节是由于 expert 个数有变化，所以还需要使用 tuning fused moe 脚本来针对 tuning 一下。此外，相比于原始的非 fuse 实现版本，每个 tp rank 上的显存会增加，对于 FP8 dtype 来说，一个 rank 上新增一个 shared expert。一个 shared expert 按照 FP8 dtype 折算为 42MB 内存（总参数量 = hiddensize moe_intermediate_size 3 = 7168 \_ 2048 3），然后（51-3）42MB=2016MB/1024=1.96GB。也就是说这个优化将在 TP Rank 上多使用 1-2GB 的内存才可以获得最佳性能收益。

另外，这个优化在 fused moe gate 的加持下会更有用一些，期待后续进展。

END

作者：BBuf
来源：GiantPandaLLM

推荐阅读

欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。