整体说明

  • 模型增长初始化（Model Growth Initialization，MGI）是一种让大模型在不从头开始训练的前提下迅速“长大”并具备良好初始性能的技术
  • Model Growth Initialization 的核心思想是：先训练一个小模型，然后在深度或宽度上扩展成更大的模型（如增加层数、宽度或专家数量），并把小模型已学到的知识完整复用到大模型里 ，从而显著节省训练成本、提升收敛速度
  • TLDR：Model Growth Initialization 就是“把小模型当预制件，复制粘贴成大模型，再微调”，用最小的算力让大模型站在小模型的肩膀上起跑
  • 一个通俗的比喻：可以把模型训练比作盖楼：
    • 传统做法 ：平地起高楼，从地基开始逐层盖（随机初始化）
    • MGI 做法 ：先盖一栋“小楼”并装修完毕，然后把整栋小楼连同装修一起复制+堆叠 ，瞬间变成一栋“大厦”，再只对新增部分做微调

Background

  • 直接训练千亿参数的大模型成本极高（如 Llama-3 需消耗 770 万 GPU 小时），而 Model Growth Initialization 通过复用小模型的知识，显著降低计算成本
    • 例如，使用已训练的 7B 模型初始化 30B 模型时，可减少约 50% 的训练时间

结构扩展策略

  • 模型扩展通常分为三类：
    • 深度扩展 ：增加 Transformer 层的数量
    • 宽度扩展 ：增加神经元数量、头数或 FFN 维度（如 Net2Net 通过复制神经元并均分权重）
    • 混合扩展 ：同时调整深度和宽度（在扩展层数的同时增加隐藏层维度）

参数初始化原则

  • 扩展后的模型需保持与原模型行为一致 ，避免训练震荡
  • Model Growth Initialization 的初始化原则是确保模型在结构扩展（如增加层数、宽度或专家数量）后，既能继承原小模型的知识，又能保持训练稳定性（避免梯度爆炸/消失或性能骤降），同时为新参数提供合理的学习起点
  • 这些原则的核心逻辑是：在“继承知识”与“学习新能力”之间找平衡
    • 通过功能保留和身份映射确保模型初始稳定
    • 通过部分保留和跨层传递实现精准知识迁移
    • 通过优化器状态一致保证训练连续性

Function-Preserving Initialization, FPI 原则

  • 让扩展后的大模型在初始状态下，对任意输入的输出与原小模型完全一致（或近似一致），实现“无损知识迁移”
  • 避免扩展后模型“忘记”原有的能力，为后续训练提供稳定起点
  • 通过精确的参数复制与调整，确保扩展后的模型计算逻辑与原模型等效
  • 示例：
    • 宽度扩展示例 ：若原模型某层有2个神经元（h1, h2），输出为 y = w1*h1 + w2*h2，扩展到3个神经元时，新增神经元h3的权重复制h2的参数，同时将w2拆分为w2/2和w2/2，使新输出 y = w1*h1 + (w2/2)*h2 + (w2/2)*h3 与原输出完全一致
    • 深度扩展示例 ：新增 Transformer 层时，将其参数初始化为“恒等映射”（如自注意力的输出投影矩阵设为单位矩阵，偏置设为 0），确保新增层对输入不做任何修改，等效于原模型的计算流程（实际上就是后面要介绍的 IMI 方法）

Identity Mapping Initialization, IMI 方法

  • 让新增的层/参数在初始状态下“不干扰”模型原有计算，仅在训练中逐步学习新功能
  • 防止新增结构破坏原模型的优化状态（如损失函数突增），降低训练震荡风险
  • 将新增组件的参数初始化为“中性值”，使其对模型输出的影响为零或极小
  • 示例：
    • 新增 Transformer 层 ：将多头注意力的输出权重矩阵初始化为单位矩阵（确保输入=输出），前馈网络（FFN）的中间层权重设为 0（使 FFN 等效于“跳过连接”）
    • MoE 模型新增专家 ：将新专家的输入/输出投影权重初始化为0，使其在初始阶段不参与计算，仅通过训练逐步被激活
  • 注：“自注意力输出投影矩阵设为单位矩阵、偏置设为 0” 的做法，本身是 Identity Mapping Initialization 的具体操作，因为它直接让该组件成为 “恒等变换”
    • 但当这种操作被用于确保 “扩展后的整体模型与原模型功能一致” 时，它成为实现 FPI 的手段之一
  • 注：FPI（是一个原则） 保证新旧模型输出完全一样，而 IMI（是一种方法） 只保证新增部分是恒等映射，不保证整体输出不变
    • IMI 只保证新增部分是恒等映射，原始参数往往可能也会被修改

Optimizer State Consistency

  • 确保扩展后的模型优化器（如Adam）状态与原模型兼容，避免训练进程中断
  • 使扩展后的模型训练能“无缝衔接”原训练过程，减少重新收敛的时间
  • 复用原模型的优化器参数（如动量、二阶矩估计），并对新增参数初始化合理的优化器状态
    • 对复用的参数，直接继承原优化器的动量值，确保梯度更新方向与原训练一致；
    • 对新增参数，将优化器的动量初始化为 0（或小值），避免其初期干扰整体更新节奏

附录：一些典型方法与技术细节

Net2Net：开创性的结构扩展框架

  • 深度扩展（Net2DeeperNet） ：直接复制Transformer层（如将L层模型扩展为2L层），确保每层输入输出形状一致
  • 宽度扩展（Net2WiderNet） ：新增神经元的权重复制相邻神经元，并调整输出权重使总和不变。例如，原输出为y = e*h1 + f*h2，扩展后变为y = e*h1 + (f/2)*h2 + (f/2)*h3

Stacking Your Transformers：深度堆叠优化

  • G_stack操作符 ：通过堆叠多个小模型（如 7B 到 70B），使大模型在 194B tokens 即可收敛到传统 300B tokens 的损失，速度提升 54.6%
  • 增长规划（Growth Schedule） ：分阶段扩展模型，例如先训练 16B 模型，再逐步扩展至 101B，同时调整学习率和优化器状态

MoE 模型的扩展

  • Mixtral-8x7B ：从 Mistral-7B 初始化，直接复用其 FFN 层作为专家，并通过微调不同任务的 FFN 生成多样化专家
  • 参数共享与隔离 ：专家层共享底层编码器参数，但各自保留独立的前馈网络，平衡效率与多样性

初始化策略的精细化设计

  • 部分保留初始化（Partial Preservation Init） ：保留原模型部分层的参数，随机初始化新增层
  • 交叉层知识传递（AKI） ：不仅考虑当前层参数，还结合下一层参数进行初始化

附录：MGI 中的 FPI 原则详细介绍

  • 功能保留初始化（Function-Preserving Initialization，FPI） 特指通过精确的参数复制与调整，确保扩展后的大模型在初始状态下对任意输入的输出与原小模型完全一致（或高度近似），从而实现“无损知识迁移”
  • FPI 是模型增长初始化（Model Growth Initialization，MGI）中的关键要求

核心机制与示例

  • 宽度扩展 ：当某层神经元数量从 2 个扩展到 3 个时，新增神经元的权重复制原模型中邻近神经元的参数，并通过权重拆分（如将原权重w拆分为w/2和w/2），使新输出与原输出完全一致。例如：
    • 原模型：\( y = w_1 h_1 + w_2 h_2 \)
    • 扩展后模型：\( y = w_1 h_1 + (w_2/2) h_2 + (w_2/2) h_3 \)
  • 深度扩展 ：在新增 Transformer 层时，将其参数初始化为“恒等映射”，确保新增层对输入不做任何修改，等效于原模型的计算流程
    • 如自注意力的输出投影矩阵设为单位矩阵可实现恒等映射
  • MoE 模型扩展
    • 新增专家的输入/输出投影权重初始化为 0，使其在初始阶段不参与计算，仅通过训练逐步被激活，避免干扰原模型的优化状态

为什么 MGI 中需要 FPI？

  • 降低训练震荡风险 ：通过确保扩展后的模型初始输出与原模型一致，避免损失函数突增或优化器状态中断
  • 加速大模型收敛 ：复用小模型的知识，使大模型在训练初期即可继承成熟的特征提取能力，减少从头学习的时间成本

注：本文包含 AI 辅助创作

Paper Summary

  • 本论文首次对开源语言模型在不同任务上的性能进行了系统分析，将其性能与架构和数据等联系起来
    • 但本文的分析存在一些不太严谨的地方，因为各个公司的实现方式可能是不一样的，比如基建或者各种配方的使用等
  • 一般来说：语言模型能力的提升通常归因于模型规模或训练数据的增加
  • 但在某些情况下：
    • 使用精选数据训练的小模型或采用不同架构决策的模型 可以超越在更多 token 上训练的更大模型
    • 引出问题：这是什么原因造成的呢？
  • 为了量化这些设计选择的影响，论文对 92 个不同规模的开源预训练模型进行了元分析 (meta-analyze)，这些模型包括：
    • SOTA 开源权重模型
    • 性能较差的模型（less performant models）
    • 采用非常规设计决策的模型（less conventional design decisions）
  • 论文发现：
    • 通过纳入除模型大小和训练 token 数量之外的特征 ，论文预测下游任务性能的能力相对提高了 3-28%
      • 注：这个提升是与仅使用规模特征相比
    • 对模型设计决策的分析揭示了对数据构成的见解
      • 例如代码占比在 15-25% 时语言任务和代码任务之间的权衡，以及网络数据对真实性 (truthfulness) 的负面影响
  • 论文的框架为更系统地研究模型开发选择如何塑造最终能力奠定了基础

Introduction and Discussion

  • 语言模型训练的效果关键取决于预训练 (pretraining) 期间所做的决策
    • 例如，扩展数据 (scaling up data) 的有效性取决于其构成
    • 理解：即使处理了一万亿个 token，如果它们全部由单词 “the” 组成，那也是无效的
  • 研究发现，语言模型的性能可以通过 Scaling Laws (2020, 第 2 节) 进行相当准确的预测
    • 即基于模型参数数量和训练所用 token 数量对模型性能进行外推
  • 但仅基于这两个方面的扩展定律并不总能解释下游任务性能 (2024; 2024)
  • 研究界在理解训练决策如何影响下游性能方面已经取得了进展，特别是在数据构成方面。例如，对照研究 (controlled studies) 表明
    • 在代码数据上训练可以提高在某些推理基准测试上的性能 (2024; 2024)；
    • 数据的元特征 (meta-features)，如年龄和毒性过滤器 (toxicity filters) 的使用，会影响许多问答 (QA) 任务的性能 (2024)；
    • 多语言数据的平衡会影响英语和其他语言的性能 (2023; 2025)
  • 这些工作揭示了宝贵的见解，但它们往往只关注改变训练方案 (training recipe) 的单个方面，而保持其他方面不变
    • 尽管严谨，但这在计算和开发时间上成本高昂
  • 论文转而提出一个问题：论文能否利用过去来自开源语言模型的发现来检验训练决策如何共同影响下游性能？
  • 为此，论文首先对来自不同系列的 92 个基础预训练 LM 的模型架构和数据相关的特征进行了 编目 (catalog) (章节3)
    • 由此产生的模型特征数据库涵盖了 2019 年至 2024 年间发布的大多数主要的、原始的开源权重 Decoder-only 预训练模型
  • 然后，论文开发了方法来 预测 这些模型在广泛基准测试上的性能
    • 预测依据既包括传统的扩展因素，也包括架构决策和数据构成 (章节4)
  • 具体来说，论文训练回归模型 (regression models)
    • 模型输入：提取的特征
    • 模型输出：预测基准测试结果
  • 进一步使用模型可解释性技术 (model interpretability techniques) 来识别在做出这些预测时最显著的特征
  • 论文在预测 12 个流行 LLM 基准测试的性能上评估了这种方法，并证明决定模型性能的 不仅仅是扩展 (not just scaling)
    • 在所有基准测试上，包含所有特征的回归器 (regressor) 的性能都优于仅基于扩展模型特征的回归器 (章节5.1)
  • 论文对特征重要性 (feature importance) 的分析揭示了数据领域 (data domains) 对任务性能的潜在影响，再次证实了经验性结果，例如预训练中使用代码的最佳比例 (章节5.2)
  • 此外，论文发现从模型生成文本中提取的特征（例如问题相关词的频率或类似网络文本的比例），有助于预测各种基准测试的性能
    • 这表明模型的生成模式可以反映其预训练数据中的潜在偏差 (underlying biases) ，进而影响下游性能
  • 通过记录整个社区训练的开源模型并提取见解，论文为模型开发者提供了一个实用的资源，以从集体经验中学习
  • 论文在 (章节8) 中讨论了这一点以及未来的工作

Scaling Laws

Definition

  • 论文在此将扩展定律定义为语言模型系列的参数数量 \(N\) 和 token 数量 \(D\) 与收敛时期望的语言建模损失 \(L(N,D)\) 之间的关系
  • 重要的是，这些定律通常是在保持所有其他因素不变的情况下进行研究的：
    • 保持相同的模型架构、训练数据和模型参数
  • 最初，Kaplan 等 (2020) 表明，在广泛的基于 Transformer 的模型中，这种关系可以表示为幂律 (power law)：
    $$L(N,D)=\left(\left(\frac{N_{c} }{N}\right)^{\frac{\alpha_{N} }{\theta_{D} } }+\frac{ D_{c} }{D}\right)^{\alpha_{D} }$$
  • 后来，Hoffmann 等 (2022a) 提出了一个类似的定律，其拟合的系数不同，但同样基于幂律
  • 但扩展定律并非绝对，其确切函数形式和拟合系数可能取决于架构类型、规模范围 (Pearce and Song, 2024) 或其他考虑因素，如推理成本 (inference costs)
  • 更多讨论请参见 (章节7.2)

Maybe it’s Not Just Scaling?

  • 参数数量和训练 token 数量真的是准确预测模型下游性能所需的全部吗？直觉上答案是否定的
  • 模型训练涉及许多设计决策，所有这些都可能对模型性能产生影响
  • 模型架构细节 (Model Architecture Details)
    • 虽然大多数现代语言模型都遵循 Transformer 架构，但存在一些细节差异
      • 例如，层归一化 (layer normalization) 的种类 (2019) 和位置 (2020)，以及位置编码 (positional encoding) 的类型 (2021; 2022) 都会对模型性能产生显著差异
    • 先前的工作，例如 Gu and Dao (2023)，已经凭经验证明，在保持所有其他因素相等的情况下，做出更好架构决策的模型 (2023a) 优于做出更差决策的模型 (2017)
  • 数据构成 (Data Composition)
    • 数据构成和质量在模型的最终质量中起着重要作用
      • 例如，过去的工作表明，训练一定数量的代码可以提高英语推理任务的性能 (2023)
    • 同样，有工作表明，筛选“教育”内容可以实现更高效的学习，并在基于知识的问答任务上获得更高的性能 (2023)
  • Task Setting
    • 最后，所有上述因素与模型性能的衡量方式之间存在相互作用
    • 虽然先前关于扩展定律的工作主要测量损失值，但下游用户通常关心的是任务性能，而不是预训练数据集上的验证损失 (validation loss)
    • 尽管对于许多任务来说，两者之间通常存在相关性，但某些任务可能更难仅从模型的损失来预测 (2024)
    • 此外，某些任务表现出异常的扩展行为，例如反向扩展 (inverse scaling) 或 U 型扩展 (U-shaped scaling) (2023; 2023; 2024)，或者仅仅是更不可预测的性能 (2024)
  • 论文提问：论文能否通过设计一套新的、不仅仅依赖于基于扩展的因素的“定律”来更有效地预测 LLM 的性能？

Building a Database of Publicly-Available Language Models

  • 为了解决论文的研究问题，论文构建了一个包含 11M 到 110B 参数的公开可用语言模型的数据库（包括嵌入参数），仅限于不同的 Decoder-only 基础预训练模型
    • 注：不同是指训练数据和架构的独特组合。在去重数据集上训练的模型被单独计数，但具有不同课程/初始化的变体不计入
  • 本节描述了论文的纳入标准、模型特征化以及评估方法

Data Collection

  • 为了确保论文的分析是一致的，论文应用了以下标准：
  • Pretrained-only：
    • 仅包含从头开始预训练的基础模型，排除了微调变体、合并模型以及经过额外后训练的模型
  • Architecture：
    • 仅包含基于 Transformer 的 Decoder-only 模型以保持一致性
    • 排除了 MoE 或其他架构
  • Publicly available information：
    • 仅包含具有公开可用元数据的模型，这些元数据通过配置文件或论文记录
    • 特别是，纳入需要总参数数量和训练所用总 token 数量这两个信息
    • 模型和模型系列的完整列表可在附录A 中找到

Characterizing Models and Data

  • 论文通过每个模型所做的架构选择以及其预训练数据的选择来表示每个模型
  • 形式上，令 \(\mathcal{A}\) 为与模型架构相关的特征集合， \(\mathcal{D}\) 为与模型预训练数据集相关的特征集合
  • 对于每个任务 \(T\) ，作者希望用预测值 \(\widehat{s_{T} }\) 来近似模型 \(M\) 的真实得分 \(s_{T}\) ：
    $$\widehat{s_{T} }(M)=f_{\theta}([\mathcal{A}_{M};\mathcal{D}_{M}]).$$
  • 当 \(\mathcal{A}=\{\#\text{params}\}\) ，\(\mathcal{D}=\{\#\text{tokens}\}\) ，且 \(f_{\theta}\) 是幂律时，这就简化为典型的扩展定律
  • 论文总共记录了 92 个开放模型，涵盖模型特征、高层数据集特征以及从该模型的无上下文生成中派生出的特征等维度
  • 有关完整的特征集和定义，请参见附录B
Features from Model Documentation
  • 论文首先通过阅读源论文/博客（如有，请参见附录A了解原始引用）以及在 Hugging Face Hub (2020) 上列出的数据来收集每个模型的信息
  • Architectural Features：
    • 这些特征捕获了决定模型结构的设计决策
    • 例如，总参数（包括嵌入参数）、Transformer层数、嵌入和前馈维度，以及细节，例如使用的层归一化类型或注意力变体
  • Data Features：
    • 这些特征总结了预训练数据的组成
    • 代表性示例包括训练所用的总 token 数以及来自图2中定义的各个领域的 token 百分比细分，以及英语 token 的比例
    • 论文的预训练数据领域源自开放预训练数据集(2020; 2024)中常见的子领域
    • 论文使用顶级领域（网络、代码、书籍、参考、学术），因为这往往是论文中描述数据组成的粒度
Exploring Data Composition via Generation
  • 尽管许多模型记录了一些数据组成细节，但相对较少的模型发布了其完整的预训练语料库，导致论文研究中许多模型的这些值缺失
  • 为了填补这些空白，论文探索了一种替代方法：
    • 分析模型在无提示情况下生成的文本，以估计其训练数据的特征
    • 论文假设模型的生成风格和内容反映了其训练数据的分布
    • 对于每个模型，论文使用温度 \(T=0.8\) 和 top-p \(p=0.9\) 的核采样生成 5k-10k 个无上下文的生成文本（每个生成文本最多 256 个 token）
    • 然后，论文使用标准 NLP 工具和基于 LM 的分类器从这些生成文本中提取语言学和领域特征
    • 论文在附录E和F中验证了这种方法
  • 论文还提取了 low-level 语言特征，例如每句词数（words per sentence）、成分树深度（constituency tree depth）和依存长度（dependency length）
    • 论文的验证分析（附录G）表明：
      • 领域层级特征与实际预训练数据构成具有较强的相关性
        • 例如，网页内容相关性：\(r = 0.916\)，\(p = 7.55 \times 10^{-12}\)），
      • low-level 风格特征的相关性较弱
    • 然而，所有特征的整体 Model-level 相关性表现强劲（通常 \(r > 0.8\)）
    • 这一结果支持我们将“自由生成内容”（free-generations）用作预训练数据构成的替代指标（proxies），同时也说明不能用自由生成特征替代预训练特征
  • 注：一些关键术语补充说明：
    • constituency tree depth（成分树深度） ：句法分析中的核心概念，指“成分树”（用于表示句子句法结构的树形图，如主谓、动宾等成分的层级关系）从根节点到最深叶节点的路径长度，反映句子句法结构的复杂程度
    • dependency length（依存长度） ：依存句法分析中的指标，指句子中两个存在依存关系的词语（如中心词与修饰词）在文本序列中的距离，常用来衡量句子结构的线性复杂度

Evaluation Datasets and Metrics

  • 为评估设计选择对推理能力的影响，我们在 Open LLM 排行榜（2024）的数据集上对模型进行了测评，这些数据集涵盖了推理能力的多个不同维度（见表 1）
    • 其中，部分模型的测评结果直接从该排行榜获取；
    • 对于未列入该排行榜的模型，我们使用 Eleuther LM 评估工具包（2023），在完全相同的设置下开展测评
    • 此外，若某项任务或子任务存在多个版本，我们会对所有版本均进行测评，并通过求平均值得到该任务的整体得分
    • 有关评估数据集及测评设置的完整列表，详见附录 C
  • 对于评估数据集 \(T\) ，其中第 \(i\) 个样本是 \(y_{i}\) ，模型为 \(M\) ，论文如下定义 \(s_{T}(M)\) ：
  • 准确率(Accuracy) ：对于大多数任务，论文使用未归一化的精确匹配准确率
    $$ s_{T,\text{acc} }=\frac{1}{|T|}\sum_{i=1}^{|T|}\mathbb{I}\{y_{i}=\hat{y}_{i}\} $$
    • 对于 Humaneval，论文使用 pass@1，但为方便起见，将其与准确率任务归为一组
  • Brier分数(Brier score) 对于较小模型难以达到非零准确率的任务，论文遵循(2023)的做法，使用多类 Brier 分数作为多项选择任务的替代连续指标(1950) (注意：对于 Brier 分数来说，值越低越好，多类 Brier 分数范围在 0-2 之间)
    • 对于一个有 \(K\) 个类别的任务，令 \(p_{ik}\) 为样本 \(i\) 上类别 \(k\) 的预测概率。则
      $$ s_{T,BS}=\frac{1}{|T|}\sum_{i=1}^{|T|}\sum_{k=1}^{K}(p_{ik}-\mathbb{I}\{y_{i}=k\})^{2} $$

异质性(Heterogeneity) in Task-specific Scaling

  • 在加入其他因素之前，作者检查了所选任务之间沿 \(N\) 和 \(D\) 扩展的差异
  • 论文为每个任务拟合了一个(2020)风格的定律
  • 如图3所 示，论文看到不同的任务在遵循扩展趋势的程度以及它们各自的扩展轮廓上可能表现出显著差异
  • 例如，TruthfulQA 似乎表现出 U 形扩展，而 Humaneval 有更多的“异常值”模型
  • 任务的 \(R^{2}\) 值完整列表可在附录D 中找到

Predictive Modeling

  • 接下来，给定论文的数据库，论文拟合一个回归器来尝试预测性能
  • 在传统的扩展定律中，回归器是基于幂律拟合的
  • 然而，论文现在要处理大量特征，其中一些可能无法通过简单的参数形式很好地捕捉
  • 因此，论文遵循先前关于性能预测的工作(2020; 2021)，利用基于 XGBoost (2016)的树形回归器
    • 论文还试验了LightGBM (2017)，发现其性能相似。结果见附录K
  • 对于每个评估基准，训练一个模型 ，以基于架构特征 \(\mathcal{A}\) 和数据特征 \(\mathcal{D}\) 来预测该任务上的性能指标
    • 理解：每个评估基准都有一个单独的 XGB 模型
  • 对于每个任务设置，由于模型数量相对较少，论文执行 3 折交叉验证，并在每折的训练集上进行嵌套内部交叉验证
    • 内部交叉验证在一小组超参数上进行网格搜索，允许模型随任务略有变化。更多细节请参见附录I
  • Evaluation 为了评估预测器，论文使用所有模型和折迭的平均绝对误差(Mean Absolute Error)
    • 对于一个有 \(N\) 个模型被评估的任务
      $$ \text{MAE}_{T}=\frac{1}{|T|}\sum_{i=1}^{N}|s_{T}(M_{i})-\widehat{s_{T} }(M_{i})|$$
    • 论文将扩展定律预测器以及全特征预测器相互比较，同时也与中位数基线(median baseline)（它只是为该折迭测试集中的每个模型预测训练集中模型的中位数得分）和对数线性基线(log-linear baseline)（它将一个对数线性函数拟合到参数数量和 token 数量）进行比较
  • 迭代特征选择(Iterative Feature Selection) 由于完整的特征集非常大，论文根据哪个特征能最大程度地减少 MAE（在 5 个随机种子上平均），从完整集合中贪心地顺序选择特征
    • 不断添加特征，直到观察到的减少量不再至少为 \(1\times 10^{-4}\)
    • 论文开始时仅使用两个扩展定律特征，并将其称为扩展定律(scaling-laws) 模型，尽管它不具有传统幂律的形式
      • 注：不具有传统幂律形式的解释：由于论文使用基于树的预测器来适应多样化的特征类型（包括非数值型特征），论文的方法优先考虑在观察到的界限（10M-100B 参数，50B-3T token）内进行插值，而不是外推（探索其他预测方法仍然是未来的工作）
    • 然后，通过合并额外的架构或数据特征，我们可以直接量化这些额外特征带来的增量预测能力
      • 论文将具有该组特征的模型称为全特征(all-features) 模型
      • 在所有情况下，论文使用相同的超参数网格、相同的随机种子和分割来运行模型
  • 显著性检验(Significance Testing) 由于基线之间的相对差异很小，论文在多个种子（50个）上测试两个预测器
    • 然后，论文对每个种子的总体 MAE 值运行配对 \(t\) 检验，并使用错误发现率(1995)对跨任务的多重比较进行校正

Results

Predictor Performance

  • 加入与规模无关的特征能持续提升基准测试性能
    • 论文发现，在传统的扩展定律特征之外加入额外特征，能在多个基准测试上显著提升预测准确度，如表 2 所示
    • 在所有评估案例中，全特征预测器均优于仅使用扩展定律的预测器，相对误差减少的幅度大约从 3%（MathQA）到 28%（Lambada）不等
    • 值得注意的是，在语言建模和常识推理任务中观察到了最强的改进效果
  • 某些任务更强烈地依赖于非规模特征
    • 这种改进模式表明，架构和训练数据特征，对于预测与特定数据“类型”更紧密相关的某些任务，其表现可能更具信息量
    • 在代码生成任务（13% 的改进）和基于自然语言的推理任务（例如 Lambada，28% 的改进）上都观察到了巨大的改进
    • 即使是领域较窄的任务，如数学推理（GSM8k，+16%）或知识密集型评估（MMLU，+11-14%），也看到了一致但更温和的增强
    • 然而，使用 Brier 分数的基准测试显示出较小的改进（约 3-6%）
      • 这可能是因为 Brier 分数本身对模型性能中的涌现效应敏感性较低，特定任务的选择限制了改进空间，或者是这两个因素共同作用的结果

What Features Does Task Performance Depend On?

  • 为了理解影响任务性能的因素，作者检查了 Shapley (1953) (SHAP) 值，这些值显示了特征值如何影响预测
    • 注：SHAP Value（SHapley Additive exPlanations，沙普利可加解释）是一种基于合作博弈论的模型解释方法，核心目标是量化机器学习模型中每个特征对预测结果的贡献程度 ，让复杂模型（如随机森林、神经网络）的决策过程变得可解释
    • “SHAP Value”：每个特征对这个 “收益” 的贡献值，即该特征让预测结果偏离基准值的程度（正值表示推动预测值升高 ，负值表示推动预测值降低）
  • Arc Challenge、HumanEval、Winogrande 和 TruthfulQA 的结果如图 4 所示，其余基准测试的结果见附录 L
  • 少量代码大有裨益，但过多则对自然语言推理（NLI）有害
    • 预训练数据中代码的百分比是一个关键的非规模特征
    • 较高的代码组成有益于 Humaneval 性能，但对包括 Arc Challenge、Hellaswag、Winogrande 和 Lambada 在内的自然语言推理任务产生负面影响
    • 如图 5 所示，代码比例超过 20-25% 的模型在 Humaneval 上显示出增益，但在语言基准测试上受到惩罚
    • 15-25% 的中等代码比例似乎能平衡这些相互竞争的需求
  • 其他数据领域显示出任务特定的效应
    • 从自由生成特征中，论文观察到最近使用合成数据训练的模型（Phi (2023)、SmoILM (2024)）生成了更多疑问词，这表明训练数据中包含问答内容
    • 类似参考书或包含大量问题的生成内容与 Arc Challenge 和 Winogrande 的更好性能相关，而类似网络文本的生成内容则与更差的 TruthfulQA 性能相关（图 4）
  • 非规模的架构决策影响较小
    • 大多数高影响力特征是与数据相关的或与规模相关的架构特征（例如，维度）
    • 在某些情况下，层归一化的类型和位置嵌入都被认为具有显著影响

Validating Performance Predictions with Confirmatory Experiments(验证性实验)

  • 为了验证元分析的发现，论文还使用 Dolma 数据集上训练的 460M 参数模型进行了验证性的预训练实验
  • 论文旨在验证两个关于数据分布的发现：
    • (1) 当仅考虑自然语言推理时，约 8% 的代码比例是最优的 ，但在平衡代码和自然语言时，15-25% 可能是最佳比例；
    • (2) TruthfulQA 性能随着网络数据比例的增加而降低
      • 理解：网络数据的虚假信息多
  • 由于这是一个小规模模型，准确度差异可能不显著，论文将相关数据集转换为基于损失的评估
  • 由于计算限制，论文将每个检查点训练 10B 个 token，但使用按 100B token 运行规模调整的余弦学习率调度
  • 详细信息和精确的损失图见附录 M
  • 总体而言，在图 6 中，论文发现验证性运行在很大程度上验证了论文的元分析预测
    • 唯一不足：尽管准确度的趋势符合预期，但 TruthfulQA 的基于边际的损失在 50% 网络数据检查点上略低于 30% 检查点
  • 这提供了初步证据，表明论文的分析方法可以用于先验地智能预测语言模型训练设计决策

Empirical Data Composition Results

  • 先前的研究已经探讨了预训练中代码数据的作用 (2023; 2024) 以及领域消融 (2024)
  • 数据过滤可以在单纯扩展规模的基础上进一步提升性能 (2023; 2024)
  • 论文的结果表明，代码数据在中等比例（最佳比例 15–25%）下能增强自然语言推理能力，这修正了先前 25% 的估计 (2024)
  • 论文通过汇集现有模型的见解来识别有前景的测试方向，从而对实证消融研究进行了补充

Observational and Task-Specific Scaling Law Fitting（理解：观测性与任务特定的扩展定律拟合）

  • 任务特定的扩展定律研究表明，参数分配会影响机器翻译的结果 (2021)，而多任务处理对英语-目标语言对有益 (2023)
  • 关于下游任务的研究强调了预训练数据与下游数据之间对齐的重要性 (2021; 2024)
  • 各种研究探讨了数据重复 (2024)、多领域数据 (2024) 以及稀疏性 (2023)、精度 (2024) 和推理成本 (2022a) 等因素，而一些研究发现训练超参数具有稳定性 (DeepSeek-2024a)
  • Ruan 等 (2024) 也使用开源模型的观测数据来预测任务性能，但他们是根据模型在其他任务上的表现来预测某一任务的性能
  • 论文在识别性能的通用自然语言能力和编码能力两个方面得出了类似的结果，但论文的动机是将这些能力追溯到预训练决策

Pretraining Data Selection

  • 领域混合在预训练中已被研究，其他工作将其表述为回归问题 (2024; 2025) 或在训练过程中使用代理模型来选择领域权重 (2023; 2023; 2024b; 2025)
  • 相比之下，论文回顾性地分析了领域构成和训练决策如何影响跨任务的性能，这是在训练期间为单个模型优化数据权重的补充视角

Tracing Capabilities to Data

  • 特定的语言模型能力已被关联到预训练数据中的模式
  • 数值推理和句法规则学习的性能取决于训练数据中数字术语的频率 (2020; 2021)
  • Ruis 等 (2024) 发现，对推理有影响的数据分散在众多文档中，并且与程序性内容相关
  • 类似地，Chen 等 (2024) 观察到 “并行结构”与上下文学习能力密切相关
    • 问题：这里的并行结构是什么？
  • 论文目前关注更广泛的数据领域，但论文的框架可以通过更细粒度的任务或更精细的数据特征进行扩展

Future Work

  • 展望未来，有几个明确的方向
    • 首先，论文的数据库（章节3）可以随着新模型和基准测试的发布而进一步扩展，论文将发布代码和数据以帮助推动社区进行更系统的数据记录工作
    • 其次，作者希望论文的工作将有助于发现在更受控环境中测试的假设
      • 现有模型交织了许多设计决策，而进一步仅涉及单一变化轴的受控预训练实验可以进一步阐明每个特征的影响
    • 最后，在论文的研究中，绝大多数预训练模型专注于密集 Transformer 架构，而混合专家 (2024a; DeepSeek-2024b) 和状态空间模型 (2023) 等替代架构也引起了显著的研究兴趣
      • 如何恰当地对这些更多样化的模型架构进行特征化，并在性能预测中使用这些信息，是一个有趣的挑战，可能会揭示更多的见解
  • 尽管预训练数据分析和选择迄今为止主要集中于实证发现，但通过大规模实证研究更好地理解训练如何影响模型能力，也可以促进可解释性实验和对学习表征的可能干预，其中受控的变化轴提供了案例研究

Limitations

  • 论文当前的工作有几个局限性，可以在未来的工作中改进
  • 第一，尽管论文记录了许多开源模型，但论文的样本量仍然有限，特别是对于较大（>50B）参数的模型
    • 这限制了论文得出关于大型模型扩展行为的稳健结论的能力
    • 而且论文拥有的模型在参数数量、数据大小和数据分布上并不均匀，某些规模范围和数据分布被过度代表
      • 哪些模型被开源也可能存在选择效应，并且在不同的时间段，流行的架构决策或数据构成可能存在时间效应
  • 第二，论文的方法论也带来了一些局限性
    • 因为论文没有系统地训练所有论文自己的模型（尽管论文在附录 A 中有一些自己的模型），所以论文的分析本质上是观察性的
    • 虽然我们可以观察到设计选择与性能之间的有趣关系，但要做出因果断言需要实验验证
    • 此外，虽然基于树的回归器能有效捕捉复杂的特征交互，但它们限制了论文外推超出数据集中所见模型大小（参数和 token 数量）范围的能力
  • 第三，论文注意到论文工作的范围也有局限性
    • 论文专注于 Decoder-only 的基预训练密集 Transformer 模型，这排除了重要的架构变体，例如混合专家模型、非基于 Transformer 的架构以及经过后训练的模型
    • 此外，论文主要检查英语模型，因为论文在这项工作中不关注多语言性
    • 论文的特征集虽然广泛，但可能仍未捕捉到模型设计和训练的所有相关细节，特别是目前的优化细节
  • 这些局限性为未来的工作指明了方向：
    • 扩展数据库以包含更多样化的模型类型和语言覆盖范围；
    • 开发更具针对性的函数形式，以便在输入异构特征集的同时实现更好的外推；
    • 使用新的预训练模型进行有针对性的实验，以验证特定设计选择的影响

Ethical Considerations

  • 在这项工作中，论文专注于理解模型为何在标准基准测试上表现良好，但并未关注其他重要的考量因素，例如安全性或社会偏见
  • 而且论文的分析侧重于英语模型和基准测试
    • 这一局限性反映但也可能强化了该领域现有的对英语的偏向，可能导致对其他语言有效架构的开发投入不足

附录 A：List of all models

  • All models are listed in Table 3.

附录 B：List of all architectural and data features

B.1 Architectural Features

  • （注意，本部分的特征是从官方文档（例如 Hugging Face 的模型/数据卡片或原始论文）中收集的）
  • 总参数量 (Total parameters) ：模型中的参数总数（包括嵌入参数）
    • 注意，论文仅包含 Decoder-only 的密集模型
  • 维度 (Dimension) ：嵌入维度
  • 头数 (Num heads) ：注意力头的数量
  • MLP 比率 (MLP ratio) ：\(\frac{\text{FFN dimension} }{\text{Embedding dimension} }\) 的比率
  • 位置嵌入 (Positional Embeddings) ：位置嵌入的类型
    • 这可以是非参数的（正弦或固定嵌入）、学习的（仅作为每个位置的向量学习）、rope (rope 嵌入) 或 alibi（技术上不是嵌入，但因其功能目的而包含在此）
  • 层归一化 (LayerNorm) ：应用的层归一化类型
    • 这可以是非参数的（仅基于算术的归一化）、参数的（类似，但有一些可学习的参数，如扩展/偏置）和 RMSNorm（参数版本的简化版）
  • 注意力变体 (Attention variant) ：使用的注意力的大致类型
    • 这可以是 full（普通注意力）、local（每个 Token 位置仅关注其周围的位置）、mqa（多查询注意力）或 gqa（分组查询注意力）
  • 偏置 (Biases) ：模型的某些部分是否存在偏置项
    • 可以是 none（无偏置）、attn only（仅在注意力层中）、ln only（仅在层归一化中）
  • 块类型 (Block type) ：变压器块是否完全并行计算
    • Sequential 表示不并行，而 parallel 表示在注意力或 FFN 层中存在某种并行性
  • 激活函数 (Activation) ：使用的激活函数
    • 可以是 relu、gelu/gelu 变体、silu 或 swiglu
  • 序列长度 (Sequence length) ：序列长度
  • 批次实例数 (Batch instances) ：预训练期间使用的批次大小

B.2 Data Features

  • （注意，本部分的特征是从官方文档（例如 Hugging Face 的模型/数据卡片或原始论文）中收集的）
  • 总 Token 数 (B) (Total tokens (B)) ：预训练期间使用的 Token 总数，以十亿计（转换为对数尺度）
  • 预训练数据中网络数据百分比 (% Web in Pretraining) ：来自通用网络来源的预训练数据百分比
  • 预训练数据中代码百分比 (% Code in Pretraining) ：由代码组成的预训练数据百分比
  • 预训练数据中书籍百分比 (% Books in Pretraining) ：来自书籍的预训练数据百分比
  • 预训练数据中参考文献百分比 (% Reference in Pretraining) ：来自参考文献来源的预训练数据百分比
  • 预训练数据中学术内容百分比 (% Academic in Pretraining) ：来自学术来源的预训练数据百分比
  • 预训练数据中英文百分比 (% English in Pretraining) ：预训练数据中英文文本的百分比

B.3 Freegen-derived Features

  • 这些特征源自模型的生成文本
  • 对于每个模型，提取 5-10k 个生成文本，并聚合以下指标（通过均值和标准差）
    • 但二元组熵、教育分类器分数和领域分类是例外，因为它们是在所有生成文本上计算一次的
  • 论文使用 Stanza (2020) 在按语言对生成文本进行分类后生成基于解析的特征
    • 论文仅将 stanza 支持的语言包含在解析特征所基于的最终生成文本集中
B.3.1 生成长度和基本统计量 (Generation Length & Basic Statistics)
  • 平均字符长度 (Mean Character Length) ：每个生成文本的平均字符数（上限为 2048）
  • 平均生成 Token 数 (Mean Tokens Generated) ：每个生成文本的平均 Token 数
  • 平均句子数 (Mean Sentences) ：每个生成文本的平均句子数
  • 平均词数 (Mean Words) ：每个生成文本的平均词数
  • 平均每句词数 (Mean Words per Sentence) ：每个句子的平均词数
B.3.2 Constituency Parse Features
  • 最深解析树平均深度 (Mean Depth of Deepest Parse Tree) ：每个生成文本的平均最大选区树深度
  • 解析树平均深度 (Mean Depth of Parse Trees) ：所有句子/短语的平均选区树深度
  • 词平均深度 (Mean Word Depth) ：选区树内词的平均深度
  • 词深度变异平均 (Mean Word Depth Variation) ：跨句子/短语的词深度标准差的平均值
B.3.3 Dependency Parse Features
  • 依存头距离 90% 分位数平均值 (Mean 90th-Percentile Dependency Head Distances) ：对于每个生成文本，计算词与其依存头之间的线性距离的 90% 分位数，然后对这些值取平均
  • 最大依存头距离平均值 (Mean Maximum Dependency Head Distances) ：每个生成文本中任何词到其依存头的最大距离的平均值
  • 依存头距离中位数平均值 (Mean Median Dependency Head Distances) ：每个生成文本的依存头距离中位数的平均值
  • 最大依存根距离平均值 (Mean Maximum Dependency Root Distances) ：每个生成文本中任何词到句子根的最大距离的平均值
  • 平均依存根距离平均值 (Mean Mean Dependency Root Distances) ：每个生成文本中词到句子根的平均距离的平均值
  • 依存根距离中位数平均值 (Mean Median Dependency Root Distances) ：每个生成文本中词到句子根的距离中位数的平均值
B.3.4 Domain Classification Features
  • 生成学术类文本百分比 (% Generated Academic-like Text) ：被分类为学术类的生成文本百分比
  • 生成书籍类文本百分比 (% Generated Books-like Text) ：被分类为书籍类的生成文本百分比
  • 生成代码类文本百分比 (% Generated Code-like Text) ：被分类为代码类的生成文本百分比
  • 生成参考类文本百分比 (% Generated Reference-like Text) ：被分类为参考类的生成文本百分比
  • 生成专业文本百分比 (% Generated Specialized Text) ：被分类为专业类（例如，乐谱、象棋 PGN、生物医学数据）的生成文本百分比
  • 生成网络类文本百分比 (% Generated Web-like Text) ：被分类为网络类的生成文本百分比
B.3.5 Classifier and Language Metrics
  • 教育分类器分数平均值 (Mean Educational Classifier Score) ：教育分类器给出的平均分数
  • 生成英文文本百分比 (% Generated English Text) ：生成的英文文本的平均百分比

B.3.6 Lexical Diversity and Entropy Metrics

  • 平均二元组熵 (Mean Bigram Entropy) ：跨生成文本计算二元组的平均熵
  • 型符比 (Type-Token Ratio) ：唯一 Token 数与总 Token 数的平均比率
  • 唯一 Token 数 (Unique Tokens) ：每个生成文本的平均唯一 Token 数
B.3.7 Lexical and Stylistic Features
  • 实词-功能词比率 (Content-Function Ratio) ：实词（名词、动词、形容词、副词）与功能词的比率
  • 疑问词比率 (Question Words Ratio) ：每 10 万个词中疑问相关词（例如 how, what, why, when, where, who, which, whose）的比率
  • 祈使词比率 (Imperative Words Ratio) ：每 10 万个词中祈使词（例如 do, make, consider, take, use, ensure, check, build, apply, run, create, find, go, try, turn, start, stop, put, keep, leave, get, move）的比率
  • 连词比率 (Conjunctions Ratio) ：每 10 万个词中连词（例如 and, but, or, so, because, although, however, therefore, yet）的比率
  • 指令词比率 (Instruction Words Ratio) ：每 10 万个词中指令导向短语（例如 “Question:”, “Answer:”, “Instruction:”, “User:”, “Assistant:”, “Q:”, “A:”）的比率
  • 数字比率 (Numbers Ratio) ：生成文本中数字 Token 的比率

附录 C：List of all evaluations and settings

  • 尽管论文理想情况下会评估模型和任务的全部组合，但论文发现由于一些模型与 LM Evaluation Harness 不兼容以及计算限制，论文无法在每個数据集上评估所有 92 个模型
  • 论文在表 4 中列出了论文目前每个基准测试的评估数量，并将在数据库中继续补充评估结果

附录 D：Task Deviations from Kaplan-style Scaling Laws

  • 表 5 记录了针对每个模型性能拟合幂律分布所得到的决定系数（\(R^2\) value）

附录 E：Free-generation Domain Classification

  • 论文使用 GPT-4o-mini 将模型生成文本分类到顶级领域
  • 论文发现这种多阶段提示（清单 1，清单 2）在 Dolma 按领域采样的样本上具有合理的精确度 (2024)，因此使用它对自由生成文本进行分类

附录 F：Domain Classifier Validation

  • 为了验证基于 4o-mini 的分类器的可靠性，论文请论文的一位作者根据附录 E 中使用的相同标注标准，对来自三个预训练数据集（the Pile, the SmoILM corpus, 和 RefinedWeb）的 300 个选定样本进行标注
  • 被模型或人类标注者标注为“unknown”或“incoherent”的样本被排除，因为这些样本不包含在领域混合的计算中
  • 过滤后，论文分析了 258 个文本样本，发现人类标注者和模型的绝对一致率为 85.8%，Cohen’s \(\kappa\) 为 0.746，表明人类分类和模型分类之间具有高度一致性

附录 G：Free-generation Validation

  • 为了验证论文的自由生成方法作为预训练数据组成的代理，论文分析了模型的自由生成特征与其预训练数据之间的相关性
  • 对于在三个开放预训练数据集（the Pile, the SmoILM corpus, 和 Refinedweb）上训练的模型，论文比较了它们的自由生成特征与相同标注器和基于 LM 的分类器（附录 E）在预训练语料库随机采样的 100 万文档子集上产生的特征
  • 由于成本原因，对于领域分类，每个语料库使用了 100 万中的 5000 个示例
    • 这 100 万文档是通过水库采样均匀采样的
  • 此外，论文计算了两个整体模型层面的相关性，它们衡量了每个模型的完整生成配置文件与其训练数据的匹配程度：
    • 1）领域层面相关性 (Domain-level correlations): 对于每个领域类别（web, code, academic, books, reference），论文计算了模型文档记录的预训练数据中该领域的百分比与被分类到该类别的自由生成文本百分比之间的相关性
    • 2）特征层面相关性 (Feature-level correlations): 对于语言特征（连词比率、祈使动词比率、疑问词比率、数字比率、字符长度），论文评估了自由生成中的每个特征与预训练数据中相应特征的相关性
  • 这些模型层面的相关性明显高于单个特征的相关性，表明模型比任何单一维度都更好地保留了其训练数据的整体“指纹”
  • 表6 和 表7

附录 H：Free-generation Examples

  • Llama-2-70B，网页风格(web-like)：

    • 给定斯蒂芬妮在与布雷克谢尔顿一起出席的活动上无法认出自己的第一首歌曲你创作和/或录制的第一首歌是什么？如果你再也没听过它，直到在参加一个大型乡村音乐活动时与数十万人进行Facetime视频通话，你会有什么感觉？这正是格温斯蒂芬妮本周在乡村电台研讨会活动上被授予荣誉时发生的事情。在纳什维尔，可以想象，当你拥有像她那样疯狂的职业生涯时，可能很难记住这些年所有的里程碑，但试想一下还要当场为所有人做到这一点。嗯，结果发现，双重才华的多莉帕顿原本要介绍格温是“一个打破壁垒的艺术家”，当时她引用了斯蒂芬妮第一首“乡村”歌曲的歌词……‘3000年’。外交辞令的“好声音”教练布雷克谢尔顿在她当时展示的视频中插话，问道“还有其他完成的工作吗……

    • phi-2，书籍风格(books-like)： 太阳在傍晚的天空中明亮地闪烁，将温暖的光线投射在温馨的小厨房里，一对勤奋的父母正在那里准备一顿美味的饭菜。空气中充满了期待，他们开始了一次烹饪冒险，心中充满了爱和喜悦。他们卷起袖子，灵巧地挥舞着各种锋利的刀具，每把刀都是他们技能和精确度的延伸。就像艺术家拿着画笔一样，他们精巧地在食材中穿梭，与刀具共舞，如同舞伴。刀片的移动速度之快违背了日常动作的物理规律，优雅而有目的地在空中滑行。……

    • Qwen2-72B，参考（指令）风格(reference (instruct)-like)： 儿童：家庭票（最多2名成人和2名儿童）23.00美元 (1)布莱克先生想和他的妻子以及两个儿子（他们都超过9岁）一起看电影，他应该支付多少？__ A.23美元 B.46美元 C.48美元 D.49美元 (2)布朗先生想和他5岁的双胞胎儿子一起看电影，他应该支付多少__。 A.18美元 B.23美元 C.28.5美元 D.32美元 (3)丽莎在这部电影中看不到雪。__ A.正确 B.错误 (4).简是一位数学老师，她有两个女儿（一个12岁，妹妹8岁）。她们想看电影。她应该支付多少？__ A.28.5美元 B.27美元 C.33美元 D.26.5美元 (5)迈克十六岁。今天是星期六，他想和他40岁的父母一起看电影。他们将为门票支付____。 A.22美元

    • OPT-6.7B，网页风格(web-like)： 我不认为这是一个问题。这是一个陈述。我说养只狗吧，世界有狗更美好。我对自己说了完全一样的话。我不知道你为什么被踩了。因为reddit是一群讨厌狗的人

    • Codegen-16B-mono，代码风格(code-like)：

      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      11
      12
      13
      14
      15
      16
      17
      18
      19
      20
      21
      22
      from typing import Optional
      from fastapi import FastAPI
      from pydantic import BaseModel

      app = FastAPI()

      class Item(BaseModel):
          name: str
          price: float
          is_offer: Optional[bool] = None

      @app.get("/")
      def read_root():
          return {"Hello": "World_from_ FastAPI"}

      @app.get("/items/{item_id}")
      def read_item(item_id: int, q: Optional[str] = None):
          return {"item_id": item_id, "q": q}

      @app.put("/items/{item_id}")
      def update_item(item_id: int, item: Item):
          return {"item_name": item.name, "item_id": item_id}

    • Pythia-410M，学术风格（专利）(academic-like (patent))：

      • 本发明涉及用于从发电站的废物流中去除固体废物材料的新的和改进的方法和设备，其中该固体废物材料计划在发电厂完工后进行处置。先前已经进行了各种尝试来从工厂的废物流中去除固体废物。这是真的，例如，对于污泥的排放，这些污泥通常被收集到一个污泥罐中，并在一个通常连接到工厂的敞开排水出口的清洗罐中被洗出工厂。这项现有技术由美国专利号3,623,579进行了讨论，该专利授予了G.R.Clark并描述了一种通过在罐中絮凝和絮凝并搅动固体以打破固体颗粒之间的键合来处理废物流以去除固体废物的方法。此外，美国专利号4,016,823描述了一种装置，该装置描述了一种方法，其中液体污水从废物流中和从污水处理厂中被去除，在那里要被去除的固体废物将被处理以生产用于沐浴浴缸或肥皂的氨净化水，并且其中来自废水处理厂的污水被去除到污水处理厂，在那里这些污水与水混合或作为肥料处理。…

附录 I：Appendix I XGBoost Settings

  • 对于内部网格搜索，树的最大深度在[2,3,5]中，学习率在[0.01,0.1,0.3]中，树的数量在[50,100]中

附录 J：Selected Features by Task

  • 在表8中，论文展示了每个基准测试所选的特征

附录 K：LightGBM Results

  • 表2的 LightGBM 版本可以在表9中找到
  • 注意：未对 LGBM 进行显著性检验，因此这反映了一次运行的结果，尽管对两个预测器都仍在附录I的相同值上进行了超参数搜索
    • Brier 分数扩展 ×100 以便比较
    • 这里的两个预测器都使用 LGBM

附录 L：SHAP Plots for remaining benchmarks

  • 剩余基准测试的SHAP图可以在图7-图15中找到。请注意，对于Brier分数任务（ANLI,XNLI,MathQA,LogiQA2），分数越低越好

附录 M：Details on confirmatory pretraining runs

M.1 训练(Training)

  • 对于论文的验证性实验，论文使用 Megatron-Deepspeed 库从头开始训练了 460M 参数的 Llama-2 架构模型

  • 论文将训练 token 数量上限设为 10B，同时使用设置为 100B token 长度的余弦学习率调度（意味着每个检查点大约完成了“完整”预训练运行的 10%）

  • 训练在每个检查点一个节点上进行，使用 8 个 H100 GPU

    • 每个检查点大约需要 6 小时来训练

  • 对于论文的数据混合，论文使用 Dolma v1 数据集的子集构建了各种混合

    • 在网页与其他的实验中，论文固定了所有其他数据源的相对百分比，同时改变网页的百分比

  • 训练配置如下：

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    training:
        num_layers: 14
        num_attention_heads: 12
        seq_length: 2048
        num_kv_heads: 12
        hidden_size: 1536
        ffn_hidden_size: 4128
        tune_steps: 1000
        lr: 0.00015
        min_lr: 1.0e-5
        weight_decay: 1e-2
        grad_clip: 1.0
        lr_warmup_steps: 100
        save_interval: 2000
        eval_interval: 2000
        train_epochs: 1
        tp: 1
        micro_batch_size: 16
        global_batch_size: 512
        seed: 42

  • 除了数据混合之外，所有实验都使用相同的超参数以确保公平比较

M.2 Evaluation

  • 为了评估不同数据混合对模型性能的影响，论文在以下任务上评估了论文的模型：
    • 1）自然语言推理(Natural language inference): Lambada, winogrande, arc challenge
    • 2）Code Generation: Humaneval
    • 3）Math: GSM8K
    • 4）事实性(Factuality): TruthfulQA
  • 注意，由于时间限制，论文没有选择完整的评估集
  • 由于LM eval harness没有为所有任务实现困惑度/基于损失的评估，论文手动将多项选择任务转换为基于损失的指标，并在计算所有任务的损失时屏蔽提示或问题

M.3 转换为基于损失的指标(Conversion to Loss-Based Metrics)

  • 为了确保跨不同任务和模型的一致评估，论文将各种基准测试数据集转换为基于损失的指标
  • 这种方法允许在模型之间进行更直接的比较，并更清晰地解释改进
  • 以下是论文为每种数据集类型实现损失计算的方式：
  • 多项选择任务(ARC Challenge, Winogrande, HellaSwag, TruthfulQA)： 对于这些数据集，论文计算了两个主要的基于损失的指标：
    • 平均损失(Average Loss)： 论文计算了正确答案的负归一化对数概率。对于每个问题，论文将输入格式化为“问题+答案选项”，然后为每个选项计算按token长度归一化的序列对数概率。正确答案的负对数概率被用作损失
    • 基于边际的损失(Margin-based Loss)： 特别是对于 TruthfulQA，论文计算了真实答案和非真实答案之间的边际。这被计算为最佳真实答案的对数概率与最佳非真实答案的对数概率之差的负值。损失越低表示区分真实和非真实信息的能力越好
  • 生成任务(GSM8K, HumanEval, Lambda)： 对于生成任务，论文计算：
    • 回答损失（Answer Loss）： 论文计算 solution Token 上的交叉熵损失
      • 注：对 Lambada 任务，仅使用最后一个 word
    • 所有对数概率均被序列长度归一化

M.4 Full Results

  • 表 10：代码与自然语言混合数据
  • 表 11：网络数据与其他数据混合的精确损失值

注：本文包含 AI 辅助创作

Paper Summary

  • 背景 & 问题：LLM 虽然强大，但缺乏动态调整其权重以应对新任务、知识或示例的机制
  • 论文提出了 自适应大语言模型（SEAL, Self-Adapting LLMs） 框架，通过生成自身的微调数据和更新指令，使 LLM 能够自我调整
    • 给定新输入时，模型会 produces a self-edit
      • a generation，可能以不同方式重组信息、指定优化超参数，或调用工具进行数据增强和基于梯度的更新

        Given a new input, the model produces a self-edit—a generation that may restructure the information in different ways, specify optimization hyperparameters, or invoke tools for data augmentation and gradient-based updates

    • 通过 SFT ，这些 self-edit 会带来持久的权重更新，从而实现长期适应（lasting adaptation）
    • 为了训练模型生成有效的 self-edit，论文使用强化学习循环（loop） ，将更新后模型在下游任务中的表现作为奖励信号
    • 与依赖独立适应模块或辅助网络的现有方法不同，SEAL 直接利用模型的生成能力参数化并控制其自身的适应过程
  • 在知识整合和 Few-shot 泛化的实验中，SEAL 展现了语言模型在新数据下实现自我导向适应的潜力

Introduction and Discussion

  • 在大规模文本语料库上预训练的 LLM 在语言理解和生成方面表现出卓越能力（2020; 2023; 2024; 2025）
    • 但将这些强大模型适配到特定任务（2020）、整合新信息（2020）或掌握新推理技能（2025）仍然具有挑战性，主要由于任务特定数据的稀缺性
  • 论文探讨了一个有趣的假设：LLM 能否通过转换或生成自身的训练数据和学习过程来实现自我适应？（can an LLM self-adapt by transforming or generating its own training data and learning procedure?）
  • 以人类学生学习为例，学生通常通过整理笔记来备考，这些笔记是对原始内容的重新解读和增强 ：这种将外部知识转化为更易理解形式的能力是人类学习的普遍特征
  • 但当前 LLM 的训练和部署方式与人类学习形成鲜明对比：面对新任务时，LLM 只能通过微调或上下文学习（ICL, In-Context Learning）（2022; 2024; 2023）直接利用原始数据
    • 这些数据可能并非最优格式（或数量），且现有方法无法让模型开发定制化的数据转换和学习策略
  • 为实现语言模型的高效适应，论文提出赋予 LLM 生成自身训练数据和微调指令的能力
    • 具体而言，论文引入了一种强化学习算法，训练 LLM 生成 self-edit（即指定数据和优化超参数的自然语言指令（如图1 所示）
    • 论文将此类模型称为 自适应大语言模型（Self-Adapting LLMs，SEAL）
  • 论文在两个应用中评估 SEAL
    • 首先，测试其在整合新事实知识任务中的表现：模型通过生成合成数据而非直接微调原文
      • 在无上下文版本的 SQuAD（2016）问答任务中，准确率从 33.5% 提升至 47.0%
      • 值得注意的是，SEAL 自生成的数据表现优于 GPT-4.1 生成的合成数据
    • 其次，论文在简化版 ARC-AGI 基准（2025）上评估 Few-shot 学习能力，模型通过调用工具自主选择数据增强和优化超参数（如学习率、训练轮次等）
      • 实验表明，SEAL 自动配置工具的能力显著优于标准上下文学习和未经强化学习训练的 self-edit 方法
    • 这些结果共同证明，SEAL 是实现语言模型自我适应的通用框架

Methods

  • 论文提出了 自适应性大语言模型（Self-Adapting LLMs, SEAL） ，这是一种框架，使语言模型能够通过生成自己的合成数据和优化参数（即“自我编辑”）来改进自身
    • 模型通过 Token 生成直接产生这些自我编辑，生成过程以模型上下文中的数据为基础
    • 自我编辑的生成通过 RL 进行训练，模型的奖励信号来源于应用自我编辑后模型在目标任务上的性能提升
    • 因此，SEAL 可以视为一种包含两个嵌套循环的算法：
      • 外部强化学习循环优化自我编辑的生成
      • 内部更新循环利用生成的自我编辑通过梯度下降更新模型
  • 论文的方法可以看作是一种元学习（meta-learning）实例，其目标是学习如何生成有效的自我编辑

General Framework

  • 设 \(\theta\) 表示语言模型 \(\text{LM}_{\theta}\) 的参数，SEAL 在单个任务实例 \((C, \tau)\) 上运行，其中 \(C\) 是包含任务相关信息的上下文，\(\tau\) 是用于评估模型适应能力的下游任务
    • 在知识整合任务中，\(C\) 是需要整合到模型内部知识中的段落，\(\tau\) 是关于该段落的一组问题和答案；
    • 在 Few-shot 学习中，\(C\) 包含新任务的少量示例，\(\tau\) 是查询输入和真实输出
  • 给定 \(C\)，模型生成一个自我编辑（SE），其形式因领域而异（见原论文 3.2 节），并通过 SFT 更新参数：
    $$\theta’ \leftarrow \text{SFT}(\theta, \text{SE})$$
    • 问题：SE 是什么，如何理解 SE？
    • 回答：SE 是一些指令，给定一个新输入（如一段文本或少量示例），模型会生成一组数据或优化参数（即 self-edit），这些输出 SE 会被用于更新模型自身参数，从而提升模型
  • 论文使用强化学习优化自我编辑的生成过程：模型采取动作（生成 SE），根据 \(\text{LM}_{\theta’}\) 在 \(\tau\) 上的表现获得奖励 \(r\)，并更新其策略以最大化期望奖励：
    $$
    \mathcal{L}_{\text{RL} }(\theta_t) := -\mathbb{E}_{(C,\tau)\sim\mathcal{D} } \left[ \mathbb{E}_{\text{SE}\sim\text{LM}_{\theta_t}(:C)} \left[ r(\text{SE}, \tau, \theta_t) \right] \right]. \tag{1}
    $$
    • 与标准强化学习设置不同，论文的奖励取决于模型参数 \(\theta\)（因为 \(\theta\) 会更新为 \(\theta’\) 并随后被评估）
    • 因此，强化学习的状态必须包含策略的参数，即 \((C, \theta)\)，尽管策略的观察仅限于 \(C\)（将 \(\theta\) 直接放入上下文不可行）
    • 这意味着从旧模型 \(\theta_{\text{old} }\) 收集的（状态、动作、奖励）三元组可能与当前模型 \(\theta_{\text{current} }\) 不匹配
    • 为此，论文采用同策略（on-policy）方法，即自我编辑从当前模型中采样，并且奖励也基于当前模型计算
  • 论文尝试了多种同策略方法，如 GRPO 和近端策略优化（PPO），但发现训练不稳定
    • 最终，论文采用了 ReST\(^{EM}\)（2023），这是一种基于过滤行为克隆的简化方法，也称为“拒绝采样 + SFT”
    • ReST\(^{EM}\) 可以视为一种期望最大化（EM）过程：
      • E 步 ：从当前策略中采样候选输出；
      • M 步 ：仅对获得正奖励的样本进行监督微调
  • 这种方法在二元奖励下优化了目标函数（1）的近似：
    $$
    r(\text{SE}, \tau, \theta_t) = \begin{cases}
    1 & \text{If on } \tau \text{ adaptation using SE improves } \text{LM}_{\theta_t} \text{‘s performance}, \\
    0 & \text{Otherwise}.
    \end{cases} \tag{2}
    $$
    • 具体来说，在优化（1）时，论文需要计算梯度 \(\nabla_{\theta_t} \mathcal{L}_{\text{RL} }\)
    • 但由于奖励项 \(r(\text{SE}, \tau, \theta_t)\) 依赖于 \(\theta_t\) 且不可微，论文将其视为固定值
    • 在这种近似下，对于包含 \(N\) 个上下文和每个上下文 \(M\) 个采样自我编辑的小批量，蒙特卡洛估计为：
      $$
      \nabla_{\theta_t} \mathcal{L}_{\text{RL} } \approx -\frac{1}{NM} \sum_{i=1}^{N} \sum_{j=1}^{M} r_{ij} \nabla_{\theta_t} \log p_{\theta_t}(\text{SE}_{ij} \mid C_i), \tag{3}
      $$
      • 其中 \(p_{\theta_t}\) 表示模型的自回归分布
      • \(y_s^{(i,j)}\) 是自我编辑 \(\text{SE}_{ij}\) 的第 \(s\) 个 Token
      • 由于 \(r=0\) 的序列可以忽略，因此 ReST\(^{EM}\) 通过简单的“对优质自我编辑进行 SFT”优化了目标（1）
  • SEAL 的训练流程总结在算法1 中
  • 此外，尽管论文使用单一模型同时生成自我编辑和学习这些编辑，但也可以将这两个角色解耦
    • 在这种“教师-学生”框架中，学生模型通过教师模型提出的编辑进行更新，而教师模型则通过强化学习训练，以生成能最大化学生改进的编辑

Domain Instantiations

  • 论文在两个领域实现了 SEAL 框架：知识整合（Knowledge Incorporation）和 Few-shot 学习。这两个领域展示了模型适应的两种互补形式：
    • 1）将新信息整合到模型权重中，使其无需依赖上下文即可回忆（通过无上下文 SQuAD 评估）；
    • 2）在少量示例后泛化到新任务（通过 ARC 评估）
Knowledge Incorporation
  • 论文的目标是高效地将段落中的信息整合到模型权重中
  • 一种有前景的方法是让语言模型生成基于段落的内容，然后对原始段落和生成内容进行微调（2022; 2023; 2024）
  • 论文采用典型格式：段落的逻辑推论
    • 这种方法将上下文 \(C\) 转换为一组推论 \(\text{SE} = \{s_1, s_2, \ldots, s_n\}\)，并通过 Prompt 模型“列出从内容中衍生的几条推论”实现
    • 输出可能包括原始段落的推论、逻辑结果或重述
    • 在 附录C 中，论文还探索了其他 Prompt ，如“以不同方式重写段落”或“以问答格式重写”，并发现无论基础 Prompt 如何，论文的方法都能带来相似的性能提升
  • 这些自我生成的语句构成 SFT 的训练数据：论文对每个序列 \(s_i\) 计算标准因果语言建模损失，并更新模型参数，得到 \(\theta’\)
    • 由于每次更新的数据量较小且总更新次数较多，论文使用低秩适配器（LoRA）进行高效轻量级微调
    • 最后，适应后的模型 \(\text{LM}_{\theta’}\) 在任务 \(\tau\) 上评估。这一过程如图2 所示
  • 在强化学习训练中，适应模型在 \(\tau\) 上的准确率定义了奖励 \(r\)，驱动外部强化学习优化。这训练模型以最适合微调的方式重组段落
Few-Shot Learning
  • 抽象与推理语料库（Abstraction and Reasoning Corpus，ARC）是一个测试抽象推理和从极少示例中泛化的基准
    • 每个任务包括少量输入-输出示例和一个需要预测正确输出的保留测试输入
  • 论文采用测试时训练（TTT）协议，其中少量示例的增强用于执行基于梯度的适应。论文训练 SEAL 学习这些决策，而不是依赖手动调整的启发式方法。为此，论文定义了一组工具 ，每个工具是预定义的函数，用于转换数据或指定训练参数，包括：
    • 数据增强（Data augmentations） ：旋转、翻转、反射、转置、调整大小操作（如更改网格分辨率）以及链式或重复变换
    • 优化参数（Optimization parameters） ：学习率、训练周期数以及损失是否计算所有 Token 或仅输出 Token
  • 模型通过任务的少量示例 Prompt 生成自我编辑（即指定调用哪些工具及如何配置的指令），如图3 所示
    • 然后通过基于 LoRA 的微调应用这些自我编辑来适应模型
    • 适应后的模型在保留测试输入上评估，结果决定训练自我编辑生成策略的奖励

Results

  • 本节论文通过两个适应领域实证评估 SEAL： Few-shot 学习和知识整合。完整的训练、超参数和评估细节见 附录A 和 附录B

Few-Shot Learning

  • 论文使用 Llama-3.2-1B-Instruct（2023）进行实验，这是一个未针对 ARC 进行预训练的小型开源模型
    • 由于大多数 ARC 任务对未经预训练的模型具有挑战性，论文从 ARC 训练集和评估集中分别筛选了 11 个和 8 个任务，确保它们在最优 TTT 配置下可解
  • 模型通过 ReST\(^{EM}\) 训练，每个训练任务采样 15 个自我编辑
    • 每个自我编辑独立生成 15 个更新模型，随后在对应的保留测试示例上评估
    • 论文仅强化那些导致模型正确适应的自我编辑（即适应后能正确输出测试输入的模型）
  • 训练后，论文通过为每个保留评估任务生成 5 个自我编辑并独立应用它们来评估模型
    • 然后报告导致正确输出的自我编辑百分比，反映学习到的自我编辑生成策略的质量
  • 论文对比了以下基线：
    • 1）ICL（上下文学习） ：Llama-3.2-1B-Instruct 直接通过少量示例 Prompt 并查询测试输入
    • 2）TTT + 自我编辑（无 RL 预训练） ：模型使用少量示例和合成增强进行测试时训练，但未通过 RL 优化增强或训练配置
    • 3）Oracle TTT ：模型使用人工设计的最优配置进行测试时训练，作为方法的上限
  • 结果如表1 所示
    • SEAL 显著提高了适应成功率：72.5% vs. 20%（无 RL 训练的自我编辑）和 0%（无适应），但性能仍低于 Oracle TTT，表明还有改进空间

Knowledge Incorporation

  • 论文在 Quen2.5-7B 上实验，整合来自 SQuAD 段落的新事实内容
  • 论文进行了两轮 ReST\(^{EM}\) 训练，每轮包含 50 个上下文（详见 附录B）。论文对比了以下基线方法：
    • 1）基础模型 ：未经适应的预训练模型，无法访问段落
    • 2）仅训练段落 ：模型直接对段落进行标准语言建模损失微调，无合成数据
    • 3）训练段落 + 合成数据 ：模型训练段落及自我生成的推论，与 SEAL 相同但无 RL 训练
    • 4）训练段落 + GPT-4.1 合成数据 ：模型训练段落及通过 OpenAI API 从 GPT-4.1 生成的推论
  • 表2 报告了两种设置下的平均无上下文 SQuAD 准确率：单段落（\(n=1\)）和持续预训练（CPT, \(n=200\)）
    • 在单段落设置中，直接微调段落仅比冻结基础模型略优（33.5% vs. 32.7%），表明原始数据不足
    • 使用 GPT-4.1 生成的合成数据将准确率提升至 46.3%，比仅段落基线高 12.8 个百分点
    • 使用 Quen-2.5-7B 生成的合成数据达到 39.7%，提高 6.2 个百分点
    • 经过强化学习后，SEAL 进一步将准确率提升至 47.0% ，显著优于 GPT-4.1 的合成数据，尽管模型规模小得多
  • 在 CPT 设置中，模型在一次持续预训练中整合了 \(n=200\) 个段落的信息，随后在所有 974 个对应问题上评估
  • SEAL 再次超越所有基线，达到 43.8% 准确率。尽管绝对性能低于单段落设置（可能由于梯度干扰增加），但相对改进一致，表明 SEAL 的编辑策略能够泛化到原始 RL 设置之外
  • 图4 跟踪了每轮 RL 迭代后的准确率。两轮迭代后 SEAL 即超越 GPT-4.1 数据；后续迭代收益递减，表明策略快速收敛为将段落提炼为易学习的原子事实（定性示例见图5）。所有结果均使用调优超参数（见附录B）

Limitations

Catastrophic Forgetting

  • 论文提出 self-edit 语言模型（SEAL）的一个关键动机是实现持续学习（Continual Learning）的终极目标——让模型能够随着时间的推移不断整合新信息，无论是通过与环境的主动交互还是通过标准训练
  • 虽然之前的实验评估了 SEAL 在独立编辑场景下的适应能力，但更雄心勃勃的目标是支持连续的编辑序列 ：模型能否在保留已有知识的同时，反复适应新信息？
  • 这一问题直接关联到灾难性遗忘（2014, 2015）的挑战，即新更新会破坏过去的学习成果
    • 当前的训练设置并未显式优化知识保留，但论文旨在建立一个基线，评估SEAL在没有专门机制的情况下处理连续 self-edit 的能力
    • 为了测试这一点，论文在知识整合领域模拟了一个持续学习场景：模型接收一系列测试段落，每个段落触发一次新的 self-edit
    • 每次更新后，论文重新评估模型在所有已见任务上的表现，以衡量其知识保留能力
  • 如图6所示，随着编辑次数的增加，模型在早期任务上的表现逐渐下降，这表明SEAL仍然容易受到灾难性遗忘的影响
    • 尽管如此，它能够在多次更新后避免完全崩溃，这表明未来仍有改进空间
    • 未来的工作可以通过奖励塑形（2020, 2024）来增强这一能力，例如惩罚对早期任务的回归，或整合持续学习策略，如零空间约束编辑（2025）或表示叠加（2019）

Computational overhead

  • TTT（Test-Time Training）奖励循环的计算成本显著高于其他用于 LLM 的 RL 方法
  • 例如，基于人类偏好的奖励信号通常只需要一次模型前向传播，而基于验证解的奖励可能仅依赖简单的模式匹配（如正则表达式）
  • 相比之下，论文的方法需要对整个模型进行微调和评估以计算奖励——每次 self-edit 评估大约需要30-45秒，带来了显著的开销（详见附录B.5）

Context-dependent evaluation

  • 当前的实例化假设每个上下文都配有一个明确的下游任务：Few-shot 示例附带一个保留的查询对，每个段落捆绑了参考问答
  • 这种耦合简化了奖励计算，但阻碍了 SEAL 的 RL 训练扩展到未标注语料库
  • 一个潜在的解决方案是让模型不仅生成 self-edit，还为每个段落生成自己的评估问题（例如草拟问答项或合成测试用例），同时保留原始内容在上下文中
  • 这些模型编写的查询可以提供强化学习所需的即时监督，从而将适用性扩展到缺乏外部问答集的通用训练领域
  • 合成数据生成（Synthetic Data Generation） ：
    • 合成数据在训练中的应用日益广泛，从大规模预训练数据集（2023; 2024; 2024）到任务特定的数据增强（2023; 2024）和指令微调集（2023; 2023）
    • Yang 等人（2025）通过基于图的 Prompt 生成合成数据
    • SEAL 在此基础上，利用强化学习训练生成策略，直接最大化合成数据在梯度更新中的下游效用，而非依赖手动调整的静态启发式方法
  • 知识更新（Knowledge Updating） ：
    • 近期研究尝试通过权重更新修改或注入事实知识
    • 部分方法直接定位与特定事实相关的参数（2022; 2022; 2023）
    • 另一些则利用上下文信息生成额外的微调数据（2024; 2024; 2025; 2025）
    • 论文采用后者，参考 Akyurek 等人（2024）提出的逻辑蕴涵生成和 Lampinen 等人（2025）展示的蕴涵微调优于上下文学习的结果
    • SEAL 通过强化学习训练模型生成更优的微调数据，进一步扩展了这些方法
    • Park 等人（2025）表明，直接生成问答对（QA）的 Prompt 优于蕴涵式 Prompt
    • 由于 SEAL 框架对 self-edit 数据的格式无关，它同样可以训练生成 QA 对或其他输出格式
  • 测试时训练（TTT, Test-Time Training） ：
    • 测试时训练基于输入临时调整模型权重（2020; 2022; 2024）。Akyurek 等人（2025）表明，TTT 与上下文学习结合可在 Few-shot 设置中超越标准 ICL
    • SEAL 在内部优化中整合了 TTT，利用其高效性执行多次更新，并奖励带来最大性能提升的数据生成策略
  • LLM 的强化学习（Reinforcement Learning for LLMs） ：
    • 强化学习在改进 LLM 行为中发挥核心作用，最初通过 RLHF（2022）实现
    • 近期研究利用可验证奖励直接优化任务成功率（2022; 2024; 2025）
    • SEAL 将强化学习应用于优化 self-edit 数据的生成，而非最终答案或推理轨迹的修订
  • 元学习与自修改系统（Meta-Learning and Self-Modifying Systems） ：
    • SEAL 通过外部优化循环学习适应策略（即如何生成有效的 self-edit），体现了元学习原则（2001; 2017; 2025），其目标是学习如何高效地从任务上下文中学习
    • 元学习同样已应用于强化学习领域在该领域中，模型通过元目标进行训练，以快速适应新任务
    • 这类工作的一个自然延伸是自指网络（self-referential networks），即模型自行修改自身参数（1992; 2022）
    • 在大型语言模型领域，近期的研究已将元学习原则应用于改进 LLM 的适应性[2024;2023]
    • 值得注意的是，Hu等人（2023）训练了一个较小的模型，使其在对语料库进行微调时输出特定于标记的权重，以解决与我们类似的知识整合任务
    • 然而，SEAL 通过利用模型现有的生成能力来参数化更新，从而在跨领域场景中展现出更强的通用性
  • 自我改进（Self-Improvement） ：
    • 近期研究涵盖自我改进或自训练的多种方法
    • RLAIF（2022; 2024）和自我奖励语言模型（2024; 2025）利用模型自身提供奖励信号，基于判断输出比生成更容易的观察（2025）
    • 其他工作通过多数投票或模型置信度作为强化学习奖励，在无真实标签的情况下提升数学任务性能（2023; 2024; 2025; 2025）
    • 然而，这些方法受限于模型的当前评估能力和自一致性，相比之下，SEAL 通过与外部数据的交互实现自我改进，为更具扩展性的路径提供了可能

Discussion and Conclusion

  • Villalobos等人（2024）预测，到2028年，前沿 LLM 将完成对所有公开人类生成文本的训练
  • 作者认为，这一迫近的“数据墙”将迫使人们采用合成数据增强 ，一旦网络规模的语料库耗尽，进展将取决于模型自主生成高效用训练信号的能力
  • 自然的下一步是元训练一个专用的 SEAL 合成数据生成模型，生成新的预训练语料库，使未来模型能够在无需额外人类文本的情况下扩展并实现更高的数据效率
  • 我们可以设想一个未来场景：大语言模型能够消化新数据（如学术论文），并利用上下文中的数据和已有知识生成大量解释和推论
    • 这种自我表达和自我优化的迭代循环可能使模型即使在缺乏外部监督的情况下，也能在罕见或代表性不足的主题上持续改进
  • 此外，尽管现代推理模型通常通过强化学习生成思维链（CoT）轨迹，但 SEAL 可以提供一种互补机制，让模型学习何时以及如何更新自身权重
    • 这两种方法可以协同作用：模型可能在推理过程中选择更新权重以引导当前轨迹，或在完成推理后将关键见解提炼到参数中——通过内部化学习提升未来推理能力
  • 这种持续优化的循环对于构建代理系统（Agentic Systems）也很有前景——这些模型在长期交互中运行，并动态适应不断变化的目标
    • 代理模型必须逐步获取和保留知识
    • 论文的方法通过支持结构化的自我修改来实现这种行为：在交互后，代理可以合成一个 self-edit，触发权重更新
    • 这可能使代理能够随时间发展，使其行为与先验经验对齐，并减少对重复监督的依赖
  • SEAL 证明了大语言模型在预训练后不必保持静态：通过学习生成自己的合成 self-edit 数据并通过轻量级权重更新应用这些数据，它们可以自主整合新知识并适应新任务
  • 展望未来，作者计划将 SEAL 框架扩展到预训练、持续学习和代理模型中，最终使语言模型能够在数据受限的世界中自我学习和扩展

附录A 实验细节：Few-shot 学习（Experimental Details: Few-shot Learning）

  • 模型与设置 ：论文使用Llama-3.2-1B-Instruct（2024）作为基础语言模型。由于该模型未针对ARC进行专门训练，其解决ARC任务的能力有限。为了进行受控评估，论文从训练集和评估集中筛选了一小部分可通过最优TTT超参数解决的 ARC 问题
    • 训练集 ：从训练集中选择了11个ARC任务作为RL优化的环境
    • 评估集 ：从评估集中选择了8个不同的ARC问题用于衡量泛化性能。这8个问题明确筛选为适合TTT的场景
  • 这些集合的选择是为了隔离 self-edit 学习的效果，而非通用ARC能力
  • 训练流程 ：论文在11个训练问题上使用单轮强化学习训练SEAL。对于每个问题，模型生成15个 self-edit 配置，包括：
    • 数据增强策略 ：控制是否包含基本示例、基于尺寸的增强、链式推理和重复输入
    • 优化配置 ：指定学习率、训练轮数和损失函数类型（例如是否对所有 Token 或仅最终输出 Token 计算损失）
  • 每个配置通过测试时训练（TTT）评估，并分配二元奖励：如果适应后的模型生成正确解则为1，否则为0
  • 评估细节 ：对于8个保留的评估任务，模型生成5个 self-edit 配置，共40个配置。成功率衡量为适应后生成正确输出的配置百分比

附录B 实验细节：知识整合（Experimental Details: Knowledge Incorporation）

  • 模型与设置 ：论文在知识整合实验中使用Quen-2.5-7B基础模型（2025）。论文将SQuAD数据集v1.1（2016）重新用于无上下文段落的问答任务
  • RL训练流程 ：论文运行2轮 ReST\(^{EM}\) 训练（2025）
    • 每轮从 SQuAD 训练集中抽取 50 个 上下文-问题-答案 三元组。对于每个上下文，采样 5 个 self-edit 生成（温度=1），并通过 3 个随机种子评估每个生成
  • 合成数据生成与微调细节 ：在所有模型中，论文通过 Prompt 生成段落的推论来生成合成数据
  • 评估细节 ：论文在 SQuAD 评估集的 200 个段落子集上评估，共974个问题
    • 使用 GPT-4.1（2025）通过 OpenAI API 进行自动评分
  • 计算资源 ：所有实验在 2×H100 或 2×H200 上运行
    • 使用 DeepSpeed ZeRO-3（2020）进行 ReST\(^{EM}\) 训练的 SFT，使用 vLLM（2023）进行高效推理

附录C Prompting

  • 近期研究表明，强化学习基线和结果对 Prompt 高度敏感。论文在知识整合设置中测试了 4 种额外的 self-edit Prompt；五种 Prompt 如下：
    • 1）推论（Implications）
    • 2）长推论（Implications-long）
    • 3）超长推论（Implications-very-long）
    • 4）重写（Rewrite）
    • 5）自问答（Self-QA）
  • 结果显示，尽管通过 Prompt 生成长响应可以提高性能，但以这些 Prompt 为基础的 RL 训练能带来更大的改进
    • 在所有情况下，ReST\(^{EM}\)将性能提升了约 6 到 11 个百分点

注：本文包含 AI 辅助创作

Paper Summary

  • 整体内容总结
    • 论文提出了一种新颖的，用于预训练大语言模型的新范式，强化预训练（Reinforcement Pre-Training, RPT）
    • 通过将下一词预测任务重构为可验证的推理任务，并应用基于正确性的强化学习奖励，RPT 使大语言模型能够在预训练期间利用扩展计算构建更强的基础推理能力
    • 实验表明：RPT 提升了下一词预测的准确性 ，在数学 和通用推理基准的 Zero-Shot 设置 中表现出色 ，并为后续强化学习微调 提供了更好的起点
    • RPT 通过从根本上重新思考预训练目标本身，为开发更强大、更通用的智能大语言模型提供了新的方向
  • 论文将 Next-Token Prediction（NTP）任务重新定义为一种通过强化学习训练的推理任务，模型在正确预测给定上下文的 Next Token 时会获得可验证的奖励
  • RPT 提供了一种可扩展的方法，能够利用海量文本数据实现 通用强化学习（general-purpose RL） ，而 无需依赖特定领域的标注答案
    • 通过激励 Next Token 推理能力，RPT 显著提升了语言模型在预测 Next Token 时的准确性
  • RPT 为后续的强化微调提供了强大的预训练基础
  • 实验结果表明，随着训练计算量的增加， Next-Token Prediction 的准确性持续提升
    • 作者认为，这些结果证明了 RPT 是一种有效且有前景的规模化范式，能够推动语言模型预训练的进步

Introduction and Discussion

  • LLM 在各种任务中展现出了卓越的能力，这主要得益于基于海量文本语料的 Next-Token Prediction 目标的可扩展性
    • 这种自监督范式已被证明是一种高效的通用预训练方法
  • RL 已成为一种强大的技术，可用于微调大语言模型，使其与人类偏好对齐或增强特定技能，例如复杂推理 (2022, 2023, 2024)
  • 但当前强化学习在大语言模型训练中的应用面临可扩展性和通用性挑战
    • RLHF 在对齐任务中表现有效，但其依赖昂贵的人类偏好数据，且学习到的奖励模型容易受到奖励破解（reward hacking）的影响，从而限制了可扩展性
    • RLVR 是一种 RLVR 方法
      • RLVR 利用客观的、基于规则的奖励（通常来自问答对），这种方法能够缓解奖励破解问题
      • 但 RLVR 通常受限于带有可验证答案的标注数据的稀缺性，因此其 应用范围仅限于特定领域的微调 ，而 非通用预训练
  • 在这项工作中，作者提出了 强化预训练（RPT） ，这是一种新颖的范式，旨在弥合可扩展的自监督预训练与强化学习能力之间的差距
    • RPT 将 Next-Token Prediction 任务重新定义为一种 Next Token 推理过程
    • 对于预训练语料中的任何给定上下文，模型会被激励在预测 Next Token 之前对其进行推理
    • 模型会根据其预测结果与语料中真实 Next Token 的匹配程度，获得一种可验证的内在奖励
    • 这种方法将通常用于 Next-Token Prediction 的海量未标注文本数据，转化为一个适用于通用强化学习的庞大数据集，而无需依赖外部标注或特定领域的奖励函数
  • 这种方法具有以下几个关键优势：
    • 1）可扩展性与通用性 ：RPT 利用了与标准 Next-Token Prediction 相同的海量未标注文本数据，将其转化为通用强化学习的大规模数据集，无需外部标注
    • 2）减少奖励破解风险 ：通过使用直接的、基于规则的奖励信号（即预测 Next Token 的正确性），RPT 从根本上降低了复杂学习奖励模型中常见的奖励破解风险
    • 3）促进深度理解与泛化 ：通过显式激励 Next Token 推理模式，RPT 鼓励模型深入理解上下文，而非简单地记忆 Next Token
      • 模型学会探索和验证关于“为何某个 Token 应该出现”的假设，从而构建更鲁棒的表示
    • 4）推理过程的计算分配 ：预训练中的内部推理过程使模型能够为每个预测步骤分配更多的“思考”或计算资源，类似于在训练时为每个 Token 应用推理时扩展，从而直接提升 Next-Token Prediction 的准确性
  • 论文的实验表明：
    • RPT 显著提升了 Next-Token Prediction 的准确性
    • RPT 为后续的强化微调提供了更鲁棒的预训练基础，从而在最终任务中表现更优
    • 扩展曲线显示，在 RPT 框架下，增加训练计算量能够持续提升 Next-Token Prediction 的准确性 ，这表明 RPT 是一种可持续的规模化策略
    • 这些结果证明了强化预训练是一种有效且有前景的新范式，能够推动大语言模型预训练的进步
  • 论文的主要贡献如下：
    • 提出了一种新的规模化范式： 强化预训练（RPT） ，将 Next-Token Prediction 重新定义为通过强化学习训练的推理任务，并利用预训练语料直接生成的内在可验证奖励
    • RPT 提供了一种可扩展且通用的强化学习预训练方法，通过基于规则的奖励减少奖励破解风险，并通过激励 Next Token 推理模式（而非机械记忆（rote memorization））促进泛化能力
    • RPT 显著提升了 Next-Token Prediction 的准确性，并展现出良好的扩展特性，即性能随着训练计算量的增加而持续提升
    • RPT 为后续的强化微调提供了更强的预训练基础，并提升了在多种下游任务中的 Zero-Shot 性能

Preliminary

Next-Token Prediction, NTP

  • Next-Token Prediction 是现代大语言模型的基本训练目标 (2022)
  • 给定训练语料中的输入序列 \(x_{0}\cdots x_{T}\)，模型的训练目标是最大化以下目标函数：
    $$
    \mathcal{J}_{\text{NTP} }(\theta)=\sum_{t=1}^{T}\log P(x_{t}\mid x_{0},x_{1},\ldots,x_{t-1};\theta),
    $$
    • 其中 \(\theta\) 表示语言模型的参数

Reinforcement Learning with Verifiable Rewards, RLVR

  • RLVR 利用强化学习目标来增强具有可验证答案的特定技能 (2023)
  • RLVR 需要一个标注的问答对数据集 \(\mathcal{D}=\{(q,a)\}\)
  • 对于特定的问答对 \((q,a)\in \mathcal{D}\)，大语言模型 \(\pi_{\theta}\) 会生成一个响应 \(o\sim \pi_{\theta}(\cdot \mid q)\)
    • 然后使用一个确定性的验证器 \(\mathcal{V}\) 计算可验证奖励 \(r=\mathcal{V}(o,a)\)，模型的训练目标是最大化期望奖励：
      $$
      \mathcal{J}_{\text{RLVR} }(\theta)=\mathbb{E}_{(q,a)\sim \mathcal{D},o\sim \pi_{\theta}(\cdot|q)}\left[r(o,a)\right].
      $$

Reinforcement Pre-Training

Pre-Training Task: Next-Token Reasoning

  • 论文提出了 Next Token 推理任务用于语言建模
  • 给定训练语料库中的输入序列 \(x_0 \cdots x_T\)，对于每个位置 \(t \in \{1, \ldots, T\}\)，前缀 \(x_{ < t}\) 被视为上下文，而真实的 Next Token 是 \(x_t\)
  • 在 Next Token 推理任务中，模型 \(\pi_\theta\) 需要在生成对 Next Token 的预测 \(y_t\) 之前生成一个思维链（chain-of-thought）推理序列，记为 \(c_t\)
  • 模型的整体响应为 \(o_t = (c_t, y_t)\)，其中 \(o_t \sim \pi_\theta(\cdot \mid x_{ < t})\)
  • 如图 2 所示， Next Token 推理的长思维链过程可能涉及多种推理模式，例如头脑风暴、自我批判和自我修正
  • Next Token 推理任务将预训练语料库重构为大量的推理问题集合，使预训练从学习表面的 Token-level 关联转向理解其背后的隐藏知识，并使强化学习的扩展成为可能

Pre-Training with Reinforcement Learning

  • RPT 通过在线策略强化学习训练 LLM 执行 Next Token 推理，如图 3 所示
  • 对于上下文 \(x_{ < t}\)，RPT 提示语言模型 \(\pi_\theta\) 生成 \(G\) 个响应（思维轨迹），记为 \(\{o^i_t\}_{i=1}^G\)
  • 每个响应 \(o^i_t = (c^i_t, y^i_t)\) 包含一个思维链推理序列 \(c^i_t\) 和一个最终的预测序列 \(y^i_t\)
  • 为了验证 \(y^i_t\) 的正确性，论文引入了前缀匹配奖励（prefix matching reward），该奖励支持验证跨越多 Token 或涉及词汇表外 Token 的预测
  • 符号定义如下：
    • \(x_{\geq t}\) 表示真实补全序列，其字节（byte）序列表示为 \(\overline{x}_{\geq t}\)
      • 问题：为什么不是单个 token? 论文每次仅预估下一个 token 吧？到底是每次生成单个 token 还是多个 token？
      • 猜测：这里是表示每次推理时，可以只看一个 token，也可以看多个 token
    • \(y^i_t\) 表示预测的序列，其字节（byte）序列表示为 \(\overline{y}^i_t\)
    • \(\overline{y}^i_t\) 的字节长度记为 \(l\)
    • 真实补全序列中 Token 的累积字节长度定义为有效边界，记为 \(\mathcal{L}_{gt}\)
      • 理解：这里的有效边界是一个整数集合，表示有效的长度值的集合
  • 形式上，对于上下文 \(x_{ < t}\) 的第 \(i\) 个输出，奖励 \(r^i_t\) 定义为：
    $$
    r^i_t = \begin{cases}
    1 & \text{if } \overline{y}^i_t = \overline{x}_{\geq t}[1:l] \text{ and } l \in \mathcal{L}_{gt} \\
    0 & \text{otherwise}
    \end{cases},
    $$
    • 如果预测的字节序列是真实补全字节序列的精确前缀且其长度 \(l\) 匹配任何有效 Token 边界 ，则奖励为 1
    • \(\overline{y}^i_t\) 表示预测的字节序列
    • \(\overline{x}_{\geq t}\) 表示真实补全的字节序列
  • 令 \(\mathcal{D}\) 为所有 \(\{x_{ < t}\}_{t=1}^T\) 的集合，模型训练的目标是最大化期望奖励：
    $$
    \mathcal{J}_{\text{RPT} }(\theta) = \mathbb{E}_{(x_{ < t}, x_{\geq t}) \sim \mathcal{D}, \{o^i_t\}_{i=1}^G \sim \pi_\theta(\cdot|x_{ < t})} \left[r^i_t \right].
    $$

Pre-Training Setup

  • 论文使用 OmniMATH 数据集（2024）进行强化预训练。OmniMATH 包含 4,428 个竞赛级数学问题及其解答，数据来自 AoPS Wiki 和 AoPS 论坛等官方网站
  • 由于许多 Token 即使无需推理也容易预测，论文在强化预训练前进行了 Token-level 数据过滤
    • 论文使用 Deepseek-R1-Distill-Queen-1.5B 作为小型代理模型，计算每个 Token 在前 16 个候选 Token 上的代理模型熵
    • 通过应用熵阈值，论文过滤掉低熵位置，优先训练那些需要更多计算努力预测的挑战性 Token（注：高熵的 token 是较难预测的）
  • 在所有实验中，论文以 Deepseek-R1-Distill-Queen-14B（2025）作为基础模型
    • R1-Distill-Queen-14B 因其基本的推理能力而成为强化学习的良好起点
  • 论文使用 verl 库（2025）实现训练框架，并使用 vllm 进行推理
  • 论文采用 GRPO 算法（2025），具体超参数详见附录 B
  • 训练时，论文采用 8k 的训练长度，学习率为 \(1 \times 10^{-6}\)，KL 惩罚为零，批次大小为 256 个问题，每个问题采样 \(G=8\) 个响应，在 rollout 过程中使用温度为 0.8
  • 从每个响应中，论文直接提取最后一个 \(\backslash\)boxed{ } 内的完整序列作为模型对 Next Token 的预测
  • 从第 500 步开始 ，论文使用动态采样以提高训练效率（2025），主实验的总训练步数为 1,000
    • 补充：这里的动态采样是 DAPO 中的动态采样技术，把奖励全为 0 或者全为 1 的 Prompt/样本 丢弃掉
  • 提示模板及其变体在附录 D 中讨论

Evaluation of Pretrained Models

  • 模型预训练完成后，我们可以直接在下游任务上进行 Next-Token Prediction 和强化微调
  • 论文通过以下设置展示强化预训练如何提升大语言模型的语言建模能力和推理能力
  • 语言建模（Language Modeling）
    • 基于 Next Token 推理目标，论文的模型可以自然地用于语言建模
    • 论文报告 Next-Token Prediction 准确率，以评估 RPT 的语言建模性能和扩展性
  • 下游任务的强化微调（Reinforcement Fine-Tuning on Downstream Tasks）
    • 论文以预训练后微调的方式对 RPT 模型进行持续的强化微调
    • 由于 RPT 将预训练过程与强化学习对齐，预训练与后续强化微调之间的目标差距被最小化
    • 论文评估强化预训练过程是否进一步提升了最终任务的性能

Experiments

Language Modeling

  • 论文在 OmniMATH 的 200 个验证集样本上评估语言建模性能
  • 根据第 3.3 节描述的基于熵的数据过滤策略，论文根据难度对验证集中的 Token 位置进行分类
    • 论文使用 R1-Distill-Queen-14B 计算每个 Token 位置的熵，并根据熵是否超过阈值 0.5、1.0 和 1.5 将位置划分为简单、中等和困难三类
    • 为了比较，论文报告了 R1-Distill-Queen-14B 在两种评估方式下的性能：
      • (1) 标准 Next-Token Prediction ，选择概率最高的 Token ；
      • (2) Next Token 推理，生成思维链后再进行最终预测
    • 论文还包含了 Qwen2.5-14B 的结果，因为它是 R1-Distill-Queen-14B 的基础模型
  • 如表 1 所示，RPT-14B 在所有难度级别上的 Next-Token Prediction 准确率均高于 R1-Distill-Queen-14B
    • 值得注意的是，它的性能与显著更大的模型 R1-Distill-Queen-32B 相当（图 4）
  • 这些结果表明，强化预训练能有效捕捉 Token 生成背后的复杂推理信号，并具有提升大语言模型语言建模能力的强大潜力

Scaling Properties of Reinforcement Pre-Training

  • 本节论文研究强化预训练的扩展性
  • 自然语言语料库上的 Next Token 预训练损失在模型大小、训练 Token 数量和训练计算量方面通常遵循幂律衰减（2020, 2022）
  • 论文使用以下幂律形式建模训练计算量 \(C\) 与性能的关系：
    $$
    P(C) = \frac{A}{C^\alpha} + P^*, \tag{5}
    $$
    • 其中 \(P(C)\) 表示验证集上的 Next-Token Prediction 准确率，\(P^*\)、\(\alpha\) 和 \(A\) 是待估计的参数
  • 论文在不同训练步数（100、200、400、800、1000 和 1200）下评估 RPT 的 Next-Token Prediction 准确率，并将其转换为相应的训练计算量
  • 为了评估数据难度的影响，论文考虑了基于熵阈值 0.5（简单）、1.0（中等）和 1.5（困难）过滤的验证集分割
    • 更高的阈值对应更具挑战性的输入
    • 对于每个难度级别，论文根据公式 (5) 拟合结果，并使用决定系数 \(R^2\) 衡量拟合优度（Goodness of fit）
    • 理解：按照 15W 词表算：
      • 熵为 0.5 对应这最大的概率差不多是 0.970；
      • 熵为 1.0 对应这最大的概率差不多是 0.936；
      • 熵为 1.5 对应这最大的概率差不多是 0.901；
      • 更多可视化详情见附录
  • 如图 5 所示，随着训练计算量的增加，RPT 的 Next-Token Prediction 准确率持续提升
  • 所有难度级别的高 \(R^2\) 值表明拟合曲线能准确捕捉性能趋势（理解：说明在不同难度上，均能很好的拟合到公式 (5) 上）

Reinforcement Fine-Tuning with RPT

  • 为了研究 RPT 模型是否能通过 RLVR 更有效地微调，论文从 Skywork-OR1（2025）随机采样带有可验证答案的问题进行进一步训练
    • 论文使用 256 个样本进行训练，200 个样本进行测试
    • 遵循 Skywork-OR1 的数据过滤流程（2025），论文使用 R1-Distill-Queen-32B 识别训练中的挑战性实例
    • 论文将训练批次大小和 PPO 小批次大小均设为 64，并训练模型 15 个周期
    • 评估时，验证的最大 Token 数设为 32,000，温度为 0.6
  • 如表 2 所示：
    • 强化预训练模型在使用 RLVR 进一步训练时达到了更高的上限
    • 当使用相同的 Next-Token Prediction 目标持续训练相同数据时，模型的推理能力显著下降，而后续的 RLVR 仅带来缓慢的性能提升
      • 理解：可以观察到，直接进行普通 NPT 的持续预训练（目标和 RPT 相同）会导致推理能力大幅下降；猜测这里是训练太多次，发生了过拟合了！
        • 引申问题：这里的 普通 NPT 和 RPT 训练轮次是相同的吗？
  • 这些结果表明，在数据有限的情况下，强化预训练可以快速将从 Next Token 推理中学到的强化推理模式迁移到最终任务中

Zero-Shot Performance on End Tasks

  • 论文评估了 RPT-14B 在下游任务上的 Zero-Shot 性能
  • 作为比较，论文评估了 R1-Distill-Queen-14B 和 R1-Distill-Queen-32B 的 Next-Token Prediction 性能，以及 RPT-14B 与 R1-Distill-Queen-14B 的推理性能
  • 论文的评估涉及两个广泛认可的基准：
    • MMLU-Pro（2020），一个综合性多任务理解基准，评估大语言模型在不同领域的表现；
    • SuperGPQA（2025），一个涵盖 285 个学科的研究生级推理问题的大规模基准
  • 在推理设置下，论文将最大 Token 数设为 12,288，温度为 0.8
  • 遵循先前工作（2024, 2025），论文使用多项选择题格式进行评估并报告准确率
  • 如表 3 所示
    • RPT-14B 在所有基准上均优于 R1-Distill-Queen-14B（无论是标准 Next-Token Prediction 还是作为推理模型评估）
    • RPT-14B 还超越了显著更大的 R1-Distill-Queen-32B（在 Next-Token Prediction 模式下），在 SuperGPQA 上提升了 7 分，在 MMLU-Pro 上提升了约 22 分
  • 每个基准的详细分科结果见附录 C

Next-Token Reasoning Pattern Analysis

  • 论文分析了 Next Token 推理与显式问题解决（explicit problem solving）在推理模式上的差异
    • 根据先前研究（2024, 2025），论文统计了模型响应中包含推理关键词（如“break down”、“alternatively”）的比例
  • 论文的分析比较了两种模型在 OmniMATH 数据集上的思维过程：
    • R1-Distill-Queen-14B 用于问题解决
    • RPT-14B 用于 Next Token 推理
    • 对每个模型，采样 200 个响应
  • 论文将推理模式分为六类：
    • 转换（切换策略）、反思（自我检查）、分解（分解问题）、假设（提出并验证假设）、发散思维（探索可能性）和演绎（逻辑推理）

      transition (switching strategies), reflection (self-checking), breakdown (decomposing the problem), hypothesis (proposing and verifying assumptions), divergent thinking (exploring possibilities), and deduction (logical inference).

  • 如图 6 所示，RPT-14B 的 Next Token 推理过程与 R1-Distill-Queen-14B 的问题解决过程显著不同
    • RPT-14B 相对 R1-Distill-Queen-14B： 假设模式的使用量增加了 161.8%
    • RPT-14B 相对 R1-Distill-Queen-14B：演绎模式的使用量增加了 26.2%
    • 问题解决过程（R1-Distill-Queen-14B）更依赖分解模式（breakdown） ，这表明 Next Token 推理引发的推理过程在性质上与结构化问题解决不同
  • 表 4 展示了一个推理模式的例子
    • 该例子揭示了模型参与的是一个深思熟虑的过程，而非简单的模式匹配
    • 它分析了更广泛的语义上下文（“calculating vector magnitude”），识别关键短语（“go over some…”），然后进行头脑风暴并权衡多个可能的延续
    • 这涉及假设生成（“the next part is likely going to be…”）、替代方案考虑（“Alternatively, it could be…”）以及对结构线索（“markdown with headers”）甚至细粒度 Token-level 细节（“could have a space”）的反思
    • 这种多方面的推理，既包含高级语义理解，又包含低级文本特征，展示了模型通过推理探索推断 Next Token 的努力，与 RPT 培养超越表面关联的更深层次理解的目标一致
  • 更多例子见附录 F

Scaling Paradigms of Large Language Models

  • LLM 的进步主要由两个扩展维度驱动：
    • 训练时计算（training-time compute）（2022a）：通过大幅增加模型参数和训练数据，以下一词预测（next-token prediction）作为预训练任务
    • 测试时计算（test-time compute）（2025a）：测试时扩展（2024）通过延长推理计算时间提升大语言模型的推理能力
  • RPT 独特地整合了上述原则，超越现有扩展范式，将每一词预测任务重构为推理任务

Reinforcement Learning for Large Language Models

  • 强化学习在大语言模型的后训练阶段发挥了关键作用
    • RLHF（2022）通过人类偏好数据微调预训练语言模型以提升对齐性
    • 除对齐外，大规模强化学习还被用于增强语言模型的推理能力（2025）
  • 最相关的工作（2025）鼓励语言模型为下一词预测生成有帮助的推理过程
    • 基于帮助性的奖励容易被生成的推理中重复目标词所“破解” ，这种捷径可能损害模型性能
    • 相比之下，论文使用下一词预测的正确性作为基于规则的奖励信号，以最小化奖励破解风险

Future Work

  • RPT 的初步探索仍存在一些局限性
    • 论文的实验主要基于 14B 参数的模型，没有在更大的模型进行测试
    • 虽然 RPT 方法设计为通用，但当前预训练语料库 主要由数学文档组成；
      • 未来工作将探索其在更广泛的通用领域文本上的有效性
    • RPT 训练是从一个具备基础推理能力的模型初始化（R1-Distill-Qwen-14B）的；
      • 后续可以研究从标准基础语言模型开始的 RPT 训练
      • 这将为 RPT 基础性影响提供进一步分析和结论
  • 未来工作可从以下方向推进：
    • 扩展训练语料库的规模和领域覆盖范围，利用大规模通用互联网数据进行强化预训练
    • 增加训练计算资源以突破性能边界
    • 建立强化预训练的扩展定律（scaling laws），指导大语言模型的扩展
    • 探索将混合思维（hybrid thinking）（2025）与 RPT 结合，通过自适应触发下一词推理实现细粒度的适应性思考

附录 A Design Choices of Reward

  • 除了第 3 节描述的基于前缀匹配的奖励机制外，论文还研究了其他几种奖励函数变体以评估其对强化预训练的影响
    • 变体一：首词匹配（first-token matching）
      • 在此设置中，奖励仅反映模型预测 \( y_t^i \) 的首词是否与真实下一词 \( x_t \) 匹配，忽略预测中首词之后的所有词
    • 变体二：探索了“密集奖励”（dense reward）方案：
      • 正确预测的下一词（即 \( y_t^i[0] = x_t \)）获得满分奖励（如 1）；
      • 对于错误预测（\( y_t^i[0] \neq x_t \)），奖励为一个较小的正值 ，具体为语言模型生成该错误词的概率 \( P(y_t^i[0] \mid x_{ < t}; \theta) \)
        • 问题：为什么是错误词的概率？岂不是错误词的概率越大，奖励越大，应该是正确词的概率吧
      • 这提供了比二元奖励更密集的反馈信号
    • 变体三：条件性应用密集奖励结构：
      • 仅当给定前缀 \( x_{ < t} \) 的 \( G \) 次采样中至少有一次正确预测下一词时 ，才使用密集奖励；
      • 若所有 \( G \) 次采样均错误，则应用其他奖励方案（如零奖励或统一的小惩罚）
  • 实验表明，这些替代奖励与前缀匹配奖励相比，性能相当
    • 表明强化预训练框架对这些奖励信号的修改具有较强的鲁棒性 ，其核心优势可能对这些特定选择不敏感，至少在测试的变体范围内如此

附录 B Hyperparameters Used for Reinforcement Pre-Training

  • 表 5 展示了第 4 节中强化预训练的详细超参数
  • 论文遵循精确策略强化学习（2025）的设置，将熵损失系数设为 0

附录 C Detailed Results on End Tasks

  • 表 6 和表 7 展示了通用终端任务基准的详细分类性能
  • RPT-14B 模型在大多数类别中表现优于 R1-Distill-Qwen-14B 和 R1-Distill-Qwen-32B

附录 D Impact of Prompt Templates

  • 论文探索了不同提示模板对初始下一词推理性能的影响
  • 表 10 展示了七种模板变体，这些模板使用不同指令片段，并以不同形式包装上下文
  • 如表 8 所示，清晰的 Prompt 能很大程度提升初始表现的准确性
    • 第 4 节的强化预训练实验使用了“v0”模板，其他模板变体的优化留待未来工作

附录 E Keywords for Reasoning Pattern Analysis

  • 表 9 列出了第 4.5 节中用于推理模式分析的模式组和关键词

附录 F Case Studies

  • 表 11 展示了 RPT-14B 在下一词推理任务中的三个案例，包括模型对数学问题和文本上下文的推理过程
  • 这些案例揭示了模型如何通过多角度思考生成最终预测

附录：概率和熵的关系图

  • 关键词：entropy curve；熵和概率；概率和熵；曲线图；

  • 假定只有一个 token 的值较大，其他 token 概率相同，此时的熵和最大概率的关系是如何的？

  • 可视化最大概率和熵的关系的代码

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    import numpy as np
    import matplotlib.pyplot as plt

    # 设置中文显示
    plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
    plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

    def calculate_entropy(p, n=150000):
        remaining_p = (1 - p) / (n - 1)

        probabilities = np.full(n, remaining_p)
        probabilities[0] = p

        entropy = -np.sum(probabilities * np.log(probabilities))

        return entropy

    p_values = np.linspace(0.0001, 0.9999, 9999)

    entropy_values = [calculate_entropy(p) for p in p_values]

    plt.figure(figsize=(10, 6))
    plt.plot(p_values, entropy_values, 'b-', linewidth=2)

    max_entropy_idx = np.argmax(entropy_values)
    plt.scatter(p_values[max_entropy_idx], entropy_values[max_entropy_idx], color='red', s=50, zorder=5)
    plt.annotate(f'最大熵: p={p_values[max_entropy_idx]:.2f}, H={entropy_values[max_entropy_idx]:.4f}',
                xy=(p_values[max_entropy_idx], entropy_values[max_entropy_idx]),
                xytext=(p_values[max_entropy_idx]+0.1, entropy_values[max_entropy_idx]+0.2),
                arrowprops=dict(facecolor='black', shrink=0.05, width=1.5, headwidth=8))

    # print(entropy_values)
    points = [0.5, 1.0, 1.5]
    for point in points:
        for index, entropy in enumerate(entropy_values):
            if entropy <= point:
                print((p_values[index], entropy_values[index]))
                plt.scatter(p_values[index], entropy_values[index], color='red', s=50, zorder=5)
                break

    plt.title('概率值p与熵的关系图 (15W个候选值)')
    plt.xlabel('概率值p (第一个候选值的概率)')
    plt.ylabel('熵 (nats)')
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.xlim(-0.05, 1.05)
    plt.ylim(0, max(entropy_values) * 1.1)

    plt.tight_layout()
    plt.show()

    # (0.9695, 0.4999866540719247)
    # (0.9361, 0.9991510505356568)
    # (0.9012, 1.4999746052311926)

  • 示意图：

整体说明

  • 本文讨论如何将深度强化学习技术移植到 RLHF/RLAIF 和使用工具的 LLM 智能体中
    • 具体设计到的技术包括 \(n\) 步回报、TD(\(\lambda\))、不确定性、安全性以及其他
  • 本文收集了著名的深度强化学习技术（前 LLM 时代），并将其适配到现代 LLM 的训练/推理中
  • 以下方法并非唯一途径（可将其视为实用的起点），可以根据自己的技术栈进行优化

为何要将深度强化学习理念引入 LLM？

  • 生成过程就是一条轨迹：隐藏状态 \(h_t\) 和动作 \(y_t\) （Token 或工具调用）
  • 奖励可能是稀疏的（仅在序列末尾有一个分数）或 Dense 的（规则、自我批判、任务进度）
  • 这正是经典深度强化学习所针对的场景：长 horizon 信用分配、带噪声的目标以及安全约束

用于长文本的多步回报、TD(\(\lambda\)) 和 GAE

  • 设 \(h_t\) 为 Token \(y_{1:t-1}\) 之后的解码器状态
  • 训练一个价值头 \(V_\phi(h)\) 用于预测未来奖励
  • \(n\)步回报
    $$
    G_t^{(n)}=\sum_{i=0}^{n-1}\gamma^i r_{t+i}+\gamma^n V_\phi(h_{t+n})
    $$
  • TD 更新
    $$
    V_\phi(h_t)\leftarrow V_\phi(h_t)+\alpha\big(G_t^{(n)}-V_\phi(h_t)\big)
    $$
  • TD(\(\lambda\))/GAE（Actor-critic）
    $$
    \hat A_t=\sum_{l=0}^{\infty}(\gamma\lambda)^l\delta_{t+l}
    $$
    • 其中 \(\delta_t=r_t+\gamma V_\phi(h_{t+1})-V_\phi(h_t)\)
  • 带 KL 控制的策略梯度
    $$
    \nabla_\theta J \approx \mathbb{E}\left[\sum_t \hat A_t \nabla_\theta \log \pi_\theta(y_t|h_t)\right] -\beta\nabla_\theta \mathrm{KL}\big(\pi_\theta|\pi_{\text{ref} }\big)
    $$

应用场景

  • 仅带有 end-of-sequence scores 的长文本生成/摘要（Long-form generation/summarization）：
    • 通过 GAE 在 Token 间传播信用（propagate credit via GAE across tokens.）
  • 多轮对话助手（Multi-turn assistants）：
    • 将每一轮视为一个步骤；
    • 设置 \(\gamma<1\) 以减少不必要的轮次（理解：对于最后一步给奖励的场景，\(\gamma<1\) 相当于鼓励缩短决策轮次）
  • 工具/代码智能体（Tool/Code agents）：
    • 每次工具调用或单元测试结果作为一个步骤；
    • \(n\) 步回报可快速利用中间反馈
  • RAG/问答：
    • 将检索质量/格式检查作为 Dense 奖励，以解决稀疏信用问题（sparse-credit problems.）

Off-policy multi-step without explosions

  • 混合日志数据和新轨迹是常见做法（需结合修正进行多步学习）
  • V-trace（适用于分布式采样，具有稳定性） ：
    $$
    \rho_t=\min\left(\bar\rho,\frac{\pi_\theta(y_t|h_t)}{\mu(y_t|h_t)}\right),\quad c_t=\min\left(\bar c,\frac{\pi_\theta(y_t|h_t)}{\mu(y_t|h_t)}\right)
    $$
    • 使用 \(c_t\) 截断 eligibility 并构建修正后的优势
  • Retrace(\(\lambda\))
    $$
    c_t=\lambda \min\left(1,\rho_t\right)
    $$
  • Tree-Backup(\(\lambda\)) ：
    • 期望回溯（无重要性权重），方差更低，偏差较小
  • 典型截断值： \(\bar\rho\in[1,2]\) ， \(\bar c=1\)

附录：V-trace 介绍

  • 原始论文 IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures, ICML 2018, Google
    • IMPALA：Importance Weighted Actor-Learner Architectures
  • V-trace 是一种用于强化学习的 Off-policy 修正方法，主要用于解决分布式强化学习中 Actor 和 Learner 策略不一致的问题
  • 问题引入：
    • 在分布式强化学习框架 IMPALA 中，Actor 负责与环境交互生成经验轨迹，Learner 则根据这些轨迹来更新策略
    • 由于 Learner 的更新速度通常比 Actor 快，导致用于生成轨迹的策略滞后于 Learner 的当前策略，使得学习过程变成了 Off-policy 学习
    • V-trace 就是为了校正这种策略差异带来的不良影响而设计的
  • V-trace 核心思想 ：
    • V-trace 通过引入重要性采样（Importance Sampling）来修正 Off-policy 学习中的偏差
    • V-trace 对传统的策略梯度公式进行了改进，通过裁剪重要性系数来稳定方差，使得训练更加稳定
    • 具体来说，V-trace 在计算值函数更新时，使用了经过裁剪的重要性采样比率，限制了重要性系数的最大值，以防止因策略分布差异过大导致的估计偏差过大
  • V-trace 数学公式 ：
    • 在 V-trace 算法中， \(n\) 步下的目标价值函数可以表示为：
      $$ V_s = V(x_s) + \sum_{t=s}^{s+n-1} \gamma^{t-s}(\prod_{i=s}^{t-1} c_i )\delta_t V$$
      • \(\delta_t V\) 的表达式为：
        $$ \delta_t V = \rho_t(r_t + \gamma V(x_{t+1}) - V(x_t))$$
        • 其中：
          $$
          \rho_t = min(\bar{\rho},\frac{\pi(a_t|x_t)}{\mu(a_t|x_t)}) \\
          c_i = min(\bar{c},\frac{\pi(a_i|x_i)}{\mu(a_i|x_i)}) \\
          $$
      • \(\pi\) 是目标策略（Learner 的策略）
      • \(\mu\) 是行为策略（Actor 的策略）
      • \(\bar{\rho}\) 是重要性系数的裁剪阈值
      • \(c_i\) 是用于控制模型收敛速度的参数
  • V-trace 的作用 ：
    • V-trace 使得 IMPALA 能够在高吞吐量的情况下，通过结合解耦的 Acting 和 Learning 过程，实现相当稳定的学习
    • V-trace 有效地解决了分布式强化学习中因策略时间差导致的训练不稳定问题，提高了算法的效率和鲁棒性

应用场景

  • 带大量历史日志的 RLHF/RLAIF：减少行为-目标不匹配导致的偏差
  • Distributed/asynchronous sampling（IMPALA-style）：学习器和执行者不同步
    • 分布式/异步采样场景中
  • 离线+在线混合训练：安全复用旧数据，同时保持稳定性

Uncertainty and risk: optimize for reliability, not just averages

  • 奖励模型不确定性（集成/贝叶斯头） （Reward-model uncertainty (ensembles/Bayesian head)）
    • 设 \(\hat r\) 为均值， \(\sigma^2\) 为方差
    • 对高风险样本进行降权：
      $$
      \tilde A_t=\frac{\hat A_t}{1+\alpha \sigma_t}
      $$
  • 条件风险价值（CVaR，聚焦尾部） (tail-focus)
    $$
    \mathrm{CVaR}_\alpha(R)=\mathbb{E}[R \mid R \le q_\alpha]
    $$
    • 通过奖励的分位数回归（quantile regression）实现，并在更新时对低分位数值进行掩码
  • 偏好可靠性 (Preference reliability)
    • 通过 Bradley-Terry 置信度对成对偏好进行加权，以减少噪声标签导致的过拟合
  • 参数不确定性 (Parameter uncertainty)
    • 策略/价值上的 Laplace-LoRA 或小型集成可提供每个状态的方差，用于控制步长或触发“重新生成与重新评分”

附录：条件风险价值（CVaR）

  • 条件风险价值（Conditional Value at Risk，CVaR），常被称为 Expected Shortfall
  • CVaR 含义是：
    • 在给定置信水平 \(\alpha\)（通常是一个较小的值，如 0.05 或 0.1）的情况下，所有风险超过 \(\alpha\) 的情况对应损失的期望值
    • 简单来说，CVaR 关注的是 “最坏情况中的平均损失”
      • 例如，当 \(\alpha=0.05\) 时，CVaR 表示在所有可能发生的结果中，损失最大的 5% 的情况的平均损失
  • CVaR 常被用于风险敏感型决策，尤其适用于需要规避极端不利结果的场景

应用场景

  • 安全性/合规性要求高的领域（医疗、金融、教育）：减少罕见的灾难性失败
  • 噪声或主观的人类反馈：不确定性加权更新可稳定学习
  • 领域偏移/检索质量波动：检测分布外（OOD）情况并重新路由至重新评估

Safety as constraints, not just penalties

  • 设 \(c(h_t,y_t)\) 为安全成本（毒性、个人身份信息（PII）、事实风险）。通过拉格朗日方法进行优化：
    $$
    \max_\theta\ \mathbb{E}[R] - \lambda\big(\mathbb{E}[c]-\tau\big) \ -\ \beta\mathrm{KL}\left(\pi_\theta|\pi_{\text{ref} }\right)
    $$
  • 通过对偶上升更新 \(\lambda\)
  • 训练一个成本价值头 \(C_\psi(h)\)
  • 推理时添加屏蔽（分类器/正则表达式/规则）以过滤不安全 Token，或路由至拒绝模板
  • 训练时的约束与解码时的屏蔽结合使用效果最佳

应用场景

  • 企业/公共部门（Enterprise/public-sector）：严格的个人身份信息（PII）/合规控制
  • 开放式对话/内容（Open-ended chat/content）：减少毒性/偏见
  • 高事实性要求（High factuality requirements）：将幻觉信号视为成本

Robustness to distribution shift and prompt attacks

  • 分布鲁棒优化（distributionally robust optimization，DRO）：在训练提示分布 \(\hat P\) 周围的 \(f\)-divergence ball \(\mathcal{Q}\) 内最大化最坏情况下的奖励
    $$
    \max_\theta \ \min_{Q\in \mathcal{Q} } \ \mathbb{E}_{x\sim Q}\big[R(\pi_\theta; x)\big]
    $$
    • 问题：待补充理解
  • 实用方案（Practical recipe）：
    • 对抗性重加权和对抗性提示生成
    • 添加红队测试循环和领域随机化（检索噪声、工具延迟、系统提示）

附录：DRO

  • 在强化学习中，分布鲁棒优化（Distributionally Robust Optimization, DRO） 是一种旨在提高策略对环境分布不确定性鲁棒性的方法
  • DRO 结合了分布鲁棒优化的核心思想与强化学习的框架，解决了传统 RL 方法在环境分布未知（真实的状态转移概率 \( p(s’|s,a) \) 未知）、存在扰动或偏移时性能下降的问题
  • DRO 通过建模环境分布的不确定性集合，在最坏情况下优化策略性能，从而提高策略对分布扰动、估计误差或偏移的鲁棒性
  • DRO 是解决 RL 中“分布不确定性”问题的重要框架，尤其适用于非平稳环境、安全关键场景或模型误差较大的场景，但需在鲁棒性与计算效率、性能之间进行精细权衡
DRO 背景：强化学习中的分布不确定性
  • 强化学习的核心是在马尔可夫决策过程（MDP）中学习最优策略，而 MDP 的关键组件（如状态转移分布、奖励分布、初始状态分布）往往存在不确定性：
    • 状态转移分布 ：真实的状态转移概率 \( p(s’|s,a) \) 通常未知，只能通过有限样本估计，可能存在偏差；
    • 奖励分布 ：奖励函数可能受噪声影响，或随环境动态变化；
    • 分布偏移 ：实际部署时，环境分布可能与训练时不同（如领域自适应、非平稳环境）
  • 传统 RL 方法（如 Q-learning、PPO 等）通常假设可以通过采样准确估计真实分布，或直接使用经验分布，但这种假设在分布不确定时会导致策略不稳定、泛化能力差
DRO 的基本思想
  • DRO 的核心是 “在不确定性中求稳健” ：
    • 当无法获知真实分布时，按如下步骤进行：
      • 第一步：定义一个包含真实分布的 “不确定性集合”（uncertainty set） （即所有可能的候选分布构成的集合）
      • 第二步：在这个集合中 针对最坏情况（worst-case）优化目标函数
      • 最终得到的解对集合内的所有分布都能保持较好性能，从而提高鲁棒性
  • 简单来说：传统优化是 “针对真实分布求最优”，而 DRO 是“针对最坏可能的分布求最优”
DRO：RL 问题形式化
  • 在 RL 中，DRO 的目标是学习一个策略，使其在不确定性集合内的所有可能分布下都能最大化（或保证）累积奖励
  • 以 MDP 为例，其核心形式化为：
    $$
    \max_{\pi} \min_{p \in \mathcal{U}} V^{\pi}(p)
    $$
  • 其中：
    • \( \pi \) 是待优化的策略；
    • \( \mathcal{U} \) 是不确定性集合（包含真实分布的候选分布集合）；
    • \( V^{\pi}(p) \) 是策略 \( \pi \) 在分布 \( p \) 下的价值函数（累积奖励期望）
  • 目标是找到一个策略 \( \pi \)，使得在不确定性集合 \( \mathcal{U} \) 中最坏的分布 \( p \) 下，价值函数 \( V^{\pi}(p) \) 尽可能大
DRO 中不确定性集合如何构造？
  • 不确定性集合 \( \mathcal{U} \) 的定义是DRO的核心，它直接决定了鲁棒性的范围和程度。构造方式通常基于统计距离（衡量分布间的差异），常见的包括：
    统计距离 定义与应用场景 特点
    KL散度（KL-divergence） \( \mathcal{U} = \{ p \mid \text{KL}(p \parallel \hat{p}) \leq \epsilon \} \)，其中 \( \hat{p} \) 是经验分布，\( \epsilon \) 是不确定性预算。 适用于高维分布，计算相对简单，但不对称（\(KL(p||q) \neq KL(q||p)\)）
    Wasserstein距离 \( \mathcal{U} = \{ p \mid W_c(p, \hat{p}) \leq \epsilon \} \)，基于最优传输理论，衡量分布间的“运输成本”。 对异常值更稳健，适用于连续分布，但高维下计算复杂
    总变差距离 \( \mathcal{U} = \{ p \mid \text{TV}(p, \hat{p}) \leq \epsilon \} \)，衡量分布最大差异。 直观但严格，导致集合较小，解可能过于保守
  • 不确定性集合的大小由参数 \( \epsilon \) 控制：
    • \( \epsilon \) 越大，集合包含的分布越多（鲁棒性越强，但可能过于保守，牺牲性能）；
    • \( \epsilon \) 越小，集合越接近经验分布（性能可能更好，但鲁棒性下降）
DRO 在强化学习中的典型应用
  • Model-based RL中的鲁棒策略学习
    • 在 Model-based RL 中，若环境模型（如状态转移函数）存在误差，DRO 可通过构造模型分布的不确定性集合，优化最坏情况下的策略，避免模型误差导致的性能崩溃
  • Model-free RL中的分布偏移处理
    • Model-free 方法依赖采样数据，若采样分布与真实分布存在偏移（如探索不足、噪声干扰），DRO 可通过对采样分布的不确定性建模，使策略对偏移更稳健
  • 安全强化学习（Safe RL）
    • 在需要满足安全约束（如避免碰撞、能耗上限）的场景中，DRO 可确保策略在约束分布的最坏情况下仍不违反安全条件
DRO 相关的挑战与权衡
  • 计算复杂度 ：DRO 需求解 min-max 优化问题（先最小化最坏分布，再最大化策略性能），比传统 RL 的单目标优化更复杂，尤其在高维状态/动作空间中难以高效求解
  • 鲁棒性与保守性的权衡 ：不确定性集合过大可能导致策略过于保守（为了稳健牺牲了最优性能）；过小则可能无法覆盖真实分布，失去鲁棒性
  • 不确定性集合的合理性 ：如何基于有限数据构造“既包含真实分布，又不过大”的集合，是 DRO 的核心难点（通常依赖统计理论或领域知识）

应用场景

  • 公共 LLM 端点：抵御越狱/攻击提示
  • RAG 系统：证据质量和风格存在差异
  • 跨领域泛化：训练-部署偏移（training–serving shift）

Model-based flavor: value-guided decoding

  • 采样时使用价值头作为短 horizon 前瞻：
    $$
    \log \tilde \pi(y_t|h_t) = \log \pi_\theta(y_t|h_t) + \eta\big( V_\phi(h_{t+1}) - V_\phi(h_t) \big)
    $$
    • 公式理解：朝前面看一步，看看下一步 \(V_\phi(h_{t+1})\) 相对当前这步 \(V_\phi(h_{t})\) 带来的价值提升，提升越大的动作 \(y_{t}\)，则赋予更大的采样概率（直接将增益加到原始对数概率上）
    • 问题：是不是在于原始策略输出概率增加前或后，增加一个归一化步骤会更好？
  • 对于工具智能体，展开 \(n\) 步工具计划并引导至 \(V_\phi\)
  • 这可在不重新训练整个策略的情况下，推动生成过程向高下游价值方向发展

应用场景

  • 代码/测试驱动生成：偏向通过测试/完成子任务
  • 冗长推理/约束写作：解码时更严格地遵循目标
  • 低重训练预算：无需完整强化学习周期即可获得“轻量级规划”收益

Offline and conservative RL from logs

  • IQL/AWAC 风格的优势加权更新
    $$
    \pi_\theta \leftarrow \arg\max_\theta \ \mathbb{E}\left[\exp\left(\frac{A_\beta(s,a)}{\lambda}\right)\log \pi_\theta(a|s)\right]
    $$
  • CQL 风格的抑制 ：
    • 通过压低未见过的文本/动作上的 \(Q\) 值，抑制过度乐观的分布外（OOD）动作
    • 在偏好空间中，添加行为正则化器 \(\log \pi_\beta(y|h)\) 以保持策略接近日志行为
  • 通过双重鲁棒异策略估计器进行评估：
    • 建模奖励 \(\hat r\) ，并将重要性权重与控制变量结合使用

应用场景

  • 日志丰富，在线数据有限：在风险探索前充分利用历史数据的价值
  • 高风险领域：先进行保守改进，再逐步扩展
  • 新领域冷启动：初始时保持接近 \(\pi_\beta\)

Exploration and diversity (without going off the rails, 不偏离正轨)

  • 熵/温度：SAC 风格的熵奖励或受控采样温度
  • 内在动机：基于 \(h_t\) 的分歧/RND，以鼓励新颖的语义或新工具路径
    *
  • 多样性正则化器：去重；与提示的互信息

应用场景

  • 创意写作/广告/教育内容：在安全范围内实现风格/结构多样性
  • 工具链发现：找到可靠的新序列
  • 覆盖导向的评估：扩大提示集群覆盖范围

Hierarchy and skills: plan–act–verify

  • 慢规划器输出子目标（工具计划、大纲），快执行器实现子目标
  • 训练类选项策略，或通过模仿/离线强化学习预训练技能库，并通过高层控制器调用技能

应用场景

  • 多工具/多步骤工作流（检索 -> 规划 -> 执行 -> 验证）
  • 可分解的大型任务（数据 Pipelines、无人机调度（UAV scheduling）、城市分析(urban analytics)）
  • 跨任务/领域的技能复用

Common pitfalls（陷阱）

  • Only end-of-sequence reward + weak value -> unstable advantages. Densify or invest in \(V\).
    • 理解：仅序列最后奖励+奖励模型比较弱，会造成不稳定的优势，需 Dense 化奖励或优化 \(V\)
  • Off-policy drift without correction -> biased updates. Use V-trace/Retrace.
    • 理解：无修正的 Off-policy 偏移是有偏更新，需使用 V-trace/Retrace 修正
  • Single deterministic reward model -> brittle. Prefer ensembles/quantiles.
    • 单一确定性奖励模型是脆弱的（理解：不稳定，方差大），建议优先选择集成/分位数奖励
  • Safety only at decoding time -> model still learns unsafe regions. Train with constraints too.
    • 仅在解码时考虑安全意味着模型仍会学习不安全区域，建议训练时也需添加约束

Closing Remarks

  • 以上这些是将“传统技巧”切实转化为 LLM 实践的方法
  • 如果只选择一种技术组合，可从 PPO + GAE + V-trace + Lagrangian safety 开始，然后添加不确定性感知加权（uncertainty-aware weighting）和价值引导解码（value-guided decoding）
  • 作者很想了解反例和更好的方案，欢迎调整并分享在你的系统中有效的（或无效的）方法

注：本文包含 AI 辅助创作

Paper Summary

  • 整体总结：
    • MHA2MLA 包含贡献感知的部分 RoPE 移除和 SVD 驱动的低秩投影
      • 论文实现了 KV 缓存的近乎无损压缩
    • 本文创新：通过针对性的参数重用和数据高效的微调来实现大语言模型架构迁移的可行性
  • 问题提出：
    • 基于 MHA 及 MHA 的变体（GQA 等）的标准 LLM 在成本上存在显著劣势
  • DeepSeek 的解法：MLA
    • Multi-head Latent Attention（MLA，多头潜在注意力）是 DeepSeek 提出的一种创新架构
    • MLA 通过将键值（Key-Value, KV）缓存显著压缩为一个潜在向量，从而实现高效且经济的推理
  • 本文核心解决的现实问题及挑战：
    • 使训练良好的大语言模型（例如 Llama）能够快速适配 MLA 而无需从头预训练，既具有重要意义又充满挑战
  • 论文提出了首个数据高效的精调方法，用于从 MHA 迁移至 MLA（MHA2MLA），该方法包含两个关键组件：
    • 对于 部分旋转位置编码（partial-RoPE） ，论文从对注意力分数贡献较小的查询和键的维度中移除 RoPE
    • 对于 低秩近似（low-rank approximation） ，论文基于预训练的键和值参数引入联合奇异值分解（SVD）近似
  • MHA2MLA 仅需使用一小部分（3% 到 6%）数据即可恢复性能，显著降低了推理成本，同时可与 KV 缓存量化等压缩技术无缝集成
    • 举例：Llama2-7B 的 KV 缓存大小减少了 92.19%，而在 LongBench 上的性能仅下降 0.5%

Introduction and Discussion

  • LLM 的快速发展显著加速了通往通用人工智能（Artificial General Intelligence, AGI）的进程，模型能力随参数数量增加而呈现可预测的扩展 (2020)
    • 然而，这些收益伴随着高昂的代价：训练的计算需求不断攀升，推理吞吐量下降，导致巨大的能源消耗和碳排放 (2019)
  • 随着下游任务日益复杂，长上下文处理和高计算量的推理已成为大语言模型应用的核心 (2024)
    • 一个关键瓶颈在于 MHA (2017) 机制固有的键值（Key-Value, KV）缓存的内存占用，其随序列长度和模型大小线性增长
    • 为缓解此问题，研究者探索了诸如分组查询注意力（Grouped-Query Attention, GQA）(2023) 和多查询注意力（Multi-Query Attention, MQA）(2019) 等变体
    • 然而，这些方法不仅减少了 KV 缓存大小，也减少了注意力中的参数数量，导致性能下降
  • DeepSeek 引入了 MLA (2024)
    • MLA 是一种配备低秩键值联合压缩的注意力机制
    • 经验上，MLA 实现了优于 MHA 的性能，同时显著减少了推理期间的 KV 缓存，从而提升了推理效率
  • 一个关键但尚未探索的问题随之产生：原本为 MHA 良好训练的大语言模型能否适配 MLA 以进行推理？
    • MHA 与 MLA 之间固有的架构差异使得零样本迁移不切实际，而从头预训练的惊人成本使得这种转变在技术上具有挑战性且在现有研究中探索不足
  • 为填补这一空白，论文提出了首个精心设计的 MHA2MLA 框架，该框架最大化地复用了预训练 MHA 网络的参数 ，同时将 KV 缓存存储和推理过程与 MLA 的范式对齐（图 1）
  • 论文的框架具有两项关键的技术创新：
    • 部分旋转位置编码（partial rotary position embedding, partial RoPE）
    • 低秩近似（low-rank approximation）
  • MHA2MLA 的主要目标是实现数据高效的性能恢复，即使用最少的精调数据来恢复由架构变更引起的能力下降
  • MLA 的推理加速机制与 RoPE 之间固有的不兼容性 necessitates 架构上的折衷
  • DeepSeek 的解决方案是在有限维度中保留位置编码（PEs）同时压缩其他维度，这需要在 MHA 中策略性地移除 RoPE 维度（将其转换为 NoPE）以实现与 MLA 的对齐
  • 虽然更高的移除比率提升了压缩效率，但也加剧了性能下降，形成了效率与能力之间的权衡
  • 通过系统性地探索 RoPE 移除策略，论文发现基于贡献的维度选择（保留按注意力分数影响排序的前 k 个维度）能最优地平衡这些竞争目标
  • 尽管先前的研究已经探索了从头训练部分 RoPE 大语言模型 (2024; 2021)，但论文的工作开创了在大语言模型中进行全 RoPE 到部分 RoPE（Partial RoPE） 转换的数据高效精调方法
  • MLA 通过将键和值投影到低秩潜在表示空间（存储在 KV 缓存中）来减少内存占用
  • MHA2MLA 也可以对剥离了 RoPE 的值和键（NoPE 维度）应用低秩近似
  • 通过对对应于 NoPE 子空间的预训练参数矩阵 \(\boldsymbol{W}_{v}\) 和 \(\boldsymbol{W}_{k}\) 执行奇异值分解（Singular Value Decomposition, SVD），论文将这些组件压缩到一个潜在空间中，同时最大限度地保留原始模型学到的知识
  • 论文的主要贡献是：
    • 论文提出了 MHA2MLA，这是首个参数高效的精调框架，能够仅使用 \(3%\) 到 \(6%\) 的训练数据将预训练的基于 MHA 的大语言模型适配到 MLA 架构，而无需从头训练
    • 论文证明了 MHA2MLA 架构可以与 KV 缓存量化技术集成，以实现更经济的推理（最高减少 96.87%）
    • 论文在四种模型规模（从 135M 到 7B，涵盖 MHA 和 GQA）上进行了实验，并进行了详细的消融研究，为 MHA2MLA 提供了指导和见解

Preliminary

多头注意力机制 (Multi-Head Attention, MHA)

  • 给定一个输入序列 \(\{\boldsymbol{x}_{1},\ldots,\boldsymbol{x}_{l}\} \in \mathbb{R}^{l \times d}\)，标准的 MHA (2017) 将每个 token \(\boldsymbol{x}_{i}\) 投影为查询向量 \(\boldsymbol{q}_{i}^{(h)} = \boldsymbol{x}_{i}\boldsymbol{W}_{q}^{(h)}\)、键向量 \(\boldsymbol{k}_{i}^{(h)} = \boldsymbol{x}_{i}\boldsymbol{W}_{k}^{(h)}\) 和值向量 \(\boldsymbol{v}_{i}^{(h)} = \boldsymbol{x}_{i}\boldsymbol{W}_{v}^{(h)}\)，其中对于每个头 \(h \in \{1,\ldots,n_{h}\}\)，有 \(\boldsymbol{W}_{q}^{(h)}, \boldsymbol{W}_{k}^{(h)}, \boldsymbol{W}_{v}^{(h)} \in \mathbb{R}^{d \times d_{h} }\)。旋转位置编码（Rotary Positional Encoding, RoPE）(2021) 被应用于查询和键（例如，\(\boldsymbol{q}_{i,\text{rope} }^{(h)} = \text{RoPE}(\boldsymbol{q}_{i}^{(h)})\)），随后进行缩放点积注意力计算：
    $$
    \boldsymbol{o}_{i}^{(h)} = \text{Softmax}\left( \boldsymbol{q}_{i,\text{rope} }^{(h)} \boldsymbol{k}_{\leq i,\text{rope} }^{(h)\top} \right) \boldsymbol{v}_{\leq i}^{(h)}, \\
    \text{MHA}(\boldsymbol{x}_{i}) = \left[ \boldsymbol{o}_{i}^{(1)}, \ldots, \boldsymbol{o}_{i}^{(n_{h})} \right] \boldsymbol{W}_{o},
    $$
    • 其中 \(\boldsymbol{W}_{o} \in \mathbb{R}^{(n_{h}d_{h}) \times d}\)
    • \([\cdot,\cdot]\) 表示向量拼接
    • 注：为简化符号，此处忽略了 \(\frac{1}{\sqrt{d} }\) 缩放因子
  • 在自回归推理过程中，MHA 需要存储大小为 \(O(2ln_{h}d_{h})\) 的键值（KV）缓存 \(\{\boldsymbol{k}_{\text{rope} }^{(h)}, \boldsymbol{v}^{(h)}\}_{h=1}^{n_{h} }\)，该大小随序列长度 \(l\) 线性增长，造成了内存瓶颈

分组查询注意力 (Grouped-Query Attention, GQA)

  • GQA (2023) 通过在 \(n_{g}\) 个组（\(n_{g} \ll n_{h}\)）之间共享键/值来减少 KV 缓存。对于每个头 \(h\)，它映射到组 \(g = \lfloor h / n_{g} \rfloor\)：
    $$
    \boldsymbol{o}_{i}^{(h)} = \text{Softmax}\left( \boldsymbol{q}_{i,\text{rope} }^{(h)} \boldsymbol{k}_{\leq i,\text{rope} }^{(g)\top} \right) \boldsymbol{v}_{\leq i}^{(g)}, \\
    \text{GQA}(\boldsymbol{x}_{i}) = \begin{bmatrix} \boldsymbol{o}_{i}^{(1)}, \ldots, \boldsymbol{o}_{i}^{(n_{h})} \end{bmatrix} \boldsymbol{W}_{o}. \tag{2}
    $$

多头查询注意力 (Multi-Query Attention, MQA)

  • MQA (2019) 是 GQA 的一个特例，其中 \(n_{g} = 1\)，即所有头共享一个全局的键/值
  • 虽然 GQA 和 MQA 方法将 KV 缓存减少到 \(O(2ln_{g}d_{h})\)，但由于参数剪枝，它们会导致性能下降

多头潜在注意力 (Multi-Head Latent Attention, MLA)

  • MLA (DeepSeek-AI, 2024) 引入了一种混合架构，将位置编码（PE）与潜在 KV 压缩解耦
  • 对于每个头 \(h\)，输入 \(\boldsymbol{x}_{i}\) 被投影为两个互补的分量：
  • 位置感知分量 (Position-Aware Component) ：一部分维度保留 PE 以保持位置敏感性：
    $$
    \boldsymbol{q}_{i,\text{rope} }^{(h)}, \boldsymbol{k}_{i,\text{rope} } = \text{RoPE}\left( \boldsymbol{x}_{i}\boldsymbol{W}_{dq}\boldsymbol{W}_{qr}^{(h)}, \boldsymbol{x}_{i}\boldsymbol{W}_{kr} \right),
    $$
    • 其中 \(\boldsymbol{W}_{dq} \in \mathbb{R}^{d \times d_{q} }\)，\(\boldsymbol{W}_{qr}^{(h)} \in \mathbb{R}^{d_{q} \times d_{r} }\)，\(\boldsymbol{W}_{kr} \in \mathbb{R}^{d \times d_{r} }\) 将查询/键投影到保留 RoPE 的 \(d_{r}\) 维分量中
  • 位置无关分量 (Position-Agnostic Component) ：剩余的 \(d_{c}\) 个维度被移除 PE（即 NoPE），并将 \(\boldsymbol{k}_{i,\text{nope} }^{(h)}\) 和 \(\boldsymbol{v}_{i}^{(h)}\) 压缩成一个共享的潜在向量 \(\boldsymbol{c}_{i,kv}^{(h)}\)：
    $$
    \boldsymbol{q}_{i,\text{nope} }^{(h)} = \boldsymbol{x}_{i}\boldsymbol{W}_{dq}\boldsymbol{W}_{qc}^{(h)}, \\
    \boldsymbol{c}_{i,kv} = \boldsymbol{x}_{i}\boldsymbol{W}_{dkv}, \\
    \boldsymbol{k}_{i,\text{nope} }^{(h)}, \boldsymbol{v}_{i}^{(h)} = \boldsymbol{c}_{i,kv}\boldsymbol{W}_{uk}^{(h)}, \boldsymbol{c}_{i,kv}\boldsymbol{W}_{uv}^{(h)},
    $$
    • 其中 \(\boldsymbol{W}_{qc}^{(h)} \in \mathbb{R}^{d_{q} \times d_{c} }\)，\(\boldsymbol{W}_{dkv} \in \mathbb{R}^{d \times d_{kv} }\)，\(\boldsymbol{W}_{uk}^{(h)} \in \mathbb{R}^{d_{kv} \times d_{c} }\)，\(\boldsymbol{W}_{uv}^{(h)} \in \mathbb{R}^{d_{kv} \times d_{h} }\)
  • 注意 \(d_{r} + d_{c} = d_{h}\)。MLA 的注意力输出结合了两个分量：
    $$
    \boldsymbol{o}_{i}^{(h)} = \text{Softmax}\left( \boldsymbol{q}_{i,\text{rope} }^{(h)} \boldsymbol{k}_{\leq i,\text{rope} }^{(h)\top} + \boldsymbol{q}_{i,\text{nope} } \boldsymbol{k}_{\leq i,\text{nope} }^{(h)\top} \right) \cdot \boldsymbol{v}_{\leq i}^{(h)} \\
    \text{MLA}(\boldsymbol{x}_{i}) = \begin{bmatrix} \boldsymbol{o}_{i}^{(1)}, \ldots, \boldsymbol{o}_{i}^{(n_{h})} \end{bmatrix} \cdot \boldsymbol{W}_{o}. \tag{3}
    $$
  • 与 MHA 及其变体不同，MLA 存储潜在向量 \(\boldsymbol{c}_{kv}\) 和 \(\boldsymbol{k}_{i,\text{rope} }^{(h)}\)（\(\mathcal{O}(ld_{r} + ld_{kv})\)）而不是全秩的 \(\boldsymbol{k}_{i}\), \(\boldsymbol{v}_{i}\)（\(\mathcal{O}(2ln_{h}d_{h})\)），其中 \((d_{r} + d_{kv}) \ll 2n_{h}d_{h}\)
  • 为什么 MLA 需要分离 RoPE 和 NoPE？
    • MLA 在推理过程中对 NoPE 部分引入了矩阵合并技术，有效减少了内存使用
    • 对于点积操作 \(\boldsymbol{q}_{i,\text{nope} }^{(h)} \boldsymbol{k}_{j,\text{nope} }^{(h)\top}\)，可以应用以下恒等变换：
      $$
      \boldsymbol{q}_{i,\text{nope} } \boldsymbol{k}_{j,\text{nope} }^{\top} = (\boldsymbol{x}_{i}\boldsymbol{W}_{dq}\boldsymbol{W}_{qc}) (\boldsymbol{c}_{j,kv}\boldsymbol{W}_{uk})^{\top} = \boldsymbol{x}_{i} (\boldsymbol{W}_{dq}\boldsymbol{W}_{qc}\boldsymbol{W}_{uk}^{\top}) \boldsymbol{c}_{j,kv}^{\top}
      $$
      • 注：为简化符号，论文省略了上标 \({}^{(h)}\)。矩阵 \(\boldsymbol{W}_{uv}\) 和 \(\boldsymbol{W}_{o}\) 也可以合并，请参阅 DeepSeek-AI 等人 (2024) 的附录 C
      • 其中 \((\boldsymbol{W}_{dq}\boldsymbol{W}_{qc}\boldsymbol{W}_{uk}^{\top})\) 可以预先合并为单个矩阵，而 \(\boldsymbol{c}_{j,kv}\) 已经存储在 KV 缓存中。对于 RoPE 部分，RoPE() 函数将输入向量乘以旋转矩阵（例如，\(\text{RoPE}(\boldsymbol{q}_{i}) = \boldsymbol{q}_{i}\boldsymbol{R}_{i}\)，\(\boldsymbol{R}_{i}\) 的具体形式将在第 3.1 节介绍）
  • 因此，恒等变换变为：
    $$
    \boldsymbol{q}_{i,\text{rope} } \boldsymbol{k}_{j,\text{rope} }^{\top} = (\boldsymbol{x}_{i}\boldsymbol{W}_{dq}\boldsymbol{W}_{qr} \boldsymbol{R}_{i}) (\boldsymbol{x}_{j}\boldsymbol{W}_{kr} \boldsymbol{R}_{j})^{\top} = \boldsymbol{x}_{i} (\boldsymbol{W}_{dq}\boldsymbol{W}_{qr} \boldsymbol{R}_{j-i} \boldsymbol{W}_{kr}^{\top}) \boldsymbol{x}_{j}^{\top}
    $$
  • 由于 \((\boldsymbol{W}_{dq}\boldsymbol{W}_{qr} \boldsymbol{R}_{j-i} \boldsymbol{W}_{kr}^{\top})\) 与相对位置 \(j-i\) 相关，它不能被合并成一个固定矩阵。考虑到 LLM 中的相对距离可能非常长（例如 128K），RoPE 部分更适合使用原始形式进行计算

MHA2MLA

部分旋转位置编码（Partial-RoPE）

  • 为实现从标准 MHA 到 MLA 的迁移，论文提出了部分旋转位置编码微调（partial-RoPE finetuning）策略，该策略从目标比例维度中移除 RoPE，并将其转换为 NoPE
  • 关键的是，尽管先前的工作已经探索了从头开始训练具有部分 RoPE 的 LLM（实现了比全 RoPE 略好的困惑度 (2021; 2024)），但现有方法均未解决如何高效地将预训练的全 RoPE 模型（例如 Llama）适配到部分 RoPE，而无需昂贵的重新训练
  • 论文的工作通过系统评估部分 RoPE 的变体，以确定最数据高效的微调方案来恢复模型在适配后的性能，从而弥补了这一空白
  • RoPE 的数学形式化表示如下：
    • 对于维度为 \(d_h\) 的查询或键向量，RoPE 将向量划分为 \(\frac{d_h}{2}\) 个子空间，每个子空间包含两个连续维度（例如，第 \(k\) 个子空间包含维度 \(2k\) 和 \(2k+1\)）
    • 每个子空间以不同的旋转速度（频率）旋转，其中第 \(k\) 个子空间的旋转角为 \(\theta_k = b^{-2k/d_h}\)，\(b\) 是预定义的基数（例如，Llama 使用 \(b=10000\)）
    • 因此，对查询和键应用 RoPE 变为：
      $$
      \boldsymbol{q}_{i,rope} =\left[\boldsymbol{R}_{i}^{[2k,2k+1]}(\theta_{k})\boldsymbol{q}_{i }^{[2k,2k+1]}\right]_{0\leq k < \frac{d_{h} }{2} }, \\
      \boldsymbol{k}_{i,rope} =\left[\boldsymbol{R}_{i}^{[2k,2k+1]}(\theta_{k})\boldsymbol{k}_{i }^{[2k,2k+1]}\right]_{0\leq k < \frac{d_{h} }{2} }.
      $$
全 RoPE 到部分 RoPE 的策略（Full-RoPE to Partial-RoPE Strategies）
  • 给定保留的旋转子空间数量 \(r\)（\(r=\frac{d_r}{2} \ll\) 总子空间数 \(\frac{d_h}{2}\)），论文提出了四种策略（如图 2 所示）来选择哪些 \(r\) 个子空间保留 RoPE 编码
  • 高频保留（High-Frequency Preservation） 保留 \(r\) 个旋转最快（高频）的子空间：
    $$
    \mathcal{S}_{\text{high} }=\left\{k,|,0\leq k<r\right\}.
    $$
    • 这与 Barbero 等人 (2024) 提出的 p-RoPE 方法一致，他们探索了 \(r\) 占总子空间数 25%、50% 和 75% 的设置，并观察到在从头训练的 LLM 中比全 RoPE 略有优势
  • 低频保留（Low-Frequency Preservation） 保留 \(r\) 个旋转最慢（低频）的子空间：
    $$
    \mathcal{S}_{\text{low} }=\left\{k,\Big{|},\frac{d_{h} }{2}-r\leq k<\frac{d_{h} } {2}\right\}.
    $$
    • 选择该策略作为高频策略的对照实验
  • 均匀采样（Uniform Sampling） 以等间隔选择 \(r\) 个子空间：
    $$
    \mathcal{S}_{\text{uniform} }=\left\{\left.\left\lfloor k\frac{d_{h} }{2r}\right \rfloor,\right|0\leq k<r\right\}
    $$
    • 这通过几何间距平衡了高频和低频分量。实践中，\(2r\) 通常能整除 \(d_h\)。这与 GPT-Neo (2021) 中使用的部分 RoPE 类似
  • 头部感知 2-范数贡献度（Head-wise 2-norm Contribution） Barbero 等人 (2024) 首次提出了 2-范数贡献度来研究这些频率是否被使用以及它们如何发挥作用。该方法基于这样的观察：根据柯西-施瓦茨不等式，第 \(k\) 个频率子空间对注意力对数几率（logits）的影响受相应查询和键分量的 2-范数上界限制，即 \(\left|\left\langle\mathbf{q}_{i}^{[2k,2k+1]},\mathbf{k}_{j}^{[2k,2k+1]}\right \rangle\right|\leq\left|\mathbf{q}_{i}^{[2k,2k+1]}\right|\left|\mathbf{k}_{j }^{[2k,2k+1]}\right|\)。对于每个头 \(h\)，论文在长序列上计算 LLM 中每个子空间的平均 2-范数得分 \(^4\)。然后，论文提出按它们的 2-范数得分对所有子空间进行排序，并选择前 \(r\) 个：
    $$
    \mathcal{S}_{2\text{-norm} }=\operatorname*{top}\nolimits_{r} \left(\left|\mathsf{q}^{[2k,2k+1]}_{*}\right|\left|\mathsf{k}^{[2k,2k+1]}_{*}\right|\right).
    $$
    • 这种头部特定的选择自适应地保留了旋转关键的子空间
    • 图 3 可视化了 Llama2-7B 四个头部的 2-范数
  • 论文将在第 4.3 节分析这四种策略的有效性，并在附录 D 中对关键超参数 \(r\) 进行消融研究
  • 对于所有策略，未选择的子空间（\(k \notin \mathcal{S}\)）变为 NoPE 维度，从而实现与 MLA 潜在压缩的无缝集成

Low-rank Approximation, Low-rank Approximation

  • 在从全 RoPE 转换为部分 RoPE 后，论文得到了 MLA 中 KV 缓存的第一个分量，表示为：
    $$ \boldsymbol{k}_{i,rope}=\left[\boldsymbol{R}^{[2k,2k+1]}_{i}(\theta_{k})\boldsymbol{k}^{[2k,2k+1]}_{i} \right]_{k\in\mathcal{S} } $$
  • 论文的下一个目标是推导第二个分量
    $$ \boldsymbol{c}_{i,kv} \in \mathbb{R}^{d_{kv} } $$
    • 它作为 \(\boldsymbol{k}_{i,\text{nope} }\) 和 \(\boldsymbol{v}_{i}\) 的低秩表示
  • 给定 MHA 中的键 \(\boldsymbol{k}_{i}=\boldsymbol{x}_{i}\boldsymbol{W}_{k}\) 和值 \(\boldsymbol{v}_{i}=\boldsymbol{x}_{i}\boldsymbol{W}_{v}\)，论文首先提取 \(\boldsymbol{W}_{k}\) 中对应于 \(\boldsymbol{k}_{i,\text{nope} }\) 的子空间，即未包含在 \(\mathcal{S}\) 中的维度，得到：
    $$ \boldsymbol{k}_{i,\text{nope} }=\boldsymbol{x}_{i}\boldsymbol{W}_{k,\text{nope} } $$
  • 论文提出了两种基于奇异值分解（SVD）的策略（如图 4 所示）来在实现降秩的同时保留预训练知识：
  • 解耦 SVD（Decoupled SVD, SVD\({}_{\text{split} }\)） 分别将 \(\boldsymbol{W}_{k,\text{nope} }\) 和 \(\boldsymbol{W}_{v}\) 分解为截断 SVD，各分配 \(d_{kv}/2\) 个维度：
    $$
    \boldsymbol{W}_{k,\text{nope} }=\boldsymbol{U}_{k}\boldsymbol{\Sigma}_{k}\boldsymbol{V}^{\top}_{k}, \quad \boldsymbol{W}_{v}=\boldsymbol{U}_{v}\boldsymbol{\Sigma}_{v}\boldsymbol{V}^{\top}_{v},
    $$
    其中 \(\boldsymbol{U}_{k},\boldsymbol{U}_{v},\boldsymbol{V}_{k},\boldsymbol{V}_{v} \in \mathbb{R}^{d_{h} \times \frac{d_{kv} }{2} }\)，\(\boldsymbol{\Sigma}_{k},\boldsymbol{\Sigma}_{v} \in \mathbb{R}^{\frac{d_{kv} }{2} \times \frac{d_{kv} }{2} }\)。下投影矩阵 \(\boldsymbol{W}_{d\cdot}\) 和上投影矩阵 \(\boldsymbol{W}_{u\cdot}\) 变为：
    $$
    \boldsymbol{W}_{dk} =\boldsymbol{U}_{k}\boldsymbol{\Sigma}^{1/2}_{k}, \quad \boldsymbol{W}_{uk} =\boldsymbol{\Sigma}^{1/2}_{k}\boldsymbol{V}^{\top}_{k},
    $$
    $$
    \boldsymbol{W}_{dv} =\boldsymbol{U}_{v}\boldsymbol{\Sigma}^{1/2}_{v}, \quad \boldsymbol{W}_{uv} =\boldsymbol{\Sigma}^{1/2}_{v}\boldsymbol{V}^{\top}_{v}.
    $$
    低秩表示 \(\boldsymbol{c}_{i,kv}\) 可以使用 \(\boldsymbol{c}_{i,kv}=[\boldsymbol{x}_{i}\boldsymbol{W}_{dk}, \boldsymbol{x}_{i}\boldsymbol{W}_{dv}]\) 构建
  • 联合 SVD（Joint SVD, SVD\({}_{\text{joint} }\)） 为保留 \(\boldsymbol{K}_{\text{nope} }\) 和 \(\boldsymbol{V}\) 之间的相互作用，论文联合分解拼接后的矩阵：
    $$
    [\boldsymbol{W}_{k,\text{nope} }, \boldsymbol{W}_{v}] = \boldsymbol{U}_{kv}\boldsymbol{\Sigma}_{kv}\boldsymbol{V}^{\top}_{kv},
    $$
    其中 \(\boldsymbol{U}_{kv}, \boldsymbol{V}_{kv} \in \mathbb{R}^{d_{h} \times d_{kv} }\)，\(\boldsymbol{\Sigma}_{kv} \in \mathbb{R}^{d_{kv} \times d_{kv} }\)。潜在投影则为：
    $$
    \boldsymbol{W}_{dkv} = \boldsymbol{U}_{kv}\boldsymbol{\Sigma}^{1/2}_{kv},
    $$
    $$
    \boldsymbol{W}_{uk} = \boldsymbol{\Sigma}^{1/2}_{kv}\boldsymbol{V}_{kv}[:, :-d_{v}], \quad \boldsymbol{W}_{uv} = \boldsymbol{\Sigma}^{1/2}_{kv}\boldsymbol{V}_{kv}[:, d_{v}:].
    $$
    这联合优化了键和值的潜在空间，即 \(\boldsymbol{c}_{i,kv} = \boldsymbol{x}_{i}\boldsymbol{W}_{dkv}\)，保留了对自回归生成至关重要的跨参数依赖性 \(^5\)。第 4.3 节显示 SVD\({}_{\text{joint} }\) 优于 SVD\({}_{\text{split} }\) ，验证了联合分解能更好地保留预训练知识

Experiment

  • 论文在不同规模（SmoILM-135M/360M/1B7, Llama2-7B）且使用 MHA 或 GQA 预训练的 LLM 上评估了论文的方法
    • 选择 SmoILM 系列是因为其预训练数据和框架都是开源的，这可以最大程度地减少微调数据和过程上的差异
    • 选择 Llama2-7B 是因为它是广泛使用的开源大语言模型之一（但其预训练数据未开源，微调数据可能存在潜在差异）
  • 论文分别使用 MHA2MLA 和 GQA2MLA 来表示架构迁移
    • 两者均采用数据高效的全参数微调（data-efficient full-parameter fine-tuning）
    • 默认使用基于头部的 2-范数贡献度选择（\(\mathcal{S}_{2\text{-norm} }\)，\(r=\frac{d_{h} }{16}\)）作为部分旋转位置编码（Partial-RoPE）策略
    • 联合奇异值分解（SVD\({}_{\text{joint} }\)）作为低秩近似策略
  • 论文的实验旨在回答三个关键问题：
    • 1）MHA2MLA 如何最小化由架构转变引起的准确性下降？
    • 2）MHA2MLA 在 KV 缓存减少比率方面取得了什么成果？
    • 3）MHA2MLA 能否与 KV 缓存量化技术结合以实现复合收益？

Commonsense Reasoning Tasks

Main Results
  • 如表 1 所示，论文的方法在四种模型规模（135M 到 7B）和不同的 KV 缓存压缩比（通过潜在维度 \(d_{kv}\) 控制）下均实现了高效的架构迁移
  • 当比较论文的微调方法与原始大语言模型的性能时
    • 论文观察到四个基础模型的性能仅有微小变化：
      • 135M 模型下降 -0.25%
      • 360M 模型上升 +0.03%
      • 1B7 模型上升 +0.03%
      • 7B 模型上升 +0.37%
    • 这表明微调数据并未显著降低或提高原始模型的性能，为 MHA2MLA 框架提供了一个合适的实验环境
  • 随着 \(d_{kv}\) 减小（例如从 32 到 16 再到 8），KV 缓存减少量增加（即从 -68.75% 到 -81.25% 再到 -87.5%），但通过微调恢复性能损失变得更具挑战性
    • 图 5 显示了 135M（代表 GQA）和 7B（代表 MHA）在不同压缩比下的微调损失曲线
    • 随着压缩比增加，与基线的损失差异变大
    • 论文还观察到损失曲线的波动趋势几乎一致，这表明论文的架构迁移并未显著损害模型的内部知识
  • 更大的模型在迁移到 MLA 架构时经历的性能下降更小
    • 例如，压缩至 18.75% 时，性能下降分别为：
      • 135M 下降 2.41%
      • 360M 下降 2.69%
      • 1B7 下降 1.28%
      • 7B 下降 0.61%
    • 这揭示了 MHA2MLA 的潜在缩放定律 (potential scaling law of MHA2MLA)
  • 最后，从 135M 模型到 7B 模型，微调所需的 token 数量仅占预训练 token 的约 0.3% 到 0.6%，证明了论文方法的数据效率
  • 总体而言，无论是使用 GQA2MLA 还是 MHA2MLA，架构迁移都以极小的成本实现，从而带来高效且经济（economical）的推理
  • 表 1： 使用 MHA2MLA 或 GQA2MLA 的四个大语言模型的常识推理能力
    • 六个基准测试包括 MMLU (2021)、ARC 简单和挑战集 (ARC, 2018)、PIQA (2020)、Hellaswag (HS, 2019)、OpenBookQA (OBQA, 2018)、Winogrande (WG, 2021)

Long Context Tasks

Settings
  • 为评估模型的生成能力，论文采用 LongBench (2024) 作为生成性能的基准
  • 所有模型均使用贪心解码策略进行测试
  • 上下文窗口大小根据模型微调时使用的序列长度确定
  • 使用 HQQ ( 2023) 和 Quanto 以不同精度级别设置缓存，以评估原始模型的性能作为基线
  • 由于论文的方法与 KV 缓存量化兼容，论文还进行了额外实验来评估两种方法结合的效果
Main Results
  • 如表 2 所示，在 LongBench 上，与训练后量化方法相比，MHA2MLA 实现了具有竞争力或更优的效率-准确性曲线
    • 原生的 4 位量化在可比压缩比下仅带来适度的性能下降（-0.2% 到 -0.4%）
    • 原生的 2 位量化实现了 87.5% 的 KV 缓存减少，但出现了严重的性能崩溃（-6.2% 到 -9%）
    • 相比之下
      • MHA2MLA 在达到 87.5% 压缩（\(d_{kv}=16\)）时仅造成 3% 的准确性损失
      • 进一步与 4 位量化协同作用，实现了 92.19%/96.87% 的压缩（\(d_{kv}=64/16\)+Int4HQQ），同时将性能下降限制在 -0.5%/-3.2%，优于所有 2 位基线
    • 这突显了 MHA2MLA 的潜在空间设计与数值精度降低是正交的，从而能够实现复合效率增益 (compound efficiency gains) 而不会产生破坏性干扰
  • 表 2： Llama2-7B 和 MHA2MLA 在 LongBench 上的评估结果。粗体表示压缩比大于或等于 Int2 量化，同时性能也高于 Int2
  • 图 5： 不同 KV 缓存存储比率下的微调损失曲线（颜色从浅到深代表 12.5%, 18.75%, 31.25%, 和 100%）

Ablation Study

四种部分旋转位置编码策略：\(\mathcal{S}_{\text{high} }\), \(\mathcal{S}_{\text{low} }\), \(\mathcal{S}_{\text{uniform} }\), \(\mathcal{S}_{\text{2-norm} }\)
  • 表 3 展示了四种将完整旋转位置编码（full-RoPE）转换为部分旋转位置编码（partial-RoPE）的策略结果
    • 当将这四种策略与完整旋转位置编码进行比较时
      • 低频保留策略 \(\mathcal{S}_{\text{low} }\) 遭受了最大的性能损失（135M 减少 -6.49%，1B7 减少 -1.21%）
      • 高频保留策略 \(\mathcal{S}_{\text{high} }\) 的性能下降显著较小（135M 减少 -0.85%，1B7 减少 -0.76%）
      • 强调了高频子空间的重要性
    • \(\mathcal{S}_{\text{uniform} }\) 和 \(\mathcal{S}_{\text{2-norm} }\) 都产生了更好的性能，\(\mathcal{S}_{\text{uniform} }\) 保留了跨频率谱的子空间
    • \(\mathcal{S}_{\text{2-norm} }\) 则根据子空间对注意力分数的贡献来保留子空间
    • 论文选择 \(\mathcal{S}_{\text{2-norm} }\) 作为默认配置，因为被移除的子空间（即 NoPE）更适合（基于 SVD 的）低秩近似
两种基于 SVD 的低秩近似：\(\text{SVD}_\text{split}\), \(\text{SVD}_\text{joint}\)
  • 表 3 中每个组的最后两行比较了两种 SVD 方法的效果
    • 在两个大语言模型上，\(\text{SVD}_\text{joint}\) 方法 consistently 优于 \(\text{SVD}_\text{split}\)，在 135M 模型上平均性能提升 0.92%，在 1B7 模型上平均提升 0.74%
    • 这表明 \(\text{SVD}_\text{joint}\) 成为明确的默认选择

Efficient Attention Architectures

  • 标准的多头注意力机制（Multi-Head Attention, MHA）(2017) 在上下文长度上具有二次复杂度，这促使了众多效率创新
  • MHA 变体，如多头查询注意力（Multi-Query Attention, MQA）和分组查询注意力（Grouped-Query Attention, GQA）(2023)
    • 通过在不同头之间共享键/值来减少内存开销
    • 但这是以参数剪枝和性能下降为代价的
  • 其他并行的工作，如线性 Transformer (2019; 2020; 2021)、RWKV (2023) 和 Mamba (2023)
    • 用线性循环或状态空间模型替代了 softmax 注意力，但在自回归生成中难以匹配标准注意力的表达能力
  • 多头潜在注意力（Multi-Head Latent Attention, MLA）(2024) 通过将 KV 缓存压缩为低秩潜在向量而无需剪枝注意力参数，从而脱颖而出
  • 论文的工作将 MLA 与主流架构（MHA/GQA）连接起来，通过数据高效的微调实现无缝迁移
  • 许多线性注意力变体放弃了 softmax 查询-键交互（例如，通过核近似），但保留查询-键点积结构（即使是分解形式）的架构仍然与论文的 MHA2MLA 框架兼容

Economical Key-Value Cache（经济的键值缓存）

  • KV 缓存的内存占用已成为长上下文推理的关键瓶颈。最近的进展分为三类：
    • 创新架构方法 ，如 MLA (DeepSeek-2024)、MiniCache (2024a) 和 MLKV (2024)，跨层或头共享或压缩 KV 表示
      • 虽然有效，但跨层共享可能混淆不同的注意力模式，可能损害特定任务的性能
      • 只有 MLA 在 DeepSeek 的 LLM 中得到了成功验证
    • 量化技术 ，如 GPTQ (2022)、FlexGen (2023) 和 KIVI (2024b)
      • 以低比特格式（例如 2 比特）存储 KV 缓存，以精度损失为代价实现内存节省
    • 动态剪枝方法
      • A2SF (2024) 和 SnapKV (2024) 从 KV 缓存中剪枝“不太重要”的 Token
        • 但 Token 剪枝可能丢弃关键的长距离依赖
      • 头剪枝（例如 SliceGPT (2024)、Sheared (2024) 和 Simple Pruning (2024)）则不可逆地降低了模型容量
  • 论文的 MHA2MLA 方法实现了标准基于 Transformer 的大语言模型向更经济的 MLA 架构的迁移，并已证明其能够与 KV 量化技术集成以实现约 97% 的缓存节省
    • 它在理论上也与其他方法（如剪枝）兼容

注：本文包含 AI 辅助创作

Paper Summary

  • 评价：
    • 本文是 2020 年的文章，是非常重要的 Scaling 方向的作品，对模型规模的扩展意义重大
    • 核心思路：
      • 条件计算（Conditional Computation）：采用了稀疏门控混合专家（Sparsely-Gated MoE）架构
      • 自动分片（Automatic Sharding）：通过轻量级标注 API（replicate/split/shard）和 XLA 编译器扩展，实现张量的自动分布式划分
      • SPMD 编译优化：提出单程序多数据（SPMD）编译范式（替代传统多程序多数据（MPMD）模式），为所有设备生成单一程序，编译时间与设备数量无关
    • Google 出品，必属精品
  • GShard 是一个由一组轻量级标注接口（API）和 XLA 编译器扩展组成的模块
    • Gshard 提供了一种简洁的方式，只需对现有模型代码进行少量修改，就能实现多种并行计算模式
    • Gshard 解决了神经网络扩展过程中的 计算成本、编程便捷性以及在并行设备上的高效实现等问题
  • 作者使用 GShard 将带有稀疏门控混合专家层的多语言神经机器翻译 Transformer 模型扩展到了 600B 参数以上（应该是当时第一次有这么大的模型）
  • 论文的实验表明：
    • 这样一个巨型模型能够在 2048 个 TPU v3 加速器上高效训练 4 天，在 100 种语言到英语的翻译任务中，取得了远超现有技术水平的翻译质量

Introduction and Discussion

  • 在各类机器学习问题中，神经网络扩展都带来了显著的质量提升
    • 在计算机视觉领域，提升模型容量使得多种计算机视觉架构在图像分类和检测精度上有了更好的表现
    • 在自然语言处理领域，Transformer 模型的扩展在语言理解任务（2018, 2019）、跨语言下游迁移任务（2018, 2019）以及（大规模）多语言神经机器翻译任务（2019）中都实现了稳定的性能提升
  • 这种普遍趋势促使近期研究开始深入探究影响扩展成功的关键因素（2017, 2019, 2020），其中包括过往研究发现的训练数据量、模型大小和计算资源利用率等
  • 尽管研究发现最终模型质量与数据量、计算资源和模型大小呈幂律关系（2017, 2020），但更大模型带来的显著质量提升也伴随着各种实际挑战
  • 训练效率便是其中至关重要的挑战之一，论文将其定义为：为获得优于现有最佳系统的模型质量所消耗的计算资源和训练时间 ，而这一指标往往被忽视
  • 图1：
    • 多语言翻译质量（与双语基线模型相比的平均 \(\delta\) BLEU 值）随着 MoE 模型规模增长到 600B 参数而不断提升，而端到端训练成本（以 TPU v3 核心年为单位）仅呈亚线性增长
    • 将模型规模从 37.5B 参数扩大到 600B 参数（16 倍），计算成本仅从 6 个核心年增加到 22 个核心年（3.6倍）
    • 达到最佳翻译质量的 600B 参数模型，使用 2048 个 TPU v3 核心训练了 4 天，总成本为 22 个 TPU v3 核心年
    • 相比之下，训练所有 100 个双语基线模型则需要 29 个 TPU v3 核心年
    • 论文训练的最佳质量密集型单一 Transformer 模型（2.3B 参数）， \(\delta\) BLEU 值为 6.1，该模型使用 GPipe（2019）在 2048 个 TPU v3 核心上训练了 6 周，总成本达 235.5 个 TPU v3 核心年

Practical Challenges for Scaling（扩展面临的实际挑战）

  • 本节将列举在训练超大规模模型时面临的主要实际挑战，这类模型的规模远超单个加速器内存（如 GPU 或 TPU）的容量限制
  • 特定架构的模型并行支持问题 ：在 TensorFlow（2016）和 PyTorch（2017）等常用深度学习框架中，缺乏对高效模型并行算法的支持
    • 框架虽支持基于图分割的简单模型并行，但由于网络的顺序依赖性和基于梯度的优化方式，这种并行方式会导致设备利用率严重不足
    • 为了高效扩展现有模型，用户通常需要投入大量工程工作，例如将模型代码迁移到专门的框架中（2018, 2019）
  • 计算成本与模型规模的超线性扩展问题 ：通过增加网络深度或宽度来直接扩展模型规模（2018, 2019），通常会导致训练步长时间至少呈线性增长
    • 为解决这一问题，通常需要通过在多个设备间分割层权重和计算过程来实现模型并行，但这会带来网络通信开销和设备利用率不足的问题
    • 设备利用率不足源于神经网络底层计算任务分配不均衡以及存在顺序依赖关系
    • 这种计算成本与模型规模之间的超线性关系，无法通过简单增加设备数量来解决，使得训练超大规模模型变得不切实际
  • 巨型模型表示的基础设施扩展性问题 ：对于分布在数千个设备上的超大规模模型，简单的图表示可能会成为深度学习框架及其优化编译器的瓶颈
    • 例如，通过操作间分割（inter-op partitioning）增加 D 倍的层数，或通过操作内分割（intra-op partitioning）在 D 个设备上增加模型维度，都可能导致图节点数量大幅增加
    • 设备间的通信通道可能会进一步使图大小增加（例如，分割聚合或转置操作时）
    • 对于超大规模模型而言，图大小的这种增长会导致图构建和编译时间达到无法实现的程度
  • 实现分片策略的复杂工作 ：将模型高效地分割到多个设备上运行具有挑战性，因为这需要协调设备间的通信
    • 在图级分割方面，需要复杂的算法（2019, 2018）来减少因不同设备上分配的图分割部分之间存在顺序依赖而引入的开销
    • 在算子级并行方面，不同分割后的算子会有不同的通信模式，这取决于算子的语义，例如是否需要累积部分结果或重新排列数据分片
    • 根据论文的经验，由于 TensorFlow 等框架包含大量具有特殊语义的算子，在模型中手动处理这些问题需要耗费大量精力
    • 在所有情况下，实现模型分片对研究人员和工程师来说都是一项负担，因为改变模型架构就需要修改底层的设备通信逻辑，从而产生连锁反应

Design Principles for Efficient Training at Scale

  • 论文展示了如何通过构建一个拥有 600B 参数、带有稀疏门控混合专家层（Sparsely-Gated Mixture-of-Experts layers）的 sequence-to-sequence Transformer 模型，来克服上述挑战
    • 该模型的计算成本呈亚线性增长，编译时间为 \(O(1)\)
    • 论文在 2048 个 TPU v3 设备上，针对多语言机器翻译任务训练该模型 4 天，最终实现了单个非集成模型在 100 种语言到英语翻译任务中远超现有技术的翻译质量
    • 论文对不同规模的模型进行了实验，结果发现，随着模型规模增大，翻译质量不断提升，而训练总耗时（wall-time）相对于模型规模仅呈亚线性增长，如图1所示
  • 为构建如此庞大的模型，论文做出了以下关键设计选择
    • 亚线性扩展（Sub-linear Scaling） ：首先，模型架构的设计应确保计算和通信需求相对于模型容量呈亚线性增长
      • 条件计算（2015, 2019, 2020, 2020）通过在每个输入样本的基础上激活一个子网络，使论文能够兼顾训练和推理效率
      • 通过在基于循环神经网络（RNN）的机器翻译和语言模型中添加位置感知稀疏门控混合专家层（Position-wise Sparsely Gated Mixture-of-Experts, MoE）（2019），可以在实现最先进性能的同时，使计算成本呈亚线性增长
      • 因此，论文将在第2节中详细介绍如何用混合专家层扩展 Transformer 架构
    • 抽象的力量（The Power of Abstraction） ：其次，模型描述应与分片实现和优化相分离
      • 这种关注点分离使模型开发人员能够专注于网络架构，并灵活地改变分片策略，而底层系统则负责执行语义保持转换并实现高效的并行执行
      • 为此，论文提出了 GShard 模块，用户只需在模型中对少数关键张量标注分片策略即可
      • 该模块包含一组简单的标注接口，以及一个用于自动并行化的 XLA（2019）编译器扩展
      • 模型开发人员可以像在一个拥有超大内存和计算能力的单一设备上编写模型，编译器会根据标注信息和自身的启发式算法，自动为目标设备分割计算任务
      • 论文将在3.2节中提供更多标注示例
    • 可扩展编译器（Scalable Compilers） ：第三，包括计算表示和编译在内的系统基础设施，必须能够支持数千个设备的并行执行
      • 例如，图2 展示了在 4 个设备上分割点积运算（用颜色编码）的两种不同方式
      • 需要注意的是，图2a 中常用的 MPMD（多程序多数据，Multiple Program Multiple Data）方法在扩展性方面面临挑战
        • 因为图中的节点数量会随着设备数量的增加而线性增长
      • 相反，论文开发了一种用于 SPMD（单程序多数据，Single Program Multiple Data）转换的编译器技术
        • 该技术生成一个可在所有设备上运行的单一程序，使得编译时间与设备数量无关，保持恒定，如图2b所示
        • 论文将在3.3节中详细讨论 SPMD 框架
  • 论文其余部分的结构如下：
    • 第2节详细描述带有稀疏门控混合专家层的 Transformer 架构；
    • 第3节介绍论文开发的 GShard 模块；
    • 第4节展示混合专家模型在100个语言对的多语言机器翻译任务中的应用；
    • 第5节对实现的性能和内存使用情况进行评估；
    • 第6节讨论相关工作
  • 图2：
    • 在4个设备上对点积算子（ \([M, K] \times [K, N] = [M, N]\) ）进行 MPMD 分片与论文提出的 SPMD 分片的对比
    • 在该示例中，两个操作数均沿收缩维度 K 进行分片，每个设备计算本地结果后，通过 AllReduce 操作进行全局合并
    • MPMD 分片会为每个设备生成独立的算子，限制了其扩展性；
    • 而 SPMD 分片则生成一个可在所有设备上运行的程序
    • 需要注意的是，使用论文提出的 SPMD 分片时，编译时间与所使用的设备数量无关

Model

Sparse scaling of the Transformer architecture

  • Transformer架构（2017）已被广泛应用于自然语言处理领域，成为许多 sequence-to-sequence 任务（如机器翻译）的事实标准
  • Transformer 包含两个计算模块，即 Encoder 和 Decoder，两者均通过堆叠多个 Transformer 层实现
  • Transformer Encoder 层由两个连续的层组成，即自注意力层（self-attention layer）之后紧跟一个位置感知前馈层（position-wise feed-forward layer）
  • Decoder 则额外增加了第三个交叉注意力层（cross-attention layer），该层会关注 Encoder 的输出
  • 论文通过条件计算（conditional computation）对 Transformer 进行稀疏扩展，在 Encoder 和 Decoder 中，每隔一个前馈层就用一个位置感知混合专家层（Position-wise Mixture of Experts, MoE）（2019）替代，该混合专家层采用了一种改进的 top-2 门控机制（图3）
  • 论文通过改变 Transformer 层的数量和每个混合专家层中专家的数量来调整模型容量
  • 每个训练样本由一对子词 Token（subword tokens）序列组成
  • 在训练和推理过程中，每个 Token 都会激活混合专家 Transformer 的一个子网络
  • 子网络的大小大致与每个混合专家层中的专家数量无关，这使得计算成本能够像前一节所述的那样呈亚线性增长
  • 3.1节将进一步分析计算复杂度，5节将分析训练性能

Position-wise Mixture-of-Experts Layer

  • 论文模型中使用的 MoE 层基于文献（2019）的设计，并在稀疏门控函数和辅助损失函数方面进行了改进
  • Transformer 的混合专家层由 E 个前馈网络（FFN）组成，其计算过程如下：

$$ g_{s,e} = \text{GATE}(x_s) \quad \tag{1} $$

$$ \text{FFN}_e(x_s) = w o_e \cdot \text{ReLU}(w i_e \cdot x_s) \quad \tag{2} $$

$$ y_s = \sum_{e=1}^{E} g_{s,e} \cdot \text{FFN}_e(x_s) $$

  • 其中：
    • \(x_s\) 是混合专家层的输入 Token
    • \(w i_e\) 和 \(w o_e\) 分别是前馈层（即一个专家）的输入投影矩阵和输出投影矩阵
    • 向量 \(g_{s,e}\) 由门控网络计算得到
    • \(g_{s,e}\) 中每个专家对应一个非负值，其中大部分值为零，这意味着该 Token 不会被分配给对应的专家
      • 每个 Token 仅被分配给极少数专家，论文设定每个 Token 最多被分配给两个专家
    • \(g_{s,e}\) 中对应的值非零，表示该专家对最终网络输出的贡献程度
    • 每个专家 \(\text{FFN}_e\) 都采用含 ReLU 激活函数（2010）的两层全连接网络对 \(x_s\) 进行处理
    • 混合专家层的输出 \(y_s\) 是所有被选中专家输出的加权平均值
  • 门控函数 \(\text{GATE}(\cdot)\) 是混合专家层的核心，它采用 softmax 激活函数来表示每个专家在处理输入 Token 时的权重，即表示某个专家处理该输入 Token 的适合程度
  • 此外，门控函数必须满足以下两个目标：
    • 负载均衡（Balanced load） ：对于给定的 Token ，混合专家层最好能稀疏地激活专家
      • 一种简单的方案是根据 softmax 概率分布选择前k个专家，但研究表明这种方法会导致训练中的负载不均衡问题（2019）：训练过程中看到的大多数 Token 会被分配给少数几个专家，导致这几个（繁忙的）专家的输入缓存变得极大，而其他专家则处于未充分训练的状态，从而减慢训练速度
        • 问题：分配给少数专家也不一定会导致训练速度变慢吧，除非有 EP （专家并行）？
      • 同时，许多其他专家根本无法得到充分训练
      • 因此，门控函数需要设计得更合理，以实现所有专家间处理负载的更均匀分配
    • 大规模下的效率（Efficiency at scale） ：如果门控函数采用顺序执行方式，实现负载均衡相对容易
      • 但对于输入批次中所有 N 个 Token 和 E 个专家 ，仅门控函数的计算成本就至少为 \(O(N \cdot E)\)
      • 在论文的研究中，N 达到数百万量级，E达到数千量级，若门控函数采用顺序实现，会导致大部分计算资源在大部分时间处于空闲状态
      • 因此，论文需要一种高效的并行门控函数实现方式，以充分利用多个设备的计算能力
  • 为满足上述要求，论文在门控函数 \(\text{GATE}(\cdot)\) 中设计了以下机制（细节如算法1所示）：
    • 专家容量限制（Expert capacity） ：为确保负载均衡，论文强制每个专家处理的 Token 数量不超过某个统一的阈值，论文将该阈值定义为专家容量
      • 假设一个训练批次中的 Token 总数为 N，每个 Token 最多被分配给两个专家，那么专家容量设置为 \(\frac{2N}{E}\)
      • 门控函数 \(\text{GATE}(\cdot)\) 会为每个专家维护一个计数器 \(c_e\) ，用于记录分配给该专家的 Token 数量
      • 当一个 Token 选中的两个专家都已超过其容量限制时，该 Token 被视为溢出 Token ，此时 \(g_{s,e}\) 退化为零向量
      • 这类 Token 的表示会通过残差连接传递到下一层
        • 问题：这里是说专家容量限制也传递到下一层？
    • 本地组分配（Local group dispatching） ：
      • 门控函数 \(\text{GATE}(\cdot)\) 将训练批次中的所有 Token 均匀划分为 G 个组，即每个组包含 \(S = \frac{N}{G}\) 个 Token ，所有组独立并行处理
      • 每个组被分配到每个专家的部分容量为 \(\frac{2N}{G \cdot E}\)
      • 每个组会确保分配给每个专家的 Token 数量不超过该部分容量
      • 通过这种方式，既能保证专家容量限制得到遵守，又能实现整体负载均衡
    • 辅助损失（Auxiliary loss） ：门控函数不应总是选择相同的几个专家，否则会导致仅少数专家出现容量溢出，而其余专家利用率不足（这一点至关重要）
      • 借鉴文献（2019）的方法，论文定义了一个辅助损失项来强制执行这一约束
      • 该辅助损失项以一个常数系数 k 添加到模型的总损失函数中
      • 算法1第13行中辅助损失项的具体形式基于以下考虑：
        • \(\frac{c_e}{S}\) 表示分配给每个专家的输入 Token 比例，作者希望最小化 \(\frac{c_e}{S}\) 的均方值
        • 但由于 \(c_e\) 是通过 top-2 操作得到的，不具有可微性，因此论文使用每个专家的平均门控值作为可微近似，并用 \(m_e \cdot \frac{c_e}{S}\) 替代 \((\frac{c_e}{S})^2\) ，这样就可以通过梯度下降法对其进行优化
    • 随机路由（Random routing） ：
      • 直观上，由于 \(y_s\) 是所选专家输出的加权平均值，如果第二个专家的权重非常小，我们可以直接忽略该专家 ，以节省整体专家容量
      • 因此，除了遵守专家容量约束外，门控函数 \(\text{GATE}(\cdot)\) 还会以与第二个最佳专家权重 \(g_2\) 成比例的概率，将 Token 分配给该专家

Highly Parallel Implementation using GShard（高度并行化）

  • 本节将介绍第2节中模型在张量处理单元（TPU）设备集群上的高效运行实现方案
  • 实现的第一步是将模型用线性代数运算表示，因为论文的软件栈（TensorFlow（2016））和硬件平台（TPU）在这类运算上经过了高度定制和优化
    • 与原始 Transformer 模型类似，将模型的大部分部分用线性代数运算编写非常容易
    • 但由于混合专家层（MoE Layer），尤其是算法1 中提出的 \(\text{GATE}(\cdot)\) 函数具有顺序执行特性，要用线性代数运算表示该层则需要额外付出努力，论文将在3.1节详细介绍相关细节
  • 接下来，论文对线性代数计算进行标注，以体现并行性
  • 通过3.2节中的分片接口（sharding APIs），可以对计算中的每个张量进行标注，指定其在设备集群中是采用复制（replication）还是分布（distribution）方式存储
  • 分片标注能够实现模型描述与高效并行实现的关注点分离，让用户可以灵活地表达各种并行化策略；例如：
    • （1）注意力层（attention layer）通过沿批次维度（batch dimension）分割并将权重复制到所有设备上来实现并行化；
    • （2）由于混合专家层中专家（expert）的规模极大，无法在所有设备上进行复制，因此唯一可行的策略是将专家分片存储到多个设备中
    • 此外，整个模型会在这两种模式（1）-（2）之间切换
    • 通过标注，模型开发人员无需关注系统优化工作，也不必将并行实现细节和底层细节融入模型代码中
  • 最后，编译器基础设施会接收（部分）标注后的线性代数计算，并生成可在数千个设备上高效运行的并行程序
  • 如3.3节所述，编译器会应用单程序多数据（SPMD，Single Program Multiple Data）分片转换来表示每个设备的计算任务，插入必要的跨设备通信操作，处理非规则模式（如非均匀分片），最终生成一个可在所有设备上启动并执行并行计算的单一程序

Positions-wise Mixture-of-Expert Layer Expressed in Linear Algebra（线性代数）

  • 论文的模型实现（算法2）将整个加速器集群视为一个单一设备，并用少量与集群具体配置无关的张量运算来表示其核心数学算法

    • 爱因斯坦求和符号（Einstein summation notation）（1923）（即tf.einsum）是一种能够简洁表示模型的强大工具，论文在实现中大量使用了该符号
    • softmax 门控计算可以通过一个爱因斯坦求和运算（einsum）后紧跟 softmax 函数来轻松表示；
    • 将输入分配给选定专家的操作，可以通过分配掩码（dispatching mask）与输入之间的单次爱因斯坦求和运算来表示；
    • 所有 \(\text{FFN}_e\) 的权重被组合成单个三维张量 \(wi\) 和 \(wo\)
    • \(\text{FFN}_1 \dots \text{FFN}_E\) 的计算则通过三个算子（两个爱因斯坦求和运算和一个relu运算）来表示；
    • 最后，对所有专家输出进行加权平均以得到最终输出的操作，也可以通过另一个爱因斯坦求和运算来表示

  • 算法2中的 Top2Gating 计算了算法1中所有组本地（group-local）门控决策的并集

    • combine_weights是一个四维张量，形状为[G, S, E, C]
    • 其中，combine_weights[g, s, e, c] 非零时，表示组 g 中的输入 Token s被发送到专家 e 的输入缓存中，且位于缓存位置 c
    • 对于特定的 g 和 s，combine_weight 切片中最多包含两个非零值
    • 通过将所有非零值设为 1，可由 combine_weights 生成二进制分配掩码（binary dispatch_mask）

  • 论文需要合理选择组数（G）和专家数量（E），以确保该算法能够在包含D个设备的集群上实现扩展

    • 有必要分析在一个包含 N 个 Token 的训练批次中，该算法在一个训练步骤内的整体计算复杂度（浮点运算总数）

  • 算法2 ：位置感知混合专家层的前向传播（Forward pass of the Positions-wise MoE layer），下划线字母（如G和E）表示张量将沿该维度进行分片

    1
    2
    3
    4
    5
    6
    7
    1: gates = softmax(einsum("GSM, ME -> GSE", inputs, wg))
    2: combine_weights, dispatch_mask = Top2Gating(gates)
    3: dispatched_expert_inputs = einsum("GSEC, GSM -> EGCM", dispatch_mask, reshaped_inputs)
    4: h = einsum("EGCM, EMH -> EGCH", dispatched_expert_inputs, wi)
    5: h = relu(h)
    6: expert_outputs = einsum("EGCH, EHM -> GECM", h, wo)
    7: outputs = einsum("GSEC, GECM -> GSM", combine_weights, expert_outputs)
    • 在分析算法2的计算复杂度随设备数量D的扩展情况时，论文做出以下假设：
      • a）每个设备上的 Token 数量 \(\frac{N}{D}=O(1)\) (在实际应用中，为避免设备内存溢出，这一假设通常是必要的)，且保持恒定；
      • b） \(G=O(D)\) 、 \(S=O(1)\) ；
      • c） \(H=O(1)\) ；
      • d） \(E=O(D)\) ；
      • e） \(C=O(\frac{2S}{E})=O(\frac{1}{D})\) ，且 \(D< S\) ，D 为正整数

  • 算法2中的浮点运算总数计算如下：
    $$
    \begin{align}
    \text{FLOPS}_{\text{Softmax} } + \text{FLOPS}_{\text{Top2Gating} } + \text{FLOPS}_{\text{Dispatch/Combine} } + \text{FLOPS}_{\text{FFN} } &= \\
    O(G S M E) + O(G S E) + O(G S M E C) + O(G E C H) &= \\
    O(D \cdot 1 \cdot 1 \cdot D) + O(D \cdot 1 \cdot D) + O(D \cdot 1 \cdot 1 \cdot D \cdot \frac{1}{D}) + O(D \cdot D \cdot \frac{1}{D} \cdot 1) &= \\
    O(D^2) + O(D^2) + O(D) + O(D)
    \end{align    }
    $$

    • 因此，每个设备的浮点运算量为 \(O(D)\)
    • 每个设备上 softmax 的计算复杂度 \(\text{FLOPS}_{\text{Softmax} } / D = O(D)\) 与设备数量呈线性关系，但实际上，由于 \(D< S\) ，该复杂度会被其他项主导，因此可将整体复杂度视为 \(O(1)\) ，满足亚线性扩展的设计要求
    • 5节通过实验验证了这一分析结果

  • 除计算成本外，跨设备通信成本并非恒定，但如5节所述，当设备数量增加时，通信成本仅以 \(O(\sqrt{D})\) 的温和速率增长

GShard Annotation API for Parallel Execution

  • 由于算法1中张量的规模和计算需求极大，论文必须在多个设备上对该算法进行并行化处理
  • 算法2中带下划线的字母展示了对每个张量进行分片的一种直接方案
  • GShard 中的分片接口允许论文对程序中的张量进行标注，选择性地指定其分片方式
  • 这些信息会传递给编译器，以便编译器自动应用转换以实现并行执行
  • 在论文的研究中，使用了 TensorFlow/Lingvo（2019）中的以下接口
    • replicate(tensor)：对张量进行标注，使其在各个分片（partition）中复制，并返回标注后的张量。该接口通常用于模型中的非混合专家层（non-MoE layers），以实现权重复制
    • split(tensor, split_dimension, num_partitions)：对张量进行标注，使其沿split_dimension维度进行分片，并返回标注后的张量。第i个分片会被放置在第i个设备上，且num_partitions（分片数量）不得超过系统中的设备数量
    • shard(tensor, device_assignment)：是split()接口的泛化形式，支持对多个维度进行分片，并指定每个分片的放置位置。附录A.3对该接口进行了更详细的描述
  • 需要注意的是，调用split或shard接口仅会添加标注，不会改变用户程序中张量的逻辑形状。用户仍可使用完整形状的张量进行操作，无需担心非均匀分片等问题
  • GShard 具有良好的通用性，其简单接口可同样应用于所有维度
  • 根据具体应用场景，分片维度可以包括批次维度（数据并行）、特征维度、专家维度，甚至图像模型中的空间维度
  • 此外，由于分片标注是基于每个张量单独进行的，模型的不同部分可以采用不同的分片方式
  • 这种灵活性使论文能够对巨型混合专家层权重进行分片，并在混合专家层和非混合专家层之间切换分片模式，同时也支持论文未涉及的其他应用场景，例如对大型图像进行空间分片（2019）（附录A.4）
  • 通过上述分片接口，可以将算法2中所示的分片策略表示如下：
    • 输入张量沿第一个维度（组维度G）进行分片，门控权重张量（wg）采用复制方式存储
    • 计算得到分配后的专家输入（dispatched expert inputs）后，应用split接口将分片维度从组维度（G）切换为专家维度（E）。其中，D为设备数量 
      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      1: # 沿组维度（G）对输入进行分片
      2: inputs = split(inputs, 0, D)
      3: # 复制门控权重
      4: wg = replicate(wg)
      5: gates = softmax(einsum("GSM, ME -> GSE", inputs, wg))
      6: combine_weights, dispatch_mask = Top2Gating(gating_logits)
      7: dispatched_expert_inputs = einsum("GSEC, GSM -> EGCM", dispatch_mask, reshaped_inputs)
      8: # 沿专家维度（E）对分配后的输入进行分片
      9: dispatched_expert_inputs = split(dispatched_expert_inputs, 0, D)
      10: h = einsum("EGCM, EMH -> EGCH", dispatched_expert_inputs, wi)
Per-tensor sharding assignment
  • 如上述示例所示，用户无需对程序中的每个张量都进行标注
  • 通常只需对模型中的少数重要算子（如爱因斯坦求和算子）进行标注，编译器会通过自身的启发式算法推断出其余张量的分片方式
    • 注：由于反向传播计算通常由前端框架自动生成，用户无法访问这些张量，因此编译器推断缺失的分片信息也非常重要
    • 例如，由于输入张量沿组维度（G）进行分片，而权重张量采用复制方式存储，编译器会选择沿相同维度（组维度G）对爱因斯坦求和运算的输出进行分片（第5行）
    • 类似地，由于输入分配的爱因斯坦求和运算（第7行）的两个输入均沿组维度（G）进行分片，因此输出的分片方式会被推断为沿组维度（G）进行分片，之后论文通过添加split标注将输出的分片维度切换为专家维度（E）
    • 上述示例中的某些标注（如replicate(wg)）也可由编译器自动确定，但建议对计算的初始输入张量和最终输出张量进行标注
  • 目前，编译器采用迭代数据流分析（iterative data-flow analysis），从用户标注的算子开始，将分片信息传播到其相邻算子（操作数和使用者）
    • 该分析通过对齐相邻算子的分片决策，尽量减少重分片（resharding）的需求
    • 虽然还可以采用整数规划或机器学习等其他方法，但改进自动分片分配并非论文的重点，论文将其留作未来的研究工作
Mixing manual and automatic sharding（手动分片与自动分片结合）
  • 在常见情况下，通过分片标注实现自动分片通常已足够，但GShard也支持灵活地将手动分片算子与自动分片算子结合使用
  • 这让用户能够更好地控制算子的分片方式，例如当用户掌握算子语义之外的运行时知识时
  • 例如，XLA 和 TensorFlow 的 Gather 算子定义均未包含输入中不同范围的索引边界信息，但用户可能知道某个特定的 Gather 算子仅在每个分片内对数据进行重排
  • 在这种情况下，用户只需缩小维度大小并执行本地 Gather 操作，即可轻松对该算子进行分片；
  • 否则，编译器需要对索引范围采取保守处理，并添加不必要的通信开销
  • 例如，算法2中使用独热矩阵（one-hot matrix）分配输入的爱因斯坦求和算子（第3行），也可以通过手动分片的 Gather 算子来实现，而模型的其余部分仍采用自动分片方式
  • 以下伪代码展示了该应用场景 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    1: # 输入形状为[G, S, M]，split()不改变逻辑形状
    2: input = split(input, 0, num_devices)
    3: # s_indices形状为[E, G, C, 1]，值为输入中S维度的索引
    4: s_indices = split(s_indices, 1, num_devices)
    5: 
    6: # 开始手动分片
    7: # partitioned_input形状为[G/num_devices, S, M]
    8: partitioned_input = auto_to_manual_spmd_partition(input)
    9: # partitioned_s_indices形状为[E, G/num_devices, C, 1]
    10: partitioned_s_indices = auto_to_manual_spmd_partition(s_indices)
    11: # 在partitioned_input中拼接G维度的索引：在G维度上生成Iota张量
    12: partitioned_gs_indices = concat(
    13:     iota([E, G/num_devices, C, 1], 1), 
    14:     partitioned_s_indices, 
    15:     3
    16: )
    17: # partitioned_data形状为[E, G/num_devices, C, M]
    18: partitioned_data = gather(partitioned_input, partitioned_gs_indices)
    19: 
    20: # 切换回自动分片
    21: # data形状为[E, G, C, M]
    22: data = manual_to_auto_spmd_partition(partitioned_data)
    23: ...

The XLA SPMD Partitioner for GShard

  • 本节将介绍基于分片标注自动对计算图进行分片的编译器基础设施
  • 分片标注告知编译器每个张量应如何在设备间分布
  • SPMD（单程序多数据，Single Program Multiple Data）分片器（简称“分片器”）是编译器的一个组件，它将计算图转换为可在所有设备上并行执行的单一程序
    • 注：另一种方案是MPMD（多程序多数据，Multiple Program Multiple Data），如图2所示，该方案的扩展性较差
  • 这使得编译时间几乎与分片数量无关，从而支持论文将模型扩展到数千个分片⁴
  • 论文在 XLA 编译器（2019）中实现了该分片器
  • TensorFlow、JAX、PyTorch 和 Julia 等多个前端框架已具备将其图表示转换为 XLA HLO 图的降级逻辑（lowering logic）
  • 与 TensorFlow 等流行前端框架相比，XLA 的算子数量要少得多，这在不影响通用性的前提下降低了实现分片器的负担，因为现有前端的降级过程已实现了强大的表达能力
    • 尽管论文在 XLA 中开发了该基础设施，但论文介绍的技术也可应用于其他机器学习框架的中间表示（如 ONNX（2019）、TVM Relay（2018）、Glow IR（2018））
  • XLA 将计算表示为数据流图（dataflow graph），其中节点表示算子，边表示在算子之间流动的张量
  • 分片器的核心是对每个算子进行处理，根据输入和输出指定的分片方式，将全尺寸算子转换为分片尺寸算子
  • 对计算进行分片时，会引入各种跨设备数据传输模式
  • 为了在大规模场景下最大化性能，定义一组核心通信原语（communication primitive）并针对目标平台对其进行优化至关重要
Communication Primitives，通信原语
  • 由于分片器强制所有设备运行相同的程序，通信模式也具有规律性，XLA 定义了一组集合算子（collective operator）来执行类似 MPI（2009）的通信操作
  • 下面列出了论文在 SPMD 分片器中使用的常见通信原语
    • CollectivePermute ：该算子指定一组源-目标设备对（source-destination pairs），源设备的输入数据会发送到对应的目标设备
      • 它主要用于两个场景：改变分片张量在各分片间的设备顺序，以及本节后续讨论的 Halo Exchange
    • AllGather ：该算子按照指定顺序将所有参与设备的张量连接起来，用于将分片张量转换为复制张量
    • AllReduce ：该算子对所有参与设备的输入执行按元素归约（如求和）操作，用于合并来自不同分片的部分归约中间张量
      • 在 TPU 设备网络中，当分片数量增加时，AllReduce 的成本保持恒定（5.2节）
      • 它也是一种在其他类型网络拓扑中具有高效实现的常用原语（2019）
    • AllToAll ：该算子从逻辑上沿某个维度对每个参与设备的输入进行分割，然后将每个数据片段发送到不同的参与设备
      • 每个设备在接收来自其他设备的数据片段后，会将这些片段连接起来生成结果
      • AllToAll 用于将分片张量从一个维度重分片到另一个维度
      • 在 TPU 设备网络中，AllToAll 是实现此类重分片的高效方式，其成本随分片数量的增加呈亚线性增长（5.2节）
Per-Operator SPMD Partitioning
  • 分片器的核心是根据指定的分片方式，对每个算子进行从全尺寸到分片尺寸的转换
  • 虽然某些算子（如按元素算子）的分片支持非常简单，但论文将讨论需要跨分片通信的几种常见情况
  • 一般情况下，存在一些重要的技术挑战，论文将在3.3.3节中介绍
  • 为了使讨论与混合专家模型更相关，本节将重点关注爱因斯坦求和算子（Einsum）的分片，以展示几种通信模式
  • 为简化讨论，此处假设所有张量均采用均匀分片方式，即待分片维度的大小是分片数量的整数倍
Einsum Case Study
  • 爱因斯坦求和算子是实现混合专家模型的最关键算子
  • 在 XLA HLO 中，它们被表示为点积（Dot）运算，其中每个操作数（左操作数 LHS 或右操作数 RHS）包含三种类型的维度：
    • 批次维度（Batch dimensions） ：是易并行化维度（embarrassingly parallel dimensions）。所有左操作数、右操作数和输出都必须包含相同的批次维度集合，且输出中的每个元素仅依赖于左操作数和右操作数中对应的批次元素
    • 收缩维度（Contracting dimensions） ：仅存在于操作数中。左操作数和右操作数必须包含相同的收缩维度集合，这些维度在输出中会被求和并压缩
    • 非收缩维度（Non-contracting dimensions） ：也是并行维度，存在于某个操作数和输出中。左操作数和右操作数分别具有自己的非收缩维度集合，这些维度会被输出继承
  • 分片传播（Sharding propagation）优先选择在左操作数、右操作数和输出的批次维度上采用相同的分片方式，因为这样可以避免任何跨分片通信
  • 但在实际情况中，这并不总是可行的，以下三种情况需要跨分片通信：
    • 1）重分片（Resharding） ：在论文构建的混合专家模型中，专家分配逻辑（算法2第3行）需要在爱因斯坦求和运算后切换分片维度。由于使用AllToAll可以高效实现重分片（5.2节），论文首先在本地执行爱因斯坦求和运算，然后将结果重分片到目标维度，如图4a所示
    • 2）累积部分结果（Accumulating partial results） ：如果输入沿收缩维度进行分片，本地计算得到的结果只是部分结果，需要使用AllReduce将这些部分结果合并以生成最终结果，如图4b所示
    • 3）循环切片（Slicing in a loop） ：在某些场景下，论文还实现了一种类似于Cannon算法（1969）的方法，以限制每个分片上张量的大小。例如，如果两个操作数均沿非收缩维度进行分片，由于操作数的非收缩维度不同，无法直接在本地执行爱因斯坦求和运算。复制其中一个操作数不会导致冗余计算，但需要确保被复制的操作数能够放入设备内存
      • 因此，如果操作数的规模过大，论文会保持两个操作数的分片状态，通过循环迭代计算结果的每个切片，并使用 CollectivePermute 来传输输入切片（图4c）
  • 图4：带有跨设备通信的爱因斯坦求和算子分片示例：
    • （a）分片爱因斯坦求和算子
      • 彩色字母（G和E）表示每个张量的分片维度
      • 分片器决定首先沿G维度执行批次并行爱因斯坦求和运算，然后将结果重分片到E维度。（注：图中省略了S和M维度）
    • （b）在收缩维度上分片的简单爱因斯坦求和算子（矩阵乘法）
      • 每个分片计算部分结果，然后通过AllReduce合并得到完整结果
    • （c）在循环中使用CollectivePermute的爱因斯坦求和算子（矩阵乘法）
      • 通过循环每次计算一个切片，整个过程中不会出现全尺寸张量
Supporting a Complete Set of Operators
  • 为了使 SPMD 分片器能够支持完整的算子集合，且不对张量形状或算子配置施加额外限制，论文解决了若干额外挑战
  • 这些挑战通常涉及分片间的非对称计算或通信模式，而由于单一程序需要对所有分片都具有通用性，这些模式在 SPMD 中尤其难以表达
  • 论文不能简单地根据运行时设备ID在单一程序中创建多个分支，因为这会导致程序规模激增

Static shapes and uneven partitioning

  • XLA要求张量形状是静态的（注：中间表示中有限的动态性对于高效适配加速器通常是必要的）
  • 然而，对计算进行分片时，由于维度大小可能无法被分片数量整除，并非所有分片的输入/输出形状都相同
  • 在这种情况下，会将形状大小向上取整到分片数量的下一个整数倍，填充区域（padded region）中的数据可以是任意值
  • 在计算算子时，为了保证正确性，可能需要在填充区域中填充已知值
    • 例如，对 Reduce-Add 算子进行分片时，需要使用单位元0作为填充值
    • 考虑一个示例：待分片维度大小为 15，无法被 2（分片数量）整除，因此分片1比所需大小多一列
    • 论文生成一个范围为 [0, 8) 的 Iota 算子，加上分片偏移量（由分片 ID×8 计算得到），并与全尺寸偏移量（15）进行比较。根据比较得到的谓词值（predicate value），选择从操作数中取值或从 0 中取值，最终得到掩码操作数（masked operand）

Static operator configurations

  • XLA算子具有静态配置，例如卷积（Convolution）中定义的填充（padding）、步幅（stride）和扩张（dilation）
  • 但不同分片可能不会使用相同的算子配置
    • 例如，对于卷积算子，最左侧的分片会在其左侧应用填充，而最右侧的分片会在其右侧应用填充
    • 在这种情况下，分片器可能会选择让某些分片生成略多于所需的数据，然后切分出无关部分
  • 附录A.4讨论了卷积和类似算子的示例

Halo Exchange

  • 某些算子具有一种通信模式，需要与相邻分片交换部分数据，论文称之为 Halo Exchange
  • 论文使用 CollectivePermute 算子在分片间交换 Halo 数据
  • Halo Exchange 最典型的应用场景是对基于窗口的算子（如卷积、ReduceWindow）进行分片，因为相邻分片可能需要重叠的输入数据（图5a）
  • 在实际应用中，由于窗口配置（扩张、步幅和填充）的复杂使用以及非均匀 Halo 大小，这些算子的 Halo Exchange 通常需要结合适当的填充、切片和掩码操作
    • 论文在附录A.4中描述了各种场景
  • Halo Exchange 的另一个应用场景是改变形状大小的数据格式化算子（data formatting operator）
    • 例如，经过 Slice 或 Pad 算子后，张量的形状会发生变化，分片间的边界也会随之改变
    • 这需要论文重新对齐不同分片上的数据，而这可以通过 Halo Exchange 的形式来处理（图5b）
  • 其他数据格式化算子（尽管在逻辑上不改变形状大小）也可能需要 Halo Exchange ，这主要是由于静态形状约束和非均匀分片
    • 例如，Reverse 算子会反转张量中元素的顺序，但如果对其进行非均匀分片，需要在分片间移动数据，以确保填充在逻辑上位于结果张量的右侧
    • 另一个示例是Reshape算子：
      • 考虑将形状为 [3, 2] 的张量重塑为 [6]，其中输入在第一个维度上采用2种非均匀分片方式（分片形状为 [2, 2]），输出也采用2种分片方式（分片形状为[3]）
      • 由于非均匀分片，输入中存在填充，但重塑后输出张量不再有填充；因此，需要采用与Slice类似的 Halo Exchange 方式（图5c）

Compiler optimizations

  • SPMD 分片器会创建各种数据格式化算子，以执行切片、填充、连接、掩码和 Halo Exchange 操作
  • 为解决这一问题，论文利用了 XLA 在 TPU 上的融合能力（fusion capabilities），以及针对切片和填充的代码移动优化（code motion optimizations），在很大程度上隐藏了数据格式化的开销
  • 因此，即使在大量使用掩码和填充的卷积网络中，运行时开销通常也可以忽略不计

Massively Multilingual, Massive Machine Translation, M4

  • 待补充

Performance and Memory Consumption

  • 待补充

注：本文包含 AI 辅助创作

Paper Summary

  • 整体内容总结：
    • 论文提出了 序列并行（Sequence Parallelism） 与 减少激活重计算（Reducing Activation Recomputation）
    • 论文证明了将 SP (Sequence Parallelism) 与 TP (Tensor Parallelism) 结合使用可以显著减少所需的激活内存
    • 结合选择性激活重计算 (Selective Activation Recomputation)，论文证明了可以实现内存占用减少 \(5\times\)，并且恢复了使用完全激活重计算 (Full Activation Recomputation) 所引入的计算开销的 90% 以上
      • 传统上实现：通过不存储激活值，而是在反向传播时按需要重新计算它们，节省了内存但增加了冗余计算
      • 论文表明大部分这种冗余计算是不必要的，因为我们可以有选择的保存部分激活值
  • 补充：结合 张量并行 (tensor parallelism， TP)，这些技术几乎消除了重计算激活的需要
  • 论文在规模 高达一万亿参数 的语言模型上评估了论文的方法，结果表明：
    • 论文的方法将激活内存减少了 \(5\times\)，同时将激活重计算带来的执行时间开销降低了 90% 以上
  • 在 2240 个 NVIDIA A100 GPU 上训练一个 530B 参数的 GPT-3 风格模型 (2022) 时，论文实现了 54.2% 的模型浮点运算利用率 (Model Flops Utilization, MFU)，这比使用重计算实现的 42.1% 快了 29%

Introduction and Discussion

  • 随着 Transformer 模型规模向万亿参数扩展，需要模型并行 (model parallelism) 来将模型参数、激活值和优化器状态分布到多个设备上，以便它们能够放入设备内存并在现实的时间量内可训练
  • 虽然模型并行线性地减少了每个设备的参数数量（例如，当模型并行规模加倍时，每个设备的参数数量减半），但模型并行的扩展是有限度的
    • 张量级模型并行 (Tensor-level model parallelism, TP) 增加了通信需求并引入了更小且性能较差的矩阵乘法，使得将模型分割到大量设备上效率低下
      • 因此，TP 通常仅限于通过高速带宽连接（例如 DGX 服务器内通过 NVLink 连接的 GPU）的相对较小的 GPU 组
    • 流水线并行 (Pipeline parallelism， PP) 需要存储几个微批次 (microbatches) 的激活值以减少流水线气泡 (pipeline bubble) (2021)
      • 因此， PP 只能帮助减少存储模型参数和优化器状态所需的内存，无法在保持高设备利用率的同时减少激活值所需的内存
    • 因此，激活值的存储迅速成为扩展大型 Transformer 模型的关键问题
      • 问题：实际上，对于不适用 Activation Recomputation 的场景，PP 也减少激活存储了；但是对于 Activation Recomputation 的场景，每一层存储的激活基本都还在，所以相当于此时 PP 对激活存储是没有节省的
  • 为了量化这一点，图 1 显示了从 22B 参数到 1 万亿（1T）参数的四种模型配置所需的内存（模型配置的详细信息在表 3 中提供）
    • 可以看出，对于所有这些情况，基线情况所需的内存都超过了 NVIDIA A100 GPU 提供的 80GB 内存
    • 缓解这种内存压力的标准方法是简单地不存储大部分激活值，并在反向传播 (backward pass) 期间根据需要重新计算它们以计算梯度 (2016)
      • 这种方法通常称为“梯度检查点 (gradient checkpointing)”或“激活重计算 (activation recomputation)”，会带来严重降低训练效率的代价
    • 对于 Transformer 架构，大多数先前的工作在 Transformer 层边界处检查点（或存储）激活值，并在反向传播中重新计算其余必要的激活值
    • 在论文中，论文将这种方法称为“完全激活重计算 (full activation recomputation)”
    • 在论文的训练运行中，论文观察到当使用完全激活重计算时 ，会产生 \(30-40\%\) 的执行时间开销
  • 论文提出了新技术，帮助缓解存储激活值的内存压力，从而减少重计算激活值的需求
    • 这些技术特定于 Transformer 架构，并且易于实现，对计算效率没有影响或影响非常小
  • 正如论文在第 2 节中详细说明的那样，还有其他几种技术可以减少训练大型模型的内存需求，例如：
    • 跨数据并行 Rank (data parallel ranks) 分区各种数据
    • 将数据卸载到 CPU 内存 (2020; 2021)
  • 以上这些技术与论文提出的技术是互补的，可以额外采用以获得更大的内存节省；但通常这些其他技术比论文提出的技术具有更高的实现成本和对计算效率的更大影响
  • 将这些技术与论文的技术进行比较的分析超出了论文的范围，留待未来工作
  • 论文工作内容包括：
    • 首先论文简要回顾 Transformer 架构，然后建立一个近似公式，用于存储 Single Stack Transformer 模型前向传播 (forward pass) 中激活值所需的内存
      • 使用这个公式，我们可以研究不同形式的模型并行如何影响激活内存需求
    • 然后论文引入序列并行 (sequence parallelism) 与 TP (tensor parallelism) 结合，以防止在不适合标准 TP 的区域中冗余存储激活值
      • 论文表明，通过有选择地保存哪些激活值以及重新计算哪些激活值，我们可以消除大部分重计算成本，同时仅使用不进行重计算时所需内存的一小部分
    • 最后论文提出了几个实验，测量这些技术对训练各个组成部分以及完整训练吞吐量的改进
  • 模型并行 (model parallelism) 使得能够跨多个 GPU 训练非常大的模型
    • 模型参数以及这些模型相关的优化器状态需要大量内存，无法容纳在单个 GPU 上
    • 即使能够将模型放入单个 GPU（例如，通过在主存和设备内存之间交换参数 (2021)），所需的高计算操作数量可能导致不切实际的长训练时间
    • 结论：需要并行化
  • 两种常用的模型并行形式用于将模型参数分布到 GPU 上：
    • 1） TP (tensor parallelism)，其中每层的参数分布在许多设备上 (2018; 2019; 2021)；
    • 2） PP (pipeline parallelism)，其中模型沿网络的层维度分割 (2019; 2021a; 2019)
  • 一些最近的方法结合了两种类型的模型并行，以支持训练高达 1T 参数的大型模型 (2021)
    • 模型并行的替代方案是结合一些训练技术与数据并行 (data parallelism)，以实现大规模模型训练 (2020; 2021;)
      • 这种方法基于将优化器状态、梯度和参数跨数据并行 Rank 进行分片 (sharding)
    • 最近的一个扩展 (2021) 使用 CPU 卸载技术，使得能够在少量 GPU 上训练数万亿参数的模型
      • 与模型并行相比，这些基于数据并行的技术效率较低，并且不能很好地扩展到大量 GPU (2021)，因此更适合在资源受限的环境中进行模型微调 (finetuning)
    • 论文仅关注模型并行优化，将这些技术与论文的技术进行比较的分析超出了论文的范围
  • 此外，Megatron-LM (2019) 中引入的 TP 在一定程度上帮助减少了激活内存
    • 在这种方法中，Transformer 的某些部分激活值没有在 TP Rank 之间分割，增加了激活内存开销
  • Li 等 (2021a) 提出的序列并行 (sequence parallelism)，其中激活值沿序列维度分区贯穿整个网络可以缓解这个问题
    • 但他们的方法类似于数据并行，要求在所有设备上复制参数和优化器状态，这使得它不适合大型模型训练
  • Sagemaker (2021) 和 GSPMD (2021) 提出了内存高效的 TP 版本，这些版本在整个网络中沿隐藏维度 (hidden dimension) 在设备之间分割激活值
    • 这些方法的主要缺点是它们包含多设备层归一化 (multi-device layer normalization)，这在计算/通信上非常低效
      • LayerNorm 操作需要沿隐藏维度顺序计算均值和方差
      • 如果 LayerNorm 沿隐藏维度分割（即使用 TP），它将增加两个额外的全规约 (all-reduce) 操作
      • 而且 LayerNorm 反向传播需要两个额外的 sequential reductions（即另外两个 all-reduce）
    • 与 LayerNorm 计算时间相比，每个 Transformer 层的这 4 次通信将引入显著的开销
    • LayerNorm 操作的这种通信开销可以通过 SP 来避免
  • 在论文中，论文提出了一种新技术，它利用了 TP 和 SP 的优点，而没有先前方法的任何缺点
    • 换句话说，论文的技术结合了 TP 和 SP ，显著减少了激活内存，而没有任何额外的计算、通信或内存开销
  • 论文的论文特别侧重于研究 Transformer 架构，因为作者相信它是一个重要且广泛使用的架构，足以证明手动检查各个层以了解其内存和计算贡献并推导最佳并行化方案是合理的
  • 通过分析建模，论文表明内存（包括激活值和参数）在设备之间均匀分区
    • 换句话说，论文的并行化策略在内存需求方面是最优的
    • Transformer 网络被广泛使用，足以高度利用这种手动搜索和验证最佳并行化策略的方法
    • 自动搜索 (2019; 2019) 有可能找到这种最优策略，但据论文所知，论文尚未看到这些方法应用于 Transformer 网络的任何已发表结果
  • 与并行化方案类似，论文提出的选择性激活重计算 (selective activation recomputation) 是基于手动搜索 Transformer 模型在激活内存需求和重计算之间的最佳权衡
  • 探索自动方法 (2021) 是否可以补充并可能改进这种权衡是未来工作的一个有趣方向

Transformer Architecture

  • 论文考虑一个具有 \(L\) 层的 Single Stack Transformer 编码器或解码器，如图 2 所示
  • 符号定义：\(s\) 是序列长度 (sequence length)，\(h\) 是隐藏维度大小 (hidden dimension size)，\(v\) 是词汇表大小 (vocabulary size)
  • 在网络的开始，输入 Tokens 被送入一个词嵌入表(word embedding table)：
    $$v\times h$$
  • Token 嵌入与学习到的位置嵌入(positional embeddings) 相结合：
    $$s\times h$$
  • 嵌入层的输出，即 Transformer 块的输入，是一个 3 维张量：
    $$s\times b\times h$$
    • 其中 \(b\) 是微批次大小 (microbatch size)
  • 每个 Transformer 层由一个具有 \(a\) 个注意力头 (attention heads) 的自注意力块 (self-attention block) 和一个具有两层的多层感知器 (MLP) 组成，该 MLP 将隐藏大小增加到 \(4h\)，然后将其减小回 \(h\)
  • 每个 Transformer 层的输入和输出具有相同的大小
    $$ s\times b\times h $$
  • 最后一个 Transformer 层的输出被投影回词汇表维度以计算交叉熵损失 (cross-entropy loss)
    • 论文假设词嵌入和输出层权重是共享的
  • 论文相关变量名称列于表 1 中以供参考：

Activation Memory

  • 论文推导出一个近似公式，用于计算如图 2 所示的 Single Stack Transformer 模型在前向传播中存储激活所需的内存
  • 请注意，论文中的 “激活（activations）” 指的是 在前向传播中创建并在反向传播期间梯度计算所必需的任何张量
    • 注：不包括模型的主要参数和优化器状态，但包括 Dropout 操作使用的掩码
  • 论文只考虑内存的主要贡献者，而忽略小的缓冲区
    • 小的缓冲区包括 LayerNorm 输入的均值和方差（\(2sb\)）以及 GEMM 操作的偏置（\(O(h)\)）
    • 总的来说，这些缓冲区占激活内存的比例远小于 1%，因为隐藏维度（\(h\)）和序列维度（\(s\)）的数量级都是数千，因此与 GEMM 和 LayerNorm 激活的 \(O(sbh)\) 大小相比，\(2sb\) 和 \(O(h)\) 都缺少了其中一个因子
  • 论文还假设网络和激活以 16-bit 浮点格式存储，因此每个元素需要 2 字节的存储空间
    • 唯一的例外是 Dropout 掩码，每个元素只需要 1 个字节
  • 请注意，除非明确提及，本节中报告的所有大小均以字节（Bytes）为单位，而不是元素数量

Activations Memory Per Transformer Layer

  • 如图 2 所示，每个 Transformer 层包含一个注意力块和一个 MLP 块，通过两个 LayerNorm 连接
  • 下面，论文推导存储这些元素激活所需的内存：
  • 注意力块 (Attention block)：
    • 包括自注意力操作，后跟一个线性投影和注意力 Dropout
    • 线性投影存储其输入激活，大小为 \(2sbh\)，注意力 Dropout 需要一个大小为 \(sbh\) 的掩码
    • 图 3 所示的自注意力包含几个元素：
      • Query (\(Q\)), Key (\(K\)), and Value (\(V\))矩阵乘法 (matrix multiplies)： 论文只需要存储它们共享的输入，大小为 \(2sbh\)
      • \(QK^{T}\) 矩阵乘法 (matrix multiply)： 需要存储 \(Q\) 和 \(K\)，总大小为 \(4sbh\)
      • Softmax： 反向传播需要存储 Softmax 输出，大小为 \(2as^{2}b\)
      • Softmax Dropout： 只需要一个大小为 \(as^{2}b\) 的掩码
        • 问题：Softmax 也有 Dropout 吗？
        • 理解：在 Softmax 之后施加 Dropout，随机将部分注意力权重置为 0，目的是防止模型过度依赖某些特定的注意力权重，减少过拟合风险，提高模型的泛化能力，但原始 Transformer 论文中，没看到这部分的 dropout，只有 Residual Dropout（\(P_\text{drop} = 0.1\)）
      • 注意力作用于值（\(V\)）(Attention over Values (\(V\)))： 论文需要存储 Dropout 输出（\(2as^{2}b\)）和值（\(2sbh\)），因此需要 \(2as^{2}b+2sbh\) 的存储空间
    • 将上述值相加，注意力块总共需要 \(11sbh+5as^{2}b\) 字节的存储空间
  • MLP：
    • 两个线性层存储它们的输入，大小分别为 \(2sbh\) 和 \(8sbh\)
    • GeLU 非线性操作也需要其输入（大小为 \(8sbh\)）用于反向传播
    • 最后，Dropout 存储其掩码，大小为 \(sbh\)
    • MLP 块总共需要 \(19sbh\) 字节的存储空间
  • LayerNorm：
    • 每个 LayerNorm 存储其输入，大小为 \(2sbh\)，因此论文总共需要 \(4sbh\) 的存储空间
  • 最终，将注意力、MLP 和 LayerNorm 所需的内存相加，存储 Transformer 网络单层激活所需的内存为：
    $$\text{Activations memory per layer}=sbh\left(34+5\frac{as}{h}\right) \tag{1}$$
  • 上述方程适用于未应用任何形式模型并行的情况

Model Parallelism

  • 论文先量化 TP 对每层所需激活内存的影响，然后介绍一种将 SP 与 TP 相结合的新方法，以进一步减少每层激活所需的内存
  • 在本小节的最后，论文还讨论了 PP 对激活内存的影响，并推导了激活所需总内存的公式
TP (Tensor Parallelism)
  • 论文使用 Shoeybi 等人 (2019) 开发的 TP ，并按照图 4 所示并行化注意力块和 MLP 块
    • 这种形式的并行性引入了两个额外的通信操作 \(f\) 和 \(\bar{f}\)
    • 更多细节，请参阅该论文 (2019)
  • TP 不仅并行化了注意力块和 MLP 块内的模型参数和优化器状态，还并行化了这些块内部的激活
    • 这些块的输入激活（例如，输入到 \(Q\)、\(K\) 和 \(V\) 矩阵乘法的输入，或输入到 \(h \to 4h\) 线性层的输入）没有被并行化，只有每个块内的激活在 TP 组内被划分
    • 假设采用 \(t\) 路 TP ，存储激活所需的每层内存从公式 1 减少到：
      $$\text{Activations memory per layer}=sbh\left(10+\frac{24}{t}+5\frac{as}{ht}\right) \tag{2}$$
SP (Sequence Parallelism)
  • 如图 4 所示， TP 并行化了 Transformer 层中训练过程中最耗时的部分，因此它在计算上是高效的
  • 但它保留了 LayerNorm 以及注意力块和 MLP 块之后的 Dropout 操作不变，因此 LayerNorm 和 Dropout 在 TP 组内是复制的（单卡上需要完整保留所有参数和激活）
  • 这些元素不需要大量计算，但需要相当数量的激活内存
  • 定量地说，公式 2 中的 \(10sbh\) 部分就是由于这些复制操作造成的，因此它们没有除以 TP 大小 \(t\)
  • 论文注意到，在 Transformer 层的非 TP 区域中，操作在序列维度上是独立的
    • 这一特性允许论文沿序列维度 \(s\) 划分这些区域
    • 沿序列维度划分减少了激活所需的内存
  • 这种额外的并行性级别在 \(f\) 之前和 \(\bar{f}\) 之后引入了新的通信集合，这些操作将充当 SP 区域和 TP 区域之间的转换器
    • 例如，在前向传播中，论文需要在图 4 中的操作符 \(f\) 之前进行一次额外的 All-Gather
    • 这些额外的通信会引入开销并减慢训练速度
  • 为了避免这些额外的通信，论文将这些操作与 \(f\) 和 \(\bar{f}\) 操作符结合起来，并引入新的操作 \(g\) 和 \(\bar{g}\)，如图 5 所示
    • 可以看出，\(g\) 和 \(\bar{g}\) 是 SP 区域和 TP 区域之间的转换器
    • 论文在本小节的剩余部分推导这些操作
  • 论文使用 MLP 块详细说明 \(g\) 和 \(\bar{g}\) 的推导。在非并行形式中，如图 2 所示，LayerNorm 后接 MLP 块可以表述为：
    $$
    \begin{align}
    Y &=\text{LayerNorm}(X), \\
    Z &=\text{GeLU}(YA), \\
    W &=ZB, \\
    V &=\text{Dropout}(W),
    \end{align}
    $$
    • 其中：
      • \(X\) 是 LayerNorm 的输入，大小为 \(s\times b\times h\)
      • \(A\) 和 \(B\) 是线性层的权重矩阵，大小分别为 \(h\times 4h\) 和 \(4h\times h\)
  • 上述操作的张量和 SP 组合形式如图 6 所示
  • 下标表示在加速器间的划分（注：理解为 Rank），上标表示划分所沿的维度 ，例如：
    • \(X_{1}^{s}\) 是 \(X\) 在第一个加速器上的部分，沿 \(s\) 维度（序列维度）划分
    • \(Z_{2}^{h}\) 是 \(Z\) 在第二个加速器上的部分，沿 \(h\) 维度（隐藏维度）划分
  • LayerNorm 的输入沿序列维度并行化
    $$ X=[X_{1}^{s},X_{2}^{s}]$$
  • 因此，LayerNorm 的输出也将沿序列维度并行
    $$Y=[Y_{1}^{s},Y_{2}^{s}]$$
  • 带有 GeLU 非线性的线性层需要完整的输入 \(Y\)，因此论文需要执行一次 All-Gather
    • 这意味着 \(g\) 在前向传播中是一个沿序列维度的 All-Gather 操作
  • 通过将 \(A\) 沿其列划分（\(A_{1}^{c}\) 和 \(A_{2}^{c}\)），将 \(B\) 沿其行划分（\(B_{1}^{r}\) 和 \(B_{2}^{r}\)），论文避免了通信（更多细节请参见 (2019)）并得到 \(W_{1}\) 和 \(W_{2}\)
    • 这两个张量不再是并行的，需要在输入到 Dropout 层之前求和为 \(W=W_{1}+W_{2}\) (对应 Reduce 操作)
    • 但 Dropout 需要其输入在序列维度 \(s\) 上是并行的 （对应 Scatter 操作）
  • 论文不是先求和再在序列维度上并行化，而是将这两个操作合并为一个 Reduce-Scatter 操作
    • 因此，\(\bar{g}\) 在前向传播中可以是一个单一的 Reduce-Scatter 操作
  • 综上所述，论文得到：
    $$
    \begin{align}
    [Y_{1}^{s},Y_{2}^{s}] &=\text{LayerNorm}([X_{1}^{s},X_{2}^{s}]) \\
    Y &=g(Y_{1}^{s},Y_{2}^{s}), \\
    [Z_{1}^{h},Z_{2}^{h}] &=[\text{GeLU}(YA_{1}^{c}),\ \text{GeLU}(YA_{2}^{c})],\\
    W_{1} &=Z_{1}^{h}B_{1}^{r}\ \text{ and }\ W_{2}=Z_{2}^{h}B_{2}^{r},\\
    [W_{1}^{s},W_{2}^{s}] &=\bar{g}(W_{1},W_{2}),\\
    [V_{1}^{s},V_{2}^{s}] &=[\text{Dropout}(W_{1}^{s}),\ \text{Dropout}(W_{2}^{s})]
    \end{align}
    $$
  • 如果论文对反向传播进行类似的分解，论文会发现 \(g\) 和 \(\bar{g}\) 是彼此共轭的（Conjugate）
    • \(g\) 在前向传播中是 All-Gather，在反向传播中是 Reduce-Scatter；
    • \(\bar{g}\) 在前向传播中是 Reduce-Scatter，在反向传播中是 All-Gather
  • 对 Transformer 层的 LayerNorm 后接注意力部分进行类似的分解，得到图 5
  • TP 在单次前向和反向传播中需要四次 All-Reduce，而 TP 与 SP 结合在单次前向和反向传播中需要四次 All-Gather 和四次 Reduce-Scatter
    • 乍一看，似乎张量与 SP 比 TP 需要更多的通信
    • 但论文注意到，Ring All-Reduce 由两个步骤组成：
      • 一个 Reduce-Scatter 后接一个 All-Gather
    • 因此，TP 与张量加 SP 所使用的通信带宽是相同的
      • 问题：Reduce-Scatter 和 All-Gather 只有在无缝连着的时候才能说等价于 Ring All-Reduce 吧
      • 理解：是的，但是无论是否无缝连着，Reduce-Scatter + All-Gather 的通信都等于 Ring All-Reduce，因为都是 Reduce，然后 Scatter，继而 Gather
        • 补充：论文后面的 6.2 章节也提到了：reduce-scatter 和 all-gather 组合的执行速度比单独的 all-reduce 慢
    • 所以， SP 不会引入任何通信开销！
  • 根据公式 3， SP 与 TP 结合，将反向传播所需的所有激活沿并行维度划分
    • 注：第一个线性操作所需的张量 \(Y\)）不会被划分
    • 为了缓解这个问题，论文不存储完整的张量 \(Y\) 用于反向传播，只在第 \(i\) 个 TP Rank 上存储 \(Y_{i}^{s}\) 部分，并在反向传播中执行一次额外的 All-Gather
    • 为了消除这个额外 All-Gather 引入的延迟，论文将此通信与计算 \(Y\) 梯度所需的计算重叠，从而减少了开销
  • 使用 SP 与 TP 结合，存储每个 Transformer 层激活所需的内存从公式 2 减少到：
    $$
    \begin{align}
    \text{Activations memory per layer} &= sbh\left(\frac{10}{t}+\frac{24}{t}+5\frac{as}{ht}\right) \\
    &=\frac{sbh}{t}\left(34+5\frac{as}{h}\right) \tag{4}
    \end{align}
    $$
  • 上述方程现在是公式 1 除以 TP 大小
    • 这意味着使用张量和 SP ，我们可以在 TP 组内分布激活，并将所需内存减少 TP 大小 \(t\) 倍
PP (Pipeline Parallelism)
  • PP 简单地将 Transformer 的 \(L\) 层划分为 \(\frac{L}{p}\) 组层，其中 \(p\) 是 PP 大小
    • 但 PP 并不会将激活所需的总内存均匀地除以 \(p\)
    • 这是由于 PP 调度引入了重叠（overlapping）以减少流水线气泡 (2021)
      • 注：overlap 意味着对应的激活需要同时在内存中
  • 为了量化这一点，论文考虑 PipeDream (2020) 中开发的 1F1B 流水线调度
    • 具有最小化流水线气泡的调度对流水线的第一个阶段施加最大的内存压力
      • 流水线的第一阶段指的是第一个 \(\frac{L}{p}\) 层组，也包括输入嵌入
    • 激活内存随流水线阶段变化的可视化图见附录 B
    • 为了保持流水线加压并避免额外的空闲时间，第一阶段必须存储 \(p\) 个微批次的激活（更多细节参见 (2021) 的图 4-top）
    • 每个阶段包含 \(\frac{L}{p}\) 层，因此第一阶段必须存储激活层为：
      $$p \times \frac{L}{p} = L$$
      • 注：（与 PP 大小 \(p\) 无关）
    • 因此，第一阶段存储激活所需的总内存为：
      $$\text{Total activations memory of 1st stage}=\frac{sbhL}{t}\left(34+5\frac{as}{h}\right) \tag{5}$$
  • 对于其他流水线调度，所需的总内存会略有不同
    • 例如，Megatron-LM (2021) 中开发的交错调度需要存储 \(L(1+\frac{p-1}{pm})\) 层的激活，其中 \(m\) 是交错阶段的数量
      • 因此，如果使用交错调度，则总激活内存应乘以 \((1+\frac{p-1}{pm})\)

Total Activations Memory

  • 总的激活内存的大部分由公式 5 提供，但该公式未包括输入嵌入、最后一个 LayerNorm 和输出层所需的激活内存，如图 2 所示
  • 位置和词嵌入不需要存储任何大量的激活用于反向传播，但 Dropout 需要存储
    • 嵌入层中的 Dropout 也沿序列维度并行化
    • 因此，它将需要 \(\frac{sbhp}{t}\) 的存储空间
    • 注意，因子 \(p\) 来自 PP 以及论文需要存储 \(p\) 个微批次的事实（见第 4.2.3 节）
  • 输出层之前的 LayerNorm 也使用 SP ，因此需要 \(\frac{2sbh}{t}\) 的存储空间
  • 输出层投影到词汇维度需要存储其输入，大小为 \(\frac{2sbh}{t}\)
  • 交叉熵损失需要存储以 32 位浮点数计算的 logits，因此需要 \(\frac{4sbv}{t}\) 的存储空间
  • 请注意，由于论文只考虑流水线第一阶段的激活，上述激活（即总共 \(\frac{4sbh}{t} \cdot (1+\frac{v}{h})\)）仅在没有 PP （\(p=1\)）的情况下被计入
  • 加上上述内存，由于输入嵌入、最后一个 LayerNorm 和输出层产生的额外内存为：
    $$\frac{sbhL}{t}\left(\frac{p}{L}+\delta_{p=1}\frac{4}{L}\left(1+\frac{v}{h}\right)\right)$$
    • 其中 \(\delta_{p=1}\) 是一个 0-1 函数，在 \(p=1\) 时为 1，否则为 0
  • 论文注意到，与公式 5 中的项 \(34+5\frac{as}{h}\) 相比，\(\frac{p}{L}\) 和 \(\frac{4}{L}\cdot (1+\frac{v}{h})\) 都是可以忽略的
    • 例如，对于一个具有 22B 参数的模型，这些额外项占总激活内存需求的比例小于 \(0.01\%\)
    • 问题：\(\frac{v}{h}\) 对不同配置也不算小吧，也有几十了
  • 因此，公式 5 是总所需激活内存的一个良好近似，论文将在论文的其余部分使用它

选择性激活重计算，Selective Activation Recomputation

  • 公式 (5) 给出的总激活内存需求对于大型模型而言仍然相当可观
  • 激活重计算 (2016) 通过存储（或 checkpoint）一组层的输入激活，并在反向传播期间使用额外的前向传递重新计算其他所需的激活，来克服这一内存限制（论文中称之为完全激活重计算）
  • 假设检查组仅包含单个层，并忽略 Transformer 层之外的激活，该方法将激活所需的总内存减少到
    $$ 2sbhL $$
    • 论文也注意到，如果论文只在每个 TP Rank 上存储部分激活，这个所需内存可以进一步减少到
      $$ \frac{2sbhL}{t} $$
      • 但这种方法需要每层额外的 all-gather 操作 ，会增加通信开销，因此论文不考虑这种方法
  • 与存储所有激活（公式 (5)）相比，检查（checkpointing）所有 Transformer 层显著减少了训练模型所需的内存量
    • 这种减少确实是以重计算（一次额外的前向传递）为代价的，这可能会带来高达 \( 30-40\% \) 的计算时间开销
  • 为了平衡内存节省和计算开销，理想的做法是：仅仅 Checkpointing 足够的激活，使得给定的模型并行配置能够在设备内存的限制下进行训练
  • SP 提供的内存节省使得比以往更多的配置可以在不进行重计算的情况下进行训练，但大型模型的最佳模型并行配置通常仍然需要存储和重计算部分激活
  • 选择存储与重计算的激活量的一种简单方法是仅 checkpoint 部分 Transformer 层，并存储其他所有层的激活
    • 这种方法对于大型模型的扩展性并不好；
      • 例如，在训练 MT-NLG 时，每个设备只有三层，限制了平衡内存与计算的粒度
    • 此外，论文注意到并非所有激活都需要相同数量的操作来重新计算，因此在选择存储哪些激活和重计算哪些激活时更智能是有益的
  • 论文不是 checkpoint 和重计算完整的 Transformer 层，而是建议仅 checkpoint 和重计算每个 Transformer 层中占用大量内存但重新计算计算成本不高的部分 ，即选择性激活重计算（Selective Activation Recomputation）
  • 为此，论文注意到公式 (5) 中的项 \( \frac{5as}{h} \) 是由于在通过计算 Q、K 和 V 值的线性层增加了网络宽度之后的注意力操作（after the width of the network is increased by the linear layer calculating the Q, K, and V values）；
    • 即如图 3 所示的 \( QK^{T} \) 矩阵乘法、softmax、softmax dropout 和注意力作用于 \( V \)
    • 这些操作通常具有大的输入尺寸，因此激活也很大，然而，每个输入元素的浮点操作数非常低
    • Transformer 层的其余部分对应公式 (5) 中的常数项 34
    • 因此，对于 \( \frac{5as}{h}>34 \) 的大型模型，如果论文 checkpoint 并重计算 Transformer 层的这一部分，论文存储的激活不到一半，并且重计算那些未存储的激活只需付出适度的成本
  • 为了量化这一点，让论文考虑 GPT-3 (2020) 和 MT-NLG (2022) 模型，这些是迄今为止训练过的一些最大模型
    • 对于 GPT-3，\( a=96 \)，\( s=2048 \)，\( h=12288 \)，因此 \( \frac{5as}{h}=80 \)
    • 对于 MT-NLG，\( a=128 \)，\( s=2048 \)，\( h=20480 \)，所以 \( \frac{5as}{h}=64 \)
  • 将这些数字与常数项 34（即该层其余部分的因子）进行比较，我们可以看到这些激活占总激活的很大一部分
    • 因此，通过使用选择性激活重计算，我们可以分别为 GPT-3 和 MT-NLG 模型节省 \( 70\% \) 和 \( 65\% \) 的激活内存需求
    • 这些激活的重计算仅为这两个模型引入了 \( 2.7\% \) 和 \( 1.6\% \) 的 FLOPs 开销
    • 有关 FLOPs 计算的更多细节，请参见附录 A
  • 随着序列长度的增加，序列长度与隐藏大小的比率 \( (\frac{s}{h}) \) 将增加，执行选择性部分重计算的相对成本也会增加，从而使选择性激活重计算更加昂贵
  • 作者认为，对于任何实际的 \( \frac{s}{h} \) 比率，选择性激活重计算仍然是平衡内存节省和计算的最佳重计算策略
  • 使用这种形式的选择性激活重计算，存储激活所需的内存从公式 (5) 减少到：
    $$
    \text{Total required memory}=34\frac{sbhL}{t}.
    $$
    • 上述公式表明，使用选择性激活重计算使得所需的激活内存随序列长度线性缩放，并且与注意力头的数量无关
    • 如第 4.2.3 节所述，在使用交错流水线调度的情况下，上述公式需要乘以 \( (1+\frac{p-1}{pm}) \)
  • 当使用 PP 时，如第 4.2.3 节所述，尽管给定设备只有 \( \frac{L}{p} \) 层，但第一阶段仍然必须存储相当于 \( L \) 层的激活，因为它必须存储 \( p \) 个微批次的激活以保持流水线压力
    • 在这种情况下，可以采用一种额外技术来降低重计算成本，即根据可用设备内存尽可能多地存储所有微批次的激活，并对其余部分进行完全或选择性重计算
    • 实际上，论文发现应用 SP 和选择性激活重计算后，重计算开销已经足够小，以至于这种额外技术带来的改进非常有限
    • 该技术在附录 C 中有更详细的描述和分析

Evaluations

  • 在本节中，论文评估了论文提出的方法对内存使用和训练执行速度的影响
  • 表 3 列出了评估中使用的模型配置
  • 论文考虑了参数规模高达一万亿的模型，对于所有这些模型，TP 大小设置为 8
  • 对于 175B 和 530B 模型，论文使用具有三个交错阶段 (\( m=3 \)) 的交错调度
  • 在所有情况下，序列长度设置为 \( s=2048 \)，词汇表大小设置为 \( v=51200 \)
  • 论文还注意到，在这些评估中没有考虑数据并行，因为论文的方法与数据并行无关
  • 因此，论文分析中使用的批大小远低于端到端训练中使用的批大小
  • 论文的所有结果都是在 Selene 超级计算机上使用混合精度运行的
  • 每个集群节点有 8 个 NVIDIA 80GB A100 GPU，通过 NVLink 和 NVSwitch 相互连接
  • 每个节点有八个 NVIDIA Mellanox 200Gbps HDR Infiniband HCA 用于应用程序通信

Memory Usage

  • 论文通过借助前向和后向钩子在运行的不同时间点跟踪 GPU 活跃内存和 GPU 总内存（在 PyTorch 中可用）来验证内存消耗模型
  • 论文验证了测量的内存使用情况与论文的分析模型密切匹配
  • 表 2 的第一列总结了论文讨论的不同技术所需的内存
  • 为了量化这一点，图 7 显示了不同技术使用的激活内存，表示为在所有 TP Rank 上保持所有激活所需内存的百分比，即公式 (2)
  • 单独来看，两种技术都将内存需求减少了近一半，结合使用则提供了 5 倍的减少，将内存需求降至 20% 以下。这仅是完全激活重计算（基线水平的 10%）的约 \( \sim 2\times \)
  • 如果没有 SP 和选择性重计算共同提供的内存节省，这些模型都无法装入内存
    • 请注意，所有这些结果都包括了附录 B 中描述的内存优化

Execution Time per Layer

  • 表 2 的第二列和第三列总结了论文讨论的不同技术下每个 Transformer 层的计算量（FLOPs）和通信字节数
  • 为了量化这一点，表 4 显示了针对 22B 模型的各种实验下，执行一个 Transformer 层前向和后向传递的时间
    • 前两行显示， SP 对完成一个 Transformer 层所需的时间提供了适度的改进，将前向时间从 7.7ms 减少到 7.2ms，加速了 6%
      • 这种改进来自于 LayerNorm 和 Dropout 层在 \( \frac{1}{t} \) 的数据上执行
    • 论文还发现，尽管移动的数据量是相同的 ，但 reduce-scatter 和 all-gather 组合的执行速度比单独的 all-reduce 慢 ，这降低了 SP 带来的改进
      • 请注意，这种加速是使用 SP 主要优势之外的额外好处，其主要优势是节省内存，从而减少激活的重计算
  • 表 4 的后两行显示，如果论文在重计算哪些操作上具有选择性（由于 SP ，我们可以在更多配置中做到这一点），我们可以显著减少后向传递中重计算的开销
    • 选择性重计算的开销是 \( 1.3 \)ms，占 11.9ms 基线的 \( 11\% \)，而重计算整个层的开销是 \( 7.6 \)ms 或 \( 64\% \)
    • 对于前向和后向总时间，开销是 \( 7\% \) 对比 \( 39\% \)
    • 请注意，重计算整个层的开销为 \( 39\% \)（而非预期的 \( 33\% \)）是由于后向传递中的一项优化，论文将 all-reduce 通信与线性层权重的梯度计算重叠
    • 正如论文后面将看到的，这种好处随着模型大小的增加而增加
    • 表 4 的最后一行显示了选择性重计算和 SP 结合的好处：当两种技术一起使用时，开销降至仅 \( 4\% \)
  • 图 8 显示了所有测试案例的相同细分
    • 论文看到，随着模型规模的增长，开销的减少也在增加
    • 对于 \( 530 \)B 和 \( 1 \)T 的情况，开销仅为 \( 2\% \)，而完全重计算的开销为 \( 36\% \)

End-to-End Iteration Time

  • 表 5 列出了表 3 中列出的四种配置中每一种的完整端到端迭代时间
    • 论文发现，对于所有测试的配置，论文提出的技术相比不使用 SP 的完全重计算，在吞吐量上提供了介于 \( 29.0\% \) 和 \( 32.1\% \) 之间的改进。这些节省将直接转化为更短的训练时间
  • 论文定义了模型 FLOPs 利用率（MFU）和硬件 FLOPs 利用率（HFU） ，定义方式与 Chowdhery 等 (2022) 类似
    • 模型 FLOPs 是执行单次前向和后向传递（单次迭代）所需的浮点操作数（不包括激活值），与实现方式和硬件限制无关
      • 因此，模型 FLOPs 是硬件和实现无关的，仅取决于底层模型
      • 注：但 模型 FLOPs 利用率（MFU）的分母是理论最大峰值 FLOPs，是与硬件有关的
    • 硬件 FLOPs 表示硬件上每次迭代实际执行的浮点操作
      • 因此，如果实现需要激活重计算（例如论文的实现），那么硬件 FLOPs 将大于模型 FLOPs
      • 理解：因为增加了一些激活重计算的操作，这里的硬件主要强调的是实际计算量（包括激活重计算的操作量）
      • 注：硬件 FLOPs 利用率（HFU）的分母是也是理论最大峰值 FLOPs
    • 论文在附录 A 中提供了模型和硬件 FLOPs 的严格下界公式
    • 对于论文的方法，硬件 FLOPs 与模型 FLOPs 的比率约为 \( 1+\frac{s}{18h} \)
  • 随后，论文分别将模型和硬件每秒 FLOPs 定义为模型和硬件 FLOPs 除以迭代时间
    • 使用这些定义，MFU 和 HFU 分别定义为模型和硬件每秒 FLOPs 除以加速器的理论峰值每秒 FLOPs
      $$
      \begin{align}
      MFU &= \frac{\text{observed throughput (tokens-per-second) }}{\text{theoretical peak FLOPs of device}} \\
      HFU &= \frac{\text{observed FLOPs}}{\text{theoretical peak FLOPs of device}}
      \end{align}
      $$
    • 以上定义参考自：PaLM: Scaling Language Modeling with Pathways, 2022, Google
  • 表 5 提供了所有四种配置的 MFU 和 HFU
    • 随着模型规模的增加，论文实现了更好的 GPU 利用率，对于一万亿参数模型，论文分别达到了 \( 56.3\% \) 和 \( 57.0\% \) 的 MFU 和 HFU
  • 虽然论文的分析中没有考虑初始化、评估、检查点等时间，但这些时间与迭代时间乘以迭代次数（这主导了端到端训练时间）相比可以忽略不计
    • 用于训练大型语言模型的迭代次数根据训练目标而有很大差异
    • 因此，作者认为论文报告的迭代时间（即吞吐量），包括了所有必要的操作，如数据加载和优化器步骤，是大型语言模型端到端训练时间的一个良好代理
    • 论文还发现迭代时间在整个训练过程中是一致的，因此对吞吐量进行采样可以准确衡量预期结果
  • 特别说明：表 5 中的结果没有使用任何数据并行
    • 数据并行会由于数据并行组之间所需的梯度 all-reduce 而引入一些开销
    • 但对于大型 Transformer 模型，这种开销并不大
      • 例如，如果论文将 530B 模型扩展到 8 路数据并行（2240 个 GPU），同时保持每个模型实例的批大小不变（即批大小也乘以数据并行大小）每次迭代的时间从 \( 37.83 \) 秒略微增加到 \( 39.15 \) 秒
      • 这导致 MFU 从 \( 56.0\% \) 下降到 \( 54.2\% \)，下降幅度不大
    • 论文注意到论文没有使用任何梯度 all-reduce 与反向传播的重叠，而高效的重叠几乎可以完全消除迭代时间的增加

Future Work

  • 作者未来计划通过解决由大批次 (Large Microbatches) 导致的内存碎片化问题以及由 PP (Pipeline Parallelism) 导致的非均匀内存分配问题，来进一步减少激活内存
  • 作者准备研究能够减轻流水线第一级 (First Stage of the Pipeline) 内存压力的方法

附录：论文补充材料

注：本文包含 AI 辅助创作

Paper Summary

  • 整体总结：
    • 本文是 Megatron-LM原始论文解读-第一篇
    • 论文的核心贡献：提出 张量并行（Tensor Parallelism） ，将 Transformer 的 Attention 头与 FFN 权重按列/按行切分，实现层内模型并行，首次在 GPU 集群上训练出 8.3 B 参数的 GPT-2 模型
    • 论文通过对现有 PyTorch Transformer 实现仅进行少量修改来实现模型并行，成功突破了传统的单 GPU 单模型训练的限制
    • 论文在 512 块 NVIDIA V100 GPU 上使用 8 路模型并行高效训练了高达 8.3B 参数的基于 Transformer 的模型，并在整个应用中实现了高达 15.1 PetaFLOPs 的持续性能
    • 论文表明，对于 BERT 模型，随着模型规模的增长，类 BERT 模型中 Layer Normalization 的位置对准确率的提升至关重要
    • 论文研究了模型规模对下游任务准确率的影响，在下游任务上取得了远超现有水平的结果，并在 WikiText103、LAMBADA 和 RACE 数据集上建立了新的最佳结果（SOTA）
    • 非常值得称赞的是：论文开源了论文的代码，以支持未来基于模型并行 Transformer 的研究工作
  • 背景 & 问题提出：
    • Transformer 模型推动了 NLP 技术发展，但内存限制，导致超大型模型的训练难度极大
  • 方法内容：
    • 论文提出了训练超大型 Transformer 模型的技术，并实现了一种简单、高效的层内模型并行（intra-layer model parallel）方法，该方法能够训练具有数十亿参数的 Transformer 模型
    • 本方法不需要新的编译器或库修改，与流水线模型并行（pipeline model parallelism）相互独立且互补，完全可以通过在原生 PyTorch 中插入少量通信操作来实现
  • 论文通过使用 512 个 GPU 训练出高达 8.3B 参数的 Transformer 模型验证了该方法的有效性
    • 与单个GPU基准（可持续 39 TeraFLOPs，达到峰值浮点运算的 30%）相比，论文在整个应用中实现了 15.1 PetaFLOPs 的性能，扩展效率为76%
    • \({39 \times 512 / 1000}{39} \approx 76%\)
  • 为了证明 LLM 能进一步推动技术发展，论文训练了一个与 GPT-2 类似的 8.3B 参数 Transformer 语言模型和一个与 BERT 类似的 3.9B 参数模型
    • 论文发现，对于类 BERT 模型，随着模型规模的增长， Layer Normalization的位置对性能提升至关重要
    • 使用 GPT-2 模型，论文在 WikiText103 数据集和 LAMBADA 数据集上取得了最佳结果
      • WikiText103 数据集：困惑度为 10.8（之前最佳困惑度为 15.8）
      • LAMBADA 数据集：准确率为 66.5%（之前最佳准确率为63.2%）
    • 使用 BERT 模型，论文在在 RACE 数据集上取得了最佳结果
      • RACE 数据集：准确率为 90.9%（之前最佳准确率为 89.4%）

Introduction and Discussion

  • NLP 正在快速发展，部分原因是可用计算资源和数据集规模的增加
    • 丰富的计算资源和数据使得通过无监督预训练（unsupervised pretraining）训练越来越大的语言模型成为可能（2018；2019）
  • 实证证据表明，更大的语言模型在文章补全、问答和自然语言推理等 NLP 任务中更有用（2019；2019）
    • 近期研究显示，在下游自然语言任务上对这些预训练语言模型进行微调（finetuning），可以取得最佳结果（2018；2017；2016；2019b；）
  • 模型规模的增长超出了现代处理器的内存限制，需要额外的内存管理技术，如激活检查点（activation checkpointing）（2016）
    • ADAM 等广泛使用的优化算法（optimization algorithms）需要为每个参数额外分配内存来存储动量（momentum）和其他优化器状态，这限制了可有效训练的模型规模
    • 几种模型并行（model parallelism）方法通过对模型进行分区来克服这一限制，使得权重及其相关的优化器状态无需同时驻留在处理器上，比如：
      • GPipe（2018）和 Mesh-Tensorflow（2018）提供了不同类型的模型并行框架
      • 但它们需要重写模型，并依赖仍在开发中的自定义编译器和框架
  • 论文使用层内模型并行（intra-layer model-parallelism）实现了一种简单高效的模型并行方法
  • 论文利用基于 Transformer 的语言模型的固有结构，实现了一种简单的模型并行方案，该方案能在 PyTorch 中高效训练，无需自定义 C++ 代码或编译器
    • 这种方法与 GPipe（2018）等方法所倡导的基于流水线的模型并行相互独立
  • 为了证明论文方法的可扩展性，论文建立了一个基准：在单个 NVIDIA V100 32GB GPU 上训练 1.2B 参数的模型，可持续 39 TeraFLOPs
    • 这相当于 DGX-2H 服务器中配置的单个 GPU 理论峰值浮点运算的 30%，因此是一个很强的基准
    • 通过 8路 模型并行在 512 个 GPU 上训练 8.3B 参数的模型，论文在整个应用中实现了高达 15.1 PetaFLOPs 的持续性能
      • 与单GPU情况相比，这一扩展效率为76%
    • 图1显示了更详细的扩展结果
  • 为了分析模型规模扩展对准确率的影响：
    • 论文训练了从左到右的 GPT-2（2019）语言模型和 BERT（2018）双向 Transformer，并在多个下游任务上对它们进行了评估
      • 论文发现，现有 BERT 架构会随着规模的增加而导致模型性能下降
      • 论文通过重新排列 Transformer 层中的 Layer Normalization 和残差连接（residual connection）克服了这一挑战，并表明通过这种修改，下游任务在开发集上的结果会随着模型规模的增加而单调提升
    • 论文的模型在 WikiText103 测试集、LAMBADA 的完形填空预测准确率（cloze-style prediction accuracy）和 阅读理解 RACE 数据集上取得了最佳结果
  • 总结来看：论文的贡献如下：
    • 通过对现有 PyTorch Transformer 实现仅进行少量有针对性的修改，实现了一种简单高效的模型并行方法
    • 对模型和数据并行技术进行了深入的实证分析，并证明了在 512 个 GPU 上的扩展效率高达 76%
    • 论文实验表名：对于类 BERT 模型，随着模型规模的增长， Layer Normalization 的位置对准确率的提升至关重要
    • 论文证明：模型规模的扩展会提高 GPT-2（研究到 8.3B 参数）和 BERT（研究到 3.9B 参数）模型的准确率
    • 展示了论文的模型在测试集上取得了最佳结果：WikiText103的困惑度（10.8 ppl）、LAMBADA 的准确率（66.5%）和 RACE 的准确率（90.9%）
    • 开源了论文的代码以及训练和评估流水线，地址为 github.com/NVIDIA/Megatron-LM

Background and Challenges

Neural Language Model Pretraining

  • 预训练语言模型已成为 NLP 研究人员工具包中不可或缺的一部分
    • 利用大型语料库预训练来学习稳健的语言神经表征是过去十年中一个活跃的研究领域
  • 预训练和迁移语言神经表征的早期例子表明，预训练的词嵌入表（word embedding tables）比从头学习的词嵌入表能更好地提升下游任务结果（2013；2014；2010）
  • 后来的研究通过学习和迁移能捕捉单词上下文表征（contextual representations）的神经模型推进了这一领域（2016；2018；2017；2019）
  • 近期的相关研究（2016；2018；2019b；）进一步基于这些思想，不仅迁移语言模型来提取上下文单词表征，还在下游任务上对语言模型进行端到端微调
  • 通过这些工作，技术发展已经从仅迁移词嵌入表发展到迁移整个数十亿参数的语言模型
  • 这种方法的发展使得需要能够高效大规模运行并满足日益增长的计算需求的硬件、系统技术和框架
  • 论文的工作旨在提供必要的工具，以推动这一趋势向前迈出又一步

Transformer Language Models and Multi-Head Attention

  • 当前 NLP 研究趋势是使用 Transformer 模型（2017），因为它们具有更高的准确率和计算效率
  • 原始 Transformer 结构被设计为机器翻译架构，它使用 Encoder 和 Decoder 两部分将输入序列转换为另一个输出序列
  • 近期利用 Transformer 进行语言建模的研究（如BERT（2018）和GPT-2（2019））根据其需求仅使用 Encoder 或 Decoder
  • 本研究探讨了 Decoder 架构 GPT-2 和 Encoder 架构 BERT
  • 图2显示了论文使用的模型示意图
    • 有关模型架构的详细描述，论文建议读者参考先前的研究（2017；2018；2019）
    • 特别地，GPT-2 和 BERT都使用 GeLU（2016）非线性激活函数和 Layer Normalization （2016），并 将其应用于多头注意力（multi-head attention）和前馈层（feed forward layers）的输入，而原始 Transformer （2017）使用 ReLU 非线性激活函数，并将 Layer Normalization 应用于输出

Data and Model Parallelism in Deep Learning

  • 将深度神经网络训练扩展到多个硬件加速器有两种核心范式：数据并行（data parallelism）（1990）和模型并行（model parallelism）
    • 数据并行是将训练小批量（minibatch）分配到多个工作节点（workers）
    • 模型并行是将模型的内存使用和计算分配到多个工作节点
  • 通过与可用工作节点数量成比例地增加批量大小（即弱扩展），可以观察到训练数据吞吐量接近线性扩展
    • 但大批量训练会给优化过程带来复杂性，可能导致准确率下降或收敛时间延长，从而抵消训练吞吐量增加带来的好处（2017）
    • 进一步的研究（2017；2017；2019）开发了缓解这些影响的技术，并缩短了大型神经网络的训练时间
    • 为了进一步扩展训练，相关研究（2016）将数据并行与激活检查点相结合 ：
      • 在反向传播（backward pass）中重新计算激活（activations），而不在前向传播（forward pass）中存储它们，以减少内存需求
  • 以上这些技术在处理问题规模上有一个根本限制：模型必须完全适合单个工作节点
    • 随着 BERT 和 GPT-2 等语言模型规模和复杂性的增加，神经网络已经接近现代硬件加速器的内存容量
    • 解决这一问题的一种方法是使用参数共享（parameter sharing）来减少模型的内存占用（2019），但这会限制模型的整体容量
    • 论文的方法是利用模型并行将模型分配到多个加速器上
      • 这不仅减轻了内存压力，还独立于微批量（microbatch）大小增加了并行性
  • 在模型并行中，还有两种进一步的范式：分层流水线并行（layer-wise pipeline parallelism）和更通用的分布式张量计算（distributed tensor computation）
    • 在流水线模型并行中，一组操作在一个设备上执行，然后将输出传递到流水线中的下一个设备，在那里执行另一组操作
      • 一些方法（2018；2018）将参数服务器（parameter server）（2014）与流水线并行结合使用，但这些方法存在不一致问题
      • TensorFlow 的 GPipe 框架（2018）通过使用同步梯度下降（synchronous gradient decent）克服了这一不一致问题
        • GPipe 需要额外的逻辑来处理这些通信和计算操作的高效流水线，并且会受到流水线气泡（pipeline bubbles）的影响而降低效率，或者需要修改优化器本身，从而影响准确率
    • 分布式张量计算是一种独立且更通用的方法，它将张量操作分配到多个设备上，以加速计算或增加模型规模
      • FlexFlow（2018）是一个协调这种并行计算的深度学习框架，它提供了一种选择最佳并行策略的方法
      • Mesh-TensorFlow（2018）引入了一种在 TensorFlow（2015）中指定通用分布式张量计算类别的语言
        • 用户在该语言中指定并行维度，然后使用适当的集合原语（collective primitives）编译生成的图
  • 论文利用与 Mesh-TensorFlow 中类似的见解，并利用 Transformer 注意力头（attention heads）计算中的并行性来并行化论文的 Transformer 模型
    • 论文没有为模型并行实现框架和编译器，而是仅对现有 PyTorch Transformer 实现进行了少量有针对性的修改
    • 论文的方法简单，不需要任何新的编译器或代码重写，并且可以通过插入一些简单的原语来完全实现，如下一节所述

Model Parallel Transformer s

  • 论文利用 Transformer 网络的结构，通过添加一些同步原语（synchronization primitives），实现了一种简单的模型并行方案

  • 如图2所示：一个 Transformer 层由一个自注意力块（self attention block）和一个两层的多层感知器（multi-layer perceptron, MLP）组成

    • 论文分别在这两个块中引入模型并行

  • 首先详细介绍 MLP 块

    • MLP 块的第一部分是一个通用矩阵乘法（general matrix multiplication, GEMM），后跟一个 GeLU 非线性激活函数：
      $$Y=GeLU(X A) \tag{1}$$
    • 并行化 GEMM 的一种方法是将权重矩阵 \(A\) 按行拆分，将输入 \(X\) 按列拆分，如下所示：
      $$X=\left[X_{1}, X_{2}\right], A=\left[\begin{array}{l}A_{1} \\ A_{2}\end{array}\right] \tag{2}$$
    • 这种划分将导致
      $$ Y=GeLU(X_{1} A_{1}+X_{2} A_{2}) $$
    • 由于 GeLU 是一个非线性函数
      $$ GeLU(X_{1} A_{1}+X_{2} A_{2}) \neq GeLU(X_{1} A_{1})+GeLU(X_{2} A_{2})$$
    • 因此这种方法需要在 GeLU 函数之前设置一个同步点

  • 另一种方法是将 \(A\) 按列拆分 \(A=[A_{1}, A_{2}]\)

    • 这种划分允许 GeLU 非线性激活函数独立应用于每个划分的 GEMM 输出：
      $$\left[Y_{1}, Y_{2}\right]=\left[GeLU\left(X A_{1}\right), GeLU\left(X A_{2}\right)\right]$$

    • 这一方法的优势在于消除了同步点

    • 论文以这种列并行的方式划分第一个 GEMM，并将第二个 GEMM 按行拆分，使其可以直接接收 GeLU 层的输出，而无需任何通信（如图3a所示）

    • 然后，第二个 GEMM 的输出在多个 GPU 之间进行归约（reduce），之后再传递到 dropout 层

    • 这种方法将 MLP 块中的两个 GEMM 都拆分到多个 GPU 上，并且在正向传播中只需要一个 \(g\) 算子，在反向传播中只需要一个全归约（all-reduce）（ \(f\) 算子）

      • 这两个算子是共轭的，可以用 PyTorch 中的几行代码实现

    • 例如， \(f\) 算子的实现如下：

      1
      2
      3
      4
      5
      6
      class f(torch.autograd.Function):
          def forward(ctx, x):
              return x
          def backward(ctx, gradient):
              all_reduce(gradient)
              return gradient

    • \(g\) 与 \(f\) 类似，只是在正向传播中执行全归约，在反向传播中执行恒等操作

  • 如图3b所示

    • 对于自注意力块，利用多头注意力操作中固有的并行性 ，将与 \(K\)、 \(Q\) 和 \(V\) 相关的 GEMM 以列并行的方式划分 ，这样每个注意力头对应的矩阵乘法都在单个 GPU 上本地完成
    • 这使得我们可以将每个注意力头的参数和计算负载分配到多个 GPU 上，并且不需要立即通信就能完成自注意力计算
    • 自注意力之后的输出线性层（output linear layer）的 GEMM 按行并行化，可以直接接收并行注意力层的输出，无需 GPU 之间的通信
    • 这种针对 MLP 层和自注意力层的方法融合了两组 GEMM，消除了它们之间的同步点，从而实现了更好的扩展性
    • 这使得论文能够在一个简单的 Transformer 层中，仅在正向传播和反向传播中各使用两次全归约来执行所有 GEMM（见图4）

  • Transformer 语言模型的 Output Embedding 维度为隐藏层大小（ \(H\) ）乘以词汇表大小（ \(v\) ）

    • 现代语言模型的词汇表大小通常在数万个 token 量级（例如，GPT-2 使用的词汇表大小为 50257），因此并行化 Output Embedding 的 GEMM 是有益的
    • 但在 Transformer 语言模型中， Output Embedding 层与 Input Embedding 层共享权重，这需要对两者都进行修改
    • 对于 Input Embedding 权重矩阵 \(E_{H \times v}\) 沿词汇表维度按列拆分：
      $$ E=[E_{1}, E_{2}]$$
      • 由于每个分区现在只包含一部分嵌入，因此在 Input Embedding 之后需要一个 \(g\) 算子
    • 对于 Output Embedding ：
      • 一种方法是执行下面的并行 GEMM 以获得 logits
        $$ [Y_{1}, Y_{2}]=[X E_{1}, X E_{2}] $$
        • 然后添加一个全收集（all-gather）操作 $$ Y = \text{all-gather}([Y_{1}, Y_{2}])$$
        • 再将结果传递给交叉熵损失函数（cross-entropy loss function）
        • 但这种情况下，全收集操作将传递 \(b \times s \times v\) 个元素（ \(b\) 是批大小， \(s\) 是序列长度），由于词汇表大小很大，这一数据量非常庞大
      • 为了减少通信量，论文将并行 GEMM 的输出 \([Y_{1}, Y_{2}]\) 与交叉熵损失融合，这将维度减少到 \(b \times s\)
      • 传递标量损失而不是 logits 大大减少了通信量，提高了论文模型并行方法的效率

  • 论文的模型并行方法在很大程度上可以概括为旨在减少通信并使 GPU 保持计算受限（compute bound）的技术

  • 论文选择在多个 GPU 上复制计算，而不是让一个 GPU 计算 dropout、 Layer Normalization 或残差连接的一部分，然后将结果广播到其他 GPU

  • 具体来说：

    • 论文在每个 GPU 上维护 Layer Normalization 参数的副本
    • 并对模型并行区域的输出执行 dropout 和 残差连接
    • 然后将其作为输入馈送到下一个模型并行区域
    • 问题：如何理解这里的每个 GPU 独立维持 Layer Normalization 参数的副本？

  • 为了优化模型，论文允许每个模型并行工作节点优化自己的一组参数

    • 由于所有值要么在 GPU 上本地存储，要么在 GPU 上复制，因此在这种形式中不需要通信更新后的参数值

  • 关于混合模型并行与数据并行以及随机数生成的处理，论文在附录B中提供了更多细节

  • 总之，如上所述，论文的方法易于实现，只需要在正向传播和反向传播中添加少量额外的全归约操作

    • 它不需要编译器，并且与 GPipe（2018）等方法所倡导的流水线模型并行相互独立且互补

Setup

  • 预训练语言理解模型是自然语言处理和语言理解中的核心任务
  • 语言建模有多种形式，论文重点关注 GPT-2（2019）（一种基于 Transformer 的从左到右生成式语言模型）和 BERT（2018）（一种基于掩码语言模型的双向 Transformer 模型）
  • 论文在以下部分解释这些模型的配置，并建议读者参考原始论文以获取更多细节

Training Dataset

  • 为了收集具有长期依赖关系的大型多样化训练集，论文聚合了几个最大的语言建模数据集
  • 论文创建的聚合数据集包括 Wikipedia（2018）、CC-Stories（2018）、RealNews（2019）和 OpenWebtext（2019）
  • 为了避免训练集数据泄露到下游任务中
    • 移除了 WikiText103 测试集（2016）中存在的 Wikipedia 文章
    • 还移除了 CC-Stories 语料库中因预处理 artifacts 引入的不必要换行
  • 对于 BERT 模型，论文在训练数据集中包含了 BooksCorpus（2015），但由于该数据集与 LAMBADA 任务重叠，因此在 GPT-2 训练中排除了该数据集
  • 论文合并了所有数据集，然后做了以下操作：
    • 从聚合数据集中过滤掉内容长度小于 128 个 token 的所有文档
    • （聚合数据集中可能存在重复的相似内容），论文使用局部敏感哈希（locality-sensitive hashing, LSH）对 Jaccard 相似度大于 0.7 的内容进行去重
  • 最终的聚合语料库包含 174 GB 的去重文本

Training Optimization and Hyperparameters

  • 为了高效训练论文的模型，论文利用混合精度训练（mixed precision training）和动态损失缩放（dynamic loss scaling），以充分利用 V100 的Tensor Core（2017；2018）
  • 使用简单的正态分布初始化权重 \(W\) ：
    $$ W \sim N(0,0.02)$$
  • 在残差层（residual layers）之前立即按下面的比例缩放权重：
    $$ \frac{1}{\sqrt{2 N} } $$
    • 其中 \(N\) 是由自注意力和 MLP 块组成的 Transformer 层数
  • 使用带权重衰减（weight decay）（2019）的 Adam 优化器（2014），其中 \(\lambda=0.01\)
  • 使用全局梯度范数裁剪（global gradient norm clipping）为 1.0，以提高大型模型训练的稳定性
  • dropout 率均设置为 0.1
  • 在每个 Transformer 层之后使用激活检查点（2016）
  • 对于 GPT-2 模型
    • 所有训练都使用 1024 个子词单元（subword units）的序列
    • 批大小为 512
    • 共训练 300K 次迭代
    • 初始学习率为 1.5e-4（3k 次迭代的预热期（warmup period））
    • 剩余的 297K 次迭代中遵循单周期余弦衰减（single cycle cosine decay）至最小学习率 1e-5
  • 对于 BERT 模型
    • 在很大程度上遵循（2019）中描述的训练过程
    • 论文使用原始 BERT 词典，词汇表大小为 30522
    • 按照（2019）的建议，用句子顺序预测（sentence order prediction）替换了下一句预测（next sentence prediction）头，并使用（2019）的全词 n-gram 掩码（whole word n-gram masking）
    • 批大小设置为 1024
    • 学习率为 1.0e-4
    • 在 10k 次迭代中预热，然后在 2M 次迭代中线性衰减
    • 其他训练参数与（2018）保持一致

Experiments

  • 论文所有的实验都使用了多达 32 台 DGX-2H 服务器（总共 512 块 Tesla V100 SXM3 32GB GPU）
  • 论文的基础设施针对多节点深度学习应用进行了优化，服务器内部的 GPU 之间通过 NVSwitch 实现 300 GB/秒 的带宽，服务器之间通过每台服务器配备的 8 个 InfiniBand 适配器实现 100 GB/秒 的互联带宽

Scaling Analysis

  • 为了测试论文实现方案的可扩展性，论文考虑了四组参数的 GPT-2 模型，详情见表1
  • 为了使自注意力层中的通用矩阵乘法（GEMM）大小保持一致，每个注意力头的隐藏层大小固定为 96，通过调整注意力头数和层数来获得参数规模从 1B 到 8B 的配置
    • 1.2B 参数的配置可以在单个 GPU 上运行，而 8B 参数的模型则需要 8 路模型并行（8 块 GPU）
    • 将原始词汇表大小从 50257 填充为 51200，原因如下：
      • 为了使 logit 层的通用矩阵乘法（GEMM）更高效，每个 GPU 上的词汇表大小最好是 128 的倍数
      • 论文的模型并行最多为 8 路，因此论文将词汇表填充为可被 \(128 \times 8 = 1024\) 整除的大小，最终填充后的词汇表大小为 51200
    • 论文研究了模型并行和模型+数据并行两种扩展性
      • 对于模型并行扩展性，所有配置都使用固定的批大小 8
      • 数据并行扩展对于训练许多 SOTA 模型是必要的，这些模型通常使用更大的全局批大小
      • 在模型+数据并行的情况下，论文所有实验的全局批大小都固定为 512，这对应于 64 路数据并行
Model and Data Parallelism
  • 本节展示模型并行和模型+数据并行情况下，相对于模型参数的弱扩展性
  • 弱扩展性通常通过扩展批大小来实现，但这种方法无法解决训练无法在单个 GPU 上运行的大型模型的问题，而且会导致大批量训练的收敛性下降
  • 论文不同，在这里使用弱扩展性来训练原本无法实现的更大模型
  • 所有扩展性数据的基准是表1中的第一种配置（1.2B 参数）在单个 GPU 上的运行情况
    • 这是一个很强的基准，因为它在整个训练过程中达到了 39 TeraFLOPs，相当于 DGX-2H 服务器中单个 GPU 理论峰值浮点运算的 30%
  • 图5显示了模型并行和模型+数据并行的扩展性数值：论文在两种设置下都观察到了优异的扩展性。例如，
    • 8.3B 参数的模型在 8 路（8 块 GPU）模型并行下实现了 77% 的线性扩展效率
    • 模型+数据并行由于需要额外的梯度通信，扩展性数值略有下降
    • 但即使是在 512 块 GPU 上运行的最大配置（8.3B 参数），相对于 1.2B 参数在单个 GPU 上的强基准配置，论文也实现了 74% 的线性扩展效率
    • 更多扩展性分析见附录D

Language Modeling Results Using GPT-2

  • 为了证明大型语言模型能进一步推动技术发展，论文考虑训练表2中列出的不同规模和配置的 GPT-2 模型

    • 355M 参数的模型与 BERT-Large 模型（2018）的规模和配置相当
    • 2.5B 参数的模型比之前最大的 GPT-2 模型更大，而据论文所知，8.3B 参数的模型比任何已训练的从左到右的 Transformer 语言模型都要大

  • 论文使用第4节中描述的流程来训练和评估论文的语言模型

  • 表2还列出了完成一个 epoch 所需的时间，一个 epoch 相当于 68507 次迭代

    • 例如， 8.3B 参数的模型在 512 块 GPU 上，每个 epoch 大约需要两天
    • 与表1中用于扩展性研究的配置相比：
      • 2.5B 参数的模型是相同的
      • 8.3B 参数的模型有 24 个注意力头（而不是 32 个）
      • 355M 参数的模型比之前的任何模型都小，但仍使用 64 块 GPU 进行训练，因此每个 epoch 的时间短得多

  • 图6显示了验证集困惑度（perplexity）随迭代次数的变化

    • 随着模型规模的增加，验证集困惑度降低， 8.3B 参数的模型达到 9.27 的验证集困惑度

  • 表3报告了训练后的模型在 LAMBADA 和 WikiText103 数据集上的零样本评估结果

    • 有关评估方法的更多细节，请参见附录E

  • 论文观察到，随着模型规模的增加，WikiText103 上的困惑度降低，LAMBADA 上的完形填空准确率（cloze accuracy）提高

    • 论文的 8.3B 参数模型在 WikiText103 测试集上实现了 10.81 的调整后困惑度，达到了最佳水平
    • 8.3B 参数的模型在 LAMBADA 任务上的准确率为 66.51%，同样超过了之前的完形填空准确率结果
    • 附录C 中包含了 8.3B 参数模型生成的样本

  • 最近，微软的研究人员与 NVIDIA 合作，使用 Megatron 训练了一个 17B 参数的 GPT-2 模型，称为 Turing-NLG（微软，2020），并表明随着模型规模的扩大，准确率进一步提高，凸显了更大模型的价值

  • 为了确保论文的训练数据中不包含任何测试集数据

    • 论文计算了测试集中的 8-grams 在训练集中出现的百分比，就像之前的工作（2019）所做的那样
    • WikiText103 测试集的重叠率最多为 10.8%，LAMBADA 测试集（2016）的重叠率最多为 1.4%
    • WikiText103 测试集与 WikiText103 训练集的重叠率已经达到 9.09%（2019）
    • 由于这些结果与之前的工作一致，论文有信心训练数据中没有无意中包含任何测试数据的文档

Bi-directional Transformer Results Using BERT

  • 在本节中，论文将论文的方法应用于类 BERT（BERT-style）的 Transformer 模型，并研究模型扩展对多个下游任务的影响
  • 先前的工作（2019）发现，当模型规模超过 336M 参数的 BERT-large 时，会出现意想不到的模型性能下降
    • 为了解决这种下降，该工作的作者（2019）引入了参数共享，并表明他们的模型比原始 BERT 模型的扩展性好得多
  • 论文进一步研究了这种现象，并通过实证证明，如图7所示，重新排列 Layer Normalization 和残差连接（residual connections）的顺序，对于使类 BERT 模型能够扩展到超过 BERT-Large 的规模至关重要
  • 图7中的架构（b）消除了使用原始BERT架构（a）时观察到的不稳定性，并且具有更低的训练损失
  • 据论文所知，论文是第一个报告这种修改能够训练更大 BERT 模型的研究
  • 使用图7（b）中的架构修改，论文考虑了表4中详细列出的三种情况
    • 336M 参数的模型与 BERT-large 的规模相同
    • 1.3B 参数的模型与之前被证明比 336M 参数的 BERT-large 模型结果更差的 BERT-xlarge 配置相同（2019）
    • 论文通过同时增加隐藏层大小和层数，将 BERT 模型进一步扩展到 3.9B 参数的情况
    • 在所有情况下，每个注意力头的隐藏层大小都固定为 64
    • 336M 和 1.3B 参数的模型训练了 2M 次迭代，而 3.9B 参数的模型训练了 1.5M 次迭代，目前仍在训练中
  • 在 3% 的保留集上，336M、1.3B和 3.9B 参数的模型分别实现了 1.58、1.30 和 1.16 的验证集困惑度 ，随着模型规模的增加而单调下降
  • 论文在多个下游任务上对训练后的模型进行微调，包括 GLUE 基准（2019）中的 MNLI 和 QQP，斯坦福问答数据集（2016；2018）中的 SQuAD 1.1 和 SQuAD 2.0，以及阅读理解 RACE 数据集（2017）
    • 微调时，论文遵循（2019b）中的相同流程
    • 论文首先对批大小和学习率进行超参数调优
    • 获得最佳值后，论文报告 5 个不同随机种子初始化的开发集结果的中值
    • 每个模型和任务使用的超参数见附录A
  • 表5显示了 MNLI、QQP、SQuAD 1.1 和 SQuAD 2.0 的开发集结果，以及 RACE 的测试集结果
    • 对于 RACE 的测试集结果，论文首先使用开发集找到在 5 个随机种子上给出中值分数的检查点，然后报告该检查点在测试集上的结果
    • 论文还报告了 SQuAD 开发集和 RACE 测试集的 5 路集成（ensemble）结果
  • 从表5中论文观察到：
    • （a）随着模型规模的增加，所有情况下的下游任务性能都有所提高；
    • （b）论文的 3.9B 参数模型在开发集上建立了比其他基于 BERT 的模型更好的最佳结果；
    • （c）论文的 3.9B 参数模型在 RACE 测试集上同时实现了单模型和集成模型的最佳结果

Future Work

  • 未来的工作有几个方向
    • 继续扩大预训练规模是一个很有前景的研究方向，这将进一步考验现有的深度学习硬件和软件
      • 为了实现这一点，需要提高优化器的效率和减少内存占用
      • 此外，训练超过16B 参数的模型将需要比 DGX-2H 服务器的16块GPU更多的内存
      • 对于此类模型，混合层内和层间模型并行以及节点间模型并行将更合适
    • 其他三个研究方向包括：
      • （a）预训练不同的模型系列（XLNet、T5）；
      • （b）评估大型模型在更多更困难和更多样化的下游任务（如生成式问答、摘要和对话）上的性能；
      • （c）使用知识蒸馏（knowledge distillation）从这些大型预训练教师模型（teacher models）中训练小型学生模型（student models）

A. BERT Finetuning Hyperparameters

  • 表6展示了每个模型和任务在微调过程中使用的超参数

B. Model Parallel Supplementary Material

在本节中，论文将详细介绍混合模型并行与数据并行以及随机数生成的处理方法

B.1 Hybrid Model and Data Parallelism

  • 模型并行与数据并行是相互独立的，因此我们可以同时使用这两种方法，在合理的时间内训练大型模型
  • 图8展示了用于混合模型并行与数据并行的 GPU 分组方式
  • 同一服务器内的两块或多块 GPU 组成模型并行组（例如图8中的 GPU 1 至 8），这些 GPU 共同分布运行一个模型实例
  • 其余的 GPU（可能位于同一服务器内，但更常见的是位于其他服务器中）运行额外的模型并行组
  • 每个模型并行组中处于相同位置的 GPU（例如图8 中的 GPU 1、9、……、505）组成数据并行组，因此数据并行组内的所有GPU都持有相同的模型参数
  • 在反向传播过程中，论文并行运行多个梯度全归约（gradient all-reduce）操作，以在每个不同的数据并行组内归约权重梯度
  • 所需的 GPU 总数是模型并行组数与数据并行组数的乘积
  • 例如，对于 8.3B 参数的模型，论文每个模型并行组使用 8 块 GPU，并采用 64 路数据并行，因此总共需要 512 块 GPU
  • 所有通信都是在 PyTorch 中通过 Python 调用 NCCL 实现的
  • 每个模型并行组内的 GPU 在组内所有 GPU 之间执行全归约操作
  • 对于数据并行，每个全归约操作都在每个模型并行组中的一块 GPU 之间进行

B.2 Model Parallel Random Number Generation

  • 诸如 dropout 等利用随机数生成的技术是现代深度学习训练中的重要组成部分
  • Transformer 在模型并行区域之外的残差连接之前以及模型并行区域内的自注意力块中都有 dropout 层
  • 由于一些 dropout 层位于模型并行区域内，而另一些则不在，因此论文需要谨慎处理随机数生成，以确保 dropout 正常工作
  • 为了在模型并行工作节点之间同步残差连接 dropout，论文在训练开始时使用相同的种子初始化随机数生成器
  • 这使得所有模型并行工作节点的 dropout 模式完全相同
  • 然而，模型并行区域内的 dropout 应该为每个工作节点生成不同的随机模式，以实现整个操作的随机性
    • 为了实现这一点，论文为模型并行区域内的 dropout 维护一个单独的随机数生成器
    • 每个模型并行工作节点的这个随机数生成器都有唯一的种子

C. Text Samples

  • 本节展示了 Megatron-LM 使用上下文提示生成的一些文本样本

D. Further Scaling Analysis

  • 在本节中，论文研究注意力头数对扩展性结果的影响
  • 论文还展示了 1.2B 参数模型的强扩展性结果

D.1 Attention Heads and Scaling

  • 本节研究注意力头数对模型并行扩展性的影响
  • 为此，论文考虑 8.3B 参数配置（采用8路模型并行），并将注意力头数从 16 调整到 32
  • 结果如表7所示
    • 随着注意力头数的增加 ，自注意力层内的一些通用矩阵乘法（GEMMs）会变小 ，自注意力 softmax 中的元素数量也会增加
      • 这导致扩展性效率略有下降
    • 未来的研究在设计平衡模型速度和准确率的大型 Transformer 模型时，应注意这个超参数

D.2 Strong Scaling

  • 论文的模型并行主要是为了能够训练规模超过单个 GPU 内存的模型，但它也可以在不增加批大小的情况下加速较小模型的训练
  • 为了衡量这种加速效果，论文训练了一个固定的 1.2B 参数模型
  • 论文使用每次迭代 8 个样本的固定批大小，并通过模型并行增加 GPU 数量
  • 结果如表8所示
    • 使用两块GPU可使训练速度提高 64%
    • 超过这个数量后，论文看到收益递减，因为每个 GPU 的计算量减少，内存带宽和通信开销开始占据主导地位

E. 使用 WikiText103 和 LAMBADA 评估语言模型（Evaluating Language Models Using WikiText103 and LAMBADA）

  • 在本节中，论文详细介绍了 WikiText103 数据集（2016）和 LAMBADA 数据集（2016）上完形填空式预测准确率的评估方法

E.1 Wikitext103 困惑度（Wikitext103 Perplexity）

  • WikiText103 困惑度是过去几年中被广泛研究的评估标准，困惑度是语料库平均交叉熵的指数（2011）。是语言模型的自然评估指标，因为语言模型表示整个句子或文本的概率分布：
    $$\text{PPL}=\exp \left(-\frac{1}{T_{o} } \sum_{t}^{T} \log P(t | 0: t-1)\right) \tag{4}$$
  • 为了计算式（4）中的困惑度
    • 论文根据子词词汇表对 WikiText103 测试语料库进行分词，并对每个 token \([0, T]\) 的交叉熵损失求和
    • 然后，论文通过原始分词方案中的 token 数量 \(T_{0}\) 对交叉熵损失进行归一化
  • WikiText103 测试语料库已经过预处理，采用单词级分词，先前的研究使用这种分词来计算困惑度
  • 为了在公平的环境下与先前的研究比较论文模型的困惑度，论文必须通过原始 token 数量 \(T_{0}\) 进行归一化，而不是输入到论文模型中的实际分词数据的 token 数量 \(T\)
  • 这种预分词还会在文本中引入训练数据中不存在的 artifacts
  • 为了缓解这种分布不匹配，论文首先使用可逆的去分词器对 WikiText103 测试数据集进行预处理，以去除与标点符号和空格相关的各种 artifacts
  • \(T_{0}\) 的值是在这种预处理之前计算的
  • 对于WikiText103的测试集， \(T_{o}=245566\) ， \(T=270329\)
  • 论文还必须对困惑度计算进行一项特定于 Transformer 的修改
    • 与基于 RNN 的语言模型不同， Transformer 在固定的窗口输入大小上运行
    • 因此，它们无法完全计算 \(P(t | 0: t-1)\) ，只能计算 \(P(t | t-w: t-1)\) ，其中 \(w\) 是论文的上下文大小：1024 个 token
    • 但为数据集中的每个 token 计算这个值的成本过高，因为论文必须对 \(w\) 大小的上下文进行大约 \(T\) 次评估
    • 为了高效地评估论文的模型，论文采用了一种称为重叠评估（overlapping evaluation）的折中方法，每次滑动窗口前进一定的重叠量 \(o\) ，只计算窗口最后 \(o\) 个 token 的交叉熵损失
    • 在论文的实验中，论文使用 32 的重叠量 \(o\) ，并以这种方式对所有滑动窗口计算损失

E.2 LAMBADA Cloze Accuracy

  • 处理长期上下文的能力对于 SOTA 语言模型至关重要，也是长文本生成和基于文档的问答等问题的必要前提
  • 像 LAMBADA 这样的完形填空式数据集旨在衡量模型在这些类型的长期上下文中运行和推理的能力
  • 完形填空式阅读理解使用单词 token 的上下文 \(x=x_{1: t}\) ，其中一个 token \(x_{j}\) 被掩码；
    • 模型的目标是正确预测缺失的第 \(j\) 个 token
  • 为了准确预测缺失的 token，模型需要深入理解周围的上下文以及在这种上下文中应如何使用语言
  • LAMBADA 使用完形填空式阅读理解来测试生成式从左到右语言模型，它构建了 45 个句子的示例，其中上下文 \(x_{t}\) 的最后一个单词被掩码
  • 论文的模型使用子词单元，因此在 LAMBADA 评估中，论文使用原始的、未处理的 LAMBADA 数据集，并要求论文的模型预测组成单词 token 的多个子词token
  • 论文使用教师强制（teacher forcing），只有当所有输出预测都正确时，才认为答案正确
    • 这种形式与原始的单词 token 预测任务等效

注：本文包含 AI 辅助创作

  • 参考链接：

Paper Summary

  • 整体说明：
    • 本文是 Megatron-LM原始论文解读-第二篇
    • 论文展示了如何将 PTD-P（节点间 PP、节点内 TP 和 DP ）组合起来，以在训练具有万亿参数的大模型时实现高聚合吞吐量
    • 论文首次实现了端到端训练能够在合理的时间内完成（万亿参数模型的估计时间约为 3 个月）
    • 论文通过分析讨论了与每种并行类型相关的各种权衡，以及在组合使用时需要仔细考虑它们之间的相互作用
    • 论文中的很多思想是与加速器无关：
      • 1）智能地划分模型训练计算图以在保持设备活跃的同时最小化通信量的思想
      • 2）通过算子融合和仔细的数据布局来最小化内存受限的内核数量
      • 3）其他领域特定的优化（例如，散射-聚集优化）
  • 背景 & 问题：
    • 高效地训练 LLM 模型具有挑战性，原因在于：
      • 1）GPU 内存容量有限，使得即使是在多 GPU 服务器上也无法容纳大型模型；
      • 2）所需的计算操作数量可能导致不切实际的长训练时间
  • 之前的工作：
    • 之前已经提出了 TP （Tensor Parallelism, TP）和 PP （Pipeline Parallelism, PP），但这些方法的简单使用会在数千个 GPU 上导致扩展性问题
  • 本文的核心方法贡献：
    • 论文展示了如何组合使用 TP 、PP 和 DP （Data Parallelism, DP），以扩展到数千个 GPU
    • 论文提出了一种新颖的交错流水线调度方法，在内存占用与现有方法相当的情况下，可以将吞吐量提高 10% 以上
  • 特别亮眼的表现：论文的方法使得论文能够在 3072 个 GPU 上以 502 petaFLOP/s 的速度对具有 1 万亿参数的模型执行训练迭代（每 GPU 吞吐量达到理论峰值的 52%）
    • 注意：这是在 2021 年实现的，而本文的方法已经是家喻户晓

Introduction and Discussion

  • NLP 中基于 Transformer 的语言模型近年来推动了快速进展，因为大规模计算变得更加可用且数据集变得更大
  • 最近的工作 (2020) 表明，大语言模型是有效的零样本或少样本学习者，在许多 NLP 任务和数据集上具有高准确率
  • 这些大语言模型有许多令人兴奋的下游应用，例如客户反馈摘要、自动对话生成、语义搜索和代码自动补全 (2021;)
  • SOTA NLP 模型中的参数数量呈指数级增长（图 1）
  • 训练此类模型具有挑战性，原因有二：
    • (a) 即使使用最大的 GPU（NVIDIA 最近发布了 80GB A100 卡），也无法将这些模型的参数容纳在其主内存中；
    • (b) 即使论文能够将模型装入单个 GPU（例如，通过在主机和设备内存之间交换参数 (2021)），所需的大量计算操作可能导致不切实际的长训练时间（例如，使用单个 NVIDIA V100 GPU 训练具有 175B 参数的 GPT-3 (2020) 将需要大约 288 年）
  • 所以需要并行化， DP 扩展通常效果良好，但受到两个限制：
    • a）超过某个点后，每 GPU 批大小变得太小，降低了 GPU 利用率并增加了通信成本；
    • b）可使用的最大设备数量等于批大小，限制了可用于训练的加速器数量
  • 为了应对这两个挑战，已经提出了各种模型并行技术
    • 一些工作 (2019; 2020) 展示了如何使用张量（层内）模型并行（即将每个 Transformer 层内的矩阵乘法拆分到多个 GPU 上）来克服这些限制
    • 尽管这种方法对于在 NVIDIA DGX A100 服务器（配备 8 个 80GB A100 GPU）上训练大小达 20B 参数的模型效果良好，但对于更大的模型则会失效
    • 更大的模型需要拆分到多个多 GPU 服务器上，这会导致两个问题：
      • (a) TP 所需的全局归约（All-Reduce）通信需要通过服务器间链路进行，这些链路比多 GPU 服务器内可用的高带宽 NVLink (2020) 慢；
      • (b) 高度的模型并行可能会产生小的矩阵乘法（GEMM），可能降低 GPU 利用率
  • PP (2019; 2020; 2021) 是另一种支持训练大模型的技术，它将模型的层分布到多个 GPU 上
  • 一个批次被分割成更小的微批次（microbatch），执行过程在这些微批次之间进行流水线化
    • 层可以以各种方式分配（assigned）给工作节点，并且可以使用各种输入的前向传播和后向传播调度策略
    • 层分配（layer assignment）和调度策略会导致不同的性能权衡
  • 无论采用何种调度，为了保持严格的优化器语义，优化器步骤需要在设备间同步，这导致在每个批次结束时进行一次 流水线刷新（pipeline flush） ，此时允许微批次完成执行（并且不注入新的微批次）
    • flushing the pipeline 可能会花费多达 50% 的时间（取决于注入流水线的微批次数量）
    • 微批次数量与流水线规模（阶段数）的比率越大，花费在流水线刷新上的时间就越少
    • 为了实现高效率，通常需要更大的批大小
    • 在这项工作中，论文还引入了一种新的流水线调度方案，该方案提高了小批大小下的效率
  • 这些技术可以组合使用，但组合这些技术会导致 non-trivial 的相互作用，需要仔细推理才能获得良好的性能
  • 在论文中，论文解决了以下问题：
    • 在给定批大小并保持严格优化器语义的前提下，应如何组合并行技术以最大化大模型的训练吞吐量？
  • 具体来说，论文展示了如何组合 PP 、TP 和 DP，论文称之为 PTD-P（pipeline, tensor, and data parallelism） 的技术，以在数千个 GPU 上以良好的计算性能（达到设备峰值吞吐量的 52%）训练大语言模型
  • 论文的方法利用跨多 GPU 服务器的 PP 、多 GPU 服务器内的 TP 以及 DP 的组合，在实际中训练具有万亿参数的模型，并在具有服务器内和服务器间 GPU 高速链路的优化集群环境中实现优雅的扩展
  • 给定更多的训练资源，也可以使用类似的思想来训练更大的模型
  • 在论文的实验中，论文在一个具有万亿参数的 GPT 模型 (2020) 上，使用混合精度，展示了接近线性的扩展到 3072 个 A100 GPU，实现了每 GPU 163 teraFLOP/s 的端到端训练吞吐量（包括通信、数据处理和优化），以及 502 petaFLOP/s 的总吞吐量
    • 这个吞吐量使得实际的训练时间成为可能：论文估计该模型的端到端训练大约需要 3 个月
  • 作者相信这是针对该规模模型实现的最快训练吞吐量：过去的系统 (2019; 2020) 无法训练如此大的模型，因为它们没有结合 PP 和 TP
  • 论文还将论文的方法与 ZeRO (2019) 进行了比较，发现由于跨节点通信更少，论文的方法对于 175B 和 530B 参数的模型比 ZeRO-3 性能高出 70%
    • 这些模型太大，无法容纳在单个多 GPU 服务器上
  • 在大规模下实现这种吞吐量需要在多个方面进行创新和精心的工程：
    • 高效的内核实现使得大部分计算是计算受限而非内存受限的
    • 在设备间智能划分计算图以减少通过网络链路发送的字节数同时限制设备空闲时间
    • 领域特定的通信优化以及快速的硬件（ SOTA GPU 以及服务器内和服务器间的高速链路）
  • 作者希望论文开源的软件（可在 github.com/nvidia/megatron-lm 获取）将使其他团队能够高效地大规模训练大型 NLP 模型
  • 论文通过实验和尽可能的分析，研究了影响吞吐量的各个组件之间的相互作用
  • 基于这些研究，论文提供以下关于如何配置分布式训练的指导原则 ：
    • 不同形式的并行化以 non-trivial 的方式相互作用：
      • 并行化策略会影响通信量 和 内核执行的计算效率，还会影响工作节点上的气泡时间（因流水线刷新（气泡）所花费的空闲时间）
        • 例如，在论文的实验中，论文发现，即使服务器间有高带宽网络链路，次优的 TP 和 PP 组合也可能导致吞吐量降低多达 2 倍；
        • TP 在多 GPU 服务器内是有效的 ，但对于更大的模型必须使用 PP
    • 用于 PP 的调度方案会影响通信量、流水线气泡大小以及用于存储激活值的内存
      • 论文提出了一种新颖的交错调度方案，与先前提出的调度方案 (2019; 2020) 相比，在内存占用相当的情况下，可以将吞吐量提高多达 10%
    • 微批次大小等超参数的值会影响内存占用、工作节点上执行的内核的算术效率以及流水线气泡大小
      • 在论文的实验中，微批次大小的最优值是问题相关的，并且可以将吞吐量提高 15%
    • 在大规模下，分布式训练是通信密集型的
      • 当在 3072 个 GPU 上训练一个万亿参数模型时，论文的实现用于 PP 通信的有效二分带宽为 892 GB/s，用于 DP 通信的有效二分带宽为 13 TB/s
      • 使用较慢的节点间互连或通信更密集的划分会阻碍扩展性能
  • 论文不会自动探索并行策略的搜索空间（例如 FlexFlow (2018), PipeDream (2019), Tarnawski 等 (2020), 和 DAPPLE (2021)），而是提出论文在实践中发现效果很好的启发式方法（在 章节3 中）

Models of Parallelism

  • 论文关注单个 GPU 无法装下的大型模型进行高效训练的并行技术
  • 论文将 PP (pipeline model parallelism) 和 TP (tensor model parallelism)（组合如图 2 所示）与 DP (data parallelism) 结合起来（论文将其简称为 PTD-P）

DP (Data Parallelism)

  • 使用 DP (2021) 时，每个工作节点 (worker) 都拥有完整模型的一个副本，输入数据集被分片，工作节点定期聚合它们的梯度，以确保所有权重版本一致
  • 对于无法放在单个工作节点上的大型模型，可以在较小的模型分片上使用 DP

PP (Pipeline Model Parallelism)

  • 使用 PP 时，模型的层被分片到多个设备上
  • 当用于具有重复相同 Transformer 块的模型时，可以为每个设备分配相同数量的 Transformer 层
  • 论文不考虑更不对称的模型架构 ，因为将层分配给流水线阶段更为困难；论文将这个问题留给相关工作 (2021;) 来解决
  • 一个批次 (batch) 被分割成更小的微批次 (microbatch)；然后跨微批次进行流水线执行
    • 流水线方案需要确保输入在前向传播和后向传播中看到一致的权重版本，以获得明确定义的同步权重更新语义
  • 具体来说，简单的流水线可能导致一个输入在后向传播中看到在前向传播中未见的权重更新
    • 为了精确保持严格的优化器语义，论文引入了周期性的流水线刷新 (pipeline flush)，以便跨设备同步优化器步骤
    • 在每个批次的开始和结束时，设备处于空闲状态，论文将此空闲时间称为流水线气泡 (pipeline bubble)，并希望使其尽可能小
  • 诸如 PipeMare、PipeDream 和 PipeDream-2BW (2019; 2020; 2021) 之类的异步和有界陈旧度 (bounded-staleness) 方法完全取消了刷新，但放宽了权重更新语义
    • 论文将对此类方案的考虑留待未来工作
  • 有多种可能的方式来跨设备调度前向和后向微批次；每种方法在流水线气泡大小、通信量和内存占用之间提供了不同的权衡
  • 论文将在本节讨论两种这样的方法
Default Schedule
  • GPipe 提出了一种调度方案，即首先执行一个批次中所有微批次的前向传播（如图 3 所示），然后是所有微批次的后向传播
  • 我们可以量化 GPipe 的流水线气泡大小 (\(t_{pb}\))
    • 用一个批次中的微批次数量记为 \(m\)，流水线阶段数（用于 PP 的设备数量）记为 \(p\)，每次迭代的理想时间为 \(t_{id}\)（假设完美或理想缩放），执行单个微批次的前向和后向传播的时间记为 \(t_{f}\) 和 \(t_{b}\)
    • 在此调度中，流水线气泡包括批次开始时的 \(p-1\) 次前向传播和批次结束时的 \(p-1\) 次后向传播
      • 理解：图 3 中，对于 PP=4 的场景，即每个设备 都要有 3 个完整的前向和后向传播气泡
    • 花费在流水线气泡中的总时间为：
      $$t_{pb}=(p-1)\cdot(t_{f}+t_{b})$$
    • 该批次的理想处理时间为：
      $$t_{id}=m\cdot(t_{f}+t_{b})$$
    • 因此，理想计算时间中花费在流水线气泡中的比例为：
      $$\text{Bubble time fraction (pipeline bubble size)} =\frac{t_{pb} }{t_{id} }=\frac{p-1}{m}.$$
      • 理解：这里分母上算少了，分母上应该再加上分子才对
  • 为了使气泡时间比例变小，一般需要 \(m\gg p\)，但对于如此大的 \(m\)，这种方法具有很高的内存占用，因为它需要在一次训练迭代的整个生命周期内，将所有 \(m\) 个微批次的暂存中间激活值（或者在使用激活重计算时，仅为每个流水线阶段的输入激活值）保存在内存中
  • 论文使用 PipeDream-Flush 调度 (2021)
    • 论文首先进入一个预热阶段，工作节点执行不同数量的前向传播，如图 4（顶部）所示
      • 此调度将进行中的微批次（其后向传播尚未完成且需要维护激活值的微批次）的数量限制为流水线的深度，而不是一个批次中的微批次数量
    • 预热阶段之后，每个工作节点进入稳定状态，每个工作节点执行一次前向传播，然后执行一次后向传播（简称 1F1B）
    • 最后，在批次结束时，论文完成所有剩余进行中微批次的后向传播
    • 这种新调度花费在气泡中的时间是相同的，但对于 PipeDream-Flush 调度，未完成的前向传播数量最多为流水线阶段数
    • 因此，此调度需要为 \(p\) 个或更少的微批次暂存激活值（相比之下，GPipe 调度需要为 \(m\) 个微批次暂存）
    • 因此，当 \(m\gg p\) 时，PipeDream-Flush 比 GPipe 内存效率高得多
Schedule with Interleaved Stages
  • 为了减小流水线气泡的大小，每个设备可以执行多个层子集（称为模型块, model chunk）的计算，而不是单个连续的层集
  • 例如，如果之前每个设备有 4 层（即设备 1 有第 1-4 层，设备 2 有第 5-8 层，依此类推），我们可以让每个设备执行两个模型块（每个有 2 层）的计算，即设备 1 有第 1, 2, 9, 10 层；设备 2 有第 3, 4, 11, 12 层；依此类推
    • 使用这种方案，流水线中的每个设备被分配多个流水线阶段（每个流水线阶段的计算量比之前少）
    • 像之前一样，我们可以使用此调度方案的“全部前向，全部后向”版本，但这具有很高的内存占用（与 \(m\) 成正比）
  • 论文开发了一种交错调度 (interleaved schedule) ，它适配了之前的内存高效 1F1B 调度
    • 这种新调度如图 4 所示，并要求一个批次中的微批次数量是 PP 度（流水线中的设备数量）的整数倍
    • 例如，对于 4 个设备，一个批次中的微批次数量必须是 4 的倍数
  • 如图 4 所示，相同批次大小的流水线刷新在新调度中发生得更早
    • 如果每个设备有 \(v\) 个阶段（或模型块），那么每个阶段或块的微批次的前向和后向时间现在将是 \(t_{f}/v\) 和 \(t_{b}/v\)
  • 因此，流水线气泡时间减少到
    $$t_{pb}^{\text{int.} }=\frac{(p-1)\cdot(t_{f}+t_{b})}{v}$$
    • 气泡时间比例则为：
      $$\text{Bubble time fraction (pipeline bubble size)} =\frac{t_{pb}^{\text{int.} } }{t_{id} }=\frac{1}{v}\cdot\frac{p-1}{m}.$$
  • 这意味着新调度将气泡时间减少了 \(v\) 倍
    • 但这种减小的流水线气泡大小并非没有代价：此调度需要额外的通信
    • 定量地说，通信量也增加了 \(v\) 倍
  • 在下一节中，论文将讨论如何利用多 GPU 服务器（例如 DGX A100 节点）中的 8 个 InfiniBand 网络卡来减少这种额外通信的影响

TP (Tensor Model Parallelism)

  • 使用 TP 时，模型的各个层在多个设备上进行划分
  • 在论文中，论文使用 Megatron (2019) 针对语言模型基础 Transformer 层所使用的特定划分策略
  • 也可以将类似的思想应用于其他类型的模型，如 CNN
  • 论文简要概述此策略，如图 5 所示
  • 一个 Transformer 层由一个自注意力块 (self-attention block) 和一个两层的多层感知机 (MLP) 组成
  • MLP 块：MLP 块包含两个 GEMM（通用矩阵乘法，General Matrix Multiply）和一个 GeLU 非线性：
    $$Y=\text{GeLU}(XA).\quad Z=\text{Dropout}(YB).$$
    • 第一个权重矩阵 \(A\) 按列划分，即 \(A=[A_{1}, A_{2}]\)
      • 这种划分允许将 GeLU 非线性独立应用于每个划分后 GEMM 的输出：
        $$[Y_{1},Y_{2}]=\left[\text{GeLU}(XA_{1}),\text{GeLU}(XA_{2})\right].$$
      • 这样做的好处是它消除了同步的需要（如果 \(A\) 沿其行划分则需要同步，因为 GeLU 是非线性的）
    • 第二个权重矩阵 \(B\) 的按行划分，以消除 GEMM 之间的任何通信需要（如图 4(a) 所示），如下所示：
      $$B=\begin{bmatrix}B_{1}\ B_{2}\end{bmatrix}, Y=[Y_{1},Y_{2}].$$
    • 第二个 GEMM 的输出然后在 Dropout 层之前跨 GPU 进行归约 (reduce)
  • 自注意力块：利用多头注意力操作中固有的并行来划分自注意力块（如图 4(b) 所示）
    • 键 (\(K\))、查询 (\(Q\)) 和值 (\(V\)) 矩阵可以以列并行 (column-parallel) 的方式进行划分（每个设备上负责不同的注意力头，并行计算）
    • 输出线性层然后可以直接在注意力操作的划分输出上运行（权重矩阵跨行划分，不需要 GPU 同步通信）
  • TP 将 MLP 和自注意力块中的 GEMM 拆分到多个 GPU 上，同时在前向传播中仅需要两次全归约 (All-Reduce) 操作（\(g\) 运算符），在后向传播中需要两次全归约操作（\(f\) 运算符）
  • 论文用几行代码实现了 \(f\) 和 \(g\)

Performance Analysis of Parallelization Configurations

  • 考虑将 PP (pipeline parallelism) 和 TP (tensor model parallelism) 与 DP (data parallelism) 结合使用的性能影响
  • 给定固定的 GPU 预算和批次大小 (batch size)，可以使用 PTD-P 中不同维度的并行类型来训练模型；
  • 每个维度都在内存占用 (memory footprint)、设备利用率 (device utilization) 和通信量之间进行权衡
  • 论文将在本节剩余部分讨论这些权衡，并在 5.4 节展示实证结果
    • 论文在分析流水线气泡 (pipeline bubble) 大小时会提供分析模型
    • 论文定性地描述通信时间的行为，并提供通信量的成本模型；但是，论文不提供通信时间的直接成本模型，因为对于层次化网络拓扑（同一服务器上 GPU 之间的互连带宽高于服务器之间的互连带宽）来说，通信时间更难建模
  • 据论文所知，这是首个分析这些并行化维度性能交互作用的工作

Notation

  • 论文在本节中使用以下符号表示：
    • \( (p, t, d) \)：并行化维度
      • \( p \) 表示 PP 大小 (pipeline-model-parallel size)，\( t \) 表示 TP 大小 (tensor-model-parallel size)，\( d \) 表示 DP 大小 (data-parallel size)
    • \( n \)：GPU 数量，论文要求 \( p \cdot t \cdot d = n \)
    • \( B \)：全局批次大小 (global batch size)（作为输入提供）
    • \( b \)：微批次大小 (microbatch size)
    • \( m = \frac{1}{b} \cdot \frac{B}{d} \)：每个流水线 (per pipeline) 中一个批次内的微批次数量

TP 和 PP（Tensor and Pipeline Model Parallelism）

  • TP 和 PP 均可以用于在多个 GPU 对模型的参数进行分区
  • 正如前文所述，使用带有周期性刷新（periodic flushes）的 PP 会产生下面大小的流水线气泡（pipeline bubble）：
    $$\frac{p-1}{m}$$
  • 假设 DP 规模 \(d=1\)（ DP 大小，data-parallel size），那么此时总 GPU 数量满足
    $$ t \cdot p = n; \quad \text{s.t.} \ d=1 $$
  • 基于 TP 规模 \(t\) 的流水线气泡大小可表示为：
    $$
    \frac{p-1}{m} = \frac{n/t - 1}{m}
    $$
    • 说明：在固定批次大小 \(B\)、微批次大小 \(b\) 以及固定 \(d\)（此时 \(m = \frac{B}{b \cdot d}\) 也保持固定）的前提下，随着 TP 规模 \(t\) 的增大，流水线气泡会逐渐减小
  • 不同 GPU 之间的通信量同样会受到 PP 规模 \(p\) 和 TP 规模 \(t\) 的影响
    • PP 的特点是点到点通信（point-to-point communication）成本更低；TP 则需要使用 All-Reduce communication
      • TP：正向传播（forward pass）需执行两次归约操作；反向传播（backward pass）也需执行两次归约操作（详见第2.3节）
    • 在 PP 中，对于每个微批次，每对 PP 相邻设备之间需要执行的PP 通信总量为（一次 点对点通信）：
      $$ bsh $$
      • 其中 \(s\) 代表序列长度（sequence length），\(h\) 代表隐藏层大小（hidden size）
      • 正向传播和反向传播各需要一次点对点通信，所以总共是 $$ 2bsh $$
    • 在 TP 中，对于每个层，总大小为 \(bsh\) 的张量需在 \(t\) 个模型副本（model replicas）间，正向传播和反向传播各两次通信，这使得每个设备、每个微批次、每个层的 TP 通信总量为
      $$8bsh \cdot \frac{t-1}{t}$$
      • 理解：TP 参数通信量估计详情
        • MLP 正向和反向各需要一次 TP 通信（All-Reduce），一次 All-Reduce 通信是 \(2\Phi_\text{TP} = 2bsh\)（使用 Ring All-Reduce，一次通信使用 \(2\Phi\frac{t-1}{t}\) 的通信量），单层模型上， MLP TP 累计通讯量是：
          $$ 4\Phi\frac{t-1}{t} = 4bsh\frac{t-1}{t} $$
        • Attention 部分正向和反向也各需要一次 TP 通信（All-Reduce），单层模型上，Attention TP 累计通讯量是：
          $$ 4\Phi\frac{t-1}{t} = 4bsh\frac{t-1}{t} $$
        • 注：上述通信量评估详情见：图解大模型训练之：张量模型并行(TP)，Megatron-LM - 猛猿的文章 - 知乎
    • 通常每个设备会负责多个层的计算（TP 中，每个层都需要分别与其他设备交互），因此对于每个设备、每个微批次， TP 的总通信量可表示为
      $$l^{stage} \cdot (8bsh \cdot \frac{t-1}{t})$$
      • 其中 \(l^{stage}\) 代表一个流水线阶段（pipeline stage）中包含的层数，即一个设备负责的层数
  • 由此可见，TP 会增加设备间的通信量
    • 因此，当 TP 规模 \(t\) 大于单个节点（node）中的 GPU 数量时，跨节点链路（inter-node links）的传输速度较慢，此时在这类链路上执行 TP 的开销会变得难以承受
    • 这一结果已在第5.4节的实验中得到验证
  • Takeaway #1 ：TP 不跨 节点
    • 对于使用每个节点 \(g\) 个 GPU 的服务器时， TP 的规模通常应不超过 \(g\)；若要在多台服务器间扩展以训练更大的模型，则可使用 PP
    • 理解：这里指的是 单台机器的 GPU 数为 \(g\)，TP 规模 \(t\) 大于 \(g\) 时会导致跨节点的 TP 通信，开销难以接受

DP Data and Model Parallelism

  • 注：这里的模型并行（Model Parallelism）包括了 PP（Pipeline Model Parallelism） 和 TP（Tensor Model Parallelism）
  • 论文还需要考虑 DP （data parallelism）与两种模型并行（ TP 和 PP ）之间的相互作用
    • 为简化分析，本节将分别对这些相互作用进行讨论
PP （Pipeline Model Parallelism）
  • 假设 TP 规模 \(t=1\)（tensor-model-parallel size），则每个流水线的微批次数量为
    $$ m = \frac{B}{d \cdot b} = \frac{b’}{d}$$
    • 其中：
      • \( b \) 是微批次大小 (microbatch size)
      • \(b’ = \frac{B}{b}\) 表示每个 Global Step 的 微批次数量；
    • 注：每个流水线的微批次数量 \(m\) 与 DP 大小 \(d\) 成反比
  • 在总 GPU 数量为 \(n\) 的情况下，流水线阶段的数量为
    $$ p = \frac{n}{(t \cdot d)} = \frac{n}{d} $$
  • 此时流水线气泡大小可表示为：
    $$
    \frac{p-1}{m} = \frac{n/d - 1}{b’/d} = \frac{n - d}{b’}
    $$
    • 随着 DP 规模 \(d\) 的增大，\(n - d\) 会逐渐减小，因此流水线气泡也会随之变小
  • 图6展示了在不同 \(d\)、\(n\) 和 \(b’\)（批次大小与微批次大小的比值，\(b’ = \frac{B}{b}\)）下，流水线气泡大小的变化趋势
    • 需要注意的是，并非所有模型都能将 \(d\) 增大到 \(n\)，因为部分模型的完整训练内存占用（full training memory footprint）可能超过单个加速器（accelerator）的内存容量
  • 若 DP 所需的归约通信量未随 \(d\) 的增大而大幅增加，那么整体吞吐量（throughput）会随之提升
    • 这一假设是成立的，因为基于环形实现（ring-based implementation， 即 Ring All-Reduce）的通信时间与 \(\frac{d-1}{d} = 1 - \frac{1}{d}\) 呈正相关
  • 论文还可以分析批次大小 \(B\) 增大带来的影响
    • 在特定的并行配置下，随着批次大小 \(B\) 的增加，\(b’ = \frac{B}{b}\) 会增大，进而导致流水线气泡大小 \(\frac{n - d}{b’}\) 减小，最终使吞吐量提升
    • 此外，DP 所需的归约通信频率会随 \(B\) 的增大而降低，这也会进一步提升吞吐量
      • 理解：这里是指相同的数据量和 epoch 下，\(B\) 越大，需要更新的总步数越少？
DP Data and Tensor Model Parallelism
  • TP 中，每个微批次都需要执行归约通信，而跨多 GPU 服务器执行此类通信的成本较高
  • DP 仅需在每个批次执行一次高成本的归约通信
  • 而且，在 TP 中，每个模型并行进程（model-parallel rank）仅负责模型每层计算中的一部分；
    • 若层的规模不够大，现代 GPU 执行这些子矩阵（sub-matrix）计算时可能无法达到峰值效率（peak efficiency）
  • Takeaway #2 ：
    • 在同时使用 DP 和模型并行时，应将模型并行的总规模设为
      $$M = t \cdot p$$
      • 其中 \(t\) 为 TP 规模，\(p\) 为 PP 规模，以确保模型的参数和中间元数据（intermediate metadata）能够放入 GPU 内存；
      • 问题：模型并行的总规模本来就是：\(M = t \cdot p\) 吧，还需要什么特殊处理吗？
    • 而 DP 则可用于扩展训练规模，以适配更多的 GPU

Microbatch Size

  • 微批次大小 \( b \) 的选择也会影响模型训练吞吐量
  • 例如，论文在图 7 中看到，在单个 GPU 上，使用较大的微批次大小，每 GPU 吞吐量最多可提高 1.3 倍
  • 本节回答：确定给定并行配置 \( (p, t, d) \) 和批次大小 \( B \) 时，最优微批次大小 \( b \) 是多少呢？
  • 无论微批次大小如何， DP 通信量将是相同的
  • 给定函数 \( t_{f}(b) \) 和 \( t_{b}(b) \) 它们将微批次大小映射到单个微批次的前向和后向计算时间，则处理一个微批次的总计算时间（忽略通信成本）为（如前所述，定义 \( b’ \) 为 \( \frac{B}{d} \)）：
    $$ \left(\frac{b’}{b} + p - 1\right) \cdot \left(t_{f}(b) + t_{b}(b)\right). $$
    • 因此，微批次大小既影响操作的算术强度 (arithmetic intensity)，也影响流水线气泡大小（通过影响 \( m \)）
      • 理解：随着 microbatch size \(b\) 增大，气泡占比变小，处理一个微批次的总时间减少，但总微批次数量增加
    • 注：进一步分析总的时间为：
      $$ b \cdot \left(\frac{b’}{b} + p - 1\right) \cdot \left(t_{f}(b) + t_{b}(b)\right). $$
      • 显然，上述式子存在一个处于最大值和最小值中间的最优点（类似二次函数）
  • 图 8 显示了一个具有十亿参数且 \( (p, t) = (8, 8) \) 的 GPT 模型的估计吞吐量（使用方程 (1) 估计处理时间）
    • 对于两种批次大小，最优的 \( b \) 都是 4
  • Takeaway #3：
    • 最优微批次大小 \( b \) 取决于模型的吞吐量和内存占用特性，以及流水线深度 \( p \)、 DP 大小 \( d \) 和批次大小 \( B \)

激活重计算，Activation Recomputation

  • 激活重计算 (activation recomputation) (2016; 2000; 2019; 2020) 是一种可选技术，通过 仅存储给定流水线阶段的输入激活，而不是存储整个中间激活集，后者要大得多，来权衡增加执行的计算操作数量以换取更少的内存占用
    • 注：成本是后向传递之前需要按需再次运行前向传递
  • 为了在可接受的低内存占用下训练合理的大模型，需要使用激活重计算
  • 之前的工作如 PipeDream-2BW (2021) 已经研究了激活重计算的性能影响
    • 激活检查点 (activation checkpoint) 的数量不影响吞吐量，但影响内存占用
    • 设 \( A^{\text{input} } \) 为一层的输入激活大小，\( A^{\text{intermediate} } \) 为每层的中间激活大小
    • 如果一个模型阶段有 \( l \) 层，并且有 \( c \) 个检查点，则总内存占用将为
      $$ c \cdot A^{\text{input} } + \frac{l}{c} \cdot A^{\text{intermediate} } $$
      • \(\frac{l}{c} \cdot A^{\text{intermediate} }\) 为每个时刻在使用的激活大小
    • 当 \( c = \sqrt{l \cdot (A^{\text{intermediate} }/A^{\text{input} })} \) 时，该函数取得最小值
    • 在实践中，论文通过经验测量 \( A^{\text{intermediate} } \)
    • 在大多数情况下，每 1 或 2 个 Transformer 层设置一个检查点是最优的
  • 其他技术，如激活分区 (activation partitioning) (2020)，也可以与 TP 结合使用，以进一步减少由激活引起的内存占用

Implementation

  • 论文将 PTD-P 作为 Megatron-LM 代码库的一个扩展来实现
  • 论文的实现基于 PyTorch (2017)
  • 论文使用 NCCL (2018) 进行设备间的通信
  • 为了获得良好的性能，论文实施了针对通信和计算的优化，下面将概述这些优化

Communication Optimizations

  • 当使用 PP 时，作者希望并行地发送和接收前向和后向传播的张量
  • 每个 DGX A100 节点配备了 8 个 InfiniBand (IB) 网络卡
    • 不幸的是，发送和接收是点对点的，并且只发生在两个服务器上的一对 GPU 之间，这使得很难在流水线内的单个通信调用中利用所有 8 张卡
  • 然而，我们可以利用同时使用 TP 和 PP 这一事实来降低跨节点通信的开销
  • 特别地，论文注意到每个 Transformer 层的输出在 TP Rank 之间是复制的（在 MLP 块中的 \(g\) 操作之后，见图 5(a)）
    • 因此，在执行 TP 的两个连续流水线阶段中的 Rank ，会发送和接收完全相同的张量集合（图 5(a)）
  • 对于足够大的模型，论文使用大小为 8 的 TP
    • 这意味着论文在相邻的多 GPU 服务器上的对应 GPU 之间发送相同的张量集合 8 次
    • 为了减少这种冗余，我们可以在发送端将张量分割成大小相等的块，然后只将一个块发送到下一个节点上的对应 Rank ，使用该 Rank 自己的 InfiniBand 卡（例如，在图 9 中， Rank 1 发送给 Rank 3， Rank 2 发送给 Rank 4）
  • 对于 8 个 TP Rank ，每个块的大小将减小为原来的八分之一
  • 然后，在接收端，我们可以通过 NVLink 执行一个 all-gather 操作（这比 InfiniBand 互连快得多）来重新构建完整的张量（如图 5(b) 所示）
  • 论文称之为 分散/聚集通信优化 (scatter/gather communication optimization)
    • 这种优化有助于更好地利用 DGX A100 服务器上的多个 IB 卡，并使诸如交错调度这样通信密集的调度变得可行
  • 量化来看，通过分散-聚集通信优化，每对连续阶段之间需要执行的总通信量减少到 \(\frac{bsh}{t}\)，其中 \(t\) 是 TP 大小，\(s\) 是序列长度，\(h\) 是隐藏层大小（在论文的实验中 \(t=8\)）

Computation Optimizations

  • 论文对计算图实施了三个模型特定的优化以获得高性能
  • 第一，论文改变了 Transformer 层中的数据布局，以避免内存密集的转置操作，并启用跨步批处理 GEMM 内核
    • 具体来说，论文将数据布局从 \([b, s, a, h]\) 改为 \([s, b, a, h]\)，其中 \(b\)、\(s\)、\(a\) 和 \(h\) 分别是批处理大小、序列长度、注意力头数和隐藏层大小的维度
  • 第二，论文使用 PyTorch JIT (2016) 为一序列逐元素操作（偏置 + GeLU 以及偏置 + Dropout + 加法）生成了融合内核
  • 第三，论文创建了两个自定义内核来实现缩放、掩码和 Softmax（归约）操作的融合：
    • 一个支持通用掩码（用于如 BERT 的模型），另一个支持隐式因果掩码（用于如 GPT 的自回归模型）
  • 论文将在下一节量化这些优化的效果

Evaluation

  • 在本节中，论文试图回答以下问题：
    • PTD-P 的性能如何？它是否能够满足现实的端到端训练时间？
    • 对于给定的模型和批处理大小， PP 的扩展性如何？交错调度对性能有多大影响？
    • 不同的并行化维度如何相互作用？微批处理大小等超参数的影响是什么？
    • 分散-聚集通信优化的影响是什么？在规模上运行训练迭代时，论文对硬件施加了哪些类型的限制？
  • 论文所有的结果都是在 Selene 超级计算机 (Selene, 2000) 上使用混合精度运行的
    • 每个集群节点有 8 个 NVIDIA 80-GB A100 GPU (2018)，通过 NVLink 和 NVSwitch (2018) 相互连接
    • 每个节点有八个 NVIDIA Mellanox 200Gbps HDR InfiniBand HCA 用于应用通信，另外每个节点还有两个 HCA 用于专用存储
    • 节点通过具有 850 个交换机的三级（叶子、脊柱、核心）胖树拓扑连接
    • 这种拓扑支持高效的 All-Reduce 通信（深度学习训练中的主导通信模式）
    • 集群使用全 NVME 共享并行文件系统进行高性能数据访问和存储
  • 一个具有 16 位精度的 A100 GPU 的峰值设备吞吐量为 312 teraFLOP/s
  • 对于论文的大部分结果，论文报告每个 GPU 的吞吐量
  • 总吞吐量可以通过乘以使用的 GPU 数量来计算
  • 在论文的实验中，论文使用适当大小的 GPT 模型
    • 特别地，对于任何给定的微基准测试，模型需要能够适应实验中使用的模型并行 GPU 的数量
    • 适当的时候，论文使用标准的模型架构，如 GPT-3 (2016)

End-to-End Performance

  • 论文考虑了论文的系统在参数量从十亿到一万亿的 GPT 模型上的端到端性能，使用了张量、流水线和 DP （使用第 3 节中描述的启发式方法选择维度）
  • 特别地，论文使用了启用了分散/聚集优化的交错流水线调度
  • 所有模型使用词汇表大小（表示为 \(V\)）为 51,200（1024 的倍数）和序列长度（表示为 \(s\)）为 2048
    • 论文改变隐藏层大小（\(h\)）、注意力头数和层数（\(l\)）
    • 模型中的参数量 \(P\) 可以计算为：
      $$P=12lh^{2}\left(1+\frac{13}{12h}+\frac{V+s}{12lh}\right). \tag{2}$$
  • 随着模型大小的增加，论文也增加批处理大小（\(B\)）和 GPU 的数量（\(n\)）
  • 模型中大部分的浮点运算是在 Transformer 层和 logit 层中的矩阵乘法（GEMMs）中执行的
    • 仅考虑这些 GEMMs，每次迭代的 FLOPs 数量为（更多细节见附录）：
      $$F=96Bslh^{2}\left(1+\frac{s}{6h}+\frac{V}{16lh}\right). \tag{3}$$
    • 这是真实 FLOP 计数的下限，但应接近实际值
    • 论文将 FLOP 计为浮点运算，无论精度如何
    • 方程 (3) 假设了激活重计算，并考虑了与额外前向传播相关的浮点运算
  • 表 1 显示了模型配置以及实现的 FLOP/s（包括每个 GPU 的和所有 GPU 的总和）
    • 论文看到在 3072 个 A100 GPU（384 个 DGX A100 节点）上实现了超线性扩展，因为随着模型变大（更大的矩阵乘法），GPU 利用率提高，而通信时间相对于计算时间没有显著增加
    • 吞吐量是针对端到端训练测量的，即包括所有操作，包括数据加载、优化器步骤、通信和日志记录
    • 对于最大的模型，论文达到了峰值设备吞吐量的 52%，对于最小的模型，达到了峰值设备吞吐量的 44%
训练时间估算
  • 给定这些吞吐量，论文还可以估算在 \(T\) 个 Token 上进行端到端训练所需的总时间
    • 训练需要 \(I=\frac{T}{B\cdot s}\) 次迭代（\(s\) 为序列长度）
    • 使用方程 (3) 中的 \(F\) 值和表 1 中的经验端到端吞吐量（表示为 X），我们可以估算总训练时间
    • 对于表 1 中的配置，论文有 \(12lh \gg (V+s)\) 和 \(16lh \gg V\)
  • 将这些观察结果与方程 (2) 和 (3) 结合，论文得到：
    $$\text{End-to-end training time}\approx\frac{8TP}{nX}.$$
    • 其中：
      • \(n\) 为 GPU 数量；
      • \(X\) 为每个 GPU 的吞吐量；
      • \(T\) 为 Token 数量
      • \(P\) 为模型参数量
  • 让论文以具有 \(P=\) 175B 参数的 GPT-3 模型为例
    • 该模型在 \(T=3000\) 亿个 Token 上进行了训练
      • 在 \(n=1024\) 个 A100 GPU 上使用批处理大小 1536，论文实现了每个 GPU \(X=140\) teraFLOP/s 的吞吐量
      • 因此，训练该模型所需的时间为 34 天
    • 对于 1 万亿参数模型，论文假设端到端训练需要 450B 个 Token
    • 使用 3072 个 A100 GPU，我们可以实现每个 GPU 163 teraFLOP/s 的吞吐量，端到端训练时间为 84 天
    • 作者认为这些训练时间（使用合理数量的 GPU）是可行的

Comparison to ZeRO-3

  • 论文在表 2 和图 10 中将 PTD-P 与 ZeRO-3 (2020; 2021) 进行了比较（针对标准的 GPT-3 模型架构以及表 1 中的 530B 参数模型）
    • 这些结果提供了一个与不使用模型并行的方法的对比点
    • 论文使用 DeepSpeed Python 库 (2020) 将 ZeRO 集成到论文的代码库中
    • 当论文增加 GPU 数量时，论文保持全局批处理大小不变
    • 在 GPU 数量较少且微批处理大小为 4 的情况下，PTD-P 对于 175B 和 530B 参数模型的吞吐量分别高出 6% 和 24%
    • 随着论文增加 GPU 数量，PTD-P 比单独使用 ZeRO-3 扩展得更优雅（见图 10）
      • 例如，通过将 GPU 数量加倍（保持批处理大小不变），由于跨节点通信更少，PTD-P 对两个模型的性能均优于 ZeRO-3 70%
    • 论文只考虑了不使用 TP 的 ZeRO-3
    • ZeRO-3 可以与模型并行结合，以潜在地改善其扩展行为

PP (Pipeline Parallelism)

  • 论文现在单独评估 PP 的弱扩展性能，并将非交错调度与交错调度的性能进行比较
Weak Scaling
  • 论文使用弱扩展设置、一个具有 128 个注意力头和隐藏层大小为 20480 的 GPT 模型以及微批处理大小为 1 来评估默认非交错 PP 调度的扩展性
  • 随着论文增加流水线阶段的数量，论文也按比例增加模型中的层数来增加模型的大小
    • 例如，当 PP 大小为 1 时，论文使用具有 3 个 Transformer 层和 15B 参数的模型，当 PP 大小为 8 时，论文使用具有 24 个 Transformer 层和 121B 参数的模型
  • 论文对所有配置使用 TP 大小为 8，并改变使用的 A100 GPU 总数，从 8 到 64
  • 图 11 显示了两种不同批处理大小下每个 GPU 的吞吐量，以说明流水线气泡的影响，其行为符合 \(\frac{p-1}{m}\)（第 2.2.1 节）
    • 正如预期的那样，较高的批处理大小扩展性更好，因为流水线气泡被分摊到更多的微批处理上
Interleaved versus Non-Interleaved Schedule
  • 图 12 显示了在具有 175B 参数（96 层，96 个注意力头，隐藏层大小为 12288）的 GPT-3 (2020) 模型上，交错和非交错调度每个 GPU 的吞吐量
    • 带有分散/聚集通信优化的交错调度比非交错（默认）调度具有更高的计算性能
    • 随着批处理大小的增加，这个差距会缩小，原因有二：
      • (a) 随着批处理大小的增加，默认调度中的气泡大小减小
      • (b) 流水线内的点对点通信量与批处理大小成正比，因此随着通信量的增加，非交错调度会赶上（交错调度每个样本的通信量更多）
    • 在没有分散/聚集优化的情况下，默认调度在较大批处理大小下表现优于交错调度（未显示）

Comparison of Parallel Configurations

  • 在本小节中，论文展示了结合不同并行化维度所带来的各种权衡
  • 特别地，论文展示了对于给定模型和多个批处理大小，使用相同数量 GPU 的并行配置的性能
TP 与 PP (Tensor versus Pipeline Parallelism)
  • 论文评估了对于给定模型和批处理大小，流水线和 TP 对性能的影响
  • 图 13 中的实证结果显示了结合使用张量和 PP 来训练一个 161B 参数 GPT 模型（32 个 Transformer 层以支持 PP 大小为 32，128 个注意力头，隐藏层大小为 20480）的重要性，以实现低通信开销和高计算资源利用率
  • 论文观察到
    • TP 在一个节点（DGX A100 服务器）内效果最好，因为其 All-Reduce 通信开销大
    • PP 使用成本低得多的点对点通信，可以在节点之间进行，而不会成为整个计算的瓶颈
    • 使用 PP 时，可能会在流水线气泡中花费大量时间：
      • 因此，流水线阶段的总数应受到限制，使得流水线中的微批处理数量是流水线阶段数量的合理倍数
    • 当 TP 大小等于单个节点中的 GPU 数量（对于 DGX A100 节点为 \(8\)）时 ，达到了峰值性能
  • 这一结果表明，无论是单独使用 TP （由 Megatron (2019) 使用）还是单独使用 PP （由 PipeDream (2019) 和其他人使用），都无法与结合使用这两种技术的性能相媲美
PP 与 DP (Pipeline versus Data Parallelism)
  • 论文在图 14 中评估了数据和 PP 对一个具有 59 亿参数（32 个 Transformer 层，32 个注意力头，隐藏层大小为 3840）的 GPT 模型性能的影响
  • 论文使用比之前更小的模型，因为论文想展示当模型并行大小仅为 \(2\) 时模型能够容纳的性能
  • 为简单起见，论文在这些实验中保持微批处理大小等于 \(1\)
  • 论文看到，对于每个批处理大小，吞吐量随着 PP 大小的增加而降低，这与论文在第 3.3 节中的分析模型相符
  • PP 应主要用于支持训练无法容纳在单个工作器上的大模型，而 DP 应用于扩展训练规模
TP 与 DP (Tensor versus Data Parallelism)
  • 论文在图 15 中评估了数据和 TP 对同一个 59 亿参数 GPT 模型性能的影响（使用较小模型的原因同上）
  • 如前所述，论文最初保持微批处理大小等于 \(1\)
  • 在较大的批处理大小和微批处理大小为 1 的情况下， DP 通信不频繁；
  • TP 中所需的 all-to-all 通信需要对批次中的每个微批处理执行
  • 这种 TP 的 all-to-all 通信主导了端到端训练时间，特别是当通信需要在多 GPU 节点之间执行时
  • 随着 TP 大小的增加，论文在每个 GPU 上执行更小的矩阵乘法，降低了每个 GPU 的利用率
  • 尽管 DP 可以导致高效的扩展，但不能单独使用 DP 来处理具有有限训练批处理大小的非常大的模型，原因是
    • a） 内存容量不足
    • b） DP 的扩展限制
      • 例如，GPT-3 是以 1536 的批处理大小训练到收敛的
      • 因此， DP 仅支持扩展到 1536 个 GPU；然而，大约使用了 \(10,000\) 个 GPU 来在合理的时间内训练该模型

Microbatch Size

  • 论文在图 16 中评估了微批处理大小对结合流水线和 TP 的并行配置性能的影响（针对一个具有 91B 参数（\((t,p)=(8,8)\)）的模型）
  • 论文看到：
    • 对于这个模型，最佳微批处理大小是 \(2\)；
    • 对于其他模型，最佳微批处理大小是不同的（图中未显示）并且是模型依赖的
  • 对于给定的批处理大小，
    • 增加微批处理大小会减少流水线中的微批处理数量（\(m\)），导致更大的流水线气泡；
    • 但增加微批处理大小也可以通过增加执行内核的算术强度来提高 GPU 利用率
    • 以上这两个因素是相互矛盾的，这使得选择最佳微批处理大小具有挑战性
  • 论文来自第 3.3 节的分析模型合理地近似了真实性能，并且可以作为代理来确定如何为各种训练配置和模型选择此超参数值

Activation Recomputation

  • 图 17 显示了对于具有 145B 参数（80 个 Transformer 层，96 个注意力头，隐藏层大小为 12288）的 GPT 模型，在使用 128 个 A100 GPU、\((t,p)=(8,16)\) 以及一系列批处理大小的情况下，使用和不使用激活重计算的吞吐量
    • 对于小批处理大小，由于在反向传播期间需要执行额外的前向传播，激活重计算导致吞吐量（以每秒序列数计）降低高达 33%
    • 但激活重计算是支持更大批处理大小所必需的
  • 由于流水线气泡更小，使用激活重计算的大批处理大小的吞吐量比不使用激活重计算（对于较小批处理大小）实现的最佳吞吐量高出 up to 2\(\times\)

Scatter-Gather Optimization

  • 图 18 显示了对于具有 175B 参数的 GPT-3 模型，使用和不使用（未优化）分散/聚集通信优化时每个 GPU 的吞吐量
  • 论文看到，通过减少跨节点链路上的通信量，对于通信密集的调度（大批处理大小带交错），吞吐量提高了 up to 11%

Fused Operators（融合操作）

  • 论文还评估了第 4.2 节中描述的操作符融合对性能的影响
  • 对于 GPT-3 模型（175B 参数），通过融合，吞吐量提高了 19%（从每个 GPU 113 teraFLOP/s 到每个 GPU 135 teraFLOP/s）
  • 对于更大的 530B 参数 GPT 模型（图 1 中的模型配置），吞吐量提高了 11%（从每个 GPU 133 teraFLOP/s 到每个 GPU 148 teraFLOP/s）

Inter-Node Communication Bandwidth

  • 论文强劲的结果是优化软件和硬件栈共同使用的副产品
  • 特别地，论文利用了同一服务器内和跨服务器的 GPU 之间的高带宽通信链路
  • 在具有 3072 个 GPU 的万亿参数模型上，论文观察到流水线阶段之间点对点通信的有效二分带宽为 892 GB/s，而 DP 副本之间 All-Reduce 操作的有效二分带宽为 12.9 TB/s
  • 跨设备的操作符分区若优化不足，会导致更多的节点间通信，从而阻碍扩展性能

Checkpoint Loading and Saving

  • 训练大模型的一个重要实际考虑是加载和保存模型检查点，对于论文中考虑的模型，检查点尤其大
  • 例如，万亿参数模型的检查点大小为 13.8 TB
    • 问题：如果仅考虑参数，换算后不太对，1000000000000*4/1000/1000/1000/1000 = 4 TB，若考虑激活值才会到 12 TB
  • 所有 384 个节点（3072 个 GPU）对万亿参数模型的检查点初始加载达到了 1TB/s 的峰值读取带宽，这是并行文件系统可能的最大读取吞吐量
  • 检查点保存达到了峰值写入带宽的 40%（273 GB/s）

Parallelism for Large Models

  • PP（Pipeline model parallelism）是一种用于训练大模型的常用技术
  • PP 有几种不同的模式：
    • 论文讨论的模式使用流水线刷新（flush）来确保严格（strict）的优化器语义
    • TeraPipe (2021) 为像 GPT 这样的自回归模型（auto-regressive models）在单个训练序列的 token 之间暴露了细粒度的 PP
    • PipeTransformer (2021) 通过冻结具有“稳定”权重的层，弹性地调整 PP 和 DP 的程度，并将资源专用于训练剩余的“活跃”层
    • HetPipe (2020) 在一组异构加速器上结合使用了流水线和 DP
    • PP 也可以使用松弛的语义来实现：
      • PipeDream-2BW (2021) 维护两个权重版本，并保证权重更新延迟为 1，而无需昂贵的流水线刷新；
      • PipeMare (2021) 和 Kosson 等人 (2021) 使用异步 PP
    • 与论文考虑的带有流水线刷新的技术相比，这些技术提高了吞吐量，但可能以收敛速度或最终精度为代价
    • 此外，单独的 PP 仍然只能扩展到与模型中层数相等数量的设备，这对于某些模型架构来说是有限的
  • PipeDream (2019) 以一种原则性的方式结合了 PP 和 DP ，以减少跨设备通信
    • DeepSpeed (2021) 将 PP 与 TP 和 DP 结合起来，以训练高达万亿参数规模的模型，但吞吐量低于论文所示的结果（峰值利用率的 52% 对比 36%），原因有几个：
      • 通过算子融合使大部分算子图保持计算受限（compute-bound）、使用更高效的 PP 调度以最小化流水线气泡（pipeline bubble）大小、快速的硬件（A100 与 V100 GPU 以及同一服务器内和不同服务器间 GPU 的高带宽链路）以及扩展到更多 GPU 的能力
    • 作者希望强调，这种更高的吞吐量使得估计的训练时间更加实用（约 3 个月）；37.6 petaFLOP/s 的总吞吐量将需要大约 40 个月来训练一个同等规模的模型
      • 也可以扩展到更大的模型，但需要更多的 GPU 来保持训练时间的实用性
  • Mesh-TensorFlow (2018) 提出了一种语言，用于轻松指定结合数据和模型并行的并行化策略
    • Switch Transformers (2021) 使用 Mesh-Tensorflow 训练了一个具有 1.6 万亿参数的稀疏激活（sparsely activated）的基于专家（expert-based）的模型，其预训练速度比 T5-11B 模型 (2019) 有所提高

Sharded Data Parallelism（分片 DP）

  • 作为 MLPerf 0.6 (2019) 性能优化的一部分，引入了分片 DP （sharded data parallelism）(2019, 2020)，其中优化器状态（optimizer state）在 DP 工作节点（data-parallel workers）上进行分片
  • 这种方法有两个优点：
    • (a) 它不会在原始 DP （vanilla data parallelism）的基础上引入额外的通信；
    • (b) 它将优化器的计算和内存成本分摊到 DP 分区中
  • ZeRO (2019, 2021) 扩展了这个思想：
    • 权重参数和梯度也在 DP 工作节点上进行分片，工作节点在执行计算前从其“拥有”相应状态的工作节点获取相关状态
      • 这增加了额外的通信，但可以通过仔细重叠计算和通信来部分隐藏
      • 但如果不使用 TP 或者批大小不够大以隐藏额外的通信开销，这会变得困难（图 10）
  • ZeRO-Infinity (2021) 使用 NVMe 来高效地交换参数，使得能够在少量 GPU 上训练非常大的模型
  • 论文注意到，使用少量 GPU 训练非常大的模型会导致不切实际训练时间（例如，收敛需要数千年）

Automatic Partitioning（自动分区）

  • FlexFlow (2018)、PipeDream (2019)、DAPPLE (2021) 和 Tarnawski 等人 (2020) 都借助成本模型（cost models）在多个设备上自动分区模型训练图
  • 然而，这些方法都没有考虑论文考虑的所有并行维度：流水线和 TP、DP、微批大小（microbatch size）、以及像激活重计算（activation recomputation）这样的内存节省优化对训练大于加速器内存容量的模型的影响
  • 这些增加的维度扩大了需要探索的搜索空间
  • Gholami 等人 (2018) 展示了如何对数据和模型并行组合的通信成本进行建模

高性能计算（HPC for Model Training）

  • Goyal 等人 (2017) 和 You 等人 (2018) 都展示了使用高性能计算（High Performance Computing, HPC）技术在几分钟内训练高精度 ImageNet 模型
  • 但所考虑的图像分类模型可以轻松地放在单个加速器上，使得模型并行变得不必要；
  • 它们支持非常大的批大小（\(>32k\)），允许将 DP 扩展到大量工作节点，且通信不频繁；
    • 并且它们由紧凑的卷积层组成，这些层本身就很适合 DP 通信