专栏名称: 极市平台
极市平台是由深圳极视角推出的专业的视觉算法开发与分发平台,为视觉开发者提供多领域实景训练数据库等开发工具和规模化销售渠道。本公众号将会分享视觉相关的技术资讯,行业动态,在线分享信息,线下活动等。 网站: http://cvmart.net/
目录
相关文章推荐
人人都是产品经理  ·  普通打工人专属:如何一键解锁 3 个 AI 副业 ·  2 小时前  
乐游上海  ·  元宵节限定!错过又要等一年 ·  17 小时前  
上海本地宝  ·  微信又有新功能上线! ·  4 天前  
51好读  ›  专栏  ›  极市平台

LLama 405B 技术报告解读

极市平台  · 公众号  ·  · 2024-08-01 22:00

正文

↑ 点击 蓝字 关注极市平台
作者丨包包算法笔记
来源丨包包算法笔记
编辑丨极市平台

极市导读

本文解读了Meta发布的Llama 3 405B版本的技术报告,重点介绍了该模型在预训练、模型结构、规模定律、训练策略和多模态能力集成等方面的创新和优化措施。文章还概述了支持Llama 3训练的硬件设施和网络架构,并讨论了模型的可用性和性能表现。 >> 加入极市CV技术交流群,走在计算机视觉的最前沿

果然传的消息都是真的,meta在24号凌晨发布了llama 3的405B版本,这次还是做一个技术报告解读。值得一提的是,在技术报告的开头,meta特意强调了一个 Managing complexity ,大意是管控复杂度。

为什么没用MoE却弄个405B的dense?为什么没用PPO只用DPO?

meta给的解释是:Managing complexity,大意就是这样简单吧...

评测结果如下,这个结果跟当初网上传的那个版本还是有一定出入的,没有到摁着GPT4o锤的程度。

况且,根据GPT4o的速度来看,参数量要远远小于一个405B的dense,高下立判。不过这个无可厚非,毕竟GPT4也挺慢的

虽然如此,但是llama3 405B中间有许多实用的trick还是值得我们学习的,整个的画风有点像打比赛刷榜那种程度,做的很细,抠每一个上分点,那么我们来总结下一些亮点吧。

一、预训练

数据部分

  • PII数据(个人隐私数据)清洗、去重、去黄、做模型洗数据(fasttext),混合代码和推理数据,多语言数据

数据配比-打标签

  • 做模型做细粒度的打标签工作,然后根据标签采样,配不同的量去试,最终敲定了:50% 通用数据, 25% 数理数据, 17% 代码数据, 8% 多语言数据。实际这个过程比这个深入很多,据消息,国内某知名大模型团队能做到上千级别的细粒度标签实验。

数据配比-scalinglaw

  • 探索数据配比的scaling law,作者这里没有展开,大概的方法可以看我之前写过的文章,在不同的小模型上做不同的配比实验,然后用来预测大模型的最优配比。

数据退火

  • 作者发现在大模型训练的最后阶段,用高质量的数据学习能提高性能。于是在最后40B数据上,作者逐渐将学习率衰减到0。并且发现,数据退火方法,可以用来筛数据,量少,效果明显,实验更高效。

二、模型结构

  • GQA,8-kv head,超大参数的常规操作,具体参数如下。

  • 126 层,层数多两层,这是一个训练阶段方便流水线并行切分的技巧。
  • 长文拼接时候,使用attention mask防止不同来源的数据串味。样本间穿越在预训练阶段影响不大,以前大家也不在乎,但作者说,在扩长序列时影响很大。
  • 词表大小为128K,英语语料的压缩率有所提高,同样计算量能够过更多数据,并且增强了非英语的能力。这里面中文水平如何,还等着大家测试,不过根据下面实验部分的case study看,这次中文能力还是不错的。
  • RoPE theta 调到了500000,再次感谢苏神。

三、scalinglaw

作者说了现有的scalinglaw通常只预测loss,而不是特定的benchmark上的表现;(2)scalinglaw可能因基于小资源进行的预训练运行而变得没那么可靠。

对此,作者搞了一个两步走方法。

1.先建立计算最优模型在下游任务上的负对数似然与训练FLOPs之间的相关性。

2.利用scalinglaw模型和使用更高计算FLOPs训练的旧模型,将下游任务上的负对数似然与benchmark的准确率指标关联上。

作者在这步用了LLama2系列的模型作者在ARC上使用这个方法,能看出来拟合的还不错

四、infra/硬件/网络

  • GPU资源 16K H100 80GB with NVLink,有钱。
  • 专用集群,不是之前的Meta’s AI Research SuperCluster,船新的Meta’s production clusters,240 PB  SSD,7500 台机器,2TB-7TB/s 吞吐,如此高吞吐的存储集群是为了最小化ckpt的IO耗时。

网络部分,专业的机房看管员可以仔细研究下:

  • RoCE,单口400Gb/s

具体的拓扑结构为

  • 3层网络架构,单体24K GPU
  • 最底层1个ToR下2机器配16卡
  • 1 Pod 配192 ToR
  • 3072 张GPU 1:1 收敛比
  • 往上8 个Pod 1:7 的收敛比
  • 由于跨pod 带宽降低,所以模型并行编排和资源调度均需考虑网络架构

负载均衡

  • 两个GPU 间使用16个流,而不是1个,来降低单流流量,以更好地负载均衡
  • 在网络包的头部增加了特殊区域,通过hash 来使得流的选路更加均衡

拥塞控制

  • 在spine 上使用deep-buffer 交换机以适应集合通信导致的短时拥塞,并能够降低慢节点引发持久的拥塞和反压

整体的吞吐如图

别的不说,就这万卡集群的吞吐率,要比国内小作坊的16卡微调个小模型都要高。

模型并行策略:

  • 4D并行,TP + CP + PP + DP,多了一个上下文并行CP,或者称之为序列并行。
  • 使用了FSDP,但是model weight 只拉取一次,以减少反向梯度计算时的weight allgather 通信

PP并行策略改进

  • Batch Size 限制:当前的流水线并行策略 会限制 micro batch 个数为流水线stage 的整数倍,会导致global batch size 和 流水线 stage 相互制约
  • 显存不均衡:第1个stage 会多出很多显存占用,之后逐stage 降低
  • 计算不均衡:最后一层会额外计算 output layer 和 loss,增加了计算量和延时,首尾stage 做了padding

CP并行策略的改进

  • 和Ring CP 一样的在序列维度切分,切分为2*CP 份以负载均衡,但没有用环状通信来overlap 计算和通信,而是基于allgather 的通信。

网络架构感知的并行策略

  • [TP, CP, PP, DP] 的并行策略是针对网络通信优化做专门设计的
  • 开发了显存预估和性能探查工具,用以平衡显存开销和通信性能

数值稳定性

  • FP32 的梯度累加和通信

集合通信

  • 基于NCCL 开发了 NCCLX,在高延迟网络下显著提升性能
  • [TP, CP, PP, DP]  并行策略可能导致PP 和DP 通信跨Pod:原 allgather 和 reducescatter 实现依赖数据chunk 和拷贝,需要大量的小的控制信息的通信,进行额外的拷贝操作,且占用GPU 资源来做通信。对此,llama团队 优化chunk 策略,提升小的控制包的通信优先级。

可用性

不怎么样,运行54天,崩了419次(预期外训练中断),具体的分布如下:

网络问题和GPU问题各占一半。

前面提到了那么多保障和优化措施,还是平均每天挂了8次,大约是90%的可用性。只能说操盘万卡集群不容易呀!

五、预训练

  • 三阶段训练法:(1) 初始训练 initial pre-training, (2) 长文训练 long-context pre-training, and (3) 退火训练 annealing

有点像做菜的感觉...作者特意强调了LLama3增加了非英文部分的比例,增加了数理数据,提高逻辑推理能力。看来这次是铁了心的死磕GPT4了。

初始训练:

  • 余弦调度 8 × 10−5 , 8,000 steps热身, 然后在1,200,000步上降到 8 × 10−7
  • 上下文长度和BS缓慢增加,配4M的bs用4,096长度, 再配8M的bs扩展序列长度为 8,192,这个阶段大约是 252M tokens。最终16M的BS

长上下文训练

  • 从l 8K 逐步增加到 128K

退火训练

  • 最后40M token,用128K长度,逐渐线性缩减学习率到0

六、后训练(对齐)

对话数据格式

  • 多种角色的输入输出的设定和特殊token

RM建模

  • 除了好坏标注之外,增加一路人工编辑,作为最优秀的质量样本 edited > chosen > rejected

SFT

  • 拒绝采样
  • 1e-5 学习率训 8.5K to 9K steps

DPO

  • token屏蔽,一些特殊的字符不参与学习
  • DPO+NLL的loss,在IRPO论文里提到的方法,类似正负样本的精细化调权手段

模型融合

  • 权重平均,RM, SFT, or DPO各个阶段独立融合

迭代式训练

  • 同LLama2,最新的模型采样最新的偏好数据,武当总云梯,左脚踩右脚

数据方面

  • 做的特别细,常用的方法如做细粒度标签采样,做模型打各种角度的分,精细的语义去重等
  • 合成数据
  • “model-as-judge”做模型,去做数据质量筛选

数据部分大体的方法大家都用过,各种角度的数据生产在这篇技术报告中讲的非常细,这里建议大家去看看,不管什么什么O,最终还是洗数据,造数据,选数据。


下面是正文全文的部分,powered by kimi

1 引言

基础模型是设计用来支持大量人工智能(AI)任务的通用模型,它们涵盖了语言、视觉、语音等一个或多个模态。这些模型构成了许多现代AI系统的基础。

现代基础模型的发展包括两个主要阶段:(1) 预训练阶段,在这一阶段,模型通过大规模使用简单任务(如下一个词预测或字幕生成)进行训练;(2) 后续训练阶段,在这一阶段,模型被调整以遵循指令、符合人类偏好,并提高特定能力(例如编码和推理)。

在本文中,我们介绍了一套新的语言基础模型,称为Llama 3。Llama 3模型群原生支持多语言、编码、推理和工具使用。我们最大的模型是一个具有405B参数的密集Transformer,上下文窗口可达128K个token。本文对Llama 3进行了广泛的实证评估。我们发现,Llama 3在众多任务上提供了与领先的语言模型如GPT-4相当的品质。我们公开发布了Llama 3,包括405B参数语言模型的预训练和后训练版本,以及我们的Llama Guard 3模型,用于输入和输出安全。本文还介绍了我们将图像、视频和语音能力通过组合方法集成到Llama 3中的实验结果。我们观察到这种方法在图像、视频和语音识别任务上与最先进技术表现相当。由于这些模型仍在开发中,目前尚未广泛发布。

2 总体概述

Llama 3的模型架构在图1中进行了说明。我们Llama 3语言模型的开发包括两个主要阶段:

  • 语言模型预训练 。我们首先将一个大型多语言文本语料库转换为离散的token,并在生成的数据上预训练一个大型语言模型(LLM),以执行下一个token预测。在语言模型预训练阶段,模型学习语言的结构,并从它“阅读”的文本中获得关于世界大量的知识。为了有效地做到这一点,预训练是在大规模上执行的:我们在15.6T tokens上预训练一个带有405B参数的模型,使用8K tokens的上下文窗口。这个标准预训练阶段之后是一个持续的预训练阶段,它将支持的上下文窗口增加到128K tokens。见第3节详情。
  • 语言模型后训练 。预训练的语言模型对语言有丰富的理解,但尚未遵循指令或以我们期望助手表现的方式行事。我们在几轮中通过人类反馈与模型对齐,每轮都涉及在指令调整数据上进行监督式微调(SFT)和直接偏好优化(DPO;Rafaailov等人,2024年)。在这个后训练阶段,我们还整合了新的能力,如工具使用,并观察到在其他领域,如编码和推理,取得了显著的改进。见第4节详情。最后,安全缓解措施也在后训练阶段纳入模型,详细描述见第5.4节。

我们生成的模型具有丰富的能力集。它们可以用至少八种语言回答问题,编写高质量的代码,解决复杂的推理问题,并即兴或零样本地使用工具。

我们还进行了实验,通过组合方法将图像、视频和语音能力添加到Llama 3中。我们研究的方法包括图28所示的三个额外阶段:

  • 多模态编码器预训练 。我们分别训练图像和语音的编码器。我们在大量图像-文本对上训练图像编码器。这教会了模型视觉内容与自然语言描述之间的关系。我们的语音编码器采用自监督方法训练,通过遮蔽语音输入的部分并尝试通过离散token表示重建被遮蔽的部分。结果,模型学习了语音信号的结构。见第7节和第8节分别关于图像编码器和语音编码器的详细信息。
  • 视觉适配器训练 。我们在预训练的语言模型中训练一个适配器,将预训练的图像编码器整合进去。适配器由一系列交叉注意力层组成,将图像编码器的表示输入到语言模型中。适配器在文本-图像对上进行训练。这将图像表示与语言表示对齐。在适配器训练期间,我们也会更新图像编码器的参数,但我们故意不更新语言模型参数。我们还在图像适配器的基础上训练一个视频适配器,以配对视频-文本数据。这使模型能够跨帧聚合信息。见第7节详情。
  • 语音适配器训练 。最后,我们通过一个适配器将语音编码器整合到模型中,该适配器将语音编码转换为可以直接输入到微调语言模型的token表示。适配器和编码器的参数在监督微调阶段联合更新,以实现高质量的语音理解。在语音适配器训练期间,我们不改变语言模型。我们还集成了一个文本到语音系统。见第8节详情。

我们的多模态实验导致模型能够识别图像和视频的内容,并通过语音界面支持交互。这些模型仍在开发中,尚未准备好发布。

3 预训练

语言模型预训练包括:(1) 大规模训练语料库的整理和过滤;(2) 开发模型架构和相应的规模定律以确定模型大小;(3) 开发大规模高效预训练技术;(4) 制定预训练方案。我们分别介绍以下每个组成部分。

3.1 预训练数据

我们为语言模型预训练创建的数据集来源于包含直至2023年知识的多种数据源。我们对每个数据源应用了多种去重方法和数据清洗机制,以获得高质量的token。我们移除了包含大量个人可识别信息(PII)的领域,以及已知含有成人内容的领域。

3.1.1 网络数据整理

我们使用的大部分数据来源于网络,以下描述了我们的清洗流程。

  • PII和安全过滤 :除了其他缓解措施外,我们实施了过滤器以从可能含有不安全内容或大量PII数据的网站上移除数据,包括根据各种Meta安全标准被评为有害的域名,以及已知包含成人内容的域名。

  • 文本提取和清洗 :我们处理非截断网络文档的原始HTML内容,提取高质量、多样化的文本。为此,我们构建了一个自定义解析器,用于提取HTML内容,并优化了去除样板文本和内容召回的精确度。我们通过人类评估比较了我们解析器的质量,与优化文章类内容的流行第三方HTML解析器相比,发现其性能更优。我们仔细处理了包含数学和代码内容的HTML页面,以保留该内容的结构。我们保留了图像alt属性文本,因为数学内容经常以预渲染图像的形式表示,其中数学内容也提供在alt属性中。我们实验性地评估了不同的清洗配置。我们发现,与普通文本相比,markdown对主要在网络数据上训练的模型性能是有害的,因此我们去除了所有的markdown标记。

  • 去重 :我们在URL、文档和行级别应用了多轮去重:

    • URL级别去重 :我们在整个数据集上执行URL级别的去重,为每个URL对应的页面保留最新版本。
    • 文档级别去重 :我们在整个数据集上执行全局MinHash去重,以移除几乎重复的文档。
    • 行级别去重 :我们执行了与ccNet类似的积极行级别去重。我们移除了在每30M文档的桶中出现超过6次的行。尽管我们的手动定性分析显示,行级别去重不仅移除了诸如导航菜单、Cookie警告等各种网站的剩余样板文本,还移除了频繁的高质量文本,但我们的实证评估显示了强有力的改进。
  • 启发式过滤 :我们开发了启发式规则,以移除额外的低质量文档、异常值和重复过多的文档。一些启发式的例子包括:

    • 我们使用重复n-gram覆盖率比率来移除由重复内容(如日志或错误消息)组成的行。这些行可能非常长且唯一,因此无法通过行去重过滤。
    • 我们使用“脏字”计数来过滤出未被域名阻止列表覆盖的成人网站。
    • 我们使用基于Kullback-Leibler散度的token分布来过滤出与训练语料库分布相比含有过多异常token的文档。
  • 基于模型的质量过滤 :此外,我们尝试应用各种基于模型的质量分类器来选择高质量的token。包括使用fasttext训练的快速分类器,以识别给定文本是否会被维基百科引用,以及更计算密集的基于Roberta的分类器,它们在Llama 2预测上训练。为了基于Llama 2训练质量分类器,我们创建了一个由Llama 2的聊天模型确定是否满足我们描述的质量要求的清洁网络文档训练集。我们使用DistilRoberta为每个文档生成质量分数,以提高效率。我们通过实验评估了各种质量过滤配置的有效性。

  • 代码和推理数据 :类似于DeepSeek-AI等人的工作,我们构建了特定领域的管道来提取代码和与数学相关的网页。具体来说,代码和推理分类器都是基于Roberta的模型,它们在Llama 2注释的网络数据上训练。与上述一般质量分类器不同,我们进行了针对包含数学演绎、STEM领域推理和与自然语言交织的代码的网页的提示调整。由于代码和数学的token分布与自然语言大不相同,这些管道实现了特定领域的HTML提取、定制的文本特征和过滤启发式。

  • 多语言数据 :与我们对英语的处理流程类似,我们实施了过滤器,以从可能包含PII或不安全内容的网站中移除数据。我们的多语言文本处理流程具有几个独特的特点:

    • 我们使用基于fasttext的语言识别模型,将文档分类为176种语言。
    • 我们在每种语言的数据中执行文档级别和行级别的去重。
    • 我们应用语言特定的启发式和基于模型的过滤器,以移除低质量的文档。

此外,我们使用基于Llama 2的多语言分类器对多语言文档进行质量排名,以确保优先考虑高质量内容。我们通过实验确定用于预训练的多语言token数量,平衡模型在英语和多语言基准上的性能。

3.2 模型架构

Llama 3 使用的是一个标准、密集的 Transformer 架构(Vaswani 等人,2017)。与 Llama 和 Llama 2(Touvron 等人,2023a,b)在模型架构方面相比,我们的性能提升主要是由数据质量和多样性的提升以及增加的训练规模所驱动的。

与 Llama 3 相比,我们做了一些小的修改:

  • 我们使用了分组查询注意力(Grouped Query Attention, GQA;Ainslie 等人,2023)和 8 个键值对头来提高推理速度,并在解码过程中减少键值对缓存的大小。
  • 我们使用注意力掩码来防止同一序列内不同文档之间的自注意力。我们发现这个改变在标准预训练中的影响有限,但在继续预训练非常长序列时很重要。

以下是 Llama 3 的关键超参数的概览:

  • 层数:8B 模型有 32 层,70B 模型有 80 层,405B 模型有 126 层。
  • 模型维度:8B 模型为 4,096,70B 模型为 8,192,405B 模型为 16,384。
  • FFN(Feed-Forward Network)维度:8B 模型为 6,144,70B 模型为 12,288,405B 模型为 20,480。
  • 注意力头数:8B 和 70B 模型有 32 个头,405B 模型有 128 个头。
  • 键/值头数:所有模型均为 8。
  • 峰值学习率:8B 模型为 3 × 10^-4,70B 模型为 1.5 × 10^-4,405B 模型为 8 × 10^-5。
  • 激活函数:使用 SwiGLU。
  • 词汇表大小:128,000。
  • 位置嵌入:RoPE(基数频率为 500,000)。

我们还使用了 128K tokens 的词汇表。我们的 token 词汇表结合了来自 tiktoken3 分词器的 100K tokens 和额外的 28K tokens,以更好地支持非英语语言。与 Llama 2 分词器相比,我们的新分词器在样本英语数据上的压缩率从 3.17 提高到 3.94 个字符每个 token。这使得模型在相同的训练计算量下可以“阅读”更多的文本。我们还发现,添加来自特定非英语语言的 28K tokens 既提高了压缩比,也提高了下游性能,而对英语分词没有影响。

我们还增加了 RoPE 的基数频率超参数到 500,000。这使我们能够更好地支持更长的上下文;Xiong 等人(2023)表明这个值对上下文长度达到 32,768 时是有效的。

Llama 3 405B 使用的架构有 126 层,token 表示维度为 16,384,有 128 个注意力头;更详细的信息见表 3。这导致模型大小大约是我们数据上根据规模定律计算出的,对于我们的 3.8 × 10^25 FLOPs 训练预算来说是计算最优的。

3.2.1 规模定律

我们制定了规模定律(Hoffmann 等人,2022;Kaplan 等人,2020)来确定我们的旗舰模型的最佳大小,考虑到我们的预训练计算预算。除了确定最佳模型大小之外,另一个主要挑战是预测旗舰模型在下游基准任务上的性能,由于以下几个问题:(1) 现有的规模定律通常只预测下一个 token 预测损失,而不是特定的基准性能。(2) 规模定律可能是嘈杂和不可靠的,因为它们是基于小计算预算进行的预训练运行开发的(Wei 等人,2022b)。

为了解决这些挑战,我们实施了两阶段方法来开发能够准确预测下游基准性能的规模定律:

  1. 我们首先建立计算最优模型在下游任务上的负对数似然与训练 FLOPs 之间的相关性。
  2. 接下来,我们利用规模定律模型和使用更高计算 FLOPs 训练的旧模型,将下游任务上的负对数似然与任务准确率相关联。

这种方法使我们能够在给定特定数量的训练 FLOPs 的情况下,预测计算最优模型的下游任务性能。我们使用类似的方法选择我们的预训练数据混合(见第 3.4 节)。

规模定律实验。具体来说,我们通过在 6 × 10^18 FLOPs 到 10^22 FLOPs 之间的计算预算下预训练模型来构建我们的规模定律。在每个计算预算下,我们预训练大小在 4000 万到 160 亿参数之间的模型,使用每个计算预算的模型大小的一个子集。在这些训练运行中,我们使用余弦学习率计划,线性预热 2000 个训练步骤。峰值学习率根据模型大小设置在 2 × 10^-4 到 4 × 10^-4 之间。我们将余弦衰减设置为峰值值的 0.1。每个步骤的权重衰减设置为该步骤学习率的 0.1 倍。我们为每个计算规模使用固定的批量大小,范围在 250K 到 4M 之间。

这些实验产生了图 2 中的 IsoFLOPs 曲线。这些曲线上的损失是在一个单独的验证集上测量的。我们使用二度多项式拟合测量的损失值,并确定每个抛物线的最小值。我们把抛物线的最小值称为相应预训练计算预算的计算最优模型。

我们使用这种方法识别出的计算最优模型来预测特定计算预算下的最佳训练 token 数量。为此,我们假设计算预算 C 和最佳训练 token 数量 N?(C) 之间存在幂律关系:

我们使用图 2 中的数据拟合 A 和 α。我们发现 (α, A) = (0.53, 0.29);相应的拟合如图 3 所示。将得到的规模定律外推到 3.8 × 10^25 FLOPs 表明应该在 16.55T tokens 上训练一个 402B 参数模型。

一个重要的观察是,随着计算预算的增加,IsoFLOPs 曲线在最小值附近变得平坦。这意味着旗舰模型的性能对于模型大小和训练 tokens 之间权衡的小幅变化相对稳健。基于这一观察,我们最终决定训练一个具有 405B 参数的旗舰模型。

预测下游任务的性能。我们使用得到的计算最优模型来预测旗舰 Llama 3 模型在基准数据集上的性能。首先,我们在线性相关性上将基准中正确答案的负对数似然与训练 FLOPs 相关联。在这项分析中,我们只使用在上述数据混合上训练到 10^22 FLOPs 的规模定律模型。接下来,我们使用规模定律模型和使用 Llama 2 数据混合和分词器训练的 Llama 2 模型,建立对数似然与准确率之间的 S 形关系。我们在 ARC Challenge 基准上展示了这个实验的结果(见图 4)。我们发现这种两步规模定律预测,它在四个数量级上进行外推,相当准确:它只是略微低估了旗舰 Llama 3 模型的最终性能。

3.3 基础设施、扩展性和效率

我们描述了支持 Llama 3 405B 模型在大规模上进行预训练的硬件和基础设施,并讨论了几项优化措施,这些措施带来了训练效率的提升。

3.3.1 训练基础设施

Llama 1 和 2 模型是在 Meta 的 AI 研究超级集群(Lee 和 Sengupta,2022)上训练的。随着我们进一步扩展,Llama 3 的训练转移到了 Meta 的生产集群(Lee 等人,2024)。这个设置优化了生产级别的可靠性,这在我们扩展训练时至关重要。

  • 计算 :Llama 3 405B 在多达 16K 个 H100 GPU 上进行训练,每个 GPU 运行在 700W TDP 下,配备 80GB HBM3,使用 Meta 的 Grand Teton AI 服务器平台(Matt Bowman,2022)。每个服务器配备八个 GPU 和两个 CPU。在服务器内部,八个 GPU 通过 NVLink 连接。训练作业使用 MAST(Choudhury 等人,2024)调度,这是 Meta 的全球规模训练调度器。

  • 存储 :Tectonic(Pan 等人,2021),Meta 的通用分布式文件系统,被用来为 Llama 3 预训练构建存储网络。它提供了 240PB 的存储空间,由配备 SSD 的 7,500 台服务器支持,并支持持续吞吐量 2TB/s 和峰值吞吐量 7TB/s。一个主要挑战是支持高度突发的检查点写入,这些写入在短时间内饱和存储网络。检查点保存每个 GPU 的模型状态,范围从 1MB 到 4GB 不等,用于恢复和调试。我们的目标是最小化检查点期间的 GPU 暂停时间,并增加检查点频率,以减少恢复后丢失的工作量。

  • 网络 :Llama 3 405B 使用基于 Arista 7800 和 Minipack2 开放计算项目(OCP)机架交换机的 RDMA over Converged Ethernet(RoCE)网络。Llama 3 家族中较小的模型使用 Nvidia Quantum2 InfiniBand 网络进行训练。尽管这些集群的底层网络技术不同,但我们调整了它们,以提供这些大型训练工作负载的等效性能。我们进一步讨论了我们的 RoCE 网络,因为我们完全拥有它的设计。

    • 网络拓扑 :我们的基于 RoCE 的 AI 集群包括 24K GPU,由三层 Clos 网络(Lee 等人,2024)连接。在底层,每个机架托管 16 个 GPU,分布在两台服务器之间,并通过一个 Minipack2 机架顶部(ToR)交换机连接。在中间层,192 个这样的机架通过集群交换机连接,形成一个拥有 3,072 个 GPU 的 pod,具有完整的双工带宽,确保没有过订阅。在顶层,同一数据中心大楼内的八个这样的 pod 通过聚合交换机连接,形成一个 24K GPU 的集群。然而,聚合层的网络连接没有保持完整的双工带宽,而是有一个 1:7 的过订阅比率。我们的模型并行方法(见第 3.3.2 节)和训练作业调度器(Choudhury 等人,2024)都针对网络拓扑进行了优化,旨在最小化跨 pod 的网络通信。
    • 负载均衡 :LLM 训练产生难以使用传统方法(如 Equal-Cost Multi-Path(ECMP)路由)在所有可用网络路径上进行负载均衡的胖网络流。为解决这一挑战,我们采用了两种技术。首先,我们的集合库在两个 GPU 之间创建 16 个网络流,而不是仅仅一个,从而减少了每个流的流量,并提供了更多的流用于负载均衡。其次,我们的 Enhanced-ECMP(E-ECMP)协议通过在 RoCE 数据包头中散列额外的字段,有效地在不同的网络路径上平衡这些 16 个流。
    • 拥塞控制 :我们使用深缓冲交换机在脊柱(Gangidi 等人,2024)上容纳由集体通信模式引起的瞬态拥塞和缓冲。这种设置有助于限制由慢服务器引起的持续拥塞和网络背压的影响,这在训练中很常见。最后,通过 E-ECMP 进行更好的负载均衡显著降低了拥塞的可能性。通过这些优化,我们成功地运行了一个 24K GPU 集群,而无需使用传统的拥塞控制方法,如数据中心量化拥塞通知(DCQCN)。

3.3.2 模型扩展的并行性

为扩展我们最大模型的训练,我们使用 4D 并行性 - 结合四种不同类型的并行性方法 - 来分片模型。这种方法有效地在许多 GPU 上分布计算,并确保每个 GPU 的模型参数、优化器状态、梯度和激活都适合其 HBM。我们在图 5 中展示了 4D 并行性的实现。它结合了张量并行性(TP;Krizhevsky 等人(2012);Shoeybi 等人(2019);Korthikanti 等人(2023))、流水线并行性(PP;Huang 等人(2019);Narayanan 等人(2021);Lamy-Poirier(2023))、上下文并行性(CP;Liu 等人(2023a))和数据并行性(DP;Rajbhandari 等人(2020);Ren 等人(2021);Zhao 等人(2023b))。

张量并行性将个别权重张量分割成不同设备上的多个块。流水线并行性通过层将模型垂直划分为多个阶段,以便不同的设备可以并行处理模型管道的不同阶段。上下文并行性将输入上下文分割成段,减少了非常长序列输入的内存瓶颈。我们使用完全分片的数据并行性(FSDP;Rajbhandari 等人(2020);Ren 等人(2021);Zhao 等人(2023b)),它在处理数据时在多个 GPU 上并行同步每个训练步骤后的模型、优化器和梯度。我们对 Llama 3 使用 FSDP 来分片优化器状态和梯度,但对于模型分片,我们在前向计算后不重新分片,以避免在反向传递期间进行额外的全收集通信。

GPU 利用率。通过仔细调整并行性配置、硬件和软件,我们实现了在表 4 所示配置中 BF16 模型 FLOPs 利用率(MFU;Chowdhery 等人(2023))为 38-43%。在 16K GPU 上,DP=128 的 MFU 略降至 41%,而在 8K GPU 上,DP=64 的 MFU 为 43%,这是由于在训练期间需要保持全局批次中的 token 数量不变,因此每 DP 组的批量大小较小。

流水线并行性改进。我们在使用现有实现时遇到了几个挑战:

  • 批量大小限制 :当前实现对每个 GPU 支持的批量大小有限制,要求它能够被流水线阶段的数量整除。例如,在图 6 中,流水线并行性的深度优先调度(DFS)要求 N = PP = 4,而广度优先调度(BFS;Lamy-Poirier(2023))要求 N = M,其中 M 是微批次的总数,N 是同一阶段的前向或后向的连续微批次的数量。然而,预训练通常需要灵活性来调整批量大小。
  • 内存不平衡 :现有的流水线并行性实现导致资源消耗不平衡。第一阶段由于嵌入和热身微批次而消耗更多的内存。
  • 计算不平衡 :在模型的最后一层之后,我们需要计算输出和损失,这使得这个阶段成为执行延迟的瓶颈。

为解决这些问题,我们修改了流水线计划,如图 6 所示,它允许灵活地设置 N —— 在这种情况下 N = 5,可以运行每个批次中的任意数量的微批次。这使我们能够:(1) 当我们在大规模时遇到批量大小限制时,运行少于流水线阶段数量的微批次;或者 (2) 运行更多的微批次以隐藏点对点通信,找到 DFS 和 BFS 最佳通信和内存效率之间的最佳平衡点。为了平衡流水线,我们分别从第一和最后阶段减少了一个 Transformer 层。这意味着第一阶段的第一个模型块只有嵌入,最后阶段的最后一个模型块只有输出投影和损失计算。为了减少流水线气泡,我们使用了一个交错计划(Narayanan 等人,2021),在 V 个流水线阶段上使用一个流水线等级。总体流水线气泡比率是 ( (PP-1) \times \frac{V \times M}{M+1} )。此外,我们在 PP 中采用了异步点对点通信,这在训练中大大加快了速度,特别是在文档掩码引入额外计算不平衡的情况下。我们启用了 TORCH_NCCL_AVOID_RECORD_STREAMS 来减少异步点对

点通信的内存使用。最后,为了减少内存成本,基于详细的内存分配分析,我们主动释放了未来计算中不会使用的张量,包括每个流水线阶段的输入和输出张量。

上下文并行性用于长序列。我们利用上下文并行性(CP)在扩展 Llama 3 的上下文长度时提高内存效率,并支持训练长达 128K 的极长序列。在 CP 中,我们跨序列维度进行分割,具体来说,我们将输入序列分割成 2 × CP 块,以便每个 CP 等级接收两个块以实现更好的负载均衡。第 i 个 CP 等级接收第 i 个和第 (2 × CP − 1 − i) 个块。

与现有的 CP 实现不同,它们在环状结构中重叠通信和计算(Liu 等人,2023a),我们的 CP 实现采用了基于 all-gather 的方法,我们首先 all-gather 键(K)和值(V)张量,然后计算局部查询(Q)张量块的注意力输出。尽管 all-gather 通信延迟在关键路径上暴露出来,但我们仍然采用这种方法,主要有两个原因:(1) 在 all-gather 基础的 CP 注意力中,支持不同类型的注意力掩码(如文档掩码)更容易、更灵活;(2) 由于使用 GQA(Ainslie 等人,2023),通信的 K 和 V 张量比 Q 张量小得多,因此注意力计算的时间复杂度比 all-gather 大一个数量级(O(S^2) 相对于 O(S),其中 S 表示全因果掩码中的序列长度),使 all-gather 延迟可以忽略不计。

网络感知并行性配置。并行性维度的顺序 [TP, CP, PP, DP] 针对网络通信进行了优化。内部并行性需要最高的网络带宽和最低的延迟,因此通常限制在同一个服务器内。外部并行性可以跨越多跳网络,并应容忍更高的网络延迟。因此,基于对网络带宽和延迟的要求,我们按照 [TP, CP, PP, DP] 的顺序放置并行性维度。DP(即 FSDP)是外部并行性,因为它可以通过异步预取分片模型权重和减少梯度来容忍更长的网络延迟。识别具有最小通信开销的最优并行性配置,同时避免 GPU 内存溢出是具有挑战性的。我们开发了一个内存消耗估计器和性能预测工具,帮助我们探索各种并行性配置,并有效地预测整体训练性能,识别内存差距。

数值稳定性。通过比较不同并行性设置之间的训练损失,我们确定了影响训练稳定性的几个数值问题。为确保训练收敛,我们在多个微批次的反向计算中使用 FP32 梯度累积,并在 FSDP 中跨数据并行工作进行 FP32 reduce-scatter 梯度。对于在前向计算中多次使用的中间张量,例如视觉编码器输出,反向梯度也以 FP32 累积。

3.3.3 集体通信

我们的 Llama 3 集体通信库基于 Nvidia 的 NCCL 库的一个分支,称为 NCCLX。NCCLX 显著提高了 NCCL,特别是对于更高延迟网络的性能。回想一下,平行性维度的顺序是 [TP, CP, PP, DP],其中 DP 对应于 FSDP。最外层的平行性维度,PP 和 DP,可能通过多跳网络通信,延迟可能高达数十微秒。原始的 NCCL 集合体 —— FSDP 中的 all-gather 和 reduce-scatter,以及 PP 中的点对点 —— 需要数据分块和阶段性数据复制。这种方法引入了几个效率低下的问题,包括:(1) 需要在网络上交换大量的小控制消息以促进数据传输;(2) 额外的内存复制操作;(3) 使用额外的 GPU 周期进行通信。对于 Llama 3 训练,我们通过调整分块和数据传输以适应我们的网络延迟,解决了这些问题的子集,这些网络延迟可能高达大型集群的数十微秒。我们还允许小控制消息以更高的优先级穿越网络,特别是避免在深缓冲核心交换机中被阻塞。我们正在进行的 Llama 版本未来的工作涉及在 NCCLX 中进行更深入的更改,以全面解决前述所有问题。

3.3.4 可靠性和操作挑战

16K GPU 训练的复杂性和潜在故障场景超过了我们操作过的更大的 CPU 集群。此外,训练的同步性质使其对故障的容忍度更低 —— 一个 GPU 故障可能需要整个作业重新启动。尽管存在这些挑战,对于 Llama 3,我们实现了超过 90% 的有效训练时间,同时支持自动化集群维护,如固件和 Linux 内核升级(Vigraham 和 Leonhardi,2024),这导致至少每天有一次训练中断。有效训练时间衡量的是在经过时间中用于有用训练的时间。

在 Llama 3 405B 预训练的 54 天快照期间,我们经历了总共 466 次作业中断。其中,47 次是计划内的中断,由于自动化维护操作,如固件升级或操作员启动的操作,如配置或数据集更新。其余 419 次是意外的中断,如表 5 所示。大约 78% 的意外中断归因于确认或疑似硬件问题,如 GPU 或主机组件故障,或疑似与硬件相关的问题,如静默数据损坏和计划外的单个主机维护事件。GPU 问题是最大部分,占所有意外问题的 58.7%。尽管失败数量众多,但在此期间,仅需要三次重大手动干预,其余问题由自动化处理。

为了增加有效训练时间,我们减少了作业启动和检查点时间,并开发了快速诊断和问题解决工具。我们广泛使用 PyTorch 内置的 NCCL 飞行记录器(Ansel 等人,2024),这是一个在环形缓冲区中捕获集体元数据和堆栈跟踪的功能,使我们能够在大规模上快速诊断挂起和性能问题,特别是与 NCCLX 相关的。使用此功能,我们有效地记录了每个通信事件和每个集体操作的持续时间,并且在 NCCLX 看门狗或心跳超时时自动转储跟踪数据。我们通过在线配置更改启用了更多计算密集型的跟踪操作和元数据收集,并且根据需要实时在生产中进行,无需代码发布或作业重新启动。

调试大规模训练中的问题由于 NVLink 和 RoCE 在我们的网络中的混合使用而变得复杂。通过 CUDA 内核发出的 load/store 操作通常进行 NVLink 上的数据传输,远程 GPU 或 NVLink 连通性的故障通常表现为 CUDA 内核中的 load/store 操作停滞,而没有返回清晰的错误代码。NCCLX 通过与 PyTorch 的紧密设计,提高了故障检测和定位的速度和准确性,允许 PyTorch 访问 NCCLX 的内部状态并跟踪相关信息。虽然不能完全防止由于 NVLink 故障引起的停滞,但系统会在检测到此类停滞时自动超时。此外,NCCLX 跟踪每次 NCCLX 通信的内核和网络活动,并提供失败的 NCCLX 集合体的内部状态快照,包括所有等级之间完成和待定的数据传输。我们分析这些数据以调试 NCCLX 扩展问题。

有时,硬件问题可能导致仍然运行但运行缓慢的落后者,这很难检测。即使一个落后者也可以减慢数千个其他 GPU,通常表现为通信功能正常但速度慢。我们开发了工具,以优先处理选定进程组中可能存在问题的通信。通过仅调查少数主要嫌疑人,我们通常能够有效地识别落后者。

一个有趣的观察是环境因素对大规模训练性能的影响。对于 Llama 3 405B,我们注意到基于一天中的时间,吞吐量有 1-2% 的波动。这种波动是由于中午温度较高影响 GPU 动态电压和频率缩放。在训练期间,数万个 GPU 可能同时增加或减少功耗,例如,由于所有 GPU 等待检查点完成或集体通信结束,或者整个训练作业的启动或关闭。当这种情况发生时,它可能导致数据中心的功耗瞬间波动数十兆瓦,拉伸电网的极限。随着我们为未来的更大 Llama 模型扩展训练,这是我们面临的一个持续挑战。

3.4 训练方案

用于预训练 Llama 3 405B 的方案包括三个主要阶段:(1) 初始预训练,(2) 长上下文预训练,和 (3) 退火。下面分别描述这三个阶段。我们使用类似的方案来预训练 8B 和 70B 模型。

3.4.1 初始预训练

我们使用余弦学习率计划对 Llama 3 405B 进行预训练,峰值学习率为 8 × 10^-5,线性预热 8,000 步,然后在 1,200,000 个训练步骤中衰减到 8 × 10^-7。我们在训练初期使用较小的批量大小以提高训练稳定性,并随后增加它以提高效率。具体来说,我们最初使用 4M tokens 的批量大小和 4,096 的序列长度,然后在预训练了 252M tokens 后,将这些值翻倍到 8M tokens 和 8,192 的序列长度。我们再次将批量大小翻倍到 16M,达到 2.87T tokens 的预训练量。我们发现这种训练方案非常稳定:我们观察到很少的损失峰值,并且不需要干预来纠正模型训练发散。

调整数据混合。我们在训练期间对预训练数据混合进行了几次调整,以提高模型在特定下游任务上的性能。特别是,我们增加了非英语数据在预训练中的比例,以提高 Llama 3 的多语言性能。我们还上采样数学数据以提高模型的数学推理性能,我们在预训练的后期增加了更多的最新网络数据以推进模型的知识截止日期,并且我们对后来被识别为质量较低的预训练数据子集进行了下采样。

3.4.2 长上下文预训练

在预训练的最后阶段,我们训练长序列以支持长达 128K tokens 的上下文窗口。我们不在早期训练长序列,因为自注意力层中的计算在序列长度上呈二次方增长。我们以增量方式增加支持的上下文长度,直到模型成功适应增加的上下文长度。我们通过测量以下两点来评估成功的适应:(1) 模型在短上下文评估中的性能是否完全恢复;(2) 模型是否完美解决了长达该长度的“大海捞针”任务。在 Llama 3 405B 预训练中,我们从原始的 8K 上下文窗口开始,逐步增加上下文长度,分为六个阶段,最终达到 128K 上下文窗口。这个长上下文预训练阶段使用了大约 800B 训练 tokens。

3.4.3 退火

在预训练的最后 40M tokens 期间,我们线性地将学习率退火到 0,同时保持 128K tokens 的上下文长度。在退火阶段,我们还调整了数据混合,以提高选择特定领域的高质量数据的比例。最后,我们计算了退火期间模型检查点的平均值(Polyak(1991)平均),以产生最终的预训练模型。

3.4 训练方案

用于预训练 Llama 3 405B 的方案由三个主要阶段组成:(1) 初始预训练,(2) 长上下文预训练,和 (3) 退火。下面将分别描述这三个阶段。我们也使用类似的方案来预训练 8B 和 70B 模型。

3.4.1 初始预训练

我们使用余弦学习率计划对 Llama 3 405B 进行初始预训练,峰值学习率为 (8 \times 10^{-5}),线性预热 8,000 步,然后在 1,200,000 个训练步骤中衰减到 (8 \times 10^{-7})。我们在训练初期使用较小的批量大小以提高训练稳定性,并随后增加它以提高效率。具体来说,我们最初使用 4M tokens 的批量大小和 4,096 的序列长度,然后在预训练了 252M tokens 后,将这些值翻倍到 8M tokens 和 8,192 的序列长度。我们再次将批量大小翻倍到 16M,达到 2.87T tokens 的预训练量。我们发现这种训练方案非常稳定:我们观察到很少的损失峰值,并且不需要干预来纠正模型训练发散。

调整数据混合。我们在训练期间对预训练数据混合进行了几次调整,以提高模型在特定下游任务上的性能。特别是,我们增加了非英语数据在预训练中的比例,以提高 Llama 3 的多语言性能。我们还上采样数学数据以提高模型的数学推理性能,我们在预训练的后期增加了更多的最新网络数据以推进模型的知识截止日期,并且我们对后来被识别为质量较低的预训练数据子集进行了下采样。

3.4.2 长上下文预训练

在预训练的最后阶段,我们训练长序列以支持长达 128K tokens 的上下文窗口。我们不在早期训练长序列,因为自注意力层中的计算在序列长度上呈二次方增长。我们以增量方式增加支持的上下文长度,直到模型成功适应增加的上下文长度。我们通过测量以下两点来评估成功的适应:(1) 模型在短上下文评估中的性能是否完全恢复;(2) 模型是否完美解决了长达该长度的“大海捞针”任务。在 Llama 3 405B 预训练中,我们从原始的 8K 上下文窗口开始,逐步增加上下文长度,分为六个阶段,最终达到 128K 上下文窗口。这个长上下文预训练阶段使用了大约 800B 训练 tokens。

3.4.3 退火

在预训练的最后 40M tokens 期间,我们线性地将学习率退火到 0,同时保持 128K tokens 的上下文长度。在退火阶段,我们还调整了数据混合,以提高选择特定领域的高质量数据的比例。最后,我们计算了退火期间模型检查点的 Polyak 平均值,以产生最终的预训练模型。

4.1 建模







请到「今天看啥」查看全文