让我们逐步扩展:大规模 Mixture-of-Experts 的计算高效超参数迁移

TL;DR 用小型 MoE 代理模型结合 μP 宽度迁移与 token 外推，以 R²=0.95 精度预测万亿 token 预训练最优学习率，省去昂贵超参扫描。

问题

问题背景

当前大规模预训练的核心矛盾是：Mixture-of-Experts (MoE) 架构能在计算预算大致不变的前提下扩展模型容量，但超参数（尤其学习率）的调优成本随模型规模和 token 预算指数级上升。

现有方法局限

传统超参数搜索需要在目标规模上实际训练并扫描多组配置。对于百亿激活参数、万亿 token 的 MoE 模型，单次训练成本达数百万 GPU 小时，网格或贝叶斯搜索完全不可行。Maximal Update Parameterization (μP) 虽然能在宽度维度实现学习率迁移，但现有工作主要针对 dense 模型，对 MoE + Multi-head Latent Attention (MLA) + Muon 等现代组件缺少直接适配；且 μP 只解决宽度迁移，无法外推到极长 token 预算——最优学习率随训练 token 数呈现非平凡漂移。

为什么难 / 重要

大规模训练是一次性的、不可回退的；学习率选错会导致训练发散或收敛缓慢，浪费的算力无法挽回。MoE 的路由动态、专家负载不均衡等因素进一步放大学习率敏感性，使小规模经验难以直接适用。业界需要一种能在极小代理模型上可靠预测目标配置的方法，把稀缺算力留给真正训练。

行业类比

这类似于用小型风洞实验预测大型飞行器的气动特性——只是把“风洞”换成小规模 MoE 代理模型，目标换成 10T token 级别的最优学习率。

核心洞察

• 将 μP 从 Dense Transformer 扩展到 MoE + MLA + Muon，实现跨宽度学习率迁移。此前 μP 主要针对 Dense 模型与 Adam/AdamW 优化器，对 MoE 的稀疏门控、专家路由、MLA 的 latent attention 以及 Muon 的二阶优化特性的组合缺乏系统论证。本文证明最优学习率在 width-scaled MoE 上稳定迁移，使小模型筛选出的超参可直接用于大模型，省去昂贵的大规模扫描。
• 最优学习率随 token 预算呈可预测的线性外推关系，小代理模型可外推到 10T token 量级。传统 scaling laws 通常预测 loss 或数据/参数规模，本文直接对最优学习率与 token horizon 建立回归模型（R²=0.95）。结合宽度迁移，形成两步链：小预算小宽度 → 大宽度 → 巨量 token，使 155B/17B active 模型的从头训练无需学习率扫描，显著降低计算成本。

方法

方法框架

输入：小规模 MoE 代理模型在不同宽度和 token 预算下的训练配置及对应的最优学习率（通常通过小范围搜索获得）。

关键模块 1：宽度缩放迁移（μP 适配）

• 将 Maximal Update Parameterization (μP) 扩展至 MoE 架构，兼容 Multi-head Latent Attention (MLA) 和 Muon 优化器。
• 通过调整参数初始化和逐层学习率缩放，使不同宽度模型的最优学习率趋于一致，从而可将小宽度模型搜索到的最优学习率直接用于大宽度模型。

关键模块 2：token 维度外推（缩放定律）

• 在固定宽度下，训练多个小模型于不同 token 预算（如 0.1B、0.5B、1B tokens），得到各自的最优学习率。
• 观察到最优学习率与 token 预算之间呈稳定规律（近似幂律或线性），使用 线性回归 拟合该规律，并外推至 10 万亿 token。
• 外推结果与实际训练验证高度吻合（R^2 = 0.95），证明小预算代理训练足以决定大规模训练的学习率。

输出：目标模型（155B 总参数 / 17B 活跃）在 10T token 预训练中的最优学习率，即可直接用于生产训练。

跟同类方法的差异点：该方法不仅利用 μP 实现宽度迁移，还首次将学习率迁移扩展到 token 维度，通过缩放定律从短预算外推到极长预算，避免了传统超参数扫描的巨大开销。

实验

实验设计

采用两步 超参数迁移 框架。第一步针对 MoE + MLA + Muon optimizer 构建宽度缩放的小型代理模型，验证在 μP 参数化下最优学习率能否跨宽度保持一致。第二步在有限 token 预算（短 horizon）下估计各代理模型的最优学习率，再用线性回归外推到 10T token 级别。最终将预测的学习率用于自研 155B total / 17B active 基础模型从零预训练。

关键发现

• 宽度迁移：最优学习率在宽度缩放的 MoE 模型间迁移一致，证实 μP 适配有效。
• token 外推：基于短预算最优值做线性外推，对 10T token 的预测拟合优度 R²=0.95。
• 全尺度验证：预测的学习率支撑目标基础模型稳定训练，表明小模型代理训练足以确定大规模 MoE 的最优配置，消融成本极低。

与基线对比解读

传统做法是在目标规模模型上直接做学习率 grid search，算力成本随模型规模和 token 预算指数级增长，通常不可行。本方法将搜索转移到小模型和短预算，仅需极少消融实验即可完成外推。相比固定学习率或启发式选择，该方法提供了可预测、可复用的迁移路径；与单纯 μP 宽度迁移相比，进一步把预算维度纳入 scaling law，实现“宽度 + token”联合迁移。

行业影响

落地场景

该框架适用于需要从头预训练大规模 MoE 模型的产品团队，尤其当模型宽度或 token 预算超出常规调参能力时。典型场景包括：

• 电商推荐系统：训练数百亿参数 MoE 排序模型，需在有限 token 预算内确定学习率，避免全量扫描。
• 内容平台生成/理解模型：使用 MoE 降低推理成本，但训练超参优化困难。
• 金融风控大模型：企业自训 MoE 基座，希望以最小算力完成调参。

商业价值

直接降低训练成本：通过小代理模型预测最优学习率，可将超参搜索算力开销降低一个数量级。对千亿参数、万亿 token 训练任务，节省 GPU 时长转为显著财务收益。同时缩短从模型设计到上线的周期，支持更频繁的迭代实验。间接提升模型质量：稳定的最优学习率减少训练不收敛或性能次优风险，提升最终产品体验。

与现有工作流接口

该方法可嵌入现有大规模训练栈：

• 与 Megatron-DeepSpeed / FSDP 等分布式框架平行，作为前置的轻量级代理训练阶段。
• 使用 μP 和 Muon optimizer 时，可直接复用作者给出的 MoE 适配公式。
• 输出（最优学习率）可作为训练配置 YAML 中的 learning_rate 参数，无需改动主训练代码。

具体落地用例

1. 视频内容平台：训练一个 100B 活跃 10B 的 MoE 视频理解模型用于内容审核/标签，先在 1B 宽度模型上执行 100B token 的代理训练，通过线性回归外推至 10T token 的最优学习率，将扫描成本从数百万 GPU 小时降至数千小时。
2. 自动驾驶感知：车厂训练稀疏 MoE 多模态模型处理摄像头+LiDAR 数据，利用该框架在不同专家数量间迁移学习率，避免重复调参。

局限

• 论文中的学习率外推主要基于最优学习率与 token 预算之间的线性关系（可能是 log-log 线性？），但该关系在极大 token 尺度上是否持续成立尚未在完全目标尺度下直接验证，仅通过 R^2=0.95 的小规模代理实验间接支持。实际 10T token 训练中可能出现分布偏移或优化器动态变化，导致外推偏差。
• μP 适配针对 MoE + MLA + Muon 组合，但未覆盖其他常见 MoE 变体（如不同专家路由策略、不同注意力机制）或优化器（如 AdamW），通用性受限。实验仅在一个基础模型配置上验证，缺乏跨架构的鲁棒性证据。
• 论文只关注学习率这一超参数，而其他关键超参数（如批量大小、权重衰减、专家负载平衡损失系数等）的迁移与缩放规律未涉及，完整训练配置仍需额外成本确定，框架的降本效果局限于单一参数。