目录
针对长序列的上下文并行(Context Parallelism, CP)
基于网络感知的并行配置(Network-aware Parallelism Configuration)
3.3.3 集体通信(Collective Communication)
3. 预训练
语言模型的预训练包括以下几个方面:
大规模训练语料的筛选和过滤;
模型架构的开发及相应的模型规模确定(基于扩展规律/缩放定律);
大规模高效预训练技术的开发;
预训练方案的制定。
3.1 预训练数据
我们为语言模型预训练创建了一个数据集,该数据集来自多种数据源,包含截至2023年底的知识。我们对每个数据源应用了多种去重方法和数据清理机制,以获得高质量的 token(标记)。此外,我们移除了包含大量个人身份信息(PII)的域,以及已知含有成人内容的域。
3.1.1 网络数据筛选
们使用的大部分数据来自网络,以下是我们数据清理的具体流程:
PII 和安全过滤
除了其他减轻风险的措施外,我们设计了过滤器,用来移除以下来源的数据:
可能包含不安全内容或大量个人身份信息(PII)的网站;
根据 Meta 的多种安全标准被评为有害的域名;
已知含有成人内容的域名。
文本提取与清理
我们从未被截断的网页文档中提取原始 HTML 内容,以获得高质量、多样化的文本。为此,我们开发了一个定制解析器,用于提取 HTML 内容,并在模板(boilerplate)的去除精度和内容的召回率(recall)之间进行优化。
在人工评估中,我们比较了自定义解析器与流行的第三方 HTML 解析器(这些工具通常优化用于提取类似文章的内容)的性能,发现我们的解析器表现更优。
对于包含数学和代码内容的 HTML 页面,我们进行了仔细处理,以保留内容的结构。
我们保留了图片的
alt
属性文本,因为数学内容通常以预渲染图片的形式表示,而数学公式的文字说明往往嵌入在alt
属性中。我们实验性地评估了不同的清理配置。研究发现,对于主要以网络数据训练的模型来说,markdown 格式会对性能产生负面影响,因此我们移除了所有 markdown 标记。
去重(De-duplication)
我们在 URL、文档和行级别进行了多轮去重处理:
URL 级别去重(URL-level de-duplication) 我们对整个数据集进行 URL 级别的去重,对于每个 URL 对应的页面,仅保留最新版本的数据。
文档级别去重(Document-level de-duplication) 我们使用全局 MinHash 方法(Broder, 1997)对整个数据集进行去重,以移除内容相似的文档。
行级别去重(Line-level de-duplication) 我们采用类似于 ccNet(Wenzek et al., 2019)的方法进行激进的行级别去重处理。具体来说,我们移除了在每 3000 万文档的分组中重复出现超过 6 次的行。 尽管手动的质量分析表明,行级去重不仅移除了来自各种网站(例如导航菜单、cookie 警告)的残留模板内容,也移除了部分高质量的重复文本,但我们的实验证明,这种去重处理带来了显著的性能提升。
启发式过滤(Heuristic Filtering)
我们开发了启发式规则,用于移除额外的低质量文档、异常值以及包含过多重复内容的文档。一些具体的规则包括:
重复 n-gram 覆盖率比(duplicated n-gram coverage ratio) 我们移除了由重复内容(例如日志或错误信息)组成的行。这些行可能非常长且唯一,因此无法通过行级别去重过滤掉。
“敏感词”计数(dirty word counting) 我们通过计数“敏感词汇”(Raffel et al., 2020),过滤掉未被域名屏蔽列表覆盖的成人网站。
基于 token 分布的 Kullback-Leibler 发散(KL divergence) 我们使用 KL 发散来过滤掉那些包含过多异常 token 的文档,与训练语料的分布相比,这些文档的 token 分布存在显著偏差。
基于模型的质量过滤
核心目的:
通过自动化工具筛选高质量的训练数据,避免人为筛选的低效和主观性。
方法与工具:
fasttext:
快速高效的分类工具,用于判断文档是否可能被 Wikipedia 参考(一个高质量的标准)。
适合大规模数据快速初筛。
Roberta 和 DistilRoberta:
更复杂的模型,用于细化筛选,基于 Llama 2 的预测结果进行训练。
DistilRoberta 是 Roberta 的轻量级版本,运行速度快且计算成本低。
训练方式:
定义明确的质量标准,并让 Llama 2 的聊天模型标注哪些文档符合这些标准,再用这些标注数据训练质量分类器。
每个文档都生成质量分数,最终通过实验验证筛选策略的效果。
代码和数学推理数据处理
背景:
代码和数学内容的 token 分布(例如变量名、公式符号等)与普通自然语言差别很大,需要定制化处理管道。
处理方法:
领域特定分类器:
DistilRoberta 被用于从网页中提取数学推理、STEM 推导过程和代码与自然语言混合的内容。
Prompt 调优:
通过特定提示优化分类器的能力,专注于目标内容的提取。
定制管道:
领域特定的 HTML 提取:针对数学公式或代码段的网页结构优化解析器。
自定义过滤:基于代码和数学的特征过滤低质量数据,例如避免干扰性标记或无用符号。
多语言数据处理
背景:
预训练需要支持多语言,但多语言数据质量参差不齐,需要精细化的处理流程。
具体流程:
语言识别:
使用 fasttext 分类器将数据分类到 176 种语言中,确保对每种语言的单独处理。
多语言去重:
文档级别和行级别去重,避免在多语言数据中出现冗余和重复内容。
语言特定过滤器:
针对不同语言的特点,开发独立的启发式规则和基于模型的过滤器。例如:某些语言可能需要过滤掉特定的格式或符号。
质量排名:
使用多语言版本的 Llama 2 分类器,对多语言数据按质量进行排序,确保优先使用高质量内容。
平衡实验:
在预训练中通过实验调整多语言 token 的比例,确保模型在英语和多语言基准上的性能都达到最佳平衡。
3.1.2 确定数据组合
为了获得高质量的语言模型,精心确定预训练数据组合中不同数据源的比例至关重要。我们确定数据组合的主要工具是知识分类和缩放法实验。
知识分类。我们开发了一个分类器,用于对我们从网络获取的数据中包含的信息类型进行分类,从而更有效地确定数据组合。我们使用这个分类器来对那些在网络上过度代表的类别进行下采样,例如艺术和娱乐类数据。
数据组合的缩放法。为了确定最佳的数据组合,我们进行缩放法实验,其中我们训练几个小模型,使用这些小模型的数据组合来预测一个大模型在该数据组合上的性能(见3.2.1节)。我们针对不同的数据组合多次重复这个过程,以选择新的数据组合候选。随后,我们在这个候选数据组合上训练一个更大的模型,并在多个关键基准上评估该模型的性能。
数据组合总结。我们的最终数据组合大约包含50%的常识性知识标记,25%的数学和推理标记,17%的代码标记,和8%的多语言标记。
3.1.3 数据退火
什么是退火(Annealing)?
退火(Annealing)最初是从物理学中的一个概念借用来的,通常用来描述一种逐渐冷却的过程。在材料科学中,退火是指通过加热并缓慢冷却金属或其他材料,来改善其结构、降低应力和增加稳定性。这个过程有助于材料的结构更加均匀、强度提高。
在机器学习和神经网络训练中,“退火”通常指的是一种优化技术,用来在训练过程中逐步减小学习率(learning rate)或其他超参数,从而让模型更加平稳地收敛。这个过程的灵感来源于物理中的退火过程——让系统逐渐趋向最优状态。
为什么要退火?
在机器学习中,退火的主要目的是帮助模型更好地收敛,并避免陷入局部最优解。具体来说,退火有以下几个作用:
稳定训练过程:当学习率过大时,模型可能会跳跃到错误的解或收敛得很慢。通过逐渐减小学习率,模型可以在训练的后期更稳定地收敛,避免因为学习率过大而错过全局最优解。
提高泛化能力:在一些任务中,特别是处理小数据集或复杂任务时,退火可以帮助模型避免过拟合。通过逐步减少对新数据的过分依赖,模型可以学到更具有泛化能力的特征。
控制训练过程中的噪声:在开始时,较大的学习率让模型能够快速学习和调整权重,而当模型接近收敛时,通过逐渐减小学习率,模型能够精细调整,避免过多的噪声影响最终结果
遵循OpenAI(2023a)的做法,我们评估了退火在GSM8k(Cobbe等,2021)和MATH(Hendrycks等,2021b)训练集上的效果。我们发现,退火使得预训练的Llama 3 8B模型在GSM8k和MATH验证集上的表现分别提升了24.0%和6.4%。然而,在405B模型上,退火的改进微不足道,表明我们的旗舰模型具有强大的上下文学习和推理能力,并不需要特定领域的训练样本来获得强大的性能。
使用退火评估数据质量。类似于Blakeney等人(2024),我们发现退火使我们能够判断小规模领域特定数据集的价值。我们通过将Llama 3 8B模型在40B个标记上进行线性退火,将学习率从50%训练状态逐渐退火到0,来评估这些数据集的价值。在这些实验中,我们为新数据集分配了30%的权重,其余70%的权重用于默认数据组合。使用退火来评估新的数据源,比对每个小数据集进行缩放法实验更为高效。
3.2 模型架构
这段文字介绍了Llama 3模型的架构及其与Llama和Llama 2的主要区别。Llama 3仍然采用的是经典的Transformer架构,这种架构自2017年由Vaswani等人提出以来,成为了自然语言处理(NLP)任务中的主流架构。Llama 3与Llama和Llama 2的架构差异不大,主要的性能提升来自于数据质量和多样性的提高,以及模型训练规模的增加。
不过,Llama 3相较于Llama 2做了一些小的架构修改,旨在提高模型性能和推理效率:
分组查询注意力(GQA):Llama 3采用了**分组查询注意力(GQA)**方法,这是一种优化的注意力机制,旨在提高推理速度并减少解码过程中的内存消耗。具体来说,GQA使用8个键值头(key-value heads),减少了每次推理时需要缓存的键值的数量,从而加速了推理过程并降低了内存需求。
注意力掩码:Llama 3引入了注意力掩码,这个掩码限制了同一序列中不同文档之间的自注意力(self-attention)。在标准的预训练阶段,这一修改的影响并不显著,但对于处理非常长的序列时,这个修改变得非常重要,因为它可以有效地防止模型在处理长序列时产生不必要的计算负担,优化训练效率。
Llama 3 使用一个包含 128K 令牌的词汇表。我们的令牌词汇表结合了 tiktoken 分词器中的 100K 令牌,以及额外的 28K 令牌,以更好地支持非英语语言。相比于 Llama 2 的分词器,我们的新分词器在一组英语数据样本上的压缩率从 3.17 字符/令牌提升到 3.94 字符/令牌。这使得模型在相同的训练计算量下能够“阅读”更多文本。我们还发现,从一些精选的非英语语言中添加的 28K 令牌同时改善了压缩率和下游性能,并且对英语分词没有影响。
我们将 RoPE(旋转位置编码)的基频超参数增加到 500,000,这使我们能够更好地支持更长的上下文长度;根据 Xiong 等人(2023)的研究,这个值在上下文长度达到 32,768 时是有效的。
Llama 3 405B 使用了一个架构,包含 126 层,令牌表示维度为 16,384,并有 128 个注意力头。根据我们数据的缩放规律,在我们的训练预算(3.8 × 10²⁵ FLOPs)内,这种架构的模型大小是计算最优的。
PS
1.压缩率指的是模型将文本转换为令牌(tokens)时,每个令牌平均包含的字符数。简单来说,压缩率越高,每个令牌能够表达的文本信息就越多,因此模型可以处理更多的文本内容,而不需要使用太多的令牌。
2.RoPE
(1)RoPE 和传统位置编码的区别
传统的 Transformer 通常用 绝对位置编码,比如用一个固定的向量代表每个词的位置:
“第1个词”用 [0.1, 0.2, 0.3] 表示,
“第2个词”用 [0.2, 0.3, 0.4] 表示,
这种方法只能告诉模型每个词是第几位,但是不能直接告诉它词与词之间的距离(比如,第1个词和第5个词有多远)。
RoPE(旋转位置编码) 是一种更高级的编码方法,它不仅告诉模型每个词的位置,还能告诉它词与词之间的相对位置关系,并且它对非常长的文本(比如成千上万个词)表现更好。
(2)RoPE 是怎么工作的?(简单理解)
RoPE 使用的是一种数学技巧,叫做“旋转变换”。它把每个令牌的位置用一个角度表示,并通过数学上的正弦(sin)和余弦(cos)函数来编码这些角度。
类比举例:
想象你在一张纸上画一个圆,每个位置对应一个点:
第1个词是“12点钟方向”,
第2个词是“1点钟方向”,
第3个词是“2点钟方向”,
以此类推,每个位置都有一个独特的方向角。
当两个词的方向越接近时,RoPE 能够捕捉到它们之间的距离关系。
(3)RoPE(旋转位置编码)总结:
用数学方法(正弦和余弦)为每个词生成位置信息。
可以捕捉词与词之间的相对关系,而不仅是绝对位置。
非常适合长文本,因为即使句子长度增加,它仍然能保持良好的性能。
表格解读
1. FFN 是 Transformer 模型中每一层的重要组成部分,位于多头注意力机制之后,用来对每个词(令牌)的表示进行非线性变换。它是一个两层的全连接网络(MLP),作用是通过非线性激活函数丰富向量的表达能力。
2. RoPE(θ=500,000 )
旋转频率 是一个参数,用于控制向量每一维度的旋转角度变化的快慢。
3. SwiGLU 激活函数
GLU
GLU(Gated Linear Units,门控线性单元)引入了两个不同的线性层,其中一个首先经过sigmoid函数,其结果将和另一个线性层的输出进行逐元素相乘作为最终的输出:
这里使用σ \sigmaσ作为激活函数,修改改激活函数得到的变体通常能带来更好的性能表现
Swish激活函数的形式为:
SwiGLU 是一种将 Swish 激活函数与 Gated Linear Unit (GLU) 结合的激活函数
3.2.1 缩放法则 (Scaling Laws)
什么是“缩放法则”?
缩放法则(Scaling Laws)是指通过数学模型和实验规律,来研究模型的性能(如预测损失或准确率)如何随着以下几个因素变化:
模型规模(如参数数量或维度)。
训练数据量。
训练计算预算(FLOPs,浮点操作数)。
这些规律的研究帮助我们回答两个重要问题:
给定计算预算,如何确定最优的模型大小?
模型太大:会超出预算,无法完全训练,性能可能受限。
模型太小:无法充分利用预算,导致性能不足。
缩放法则的目标是找到预算和模型规模之间的最佳平衡点。
如何预测模型的性能?
缩放法则可以帮助推测模型在下游任务中的表现(如准确率或损失),即便是预算限制下的计算优化模型。
Llama 3 的挑战:缩放法则预测下游性能
在开发 Llama 3 时,团队面临两大挑战:
现有缩放法则只针对下一词预测损失,而非下游任务性能:
缩放法则通常用于衡量语言模型在语言建模任务中的“预测下一词的准确度”(用负对数似然 NLL 表示)。
但模型在实际使用中(如 QA、文本分类)更关注下游任务的准确率或分数,而不是语言建模任务的损失。
缩放法则的噪声和不确定性:
由于缩放法则通常基于小规模模型和较少的计算预算得出,在大规模模型和高 FLOPs 的训练条件下,它可能无法准确适配。
Llama 3 的开发团队设计了一个两阶段方法
第一阶段:建立 FLOPs 与负对数似然的关系
负对数似然(NLL, Negative Log-Likelihood):
是模型在任务中的损失值,值越小表示模型越擅长在任务中进行预测。
它是衡量语言模型性能的常用指标。
建立关系:
团队通过实验,分析模型在训练过程中:
给定不同的计算预算(FLOPs)。
模型的负对数似然是如何变化的。
目的:
找出计算预算和 NLL 之间的规律,作为第一步的基础。
第二阶段:建立 NLL 与任务准确率的关系
准确率 vs NLL:
团队进一步研究下游任务的负对数似然(NLL)和任务准确率之间的关系。
目的是从 NLL 推测下游任务的表现(如准确率)。
结合缩放法则模型和老模型的训练数据:
团队使用缩放法则预测的新模型。
结合 Llama 2 家族模型(已有数据,更高的训练 FLOPs)。
利用这些数据训练一个更可靠的预测模型。
结果:
得到一种方法,能够在给定特定训练 FLOPs 的情况下预测下游任务的表现。
通过这种方法,我们能够在给定计算优化模型的特定训练 FLOPs 下预测下游任务性能。此外,我们使用类似的方法来选择预训练数据的组合比例
总结
两阶段预测的关键思路:
直接预测困难:
FLOPs 和准确率之间关系复杂。
很难直接预测准确率。
引入中间状态(NLL):
NLL 是一个更容易预测的指标,因为它和 FLOPs 的关系更清晰。
NLL 也能很好地反映模型的性能。
两步推导:阶段 1:找规律,阶段 2:推导性能
FLOPs -> NLL -> 准确率。
核心思想: 利用中间状态(NLL)作为桥梁,间接推导
缩放法则实验
具体来说,我们通过预训练模型,构建了缩放法则。这些实验的计算预算范围从 6 × E18 FLOPs andE22 FLOP。对于每个计算预算,我们在 40M(4000 万)到 16B(160 亿)参数之间选择一部分模型规模进行预训练。在这些训练过程中,我们采用以下设置:
总结
目标:找到这些曲线的最低点,确定在每个计算预算下的最佳训练数据量。
进一步分析:用这些最佳训练数据量的点,画出计算预算(FLOPs)和最佳训练数据量的关系,并尝试用一个函数拟合它。
Figure 3
将 Figure 2 的最低点汇总,展示了 FLOPs 和最佳训练数据量的关系。
提供了一个整体规律,用于预测在给定预算下的最优训练数据量。
ps:
FLOPs(Floating Point Operations) 是一种衡量计算资源的指标,表示训练过程中模型需要进行的浮点运算次数。
它综合反映了模型的规模(参数数量)、训练数据量(tokens),以及计算机硬件资源(如 GPU 性能)。
步骤:在不同的计算预算(FLOPs)下,研究训练数据量(Training Tokens)与验证损失的关系曲线(IsoFLOPs 曲线)。
学习率调度:使用余弦学习率调度(cosine learning rate schedule),并在前 2000 个训练步骤中采用线性预热(linear warmup)。
峰值学习率:根据模型规模的不同,峰值学习率设置在2 × E-4 and 4 × E-4 之间。
余弦衰减:将学习率从峰值逐渐衰减至峰值的 0.1 倍。
权重衰减:每一步的权重衰减值设置为当前学习率的 0.1 倍。
批量大小:在每种计算预算下使用固定的批量大小,范围从 25 万(250K)到 400 万(4M)。
图表解读
Figure 2
横轴:Training Tokens(训练数据量,单位是 token 数)
纵轴:Validation Loss(验证集损失)
表示模型在验证集上的负对数似然(NLL,Negative Log-Likelihood)。值越小,模型在验证集上的表现越好。
线条:不同的计算预算(Compute FLOPs)
图中每条曲线对应一个固定的计算预算,例如:
每条曲线展示在该计算预算下,不同的训练数据量如何影响验证损失。
这张图展示了在不同计算预算下,训练数据量对验证损失的影响。
每条曲线的最低点对应最佳的训练数据量,帮助找到计算优化模型。
Figure 3 在对数坐标下看起来接近“线性”,但这实际上是幂函数关系的特征(在对数坐标下,幂函数是直线)。
随着预算增加,IsoFLOPs 曲线在最低点附近变得越来越平坦,意味着:
对小变化的稳健性:
当预算很高时,即便训练数据量或模型规模稍有偏差,验证损失不会显著增加。
最终决定:
基于这种稳健性,研究者决定训练一个 405B 参数的旗舰模型(略大于外推预测的 402B 参数)。
预测下游任务的性能
两步预测法:
第一步:训练 FLOPs -> NLL,利用线性关系拟合。(仅使用通过缩放法则(Scaling Laws)在 E22 FLOPs 预算内训练的模型以及上述的数据组合。)
第二步:NLL -> Accuracy,利用 S 型关系转换。
验证实验:
在 ARC Challenge 基准任务上进行验证,预测结果与实际结果非常接近,仅略微低估了最终性能。
重要结论:
这种方法不仅有效,还可以跨越多个数量级进行外推预测,为高 FLOPs 模型性能的评估提供了可靠工具。
3.3 基础设施、扩展性和效率
3.3.1 训练基础设施
1. 计算(Compute)
GPU 集群:
使用 16,000 个 NVIDIA H100 GPU,每个 GPU 功耗高达 700W,配备 80GB 高性能显存(HBM3)。
H100 GPU 是目前性能最强的 AI 专用硬件之一,能加速大规模深度学习任务。
服务器配置:
每台服务器配备 8 个 GPU 和 2 个 CPU,GPU 之间通过高速 NVLink 互联。
NVLink 是 NVIDIA 专门设计的高速通信接口,可以大幅提升多 GPU 之间的数据传输速率。
任务调度:
训练任务由 MAST 调度器管理,它是 Meta 自研的分布式任务管理系统,负责在全球范围内调度训练任务,优化资源利用率。
2. 存储(Storage)
Tectonic 文件系统:
Meta 自研的分布式文件系统,提供 240 PB 的存储空间,支持高吞吐量。
适用于深度学习任务中的高频读写操作。
检查点(Checkpointing):
每隔一段时间将模型的状态(权重、参数等)保存到存储中,用于训练中断后的恢复。
检查点数据量非常大(1 MB 到 4 GB 每 GPU),且写入过程通常是突发性的,这对存储架构是一个很大的挑战。
目标是最小化 GPU 在保存检查点时的停机时间,并通过更频繁的检查点写入减少训练中断时的工作损失。
3. 网络(Network)
RoCE(RDMA over Converged Ethernet):
一种高效的网络协议,用于在 GPU 和服务器之间传输大规模数据。
RoCE 使用传统以太网架构,但具有类似 InfiniBand 的高性能特点,允许直接内存访问,减少数据传输延迟。
400 Gbps 的高速互联:
确保 GPU 集群之间能够以极高速度传输数据,这是大规模分布式训练的关键。
多种网络架构优化:
较小的 Llama 3 模型使用 Nvidia Quantum2 Infiniband 架构,而旗舰模型使用 RoCE。
尽管底层网络技术不同,但都经过调优,确保训练性能一致。
4. 网络拓扑结构(Network topology)
Clos 网络是一种分层架构,适合大规模分布式系统,能够提供高带宽和容错能力。
Llama 3 的训练使用了一个三层 Clos 网络结构:
底层(ToR 层):机架内的 GPU 通过一个交换机连接,形成一个小型局域网络。
中层(Cluster Switch 层):将多个机架的 GPU 连接在一起,形成一个 pod(3,072 个 GPU)。
顶层(Aggregation 层):将多个 pod 连接在一起,形成整个集群(24,000 个 GPU)。
特点:
全双工带宽(Full bisection bandwidth):确保 pod 内所有 GPU 都有充足的网络带宽。
超分配(Oversubscription):顶层的带宽是 1:7,即不是所有 GPU 都能同时获得全速的网络带宽。这种设计是为节省网络成本。
大模型训练会产生大量通信开销(例如梯度同步、参数更新等)。
目标:减少不同 pod 之间的通信(因为带宽有限),尽量让计算任务在同一个 pod 内完成。
5. 负载均衡(Load balancing)
问题:为什么传统均衡方法不够?
传统的 ECMP(Equal-Cost Multi-Path 路由)方法:
会将流量均匀分配到不同的网络路径。
但在大模型训练中,数据量过大(fat network flows),一个流就可能占满整个路径,导致负载不均。
解决方法:两种优化技术
多路流量(16 条流):
通过将原本单条流分成 16 条小流,可以更高效地利用网络资源。
每条流的负载减少,减少了网络路径的拥塞。
增强版 ECMP(E-ECMP):
使用 RoCE(RDMA over Converged Ethernet)协议的包头信息进行哈希处理,将这 16 条流量分配到不同路径上。
效果:实现更细粒度的流量均衡。
6. 拥塞控制(Congestion control)
问题:拥塞的原因
在大模型训练中,集体通信模式会产生短时但高频的数据流(例如梯度同步)。
这些流量可能导致网络拥堵,尤其是在不同服务器之间传输数据时。
解决方法:两大策略
深度缓冲交换机(Deep-Buffer Switches):
使用缓冲区较大的交换机,能够存储瞬时流量高峰期间的数据包,减少丢包和网络拥堵。
减少网络拥塞源:
通过负载均衡优化(E-ECMP 和多路流量),显著降低拥塞发生的概率。
最终效果
研究者通过优化,成功运行了一个 24,000 GPU 的大规模集群,而无需依赖传统拥塞控制技术(如 DCQCN)。
表格总结
这张表格描述了 Llama 3 模型在不同阶段的训练配置和性能变化。主要内容包括:
并行策略:通过张量并行、流水线并行和数据并行来实现分布式训练,充分利用 GPU 资源。
序列长度:训练中逐渐增加输入序列长度,最终支持超长上下文(131,072 Tokens)。
GPU 利用率:训练阶段越深入,GPU 利用率略有下降,但可以支持更复杂的任务。
3.3.2 模型扩展的并行化方法
4D 并行化(4D Parallelism)
为了扩展训练以支持我们最大的模型,我们使用了 4D 并行化(4D Parallelism) 方法。这种方法结合了四种不同的并行化技术,将模型高效地分割到多个 GPU 上,并确保每个 GPU 的模型参数、优化器状态、梯度和激活值可以适配到其高带宽存储器(HBM)中。以下是 4D 并行化的组成:
数据并行(Data Parallelism, DP)
定义:将整个数据集分成多个小批次,每个 GPU 负责处理不同的批次。
特点:每个 GPU 运行相同的模型,但使用不同的数据。
例子:如果有 1000 张图片,4 个 GPU,可以让每个 GPU 处理 250 张图片,最后再将结果汇总。
张量并行(Tensor Parallelism, TP)
定义:将模型的张量(权重矩阵)切分到多个 GPU 上,每个 GPU 只负责计算其中一部分。
特点:多个 GPU 需要协同完成每一步的计算。
例子:一个 1000 × 1000 的权重矩阵可以分成 4 块,每个 GPU 计算 250 × 1000 的子矩阵。
流水线并行(Pipeline Parallelism, PP)
定义:将模型的不同层分配到不同的 GPU 上,每个 GPU 只负责一部分层的计算。
特点:不同的 GPU 像“流水线”一样,依次处理数据。
例子:一个 12 层的 Transformer 模型,可以分配 4 层给每个 GPU。
上下文并行(Context Parallelism, CP)
定义:将输入序列(如文本数据)分成多个上下文片段,每个 GPU 负责处理部分片段。
特点:适用于超长序列的情况,避免单个 GPU 因输入序列过长而超出内存限制。
例子:一个 100,000 字符的输入序列,可以分成 4 段,每个 GPU 处理 25,000 字符。
GPU 利用率
通过对并行化配置、硬件和软件的仔细调优,我们在表 4 中的配置下,实现了总体的 BF16 模型浮点运算利用率(Model FLOPs Utilization,MFU) 为 38%-43%。
当 GPU 数量从 8,000 增加到 16,000(数据并行组从 DP=64 增加到 DP=128)时,MFU 从 43% 略微下降到 41%。
原因:为了在训练中保持全局每批次 Token 数(global tokens per batch)不变,每个数据并行组的批次大小(batch size per DP group)需要减少。
流水线并行的改进
我们在现有实现中遇到了几个挑战:
批量大小限制:
当前实现对每个 GPU 支持的批量大小有约束,要求它必须能够被流水线阶段的数量整除。
例如在图 6 中,流水线并行的深度优先调度(DFS,Depth-First Schedule;Narayanan 等人, 2021)要求批次大小 N=PP=4N = PP = 4N=PP=4(即等于流水线阶段的数量)。而广度优先调度(BFS,Breadth-First Schedule;Lamy-Poirier, 2023)要求 N=MN = MN=M,其中 MMM 是总微批次数量,NNN 是同一阶段前向或后向传播的连续微批次数量。
然而,预训练通常需要更大的灵活性来调整批量大小。
内存不平衡:
当前的流水线并行实现会导致资源消耗的不平衡。
第一个阶段由于包含嵌入层(embedding layer)和初始的微批次(warm-up micro-batches),会消耗更多的内存。
计算不平衡:
在模型的最后一层之后,我们需要计算输出和损失,这使得这一阶段成为执行延迟的瓶颈。
灵活的批次大小 N
问题:之前的调度方法(DFS 和 BFS)对批次大小 N 的设置有严格限制,不够灵活。
改进:
修改流水线调度方式,允许批次大小N 更灵活。
例如,在新调度中设置 N=5,可以:
在批次大小受限时,运行比阶段数少的微批次(micro-batches),避免内存不足。
运行更多的微批次来隐藏通信开销(即在计算和通信之间找到平衡点),提高效率。
平衡流水线(减少第一和最后阶段的负担)
问题:
第一个阶段(embedding 层)和最后一个阶段(输出和损失计算)负担过重,导致内存和计算资源不平衡。
改进:
从第一个阶段和最后一个阶段分别减少一个 Transformer 层:
第一个阶段只包含嵌入(embedding)。
最后一个阶段只包含输出投影(output projection)和损失计算(loss calculation)。
这样可以让内存和计算资源在流水线各阶段之间更均衡分配。
减少流水线气泡
问题:
在流水线并行中,"气泡" 是指某些 GPU 在等待其他阶段的结果时处于闲置状态。
改进:
使用交错调度(interleaved schedule),将流水线阶段(Pipeline Stages)分配到多个 GPU 上,使得每个阶段的 GPU 不会因为等待而闲置。
气泡比率的公式为
即气泡的影响随着更多流水线阶段(PP)和微批次(M)的增加而减小。
异步点对点通信
问题:
点对点通信(GPU 之间的数据传输)可能导致训练过程的延迟,尤其是在输入数据中存在不均匀计算(例如文档掩码)的情况下。
改进:
使用异步通信方法,在传输数据的同时进行计算,从而加快训练。
启用
TORCH_NCCL_AVOID_RECORD_STREAMS
,进一步减少异步通信的内存开销。
主动内存管理
问题:
训练过程中,许多张量(tensors)可能会被长时间占用内存,但实际上并不再参与后续计算。
改进:
基于内存分析,主动释放那些不会参与未来计算的张量,包括每个流水线阶段的输入和输出张量。
这种方法有效降低了内存使用,为更长序列的训练留出空间。
针对长序列的上下文并行(Context Parallelism, CP)
我们利用上下文并行(CP)来提高 Llama 3 在扩展上下文长度时的内存效率,从而支持最长 128K Tokens 的超长序列训练。在 CP 中:
我们沿着序列维度进行分区,具体地,将输入序列划分为 2 × CP 个小块(chunks),每个 CP 排名(CP rank)会接收两个块,以实现更好的负载均衡。
第 iii 个 CP 排名同时接收第 iii 个块和第 (2×CP−1−i)(2 × CP - 1 - i)(2×CP−1−i) 个块。
不同于现有的 CP 实现(如 Liu 等人,2023a),它们采用环状结构(ring-like structure),将通信和计算重叠;我们的 CP 实现使用基于全聚合(all-gather)的方法:
首先,对键(Key, K)和值(Value, V)张量进行全聚合(all-gather)操作;
然后,对本地查询(Query, Q)张量块计算注意力输出。
尽管全聚合(all-gather)的通信延迟暴露在关键路径中,我们仍选择这种方法,原因如下:
灵活性:基于 all-gather 的 CP 注意力实现更容易支持不同类型的注意力掩码(如文档掩码)。
延迟影响可忽略:由于我们采用了 分组查询注意力(Grouped Query Attention, GQA),K 和 V 张量的大小远小于 Q 张量。对于注意力计算,时间复杂度为 O(S2)O(S^2)O(S2)(S 表示全因果掩码中的序列长度),而全聚合的时间复杂度为 O(S)O(S)O(S),因此全聚合的通信开销相比注意力计算的时间复杂度可以忽略不计。
基于网络感知的并行配置(Network-aware Parallelism Configuration)
我们对并行维度的顺序 [TP(张量并行), CP(上下文并行), PP(流水线并行), DP(数据并行)] 进行了优化,以减少网络通信开销。
并行顺序的重要性:
内层的并行(innermost parallelism)需要更高的网络带宽和更低的延迟,因此通常局限在同一个服务器内。
外层的并行(outermost parallelism)可能跨越多个网络跳数(multi-hop network),并且需要容忍较高的网络延迟。
基于这些带宽和延迟需求,我们采用 [TP, CP, PP, DP] 的顺序:
TP:最内层并行,要求最快的通信速度。
DP:最外层并行,因为它可以通过异步预取分片的模型权重和归约梯度,容忍更高的延迟。
优化并行配置的挑战:
找到一个在通信开销最小的情况下避免 GPU 内存溢出的并行配置是很有挑战的。
我们开发了一个内存消耗估算器和性能预测工具,这些工具帮助我们探索不同的并行配置,预测整体训练性能,并有效识别内存瓶颈。
数值稳定性(Numerical Stability)
通过比较不同并行配置下的训练损失,我们修复了影响训练稳定性的几个数值问题:
FP32 梯度累积:
在多微批次的反向计算过程中,我们在 FP32 精度下进行梯度累积。
在数据并行的 FSDP 中,我们也在 FP32 精度下对梯度进行 reduce-scatter 操作(即分布式梯度归约)。
中间张量的 FP32 计算:
对于正向计算中多次使用的中间张量(例如视觉编码器的输出),其反向梯度也在 FP32 精度下累积,以确保数值稳定性。
3.3.3 集体通信(Collective Communication)
我们为 Llama 3 开发的集体通信库是基于 Nvidia 的 NCCL 库的一个分支,称为 NCCLX。NCCLX 显著提升了 NCCL 的性能,特别是在高延迟网络中的表现。请回忆一下并行维度的顺序是 [TP, CP, PP, DP],其中 DP 对应于全分片数据并行(FSDP)。外层的并行维度(PP 和 DP)可能需要通过多跳网络(multi-hop network)进行通信,延迟可能高达数十微秒。
原始 NCCL 的通信操作(如 FSDP 中的全收集 all-gather 和减少-分散 reduce-scatter,以及 PP 中的点对点通信)需要进行数据分块(chunking)和分段数据复制(staged data copy)。这种方法导致了一些低效之处,包括:
需要通过网络交换大量的小型控制消息来完成数据传输;
额外的内存复制操作;
使用了额外的 GPU 计算周期进行通信。
在 Llama 3 的训练中,我们通过调整数据分块和数据传输,使其适配我们高达数十微秒延迟的大型集群网络,解决了部分低效问题。我们还允许小型控制消息以更高的优先级通过网络,避免在深缓冲核心交换机中被“队首阻塞”(head-of-line blocking)。对于未来版本的 Llama,我们的工作重点是对 NCCLX 进行更深层次的更改,以全面解决上述所有问题。
这张表格展示了在 Llama 3 405B 模型的 54 天预训练期间,发生的意外中断(unexpected interruptions)的根本原因分类及其对应的中断次数(Interruption Count)和占总中断次数的百分比(% of Interruptions)。
表格解读
Component:导致中断的主要组件,比如 GPU、网络、主机(Host)等。
Category:对组件的归类,标明中断问题是 GPU、网络还是其他依赖(Dependency)问题。
Interruption Count:该组件导致的中断次数。
% of Interruptions:该组件导致的中断次数占总中断次数的百分比。
中断原因分布
GPU相关问题占大头:
故障 GPU(Faulty GPU):148 次,占 30.1%,是最大的中断原因。
GPU 的 HBM3 内存(高速显存):72 次,占 17.2%。
GPU 的 SRAM 内存(更小的高速缓存):19 次,占 4.5%。
GPU 系统处理器(GPU System Processor):17 次,占 4.1%。
GPU 的热接口和传感器(Thermal Interface + Sensor):6 次,占 1.4%。
GPU 的数据损坏问题(Silent Data Corruption):6 次,占 1.4%。
GPU 相关问题总和:约 58.7%。
依赖性问题(Dependency):软件 Bug(12.9%)和网络相关(8.4%)问题也占了一部分。
软件 Bug:54 次,占 12.9%。
网络交换机/电缆(Network Switch/Cable):35 次,占 8.4%。
主机维护(Host Maintenance):
占 7.6%,是由于人为计划的系统维护。
硬件问题主导:
表格底部列出的如 SSD、服务器机箱(Server Chassis)、电源(Power Supply)等占比非常小,但总计约占 4%-5%。
整体来看,78% 的意外中断是由于确认或疑似的硬件问题。
3.3.4 可靠性与操作挑战
16,000 个 GPU 的训练复杂性和潜在故障场景,远远超过我们操作过的更大规模的 CPU 集群。此外,训练的同步特性使其容错能力较低——单个 GPU 的故障可能需要重新启动整个任务。尽管存在这些挑战,我们在 Llama 3 的训练中仍然实现了超过 90% 的有效训练时间,并支持自动化的集群维护操作,例如固件和 Linux 内核的升级(Vigraham 和 Leonhardi,2024),尽管这每天至少会导致一次训练中断。有效训练时间衡量了有意义训练的时间占总耗时的比例。
在 54 天的预训练期间,我们共经历了 466 次任务中断。其中:
47 次是计划内中断,例如自动化维护操作(固件升级)或操作员发起的任务(如配置或数据集更新)。
419 次是意外中断,这些中断如表 5 所示被分类。大约 78% 的意外中断归因于确认的硬件问题(如 GPU 或主机组件故障)或疑似硬件相关的问题(如静默数据损坏或未计划的主机维护事件)。GPU 问题是最大的中断类别,占所有意外问题的 58.7%。
尽管发生了大量故障,但在此期间仅有三次需要人工干预,其余问题均由自动化系统处理。
提升有效训练时间的措施
减少任务启动和检查点保存的时间。
开发快速诊断和问题解决工具:
我们广泛使用了 PyTorch 的内置 NCCL 航行记录器(NCCL Flight Recorder,Ansel 等,2024)。该功能可以将通信元数据和堆栈跟踪记录到环形缓冲区中,从而快速诊断大规模训练中的卡顿和性能问题,尤其是在 NCCLX 的通信操作中。
通过这一工具,我们能够有效记录每次通信事件和每次集体操作的持续时间,并在 NCCLX 超时或心跳超时时自动转储跟踪数据。
我们能够通过在线配置更改(Tang 等,2015)启用更高强度的跟踪操作和元数据收集,而无需重新发布代码或重新启动任务。
大规模训练中的调试难题
NVLink 和 RoCE 的混合使用带来的挑战:
NVLink 上的数据传输通常通过 CUDA 内核发出的加载/存储操作完成,而远程 GPU 或 NVLink 连接中的故障通常表现为 CUDA 内核中未返回明确错误代码的操作卡顿。
NCCLX 通过与 PyTorch 的紧密协作设计,提升了故障检测和定位的速度和准确性。PyTorch 可以访问 NCCLX 的内部状态,并跟踪相关信息。
尽管 NVLink 故障导致的卡顿无法完全避免,但系统会监控通信库的状态,并在检测到此类卡顿时自动超时。
NCCLX 还跟踪每次通信的内核和网络活动,并提供故障通信操作的内部状态快照,包括已完成和待处理的数据传输。
难以检测的硬件问题:
某些硬件问题可能导致“慢速的滞后者”(stragglers),即性能下降但未完全故障的设备。即使是一个滞后者,也可能拖慢数千个其他 GPU,通常表现为正常工作但通信速度较慢。
我们开发了工具,能够优先处理选定进程组中可能存在问题的通信。通过调查少数几个优先级最高的可疑通信,我们通常能够有效识别滞后者。
环境因素对训练性能的影响
昼夜温度变化:
在 Llama 3 405B 的训练中,我们观察到基于一天中不同时间的 GPU 吞吐量波动为 1-2%。这种波动是由于中午较高的温度影响了 GPU 的动态电压和频率调整。
电力波动:
在训练过程中,数万台 GPU 的功耗可能同时升高或降低,例如在所有 GPU 等待检查点保存、完成集体通信,或整个训练任务启动/关闭时。
这可能导致数据中心电力消耗瞬间波动,幅度可达 数十兆瓦,给电网带来压力。这是我们在扩展未来更大规模的 Llama 模型训练时面临的持续挑战。
3.4 训练配置
Llama 3 405B 的预训练配方分为三个主要阶段:
初始预训练;
长上下文预训练;
退火(Annealing)。
以下分别描述了这三个阶段。我们在预训练 8B 和 70B 模型时也使用了类似的配方。
3.4.1 初始预训练
优化器和学习率:我们使用 AdamW 优化器,峰值学习率为 8×E−5,采用线性预热(warm-up)8,000 步,并使用余弦学习率调度器,学习率在 1,200,000 步中衰减到 8×E−7。
批次大小:
初始阶段,我们使用较小的批次大小以提升训练稳定性,随后逐渐增大批次大小以提高效率。
起始批次大小为 4M(400 万)tokens,序列长度为 4,096。
在预训练 252M(2.52 亿)tokens 后,将批次大小和序列长度加倍,变为 8M 的批次和 8,192 tokens 的序列。
在预训练 2.87T(2.87 万亿)tokens 后,再次将批次大小加倍到 16M。
稳定性:我们发现这个训练配方非常稳定,几乎没有损失值的突然上升(spikes),也不需要人工干预来纠正模型训练的偏离。
调整数据组合:
在训练过程中,我们对预训练数据的组合比例进行了多次调整,以改进模型在某些下游任务上的表现。
具体调整:
增加了非英语数据的比例,以提高 Llama 3 的多语言能力;
增加了数学数据的采样权重(upsampling),以提升模型的数学推理能力;
在训练后期加入了更近期的网络数据,更新了模型的知识截止点(knowledge cut-off);
下采样了一些被认为质量较低的训练数据子集。
3.4.2 长上下文预训练
目标:在预训练的最后阶段,我们引入长序列训练,以支持最多 128K(12.8 万)tokens 的上下文窗口。
原因:
我们没有在训练早期使用长序列训练,因为自注意力(self-attention)层的计算量随序列长度的平方增长(即计算复杂度是 O(S2)O(S^2)O(S2))。
过程:
逐步增加支持的上下文长度,训练直到模型能够成功适应增加的上下文长度。
判断成功适应的标准:
短上下文评估中的模型性能完全恢复;
模型能够完美解决“needle in a haystack”(大海捞针)任务,直到该长度。
在 Llama 3 405B 的训练中,我们分 6 个阶段逐步增加上下文长度,从初始的 8K(8,192 tokens)窗口扩展到最终的 128K(128,000 tokens)窗口。
训练数据量:这一阶段使用了大约 800B(8000 亿)训练 tokens。
Ps
“大海捞针”任务是什么?
“大海捞针”任务(Needle in a Haystack)是一种极端情况下的模型测试任务,用来评估模型在超长上下文中找到特定目标信息的能力。具体来说,给模型一段非常长的上下文(比如几万甚至十几万的 tokens),在这段上下文中埋入一个很短的、很稀有的信息(像“针”一样的东西)。模型需要在不忽略大量无关信息的情况下,准确找到这个稀有信息。
例子:
假设给模型一个有 10 万个 tokens 的上下文,其中包括:
一段关于历史的内容(90% 的文本,比如 “人类的第一场工业革命发生在18世纪……”)。
一段关于天气的内容(5% 的文本,比如 “今年的温度比往年偏高……”)。
一句话嵌入在文本中的关键目标信息(比如“密码是: 1234”)。
任务要求模型在整个上下文中找到“密码是: 1234”这一关键句子,并提取出来。
测试目的:
考察模型在极长上下文下的性能。
测试模型的注意力机制(是否能把注意力聚焦在目标信息上,而不是被其他无关信息分散注意力)。
3.4.3 退火阶段(Annealing)
在预训练的最后阶段(最后 40M tokens),我们将学习率线性减小至 0,同时保持上下文长度为 128K tokens。在这个退火阶段,我们还调整了数据组合比例,加大了高质量数据源的采样比例(详见 3.1.3 节)。最后,我们在退火阶段对模型的多个检查点进行了平均(使用 Polyak 平均法,Polyak, 1991),以生成最终的预训练模型。
什么是退火阶段(Annealing)?
“退火”是指在训练的最后阶段,通过逐步减少学习率,来确保模型能够更平稳地收敛到一个好的状态。这类似于物理中的“退火过程”:通过逐步降低温度,让金属或玻璃材料的分子结构更稳定。
在这里,退火阶段是在训练的最后 40M tokens中进行。
学习率线性减小至 0 的目的是让模型的参数在训练的最后阶段微调得更精确,而不是大幅调整,避免模型跳出已经找到的“局部最优解”。
为什么要调整数据组合比例?
在退火阶段,研究者特别强调使用高质量的数据源。这是因为:
在训练的初始阶段,模型需要“大量多样化的数据”来学习基础知识。
而在最后阶段,模型更需要通过“高质量数据”进行精细优化,提升某些特定任务(如数学推理、多语言能力)的表现。
高质量数据的采样比例被加大(upsample),目的是让模型在这段高质量数据上更加专注,提升关键能力。
什么是 Polyak 平均(Polyak Averaging)?
Polyak 平均法 是一种优化技术,目的是通过对多个训练中间结果的平均值来生成更稳定的最终模型。具体来说:
在退火阶段,模型会保存多个“检查点”(Checkpoints),即不同训练步的模型状态。
Polyak 平均会对这些检查点的参数取平均,生成最终的模型。