Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts
Nayeon Kim, Hojin Lee, Yunju Bak, Jaesun Park, Boseop Kim
COLM 2026
cs.LG, cs.AI, cs.CL
2026-08-20
Kakao把μP宽度迁移和token维线性外推接成两步,用约四分之一宽度的小MoE扫学习率,预测155B总参、17B激活、10T预训练最优学习率为3.85×10⁻⁴(R²=0.95),全程无loss尖刺。
大 MoE 预训练里,学习率对模型宽度和 token 预算都敏感。宽度一变、数据一加长,最优学习率就跟着走。按常规做法,模型和 token 两维一起扫,百亿参数以上根本付不起。μP 本来是给 dense 模型做宽度零样本迁移的,MoE 还多一条稀疏轴:总专家数增加、激活专家数固定。已有 MoE 上的 μP 试验多半把专家数钉死、只扩 hidden,跟现在细粒度专家(一百多个)的扩法对不上。
Kakao 和 Upstage 把问题拆成两步:先让最优学习率在宽度(含总专家数)上可迁移,再用小模型短预算把学习率外推到万亿 token。
架构是 MLA MoE,优化器是 Muon。μP 按「宽度无穷时有几维一起膨胀」给参数分类:embedding、bias、expert 的 FC2 算 vector-like,只改初始化;FFN、attention、router、expert 的 FC1 算 matrix-like,初始化和学习率都按 faninbase/fanin 缩放。深度固定,head dim 固定,头数跟 hidden 一起涨。扩容时激活专家数和 expert 中间维钉死,只加总专家数和 hidden。稀疏变了,单个 expert 的 fan-in/fan-out 并不额外变,μP 规则还能用;代理如果直接拉到目标那么稀,算术强度太低,搜索本身会很亏。
Token 维不直接扫到 10T。代理跑 Warmup-Stable-Decay 的稳定段,不做 decay,用 EMA(α=0.6)每隔约 2B token 合一次权重,每 10B 取一个点。每个预算上,验证 loss 对 log 学习率拟合二次曲线,顶点当最优学习率。再在 log(η) 和 log(token) 上做线性回归。batch size 不当迁移对象:它同时是系统吞吐旋钮,文献里对它怎么随算力或 token 缩放也互相打架。实验里 batch 按 GPU 效率钉死,学习率单独外推。拟合只用 batch 加大之后、动力学稳住的点(255B 之后)。
标准参数化下,0.6B 总参、0.3B 激活的基模型最优学习率迁不到 2×、4×、8× 宽度(对应 2.2B/0.7B、8B/1.5B、30.7B/3.6B)。μP 下同一最优值能穿过这三档。dense MLA 上同样成立。
5.6B/1.8B 代理和 2× 的 20.7B/3.8B 在 40B–100B token 上,二次曲线顶点和曲率几乎重合,最优学习率随 token 略往下走。真正外推用 10.8B/3.3B 代理(目标 155B/17B 的 1/4 宽度)训约 500B token:255B–502B 的 log-log 拟合 R²=0.95,外推 10T 得到 3.85×10⁻⁴。用 255B–350B 拟合、在约 500B 留出验证,预测与实测最优学习率平均差约 4.4%。
| 设置 | 指标 | 结果 |
| 10.8B 代理、255B–502B | log-log R² | 0.95 |
| 外推到 10T | 最优学习率 | 3.85×10⁻⁴ |
| 约 500B 留出 | 预测相对实测 | 平均偏差约 4.4% |
| 相对只跑代理 | 再扫 1.5×/2× 宽度 | 额外 240.3 ZFLOPs(代理本身 64.8) |
| 155B 目标相对代理 | 总算力 | 约 98× |
目标模型 155B 总参、17B 激活,Stage 1 训 10T,200B 后 batch 从 8M 提到 32M。loss 全程无尖刺。数据配比 6T 前是英文 45%、Math/STEM 12.5%、代码 27.5%、多语 15%,之后改成 22.5%/27.5%/25%/25%。同一套 harness 下,MMLU-Pro 约 64%,高于 dots.llm1 和 GLM-4.5-Air,低于 Hunyuan-A13B 和 DeepSeek-V4-Flash,落在 Pareto 前沿上。论文没给出 Figure 7 各棒的精确读数。
百亿以上 MoE 预训练,学习率几乎没法在目标尺度上网格搜索。这套流程把二维扫描压成「μP 过宽度、小模型沿 token 外推」,代理搜索大约是目标预训练的 1/98。对正在上 MLA 加 Muon 的团队,可直接抄参数分类和 EMA 取点办法。它是工程上能落地的渐进改进,没有提出新的优化器。
作者写明:框架绑在 MLA 和 Muon 上,换结构或优化器要重做。稀疏和宽度一起扩,分不清迁移到底来自哪一条轴。top-k 路由下各 expert 看到的有效 batch 不同,按 expert 调学习率没做。10T 目标尺度不可能再扫一遍学习率来证明「最优」,稳定 loss 和竞品对比只能说明「能用」。Stage 1 数据偏多样性、质量过滤不严,benchmark 不能当最终模型成绩读。batch size 被刻意排除,换硬件吞吐之后学习率外推还成不成立,这篇没有回答。