研究背景
- 现代人工智能系统:由基础模型驱动。
- Llama 3:一组新的基础模型,原生支持多语言、编码、推理和工具使用。
模型概述
- 最大模型:具有 405B 参数和最多 128K 个 token 上下文窗口的密集 Transformer。
1 Introduction
基础模型介绍
基础模型是设计用于支持多种AI任务的语言、视觉、语音和其他模态的通用模型。它们构成了许多现代AI系统的基础。
现代基础模型的开发阶段
现代基础模型的开发包括两个主要阶段:
- 预训练阶段:模型使用简单的任务(如 next-word prediction 或 captioning)进行大规模训练。
- 后训练阶段:模型被调整以遵循指令、与人类偏好保持一致,并改进特定能力(如编码和推理)。
Llama 3 模型群
在这篇论文中,我们介绍了一组名为 Llama 3 的语言基础模型。Llama 3 模型群原生支持多语言、编码、推理和工具使用。最大的模型是一个具有 405B 参数的密集 Transformer,可以处理最多 128K 个 token 的上下文信息。
关键发展要素
高质量基础模型的开发中有三个关键要素:
- 数据:与之前的 Llama 版本相比,我们在预训练和后训练中使用的数据的数量和质量都有所提高。我们使用大约 15T 个多语言 token 的语料库对 Llama 3 进行预训练,而 Llama 2 只有 1.8T。
- 规模:我们使用 15.6T 个文本 token 对一个拥有 405B 可训练参数的旗舰模型进行了预训练。根据基础模型的规模定律,旗舰模型比使用相同训练过程的较小模型表现更好。
- 复杂性管理:我们选择标准的密集 Transformer 模型架构,并采用基于监督微调(SFT)、拒绝采样(RS)和直接偏好优化(DPO)的相对简单的后训练程序,以最大化训练稳定性。
Llama 3 模型群的性能评估
Llama 3 在广泛的基准数据集上进行了性能评估,并进行了广泛的人类评估。实验评估表明,旗舰模型在各种任务上与领先的语言模型(如GPT-4)表现相当,并接近最先进水平。较小的模型在同类参数模型中表现最佳。
2 总体概述
2.1 Llama 3 模型架构
Llama 3 的模型架构如图 1 所示。Llama 3 语言模型的开发包括两个主要阶段:
-
语言模型预训练:首先,我们将大量多语言文本语料库转换为离散 token,并在生成的数据上预训练大型语言模型,以执行下一个 token 的预测。在这个阶段,模型学习语言的结构,并从所“阅读”的文本中获得大量关于世界的知识。为了有效地进行预训练,我们在大规模上进行预训练:我们使用 8K token 的上下文窗口,在 15.6T token 上预训练一个具有 405B 参数的模型。这个标准的预训练阶段之后是一个持续的预训练阶段,将支持的上下文窗口增加到 128K token。
-
语言模型后训练:预训练的语言模型对语言有丰富的理解,但它还不遵循指令或以我们期望助手的方式行事。我们通过几轮人类反馈来对齐模型,每轮包括在指令调谐数据(instruction tuning data)上的监督微调(SFT)和直接偏好优化(DPO)。在这个后训练阶段,我们还集成了新功能,如工具使用,并在其他领域(如编码和推理)观察到强烈的改进。最后,在后训练阶段还纳入了安全缓解(safety mitigations)措施。
2.2 多模态实验
我们还进行了实验,通过组合方法向 Llama 3 添加图像、视频和语音功能。我们研究的方法包括如下三个额外阶段:
-
多模态编码器预训练:我们分别为图像和语音训练单独的编码器。我们在大量的图像-文本对上训练我们的图像编码器。这教会了模型视觉内容与自然语言描述之间的关系。我们的语音编码器使用自监督方法进行训练,该方法屏蔽了语音输入的部分,并试图通过离散 token 表示重建屏蔽部分。通过这个方式,模型学习了语音信号的结构。
-
视觉适配器训练:我们训练一个适配器,将预训练的图像编码器集成到预训练的语言模型中。适配器由一系列交叉注意力层组成,这些层将图像编码器的表示输入到语言模型中。适配器在文本-图像对上进行训练。这将图像表示与语言表示对齐。在适配器训练期间,我们还更新图像编码器的参数,但我们故意不更新语言模型的参数。我们还在图像适配器之上训练了一个视频适配器,以处理成对的视频-文本数据。这使得模型能够跨帧聚合信息。
-
语音适配器训练:最后,我们通过一个适配器将语音编码器集成到模型中,该适配器将语音编码转换为可以直接输入到微调语言模型中的 token 表示。适配器和编码器的参数在监督微调阶段联合更新,以实现高质量的语音理解。我们在语音适配器训练期间不改变语言模型。我们还集成了一个文本到语音的系统。
3 预训练
语言模型预训练包括以下四个部分:
- 大规模训练语料库的策展(curation)和筛选:对用于预训练的大量文本数据进行筛选和处理。
- 模型架构和相应的规模定律(scaling laws)的开发:确定模型大小并开发相应的模型架构。
- 大规模高效预训练技术的开发:开发在大规模数据上进行高效预训练的技术。
- 预训练配方(recipe)的开发:制定预训练过程中的具体步骤和方法。
每个部分将在下文分别介绍。
3.1 预训练数据
语言模型预训练的数据集是通过从各种数据源收集知识直到 2023 年底创建的。我们对每个数据源应用了多种去重方法和数据清洗机制,以获得高质量的 token。我们移除了包含大量个人身份信息(PII)的领域,以及已知含有成人内容的领域。
3.1.1 网络数据策展
数据来源与清理过程
- 数据来源:大部分数据来自网络,直至 2023 年底。
- 个人身份信息(PII)和安全过滤:实施过滤器,移除可能包含不安全内容或大量 PII 的网站数据,以及已知包含成人内容的网站。
- 文本提取和清理:从非截断的网络文档中提取高质量、多样化的文本。构建自定义解析器,优化模板去除和内容召回的精度。对于包含数学和代码内容的 HTML 页面,小心处理以保留内容结构。
- 去重:在 URL、文档和行级别进行多轮去重。
- URL 级别去重:保留每个 URL 对应的最新页面版本。
- 文档级别去重:使用全局 MinHash (On the resemblance and containment of documents)去除近似重复文档。
- 行级别去重:移除在每 30M 文档桶中出现超过 6 次的行。
- 启发式过滤:开发启发式方法移除低质量文档、异常值和重复过多的文档。
- 使用重复的n-gram覆盖率比例(Scaling Language Models: Methods, Analysis & Insights from Training Gopher)来删除由重复内容组成的行,例如日志或错误消息。这些行可能非常长且独特,因此无法通过行去重来过滤。
- 使用“脏词”计数(Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer)来过滤掉未被域名屏蔽列表覆盖的成人网站。
- 使用基于 token 分布的 Kullback-Leibler 散度来过滤掉包含与训练语料库分布相比具有过多异常令牌的文档。
- 基于模型的质量过滤:实验性地应用各种基于模型的质量分类器来进一步筛选高质量 token。
- 使用快速分类器,如 fasttext(Bag of Tricks for Efficient Text Classification),经过训练以识别给定文本是否会被维基百科引用(Llama: Open and efficient foundation language models)。
- 使用更计算密集的基于 Roberta 的分类器(RoBERTa: A Robustly Optimized BERT Pretraining Approach),该分类器是基于 Llama 2 的预测进行训练的。为了基于 Llama 2 训练一个质量分类器,我们创建了一个经过清理的网络文档的训练集,描述了质量要求,并指示 Llama 2 的聊天模型确定文档是否符合这些要求。
- 出于效率原因,我们使用 DistilRoberta(DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter)为每个文档生成质量分数。
-
代码和推理数据: 与 DeepSeek-AI 类似,我们构建了领域特定的流程,提取与代码和数学相关的网页。代码和推理分类器都是在 Llama 2 注释的网络数据上训练的 DistilledRoberta 模型。与上述通用质量分类器不同,我们进行了提示调整(prompt tuning),以针对包含数学推导、STEM领域的推理以及自然语言交织的代码的网页。由于代码和数学的 token 分布与自然语言大不相同,这些流程实现了领域特定的 HTML 提取、定制的文本特征和过滤启发式方法。
- 多语言数据:
- 使用基于 fasttext 的语言识别模型将文档分类为 176 种语言。
- 在每种语言的数据中进行文档级别和行级别的去重。
- 应用特定于语言的启发式和基于模型的过滤器来移除低质量文档。
- 使用基于多语言 Llama 2 的分类器对多语言文档进行质量排名,确保高质量内容优先。
通过上述过程,确保了预训练数据的质量和多样性,为后续的模型训练打下了坚实的基础。
3.1.2 确定数据混合比例
为了获得高质量的语言模型,仔细确定预训练数据混合中不同数据源的比例至关重要。我们确定这一数据混合的主要工具是知识分类(knowledge classification )和规模定律实验(scaling law experiments)。
知识分类
我们开发了一个分类器,用于对我们网络数据中包含的信息类型进行分类,以更有效地确定数据混合比例。我们使用这个分类器对网络上过多比例的数据类别(例如,艺术和娱乐)进行下采样。
数据混合的规模定律
为了确定最佳的数据混合比例,我们进行了规模定律实验,在这些实验中,我们在一个数据混合上训练几个小模型,并使用这些结果来预测大模型在该混合上的性能(见第 3.2.1 节)。我们对不同的数据混合重复这个过程多次,以选择一个新的数据混合候选。随后,我们在这个候选数据混合上训练一个更大的模型,并在几个关键基准上评估该模型的性能。
数据混合总结
我们最终的数据混合大约包含 50% 的一般知识对应的 token,25% 的数学和推理 token,17% 的代码 token,以及 8% 的多语言 token。
3.1.3 退火数据
经验上,我们发现对少量的高质量代码和数学数据进行退火(annealing,见第 3.4.3 节)可以提高预训练模型在关键基准测试上的性能。类似于 DataComp-LM: In search of the next generation of training sets for language models 的做法,我们执行退火时使用一种数据混合方法,在选定的领域中对高质量数据进行上采样。我们的退火数据中不包括任何常用基准测试的训练集。这使得我们能够评估 Llama 3 的真实少样本学习能力和跨域泛化能力。
遵循 OpenAI(GPT-4 Technical Report)的做法,我们在 GSM8k 和 MATH 训练集上评估退火的有效性。我们发现退火使得预训练的 Llama 3 8B 模型在 GSM8k 和 MATH 验证集上的性能分别提高了 24.0% 和 6.4%。
然而,在 405B 模型上的改进可以忽略不计,这表明我们的旗舰模型具有强大的上下文学习和推理能力,并且不需要特定的领域内训练样本就能获得强大的性能。使用退火评估数据质量。类似于 Does your data spark joy? Performance gains from domain upsampling at the end of training 的做法,我们发现退火使我们能够判断小型特定领域数据集的价值。我们通过在 40B 个 token 上将 50% 训练的 Llama 3 8B 模型的学习率线性降低到 0 来衡量这些数据集的价值。在这些实验中,我们为新数据集分配了 30% 的权重,其余 70% 的权重分配给默认的数据混合。与为每个小型数据集执行规模定律实验相比,使用退火评估新数据源更有效率。
3.2 模型架构
Llama 3 使用标准的密集 Transformer 架构,在模型架构上与 Llama 和 Llama 2 没有显著偏离;我们的性能提升主要来自于数据质量和多样性的改进以及训练规模的增加。
与 Llama 2 相比,Llama 3 做了一些小的修改:
- 使用分组查询注意力(GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints)和 8 个键值头,以提高推理速度并减少解码过程中的键值缓存大小。
- 使用一种注意力掩码,防止同一序列内不同文档之间的自注意力。我们发现在标准预训练期间,这一变化的影响有限,但在对非常长的序列进行持续预训练时,这一点很重要。
- 我们使用了一个包含 128K 个 token 的词汇表。我们的 token 词汇表结合了 tiktoken3 tokenizer 的 100K 个 token 和额外的 28K 个 token,以更好地支持非英语语言。与 Llama 2 tokenizer 相比,我们的新 tokenizer 在英文数据样本上的压缩率从每个 token 的 3.17 个字符提高到 3.94 个字符。这使得模型在相同的训练计算量下可以“阅读”更多的文本。我们还发现,从选定的非英语语言中添加 28K 个 token 改善了压缩率和下游性能,对英语 tokenization 没有影响。
- 我们将 RoPE 基频超参数(RoPE base frequency hyperparameter)增加到 500,000。这使我们能够更好地支持更长的上下文;Effective long-context scaling of foundation models 表明这个值对于长达 32,768 的上下文长度是有效的。
Llama 3 405B 使用了一个具有 126 层、token 表示维度为 16,384 和 128 个注意力头的架构。
3.2.1 规模定律
我们开发了规模定律(Training compute-optimal large language models、Scaling laws for neural language models),以确定我们的旗舰模型在我们预训练计算预算下的最优模型大小。除了确定最优模型大小之外,一个主要挑战是预测旗舰模型在下游基准任务上的性能,原因有几个:
- 现有的规模定律通常只预测下一个 token 的预测损失,而不是特定基准的性能。
- 规模定律可能会因为基于小计算预算的预训练运行而变得嘈杂和不可靠。
为了应对这些挑战,我们实施了两阶段方法来开发准确预测下游基准性能的规模定律:
- 我们首先建立计算最优模型在下游任务上的负对数似然与训练 FLOPs 之间的相关性。
- 接下来,我们将下游任务上的负对数似然与任务准确性相关联,利用规模定律模型和用更高计算 FLOPs 训练的旧模型。在这一步中,我们特别利用了 Llama 2 系列模型。
这种方法使我们能够预测给定训练 FLOPs 的计算最优模型在下游任务上的性能。我们使用类似的方法来选择我们的预训练数据组合(见第 3.4 节)。
规模定律实验
具体来说,我们通过使用 $ 6 × 10^{18} $ FLOPs 到 $ 10^{22} $ FLOPs 之间的计算预算预训练模型来构建我们的规模定律。在每个计算预算下,我们预训练的模型大小在 40M 到 16B 参数之间,每个计算预算下使用一部分模型大小。在这些训练运行中,我们使用余弦学习率计划,线性预热 2,000 训练步骤。峰值学习率根据模型的大小设置在 $ 2 × 10^{-4} $ 到 $ 4 × 10^{-4} $ 之间。我们将余弦衰减设置为峰值的 0.1。每一步的权重衰减设置为该步骤学习率的 0.1 倍。我们为每个计算规模设置了一个固定的批量大小,范围在 250K 到 4M 之间。
这些实验产生了图 2 中的 IsoFLOPs 曲线。这些曲线上的损失是在单独的验证集上测量的。我们使用二阶多项式拟合测量的损失值,并确定每个抛物线的最小值。我们将抛物线的最小值称为相应预训练计算预算下的计算最优模型。
我们使用这种方式识别的计算最优模型来预测特定计算预算下的最优训练 token 数量。为此,我们假设计算预算 $ C $ 和最优训练 token 数量 $ N^{\star}(C) $ 之间存在幂律关系:
\[N^{\star}(C)=A C^{\alpha}.\]我们使用图 2 中的数据拟合 A 和 α。我们发现 (α, A) = (0.53, 0.29);相应的拟合显示在图 3 中。将得到的规模定律外推到 3.8 × 10^25 FLOPs 表明在 16.55T token 上训练一个 402B 参数模型。
一个重要的观察是,随着计算预算的增加,IsoFLOPs 曲线在最小值附近变得更平。这意味着旗舰模型的性能对于模型大小和训练 token 之间权衡的小变化相对稳健。基于这个观察,我们最终决定训练一个具有 405B 参数的旗舰模型。
预测下游任务的性能
我们使用得到的计算最优模型来预测旗舰 Llama 3 模型在基准数据集上的性能。首先,我们将基准中正确答案的(标准化的)负对数似然与训练 FLOPs 线性相关。在这项分析中,我们仅使用上述数据组合上训练至 $ 10^{22} $ FLOPs 的规模定律模型。接下来,我们使用规模定律模型和 Llama 2 模型建立对数似然与准确性之间的 sigmoidal 关系,这些模型是使用 Llama 2 数据组合和 tokenizer 训练的。我们在下图中展示了 ARC 挑战基准上的这个实验结果。我们发现这种两步规模定律预测,在四个数量级上外推,非常准确:它只是略微低估了旗舰 Llama 3 模型的最终性能。
3.3 基础设施、规模和效率
本章节介绍了支撑 Llama 3 405B 模型大规模预训练的硬件和基础设施,并讨论了几项优化措施,这些措施提高了训练效率。
- 硬件和基础设施:描述了用于支持 Llama 3 405B 模型大规模预训练的硬件和基础设施。
- 优化措施:讨论了几项优化措施,这些措施有助于提高训练效率。
3.3.1 训练基础设施
Llama 3 405B 模型的训练基础设施包括硬件、存储和网络等方面的优化。
硬件和基础设施
- 训练集群:Llama 1 和 2 模型在 Meta 的 AI 研究超级集群上进行训练,而 Llama 3 的训练则迁移到 Meta 的生产集群。
- 计算:Llama 3 405B 在多达 16K 个 H100 GPU 上进行训练,每个 GPU 的 TDP 为 700W,配备 80GB HBM3,使用 Meta 的 Grand Teton AI 服务器平台。每个服务器配备八个 GPU 和两个 CPU,GPU 之间通过 NVLink 连接。训练作业使用 MAST(Meta 的全球规模训练调度器)进行调度。
- 存储:使用 Meta 的通用分布式文件系统 Tectonic 构建存储结构,提供 240 PB 的存储空间,支持 2 TB/s 的可持续吞吐量和 7 TB/s 的峰值吞吐量。主要挑战是支持高度突发的检查点写入,这些写入会在短时间内使存储结构饱和。检查点保存每个 GPU 的模型状态,范围从 1 MB 到 4 GB,用于恢复和调试。
- 网络:Llama 3 405B 使用基于 Arista 7800 和 Minipack2 Open Compute Project OCP 机架交换机的 RoCE 网络。Llama 3 家族中的较小模型使用 Nvidia Quantum2 Infiniband 网络。两个集群都利用 400 Gbps 的 GPU 间互连。
网络拓扑和优化
- 网络拓扑:基于 RoCE 的 AI 集群包括 24K 个 GPU,通过三层 Clos 网络连接。底层每个机架托管 16 个 GPU,通过单个 Minipack2 ToR 交换机连接。中间层通过集群交换机连接 192 个这样的机架,形成一个具有全二分带宽的 3,072 GPU 节点。顶层通过聚合交换机连接同一数据中心建筑内的八个这样的 pod,形成一个 24K GPU 集群。
- 负载均衡:LLM 训练产生大量网络流量,传统方法如 ECMP 路由难以在所有可用网络路径上进行负载均衡。我们采用两种技术:首先,我们的集合库在两个 GPU 之间创建 16 个网络流,而不是一个,从而减少每个流的流量并提供更多流;其次,我们的增强 ECMP(E-ECMP)协议通过对 RoCE 头包中的附加字段进行哈希,有效地在这些不同网络路径之间平衡这 16 个流。
- 拥塞控制:我们在 spine(RDMA over Ethernet for Distributed AI Training at Meta Scale)中使用深度缓冲交换机来容纳由集合通信模式引起的瞬态拥塞和缓冲。这种设置有助于限制由慢服务器引起的持续拥塞和网络背压(network back pressure)的影响,这在训练中很常见。通过 E-ECMP 实现更好的负载均衡显著减少了拥塞的机会。通过这些优化,我们成功运行了一个 24K GPU 集群,而不需要传统的拥塞控制方法,如数据中心量化拥塞通知(DCQCN)。
3.3 基础设施、扩展和效率
3.3.2 模型扩展的并行性
为了扩展我们最大模型的训练,我们使用了 4D 并行性——四种不同类型的并行性方法的组合——来分片模型。这种方法有效地将计算分布在许多 GPU 上,并确保每个 GPU 的模型参数、优化器状态、梯度和激活都适合其 HBM。我们对 4D 并行性的实现如下图所示。它结合了张量并行性(TP)、流水线并行性(PP)、上下文并行性(CP)和数据并行性(DP)。
张量并行性将单个权重张量分割成不同设备上的多个块。流水线并行性通过层将模型垂直划分为阶段,以便不同设备可以并行处理完整模型管道的不同阶段。上下文并行性将输入上下文划分为段,减少非常长序列长度输入的内存瓶颈。我们使用完全分片的数据并行性(FSDP)(ZeRO: Memory Optimizations Toward Training Trillion Parameter Models、ZeRO-Offload: Democratizing Billion-Scale Model Training、PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel),实现数据并行性的同时,分片模型、优化器和梯度,在多个 GPU 上并行处理数据,并在每个训练步骤后同步。
流水线并行性的改进
我们遇到了现有实现的几个挑战:
- 批量大小约束:当前实现对每个 GPU 支持的批量大小有限制,要求它必须能被流水线阶段的数量整除。
- 内存不平衡:现有的流水线并行性实现导致资源消耗不平衡。第一阶段由于嵌入和热身微批次而消耗更多内存。
- 计算不平衡:在模型的最后一层之后,需要计算输出和损失,导致这个阶段成为执行延迟瓶颈。
为了解决这些问题,我们修改了我们的流水线调度,如下图所示,它允许灵活设置N,本例中 N = 5,可以在每个批次中运行任意数量的微批次。这使我们能够运行:(1)比阶段数量少的微批次,当我们在大规模时有批次大小限制;或者(2)更多的微批次来隐藏点对点通信,找到 DFS 和 BFS 之间的最佳通信和内存效率的平衡点。
为了平衡流水线,我们分别从第一个和最后一个阶段中减少一个Transformer层。这意味着第一个阶段上的第一个模型块只有嵌入层,而最后一个阶段上的最后一个模型块只有输出投影和损失计算。
为了减少流水线泡沫,我们使用交错调度(Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM)在一个流水线排上具有 V 个流水线阶段。此外,我们在 PP 中采用异步点对点通信,大大加快了训练速度,特别是当文档掩码引入额外的计算不平衡时。
我们启用 TORCH_NCCL_AVOID_RECORD_STREAMS 以减少来自异步点对点通信的内存使用。最后,为了减少内存成本,基于详细的内存分配分析,我们主动释放将不用于未来计算的张量,包括每个流水线阶段的输入和输出张量。通过这些优化,我们可以在 8K 标记序列上对 Llama 3 进行预训练,而无需激活检查点。
长序列的上下文并行性
我们利用上下文并行性(CP)来提高扩展 Llama 3 上下文长度时的内存效率,并实现最长 128K 长度的序列训练。在 CP 中,我们跨序列维度进行划分,具体来说,我们将输入序列划分为 2 × CP 块,以便每个 CP 等级接收两个块以实现更好的负载平衡。
与现有的CP实现(环形结构中重叠通信和计算)(Ring attention with blockwise transformers for near-infinite context)不同,我们的 CP 实现采用了基于全收集(all-gather)的方法,首先全收集键(K)和值(V)张量,然后计算本地查询(Q)张量块的注意力输出。尽管全收集通信延迟暴露在关键路径中,但我们仍然采用这种方法主要有两个原因:(1)基于全收集的 CP 注意力更容易且更灵活地支持不同类型的注意力掩码,例如文档掩码;(2)由于使用了 GQA,所以传输的 K 和 V 张量比 Q 张量要小得多,因此暴露的全收集延迟很小。因此,注意力计算的时间复杂度比全收集($ O(S^2) $ 与 $ O(S) $,其中$ S $表示完全因果掩码(full causal mask)中的序列长度)大一个数量级,使得全收集开销可以忽略不计。
网络感知的并行性配置
并行性维度的顺序[TP, CP, PP, DP]针对网络通信进行了优化。最内层的并行性需要最高的网络带宽和最低的延迟,因此通常限制在同一服务器内。最外层的并行性可能分布在多跳网络上,可容忍更高的网络延迟。因此,基于对网络带宽和延迟的要求,我们按照[TP, CP, PP, DP]的顺序安排并行性维度。DP(即 FSDP)是最外层的并行性,因为它可以通过异步预取分片模型权重和减少梯度来容忍更长的网络延迟。
数值稳定性
通过比较不同并行性设置下的训练损失,我们解决了影响训练稳定性的一些数值问题。为了确保训练收敛,我们在多个微批次的反向计算中使用 FP32 梯度累积,并在FSDP中在数据并行工作器之间使用 FP32 进行梯度减少散播(reduce-scatter)。对于在前向计算中多次使用的中间张量(例如视觉编码器输出),反向梯度也以 FP32 进行累积。
3.4 训练方法
Llama 3 405B 模型的预训练方法包括三个主要阶段:
- 初始预训练:这是预训练的第一步。
- 长上下文预训练:在此阶段,模型会进行更长上下文的训练。
- 退火:这是预训练的最后一步,用于优化模型的性能。
对于 8B 和 70B 模型,我们也使用了类似的预训练方法。
3.4.1 初始预训练
Llama 3 405B 的预训练使用了余弦学习率计划,峰值学习率为 $ 8 × 10^{-5} $,线性预热 8,000 步,然后在 1,200,000 训练步内衰减到 $ 8 × 10^{-7} $。我们在训练初期使用较低的批量大小以提高训练稳定性,随后增加批量大小以提高效率。具体来说,我们初始的批量大小为 4M 个 token 和长度为 4,096 的序列,在预训练 252M 个 token 后,我们将这些值翻倍为批量大小 8M 和序列长度 8,192 的 token。在预训练 2.87T 个 token 后,我们再次将批量大小翻倍至 16M。我们发现这种训练配方非常稳定:我们观察到很少的损失尖峰,并且不需要干预来纠正模型训练发散。
调整数据混合。我们在训练过程中对预训练数据混合进行了几次调整,以提高模型在特定下游任务上的性能。特别是,我们在预训练期间增加了非英语数据的百分比,以提高 Llama 3 的多语言性能。我们还对数学数据进行了上采样,以提高模型的数学推理能力,在预训练的后期阶段添加了更多最新的网络数据,以推进模型的知识截止点,并对后来被识别为质量较低的预训练数据子集进行了下采样。
3.4.2 长上下文预训练
在预训练的最后阶段,我们在长序列上进行训练,以支持最多 128K 个 token 的上下文窗口。我们没有在早期就在长序列上进行训练,因为在自注意力层中的计算量随着序列长度的增加呈二次方增长。我们逐步增加支持的上下文长度,并进行预训练,直到模型成功适应了增加的上下文长度。我们通过以下两个标准来评估成功的适应:
- 模型在短上下文评估上的性能是否完全恢复
- 模型是否能够完美解决达到该长度的“大海捞针”(needle in a haystack)任务
在 Llama 3 405B 的预训练中,我们逐步增加了六个阶段的上下文长度,从最初的 8K 上下文窗口开始,最终达到 128K 上下文窗口。这个长上下文预训练阶段使用了大约 800B 的训练 token。
3.4.3 退火
在最后的4000万个标记的预训练过程中,我们线性退火学习率至0,同时保持上下文长度为128K个标记。在这个退火阶段,我们还调整了数据混合,对质量非常高的数据源进行了上采样。最后,在退火过程中,我们计算模型检查点的平均值(Polyak(1991)平均法)以生成最终的预训练模型。