""" ============================================================================= CS349E · 第 3 讲 · 05_moe_switch.py —— 混合专家（MoE）：Switch 路由、容量因子、负载均衡，以及 MoE 的 roofline ============================================================================= 【这个脚本要回答的问题】 Switch Transformer（Fedus, Zoph, Shazeer 2021/2022，本讲阅读材料）把 MLP 换成 E 个"专家"，每个 token 只走其中一个。 参数量涨 E 倍，每个 token 的运算量不变——这是**条件计算**（conditional computation），也是一种训练出来的稀疏。 演示 1：一个 Switch 层：路由器、top-1、门控值；和逐 token 计算的参照对比；参数 vs FLOPs（论文的模型尺寸）。 演示 2：专家容量与容量因子：负载不均时，超出容量的 token 被"丢弃"；容量因子越大丢得越少、浪费越多。 演示 3：负载均衡损失 α·N·Σ f_i·P_i：它怎样把 token 推向空闲的专家，以及它和"任务"之间的拉锯。 演示 4：MoE 在 roofline 上：decode 时一个批会碰到多少个专家、要读多少权重；为什么 MoE 需要大得多的批才能 compute-bound。 【怎么读输出】全部是确定性的计算（固定种子），每次运行结果相同。路由的数据是合成的（有大小不一的"话题"簇）。 【运行】python3 05_moe_switch.py （约 3 秒） 【依赖】numpy；../lecture01 的 console、hw_specs。 """ import math import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_bytes, fmt_num, fmt_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import H100 # noqa: E402 def softmax(z): """数值稳定的 softmax（沿最后一维）。""" z = z - z.max(axis=-1, keepdims=True) e = np.exp(z) return e / e.sum(axis=-1, keepdims=True) # --------------------------------------------------------------------------- # Switch 层 # --------------------------------------------------------------------------- def switch_layer(x, Wr, W1, W2, capacity_factor): """ 一个 Switch（top-1 MoE）层。 【参数】 x ：[T, d]，这一批的 T 个 token Wr ：[d, E]，路由器（一个线性层） W1 ：[E, d, f]，W2：[E, f, d]，E 个专家（每个是 ReLU MLP：relu(x·W1)·W2） capacity_factor：容量因子 CF 【返回】(y [T, d], kept：每个 token 有没有被处理的布尔数组, expert：每个 token 选的专家, capacity) 【思路（论文公式 1–3）】 p(x) = softmax(x·Wr)；选概率最大的专家 i；输出 y = p_i(x) · E_i(x)（乘门控值，路由器才能得到梯度）。 每个专家最多处理 C = ⌈T / E × CF⌉ 个 token（专家容量）。超出的 token 被"丢弃"：这一层对它输出 0， 它只靠残差连接把原来的值传到下一层。按 token 在批里的顺序先到先得。 【为什么要固定容量】硬件（尤其是 TPU）需要静态的张量形状：每个专家的输入预先分配 [C, d] 的缓冲区。 """ T, d = x.shape E = Wr.shape[1] p = softmax(x @ Wr) expert = np.argmax(p, axis=1) gate = p[np.arange(T), expert] C = int(math.ceil(T / E * capacity_factor)) y = np.zeros_like(x) kept = np.zeros(T, dtype=bool) for i in range(E): idx = np.nonzero(expert == i)[0][:C] if len(idx) == 0: continue h = np.maximum(x[idx] @ W1[i], 0.0) y[idx] = gate[idx, None] * (h @ W2[i]) kept[idx] = True return y, kept, expert, C def clustered_tokens(T, d, E, zipf=True, seed=0): """ 合成的 token 表示：E 个"话题"簇，簇的大小按 1/k 分布（zipf=True）或均匀；每个 token = 强度 m × 簇中心 + 噪声。 【返回】(X [T, d], 每个 token 的簇号, 簇中心 [E, d]) 【直觉】真实文本里有些话题（代码、数学、日常对话……）比别的常见得多，所以路由器很容易把大量 token 送给少数专家。 m ~ 均匀(0, 3) 表示"这个 token 对自己话题的偏好有多强"：m 小的 token 送去哪个专家都差不多。 """ rng = np.random.default_rng(seed) w = 1.0 / np.arange(1, E + 1) if zipf else np.ones(E) w /= w.sum() c = rng.choice(E, size=T, p=w) mu = np.linalg.qr(rng.standard_normal((d, d)))[0][:E] # E 个互相正交的单位向量 m = rng.uniform(0, 3, size=T) X = m[:, None] * mu[c] + 0.3 * rng.standard_normal((T, d)) return X, c, mu # --------------------------------------------------------------------------- # 演示 1：Switch 层与参数 / FLOPs # --------------------------------------------------------------------------- def demo1_switch_layer(): print_header("演示 1：一个 Switch 层——每个 token 只走一个专家") rng = np.random.default_rng(1) T, d, f, E = 256, 32, 64, 8 x = rng.standard_normal((T, d)) Wr = rng.standard_normal((d, E)) / math.sqrt(d) W1 = rng.standard_normal((E, d, f)) / math.sqrt(d) W2 = rng.standard_normal((E, f, d)) / math.sqrt(f) y, kept, expert, C = switch_layer(x, Wr, W1, W2, capacity_factor=float(E)) # 容量足够大：不丢 ref = np.zeros_like(x) for t in range(T): # 最朴素的参照：逐 token 计算 p = softmax(x[t] @ Wr) i = int(np.argmax(p)) ref[t] = p[i] * (np.maximum(x[t] @ W1[i], 0) @ W2[i]) print(f"T = {T} 个 token、E = {E} 个专家：Switch 层（按专家分组计算）vs 逐 token 参照，最大差 {np.max(np.abs(y - ref)):.1e}；" f"每个专家分到的 token 数 {np.bincount(expert, minlength=E).tolist()}") dense_flops = 2 * 2 * d * f moe_flops = 2 * d * E + 2 * 2 * d * f print(f"每个 token 的 FLOPs：稠密 MLP {dense_flops:,}，Switch（路由器 + 一个专家）{moe_flops:,}；" f"参数：稠密 {2 * d * f:,}，Switch {E * 2 * d * f + d * E:,}（{E} 倍多）") rows = [["T5-Base（稠密）", "0.2 B", "124 B", "-"], ["Switch-Base（128 专家）", "7 B", "124 B", "35×"], ["Switch-XXL", "395 B", "6.3 T", "-"], ["Switch-C（2048 专家）", "1571 B", "890 B", "-"], ["DeepSeek-V3（256 路由 + 1 共享专家，top-8）", "671 B", "每 token 激活 37 B 参数", "18×（总 / 激活）"]] print("\n论文表 9（每个序列的前向 FLOPs）和 DeepSeek-V3 技术报告：") print_table(["模型", "参数", "FLOPs / 序列（或激活参数）", "参数倍数"], rows) print_observations([ "Switch 层的输出和逐 token 计算完全一样；实现上把同一个专家的 token 收集到一起，做一次 [C, d]×[d, f] 的矩阵乘法——" "这是 MoE 高效的关键：按专家分组之后，仍然是大矩阵乘法。", "每个 token 的 FLOPs 几乎不变（多一个很小的路由器），参数却是 E 倍。Switch-Base 和 T5-Base 每个序列的 FLOPs 相同，" "参数多 35 倍，论文报告预训练快了 7 倍（达到同样的质量）。**参数量和计算量从此脱钩**。", "代价转移到了别处：所有专家的权重都要放在显存里（容量），decode 时一个批会碰到很多专家（带宽，演示 4），" "专家分布在不同芯片上时 token 要来回搬（all-to-all 通信，第 4 讲）。", ]) # --------------------------------------------------------------------------- # 演示 2：容量因子与丢弃 # --------------------------------------------------------------------------- def demo2_capacity(): print_header("演示 2：专家容量 = (这批的 token 数 / 专家数) × 容量因子——丢弃 vs 浪费") T, d, E = 4096, 16, 8 rows = [] for label, zipf in (("话题大小均匀", False), ("话题大小 ∝ 1/k", True)): X, c, mu = clustered_tokens(T, d, E, zipf=zipf, seed=2) Wr = mu.T * 3.0 # 一个"按话题路由"的路由器 e = np.argmax(X @ Wr, axis=1) load = np.bincount(e, minlength=E) / T for cf in (1.0, 1.25, 1.5, 2.0, 4.0): C = int(math.ceil(T / E * cf)) kept = np.minimum(np.bincount(e, minlength=E), C).sum() rows.append([label, f"{load.max() * E:.2f}", f"{cf:g}", C, f"{(T - kept) / T * 100:.1f}%", f"{(E * C - kept) / (E * C) * 100:.1f}%"]) print(f"T = {T} 个 token，E = {E} 个专家；'最忙专家的负载'= 最忙专家分到的比例 × E（1.00 = 完全均衡）") print_table(["数据", "最忙专家的负载", "容量因子", "每个专家的容量", "被丢弃的 token", "空着的槽位（浪费）"], rows) print_observations([ "负载均匀时，容量因子 1.0 只丢 1.1% 的 token，1.25 就一个不丢，代价是 20% 的槽位空着。论文的实验也在 1.0–1.25 时" "效果最好，报告通常丢弃 < 1%。", "负载不均时（最忙的专家分到 2.7 倍的平均量），容量因子 1.25 要丢掉 19% 的 token；容量因子提到 4 才不丢，" "但这时 75% 的槽位是空的——算力和显存都白白分配了。", "所以光靠加大容量因子解决不了问题，必须让路由**本身**更均衡：这就是负载均衡损失（演示 3）。", ]) # --------------------------------------------------------------------------- # 演示 3：负载均衡损失 # --------------------------------------------------------------------------- def train_router(X, c, E, alpha, steps=400, lr=2.0): """ 训练路由器：损失 = 任务项 + α · E · Σ_i f_i · P_i。 【任务项】我们没有真实的语言模型，就用"把 token 送到它自己话题对应的专家"的交叉熵代替（偏好强的 token 更难被挪走）。 【负载均衡项（论文公式 4–6）】f_i = 分到专家 i 的 token 比例（来自 argmax，不可导），P_i = 专家 i 的平均路由概率（可导）。 均匀时 Σ f_i·P_i = 1/E，乘 E 后等于 1，是最小值。 【梯度（手推）】对第 t 个 token 的 logit z_tj：∂/∂z_tj = (α·E / T) · p_tj · (f_j − Σ_i f_i·p_ti)—— 分给一个"过忙"专家（f_j 大）的概率会被压低。 【返回】(训练后的路由器 W [d, E], 训练过程中最后一步的 f) """ T, d = X.shape W = np.zeros((d, E)) Y = np.eye(E)[c] for _ in range(steps): p = softmax(X @ W) f = np.bincount(np.argmax(p, axis=1), minlength=E) / T g = (p - Y) / T g += alpha * E / T * p * (f[None, :] - (p * f[None, :]).sum(axis=1, keepdims=True)) W -= lr * X.T @ g return W, f def demo3_balance_loss(): print_header("演示 3：负载均衡损失 α·E·Σ f_i·P_i——把 token 推向空闲的专家") E = 8 rows = [] for name, f in (("完全均衡", np.full(E, 1 / E)), ("一半给专家 0", np.r_[0.5, np.full(E - 1, 0.5 / (E - 1))]), ("全部给专家 0", np.r_[1.0, np.zeros(E - 1)])): rows.append([name, f"{E * np.sum(f * f):.3f}"]) print("取 P = f 时（路由器很确定），E·Σ f_i·P_i 的值：") print_table(["分配", "E·Σ f_i·P_i"], rows) T, d = 4096, 16 X, c, _ = clustered_tokens(T, d, E, zipf=True, seed=3) rows = [] for alpha in (0.0, 0.01, 0.1, 0.3, 1.0, 3.0): W, _ = train_router(X, c, E, alpha) p = softmax(X @ W) e = np.argmax(p, axis=1) f = np.bincount(e, minlength=E) / T drops = [] for cf in (1.25, 2.0): C = int(math.ceil(T / E * cf)) drops.append(np.maximum(np.bincount(e, minlength=E) - C, 0).sum() / T) aux = E * np.sum(f * p.mean(axis=0)) rows.append([f"{alpha:g}", f"{f.max() * E:.2f}", f"{aux:.3f}", f"{drops[0] * 100:.1f}%", f"{drops[1] * 100:.1f}%", f"{np.mean(e == c) * 100:.1f}%"]) print(f"\n话题大小 ∝ 1/k 的 {T} 个 token，训练路由器 400 步（任务项 + α × 负载均衡项）：") print_table(["α", "最忙专家的负载", "E·Σ f·P", "丢弃（CF 1.25）", "丢弃（CF 2）", "送到'自己话题'专家的比例"], rows) print_observations([ "E·Σ f_i·P_i 在完全均衡时等于 1（最小值），越不均衡越大（全部给一个专家时等于 E）。乘以 E 是为了让这个值不随专家数变化。", "f 来自 argmax，不可导；P 是概率的平均，可导。梯度只通过 P 起作用：哪个专家的 f 大，分给它的概率就被往下压，" "那些'去哪都行'（偏好弱）的 token 会先被挪到空闲的专家。", "α 越大越均衡：丢弃率（CF = 1.25）从 22.7% 降到 9.3%，CF = 2 时降到 0；代价是更多 token 被送离'自己话题'的专家" "（87% → 78%）——**均衡和专业化之间有拉锯**。", "在这个玩具里，α 要到 1 左右才明显起作用，而 Switch 论文用 α = 0.01（从 1e-5 到 1e-1 扫过）。原因是我们的任务项" "是直接对路由器的交叉熵，信号比真实训练强得多；真实模型里路由器只通过门控值间接得到任务梯度，很弱的均衡项就够了。", "DeepSeek-V3 换了一种做法（'无辅助损失'的负载均衡）：给每个专家的路由分数加一个偏置，过忙就调低、过闲就调高，" "偏置只影响选谁、不影响门控值，避免了均衡项对任务梯度的干扰。", ]) # --------------------------------------------------------------------------- # 演示 4：MoE 在 roofline 上 # --------------------------------------------------------------------------- def demo4_moe_roofline(): print_header("演示 4：MoE 的 decode——一个批会碰到多少专家、要读多少权重（DeepSeek-V3 的一个 MoE 层，H100）") E, k, d, f = 256, 8, 7168, 2048 shared = 1 w_bytes = 1.0 # FP8 权重 expert_params = 3 * d * f # SwiGLU：gate、up、down chip = H100 peak = chip.peak_flops("fp8") rows = [] for B in (1, 8, 32, 128, 512, 2048, 8192): touched = E * (1 - (1 - k / E) ** B) # 均匀随机路由时被碰到的专家数的期望 by = (touched + shared) * expert_params * w_bytes fl = B * (k + shared) * 2 * expert_params t_mem, t_math = by / chip.mem_bw, fl / peak per_expert = B * k / E rows.append([B, f"{touched:.1f}", f"{per_expert:.2f}", fmt_bytes(by), f"{fl / by:.1f}", fmt_time(t_mem), fmt_time(t_math), "compute" if t_math > t_mem else "memory"]) print(f"E = {E} 个路由专家（每个 {fmt_num(expert_params)} 参数）+ {shared} 个共享专家，每个 token 选 {k} 个；" f"FP8 权重和计算；屋脊点 {chip.critical_intensity('fp8'):.0f}") print_table(["批大小 B", "碰到的专家（期望）", "每个专家平均分到的 token", "读权重", "强度", "读权重时间", "计算时间", "瓶颈"], rows) b_star = E * chip.critical_intensity("fp8") / (2 * k) print(f"\n所有专家都被碰到之后，要 compute-bound 需要 B > 屋脊点 × E / (2k) ≈ {b_star:,.0f} 个 token（稠密层只要约 {chip.critical_intensity('fp8') / 2:.0f}）。" f"How to Scale Your Model 第 4 章对 TPU v5e、int8 权重给出 B > 120·E/k = {120 * E // k:,}。") print_observations([ "B = 1 时只读 9 个专家（8 个路由 + 1 个共享），396 MB、0.12 ms：MoE 的单 token decode 只读一小部分参数，" "这正是'每 token 只激活 37 B 参数'的好处。", "但批一大，碰到的专家迅速增多：B = 32 时已经碰到 163 个，B = 128 时几乎全部 256 个。读权重的字节涨了 28 倍，" "而每个专家平均只分到 B·k/E 个 token——**每个专家的'有效批'被稀释了 E/k = 32 倍**。", f"结果：稠密模型 B ≈ 300 就 compute-bound（FP8），这个 MoE 层要 B ≈ {b_star:,.0f}。decode 时根本凑不出这么大的批，" "所以 MoE 的 decode 几乎总是 memory-bound，而且要读的权重是稠密模型的很多倍。", "DeepSeek 的应对（第 4、5 讲）：把专家分散到很多块 GPU 上（专家并行，每块卡只放几个专家，它们分到的 token 来自整个集群的批），" "用 all-to-all 通信把 token 送过去——用通信换'每个专家的有效批'。这又是一个'省一处、花另一处'的取舍。", ]) if __name__ == "__main__": demo1_switch_layer() demo2_capacity() demo3_balance_loss() demo4_moe_roofline()