4 预训练

4 预训练

本文为论文《SmolLM2：紧凑高效的小语言模型家族》「中文全文翻译」板块的对应章节，与「论文原文」板块同一章节一一对应，层级保持一致。

语言模型预训练的近期趋势显示出向显著更长的训练周期转变，尤其是对于小型模型（Yang et al., 2024a; Yang et al., 2024b; AI@Meta, 2024b）。尽管这一策略偏离了 Chinchilla-最优指南（Hoffmann et al., 2022），但带来的性能提升和降低的推理成本使延长训练成为一项值得的权衡（de Vries, 2023）。例如，Qwen2-1.5B 在 7 万亿个 token 上训练，Qwen2.5-1.5B 在 18 万亿个 token 上训练，而 Llama3.2-1B 源自一个裁剪的 8B 模型，使用蒸馏在 9 万亿个 token 上训练（Yang et al., 2024a; Yang et al., 2024b; AI@Meta, 2024b）。

在构建 SmolLM2 时，我们在 11 万亿个 token 上进行训练（约在我们收集的数据集上两个时代），采用多阶段训练方法，而非在整个预训练期间使用固定的数据集混合比率。这种设计由四个关键原则指导：（1）由性能驱动的干预，即我们监控关键基准上的评估指标，并调整数据集混合比率以解决特定能力瓶颈；（2）在退火阶段对高质量数学和代码进行过采样，保留 FineMath 和部分 Stack-Edu 等数据集到最后阶段以最大化其影响（Blakeney et al., 2024; Ai2, 2024）；（3）在训练中期战略性地引入中等规模的数据集，如 OWM、InfiMM-WebMath 和 Stack-Edu，以避免早期被更大的数据集稀释；以及（4）避免过度数据重复，遵循 Muennighoff et al. (2023)，我们尽量使大多数数据集接近推荐的 4-5 个时代阈值。

虽然进行多次从头开始的训练运行来探索不同的数据混合计划可能很有成果，但预训练 SmolLM2 的高昂成本（约 25 万美元的 GPU 计算费用）促使我们采用“在线”方法。

在接下来的各节中，我们描述训练过程的每个阶段，详细说明数据集混合比率、选择背后的理由以及指导我们干预的观察。虽然某些决定由文献中已确立的发现告知，但其他决定则由训练期间收集的经验性见解驱动。四个预训练阶段的数据混合比率在图 2 中提供。

4.1 训练设置

我们的基础模型包含 17 亿个参数，遵循 LLama2（Touvron et al., 2023）架构，如附录 A 所述。我们使用 nanotron 框架在 256 块 H100 上训练模型，并使用 AdamW 优化器，采用 $(\beta,\beta_{2})=(0.9,0.95)$ 与 Warmup Stable Decay（WSD）（Hu et al., 2024; Zhai et al., 2022）学习率调度，以避免设置固定的训练时长（参见图 3，附录 A）。调度从 2000 步的预热阶段开始，维持峰值学习率 $5.0\times 10^{-4}$（稳定阶段），并在需要时可转变为衰减阶段，在总训练步骤的 10% 期间将学习率降至零（Hägele et al., 2024）。

我们使用 Allal et al. (2024) 的分词器，其词汇表大小为 49,152 个 token，在 70% FineWeb-edu、15% Cosmopedia-v2、8% OpenWebMath、5% StarCoderData 和 2% StackOverflow 的混合物上训练。

图 2：训练阶段间的数据集混合比率。详细描述在第 4 节。x 轴表示训练 token 数量。

4.2 稳定阶段：阶段 1

数据混合

在 SmolLM2 预训练的第一阶段（0 到 6T 个 token），我们根据英文网络消融结果和现有文献的洞察，设计了数据集混合比率。我们对网络数据采用 60% FineWeb-Edu 和 40% DCLM 的比例（在第 2.2 节中讨论），这在教育内容与多样的现实世界问答式数据之间提供了最佳平衡。对于代码数据，遵循 Aryabumi et al. (2024)，我们纳入了 StarCoderData，包含跨 80 种编程语言的 2500 亿个 token，并将其限制在总混合物的 10%，以确保在 11T 个 token 上约 4 个时代，并为后续阶段的过采样留出空间。由于我们的数学数据集规模相对较小，我们在阶段 1 中未包含数学数据。

发现

在 6T 个 token 的训练后，我们按表 3 所示在关键基准上评估了 SmolLM2。知识和推理性能与基于我们英文网络消融结果的预期一致。然而，我们观察到编码和数学性能总体较差。

4.3 稳定阶段：阶段 2

数据混合

对于阶段 2（6T 到 8T 个 token），我们将 OWM 以 5% 的比例加入混合物，并增加了代码数据的比例，希望在保持强知识保留的同时解决观察到的编码和数学推理差距。以较低百分比纳入 OWM 反映了该数据集的小规模（120 亿个 token）以及我们逐步纳入数学内容的方法。阶段 2 的最终混合物由 75% 英文网络数据（保持阶段 1 的 60/40 FineWeb-Edu 与 DCLM 比例）、20% 代码数据和 5% 数学数据组成，如图 2 所示。

发现

阶段 2 后，代码性能在大多数语言上提高，验证了对 StarCoderData 进行过采样的决定。OWM 的整合对数学性能没有显著影响，强调了需要在后续阶段引入更大、更高质量的数学数据集。除了代码和数学性能之外，如图 6（第 E.2 节）所示，我们观察到使用多选题表述（MCF，即明确输出“A”、“B”、“C”或“D”选项，而非像填空表述那样计算不同答案的可能性）时，MMLU 准确率高于随机水平（>25%）。这与先前工作形成对比，先前工作显示小型模型在 MCF 上挣扎（Gu et al., 2024; Du et al., 2024），并表明小型模型的长期训练可以使其获得通常与大型模型相关的能力（Blakeney et al., 2024; Gu et al., 2024; Du et al., 2024）。为进一步优化 MMLU 性能，我们重新审视了英文数据集混合，通过额外的退火消融实验发现，在此阶段将 DCLM 相对于 FineWeb-Edu 的比例提高，可略微提升 MMLU MCF。

表 3：各训练阶段后模型在不同基准类别上的平均表现。阶段 1-3 处于稳定阶段（无衰减）。完整的每基准结果在第 E 节。

4.4 稳定阶段：阶段 3

数据混合

在稳定阶段的第三个也是最后一个阶段（8T 到 10T 个 token，退火开始前），我们加入了 InfiMM-WebMath 的文本-only 英文部分与 OWM，将数学数据的总比例提升至约 10%，如图 2 所示。对于英文网络数据，我们重新审视了消融结果，并将 FineWeb-Edu 与 DCLM 的比例调整为 40/60。对于代码，我们将 StarCoderData 替换为 Stack-Edu（第 3.4 节）。对于 Stack-Edu 中少于 40 亿个 token 的语言（TypeScript、Shell、Swift、Go、Rust 和 Ruby），我们使用其 StarCoder2Data 子集。我们还加入了来自 StarCoder2（Lozhkov et al., 2024）的 Jupyter Notebooks，它提供代码与解释交织的丰富上下文示例，增强了模型在编程任务上的推理能力。

发现

虽然这些新数据集的整合在多个基准上带来了改进，但我们在此阶段观察到明显的损失峰值，即使在回滚训练并跳过与峰值相关的数据后仍然存在（Chowdhery et al., 2023; Almazrouei et al., 2023）。确切原因尚未确定，但大多数评估指标在阶段结束时恢复。

4.5 衰减阶段：阶段 4

数据混合

最终阶段包括将学习率线性降至 0，持续总训练时长的 10%（从 10T 到 11T 个 token）（Hägele et al., 2024）。遵循 Blakeney et al. (2024)，我们引入了我们最高质量的数学数据集，InfiWebMath-3+ 和 FineMath 4+。我们还将混合物的 0.08% 分配给 OWM，0.02% 分配给 AugGSM8K（Li et al., 2024a），这是 GSM8K 基准训练集的增强版本，已成为近期预训练数据集的常见组成部分（Achiam et al., 2023; Dubey et al., 2024; Ai2, 2024）。总体而言，数学内容占混合物的 14%。我们将 Stack-Edu 扩展到包含阶段 3 未覆盖的更多编程语言，并将该数据集的贡献设为混合物的 24%。我们保持编程语言的自然分布，Python 分配更高。剩余混合物由 58% 的英文网络数据组成（保持更高的 DCLM 与 FineWeb-Edu 比例）以及 Cosmopedia v2（Allal et al., 2024）的 4%，后者提供 300 亿个高质量合成教科书、博客文章和故事。

发现

虽然所有基准任务在阶段 4 后都有改进，但我们观察到编码性能和，尤其是，数学性能有实质性提升，验证了我们专门针对这些领域的数据混合物。

4.6 上下文长度扩展

为支持长上下文应用，我们遵循标准做法（Gao et al., 2024），将上下文长度从 2k 个 token 扩展到 8k 个 token，通过取阶段 4 的一个中间检查点（在最终训练的前 750 亿个 token 之前），并使用不同的数据集混合物和 130k 的 RoPE 值继续训练。混合物调整为包含 40% 长上下文文档（8k 个 token 或以上），来源于 DCLM（10%）、FineWeb-Edu（10%）和 Dolma 的书籍子集（20%）（Soldaini et al., 2024），其余 60% 遵循阶段 4 的混合物。完成此步骤后，我们得到最终的 SmolLM2 基础模型。

4.7 基础模型评估

我们与现有同等规模的最先进模型 Qwen2.5-1.5B（Yang et al., 2024b）和 Llama3.2-1B（AI@Meta, 2024a）对比评估并比较最终的 base SmolLM2 模型，覆盖广泛的基准。除非另有说明，评估使用 lighteval 并在零样本设置下进行。

表 4 的评估结果显示 base SmolLM2 的强劲表现，在 HellaSwag 和 ARC 上优于 Qwen2.5 基础模型。SmolLM2 还在训练期间未监控的 held-out 基准上提供强劲表现，如 MMLU-Pro（Wang et al., 2024c）、TriviaQA（Joshi et al., 2017）和自然问题（NQ，Kwiatkowski et al., 2019）。值得注意的是的，模型在 MMLU-Pro 上几乎超过 Qwen2.5-1.5B 6 个百分点，进一步验证了其泛化能力。在数学和编码基准上，SmolLM2 展示出竞争性表现。虽然它落后于 Qwen2.5-1.5B，但在 GSM8K、MATH 和 HumanEval 上优于 Llama3.2-1B。重要的是，上下文长度扩展后性能几乎没有下降，而 HELMET（Yen et al., 2024）和“ haystack 中的大针”（NIAH，Kamradt, 2024）结果显示强劲表现——见附录 G。这些结果凸显了我们策划的数据集、数据混合物和训练阶段的有效性。

表 4：SmolLM2 与其他 1-2B 基础模型在基准上的表现对比。SmolLM2 展示竞争性结果，突显其泛化能力。