AI News字数 3249阅读时长9 分钟

Marin 开始公开训练 535B 参数 MoE 模型

Marin 已开始一项公开的 535B 参数 MoE 训练任务,使用 18.75 万亿个词元,并公开代码、日志与规模扩展预测。

一、一次公开进行的 535B 参数训练

Marin 已开始训练 Marin 535B-A23B,这是一款混合专家语言模型,总参数量约为 5350 亿,每个词元激活 230 亿参数。该项目计划在约三个月内使用 11 套 NVIDIA GB200 NVL72 系统处理 18.75 万亿个词元。

已公布的计划将 80% 的词元预算分配给预训练,20% 分配给中期训练。随后将进行后训练,但 Marin 尚未公布最终的后训练方案或时间表。该项目估计,主训练任务约需 \(2.7 \times 10^{24}\) 次浮点运算。

每套 GB200 NVL72 都是一种机架级系统,包含 72 块 Blackwell GPU 和 36 颗 Grace CPU。11 套此类系统的硬件规模相当于 792 块 GPU,尽管 Marin 的专家并行实现报告描述的是每个训练机架内一个由 64 块 GPU 组成的专家并行域。

此次训练并非模型发布。截至 8 月 24 日,训练仍在进行,因此尚无最终权重、基准测试结果或后训练评估可供评判。它眼下的重要性在于,团队选择在这一规模的训练进行期间公开其过程,而不是仅在筛选出成功结果后才发布技术报告。

Marin 针对这次训练的公开 issue 于 8 月 18 日创建。其中包括运行计划、工程风险、规模扩展方法、上下文扩展提案和应急程序。一个关联的 Weights & Biases 报告提供了该项目的实时追踪界面。

公告明确写明为 18.75 万亿个词元,而 GitHub issue 标题将此次训练简写为“18T tokens”。更精确的数字出现在公开的训练计划中;该项目并未将较短标题表述为修订后的预算。

二、稀疏模型如何组织

Marin 535B-A23B 使用 48 个 Transformer 块。每个块都将注意力分支与稀疏 MoE 分支结合,后者包含 384 个路由专家。路由器为每个词元选择 8 个专家,另有 2 个共享专家保持本地运行,并独立处理每个词元,不经过路由路径。

模型状态宽度为 6,144 个值。在路由激活值于 GPU 之间交换之前,潜在投影会将其压缩至 3,072 个值。Marin 表示,这使通过专家并行全对全操作的激活通信宽度减半。输出会在与共享专家路径合并前投影回模型宽度。

这一传输问题的规模不小。384 个路由专家分布在一个由 64 块 GPU 组成的专家并行域中,每块 GPU 承载 6 个路由专家。若为每位专家分别发送大小不同的缓冲区,将带来棘手的内存需求和动态通信模式。

因此,Marin 为 JAX 和 XLA 开发了固定池化波次全对全实现。发送端为每个目标 GPU 创建一个固定池,而非为每位专家创建一个缓冲区。传输以三个连续波次进行,数组形状相同,专家标识符被打包进激活值负载中。这避免了单独交换词元计数或路由元数据。

该实现采用两项容量限制。发送端容量系数为 1.10,限制从一个源端发送至一个目标端的通信量;接收端容量系数为 1.15,限制分配给单个本地专家的行数。超过任一固定缓冲区容量的分配都会被丢弃并单独报告。

一次 20 步、单机架的准入测试在未发生内存溢出的情况下完成,并在第 2 步至第 19 步期间记录到每秒 250,691 个词元的中位吞吐量。Marin 明确警告,这项短测试并不能确定最终的词元丢弃率。它是一项工程资格验证结果,而非训练质量基准,也不是对完整 11 机架训练的测量。

该报告还记录了失败的配置。直接固定专家单元设计产生了 192.65 GiB 的 XLA 内存估算,并在一次 123.49 GiB 的 CUDA 分配时失败。一个由 6 位专家组成的接收端池也发生了内存溢出,而将工作拆分为 3 个波次则使该池每次只有一部分处于活跃状态。这些负面结果是该项目公开设计记录的一部分。

三、主训练前先完成了规模扩展阶梯

在启动 535B 模型之前,Marin 训练了一个由四级组成的规模扩展阶梯。它从一款总参数量为 16 亿、激活 6100 万参数、在 480 亿个词元上训练的 MoE 模型开始。最大一级包含 277 亿总参数,激活 12 亿参数,并处理了 9260 亿个词元。

这一级阶梯既用于预测,也作为诊断参考。Marin 可以将主训练任务的损失、梯度范数、词元丢弃和评估轨迹,与较小规模下观察到的模式进行比较。若出现显著偏差,便可在完整训练消耗数月算力之前触发调查。

据该项目介绍,这一阶梯的成本约为主训练任务算力的 1%。Marin 将这笔开支视为一种风险控制:它测试所选架构、数据混合、优化器设置和训练时长是否会随模型规模增大而保持一致的表现。

该项目表示,早期阶梯曾发现,随着词元训练周期延长,梯度范数会增长至 4 以上。这一发现促成了 logit z-loss 的采用。后续消融实验据称表明,某些大批量配置若不采用这一方法,可能会在训练过程中发散。

较小规模的训练并不能保证 535B 模型会遵循其预测。外推仍是该实验的核心不确定性之一。它们的实际价值在于,Marin 已公开了一项基线,外部观察者可据此评估约 100 天训练期间作出的干预。

Marin 还记录了针对基础设施延迟或模型 FLOP 利用率低于预期的应急方案。在词元预算约前四分之一期间,默认响应将是缩短词元训练周期、调整数据混合,并重新安排线性学习率衰减,使其仍能在修订后的终点降至峰值学习率的 5%。这意味着已公布的时间表仍是一项运行计划,而非不可变更的规格。

四、长上下文取决于解决词元丢弃问题

模型以 4,096 个词元的序列长度开始预训练。Marin 上一次大型训练从 8K 开始,扩展至 65K 并训练 1 万亿个词元,之后还计划进行 262K 的扩展。回到 4K 会使每个批次包含的独立序列数达到 8K 时的两倍,这应能让词元更均匀地分配给各专家。

这一选择旨在解决当前专家并行实现的一项弱点。在早期测试中,词元丢弃率从 4K 上下文时约 7% 上升至 65K 时约 40%。Marin 报告称,其较新的池化波次设计在 4K 时的丢弃率约为 3%,但预计在 65K 时该比例仍可能过高。

被丢弃的分配并不一定意味着整个词元消失。两个共享专家仍会处理每个词元,且选中的 8 个路由专家是共享路径之外的额外处理。Marin 表示,当路由分配超过可用容量时,共享专家提供了更密集的骨干路径,但较高的丢弃率仍可能削弱稀疏专家的效益,并改变训练行为。

该项目计划在训练开始约 10 至 20 天后进行一次为期一到两天的早期冷却。这一分支旨在为强化学习实验提供一个全规模检查点,并测试更长上下文会如何影响路由,同时不改变主训练轨迹。

若该实验稳定,暂定的上下文计划将在训练过半时从 4K 升至 8K,在约 95% 的训练进度时从 8K 升至 65K,并在接近末期进入目标为 262K 的阶段。这些是有条件的目标,而非已完成模型的已确认能力。

如果词元丢弃率仍然过高,Marin 列出了三种替代方案:采用无丢弃的不规则全对全实现,提高容量系数并接受随之而来的内存成本,或引入序列级均衡。最后一种方案可能迫使专家在训练过程中重新专门化,因此该项目目前将其视为后备方案。

五、此次训练为何构成开放开发的里程碑

Marin 将开放开发与训练结束后发布权重区分开来。其标准工作流程始于一个 GitHub issue,用于记录实验的假设和目标。实现以可审查代码的形式提交,执行过程关联到公开遥测数据,分析结果——包括失败尝试——会回填至该 issue。

对于这次 535B 训练,公开记录现已包括主训练 issue、实时追踪报告、源代码仓库,以及对专家传输实现的详细说明。该传输报告将测量结果与工程判断区分开来,并说明哪些结论仅建立在短时性能分析运行之上。

这份记录让研究人员能够评估规模扩展预测能否经受完整模型的检验。它也公开了最终模型卡通常会压缩成几行的决策:为何训练从 4K 上下文开始、词元分配会在哪里被丢弃、内存限制如何改变传输设计,以及团队计划在训练进度落后时采取什么措施。

开放文档并不会使复现主实验变得成本低廉。复现一次 11 机架训练仍超出大多数独立研究人员的资源范围。更易获得的产物包括可在较小规模下审查或测试的代码、较小规模扩展训练、失败分析、配置选择和实时测量数据。

因此,该项目的状态必须被严格表述:Marin 已开始一项公开记录的前沿规模训练任务。它尚未证明最终模型的质量、长上下文行为、后训练表现或最终产物发布情况。

常见问题

“535B-A23B”是什么意思?

该模型总参数量约为 5350 亿,而每个词元约有 230 亿参数处于激活状态。稀疏路由会在每个 Transformer 块的 384 个路由专家中选择 8 个。

现在可以下载完成后的模型吗?

不可以。训练仍在进行,Marin 尚未发布此次训练的最终权重或后训练评估。

为什么训练仅以 4K 上下文开始?

4K 序列长度会让每个批次包含更多独立序列,从而改善专家均衡。Marin 正在测试更长上下文,因为早期实现曾在 65K 时出现大幅升高的词元丢弃率。

研究人员如何跟踪此次训练?

Marin 提供了公开的 GitHub issue、源代码仓库、专家并行设计报告,以及关联的 Weights & Biases 追踪页面。

262K 上下文长度是否得到保证?

不保证。向 8K、65K 和 262K 的拟议扩展取决于早期冷却和词元丢弃实验的结果。

参考来源

Share

分享这篇文章