专家负载不再只求全局平均：精确分位数加局部误差注入，最忙专家从6.4倍降到4.5倍
核心概要
稀疏混合专家模型里，路由若偏爱少数专家，其余专家拿不到足够训练信号，专家并行时还会让设备互相等待。Kolibri 预训练采用精确分位数均衡（EQB）与有界负载误差注入（LEI）：前者用两趟基数选择在全批次上精确求出路由偏置，后者把每个微批次的负载误差直接注入路由梯度。在7.5B参数代理模型上训练100B token，EQB把全局最大相对过载从0.92降到0.74，再加有界LEI把局部过载从5.38降到3.52，代价是平均BPB从0.6570升到0.6720。
深度剖析
负载均衡被拆成两个尺度分别处理：整个优化器步的全局均衡保证每个专家都有足够梯度信号，单个微批次内的局部均衡避免专家并行执行时某台设备成为瓶颈。 此前的做法要么用辅助损失让路由通过梯度学均衡，要么在梯度之外按固定步长调整每个专家的路由偏置，后者不显式计入被挤掉的分配，只能靠小步长逐步逼近均衡。 在7.5B参数、每token激活0.5B参数的代理模型上训练100B token，EQB把全局最大相对过载从0.92降到0.74，并在六个评测基准上取得比按rank平均分位数的QB更低的平均BPB。
EQB把分位数均衡看成阈值路由：固定每个token的偏置后，每个专家的负载可以独立调整，因此偏置可以直接由分位数设定，不需要更新速率。 按rank分别算分位数再平均更省通信，但一般得不到全局分位数，结果还依赖token如何切分；累积直方图的做法通信量与token数无关，但分位数精度受直方图箱宽限制。 EQB把BF16的token调整后logit编码为保序的16位键，用两趟基数选择先定位高字节、再解析低字节，每层每次优化器步的归约数据量为0.5 MiB，48个MoE层合计24 MiB，与token数无关。
LEI不走辅助损失，而是把每个专家在本地微批次上的负载误差直接加到其sigmoid分数的梯度上，过载专家的分数被压低、欠载专家的分数被抬高。 GShard式辅助损失用跨专家归一化的分数作为离散负载的可微替代，改变一个分数会牵动同一token的所有归一化分数；LEI用未归一化的sigmoid分数，每个分数只收到自己专家的负载误差。 无界修正会让注意力log-sum-exp急剧上升，按共享因子整体缩放后保持稳定；有界LEI把局部最大相对过载从5.38降到3.52，最忙专家从6.4倍均分降到4.5倍，平均BPB从0.6570升到0.6720。
启示与展望
讨论范围是Kolibri预训练所用的负载均衡配方，以及在7.5B参数、每token激活0.5B参数、48个MoE层、每层256个路由专家与6个激活专家的代理模型上、100B token的六组对照。评测覆盖MMLU、ARC、HellaSwag、HumanEval、MBPP与GSM8K，报告六项基准的平均BPB；注意力LSE只作稳定性检查，不是分数。是否采用有界LEI取决于局部不均衡在给定专家并行设置中对吞吐的限制有多大。
EQB消除的是分布式分位数近似，单次QB更新本身的近似与BF16量化仍然存在；偏置由当前批次算出却用于下一批次，而下一批的logit因数据和路由参数变化而不同，因此下一批的专家负载只是近似均衡。EQB的精确性以内存为代价：全分辨率直方图在256个专家时每层占64 MiB，与批大小无关，而Kimi的1000箱直方图不到1 MiB。有界LEI在改善局部均衡的同时使平均BPB上升0.015，且EQB单独使用时平均BPB最低。