""" 混合专家：不加干预的路由器会自己均匀分工吗？在你自己的电脑上量一遍。 跑一个随机初始化的路由器：256 个专家，每个输入选 8 个。 先看不做任何干预会发生什么，再加上一个最小的负载均衡，看分布怎么变。 最后换几个种子再跑一遍，看固定例是不是特例。 先说清楚这个脚本**不是**什么： - 路由器是**随机初始化、没有训练过**的。真实模型的路由器是训练出来的， 分布不一定长这样。这里演示的是「不加干预时会不会自己均匀」这个机制。 - 它只统计「谁被选中」。**没有执行任何专家网络，没有训练模型， 也没有测量吞吐、延迟或显存。** - 均衡用的偏置是按论文提到的思路写的一个最小实现，**不等同论文里的完整算法**， 更不是复现了 DeepSeek-V3。 - 一个种子是一个固定示例。脚本最后会换 8 个种子给出范围，但 8 个样本 也不能证明所有种子、所有模型都会这样。 只需要 python3，不用装任何东西： python3 moe_routing.py # 默认就跑这个，什么参数都不用给 python3 moe_routing.py --help # 想换配置（种子、专家数、选几个、输入数）时再看 默认配置在这台机器上约 8 秒跑完（8 个种子各跑两遍）。 可复现到哪一步，是有边界的：脚本没有计时，同一个解释器跑两遍逐字节相同； 作者又在一台 arm64 macOS 上逐个试过 6 个 CPython（3.9.6、3.10.21、3.11.16、 3.12.14、3.13.15、3.14.7），39 行输出同样逐字节相同。但这不等同于任何机器上都一样： 打分那一步是浮点求和，CPython 3.12 起 sum() 改用了补偿求和，同一个输入下 256 个专家里有 158 个的分数末位已经不同，只是这个配置下排序没被这点差异改变。 换一种 CPU 架构没有试过。 固定例的数字与路由的守恒、范围、排序策略 由 tests/demo/test_moe_routing.py 复跑核对。 """ import argparse import random import unicodedata N_EXPERTS = 256 TOP_K = 8 DIM = 16 N_TOKENS = 2000 BIAS_STEP = 0.02 # 负载均衡每轮调整的幅度 BIAS_EVERY = 20 # 每多少个输入调一次偏置 DEFAULT_SEED = 0 DEFAULT_PROBE_SEEDS = 8 # 多种子探查用几个种子（含固定例那个） # 这四句是这个实验的边界。它们同时出现在脚本输出和 demo 页面上， # 因为只看一眼结果的读者最容易把「入选次数」读成「算力」或「显存」。 SCOPE_NOTES = ( "路由器是随机初始化的，没有训练过", "只统计「谁被选中」：没有执行专家网络，没有训练，也没有测吞吐、延迟或显存", "均衡用的偏置是按论文思路写的最小实现，不等同论文里的完整算法", "换一个种子结果就会变一点，所以下面额外跑了多个种子给出范围", ) def build_gate(rnd, n_experts=N_EXPERTS, dim=DIM): """路由器的打分矩阵。随机初始化，没有训练过。""" s = dim**-0.5 return [[rnd.gauss(0, s) for _ in range(dim)] for _ in range(n_experts)] def build_tokens(rnd, count=N_TOKENS, dim=DIM): """输入向量。一个「输入」对应真实模型里的一个 token，不是一个汉字。""" return [[rnd.gauss(0, 1) for _ in range(dim)] for _ in range(count)] def route(gate, x, bias=None, top_k=TOP_K): """路由器给每个专家打分，选分数最高的 top_k 个。bias 只影响选谁，不影响加权。 相同分数时的策略是写明的，不依赖排序算法碰巧稳定： 先按分数从高到低，分数完全相同时按专家编号从小到大。 """ scores = [sum(w * v for w, v in zip(row, x)) for row in gate] ranked = scores if bias is None else [s + b for s, b in zip(scores, bias)] order = sorted(range(len(gate)), key=lambda i: (-ranked[i], i)) return order[:top_k] def run(gate, tokens, balanced, top_k=TOP_K, bias_step=BIAS_STEP, bias_every=BIAS_EVERY): """跑一遍，统计每个专家被选中的次数。不修改 gate，也不修改 tokens。""" n_experts = len(gate) load = [0] * n_experts bias = [0.0] * n_experts if balanced else None for step, x in enumerate(tokens, start=1): for e in route(gate, x, bias, top_k): load[e] += 1 if balanced and step % bias_every == 0: # 按论文提到的思路做的最小实现：给过载的专家降一点分， # 给冷落的专家加一点分，只改「选谁」，不改加权。 avg = sum(load) / n_experts for i in range(n_experts): if load[i] > avg: bias[i] -= bias_step elif load[i] < avg: bias[i] += bias_step return load def summarize(load, tokens_count, top_k): """把一次运行的入选次数压成几个带分母的指标。""" n_experts = len(load) total = sum(load) ranked = sorted(load, reverse=True) quarter = max(1, n_experts // 4) never = sum(1 for v in load if v == 0) return { "experts": n_experts, # 总入选次数应恒等于「输入数 × 每次选几个」，这条守恒由测试逐配置断言 "selections": total, "expectedSelections": tokens_count * top_k, "busiest": ranked[0], "idlest": ranked[-1], "neverSelected": never, "neverSelectedShare": never / n_experts, "top10Count": min(10, n_experts), "top10Share": sum(ranked[: min(10, n_experts)]) / total, "topQuarterCount": quarter, "topQuarterShare": sum(ranked[:quarter]) / total, # 平均份额：整数分配里不要求任何一个专家恰好取到这个数 "evenShare": total / n_experts, "topQuarterEvenShare": quarter / n_experts, "busiestOverEven": ranked[0] / (total / n_experts), } def build_report( seed=DEFAULT_SEED, n_experts=N_EXPERTS, top_k=TOP_K, dim=DIM, n_tokens=N_TOKENS, bias_step=BIAS_STEP, bias_every=BIAS_EVERY, probe_seeds=DEFAULT_PROBE_SEEDS, ): """跑完固定例与多种子探查，返回一份纯数据的结果。渲染在 render_text 里。""" if not 1 <= top_k <= n_experts: raise ValueError(f"每次选中的个数要满足 1 ≤ K ≤ 专家数，收到 K={top_k}、E={n_experts}") if bias_every < 1: raise ValueError(f"调偏置的间隔至少是 1 个输入，收到 {bias_every}") if probe_seeds < 1: raise ValueError(f"多种子探查至少要 1 个种子，收到 {probe_seeds}") def one_seed(s): rnd = random.Random(s) gate = build_gate(rnd, n_experts, dim) tokens = build_tokens(rnd, n_tokens, dim) return { "seed": s, "unbalanced": summarize( run(gate, tokens, False, top_k, bias_step, bias_every), n_tokens, top_k ), "balanced": summarize( run(gate, tokens, True, top_k, bias_step, bias_every), n_tokens, top_k ), } # 探查从固定例那个种子开始往后数，所以第一个就是固定例本身，不用重跑 seeds = list(range(seed, seed + probe_seeds)) runs = [one_seed(s) for s in seeds] fixed = runs[0] def span(arm, key): values = [r[arm][key] for r in runs] return {"min": min(values), "max": max(values)} return { "config": { "seed": seed, "experts": n_experts, "topK": top_k, "dim": dim, "tokens": n_tokens, "biasStep": bias_step, "biasEvery": bias_every, }, "scopeNotes": list(SCOPE_NOTES), "fixed": {"unbalanced": fixed["unbalanced"], "balanced": fixed["balanced"]}, "probe": { "seeds": seeds, "runs": runs, "spans": { arm: { key: span(arm, key) for key in ( "busiest", "idlest", "neverSelected", "topQuarterShare", "busiestOverEven", ) } for arm in ("unbalanced", "balanced") }, }, } def _wide(text): """东亚宽字符按两列算，好让表头和数字列在等宽字体里真的对齐。""" return sum(2 if unicodedata.east_asian_width(ch) in "WF" else 1 for ch in text) def _pad(text, width): """按显示宽度左对齐。`f"{s:<8}"` 数的是字符数，中文表头会对不齐。""" return text + " " * max(0, width - _wide(text)) def _rpad(text, width): """按显示宽度右对齐。""" return " " * max(0, width - _wide(text)) + text def render_text(report): """可读摘要。所有数字都来自同一份 report，不重跑。""" cfg = report["config"] fixed, probe = report["fixed"], report["probe"] even = fixed["unbalanced"]["evenShare"] lines = [ "混合专家路由：不加干预的路由器会自己均匀分工吗", f"配置 {cfg['experts']} 个路由专家 · 每个输入选 {cfg['topK']} 个 · 一共 {cfg['tokens']} 个输入" f" · 打分维度 {cfg['dim']} · 随机种子 {cfg['seed']}", f"均衡策略 每 {cfg['biasEvery']} 个输入调一次偏置，步长 {cfg['biasStep']}；偏置只影响选谁，不影响加权", f"选中规则 分数从高到低取前 {cfg['topK']} 个；分数完全相同时按专家编号从小到大", f"平均份额 {cfg['tokens']} × {cfg['topK']} ÷ {cfg['experts']} = {even:.1f} 次/专家" "（这是平均数，整数分配不要求谁恰好取到这个数）", f"这个实验 {report['scopeNotes'][0]}", *[f" {note}" for note in report["scopeNotes"][1:]], ] for arm, title in (("unbalanced", "不做任何干预"), ("balanced", "加上最小的负载均衡")): s = fixed[arm] lines += [ "", f"【种子 {cfg['seed']} · {title}】", f" 最忙的专家被选中 {s['busiest']:>5} 次 平均份额 {s['evenShare']:.1f} 次" f"（{s['busiestOverEven']:.1f} 倍）", f" 最闲的专家被选中 {s['idlest']:>5} 次", f" 从没被选中的专家 {s['neverSelected']:>5} 个 占 {s['experts']} 个专家的" f" {s['neverSelectedShare']:.2%}", f" 最忙的 {s['top10Count']} 个专家占了 {s['top10Share']:>5.1%} 的入选次数" f" 分母 {s['selections']:,} 次 = {cfg['tokens']} × {cfg['topK']}", f" 最忙的 {s['topQuarterCount']} 个（{s['topQuarterEvenShare']:.0%}）占了" f" {s['topQuarterShare']:>5.1%} 均匀分配时应为 {s['topQuarterEvenShare']:.1%}", ] seeds = probe["seeds"] spans = probe["spans"] lines += [ "", f"【换 {len(seeds)} 个种子（{seeds[0]}~{seeds[-1]}）再跑一遍：固定例是不是特例】", " 每个种子重新初始化路由器和输入，其余配置不变。下面是这几个种子的最小值 ~ 最大值。", " " + _pad("指标", 24) + _rpad("不做任何干预", 18) + _rpad("加上负载均衡", 18), " " + "-" * 60, ] rows = ( ("最忙专家的入选次数", "busiest", lambda v: f"{v:,} 次"), ("最闲专家的入选次数", "idlest", lambda v: f"{v:,} 次"), ("从没被选中的专家", "neverSelected", lambda v: f"{v} 个"), ("最忙 25% 占的入选次数", "topQuarterShare", lambda v: f"{v:.1%}"), ("最忙专家是平均份额的", "busiestOverEven", lambda v: f"{v:.1f} 倍"), ) for label, key, fmt in rows: cells = [] for arm in ("unbalanced", "balanced"): lo, hi = spans[arm][key]["min"], spans[arm][key]["max"] cells.append(_rpad(f"{fmt(lo)} ~ {fmt(hi)}", 18)) lines.append(" " + _pad(label, 24) + "".join(cells)) ub, ba = spans["unbalanced"], spans["balanced"] lines += [ "", f"在这个玩具配置的 {len(seeds)} 个种子里，不加干预时最忙的专家拿到" f" {ub['busiest']['min']:,} ~ {ub['busiest']['max']:,} 次，" f"是平均份额 {even:.1f} 次的 {ub['busiestOverEven']['min']:.1f} ~" f" {ub['busiestOverEven']['max']:.1f} 倍；", f"加上这个最小的均衡后落到 {ba['busiest']['min']:,} ~ {ba['busiest']['max']:,} 次。", "所以「专家会自己均匀分工」在这个配置下不成立，训练里要专门处理负载均衡。", "反过来也要说清楚：不加干预时最忙的 25% 拿到的是" f" {ub['topQuarterShare']['min']:.1%} ~ {ub['topQuarterShare']['max']:.1%}，" "接近一半，不是「几乎全部」；", f"{len(seeds)} 个种子也只是 {len(seeds)} 个样本，不能推广成所有 MoE、所有种子、" "所有真实模型都会这样。", ] return "\n".join(lines) def _positive(name): def parse(text): value = int(text) if value < 1: raise argparse.ArgumentTypeError(f"{name}至少是 1，收到 {value}") return value return parse def main(argv=None): parser = argparse.ArgumentParser( description="混合专家路由的玩具实验：随机初始化的路由器会不会自己均匀分工", epilog="不带参数直接跑就是页面上那一份默认配置。", ) parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="固定例用的随机种子") parser.add_argument( "--experts", type=_positive("专家数"), default=N_EXPERTS, help="一层有几个路由专家" ) parser.add_argument( "--topk", type=_positive("每次选中的个数"), default=TOP_K, help="每个输入选几个专家" ) parser.add_argument("--dim", type=_positive("打分维度"), default=DIM, help="打分向量的维度") parser.add_argument( "--tokens", type=_positive("输入数"), default=N_TOKENS, help="一共过多少个输入" ) parser.add_argument("--bias-step", type=float, default=BIAS_STEP, help="均衡偏置每次调整的幅度") parser.add_argument( "--bias-every", type=_positive("调偏置的间隔"), default=BIAS_EVERY, help="每多少个输入调一次偏置", ) parser.add_argument( "--probe-seeds", type=_positive("探查种子数"), default=DEFAULT_PROBE_SEEDS, help="多种子探查用几个种子（含固定例那个）", ) args = parser.parse_args(argv) try: report = build_report( seed=args.seed, n_experts=args.experts, top_k=args.topk, dim=args.dim, n_tokens=args.tokens, bias_step=args.bias_step, bias_every=args.bias_every, probe_seeds=args.probe_seeds, ) except ValueError as exc: parser.error(str(exc)) print(render_text(report)) return 0 if __name__ == "__main__": raise SystemExit(main())