OLMo-core 3跑通之后我最关心的反而不是榜单上那一两个百分点而是它的路由日志。作为把MoE从论文搬进训练框架的老手我早就习惯了盯着loss曲线过日子但这次不一样模型整体指标稳步上升十个专家里却有三个几乎专职收代码类token五个收自然语言QA剩下两个在到处打杂。社区里管这个现象叫token gerrymandering说得直白点就是路由器在偷偷画选区。这篇不是论文复述是我亲手把OLMo-core 3这类开源MoE模型的训练、评测、验收流程走了一遍之后确认必须补上的一个环节MoE路由要按时间窗验收。适合正在训练或微调MoE模型、负责模型质量验收、以及被路由负载问题折磨过的同学看。1. 我是在怎么发现token gerrymandering的1.1 指标正常路由分布却很不正常那次复现实验的配置我是照着OLMo-core 3的公开配置改的一共10个专家top-2路由容量因子1.25辅助负载均衡损失系数也不大。前两万步一切正常损失在掉benchmark在涨我甚至打算直接拿这个checkpoint去做下游任务。但因为我之前被专家塌方坑过一次养成了在每个MoE层后面挂一个路由统计钩子的习惯每100个微批次记录一次专家接收token的数量。某天早上我拉出来一看发现问题了某个中间层里3个专家吃掉了代码类数据集上超过70%的token而自然语言QA数据集上又是另外5个专家占主导。更离谱的是训练损失里那项负载均衡损失完全正常甚至比参考run还低一点。如果我只盯着loss曲线这个run会被我判为健康。但路由分布图摆出来那就是一幅被画得歪七扭八的选区地图。后来我在其他几个开源MoE模型上也看到了同样的模式。这不是偶发是路由网络在token空间里学出来的系统性行为。标准指标看不出来是因为它们统计的粒度和路由实际决策的粒度不一致。1.2 gerrymandering这个说法是怎么来的gerrymandering本来是政治学里的术语指的是通过重新划分选区边界来让某一方获得不正当优势。这个词在机器学习社区里被借过来形容一种特定的路由行为token就是选民专家就是选区路由器就是那个画选区的人。在MoE里路由器学到的不是一个简单的均衡分发而是一组把token空间切开、再分配给不同专家的边界。正常的专业化是好事MoE本来就希望专家各司其职。但问题在于当路由器把边界画得极度偏心时某些token类型的簇会被整体打包进某一个专家其他专家很少见到这类token。结果就是热门专家持续过载冷门专家长期饥饿专家之间的表示空间出现单议题选区式的割裂推理阶段一旦某个token簇的流量波动对应的专家立刻变成延迟瓶颈。这个词在社区讨论里已经被广泛使用我也沿用。它比负载不均衡更准确因为后者只描述了现象而gerrymandering点出了机制这是路由器主动学会的、利用概率与离散分发之间缝隙的一种操纵。1.3 为什么OLMo-core 3把这个问题摆到了台面上OLMo系列一直以开放性著称OLMo-core 3延续了这个传统开放权重、开放训练数据、开放完整配置。这种透明度让路由审计变成可能。以前我训闭源或者半开放模型路由行为只能靠推理结果反推等于隔着黑箱猜。现在配置文件和训练细节都摊在面前我可以把感觉路由有问题变成可复现、可量化、可验收。而且说实话MoE已经是当前扩大模型规模的主流路线从预训练到微调大家都在用。但有多少团队把路由健康度当成一个正式的验收门禁我看很少。大多数人还在用loss和benchmark说话。OLMo-core 3这波提醒的意义在于当一个模型连路由细节都开放给你看时你就没有借口不去看了。2. 先搞清楚MoE路由在做什么2.1 门控路由的基本流程先把基础对齐。在MoE Transformer里原本的FFN层被替换成N个专家FFN以及一个路由器。每个token进入这一层时路由器从隐藏状态里算出一组logit经过softmax变成对N个专家的概率分布。然后按top-k选专家k一般取1或2token被送到选中的专家里计算再把结果按照门控概率加权合并。这个路由器是跟着整个模型一起训练的。梯度的路径是token经过选中专家误差反传时既更新专家参数也更新路由器的参数。所以路由器本质上是在表示空间里学一组边界这个边界的形状决定了什么样的token会被划到什么样的专家手里。容量因子是另一个关键参数。它限制每个专家在一个批次里最多能接收多少token通常是理论均值的1.1到1.25倍。超过容量的token要么被丢弃要么走旁路残差直接传到下一层。这个机制是为了兜底但它恰恰是路由异常时最先出问题的地方。2.2 负载均衡损失为什么治标不治本为了防止路由塌缩训练时通常会加一个辅助负载均衡损失最常见的是Switch Transformer那版。公式是L_aux α · N · Σ f_i · P_i其中f_i是实际分发给专家i的token比例P_i是路由器分配给专家i的平均概率α是权重系数。这个损失的直觉是如果f和P都均匀乘积的均值才最小于是模型被鼓励把token均匀分给各个专家。但这个损失有一个结构性弱点P_i是所有token的平均概率。路由器可以用一种概率铺平的方式来糊弄它。举个例子假设有两簇tokenA簇90个B簇10个。路由器对A簇一律给专家1高分0.99对B簇一律给专家2高分0.99。算下来f_10.9f_20.1而平均概率P_1和P_2会因为两簇token的数量加权而显得没那么极端。整个损失值算出来并不高。问题在于f是经过top-k离散选择之后的结果而P是连续概率的平均值两个量之间隔着一个非线性选择操作。损失函数把两者乘在一起取均值等于只考察了边缘分布没有考察token到底去了哪。所以损失值低并不代表分发均衡。这就是为什么拿训练loss里的辅助项当验收指标一定抓不到gerrymandering。后来出现的router z-loss、expert-choice路由等改进能在一定程度上缓解但都没有改变一个事实你必须单独对路由分发结果做审计而不是指望辅助损失替你兜底。2.3 三种典型的钻空子路由行为我在实际日志里见过的gerrymandering大致可以归成三类每种都对应一种糊弄指标的手段。第一种叫概率铺平。路由器在平均概率上给每个专家都留了体面的数值但实际top-1选择时某一整簇token永远落到同一个专家头上。P_i的均值看起来平滑f_i却已经严重倾斜。第二种叫过度自信打包。路由器对某些常见token簇变得非常自信拍板把它们整批送进一个专家遇到稀有的、单发的token时它又表现得犹豫把概率分散到多个专家。这样算出来的平均P_i很光滑但专家收到的真实token量完全不是那么回事。第三种叫语义切割。前面两种还算单纯的计数不平衡这种更隐蔽来自同一个数据源、同一个语义簇的token被系统性划到同一个专家。它不是数量上有多悬殊而是专家变成了单领域选区。平时看不出异常一旦推理流量里这个领域的数据突然增多那个专家立刻变成热点容量因子扛不住token开始丢。3. 为什么验收必须按时间窗来3.1 瞬时指标的第一重谎言稳态不等于全时态如果你只在某一个step拉一批数据看路由分布你看到的只是token分布的一小片切片。数据加载器通常把语料切成shard并打乱一个批次可能恰好以某个领域为主。这一刻的专家负载看起来可能非常均衡下一刻另一个shard进来就完全翻转。gerrymandering是画在整张token地图上的边界不是某一小块区域的特征。想发现它必须让采样窗口长到足以覆盖token分布里的各个区域。单点快照或者说瞬时指标天然地抓不到这种全局性的边界扭曲。3.2 第二重谎言均值掩盖了周期性偏移我见过一种更容易误导的情况整个10k步跑完所有专家收到的token总数非常接近均值漂亮得很。但如果你把窗口切成500步一段会发现专家负载在剧烈摆动。代码类shard经过时代码专家吃满容量、其他专家近乎挨饿QA类shard经过时负载完全换边。这种周期性偏移危害很大因为训练和推理用的容量因子是固定的峰值负载才是决定token丢不丢、延迟高不高的关键。均值再好看撑不住峰值也是白搭。时间窗验收要求每个滑动窗口都过阈值而不是看整段run的平均值就是为了逼你把周期性偏移暴露出来。3.3 第三重谎言token分布本身在漂移预训练阶段经常有数据混比调整微调阶段任务分布更是剧烈变化。一个在step 100k时路由状态良好的模型到step 150k未必还是好的。更微妙的是评测时的流量分布和训练时的流量分布通常不一样。这意味着路由验收不能只在训练收尾时做一次。它应该是一个持续的过程按固定的时间窗在全训练周期里反复做。这样既能发现路由漂移的拐点也能在你改了数据配比之后立刻看到路由行为是否跟着出了问题。这里说的漂移不是指专家换了分工而是指边界本身变得不稳定或者开始朝gerrymandering的方向滑。3.4 时间窗验收到底验收什么时间窗验收的定义我建议这样定在任意一个滑动窗口内路由统计必须满足一组预定义的条件才算pass。不做平均值不做看心情就是把窗口当SLO一样卡。具体验收项包括负载均衡窗口内专家接收token的最大值与最小值之比低于阈值专家饥饿每个专家在窗口内收到的最低token占比不能低于某个下限稳定性相邻窗口之间各专家份额不要出现大幅摆动容量余量窗口内单专家峰值负载与容量因子的比值必须留出余量打包检查按token类型拆分统计归一化路由熵不能低于下限。这套定义把路由健康从一句主观感受变成了可执行的验收清单。下面我讲具体怎么落地。4. 落地实操从埋点到出报告4.1 路由决策日志怎么埋原则很简单不要存per-token级别的完整路由轨迹那是天文数字。要存的是聚合结果。我建议在每个MoE层的前向函数里挂一个钩子每个微批次记录一次专家接收token的计数向量和平均门控概率同时带上当前batch的数据源id。下面是一个PyTorch风格的示例class RouterLogHook: def __init__(self, num_experts, log_interval100): self.num_experts num_experts self.log_interval log_interval self.buffer [] def on_moe_forward(self, step, layer_id, dispatch_indices, gate_probs, batch_source): counts torch.zeros(self.num_experts, dtypetorch.long) counts.scatter_add_(0, dispatch_indices.view(-1), torch.ones(dispatch_indices.numel(), dtypetorch.long)) self.buffer.append({ step: step, layer: layer_id, source: batch_source, counts: counts.tolist(), prob_mean: gate_probs.detach().mean(0).tolist(), }) if len(self.buffer) self.log_interval: self.flush() def flush(self): # 追加写入JSONL或SQLite然后清空buffer passbatch_source这个字段很关键。没有它你只能看到总体的计数看不到某个数据域是不是被某个专家垄断了。日志频率我一般设在每100到200个微批次一次够用又不至于把磁盘打爆。训练期间注意观察一下IO开销实测这个钩子整体占用训练时间通常在1%到2%以内。4.2 核心指标怎么算拿到一堆窗口内的计数矩阵之后我一般算这几个指标import math import numpy as np def routing_summary(counts_matrix, num_experts, step_window): loads counts_matrix.sum(axis0) # 每个专家总token数 total loads.sum() frac loads / total # 各专家分发比例 lir loads.max() / loads.min() # 负载不均衡比 cv loads.std(ddof0) / loads.mean() # 变异系数 entropy -(frac * np.log(frac, wherefrac 0)).sum() norm_entropy entropy / math.log(num_experts) return { lir: lir, cv: cv, norm_entropy: norm_entropy, frac: frac, step_range: (step_window[0], step_window[-1]), }每个指标的含义要说清楚。LIR是最直观的过载指标max/min一拉就能看出来有没有专家饿死。CV是归一化之后的离散程度适合在不同专家数量下对比。归一化熵衡量分发有多分散1.0代表完全均匀越低代表越集中。光看熵不看LIR也不行熵对少量极端偏差不敏感两个指标配合着用。如果日志里有batch_source再按source分组算一遍同样的指标专门检查某个数据域是不是被单一专家垄断。这一步往往是发现semantic cutting的关键。4.3 时间窗参数怎么设窗口长度不是拍脑袋定的。核心原则是窗口必须覆盖数据分布的一个完整周期。先看你的数据加载器打乱逻辑如果语料切成shard、每个shard循环一遍大约要S步那窗口长度至少要5倍的S我习惯取10倍。训练阶段我常用2000到10000步的窗口微调阶段因为数据量小、分布变化快窗口可以压到200到500步。滑动步长取窗口长度的50%重叠。验收一个checkpoint时至少要跑3到5个完整窗口也就是说检查周期要覆盖训练步数的数倍。每次检查之后更新一次路由健康仪表盘不是等到训练结束才看。抽样频率上我上面说了日志每100到200个微批次记一次对2000步的窗口来说一个窗口里就有10到20个采样点足够算出稳定的统计量。采样点太稀疏会让窗口估计方差变大太密又浪费存储这个区间是平衡点。4.4 验收流程与阈值怎么定我推荐用两级门禁warning和fail。warning触发人工排查fail直接卡住checkpoint发布或者触发训练干预。下面是我在8到16个专家的模型上常用的初始阈值供参考验收项指标建议阈值初始值负载均衡LIRmax/min token数小于等于1.58专家小于等于2.032专家以上负载离散度CV小于等于0.25专家饥饿窗口内单专家最低token占比大于等于5%路由多样性归一化熵大于等于0.85峰值余量窗口内单专家峰值除以容量因子小于0.9留10%余量稳定性相邻窗口专家占比最大偏差小于10个百分点这些阈值是起点不是真理。不同数据分布、不同专家数量下都得重新调。调法很简单先跑一个你心里认为健康的参考run把它的路由统计分布打出来取分位数作为阈值基线。后续正式run按这个基线卡比凭空定阈值靠谱得多。验收动作我建议做成训练脚本里的一个固定环节每5000步自动触发一次检查最近N个窗口输出一份路由健康报告写进训练日志和模型工件目录。fail时不仅要有邮件告警最好在训练框架里把learning rate暂时冻结或者自动降低一阵给路由器一个冷静期重新调整的机会。4.5 顺手能用的监控组合工具上不需要搞多复杂的系统。我常用的组合是训练脚本里把路由指标直接打到TensorBoard或WB和loss画在同一个面板上路由日志落SQLite便于后面按数据源、按时间范围查询写一个简单的告警脚本任何一个窗口fail就发消息连续三个窗口fail就输出严重告警并建议暂停训练。有一点容易被忽略把路由健康度作为一行固定的字段加到每天和loss、accuracy一起看的评测仪表盘上。天天看着它你才会对路由行为的变化有感觉。否则日志埋了也等于没埋。5. 踩坑实录和排查技巧5.1 专家塌方一个专家吃掉四成token最典型的故障模式。某天看报告一个专家窗口内token占比飙到40%另外两个专家份额掉到2%以下LIR直接破5。排查顺序我一般这样走先看按数据源拆分的统计确认是不是某个shard被重复采样了。数据源没问题再看路由logit的量级。如果某个专家对应的logit均值明显高于其他专家说明路由器的参数更新太快已经跑偏了。我遇到过一次是辅助负载均衡损失系数设成了0.01太低了再加上路由器学习率比专家FFN高了好几倍等于路由器跟专家赛跑时一路狂奔。修复方式辅助损失系数提到0.1级别开router z-loss必要时直接重新初始化路由器头部再继续训练。修完之后再观察几个窗口确认不是暂时压下去。5.2 路由震荡同一批token反复横跳有时候负载指标看起来还行但相邻窗口之间某一专家的占比偏差能到15个百分点以上。这种震荡的危害是训练不稳定token在同一语义簇上的表示会在不同步数之间换来换去损失曲线容易反复横跳。震荡的常见原因有三个路由器学习率太高、微批次太小、门控概率方差太大。前两个都好办路由器学习率我通常压到专家FFN的1/5到1/3微批次往上提。第三个可以用router z-loss来压低logit的绝对值让概率分布更平滑。如果你用了expert-choice路由这个现象会少很多因为分配是专家主动选的天然更稳但也有别的代价评估时需要权衡。5.3 token丢包上升容量因子不够用了容量因子1.25的意思是每个专家最多接收理论均值的1.25倍token超过的会被丢弃或走旁路。如果训练日志里loss突然出现无规律尖峰推理测试里某个专家的延迟周期性飙升先去查窗口内单专家峰值负载。我见过的情况是窗口均值很健康但代码shard经过时某专家峰值冲到容量的1.15倍直接触发丢弃。这时候先别急着改模型先看数据加载器的shard顺序是不是过于集中。把shard打乱粒度改细之后峰值马上就回落。确实改不动数据的话再把容量因子从1.25提到1.35或1.5但注意这会增加计算量。5.4 常见问题速查表症状可能原因先查什么常用解法单专家占比异常高辅助损失权重低、路由器学习率高、数据shard重复按source拆分的路由统计、router logit均值调高aux权重、加z-loss、重置路由器、修正数据采样相邻窗口专家份额摆动大路由器LR过高、batch太小、门控概率方差大窗口间专家占比偏差、门控概率的标准差降路由器LR、增大batch、开z-losstoken丢包导致loss尖峰容量因子不足、shard集中导致峰值过载窗口内单专家峰值负载时间线细化shard打乱、提容量因子到1.35或1.5辅助损失正常但分布倾斜aux loss被概率铺平糊弄f_i和P_i的对比、按类型拆分熵改用分配感知型均衡损失、把路由输出加入验收项微调后路由剧烈变化任务数据分布与预训练差异大微调前后的路由统计对比微调加上router z-loss并锁定部分专家温度6. 最后分享一点实操心得自从我把时间窗验收加进MoE训练流程前后拦下了三次明显的问题release一次专家塌方一次周期性丢包一次微调后路由偏离。每一次如果只靠loss和benchmark大概率就这么混过去了。而这套东西的成本真的很低一个钩子加一个SQLite表隔几百步算一次统计量训练时间多花1%到2%。我个人的体会是MoE模型发布时应该把路由统计作为一个标准工件一起带上。就像发布微调模型要带eval结果一样发布MoE也该带上一份路由行为说明包括窗口内的负载分布、按数据域拆分的情况、稳定性指标。下游使用者看到这份说明才知道这个模型在什么数据分布下是健康的什么场景下可能会出问题。画选区的人不是我们是路由网络自己。我们能做的不是猜它想干嘛而是用足够长的时间窗死死盯住它别让它把地图画得太难看。