备注:回顾看过的论文,对目前看过的OLMo系列进行整理在此总结。(注:笔者水平有限,若有描述不当之处欢迎大家留言后期会继续更新LLM系列,文生图系列,VLM系列,agent系列等。如果看完有收获,可以【点赞】【收藏】【加粉】)

阐述的思维逻辑:会给出论文中的核心点和核心点的描述。   

一句话总结OLMo 2 是 Allen AI 推出的下一代“完全开源”基础模型系列(1.5B/8B/70B),基于18 万亿 tokensDolma 2.0 语料库训练,架构上全面采用 Grouped-Query Attention (GQA) 并将上下文长度扩展至128K tokens,同时通过 SFT 和 DPO 实现了高性能和完全透明化的对齐。

技术亮点

1 大规模数据与架构革新: 数据方面,18 万亿 tokens 的公开、高质量 Dolma 2.0 语料库上进行训练,数据规模实现大幅飞跃。架构方面,OLMo 2 全系列模型采用了 Grouped-Query Attention (GQA) 提升推理效率,并将上下文长度史无前例地扩展至 128K tokens,同时词汇表扩大至 100K,显著增强了模型的长文本处理能力和效率。

2 后训练 【SFT + DPO 与完全透明化对齐】: OLMo 2 提供了专用的 Chat 版本,通过监督微调(SFT)和直接偏好优化(DPO)等先进技术进行对齐,显著提升了模型在指令遵循、安全性与人类偏好方面的表现。

备注:后续为机器翻译

摘要

       我们呈现 OLMo 2,这是我们下一代的“完全开放”语言模型。OLMo 2 包含一系列稠密自回归语言模型,规模涵盖 7B、13B 和 32B,并完整公开全部研发产物——模型权重、完整训练数据、训练代码与配方、训练日志以及数千个中间检查点。在本工作中,我们描述了经过修改的模型架构与训练方案,重点介绍用于实现更高训练稳定性更高每 token 训练效率的技术。我们更新后的预训练数据混合方式中引入了一种新的专门化数据混合集,称为 Dolmino Mix 1124。当该数据通过预训练后期课程式训练(即在退火阶段使用专门化数据)引入时,能够在多个下游任务基准测试中显著提升模型能力。最后,我们结合 Tülu 3 的最佳实践来开发 OLMo 2-Instruct,重点使用许可度更高的数据,并在最终阶段加入**可验证奖励(RLVR)**的强化学习。我们的 OLMo 2 基座模型在“性能–训练算力”之间形成了帕累托最优,常常能够匹敌甚至超越 Llama 3.1、Qwen 2.5 和 Gemma 2 等仅开源权重的模型,同时使用更少的 FLOPs,并提供完全透明的训练数据、代码与训练配方。我们的完全开源的 OLMo 2-Instruct 模型在与同规模的开源权重模型对比中具有竞争力,甚至能够与像 GPT-3.5 Turbo 和 GPT-4o Mini 等部分专有模型竞争。

一 介绍

图1展示了 OLMo2及其他类似模型在预训练FLOPs 下的性能(模型规模为 6 个训练标记:Kaplan 等人,2020)。我们可以看到,完全开放的 OLMo2 位于帕累托前沿,在多种规模下均优于许多其他开放程度不同的模型。完整结果见表 6。

     

        在过去的一年里,开源语言模型生态系统迅速发展。我们见证了来自成熟开发者的开源权重模型激增——Llama 3(Grattafiori 等,2024)、DBRX(Databricks,2024)、Yi 1.5(Young 等,2024)、Qwen 2(Yang 等,2024a)、Falcon(TII,2024a,b)、Mistral(Mistral,2024a)、Ministral(Mistral,2024b)、Phi(Abdin 等,2024a,b)——以及来自新贡献者的模型——Gemma(Gemma Team 等,2024a,b;Team 等,2025)、Grok(X.AI,2023)、Command R(Cohere,2024a,c,b)。这些模型在很大程度上缩小了公共可用系统与闭源系统之间的差距(Cottier 等,2024)。然而,这些开源权重模型只是复杂语言模型配方与开发流水线的最终产物,它们本身并不足以支持对语言模型行为和使用方式的多样化研究。

        作为回应,先前的工作,包括我们首次发布的 OLMo(Groeneveld 等,2024)、Pythia(Biderman 等,2023)、Amber(Liu 等,2023c)、DCLM(Li 等,2024)、MAP Neo(Zhang 等,2024a)和 SmolLM(Allal 等,2024a,b),采用了完全开放的方法,不仅发布模型权重,还发布训练数据、训练代码和文档完善的训练配方,以支持可重复性研究。完全开放语言模型的成果在研究训练动态(Land 与 Bartolo,2024;Jin 与 Ren,2024)、概念获取(Chang 等,2024)以及语言模型的记忆机制(Antoniades 等,2024;Shaib 等,2024)方面发挥了关键作用。尽管如此,在最佳性能模型与开源模型之间仍然存在差距。

        现代语言模型的开发是一个迭代过程,当前模型的局限性会推动后续版本的改进。我们之前的版本(OLMo-0424;Ai2,2024)主要通过更好的预训练数据混合与课程机制来提升关键任务(如 MMLU)的性能。在本技术报告中,我们介绍 OLMo 2——一个新的 7B、13B 和 32B 模型家族,使用最多 6T tokens 进行训练。在英语学术基准测试上,这些模型与开源权重模型家族 Llama 3.1、Qwen 2.5 和 Gemma 2 具备竞争力(图 1)。我们进一步通过应用 Tülu 3 训练方法(Lambert 等,2024)验证了我们的预训练模型作为下游后训练(post-training)基座模型的有效性。最终得到的模型家族称为 OLMo 2-Instruct,与其他强大的开源权重模型以及部分常见专有模型(如 GPT-3.5 Turbo 和 GPT 4o Mini)相比也具备竞争力。本技术报告重点讨论了我们在 OLMo 2 开发过程中关注的四大重点领域:

        • 预训练稳定性。语言模型训练常常受到训练不稳定性与 loss 猛增的困扰,这些问题代价高昂,且会严重影响模型最终效果。我们讨论了用于提升训练稳定性的技术,这些方法对确保最终模型性能至关重要(第 §3 节)。

        • 中期训练方案。OLMo-0424(Ai2,2024)、DBRX(Databricks,2024)和 Llama 3(Grattafiori 等,2024)展示了数据课程对预训练的价值,这一现象也由 Blakeney 等(2024)讨论。我们介绍了将预训练分为两个阶段的优势,其中后阶段(mid-training)用于注入新知识并弥补能力缺陷。此外,我们展示了如何通过一种称为微退火(micro-annealing)**的技术,使中期训练数据源能够被独立评估以降低实验成本(第 §4 节)。

        • 后训练流水线。成功的基座模型必须能够有效微调用于下游任务。我们介绍了基于 Tülu 3 配方构建的 OLMo 2-Instruct,并展示了基座模型改进如何转化为更强的聊天模型。我们重点关注许可度较高的数据,并将可验证奖励强化学习(RLVR)**扩展到多个阶段以获得最高性能(第 §5 节)。

        • 作为研究催化剂的基础设施。高性能且可靠的基础设施对于成功预训练至关重要,然而许多预训练论文没有讨论其训练技术栈,或略过关键细节。我们讨论了从 OLMo-0424 以来的基础设施改进,这些改进促成了 OLMo 2 的进步,并介绍了通过投资于监控与调度基础设施来降低训练失败率、提升集群利用率的方法(第 §6 节)。

        除了上述重点内容,我们还在第 §2 节中描述了完整的模型开发流程:训练数据、预训练、后训练和评估。我们适当强调了与 OLMo 1 和 OLMo-0424 的变化,并引用了相关项目,例如我们用于高效估计模型下游性能的 scaling laws 研究(Bhagia 等,2024),以及通过 OLMES 评估框架实现基准标准化的工作(Gu 等,2024)。

二  OLMo 2 系列

        本节概述 OLMo 2,并强调相较 OLMo-0424 与之前的 OLMo 模型的改进¹。OLMo 2 系列拥有更多的 tokens、更多的参数,并在下游任务上取得比 OLMo-0424 更好的结果。我们解释了实现具备竞争力结果所需的关键细节,以实现我们让最先进语言模型可访问的使命。因此,我们尽可能在 Apache 2.0 许可下完全开源发布所有训练代码、数据和训练配方;在无法使用 Apache 2.0 的情况下,则使用当前可用的最宽松许可。

2.1 模型架构

        表 1 展示了 OLMo 系列中模型架构在多次迭代中的演进情况。我们在下方提供具体细节:

表1总结了 OLMo 系列模型架构随时间演变的过程。OLMo2的最新变化源于实验结果,这些变化表明训练稳定性有所提高。完整描述见 2.1节。

        我们采用基于 Vaswani 等(2017)的 decoder-only Transformer 架构,并提供 7B、13B 与 32B 参数规模的模型(见表 3)。我们的架构与第一代 OLMo(Groeneveld 等,2024)非常相似,但进行了多项更改以提升训练稳定性(见第 §3 节)与性能。

原始 OLMo 对 decoder-only Transformer 架构(Vaswani 等,2017)进行了以下修改:

        • 无偏置(No biases):我们从架构中移除所有 bias 项(Groeneveld 等,2024;Chowdhery 等,2022 等)。
        • SwiGLU 激活函数:使用 SwiGLU 激活函数(Shazeer,2020),并将对应的隐藏维度设为约 8/3d,但提升至最接近的 128 倍数(例如 7B 模型为 11,008)以提高吞吐量。
        • 旋转位置编码(RoPE):用旋转位置编码(RoPE;Su 等,2021)替代绝对位置编码。

在构建 OLMo-0424 时,我们为了训练稳定性和下游性能进行了以下修改:

        • QKV Clipping:用于训练稳定性,类似 DBRX(Databricks,2024)。
        • 上下文长度增加:从 2048 提升到 4096。

本次工作引入了 OLMo 2,并进一步进行了以下修改:

        • RMSNorm:我们使用 RMSNorm(Zhang & Sennrich,2019)来替代无参数 LayerNorm(Ba 等,2016),用于对激活进行归一化,且不使用 bias。
        • 归一化顺序调整(Reordered norm):我们对每个 Transformer Block 内 attention 和前馈网络(MLP)的输出而非输入进行归一化。即每个 Block 的公式变为:


         其中 x 为层输入,h 为中间隐藏状态,hout 为输出。 此策略最早由 Liu 等(2021)提出以改善训练稳定性。
        • QK-norm:遵循 Dehghani 等(2023b),在计算注意力前对 key 和 query 投影应用 RMSNorm,避免 attention logits 过大导致训练 loss 发散。
        • Z-Loss:遵循 Chowdhery 等(2022)、Chameleon Team(2024)和 Wortsman 等(2023),我们采用 z-loss 正则化——其在经验上已被证明能提升训练稳定性。
        • RoPE θ = 5e5:将 RoPE 的 θ 从 10,000 提升到 500,000,提高位置编码分辨率,与 Grattafiori 等(2024)一致。

2.2 分词器(Tokenizer)

        OLMo 1 与 OLMo-0424 使用修改版 GPT-NeoX-20B tokenizer(Black 等,2022)进行训练,该 tokenizer 包含以下 special tokens:|||PHONE_NUMBER|||,|||EMAIL_ADDRESS|||, |||IP_ADDRESS|||用于屏蔽个人可识别信息(PII)。

        按照 Tao 等(2024)的建议,我们为 OLMo 2 采用更大的 tokenizer vocabulary。我们借用了 cl100k 的 pre-tokenizer 与词表,它是 GPT-3.5(OpenAI,2023a)和 GPT-4(OpenAI,2023b)的 tokenizer,并在 Apache 2.0 许可下发布²。为了保持对早期 Dolma 数据源的兼容性,我们保留前代 OLMo 模型中使用的相同 masking tokens。

表 2 比较了 OLMo 1 和 OLMo 2分词器在基于DCLM 基线模型预训练的 100B 模型上的表现。根据 Gu 等人(2024) 的研究,OLMES 和 MMLU 使用 CF 格式,这种格式对于小型模型来说信息量更大。

        我们在较小规模上比较了两种 tokenizer(表 2)。实验显示,使用新 tokenizer 可测得性能提升,尤其在 OLMES 任务上。根据 Tao 等(2024),在当前模型规模与计算预算下,更大的 tokenizer 对 OLMo 2 略有劣势;我们预计更大的词表在更大规模模型以及更多训练 tokens 的场景下会带来更明显的改进。

2.3 基础模型训练配方

        遵循之前的 OLMo 模型以及近期在课程学习(curriculum learning)方面的进展(Blakeney 等,2024;Ibrahim 等,2024),OLMo 2 基础模型采用两阶段训练方式,每个阶段分别使用不同的数据混合。第一阶段(预训练阶段)是最长的阶段(占训练 FLOPs 的 ≥ 90%),主要使用来自网络的数据。在该阶段,我们对预训练混合进行了迭代,使用高质量的网页数据,并结合近期其他开放数据集。在第二阶段(我们称其为“中期训练”,占训练 FLOPs 的 5–10%)中,我们对最高质量的网页文档与精选的非网页来源进行上采样;同时使用合成数据以增强模型数学能力。

表3 OLMo 2超参数

        阶段 1:预训练。第一阶段——预训练——是整个训练中最长的一段(占训练 FLOPs 的 90–95%)。关键的架构与训练细节在表 3 中列出。关键点包括:我们从多头注意力(MHA)切换到 grouped query attention(GQA)(Ainslie 等,2023),以支持 32B 模型的扩展,该做法受到了同期工作 Qwen 3(Yang 等,2025)的启发。OLMo 2 训练使用均值为 0、标准差为 0.02 的截断正态分布进行随机初始化,并采用以下学习率调度策略:学习率在前 2000 步内从 0 线性升至峰值,然后采用余弦衰减,在指定的最大 token 数处衰减至峰值学习率的 10%。

        阶段 2:中期训练。第二阶段较短,我们称之为“中期训练”(占训练 FLOPs 的 5–10%)。在此阶段,我们线性衰减学习率至 0,持续到训练结束。我们构建了一个更小、更聚焦的混合数据集——Dolmino Mix 1124,以通过增加模型接触高质量文本与 STEM 领域参考内容来增强其特定领域能力,同时修补模型在第一阶段预训练后仍然不足的能力(如数学求解能力)。我们对高质量网页文档与精选的非网页来源进行上采样,并使用专门设计的合成数据以增强模型数学能力。

        模型合并或“Souping”。为了最大化高质量数据的效果、并找到更优的局部最小值,我们对不同随机数据顺序重复执行此步骤多次,并对最终训练得到的模型进行平均(Matena & Raffel,2022;Wortsman 等,2022)。对于 OLMo 2 7B,我们进行了三次独立退火,每次 500 亿 tokens,数据顺序不同;最终模型为三者的平均。对于 OLMo 2 13BOLMo 2 32B,我们进行了三次 1000 亿 tokens 的训练(与 7B 拥有相同更新步数),以及第四次 3000 亿 tokens 的训练;最终模型为这四个模型的平均。

        整体规模。OLMo 2 7B:总训练 4.05 万亿 tokens(其中 3.90 万亿来自预训练阶段)。OLMo 2 13B:总训练 5.6 万亿 tokens(其中 5 万亿来自预训练阶段)。OLMo 2 32B:总训练 6.6 万亿 tokens(其中 6.06 万亿来自预训练阶段)

2.4 基础模型数据

        本节简要介绍预训练与中期训练的数据混合。

2.4.1 预训练数据:OLMo 2 Mix 1124
表4 OLMo2预训练数据的组成。OLMo21124Mix由 StarCoder(Li等人,2023b;Kocetkov等人,2022)、peS2o(Soldaini和Lo,2023)来自 DCLM 的网页文本(Li 等人,2024)和来自Dolma 1.7 的 Wiki(Soldaini等人,2024)组成。arXiv 来自 Red-Pajama(Together Al, 2023)而 OpenWebMath(Paster等人,2023)和代数栈来自 ProofPilell(Azerbayev等人,2023)。

        用于该阶段的混合数据如表 4 所示。它由大约 3.9 万亿 tokens 构成,其中超过 95% 来源于网页数据。我们将这一数据集称为 OLMo 2 Mix 1124。这是与 OLMoE(Muennighoff et al., 2024)所使用的相同预训练数据:我们结合了 DCLM(Li et al., 2024)和 Dolma 1.7(Soldaini et al., 2024)中的数据。来自 DCLM 的部分,我们使用其“baseline 1.0”混合数据。来自 Dolma 的部分,我们使用 arXiv(Together AI, 2023)、OpenWebMath(Paster et al., 2023)、Algebraic Stack、peS2o(Soldaini and Lo, 2023)以及 Wikipedia 子集。arXiv、OpenWebMath 和 Algebraic Stack 最初属于 ProofPile II(Azerbayev et al., 2023)。最后,我们还加入了来自 StarCoder(Li et al., 2023b)的代码,其来源于 GitHub 上许可宽松的代码库(Kocetkov et al., 2022)。

        为了尝试包含更高质量的代码,我们移除了任何 GitHub 上 star 数少于 2 的代码库中的文档。此外,通过对该来源进行人工检查,我们发现其中包含以二进制格式编码的文档或主要由数字内容构成的文档;为去除这些文档,我们丢弃了任意文档中出现频率最高的单词占文档超过 30%,或**出现频率最高的两个单词合计占文档超过 50%**的情况。为缓解可能的训练 loss 峰值,我们移除了包含 32 个或更多 n-gram 重复序列的文档。我们在第 §3.1 节中报告了相关细节并展示了该干预的有效性。

2.4.2 中期训练数据:Dolmino Mix 1124

        在最初主要基于网页数据的预训练阶段之后,我们进一步使用一个经过更严格质量过滤的网页数据混合集,以及一系列领域特定的高质量数据(其中大部分为合成数据)继续训练。该混合数据的目的是赋予模型以数学为中心的能力,并提供专注的 STEM 参考内容及高质量文本的曝光。我们生成了多个大小不同的该混合数据的变体,但通常将其统称为 Dolmino Mix 1124。Dolmino Mix 1124 的基础数据来源描述于表 5。我们引导读者参考第 §4 节,以深入了解我们为该混合数据进行实验与数据筛选的流程。

表5中期训练数据集(Dolmino)的组成。我们从该数据集中创建了 500 亿、1 亿和 3 亿个标记的样本,用于对 OLMo 2 进行中期训练。有关各个数据源的详细信息,请参见第 4节;有关每种退火混合物的具体组成,请参见表 13。
2.5 评估与结果

        我们通过标准语言模型基准来评估 OLMo 2。此外,我们对 OLMo 2 进行后训练,并对结果模型——OLMo 2-Instruct——在多样化任务上进行评估,以衡量我们基础模型的适配潜力。

        基础模型评估。我们使用 OLMES 评估套件(Gu et al., 2024)对 OLMo 2 和其他基线模型进行了评估,该套件包含多种用于多项选择任务和生成式任务的基准数据集,并使用标准化提示及 in-context 示例进行少样本预测。基准任务的完整描述见附录 A.1。对于多项选择任务,我们评估准确率;对于生成式任务,我们评估 F1,以计入部分匹配。
此外,为了避免将我们的训练配方过拟合到这些基准上,我们维护了一套保留任务(held-out tasks),这些任务不用于模型开发决策;我们倡导模型开发者在开发时声明开发集任务与保留评估任务之间的区别作为标准实践。

        表 6 包含整体结果。我们发现,OLMo 2 模型在与相似规模的最佳开源权重模型相比时表现具有竞争力,尽管 OLMo 2 所需的训练 FLOPs 要少得多(见图 1),并且仍保持完全开放性(例如训练数据)。我们发现,开发指标上观察到的提升大体上能够迁移到用于评估的未见任务集上,这表明训练配方具有一定的泛化能力。

表 6 在 OLMES 套件的一个子集上比较 OLMo 2 与其他基础模型的评估结果(完整套件细节和结果见附录 A.1)。训练 FLOPs 使用 Kaplan 等人(2020)中的近似方法计算,并以 1023 的幂形式表示。我们无法估算任何 Mistral 模型(Jiang et al., 2023; Mistral AI, 2024)的计算量,因为其总训练 token 数未知。由于架构差异,Qwen 3(Yang et al., 2025)(并行工作)和 Zamba 2(Glorioso et al., 2024)的训练 FLOPs 未报告。Qwen 2.5 模型(Qwen et al., 2024)是在“最多 18 万亿 tokens”上训练的;开发者拒绝披露每个模型规模的确切 token 数。OLMo 2 模型在发布前未在保留数据集上进行评估;我们指出,对于其他模型,我们无法保证同样的情况。

        总体而言,我们发现开发指标上的增益在我们未见的评估任务集中能够较好地迁移。当然,我们无法保证在 OLMo 2 开发期间我们认为未见的任务,不在我们比较的其他模型的开发集合中。然而,我们认为模型开发者应在技术报告中标准化地保留一部分评估任务作为未见任务,并明确指出这些任务。此外,我们鼓励其他开源权重模型的开发者清楚说明在模型开发期间监控的是哪些任务。

        后训练配方与评估。对于后训练阶段,我们应用了来自 Tülu 3(Lambert et al., 2024)的配方,包括监督微调、基于策略的偏好调优(on-policy preference tuning),以及具有可验证奖励的强化学习(RLVR)。所得模型——OLMo 2-Instruct——在表 7 中进行了评估,评估内容包括通用与精确指令跟随、数学、知识推理以及安全性任务,使用的数据集与 Lambert et al. (2024) 的评估套件一致。基准任务的完整描述见附录 A.2。

        表 7 展示了下游任务的结果。我们发现 OLMo 2-Instruct 模型在最佳指令微调开源权重模型中具有竞争力,甚至优于某些常用的专有模型。这显示了 OLMo 2 作为强大的基础模型的价值,它是开展完全开放后训练研究的优秀起点。完整的后训练细节见第 §5 节。

表 7 OLMo 2 Instruct 在 1B、7B、13B 和 32B 参数规模下相对于同类开源权重模型的结果。以下评估名称为缩写:Avg – 平均值,AE2 – AlpacaEval 2,BBH – BigBenchHard,IFE – IFEval,PQA – PopQA,TQA – TruthfulQA。此表中的所有模型均为指令微调(instruction tuned)版本。
对于 Qwen QwQ 32B,PopQA 和 TruthfulQA 在答案抽取时出现困难,因为模型会将答案返回在 <think> tokens 内,因此我们未报告其得分。对于 Qwen QwQ 32B,我们通过移除思考(thinking)tokens 并对其后生成的答案进行评分来完成评估。除安全性评估外,我们按照其模型卡中推荐的采样参数进行全部评估(32K 上下文长度、0.6 温度、采样、topp 0.95、minp 0、topk 30),安全性评估仅使用较短的 8K token 上下文长度。在多项选择评估、PopQA 和 TruthfulQA 中,由于模型将答案返回在 <think> tokens 内,导致答案抽取困难,因此未报告其得分。即便撇除抽取问题,推理模型生成的超长上下文也给许多开源评估工具带来了挑战,这是我们需要改进的方向。

三 深度剖析:预训练稳定性

        尽管 OLMo-0424 在其计算预算范围内达到了预期性能,但其训练动态表现出一些令人担忧的问题:

        • 在训练过程中,loss 会突然出现峰值,并且更频繁地出现在梯度范数上。在实验中,我们发现随着模型规模增加,梯度峰值出现的频率也随之增加。此外,我们的实验发现,梯度范数中更剧烈的峰值通常出现在训练 loss 峰值之前。
        • 梯度范数的大小在整个训练过程中缓慢增大。这与梯度范数(以及训练 loss)中峰值频率的增加相关。

        最终,这些问题的组合会导致训练发散,使得更大规模的训练无法进行。这种情况促使我们开始对训练稳定性进行调查,以定位这些问题的原因并提出相应的缓解措施。图 2 展示了在实施我们缓解措施前后的训练曲线,我们将在下文总结这些措施:

        • 重复 n-grams:我们过滤预训练数据以移除重复的 n-gram 序列,因为它们可能导致 loss 峰值(§3.1)。
        • 初始化:我们从缩放初始化(Zhang et al., 2019)切换为使用均值为 0、标准差为 0.02 的截断正态分布初始化所有参数(§3.2)。
        • RMSNorm:我们使用 RMSNorm 作为 LayerNorm 的变体来归一化激活,而不是使用无参数的标准 LayerNorm(§3.3.2)。
        • 重新排序的归一化(Reordered norm):我们对每个 transformer block 内注意力层和前馈(MLP)层的输出进行归一化,而不是对输入进行归一化(§3.3.2)。
        • QK-norm:在计算注意力之前,我们使用 RMSNorm 对 key 和 query 投影进行归一化(§3.3.2)。
        • Z-Loss:我们采用 z-loss 正则化,一种抑制最终输出 logits 过大增长的正则项(§3.3.3)。
        • 权重衰减:我们将 embeddings 排除在权重衰减之外(§3.4.2)。
        • AdamW 的 ϵ:我们将 AdamW 的 ϵ 从 10⁻⁵ 降低到 10⁻⁸(§3.4.1)。

        在以下内容中,我们将讨论导致我们采用这些干预措施的实验与结果。我们将我们的修订策略与 OLMo-0424 进行比较,后者是最近一个完全开源模型(公开权重、数据和文档)的版本。

3.1 重复 n-Grams

        数据可能是梯度范数和损失突增(spikes)的原因。当我们调查发生突增的训练批次时,发现其中高度集中存在包含长、重复 n-gram 序列的实例。以下是此类序列的三个示例:

        在一系列实验中,我们发现这些序列通常与突增相关,但我们指出这一关系并非确定性的:

        • 相同的 n-gram 序列可能会在较大模型中触发突增,但在使用相同数据训练的较小模型中不会出现。
        • 相同的 n-gram 序列可能会在某个数据训练顺序中触发突增,但在数据重新洗牌后不再出现。
        • 与突增相关的相同 n-gram 序列也可能出现在其他未发生突增的训练批次中。

        尽管如此,我们发现,在训练数据中广泛移除此类序列可以在平均意义上降低突增频率。在数据整理阶段(第 §2.4 节),我们应用了一个过滤器,会移除所有包含长度为 32 或以上重复 n-gram 的文档,其中 n-gram 指任意长度为 1 到 13 个 token 的片段。我们还在训练器中实施了额外的保护机制,在数据加载期间检测这些序列,并在计算损失时进行掩码处理。图 3 展示了对包含重复 n-gram 的输入序列进行损失掩码的效果。该方法带来了明显的缓解——尽管并未完全消除梯度突增。它对梯度范数的缓慢增长没有影响。

图3 比较了两次运行的梯度范数,一次未使用 n-gram 滤波器,另-次使用了 n-gram 滤波器。忽略长重复的 n-gram 序列可以消除许多峰值。

3.2 模型初始化

        图 4 展示了 OLMo 2 初始化方案带来的训练稳定性提升。在 OLMo 2 中,我们对每个参数使用均值为 0、标准差为 0.02 的正态分布进行初始化。相比之下,OLMo-0424 的初始化方案(最早由 Zhang et al., 2019 提出,并由 Gururangan et al., 2023 实现)会按以下方式进行缩放:输入投影按 1/√dmodel 缩放,输出投影按 1/(√2 ⋅ dmodel ⋅ layer_idx) 缩放。换言之,越靠后的层初始化值越小。

        我们进行了多项分析以研究初始化的影响,并显示 OLMo 2 的初始化优于 OLMo-0424 初始化。我们的实证分析表明,它能更好地保持各层激活值和梯度的尺度,使深层模型训练得更稳定,并表现出与不同宽度模型之间超参数迁移相关的性质。这两个特性使我们有信心深层模型能够稳定训练,并且小模型的初始化超参数可以迁移到更大规模。

图4在我们的测试环境中,OLMo-0424 初始化方案很快出现不稳定现象,而 OLMo 2 则保持稳定。

        梯度与激活增长。训练深度网络的一个基本问题是确保激活与梯度不会在层间爆炸或消失,否则会导致学习不稳定或停滞。我们期望激活和梯度的尺度在层与层之间大致保持不变。受近期相关工作(Cowsik et al., 2024)启发,我们从激活和梯度在层间的 2-范数变化角度评估候选初始化方案。具体做法如下,我们随机初始化一个模型,将来自 The Pile(Gao et al., 2021)的 50 个随机文档输入模型,并收集初始层与最终层的激活和梯度(损失关于激活的梯度),忽略embedding。然后在文档和时间步上对这些张量求平均,得到初始层向量 v 和最终层向量 v′,它们的长度均为 dmodel。最后,我们计算以下度量,用于衡量层间的扩张或收缩程度,称为增长指数:

        我们对激活和梯度均计算 λ。理想情况下,两者都应接近 0,表示激活和梯度不会在层间爆炸或消失。图 5 绘制了不同随机初始化模型在不同宽度(4096 对应完整的 7B 模型)下的增长指数。关键点在于:与 OLMo-0424 相比,OLMo 2 的增长指数更接近 0。这表明,在低精度训练深层模型时,OLMo 2 初始化更稳定,因为激活与梯度在层间更不易爆炸或消失。

图5显示,在不同宽度下,OLMo2初始化的增长指数比 OLMo-0424 初始化的增长指数更接近于0,这表明更深的模型将训练得更稳定。

        跨宽度的超参数迁移。新初始化方案的另一个吸引人的性质是:它使激活和梯度范数随宽度 dmodel 的变化方式,与理论上对于跨宽度超参数迁移的重要条件一致。具体来说,Yang et al.(2024b)提出:为了实现跨模型宽度的超参数迁移,一个充分条件是每个激活标量值及其更新量(学习率 × 梯度)的幅度在宽度增加时保持不变。换言之,激活和其更新向量的范数应与 √dmodel 正相关。我们在图 6 中绘制了初始化时激活与梯度范数与 √dmodel 的关系。关键在于:相较于 OLMo-0424,OLMo 2 的梯度范数与 √dmodel 的正相关性更强。结合 Yang et al.(2024b),这表明在学习率不依赖模型宽度的设置下,OLMo 2 初始化比 OLMo-0424 初始化更有利于在不同模型宽度之间进行超参数迁移。

        突增分数。由于常规可视化工具难以理解快速突增,我们使用突增分数作为客观度量。具体来说,我们将突增分数定义为:在一个时间序列中,至少有 1000 个最近值的滚动平均值上方 7 个标准差的值所占的百分比。突增分数主要用于训练损失和梯度 L2 范数,但也可用于任意时间序列。

图 6 展示了 OLMo-0424和 OLMo-2 初始化的激活值和梯度范数与 dmodel 的关系。值得注意的是,OLMo-2 的梯度范数与 dmodel 呈正相关,而OLMo-0424 初始化则不然。这表明 OLMo-2 初始化在不同宽度下具有更好的超参数迁移能力(Yang等人,2024b)。

        实验结果。为了实验模型初始化,我们首先构建了一个基线运行,以便快速复现突增。我们通过减少 warmup 时间来实现这一点。该操作带来的影响立竿见影且非常显著(图 4),并在不同模型规模与 token 数量中保持一致。在消融实验中,新初始化没有出现任何损失突增,并且梯度 L2 范数的突增分数从 0.40 降至 0.03。新初始化的收敛速度略慢;我们通过改进其他超参数设置(第 §3.4 节)弥补了这一差异。

3.3 架构改进

3.3.1 非参数化层归一化与 RMSNorm

        OLMo 2 使用 RMSNorm,这在大多数 transformer 实现中已属标准做法。OLMo-0424 出于性能考虑并为规避我们使用的库中的一些 bug,采用了非参数化层归一化。但在我们开发 OLMo 2 时,这些 bug 已不再是问题,硬件也更快,同时我们希望采用一种安全的方案。我们的消融实验表明两者之间没有差异,因此我们改回使用 RMSNorm。

3.3.2 重新排序的归一化与 QK-norm

        图 7 展示了将层归一化应用于 MLP 和注意力模块的 输出 而非 输入 的效果。我们进一步在注意力模块中的 query 和 key 上再应用一次归一化(同样是 RMSNorm)。单独来看,这两项修改都不会带来好的结果,但组合在一起时,它们改善了梯度 L2 范数的增长与“尖峰”(spikiness)。
下表总结了层归一化的位置差异:x 是层的输入,h 是中间隐藏状态,hout 是输出。

        Liu 等人(2021)首次提出了重新排序层归一化(layer norm)的想法。随后,Chameleon Team(2024)采用了这一方法。QK-norm 最早由 Dehghani 等人(2023a)提出。

3.3.3 Z-Loss

        参考 Chowdhery et al. (2022)、Chameleon Team (2024) 与 Wortsman et al. (2023),我们通过在损失函数中添加10⁻⁴ · log₂ Z的方式应用 z-loss 正则化,其中 Z 是 softmax 在 logits 上的分母。此正则化会抑制最终 softmax 中的激活值过大,从而提高模型稳定性。

        图 8 展示了流行的 Flash Attention 库(Dao, 2024)中的 z-loss 实现,与一个仅使用 Python 原语写成的实现之间的显著差异。除了广为人知的注意力机制外,Flash Attention 也提供优化过的交叉熵损失实现,其中包括一个 z-loss 版本。为了在与 Flash Attention 不兼容的场景中保持灵活性,我们另写了一个 PyTorch 的实现。虽然两个实现的正向传播结果一致,但反向传播行为不同。我们怀疑根本原因在于精度差异。在我们的实验中,这并未影响训练过程中的交叉熵损失,或模型在下游任务上的表现。然而出于谨慎考虑,我们放弃了带有自定义 z-loss 的分叉实现,并从最初分歧点重新训练。在一次训练过程中我们无法安全地更换实现,因此我们尽量避免这样做。

图8显示,Flash Attention 实现的 z-loss 与PyTorch 中的手动实现并不一致。虽然前向传播的结果相同,但反向传播的差异导致曲线出现分歧。

3.4 超参数改进

3.4.1 AdamW 中的 ϵ

        图 9 显示了将 AdamW 的 ϵ 从 10^-5 降低到 10^-8 的效果。10^-8 是 PyTorch 的默认值,但一些流行的语言模型训练代码库默认值为 10^-5。较低的值允许在训练早期进行更大的更新,并有助于模型在通常会出现较多不稳定性的阶段更快地学习。因此,梯度范数会更快地稳定下来,并保持在较低水平。

图 9 显示,将 AdamW 的e值设置为 10-可以降低并稳定训练初期梯度的范数,同时训练损失也能更快地改善。即使运行时间比图中所示更长,这种趋势依然存在。
3.4.2 嵌入层的权重衰减

        图 10 显示了在决定对嵌入层不施加权重衰减后训练动态的变化。OLMo 使用标准的权重衰减公式,其中每个参数在每一步都乘以 1 − (0.1 ⋅ lr)。这一正则化项可以抑制参数过大,但在 token 嵌入的情况下,它会过度作用,导致嵌入过小。如 Takase 等人(2024)所述,小嵌入会在早期层产生较大梯度,因为 layer_norm(x) 对 x 的雅可比与 ∥x∥ 成反比,而在早期层,残差流的范数本质上就是嵌入的范数。我们尝试了 Takase 等人(2024)讨论的所有修正方法,但发现它们会影响收敛速度。相反,我们简单地关闭了嵌入的权重衰减,观察到嵌入范数在训练过程中稳定在健康范围内。

图 10 显示,对词嵌入应用权重衰减会导致嵌入范数逐渐减小,梯度范数相应增大。衰减嵌入也会对稳定性产生轻微的负面影响,产生比未应用衰减的类似运行更多的峰值(峰值得分分别为 0.16 和 0.092)。

四 深入解析:中期训练方案

        近期研究表明,对基础模型采用多阶段训练方法可以显著提升模型能力(Blakeney 等人,2024;Ibrahim 等人,2024;Feng 等人,2024)。在之前的 OLMo 迭代中,我们也发现学习率调度(OLMo 1;Groeneveld 等人 2024)和数据混合(OLMo-0424;Ai2 2024)都起到了重要作用。我们将模型开发阶段的此类干预称为中期训练。

        总体而言,我们的方法很简单:在预训练阶段之后,我们生成特定领域的数据混合并重新启动训练,同时线性将学习率降至零。我们的目标是灌输专业知识并提升模型能力;这些改进的反馈来源于关键基准测试,例如数学专用任务 GSM8K。

4.1 学习率退火

        我们学习率实验的起点来自 Grattafiori 等人(2024)的设置。对于 7B 变体,我们将学习率在前 2000 步线性升温至峰值 3 ⋅ 10^-4,然后在 5T tokens 上使用标准余弦衰减。OLMo-0424 的经验表明,余弦衰减计划的最后部分可以被截断并用线性衰减至零代替,对性能影响很小。因此,对于 7B 变体,我们在 4T tokens 时停止调度,然后切换到第 §4 节描述的中期训练阶段。13B 从一开始使用更高的峰值学习率,因此我们决定在 5T tokens 后再进入中期训练阶段。

        图 11 显示了使用四个额外学习率值的不同训练运行:6 ⋅ 10^-4、9 ⋅ 10^-4、12 ⋅ 10^-4 和 30 ⋅ 10^-4。特别地,我们尝试了原始学习率的 2 倍、3 倍、4 倍、10 倍和 30 倍。最后的 30 ⋅ 10^-4 在学习率升温期间就显示出训练不稳定,出现若干无法完全恢复的损失峰值,因此我们迅速放弃了这一变体。其他值训练正常,并显示出有趣的模式。从训练损失来看,较高的学习率在初期表现更好(只要避免损失峰值),但最终较低的学习率设置会超过其他设置(图 11)。值得注意的是,在比较 3 ⋅ 10^-4 和 6 ⋅ 10^-4 时,交叉点远超过 200B tokens。较短的超参数实验可能得出错误结论。

图 11 显示,较高的学习率在初期表现更佳,但最终会被较低的学习率超越。然而,在 500 亿或 1000亿个 token 内将学习率线性衰减至零,会导致相同的训练损失。

        这一系列实验的动机之一是探索更高的学习率是否会使退火步骤更有效。假设是:在预训练期间更差的训练损失会在学习率衰减至零时得到补偿。为了验证这一假设,我们在四个变体中分别在训练 300B tokens 后取一个检查点,并在 50B tokens 上将学习率衰减至零。为了考虑更高学习率的效果可能需要更多步骤才能显现,我们尝试了三个较高设置,并在 100B tokens 上衰减学习率,总共进行了七次实验。结果显示,更高的学习率确实使中期训练更有效,但其效果正好抵消了预训练期间的较差表现。四个变体在该过程结束时显示相同的训练损失,尽管最低设置略微落后于其他设置。

        表 8 显示了在更长训练周期中结果的一致性。我们选择了两个变体,3 ⋅ 10^-4 和 6 ⋅ 10^-4,并在训练 1T 和 2T tokens 后重复实验。我们选择这些变体是因为 3 ⋅ 10^-4 是 Grattafiori 等人(2024)的基线,而 6 ⋅ 10^-4 在训练损失上略优。结果显示,这两种设置在训练损失和 OLMES 套件(Gu 等人,2024)中九个下游任务的混合上几乎没有差异。对下游任务的评估噪声较大,但与仅基于训练损失的发现一致。

表 8 列出了 OLMES 验证子集中 9 道多项选择题(完形填空形式)在不同峰值学习率和学习时长下的结果。所有变体的平均分差异均小于2分,且大多数分数彼此相差不到 0.5 分。

        最后,我们希望观察预训练阶段更高的学习率是否会在切换到高质量数据时使中期训练阶段更有效。为了尽可能匹配我们的训练设置并在可用计算预算内完成,我们使用相同的两个设置(3 ⋅ 10^-4 和 6 ⋅ 10^-4),并在 100B 高质量 tokens 上线性衰减学习率至零。结果再次显示差异很小。在 OLMES 评估套件上的最终得分相差不到 0.1 分。然而,查看其他指标可能仍会显示两种设置之间的有意义差异。高质量 tokens 的混合特别针对数学任务,在 GSM8K(不属于 OLMES 套件)上,高学习率设置比低学习率高 2.8 分。需要进一步研究才能将这一有趣的数据点转化为可靠结果。

        这一发现与机器学习经验法则如“更高的学习率总是更好”或“学习曲线下的面积很重要”(McCandlish 等人,2018)相矛盾。它扩展了 Wortsman 等人(2023)的结论,他们观察到较小模型在训练到余弦衰减结束时,对几个数量级的学习率变化性能基本不敏感,并进一步发现我们使用的 QK-norm(第 3.3.2 节)和 z-loss(第 3.3.3 节)增强了这一效果。我们发现,即使在更大规模的 token 和参数下,这些结果仍然成立,且对我们的训练工作至关重要,因为我们采用了修改后的学习率计划。

        由于成本限制,我们没有探索完整的学习率范围。这是这一系列实验的主要局限。进行更广泛的学习率扫描可能有助于准确界定我们训练的平台边界。

4.2 数据课程:Dolmino Mix 1124

        本节描述我们为中期训练策划数据的实验过程。我们将为中期训练阶段创建的数据集和混合统称为 Dolmino Mix 1124。该数据集内容概览见第 §2.4 节(表 5)。具体而言,我们在中期训练方案中使用以下流程:

  • 确定高质量数据源的混合,以提高整个开发基准套件的性能(第 §4.3 节)。

  • 为修补特定能力(在 OLMo 2 的情况下,特别是数学),收集并评估领域特定数据集以在中期训练中混合(第 §4.4 节)。我们发现,可以通过我们称之为微退火(microannealing,第 §4.4.2 节)的技术独立评估这些数据源;其有效性在与其他数据源混合时依然保持。

  • 根据第 §4.1 节描述的实验,我们将高质量数据源与数学专用数据在三个不同 token 预算下混合(50B、100B、300B)。较小的混合用于中期训练 OLMo 2 7B,而 OLMo 2 13B 和 32B 则在较大的混合上进行退火。对于 OLMo 2 7B、13B 和 32B,我们发现对相同混合但不同数据顺序种子训练的不同检查点进行权重平均,一致优于单个检查点(第 §4.5 节)。为了在小规模上演示这一点,我们还提供了 1B 模型在类似干预下的结果。

        表 9 总结了中期训练阶段对开发和保留评估的显著影响。OLMo 2 7B 模型平均提升 10.6 分,超过预训练阶段的 13B 大模型。OLMo 2 13B 同样从中期训练中受益,平均性能提升 10.3 分。两个模型在知识密集型、多选题(Arc challenge: 7B 72.6 → 79.8,13B 80.2 → 83.5;MMLU: 7B 59.8 → 63.7,13B 63.4 → 67.5;AGIEval: 7B 44.6 → 50.4,13B 48.2 → 54.2)、阅读理解(Natural Questions: 7B 29.0 → 36.9,13B 34.6 → 46.7;DROP: 7B 40.7 → 60.8,13B 49.6 → 70.7)和数学能力(GSM8K: 7B 24.1 → 67.5,13B 37.3 → 75.1)基准上均有所提升。

表 9 评估 OLMo 2 1B、7B、13B 和 32B 在预训练和中期训练阶段结束时的表现(设置与表 6 相同)。预训练检查点分别在 4 万亿(1B、7B)、5 万亿(13B)和 7 万亿(32B)tokens 上训练。对于 7B,我们通过对 50B Dolmino tokens 上的三次训练运行进行权重平均,获得最终的中期训练检查点;对于 13B 和 32B,我们使用 100B tokens 上的三次运行和 300B tokens 上的一次运行。对于1B,最终检查点为在 50B Dolmino tokens 上训练的结果,未进行平均。

4.3 Dolmino Mix 1124:高质量来源

        遵循前一版本 OLMo 的配方(Ai2, 2024),我们首先策划了预训练混合中的高质量子集,并通过更多学术和百科类材料进行扩展。具体而言,我们考虑以下来源(总结见表 10):

        高质量网页为了筛选用于预训练的网页子集,我们尝试了两种现有的质量分类器:
        • 来自 Li 等人(2024)的 FastText 分类器。为训练该模型,Li 等人从 ELI5(Fan 等人,2019)中的 Reddit 子集中采样正样本文档,以及 Open Hermes 2.5 的演示。负样本则从 DCLM 流水线中随机采样。
        • 来自 Penedo 等人(2024)的 FineWeb Edu 分类器。该模型基于 Arctic Embed M13 编码器(Merrick 等人,2024)进行微调,训练数据超过 40 万个网页,由 Llama 3 70B Instruct 标注。该分类器根据学术主题和内容精炼程度对文档评分,从 0 到 5。

        遵循 Li 等人(2024),我们使用 DCLM FastText 分类器,阈值为 0.03311014,保留约 65.6% 的网页子集。我们将此过滤器与 FineWeb Edu 分类器的评分结合;实验中保留评分高于 3 的文档(保留 5.8%),以及一个较宽松的阈值为 2(保留 20.3%)。

        指令数据与问答对。我们利用 Dolma 1.7(Soldaini 等人,2024)中的 FLAN Wei 等人(2021);Longpre 等人(2023)子集。我们通过从评估套件(第 §2.5 节)中的所有任务中提取训练、验证和测试实例,去除与任何任务实例 n-gram 重叠 ≥10% 的 FLAN 文档来进行去污染。

        我们从 Stack Exchange 网络收集问答对,该网络包含 186 个论坛,涵盖各种主题。Stack Exchange 网络内容采用各种商业友好型 Creative Common 许可证。我们使用撰写时的最新数据库转储(2024 年 9 月 30 日),由 Internet Archive 分发。我们仅保留具有被接受答案的问题;进一步地,移除问题投票数少于 3 或答案投票数少于 5 的问答对。筛选后,我们通过新行序列将问题与答案连接,序列中包含的换行符比问题或答案中最长换行序列多 1 个 \n。

        代码。我们评估保留预训练中使用的相同代码子集;此外,还考虑较小、经过策划的代码来源,与自然监督交错,例如 CodeSearchNet(Husain 等人,2019)中的 docstrings;上述 StackExchange 的问答对也包含代码。

        学术、百科及其他参考内容。我们从 Dolma 1.7(Soldaini 等人,2024)获取高质量非网页数据集,包括 peS2o(Soldaini 和 Lo, 2023)、Wikipedia、Wikibooks、Gutenberg 图书、arXiv 以及 StackExchange(来自 Red-Pajama v1;Together AI, 2023)、Algebraic Stack(ProofPile II;Azerbayev 等人,2023)。

        数学。在开发 Dolmino Mix 1124 的数学子集(第 §4.4 节)的同时,我们考虑初步数学子集,以评估数学文档与混合中非数学部分的结合效果。具体而言,我们使用 OpenWebMath(Paster 等人,2023)、GSM8K 训练集(Cobbe 等人,2021)、MathPile 宽松许可(“商业”)子集训练集(Wang 等人,2023b)以及 AutoMathText(Zhang 等人,2024b)。

        混合结果见表 10,总结见表 11。所有结果均对应在 500 亿 tokens 上的中期训练运行,初始化自在 4 万亿 tokens 上预训练的 7B 模型检查点。

表 10 总结了我们评估用于中期训练的高质量数据源。我们尝试将这些数据源混合成 6种不同的混合数据集,每种数据集包含 500 亿个标记。表中的百分比表示每种 500 亿混合数据集中来自相应数据源的数据比例。PT Mix 数据集(可重复采样)取自预训练阶段。
表 11 比较了表 10 中引入的中期训练混合。每一行对应一次 500 亿 tokens 的训练运行,遵循第 §4.1 节描述的学习率调度(第一行除外)。权重初始化自在 4 万亿 tokens 上预训练的 OLMo 2 检查点。我们在 OLMES 核心任务混合(多项选择格式;见表 6)、OLMES 生成任务(表 6)、MMLU(多项选择格式;Hendrycks 等人,2021a)以及我们用作开发集的 200 个 GSM8K 随机样本问题(GSM*;第 §A.1 节)上比较每次训练运行。最终中期训练混合的结果见表 9。

        如第 §4.1 节所述,仅学习率退火(PT Mix)即可在所有平均值上产生显著提升(OLMES +4.4;OLMES-Gen +1.3;MMLU +20),但在数学开发集(GSM* −1.5)上无提升。切换到包含更高质量网页数据和参考内容的混合进一步提高性能:Web FT7 FW2 在 OLMES 上比 PT Mix 提升 +1.2 分,在 MMLU 上提升 +1.3;在 OLMES-Gen 上略低(−0.4),在 GSM* 上在误差范围内(+1.5)。最后,在混合中加入指令数据和数学来源可获得最佳性能。Web FT7 FW2 + Math + Ins 混合取得最佳整体结果:OLMES 提升 +1.7,生成任务提升 +5.7,MMLU 提升 +1.3,GSM* 提升 +19.5。我们注意到 Web FT7 FW2 + Math 在数学任务上表现略优,这促使我们在第 §4.4 节中进一步研究更优数学子集与其他高质量来源的组合。

4.4 Dolmino Mix 1124:数学混合

        早期中期训练混合(仅表 11 中的 Web∗∗ 行)显示模型在数学相关基准上表现较差。因此,提高这些集合的性能是我们中期训练研究的核心关注点。我们研究了人工编写的数据和合成生成或增强的数据;后者是通过迭代过程导出的,旨在修正我们数学验证集中的常见错误。
我们在第 §4.4.1 节中描述数据来源及其生成/过滤过程;然后在第 §4.4.2 节中详细说明微退火(microanneals),这是我们用于最终确定数学数据源的实验技术。最终得到的混合总结见表 5。

4.4.1 数学数据来源

        TuluMath。我们遵循 Chan 等人(2024)提出的最近基于角色的方法生成数学合成数据。关键思想是使用不同的角色(例如“专注于神经网络的机器学习研究员”)配合数据生成提示(例如“创建一个数学问题”),引导语言模型生成对应视角的数据。具体而言,我们根据 Persona Hub(Chan 等人,2024)中可用的角色生成针对数学问题的提示,包括需要高级数学技能的问题以及小学问题。我们使用 GPT-4o 进行零样本提示生成独特且特定于给定角色输入的问题。在生成问题后,我们使用 GPT-4o 生成多步数学解答。用于生成问题和解答的具体提示见附录图 24 和 25。总共收集约 2.3 亿个合成数学 tokens。

        DolminoSynthMath。这是一个由 2800 万个合成数学 tokens 组成的集合,专门设计用于提升 GSM8K 及基础数学计算的性能。它由三部分组成:首先,我们生成 1100 万个基本数学问答对(例如“77 * 14 = 1078”),并为每个问答配备多种提示。我们发现,包含这些数据可以显著减少模型在推理时单步 CoT 错误。接下来,我们包含 7,924 个合成 GSM8K 示例,这些示例通过替换 GSM8K 训练样本中的所有数字生成,旨在提供信号让模型从文字问题中提取计算图,忽略无关语义特征。最后,我们使用 Qwen2.5-7B-Instruct(Qwen 等人,2024)生成每个 GSM8K 训练样本的 MIND 重写版本(Akter 等人,2024)。

        TinyGSM-MIND。我们从 Tiny-GSM(Liu 等人,2023a)重写版本生成约 6.5 亿个合成数学 tokens。Tiny-GSM 是由 1100 万个类似 GSM8K 的合成问题组成的集合,答案以 Python 代码形式提供。我们筛选该集合,仅保留可执行且仅包含变量赋值语句的答案。然后我们为每行赋值语句注释结果变量的数值。接着将所有带注释的示例输入 Qwen2.5-7B-Instruct,使用 “Two Students” 和 “Problem Solving” 提示以 MIND 风格(Akter 等人,2024)进行重写。

        MathCoder2-Synthetic。我们模拟 MathCoder2(Lu 等人,2024)的合成数据生成过程,从开源库中筛选现有合成数据。具体而言,我们收集 HuggingFace 用户 Ajibawa-2023 的合成教材,以及 M-A-P Matrix 数据集,并对其进行额外筛选。特别地,我们训练一个 FastText 分类器:让 GPT-4o 对 10,000 个 OpenWebMath 示例(Paster 等人,2023)标注为数学相关或非数学相关,然后将这些作为 FastText 分类器的正负样本。我们将该分类器应用于合成教材,仅保留数学相关教材。

        ProofPile OWM-Filtered。我们使用前一步生成的 OpenWebMath 筛选器,并将其应用于 Metamath(Yu 等人,2023)和 CodeSearchNet(Husain 等人,2019)。

        GSM8K-Train。最后,我们包含 GSM8K(Cobbe 等人,2021)的训练集。

4.4.2 使用微退火评估数学数据

        为了从所有可用及合成数学数据中选择最高质量的子集,我们执行了一系列微退火,这是针对小型数学子集的退火训练。微退火的一般流程如下:

  1. 确定我们希望评估数据质量的数据源或小型数学数据集合;

  2. 从一般数据混合(例如 DCLM)中收集与数学数据量大致相同的数据,以确保高质量网页文本与领域特定数学数据混合;

  3. 将该 50/50 混合训练,仿佛它是一次退火运行,确保对该小型数据集合按线性方式调整学习率。

        该方法可以在完整退火运行成本的一小部分下评估单个数据源的质量。总共进行了 19 次微退火,总 token 数为 1,300 亿,相当于不到三次完整 500 亿 tokens 的退火运行。对比来看,这 19 次微退火所需计算量少于我们 7B 模型使用的三次 500 亿 tokens 训练材料。更明确地说,它可以在更细粒度的数据源分辨率上显示改进,训练不到 100 亿 tokens 即可见效果。

        我们通过表 12 中报告的三组实验展示微退火如何生成最终数学混合。主要评估指标为 MMLU 和 GSM*,后者为 GSM8K 评估集的 200 个样本子集。注意,中期训练的一个目标是提高 GSM8K 性能,但我们仅允许检查 200 个样本以指导数据混合决策。

表 12 展示了 OLMo 2 数学能力微退火实验的结果。我们评估了数学/非数学题的混合比例、重复数学题项的影响以及不同数学数据集的影响。我们使用 200 道 GSM8K题目的随机样本(Cobbe 等人,2021)作为开发集(GSM*:A.1节),以此作为数学能力的代理指标。我们监控平均 MMLU 分数,以确保 OLMo2 在知识密集型任务上保持良好的性能。

        微退火实验 1:领域特定数据即使占比小也有帮助。实验设置:从完成预训练的 7B 模型出发,混合 TuluMath、DolminoSynthMath、Metamath、CodeSearchNet 和 GSM8K-Train,约 2 亿 tokens。我们训练 35/65 数学/DCLM 混合和 10/90 混合,并评估 MMLU 与 GSM*。
结果显示,预退火 GSM* 分数为 28.5,35/65 混合为 63.5,10/90 混合为 61。这表明退火混合中不必有大量领域特定数据,只需存在即可。

        微退火实验 2:适度重复数据有益。在相同设置下,将数学数据重复 1、2、4 次。结果:1 份数学数据 GSM* 为 61,2 份为 66,4 份为 65。这表明即使高质量领域特定数据稀缺,少量重复仍可带来收益。

        微退火实验 3:重写数据可显著提升性能。从完成预训练的 7B 模型开始,评估将 Tiny-GSM 重写为自然语言格式对 GSM* 的影响。Tiny-GSM 答案为代码形式,预训练混合仅 2% 代码。我们在带内联注释的 TinyGSM 混合上运行微退火,并与仅使用 “Problem Solving” MIND 重写版本比较。结果显示,代码版本的 TinyGSM 降低 GSM* 性能,而重写版本显著提升性能,表明重写是廉价且有效的将数据转换为训练友好形式的工具。

4.5 最终中期训练混合与检查点融合(Checkpoint Soups)

        Dolmino Mix 1124 的最终组成见表 5。如前所述,我们采样了 50B、100B 和 300B tokens 的三种混合;每种混合的组成总结见表 13。由于第 §4.3 和 §4.4.2 节的实验显示,保持各来源混合比例大致恒定是有益的,因此对于 100B tokens 混合,我们将 Stack Exchange 问答数据和中期训练数学数据重复两次,对于 300B 混合重复四次;此外,300B 混合中的 FLAN 数据重复两次,Wiki 数据重复四次。在所有混合中,经过筛选的 DCLM 基线网页数据约占总 tokens 预算的 50%。

表 13 Dolmino 混合料 1124 的成分。"“来源百分比”列表示 Dolmino 混合料中该来源所占的比例。该列中的数字大于 100 表示使用了该数据,例如 400 表示重复使用 4 次。"混合料百分比“列描述了Dolmino 混合料中该来源所占的比例,即该列的总和应为 100%。

        我们在 50B 混合上训练 OLMo 2 7B。考虑到更大的批量大小(第 §2.3 节),我们在 OLMo 2 13B 上使用 100B 混合,以确保学习率退火过程中步数相同。

        中期训练模型融合或“Soups”。对使用不同数据顺序训练的多个模型检查点进行简单平均,在计算机视觉(Wortsman 等人,2022)和语言建模(Li 等人,2024)中已被证明有效。我们在六种不同的中期训练混合上验证了这一方法(又称模型融合或“souping”)的有效性,如表 14 所示。所有实验表明,将在同一数据混合的三种排列上退火的三个检查点融合,通常产生的性能等于或优于任何单独训练运行。

表 14 比较了六种训练中期混合模型,这些模型分别使用了最佳单个检查点和在不同数据排列上训练的三个检查点的平均值(混合模型)。所有实验均从70亿个预训练检查点开始;训练中期阶段使用了500亿个 token。混合模型始终能达到或超过在相同混合模型上训练的最佳单个检查点的性能。

        基于此证据,我们广泛使用模型融合以获得最终的 OLMo 2 7B 和 13B 模型。对于 OLMo 2 7B,我们平均三次在 Dolmino Mix 1124 的 50B 样本上训练的检查点。对于 OLMo 2 13B 和 32B,我们平均四个检查点:三个在 100B 样本上训练,一个在 300B 样本上训练;我们发现这一方法在经验上优于仅平均三次 100B 训练运行。

五 深度解析:后训练流水线

        为了将 OLMo 2 适配下游生成任务,我们遵循 Tülu 3 配方(Lambert 等人,2024),更加关注宽松许可和超参数的适当调整。Tülu 3 方法包括三个训练阶段:监督微调(SFT)、使用直接偏好优化(DPO;Rafailov 等人,2024)及策略内偏好数据的偏好调优,以及最终的可验证奖励强化学习(RLVR)。我们发现 Tülu 3 配方的所有阶段都可轻松迁移到 OLMo 2 模型。本节重点介绍 7B 和 13B 模型的开发过程,其中 1B 和 32B 模型采用了非常相似的配方。

        监督微调(SFT)。OLMo 2-Instruct 的 SFT 训练依赖于选择最高质量的现有指令数据集,并使用 PersonaHub 方法(Chan 等人,2024)生成的扩展合成数据进行补充。我们开发了两种 SFT 混合:tulu-3-sft-olmo-2-mixture,用于 7B 和 13B 模型;tulu-3-sft-olmo-2-mixture-0225,包含轻微修改,用于 1B 和 32B 模型。

        对于 tulu-3-sft-olmo-2-mixture,由于 OLMo 2 并未针对多语言任务训练,我们尝试在 SFT 阶段移除所有多语言数据。当移除 Tülu 3 中的整个 Aya 分割和 Wildchat 的多语言样本时,平均性能下降约 0.5 分,这表明 Tülu 3 数据集是平衡的,不能通过移除无关子集轻易改进。该 SFT 混合总共包含 939,104 个提示。

        对于 1B 和 32B 混合 tulu-3-sft-olmo-2-mixture-0225,我们进一步过滤掉包含合成数据生成过程日期截止提及的指令,因为我们注意到这与不良行为(如虚构日期截止和以 “作为 AI 语言模型…” 开头的回答)相关。我们还使用多数投票来提高合成数学问题答案的质量,即防止在错误数学答案上进行 SFT。对于 Tülu 3 中的 Persona MATH 和 Grade School Math 数据集,我们仅保留模型在 5 个完成中达成多数投票的提示和完成内容。该 SFT 混合总共包含 866,138 个提示。

        使用 DPO 的偏好微调(PreFT)。Tülu 3 流水线中 PreFT 的核心策略是基于并扩展 UltraFeedback 流程(Cui 等人,2023),为目标领域生成合成偏好数据。我们通过从部分开发中的 OLMo 2 SFT 模型(7B 和 13B)采样响应引入策略内数据,每个模型使用独立数据集。

表 15 OLMo 2 Instruct 评估体系(改编自 Lambert 等人(2024)):评估套件中开发(上方)和未见(下方)部分的设置。CoT 表示使用链式思维提示(Wei 等人,2022)运行的评估。# shots 为评估模板中的上下文示例数量。Chat 表示在提示模型时是否使用聊天模板。Multiturn ICL 表示将每个上下文示例作为对话中的单独轮次呈现(仅在使用聊天模板且 # Shots 不为 0 时适用)。*表示在多个子评估上的平均值——安全评估的完整细节见 Lambert 等人(2024)。

        从 Tülu 3 开始,我们更新了模型池,使其仅包含具有可许可授权的模型,如附录中的表 25 所示。我们对 Tülu 3 中用于 DPO 的具体提示进行了轻微调整——我们从表 27 列出的多个来源获取提示,最终得到包含 7B 的 36.67 万条提示和 13B 的 37.77 万条提示的数据集。在获得这组提示后,我们从包含不同家族和规模的 20 个模型的模型池中生成回应。

        为了创建合成偏好数据,我们使用 GPT-4o-2024-08-06 作为语言模型评审(LM judge,Zheng 等人,2023),并提示其基于有用性、真实性、诚实性和指令遵循等方面对生成内容进行评分。随后,我们遵循 Argilla 的方法22 对这些方面的评分进行二值化处理:我们获取所有方面的平均评分,将评分最高的回答作为“选择的回应”,并从剩余回答中采样得到“被拒绝的回应”。

        1B 和 32B 的 DPO 模型也使用相同的 on-policy 方法进行训练。

        可验证奖励强化学习(RLVR)。RLVR 是一种新颖的微调技术,可用于特定领域,在这些领域中可以构造具有可验证答案的提示。例如,对于数学问题,强化学习算法 PPO(Schulman 等人,2017)仅在答案正确时才会获得奖励。更多细节参见 Lambert 等人(2024)。

        在偏好微调之后,我们使用 7B 和 13B 的 on-policy 偏好数据集训练了 7B 与 13B 的奖励模型。接下来,我们将 RLVR 应用于表现最佳的 7B 和 13B DPO 检查点,使用的组合数据集包含 GSM8K、MATH 训练集以及来自 Lambert 等人(2024)的带约束提示。对于 RLVR,我们从相应的奖励模型(RMs)初始化 PPO 的价值函数,该方法已被证明有助于提升评估中的平均分数(Lambert 等人,2024)。在对 13B 模型进行首次 RLVR 训练后,我们观察到其 GSM8K 和 MATH 表现低于此前的开发 instruct 模型。因此,我们又进行了两次额外的 RLVR 训练:首先在 GSM8K 训练集上,然后在 MATH 训练集上。在 RLVR 阶段结束时选中的模型构成了最终的 OLMo 2 Instruct 模型。

        对于 1B 和 32B 模型,我们使用 Group Relative Policy Optimization(GRPO;Shao 等人,2024)执行 RLVR,该方法不需要奖励模型。32B 模型的评估指标如图 14 所示。

        超参数选择。我们对 7B 和 13B 模型执行了以下超参数搜索。在每个阶段,我们首先使用 1 个随机种子进行试验以确定大致配置,随后最多使用 4 个随机种子对最终超参数进行验证。最终超参数用(♥)标记:

  1. SFT:7B:学习率搜索为 1 × 10⁻⁵、2 × 10⁻⁵ (♥)、3 × 10⁻⁵,13B:学习率搜索为 1 × 10⁻⁶、4 × 10⁻⁶、5 × 10⁻⁶ (♥)、7.5 × 10⁻⁶、8 × 10⁻⁶

  2. DPO:对两个模型分别搜索:5 × 10⁻⁷、6 × 10⁻⁷、7 × 10⁻⁷、8 × 10⁻⁷ (♥ - 13B)、1 × 10⁻⁶ (♥ - 7B)

  3. RM:7B 和 13B 均使用 3 × 10⁻⁶ 的学习率和 1 个随机种子进行训练。

  4. RLVR:β 值搜索:0.03、0.05、0.07 (♥ - 7B)、0.1 (♥ - 13B),13B 模型额外搜索学习率:3 × 10⁻⁷(♥ - 13B)、4 × 10⁻⁷(♥ - 7B)

        我们使用较早的开发检查点对 SFT 和 DPO 进行超参数搜索,其结果见表 17 和图 12。一个关键发现是:与 Lambert 等人(2024)描述的 Llama 3.1 训练策略相比,OLMo 2 需要显著更高的学习率。最终模型的优化超参数列于表 17 和表 18 中。

        32B 模型的后训练是在 7B 和 13B 模型发布之后进行的,因此其超参数选择独立进行。SFT:搜索 1×10⁻⁶、2×10⁻⁶、3×10⁻⁶、4×10⁻⁶、5×10⁻⁶,最佳为 4×10⁻⁶,并额外跑了一个种子验证表现。DPO:搜索 8×10⁻⁷、1×10⁻⁶、1.5×10⁻⁶、2×10⁻⁶、2.5×10⁻⁶,最佳为 2×10⁻⁶。RLVR:由于采用 GRPO,不需要奖励模型。最终模型使用 5×10⁻⁷ 的学习率、KL β = 0.1,每个提示采样 16 个样本。

表 17 列出了针对 7B SFT 检查点尝试的
超参数配置,所有配置均基于最终模型使用
的同一数据集。SFT 模型采用有效批大小 128、
线性学习率和 0.3 的预热比进行训练。
图 12 显示了基于不同学习率,在开发 SFT
检查点上训练的 DPO 检查
点的平均得分。平均值不包括安全检查点。

      

         OLMo 2-Instruct 的评估。遵循 Tülu 3(Lambert 等人,2024),我们在表 15 所列的五个类别上评估 OLMo 2-Instruct。尽管 Tülu 3 还包含代码相关任务,但我们排除了该类别,因为代码并非 OLMo 2 的目标技能。对于剩余的类别,我们使用与 Tülu 3 中相同的评估。表 15 也展示了每项评估的设置与指标,这些设置与 Lambert 等人(2024)针对非代码任务的建议完全一致。

        表 16 展示了 OLMo 2 Instruct 各训练阶段的性能。与其它相似规模的开源模型的对比结果见表 7。图 13 与图 15 分别展示了 13B 和 7B 模型的训练轨迹及关键指标。

图 13 我们针对使用 RLVR 训练的 OLMo-2-1124-13B-Instruct 在评估套件上的得分。我们首先在 GSM8K、MATH 和带约束提示的数据混合上训练 OLMo-2-1124-13B-RLVR1,但注意到 GSM8K 得分低于预期。随后,我们在 GSM8K 上继续训练 OLMo-2-1124-13B-RLVR2,并观察到更高的 GSM8K 得分。最后,我们仅在 MATH 上训练 OLMo-2-1124-13B-Instruct,并观察到 GSM8K 和 MATH 得分进一步提升。需要注意的是,在每一次 RLVR 运行中,价值函数都会从奖励模型重新初始化。每次 RLVR 运行的完整学习曲线可在附录 C.2 中找到。

        OLMo 2-Instruct 模型表现与领先的开源权重模型相当。特别是,OLMo 2 13B Instruct 的结果接近 Qwen 2.5 14B Instruct,同时在性能基准上超过了 Tülu 3 8B 和 Llama 3.1 8B Instruct。RLVR 阶段在两个模型规模上均表现稳定有效,随着训练奖励信号增强,评估结果显著提升。

        最后,我们在 Lambert 等人(2024)的未见评估套件上评估 OLMo 2-Instruct,但不包含代码相关评估。Instruct 在未见评估套件上的得分如表 24 所示。

图 14 展示了使用 RLVR 训练的OLMo-2-0325-32B-nstruct 模型在核心指标评估套件中的得分。为了提高这些得分,我们使用GSM8K、MATH 和带有约束的提示数据集混合训练OLMo-2-0325-32B-Instruct 模型
图 15 上排展示了 OLMo-2-1124-7B-Instruct 模型在可验证奖励、KL 散度和响应长度方面的训练曲线。下排的y轴分别展示了我们评估套件的平均得分以及 GSM8K、IFEval和 MATH Flex 的得分。总体而言,我们发现 RLVR 不仅提高了我们 78 模型的训练奖励,也提高了下游评估(例如 GSM8K)的得分。

六 深度解析:作为研究催化剂的基础设施

        语言模型(LM)训练众所周知计算密集。训练大型模型需要最先进的硬件,同时也需要大量工作来让其高效运行。效率提升可以转化为更高的 Token 数或更多的参数量,从而直接影响最终模型的质量。GPU 是这一基础设施的核心,但要使其发挥最高性能,还需要在其他流程和系统上投入。数据中心需要计算节点之间的高速互连,以确保昂贵的 GPU 不会因为等待数据而空闲。训练任务需要访问大量快速、可靠的存储来读取训练数据。GPU 的故障率高于大多数其他硬件,而一次训练可能需要数千块 GPU 同时参与,因此有效的监控和替换策略至关重要。本节提供了我们用于支持 OLMo 2 工作负载的硬件和软件投入的详细信息。

6.1 集群

        OLMo 2 在 Ai2 的两个集群——Jupiter 和 Augusta——上进行训练。尽管硬件和架构存在差异,但两者都提供了足够的训练吞吐量。Beaker(Ai2 的工作负载管理系统)允许研究人员在集群之间迁移任务,7B 和 13B 模型均部分在两个集群上训练,其中 7B 的主要训练在 Jupiter 上完成,而 13B 的主要训练在 Augusta 上完成。

6.1.1 Jupiter

        Jupiter 是一个位于美国德克萨斯州奥斯汀市的 128 节点 GPU 集群,由 Cirrascale Cloud Services 运营。

        Compute。集群由 1,024 个 NVIDIA H100 GPU 组成,每个 GPU 具有 80GB HBM3,功率 700W。GPU 分布在 128 台服务器上,每台服务器配置如下:2× Intel Xeon Platinum 8468 CPU,2 TB DDR5 系统内存,18 TB 本地 NVMe 存储。

        Storage。服务器通过 800 Gbps 本地网络连接至一个 WEKA 高性能存储集群(脚注 24):1 PB NVMe SSD(11 台物理服务器),5 PB HDD(12 台主机)。Jupiter GPU 服务器配备 两张绑定的 25 Gbps Mellanox 网卡,单主机总吞吐达 50 Gbps。基准测试中,使用 64 台机器可达到 761 Gbps 的读写吞吐量。

        Interconnect。跨节点 GPU 通信通过 InfiniBand 上的 RDMA 进行,并使用一个 两层铁路优化(2-Tier Rail Optimized)、平衡、全双倍带宽(full-bisected)的网络架构(Wang et al., 2023a)。每台服务器有 8× 400 Gbps InfiniBand 卡,单主机总吞吐达 3200 Gbps。此设置使 Ai2 可以在同一个集群上同时运行数十个分布式训练任务,而无需拓扑调度。

        Cooling。Jupiter 服务器装载在 Dynamic Density Cabinets(脚注 25)中,每个机柜包含 5 台服务器,具有独立的冷却与供电系统。机柜为封闭系统,通过顶部隔层循环空气并使用水冷降温,使数据中心达到 1.2 的 PUE(能源效率)。高负载下,H100 GPU 峰值温度约为 75°C,平均温度在 60°C–65°C

6.1.2 Augusta 集群

Augusta 是一个由 Google Cloud 提供的 160 节点 GPU 集群,物理服务器位于美国爱荷华州 Council Bluffs。

        Compute(计算)集群由 **A3 Mega 虚拟机(VM)**组成,每个 VM 配置:8× NVIDIA H100 GPU(脚注 26)

        Storage(存储)。Augusta 使用 Google Cloud Storage,每个 VM 的存储带宽高达 1 GB/s。我们通过抽象文件/对象 API,确保在不同存储后端之间的可移植性。

        Interconnect。每块 GPU 均配备一个独立的 Ethernet NIC。通过以下技术实现高速跨节点 GPU 通信:GPUDirect-TCPXO,gVNIC,紧凑节点布局。还利用了:Google Jupiter 数据中心网络技术,Titanium 系统:分层卸载 + 全带宽可重构光链路,其带宽类似于非阻塞网络结构。

        Cooling。Augusta 服务器通过空气冷却,其所在数据中心 12 个月滚动平均 PUE 为 1.12

6.2 Beaker

        OLMo 2 的任务使用 Beaker(Guerquin, 2022)进行调度,这是一个自定义的工作负载管理系统。Beaker 为 OLMo 2 提供两个关键优势:

        Portability。Beaker 的架构可在不同数据中心之间调度 GPU 任务,只需很少的代码改动。它可以在任何运行单个 Linux 守护进程的系统上运行,该守护进程被打包为静态链接二进制文件。通常,只需修改一行代码即可将任务从一个位置迁移到另一个位置。

Isolation。Beaker 的任务均运行在容器中,提供隔离能力,使 OLMo 2 任务能够与其他任务同时运行,且互不干扰。容器化带来:独立的软件环境与依赖,更少的冲突,重现性:可重放数月前的实验。对比之下,常见的 Slurm 环境所有任务共享同一操作系统与 CUDA 环境,导致系统升级后难以重现过去的实验。此外,Beaker 的简洁操作让小团队也能快速接入新的计算资源。

6.3 稳定性与运维

        两个集群都经历了初期测试和压力验证阶段,在此期间我们发现并修复了从电缆松动到 NCCL 节点顺序不当等问题。这需要与硬件厂商的紧密合作,并对系统稳定性至关重要。测试期后,两套集群具有大致相同的可靠性。

        GPU 健康检查。Beaker 会在每次调度任务前运行一个简单程序——尝试矩阵乘法。若 GPU 发生错误:节点将被隔离(cordon),任务重新调度,故障节点进入隔离区待修复这减少了人工干预需求,同时允许训练任务在遇到错误时自动重启。

        Cordoning。Beaker 也允许手动隔离节点。隔离节点不会再参与调度,并会被标记为需维修。整个集群用户共同参与“众包”发现故障节点。

        Active Monitoring。Beaker 执行行业标准的监控与告警,并有运维流程及时处理系统问题。

6.4 最大化硬件利用率

        Ai2 的硬件基础设施(§6.1)需要配套高效的模型训练软件,确保资源使用最大化。效率提升不仅能训练更大模型、更多 Token,还能降低环境影响(§6.5)并提高实验速度。由于 OLMo 项目不是 Ai2 唯一的项目,我们必须合理使用资源,以减少对其他团队的影响。

        以下介绍数个对 OLMo 训练时间有显著影响的 PyTorch 优化(脚注 27),并未观察到收敛速度的损失:

        利用编译。torch.compile() 是 PyTorch 中的一个函数,它会将原生的 PyTorch 模块和函数编译为优化过的内核,从而实现显著的吞吐提升与 GPU 内存节省。其方式是避免在顺序调用各个 PyTorch 操作时产生的 Python 端开销,并减少必须在 GPU 上发生的读写次数。因此,在正确使用 torch.compile() 时,它在许多情况下可以有效匹配手工编写内核的性能,而无需额外的复杂性和工程投入(Ansel et al., 2024)。

        最小化主机-设备同步。默认情况下,GPU 操作是异步的。使用 GPU 的函数调用会被放入某个设备的队列中,但不会立即执行。这允许系统并行执行更多计算,包括 CPU 上或其他 GPU 上的操作……—— PyTorch: CUDA Semantics

        任何强制主机-设备同步(host-device sync)的训练代码,都会在队列中所有操作完成之前,阻止进一步操作排入队列。这些同步点会阻碍性能,并且很容易在无意间引入。
        令人惊讶的是,有相当多的操作会导致主机-设备同步:

  1. 从 CPU 同步拷贝张量到 GPU。例如使用 tensor.to(device="cuda") 会触发一次主机-设备同步。可以通过使用异步拷贝避免此问题,例如:tensor.to(device="cuda", non_blocking=True)

  2. 从 GPU 拷贝张量到 CPU。此操作无法安全地异步进行,因此应尽可能避免 GPU → CPU 的数据传输。许多貌似无害的代码都会造成 GPU → CPU 传输,从而导致主机-设备同步,例如:打印 CUDA 张量,或依赖 CUDA 张量布尔值的 if ...: 代码块。

  3. 某些特定的 PyTorch 操作。例如 masked_select(),可能会意外地引发主机-设备同步。29

        可以通过在训练循环开始前调用torch.cuda.set_sync_debug_mode("warn")来检测主机-设备同步。这将使 PyTorch 在同步发生时发出警告。此机制是尽力而为,因此仍可能遗漏部分同步。

        使用独立后端进行异步簿记。典型的训练循环包含定期的“簿记”操作,如记录指标(logging metrics)和保存检查点(checkpointing)。虽然这些操作可能相对较快,但在整个训练过程中它们的累计成本很可观,并且通常涉及主机-设备同步。例如,训练指标(如交叉熵损失)源自 GPU 计算,首先以 CUDA 张量形式出现,因此将其输出到控制台会强制一次同步。

        这些操作大多是必要的,无法完全避免,但可以将簿记操作放到单独线程中异步进行,从而最小化阻塞训练循环的时间。然而,PyTorch 的 NCCL 后端并不是线程安全的。为解决此问题,我们设置了一个不依赖 NCCL(如 GLOO)的独立后端,并专门用于簿记操作。簿记工作流程如下:

  1. 指标收集与日志记录:先决定日志记录的间隔(例如每 10 步或每 50 步,而不是每一步)。每一步中,在原 GPU 上计算并存储指标(保存在 GPU 张量中)。当达到日志间隔时,再将指标复制到 CPU(触发一次同步),然后传递给簿记线程。该线程使用自己的 PyTorch 后端聚合指标并进行日志记录。

  2. 保存检查点:到达保存点时,将模型与优化器状态复制一份到 CPU 内存(触发同步)。然后将这份拷贝传递给簿记线程,由此线程从各计算节点存储的模型分片中组装模型并写入磁盘,而主线程可继续训练30。

        在这两种情况下,对训练时间的唯一影响是一次主机-设备同步及 GPU→CPU 的数据拷贝。事件频率可配置,因此对整体训练时间的影响可忽略不计。

        显式 Python 垃圾回收。在训练期间,默认的 Python 垃圾收集器会周期性地运行一次回收。在分布式环境中,有数千个训练进程需要保持同步,而没有任何机制保证垃圾回收会在所有进程中同时发生。由于分布式训练由最慢的进程决定,这会导致平均训练时间变长且方差变大(如图 16 所示)。随着进程数量增加,这些问题会更加严重。

        为解决此问题,OLMo 2 的训练器会关闭自动垃圾回收(例如调用 gc.disable()31),然后在所有进程中按照固定的间隔统一触发垃圾回收(例如调用 gc.collect(1)32)。

图 16 展示了在 8 个节点上运行两个 OLMo-1B 模型,经过 1000 步训练后,每个设备的每秒令牌数(TPS)的训练吞吐量。其中一个模型启用了自动垃圾回收(GC),另一个模型则按照训练代码库中设定的时间间隔手动进行垃圾回收。启用自动垃圾回收后训练吞吐量会降低目稳定性更差,并且通常会随着运行的进行而变得更糟。

6.5 环境影响

        依据我们在 Groeneveld et al. (2024) 中的分析以及先前文献(Patterson et al., 2021;Dodge et al., 2022;Luccioni et al., 2022;Li et al., 2023a),我们首先通过计算预训练期间消耗的总能量并乘以当地电网的碳强度,来估算训练最终模型所造成的环境影响中的碳排放量。此外,我们进一步将此前的分析扩展到对用水量的估算,通过将消耗的能量乘以电力生产和冷却设备的水使用效率来计算。与 Groeneveld et al. (2024) 一样,我们强调,尽管我们的报告方式符合标准做法,但并未纳入硬件在制造、运输及最终处置过程中所产生的隐含排放与用水量,以及模型生命周期后续阶段(如部署与推理)所带来的环境影响,因此我们的估算应视为下界。我们在表 19 中报告了模型的详细结果。

表 19 预训练期间的 CO₂ 排放与用水量。
我们使用来自数据中心服务商的 PUE(电源使用效率)信息、各数据中心所在地区电网的碳强度数据和 WUE(水使用效率),以及整个训练过程中记录的时间序列总功耗数据,来估算我们新模型的碳排放总量与用水总量。Llama 2(Touvron et al., 2023)、Llama 3(Grattafiori et al., 2024)以及原版 OLMo(Groeneveld et al., 2024)的数值来自其各自论文。我们同时展示了 Llama 2 与 Llama 3 的模拟用水量,范围覆盖了在 OLMo 模型中使用的最低与最高 WUE 值所产生的用水量区间。

        与 Groeneveld et al. (2024) 相同,我们通过以下方法计算每个模型的总功耗:每 25 毫秒测量单个节点的功耗,计算整个训练过程的平均功耗,再乘以节点总数。随后,我们将该数值乘以训练该模型所使用的数据中心的电源使用效率(PUE),以反映数据中心整体的能源效率。由于 OLMo 2 7B 的绝大部分训练是在 Jupiter 集群上完成的,我们在对 7B 模型进行分析时使用 Jupiter 的效率指标。而 OLMo 2 13B 则是在 Augusta 集群训练,因此我们使用其效率指标。我们估计,预训练 OLMo 2 7B 与 13B 共消耗约 391 MWh 能量。

        为计算碳排放,我们将总功耗乘以各数据中心所在地点的碳强度系数,单位为每千瓦时 kg CO₂。Jupiter 集群由 Austin Energy 供电,该机构最新报告的碳强度为 0.332 kg CO₂/kWh。³³ Augusta 集群位于爱荷华州,而该州的平均碳强度为 0.352 kg CO₂/kWh³⁴,我们据此进行计算。我们估计,训练最新模型共排放约 154 吨二氧化碳当量(tCO₂eq)。

        为计算用水量,我们将总功耗乘以电力生产以及现场冷却设备的水使用效率(WUE)。两个集群均使用高效的闭环冷却设备,因此我们假设现场冷却的 WUEₒₙₛᵢₜₑ 为 0 升/kWh。根据 Reig et al. (2020),我们假设 Jupiter 集群的电力生产 WUEₒffₛᵢₜₑ 为 1.29 L/kWh,而 Augusta 集群为 3.10 L/kWh。我们估计训练最新模型消耗了约 110 万升水。

        虽然我们希望呈现尽可能全面的模型训练环境影响分析,但我们强调,这仍是对训练大型模型真实环境成本的下界估计。在即将发表的论文(Morrison et al., 2025)中,我们将提供更全面的分析,涵盖模型开发、预训练和部署过程中的能耗、排放与用水情况。

七 结论

        我们介绍了 OLMo 2 与 OLMo 2-Instruct,这是一个包含 7B、13B 和 32B 参数规模、在最多 6T tokens 上训练的全开放语言模型家族。基础模型与指令模型均在相同参数规模的开放权重模型(如 Qwen 2.5、Gemma 2 和 Llama 3.1)中表现具有竞争力。我们详细说明了构建高性能语言模型所需的大量工作——这些工作与原版 OLMo 在许多方面都不同——包括稳定的基础设施、提升训练稳定性的架构改进、后期训练数据的创新、最新的后训练技术以及更多细节。我们发布了所有训练与评估代码、数据集、检查点及日志,以便复现与扩展这些模型。OLMo 2 标志着开源语言模型的持续进步,构建了一个新的研究生态系统,在这个生态系统中,新的训练方法与技术需要被理解与分享。

参考文献

1 OLMo2 : https://arxiv.org/pdf/2501.00656

Logo

中国智能体开发者社区,聚焦智能体与大模型开发,提供前沿资讯、实用工具链、开源项目及行业案例。通过技术沙龙、开发者大赛等活动,促进经验交流与协作,助力开发者快速构建创新智能应用。

更多推荐