第 4 章 MoE 混合专家完整源码解析(V3 核心)

4.1 671B 总参、37B 动态激活专家架构设计

4.1.1 MoE 架构概述

DeepSeek-V3 采用稀疏激活的 MoE(Mixture of Experts)架构,实现了大模型性能、小模型成本的突破。

核心参数:

参数 说明
总参数量 671B 全部专家权重总和
动态激活参数 37B 每 token 实际计算的参数
路由专家数量 256 n_routed_experts
共享专家数量 1 n_shared_experts
每 token 激活专家数 8 n_activated_experts
专家分组数 8 n_expert_groups
受限分组数 4 n_limited_groups

4.1.2 DeepSeek-V3 MoE 配置

configs/config_671B.json:

{
“dim”: 7168,
“inter_dim”: 18432,
“moe_inter_dim”: 2048,
“n_layers”: 61,
“n_dense_layers”: 3,
“n_heads”: 128,
“n_routed_experts”: 256,
“n_shared_experts”: 1,
“n_activated_experts”: 8,
“n_expert_groups”: 8,
“n_limited_groups”: 4,
“route_scale”: 2.5,
“score_func”: “sigmoid”
}

4.1.3 MoE 与 Dense 模型对比

维度 Dense (72B) MoE (671B/37B)
总参数量 72B 671B
激活参数量 72B 37B
训练成本

4.2 专家路由算法、负载均衡防倾斜源码

4.2.1 Gate 门控机制

model.py 中的 Gate 类:

class Gate(nn.Module):
def init(self, args: ModelArgs):
super().init()
self.dim = args.dim
self.topk = args.n_activated_experts # 8
self.n_groups = args.n_expert_groups # 8
self.topk_groups = args.n_limited_groups # 4
self.score_func = args.score_func # sigmoid
self.route_scale = args.route_scale # 2.5

    self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim))
    self.bias = nn.Parameter(torch.empty(args.n_routed_experts, dtype=torch.float32)) if self.dim == 7168 else None

def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    scores = linear(x, self.weight)
    
    if self.score_func == "softmax":
        scores = scores.softmax(dim=-1, dtype=torch.float32)
    else:
        scores = scores.sigmoid()
    
    original_scores = scores
    
    if self.bias is not None:
        scores = scores + self.bias
    
    if self.n_groups > 1:
        scores = scores.view(x.size(0), self.n_groups, -1)
        if self.bias is None:
            group_scores = scores.amax(dim=-1)
        else:
            group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1)
        
        indices = group_scores.topk(self.topk_groups, dim=-1)[1]
        mask = scores.new_ones(x.size(0), self.n_groups, dtype=bool).scatter_(1, indices, False)
        scores = scores.masked_fill_(mask.unsqueeze(-1), float("-inf")).flatten(1)
    
    indices = torch.topk(scores, self.topk, dim=-1)[1]
    weights = original_scores.gather(1, indices)
    
    if self.score_func == "sigmoid":
        weights /= weights.sum(dim=-1, keepdim=True)
    
    weights *= self.route_scale
    
    return weights.type_as(x), indices

4.2.2 路由流程详解

两阶段路由算法:

阶段 1:分组路由
输入 x 经过线性层得到 scores,按组划分后选择 topk_groups 个分组。

阶段 2:专家选择
从选中分组中选择 topk 个专家,获取路由权重。

4.2.3 负载均衡策略

DeepSeek-V3 采用无辅助损失的负载均衡策略:

  1. 偏置修正:通过 self.bias 参数动态调整专家得分
  2. 分组限制:限制每 token 只能选择有限数量的分组
  3. Sigmoid 归一化:防止少数专家垄断所有 token

4.3 并行专家计算、多卡 MoE 通信逻辑

4.3.1 MoE 类整体架构

model.py 中的 MoE 类:

class MoE(nn.Module):
def init(self, args: ModelArgs):
super().init()
self.dim = args.dim

    assert args.n_routed_experts % world_size == 0
    self.n_routed_experts = args.n_routed_experts
    self.n_local_experts = args.n_routed_experts // world_size  # 256/16=16
    self.n_activated_experts = args.n_activated_experts
    
    self.experts_start_idx = rank * self.n_local_experts
    self.experts_end_idx = self.experts_start_idx + self.n_local_experts
    
    self.gate = Gate(args)
    
    self.experts = nn.ModuleList([
        Expert(args.dim, args.moe_inter_dim) if self.experts_start_idx <= i < self.experts_end_idx else None
        for i in range(self.n_routed_experts)
    ])
    
    self.shared_experts = MLP(args.dim, args.n_shared_experts * args.moe_inter_dim)

def forward(self, x: torch.Tensor) -> torch.Tensor:
    shape = x.size()
    x = x.view(-1, self.dim)
    
    weights, indices = self.gate(x)
    y = torch.zeros_like(x)
    
    counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
    
    for i in range(self.experts_start_idx, self.experts_end_idx):
        if counts[i] == 0:
            continue
        
        expert = self.experts[i]
        idx, top = torch.where(indices == i)
        y[idx] += expert(x[idx]) * weights[idx, top, None]
    
    z = self.shared_experts(x)
    
    if world_size > 1:
        dist.all_reduce(y)
    
    return (y + z).view(shape)

4.3.2 专家并行策略

专家切分方式:

world_size = 16, n_routed_experts = 256

Rank 0: 专家 0-15
Rank 1: 专家 16-31

Rank 15: 专家 240-255

4.3.3 通信流程

token 路由到不同专家,各 rank 计算本地专家输出,最后通过 all_reduce 汇总。

4.4 专家层裁剪、稀疏计算优化源码改造

4.4.1 专家层裁剪策略

def prune_experts(model, keep_ratio=0.5):
for layer in model.layers:
if isinstance(layer.ffn, MoE):
moe = layer.ffn

        expert_freq = calculate_expert_frequency(moe)
        keep_count = int(moe.n_routed_experts * keep_ratio)
        keep_indices = expert_freq.argsort(descending=True)[:keep_count]
        
        new_experts = nn.ModuleList()
        for i in keep_indices:
            new_experts.append(moe.experts[i])
        
        moe.experts = new_experts
        moe.n_routed_experts = keep_count
        
        moe.gate.weight = nn.Parameter(moe.gate.weight[keep_indices])

4.4.2 裁剪前后对比

配置 专家数 总参数量 推理速度
原始 256 671B 1x
裁剪 50% 128 336B 1.8x
裁剪 75% 64 168B 2.5x

本章小结:

DeepSeek-V3 的 MoE 架构通过两阶段路由、无辅助损失负载均衡、专家并行、稀疏计算优化等技术,实现了性能与成本的最佳平衡。
如需沟通:lxb20110121

Logo

欢迎加入DeepSeek 技术社区。在这里,你可以找到志同道合的朋友,共同探索AI技术的奥秘。

更多推荐