""" ============================================================================= 第 4 讲 · 代码 5：MoE 的负载均衡 —— 辅助损失 vs "无辅助损失"的偏置法 ============================================================================= 对应讲义：lectures/lecture04_attention_alternatives_moe.md 第 6 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 训练 MoE 时有一个"马太效应"： 某个专家一开始碰巧表现好一点 → 路由器更常选它 → 它得到更多训练、变得更好 → 路由器更常选它…… 最后少数专家包揽了大部分 token，另一些专家"饿死"（从未被选中，参数白白浪费）。 负载不均衡还会让专家并行时某些 GPU 忙死、其他 GPU 闲着。 这个文件在一个小任务上训练 MoE，比较三种做法： 1. 不做均衡 2. Switch Transformer 的辅助损失： L_aux = α · E · Σ_i f_i · P_i f_i = 分给专家 i 的 token 比例，P_i = 路由器给专家 i 的平均概率 3. DeepSeek-V3 的"无辅助损失"均衡：给每个专家一个偏置 b_i，**只用于选择专家**； 每一步后，负载高于平均的专家 b_i 减小一点，低于平均的增大一点。 并测量：最大负载、"饿死"的专家数、任务损失。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 05_load_balancing.py （约 1 分钟） 依赖同目录的 04_moe_layer.py。 """ import importlib import math import torch MoE = importlib.import_module("04_moe_layer").MoE def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：一个适合"分工"的小任务 # ============================================================================= D = 32 # 输入、输出维度 NUM_CLUSTERS = 8 # 数据来自 8 个不同的"领域" NUM_EXPERTS = 8 def make_task(seed=0): """ 构造一个回归任务：输入来自 8 个簇（像 8 个不同的"领域"），每个簇有自己的非线性映射。 【为什么这样设计】理想情况下，每个专家负责一个簇 —— 这时负载也天然均衡。 看看不同的训练方法能否找到这种分工。 【返回】sample(n, generator) 函数：返回 (x, y)，形状都是 (n, D) """ g = torch.Generator().manual_seed(seed) centers = torch.randn(NUM_CLUSTERS, D, generator=g) * 3 maps = torch.randn(NUM_CLUSTERS, D, D, generator=g) / math.sqrt(D) def sample(n, generator): c = torch.randint(0, NUM_CLUSTERS, (n,), generator=generator) x = centers[c] + torch.randn(n, D, generator=generator) y = torch.tanh(torch.einsum("nd,nde->ne", x, maps[c])) return x, y return sample # ============================================================================= # 第 2 部分：负载均衡的两种方法 # ============================================================================= def switch_aux_loss(probs, idx, num_experts): """ Switch Transformer 的负载均衡辅助损失：α 之外的部分 = E · Σ_i f_i · P_i 【参数】 probs : (T, E) 路由器的 softmax 概率 idx : (T, k) 每个 token 选中的专家 【直觉】 f_i（实际分到的比例）不可导（它来自 top-k 这个离散选择），P_i（平均概率）可导。 对 P_i 求导，梯度正比于 f_i：**被选得越多的专家，它的概率被压得越狠**。 当所有专家的 f_i 和 P_i 都等于 1/E 时，这一项取得最小值 1。 """ f = torch.bincount(idx.flatten(), minlength=num_experts).float() / idx.numel() P = probs.mean(dim=0) return num_experts * (f * P).sum() def update_balance_bias(moe, u): """ DeepSeek-V3 的"无辅助损失"负载均衡：b_i ← b_i + u · sign(平均负载 - 负载_i)。 【思路】 偏置 b_i 只加在"选择专家"用的分数上，不参与门控权重的计算，也不参与梯度下降。 负载高的专家 b_i 变小 → 下一步被选中的机会变少；负载低的反之。 这是一个简单的在线反馈控制，完全不改变训练的损失函数。 【边界】u 太大会让负载来回震荡，太小则调整太慢。DeepSeek-V3 论文中还保留了一个很小的 "序列级"辅助损失作为补充（所以 CS336 课上说它"并不完全无辅助损失"）。 """ load = moe.last["load"] moe.balance_bias += u * torch.sign(load.mean() - load) # ============================================================================= # 第 3 部分：训练与评估 # ============================================================================= def train(method, seed, steps=600, alpha=0.05, u=1e-2): """ 训练一个 MoE 层（8 个专家，top-1 路由）。 【参数】 method : "none"（不均衡）/ "aux"（辅助损失）/ "bias"（偏置法） alpha : 辅助损失的系数 u : 偏置法每步的调整量 【返回】dict：任务损失、各专家负载 """ sample = make_task() torch.manual_seed(seed) moe = MoE(D, NUM_EXPERTS, k=1, d_ff_expert=64) moe.use_balance_bias = (method == "bias") opt = torch.optim.Adam(moe.parameters(), lr=3e-3) g = torch.Generator().manual_seed(1000 + seed) for _ in range(steps): x, y = sample(512, g) pred = moe(x) task_loss = ((pred - y) ** 2).mean() loss = task_loss if method == "aux": loss = loss + alpha * switch_aux_loss(moe.last["probs"], moe.last["idx"], NUM_EXPERTS) opt.zero_grad() loss.backward() opt.step() if method == "bias": update_balance_bias(moe, u) with torch.no_grad(): x, y = sample(4096, torch.Generator().manual_seed(99)) pred = moe(x) return {"loss": ((pred - y) ** 2).mean().item(), "load": moe.last["load"].tolist()} def demo_compare(): """演示 1：几种方法的比较（每种方法用 4 个随机种子）。""" print_header("演示 1：不均衡 vs 辅助损失 vs 偏置法（8 个专家，top-1）") print(f" 理想情况：每个专家分到 1/8 = {1 / NUM_EXPERTS:.3f} 的 token") print(" （'饿死'= 分到的 token 少于 1%；每种方法跑 4 个随机种子）\n") methods = [ ("不做均衡", "none", {}), ("辅助损失 α=0.05", "aux", dict(alpha=0.05)), ("辅助损失 α=0.2", "aux", dict(alpha=0.2)), ("DeepSeek-V3 偏置法 u=0.01", "bias", {}), ] print(f" {'方法':<22}{'各种子的最大负载':<26}{'各种子饿死的专家数':<16}{'平均任务损失'}") for label, method, kw in methods: results = [train(method, seed, **kw) for seed in range(4)] maxes = [max(r["load"]) for r in results] deads = [sum(v < 0.01 for v in r["load"]) for r in results] loss = sum(r["loss"] for r in results) / len(results) print(f" {label:<24}{str([round(m, 2) for m in maxes]):<32}{str(deads):<22}{loss:.4f}") print(""" - 不做均衡：有时运气好、自然就均衡了（种子 0）；但多数种子里最忙的专家承担了约 1/4 的 token， 还有专家几乎分不到 token（饿死）。饿死的专家参数完全浪费，所以任务损失也更差； 在专家并行中，最忙的那块 GPU 还决定了整体速度。 - 辅助损失：α 太小（0.05）几乎不起作用；α 调到 0.2 后负载变得均匀。它的麻烦在于 α 需要调： 太小没用，太大又会和任务目标"抢"梯度（大模型实践中的常见问题）。 - 偏置法：不改变损失函数，只在"选专家"时加一个随负载反馈调整的偏置，同样让负载变得均匀， 任务损失与调好的辅助损失相当。它不会干扰模型学习任务 —— 这是 DeepSeek-V3 采用它的主要理由。 【边界】这是一个很小的玩具实验（每种方法 4 个随机种子，结果有波动），它能说明"不均衡会发生、 两种方法都能修正"，但区分不出两种方法在大模型上的优劣。真实大模型中负载不均衡的后果 （通信、显存、训练崩溃）严重得多；另外，大多数 MoE 同时还会让路由器用 FP32 计算、并加上 z-loss 保证数值稳定。""") def main(): demo_compare() print("\n全部演示运行完毕。") if __name__ == "__main__": main()