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

0 阅读3分钟

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

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

4.1.1 MoE 架构概述

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

核心参数:

参数说明
总参数量671B全部专家权重总和
动态激活参数37B每 token 实际计算的参数
路由专家数量256n_routed_experts
共享专家数量1n_shared_experts
每 token 激活专家数8n_activated_experts
专家分组数8n_expert_groups
受限分组数4n_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)
总参数量72B671B
激活参数量72B37B
训练成本

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 裁剪前后对比

配置专家数总参数量推理速度
原始256671B1x
裁剪 50%128336B1.8x
裁剪 75%64168B2.5x

本章小结:

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