第 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 采用无辅助损失的负载均衡策略:
- 偏置修正:通过 self.bias 参数动态调整专家得分
- 分组限制:限制每 token 只能选择有限数量的分组
- 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