""" ============================================================================= 第 4 讲 · 代码 4：从零实现混合专家（MoE）层 ============================================================================= 对应讲义：lectures/lecture04_attention_alternatives_moe.md 第 5、6 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- MoE 的核心想法：把一个大前馈层（FFN）换成 E 个"专家"（每个都是一个 FFN）， 再加一个"路由器"（router）为每个 token 挑选其中 k 个专家。 - 参数量 ≈ E 个专家的总和 → 很大； - 每个 token 的计算量 ≈ k 个专家 → 很小。 "参数多、计算少"：同样的 FLOPs 下，模型容量大得多。 这个文件从零实现一个 MoE 层，并用实验回答： 1. 它和普通 FFN 是什么关系？（E = 1、k = 1 时两者完全相同） 2. 总参数、激活参数、FLOPs 分别是多少？（用 FLOP 计数器验证"计算量只与 k 有关"） 3. "细粒度专家"和"共享专家"（DeepSeekMoE）是什么？ 4. 两种常见的门控计算方式有什么区别？ 5. 专家容量与 token 丢弃：为什么"别人的请求会把你的 token 挤掉"？ 6. 真实模型：Mixtral 8x7B 的 467 亿参数和 129 亿激活参数是怎么算出来的？ ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 04_moe_layer.py 依赖第 3 讲的 code/lecture03/01_building_blocks.py（Linear、SwiGLU）。 """ import importlib import math import sys from pathlib import Path import torch from torch.utils.flop_counter import FlopCounterMode sys.path.append(str(Path(__file__).parent.parent / "lecture03")) _blocks = importlib.import_module("01_building_blocks") Linear, SwiGLU = _blocks.Linear, _blocks.SwiGLU def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) def num_params(module): """模型的参数个数。""" return sum(p.numel() for p in module.parameters()) # ============================================================================= # 第 1 部分：路由器 # ============================================================================= class TopKRouter(torch.nn.Module): """ Top-k 路由器："token 选择专家"（token-choice），几乎所有 MoE 都用它。 【思路】 1. 一个线性层给每个专家打分：logits = x W_r，形状 (T, E)； 2. 选出分数最高的 k 个专家； 3. 计算这 k 个专家的"门控权重"（gate），专家输出按它加权求和。 【两种门控计算方式】（mode 参数） "softmax_then_topk"：先对全部 E 个分数做 softmax，再取前 k 个概率作为权重。 k 个权重之和 < 1。Switch Transformer（k=1）、DeepSeek V1/V2 用这种。 "topk_then_softmax"：先取前 k 个分数，只在这 k 个上做 softmax。k 个权重之和 = 1。 Mixtral、DBRX 用这种。 【返回】 gates : (T, k) 门控权重 idx : (T, k) 选中的专家编号 probs : (T, E) 全部专家上的 softmax 概率（负载均衡损失要用，见 05 号文件） logits : (T, E) 原始分数（router z-loss 要用） """ def __init__(self, d_model, num_experts, k, mode="softmax_then_topk"): super().__init__() if not (1 <= k <= num_experts): raise ValueError("k 必须在 1 到 num_experts 之间") if mode not in ("softmax_then_topk", "topk_then_softmax"): raise ValueError("mode 只能是 'softmax_then_topk' 或 'topk_then_softmax'") self.k, self.mode = k, mode self.proj = Linear(d_model, num_experts) def forward(self, x, bias=None): """ 【参数】bias : 可选的 (E,) 偏置，**只用于选择专家**、不影响门控权重（DeepSeek-V3 的负载均衡方法，见 05 号文件） """ logits = self.proj(x).float() # 路由器用 FP32 计算：数值稳定性（讲义第 6.4 节） probs = torch.softmax(logits, dim=-1) select_scores = logits if bias is None else logits + bias idx = select_scores.topk(self.k, dim=-1).indices if self.mode == "softmax_then_topk": gates = probs.gather(-1, idx) else: gates = torch.softmax(logits.gather(-1, idx), dim=-1) return gates, idx, probs, logits # ============================================================================= # 第 2 部分：MoE 层 # ============================================================================= class MoE(torch.nn.Module): """ 混合专家层，可以直接替换 Transformer 块中的 FFN。 【参数】 d_model : 模型维度 num_experts : 路由专家个数 E k : 每个 token 激活几个路由专家 d_ff_expert : 每个专家（SwiGLU）的隐藏维度 num_shared : 共享专家个数（每个 token 都会经过它们，DeepSeekMoE 的设计），默认 0 router_mode : 见 TopKRouter capacity_factor : None 表示不限容量；否则每个专家最多处理 ceil(capacity_factor × T × k / E) 个 token，超出的被"丢弃" 【形状】x : (..., d_model) → 同形状 【实现方式】 对每个专家：找出路由到它的 token → 把这些 token 聚在一起算一次 → 按门控权重加回各自的位置。 这样每个专家只计算分给它的 token，计算量正比于 k 而不是 E。 （真实系统会把不同专家放在不同 GPU 上，用 all-to-all 通信收发 token，见第 7、8 讲的专家并行。） 【边界】被丢弃的 token 在这一层的 MoE 输出为 0（该路由的贡献缺失）； 在完整模型中它仍然通过残差连接原样传到下一层，所以不会"消失"，只是少了这一层的处理。 """ def __init__(self, d_model, num_experts, k, d_ff_expert, num_shared=0, router_mode="softmax_then_topk", capacity_factor=None): super().__init__() self.num_experts, self.k, self.capacity_factor = num_experts, k, capacity_factor self.router = TopKRouter(d_model, num_experts, k, router_mode) self.experts = torch.nn.ModuleList([SwiGLU(d_model, d_ff_expert) for _ in range(num_experts)]) self.shared = torch.nn.ModuleList([SwiGLU(d_model, d_ff_expert) for _ in range(num_shared)]) self.register_buffer("balance_bias", torch.zeros(num_experts)) # 05 号文件的"无辅助损失均衡"会用到 self.use_balance_bias = False self.last = {} # 记录最近一次前向的统计信息 def forward(self, x): shape = x.shape x = x.reshape(-1, shape[-1]) # (T, d) T = x.shape[0] bias = self.balance_bias if self.use_balance_bias else None gates, idx, probs, logits = self.router(x, bias) out = torch.zeros_like(x) capacity = math.ceil(self.capacity_factor * T * self.k / self.num_experts) if self.capacity_factor else None tokens_per_expert = torch.zeros(self.num_experts) dropped = 0 for e, expert in enumerate(self.experts): # 【PyTorch 小知识】nonzero(as_tuple=True) 返回满足条件的元素的下标： # token_ids[i] 是第几个 token，slot[i] 是它的第几个选择（0..k-1） token_ids, slot = (idx == e).nonzero(as_tuple=True) tokens_per_expert[e] = len(token_ids) if capacity is not None and len(token_ids) > capacity: dropped += len(token_ids) - capacity token_ids, slot = token_ids[:capacity], slot[:capacity] # 按 token 顺序，超出容量的被丢弃 if len(token_ids) == 0: continue y = expert(x[token_ids]) # 【PyTorch 小知识】index_add_(0, 下标, 值)：把"值"的每一行加到 out 的对应行上（同一行可以被加多次） out.index_add_(0, token_ids, gates[token_ids, slot].unsqueeze(-1).to(y.dtype) * y) for s in self.shared: out = out + s(x) self.last = dict(probs=probs, logits=logits, idx=idx, load=tokens_per_expert / (T * self.k), dropped=dropped, total=T * self.k) return out.reshape(shape) # ============================================================================= # 第 3 部分：演示 # ============================================================================= def demo_equivalence(): """演示 1：E = 1、k = 1 的 MoE 就是普通 FFN。""" print_header("演示 1：只有 1 个专家的 MoE = 普通前馈层") torch.manual_seed(0) d, d_ff = 64, 192 moe = MoE(d, num_experts=1, k=1, d_ff_expert=d_ff) dense = SwiGLU(d, d_ff) dense.load_state_dict(moe.experts[0].state_dict()) x = torch.randn(4, 10, d) print(f" 门控权重（softmax 只有一个选项）= {moe.router(x.reshape(-1, d))[0][0].item():.1f}") print(f" MoE 输出与 SwiGLU 输出的最大差异：{(moe(x) - dense(x)).abs().max().item():.1e}") def count_flops(fn): """统计矩阵乘法 FLOPs。""" with FlopCounterMode(display=False) as fc: fn() return fc.get_total_flops() def demo_params_and_flops(): """演示 2：总参数、激活参数、FLOPs。""" print_header("演示 2：参数多、计算少") torch.manual_seed(0) d, d_ff, T = 256, 704, 512 x = torch.randn(T, d) configs = [ ("普通 FFN", SwiGLU(d, d_ff), 1), ("MoE：8 个专家，每 token 选 1 个", MoE(d, 8, 1, d_ff), 1), ("MoE：8 个专家，每 token 选 2 个", MoE(d, 8, 2, d_ff), 2), ("MoE：64 个专家，每 token 选 2 个", MoE(d, 64, 2, d_ff), 2), ] dense_flops = None print(f" {'':<34}{'总参数':>12}{'每 token 激活参数':>18}{'FLOPs（相对普通 FFN）':>22}") for name, m, k in configs: total = num_params(m) active = k * num_params(SwiGLU(d, d_ff)) + (d * getattr(m, "num_experts", 0)) flops = count_flops(lambda: m(x)) dense_flops = dense_flops or flops print(f" {name:<34}{total:>12,}{active:>18,}{flops / dense_flops:>20.2f} 倍") print(""" - 专家数 E 从 8 增加到 64：总参数涨了 8 倍，FLOPs 几乎不变（只多了一点路由器的计算）。 - k 从 1 增加到 2：FLOPs 翻倍。**计算量由 k 决定，参数量由 E 决定。** - 这正是 MoE 的吸引力：大量研究（Switch Transformer、OLMoE 等）表明，在相同 FLOPs 下， 参数更多的 MoE 模型损失更低、训练更快。 【边界】"FLOPs 少"不等于"便宜"：所有专家的参数都要放进显存（推理时尤其明显）， 而且路由带来的数据搬运和通信开销在 FLOPs 里完全体现不出来（讲义第 6.3 节）。""") def demo_fine_grained_and_shared(): """演示 3：细粒度专家与共享专家（DeepSeekMoE）。""" print_header("演示 3：细粒度专家与共享专家") d, d_ff = 256, 704 coarse = MoE(d, 8, 2, d_ff) fine = MoE(d, 32, 8, d_ff // 4) # 每个专家切成 4 份，激活数也 ×4 fine_shared = MoE(d, 31, 7, d_ff // 4, num_shared=1) # 其中 1 份变成总是激活的共享专家 print(f" {'':<42}{'总参数':>12}{'激活的专家组合数':>18}") for name, m, n_routed, k in [("8 个专家选 2 个", coarse, 8, 2), ("32 个细粒度专家选 8 个", fine, 32, 8), ("1 个共享 + 31 个细粒度专家选 7 个", fine_shared, 31, 7)]: print(f" {name:<40}{num_params(m):>12,}{math.comb(n_routed, k):>18,}") print(""" 三种配置的总参数和每 token 的计算量几乎相同（每个 token 都用到"2 个完整专家"那么多的参数）， 但细粒度专家的可选组合从 28 种暴增到上千万种：每个 token 可以更灵活地"拼装"自己需要的知识。 共享专家则用来承载所有 token 都需要的通用知识，让路由专家更专注于各自的领域。 DeepSeek 的消融实验显示两者都有帮助；OLMoE 的实验则只发现细粒度有帮助、共享专家没有 —— 结论并不统一。""") def demo_router_modes(): """演示 4：两种门控计算方式。""" print_header("演示 4：先 softmax 再 top-k，还是先 top-k 再 softmax？") torch.manual_seed(0) d, E = 32, 8 x = torch.randn(1, d) r1 = TopKRouter(d, E, 2, "softmax_then_topk") r2 = TopKRouter(d, E, 2, "topk_then_softmax") r2.load_state_dict(r1.state_dict()) g1, i1, probs, _ = r1(x) g2, i2, _, _ = r2(x) print(f" 全部 {E} 个专家的概率：{[round(p, 3) for p in probs[0].tolist()]}") print(f" 先 softmax 再 top-2：选中专家 {i1[0].tolist()}，权重 {[round(g, 3) for g in g1[0].tolist()]}，和 = {g1.sum().item():.3f}") print(f" 先 top-2 再 softmax：选中专家 {i2[0].tolist()}，权重 {[round(g, 3) for g in g2[0].tolist()]}，和 = {g2.sum().item():.3f}") print(" 选中的专家相同，区别只在权重是否归一化到 1。（DeepSeek-V3 又换成了 sigmoid 打分，再对选中的归一化。）") def demo_capacity(): """演示 5：专家容量与 token 丢弃。""" print_header("演示 5：专家容量与 token 丢弃") torch.manual_seed(0) d, E, T = 64, 8, 1024 moe = MoE(d, E, 1, 128) x = torch.randn(T, d) moe(x) load = moe.last["load"] print(f" 各专家分到的 token 比例（理想是 {1 / E:.3f}）：{[round(v, 3) for v in load.tolist()]}") print(" （随机初始化的路由器本身就不均衡：有的专家多、有的少。）") for cf in [1.0, 1.25, 2.0]: moe.capacity_factor = cf moe(x) cap = math.ceil(cf * T / E) print(f" 容量因子 {cf:<4}（每个专家最多 {cap} 个 token）：丢弃了 {moe.last['dropped'] / moe.last['total']:.1%} 的 token") print(""" 为什么要限制容量？在专家并行中，每个专家在一块（或一组）GPU 上，要预先分配固定大小的缓冲区， 所有专家同步进行计算。一个专家分到的 token 太多，其他 GPU 就得等它 —— 所以超出容量的 token 被丢弃。 容量因子越大丢得越少，但浪费的缓冲区（和计算）越多。根本的解决办法是让负载均衡（05 号文件）。""") # 【边界】token 是否被丢弃，取决于同一批次里的其他 token moe.capacity_factor = None moe(x) routes = moe.last["idx"][:, 0] probe = x[:1] # 我们关心的"那个 token" target = routes[0].item() # 它想去的专家 same = x[1:][routes[1:] == target][:127] # 另外一些也想去这个专家的 token others = x[1:][routes[1:] != target][:127] # 一些去别的专家的 token moe.capacity_factor = 1.0 # 128 个 token、8 个专家 → 每个专家容量 16 for name, batch, pos in [("批次 A（和去别处的 token 一起）", torch.cat([probe, others]), 0), ("批次 B（排在 127 个同样想去专家 %d 的 token 后面）" % target, torch.cat([same, probe]), len(same))]: out = moe(batch) kept = out[pos].norm().item() > 0 print(f" 【边界】{name}：输出范数 {out[pos].norm().item():.3f}{'' if kept else ' ← 被丢弃了！'}") print(" 同一个输入，只因为同一批次里有别人的请求，结果就不同 —— 这是 MoE 推理结果带有额外随机性的一个来源。") def mixtral_params(): """ 按公开的配置计算 Mixtral 8x7B 的总参数和激活参数。 【配置】d_model = 4096，32 层，32 个查询头 / 8 个 KV 头（头维度 128）， 每个专家是 d_ff = 14336 的 SwiGLU，8 个专家选 2 个，词表 32000，嵌入与输出层不共享。 """ d, L, V, d_ff, E, k = 4096, 32, 32000, 14336, 8, 2 attn = d * (32 * 128) * 2 + d * (8 * 128) * 2 # W_Q、W_O 与 W_K、W_V（GQA） expert = 3 * d * d_ff router = d * E norms = 2 * d per_layer_total = attn + E * expert + router + norms per_layer_active = attn + k * expert + router + norms embed = 2 * V * d + d return L * per_layer_total + embed, L * per_layer_active + embed def demo_real_models(): """演示 6：真实的 MoE 模型。""" print_header("演示 6：真实模型的 MoE 配置") total, active = mixtral_params() print(f" 按配置计算 Mixtral 8x7B：总参数 {total / 1e9:.1f}B，每 token 激活 {active / 1e9:.1f}B" "（官方公布：46.7B / 12.9B）") print(" （名字叫 8x7B 但总参数不是 56B：注意力层、嵌入层是所有专家共享的，只有前馈层复制了 8 份。）") print("\n 部分 MoE 模型的路由设置（来自 CS336 课件与各模型的论文）：") rows = [ ("Switch Transformer", 64, 1, 0, ""), ("Mixtral 8x7B", 8, 2, 0, "46.7B 总 / 12.9B 激活"), ("DBRX", 16, 4, 0, ""), ("DeepSeek V1（DeepSeekMoE 16B）", 64, 6, 2, "16B 总 / 2.8B 激活"), ("Qwen1.5-MoE", 60, 4, 4, ""), ("OLMoE", 64, 8, 0, "6.9B 总 / 1.3B 激活"), ("DeepSeek V3", 256, 8, 1, "671B 总 / 37B 激活"), ("Llama 4 Maverick", 128, 1, 1, "400B 总 / 17B 激活"), ] print(f" {'模型':<30}{'路由专家':>8}{'激活':>6}{'共享':>6} 规模") for name, e, k, s, size in rows: print(f" {name:<30}{e:>8}{k:>6}{s:>6} {size}") print(" 趋势：专家越来越多、越来越细，常配 1~2 个共享专家；激活参数只占总参数的几十分之一。") def main(): demo_equivalence() demo_params_and_flops() demo_fine_grained_and_shared() demo_router_modes() demo_capacity() demo_real_models() print("\n全部演示运行完毕。") if __name__ == "__main__": main()