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

📅 2026/7/8 13:17:26
第 4 章 MoE 混合专家完整源码解析(V3 核心)
第 4 章 MoE 混合专家完整源码解析V3 核心4.1 671B 总参、37B 动态激活专家架构设计4.1.1 MoE 架构概述DeepSeek-V3 采用稀疏激活的 MoEMixture of Experts架构实现了大模型性能、小模型成本的突破。核心参数参数值说明总参数量671B全部专家权重总和动态激活参数37B每 token 实际计算的参数路由专家数量256n_routed_experts共享专家数量1n_shared_experts每 token 激活专家数8n_activated_experts专家分组数8n_expert_groups受限分组数4n_limited_groups4.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):definit(self, args: ModelArgs):super().init()self.dim args.dimself.topk args.n_activated_experts # 8self.n_groups args.n_expert_groups # 8self.topk_groups args.n_limited_groups # 4self.score_func args.score_func # sigmoidself.route_scale args.route_scale # 2.5self.weight nn.Parameter(torch.empty(args.n_routed_experts, args.dim)) self.bias nn.Parameter(torch.empty(args.n_routed_experts, dtypetorch.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, dtypetorch.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, dtypebool).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, keepdimTrue) weights * self.route_scale return weights.type_as(x), indices4.2.2 路由流程详解两阶段路由算法阶段 1分组路由输入 x 经过线性层得到 scores按组划分后选择 topk_groups 个分组。阶段 2专家选择从选中分组中选择 topk 个专家获取路由权重。4.2.3 负载均衡策略DeepSeek-V3 采用无辅助损失的负载均衡策略偏置修正通过 self.bias 参数动态调整专家得分分组限制限制每 token 只能选择有限数量的分组Sigmoid 归一化防止少数专家垄断所有 token4.3 并行专家计算、多卡 MoE 通信逻辑4.3.1 MoE 类整体架构model.py 中的 MoE 类class MoE(nn.Module):definit(self, args: ModelArgs):super().init()self.dim args.dimassert 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/1616 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(), minlengthself.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 256Rank 0: 专家 0-15Rank 1: 专家 16-31…Rank 15: 专家 240-2554.3.3 通信流程token 路由到不同专家各 rank 计算本地专家输出最后通过 all_reduce 汇总。4.4 专家层裁剪、稀疏计算优化源码改造4.4.1 专家层裁剪策略def prune_experts(model, keep_ratio0.5):for layer in model.layers:if isinstance(layer.ffn, MoE):moe layer.ffnexpert_freq calculate_expert_frequency(moe) keep_count int(moe.n_routed_experts * keep_ratio) keep_indices expert_freq.argsort(descendingTrue)[: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