万字干货！手把手教你如何训练超大规模集群下的大语言模型

分布式训练的主要难点

简单介绍一下混合并行中经典的三种并行方案。首先是数据并行，简称 DP。正如其名，数据并行是将数据分割到不同的计算设备上，然后由这些设备完成各自的计算任务。第二种是张量并行，简称 TP。张量并行是将模型中某些层的参数分散到不同的设备上，每个设备负责完成部分的计算工作。第三种是流水并行，简称 PP。流水并行是将模型的不同层切分到不同的计算设备上，类似于流水线的工作方式，各个设备协同完成整个模型的计算过程。

大模型训练在超大规模集群下

的挑战与解决方案

随着模型规模和集群规模的扩大，通信在训练过程中的占比越来越大。为了更直观地展示这一现象，我提供了两张时间线图，它们没有应用计算通信重叠技术。第一张图突出显示了在实现 DP 重叠前的数据并行通信状态，第二张图则突出显示了在实现 TP 重叠前的张量并行通信情况。

从图中我们可以看到，在端到端的训练过程中，DP 的通信占比实际上超过了 15%，而 TP 的通信时间占比也超过了 30%。因此，减少通信对训练的影响，对提升训练效率至关重要。

DP Overlap

DP overlap 的方案在理论上看起来非常吸引人，但实际应用中，我们真的能显著提升训练效率吗？在进行 DP overlap 优化时，我们遇到了三个主要问题。首先，是通信和计算资源之间的竞争问题。当通信和计算操作同时进行时，它们会争夺有限的硬件资源，这可能会影响整体的系统性能。其次，在混合并行场景下，DP overlap 还可能带来 PP bubble 的问题。第三，不同通信资源的争抢还可能导致网络拥塞。

DP overlap 引入的 PP bubble 问题。在前面，我们讨论了通信对计算效率的影响。如果我们模仿 ZeRO 的调度策略，由于 overlap 计算的时间会长于 none overlap 计算的时间，这种负载不均衡会导致 PP bubble 的产生。即图中的 Micro batch 2 的前向传播和 Micro batch 1 的反向传播较长的现象，这展示了负载不均的情况。我们提出的解决方案是通信时机的纵向对齐，这样可以极大地缓解 PP bubble 的问题。同时需要强调的是，从计算 overlap 部分移出来的通信都被放在了 PP bubble 上，因此它不会产生任何额外的影响。这种策略有助于平衡负载，减少因通信和计算不匹配而产生的效率损失。

下图展示了我们最终优化后的 timeline。在这个优化版本中，我们实现了 reduce-scatter 与反向传播的 overlap，同时 all-gather 操作与前向传播也实现了 overlap。此外，我们通过分桶通信、网络预测控制、通信 CHANNEL 调优以及通信时机的纵向对齐等方法，大幅优化了 DP 的通信开销。这些优化措施共同作用，提高了整体的训练效率，减少了因通信而产生的延迟和资源浪费。

TP Overlap

然后我们可以将这种策略推广到四个 rank 的场景中。为了简化表述，我们将计算的 stream 都合并到了一起。对于 all-gather overlap GEMM，我们会特别关注第一个 rank。第一步，我们使用自己持有的那部分输入来进行计算，同时将自己持有的内容也发送给其他 rank，并接收其他 rank 中持有的那部分输入。接下来的第二步、第三步、第四步都是按照相同的原理进行。通过这种方式，我们就可以得到一个 all-gather 的 overlap 流程。这样，每个 rank 都在进行本地计算的同时，与其他 rank 进行数据交换，实现了计算与通信的重叠。这种策略可以有效地减少等待时间，提高资源利用率，从而提升整体的并行计算效率。

Reduce scatter 的操作也是类似的。我们可以首先关注 rank 4 在整个计算结果流程中的作用。在第一步中，rank 4 的计算结果被放置在 rank 1 上。rank 1 完成自己的计算后，在第二步中，它会将这个结果发送给 rank 2。rank 2 在接收到来自 rank 1 的结果后，会将其与自己的计算结果进行累加，然后继续进行下一步的计算。接着，在第三步和第四步中，流程与前两步相同。rank 3 和 rank 4 也会按照这个顺序接收之前 rank 传递的结果，并与自己的计算结果进行累加。最终，在流程的最后， rank 4 将拿到汇总后的最终结果。

通过上述步骤，我们得到了一个完整的解决方案，适用于处理通信和计算存在依赖关系时的通信计算重叠问题。

这是 TP overlap 的整体解决方案，对于计算通信没有依赖的情况，这里是指 column-wise linear 的反向传播。由于这部分操作没有数据依赖关系，我们采用了 bulk overlap 技术。对于其余的通信和计算，因为它们之间存在依赖关系，我们采用了 split pipeline overlap 的方法。

下图展示了实现 TP overlap 后的 timeline，我们可以看到 TP 的通信和计算重叠在了一起。同时，我们进行了两项优化措施：第一项是使用了 peer-to-peer memory copy，以此来减轻通信对 SM 的消耗。第二项优化是将计算分散到不同的 stream 上，这样计算也可以实现部分的重叠。

超长文本场景解决方案

针对 TP 作为通信换显存的两大弊端——在 h 维度上切分导致的不可扩展性以及方案本身的通信量大，我们希望找到一种在 s 维度上可以切分并且通信量相比 TP 小一些的方案。为此，我们实现了上下文并行（context parallel，简称 CP）。

在 CP 场景下，整个模型的 activation 从始至终都在 s 维度上保持着切分状态。之前无法解决的问题，通过 CP=4 就可以解决。我们可以计算这个方案的通信开销，CP 引入的通信开销仅有 KV 前向时的 all-gather 和反向时的 all-gather 以及 reduce-scatter。同时，我们改变了 QKV 的计算顺序，使得 K 的通信可以与 V 的计算重叠，V 的计算可以与 Q 的计算重叠。因此，我们可以得出下述两个结论。

1. CP 的通信量与 KV 的 activation 大小成正比。在混合并行场景下，利用了 TP 可以减少 activation 大小的特点，使得 CP 的通信量相比于直接扩大 TP 可以减少 TP 倍。
2. 由于 CP 的通信可以与计算进行重叠，因此进一步减少了对训练的影响。同时，由于 CP 的切分维度在 s 上，理论上如果有足够的机器，CP 可以解决任意大小的上下文窗口问题。

CP 与其他技术结合时，会带来一些额外的好处和挑战。首先是计算负载均衡问题，这个问题的背景是大语言模型采用了 Decoder Only 架构，并且在 attention 中使用了 causal mask，这导致 CP 会引入计算负载不均的问题。从下面的左图中可以看到，rank 0 的计算负载明显低于 rank 1。

为了解决这个问题，我们采用了类似高斯求和的方法，让每个设备负责一大一小两个 attention 的计算，以此来缓解负载不均的问题。由于同一个设备上的这两个 attention 计算之间不存在依赖关系，为了进一步提升硬件利用率，我们仿照 TP overlap，使用了不同的 CUDA stream 来 launch 两个 kernel。借助 CUDA 的 runtime 调度，我们实现了更高效的并行计算。

结合 CP 还有一些额外的好处。GQA（Grouped Query Attention）是在长上下文场景下几乎必选的技术。与原来的 Multihead attention 相比，GQA 将多个 query 作为一个 group，每个 group 对应一个 K 和 V。可以发现，GQA 可以极大地减少 KV activation 的大小。正如之前提到的，CP 的通信量与 KV 的 activation 大小成正比。因此在 GQA 的场景下，我们可以进一步减少 CP 的通信量，这是结合使用 CP 和 GQA 技术的一个显著优势。

下面是关于计算换显存的方案，其中 recomputing 是一个非常经典的技术。首先，让我们对 recomputing 做一个介绍。下图展示了一个正常的模型训练过程中的数据流。由于反向传播计算对前向传播计算结果存在数据依赖，因此在前向计算完成后，计算结果并不会立即释放，而是要等到反向计算完成后才释放。

右侧的图展示了使用 recomputing 方案的情况。可以看到，在 0 到 3 层的中间结果被释放了，只有 recomputing block 的输入，也就是 layer 0 的 input 被保存了下来。在反向传播过程中，我们会使用保存下来的 input 重新计算 0 到 3 层的前向传播结果，然后再进行反向计算，从而达到节省显存的目的。这个方案从理论上看起来非常理想，但在实际应用中也会遇到一些问题。

首先，主流的框架都采用了 full computing，这导致每次反向计算都会执行一次完整的 forward pass，引入了大量的无效计算。在大模型时代，这种情况是不可接受的。其次，目前的开源框架 Megatron-LM 对 attention 部分实现了 selective recomputing。然而，在 flash attention 时代，这个方案的效率已经不如以前了。

经过观察，我发现某些 kernel，例如 GEMM，其反向计算实际上并不依赖于前向传播的输出结果。例如，对于公式 𝑌=𝑋𝑊，𝑑𝑋 和 𝑑𝑊 的计算并不依赖于前向传播的结果 𝑌。如果我们将这类算子作为 recomputing block 中的最后一个算子，就无需对它们进行重计算。

大家可以看下图右侧。假设层 3 是一个 GEMM 操作，那么 layer 3 的反向计算只依赖于层 3 的输入，而不是层 3 的输出。这样，在重计算时，我们可以省去 layer 3 的前向计算。我将这种重算策略称为 GEMM last recomputing。

我们将 GEMM last recomputing 策略实施到大语言模型的训练中，发现只需要对计算量较小的算子进行重算。相比于没有采用 recomputing 的方案，我们的策略在增加了不到 1.5% 的计算量的情况下，减少了 40% 的显存开销。这是一个在保持计算效率的同时显著减少显存需求的有效方法。

接下来是内存换显存的方案。我最初产生这个想法的原因是，在训练过程中，显存资源已经非常紧张，然而内存资源在训练状态下却几乎处于闲置状态，这为我们提供了一定的操作空间。其次，随着硬件的升级，PCIe 已经升级到第五代，每张卡分配到的 x16 带宽达到了 64GB/s。同时，由于 H2D（Host to Device）和 D2H（Device to Host）是 memory copy 操作，它们对计算的影响几乎可以忽略不计。在混合并行场景下，每次前向计算产生的 activation 并不会立即被使用，而是至少要间隔一个完整的虚拟 pipeline stage 计算，因此混合并行架构也为我们提供了足够的时间窗口。

我们的解决方案是，将每个虚拟 pipeline stage 前一个 micro batch 的 activation H2D 和 D2H 的通信操作与下一个 micro batch 的计算进行 overlap，这样可以极大减少 offload 对关键路径上计算的影响。通过这个 offload 方案，我们能够在几乎不影响计算性能的情况下实现内存换显存的效果，上下文窗口大小提升了 2.5 倍。

接下来展示的是这个解决方案的整体成果。我们在 H800 集群上进行的测试显示，在吞吐量上，与现有的最先进开源方案相比，我们在** 任意上下文窗口下都能实现超过 30% 的性能提升。** 能达到这样的性能提升主要归功于两点原因：

• 第一，我们采用了通信代价更小的 CP 来替代 TP，从而降低了为解决显存问题而引入的通信开销；
• 第二，我们采用了 GEMM last recomputing 和 pipeline aware offloading 这两种更具成本效益的显存问题解决方案，减少了以通信换取显存的需求，进一步实现了训练吞吐量的提高。

在支持的序列长度上限方面，首先，我们通过内存换显存、通信换显存、计算换显存的方法，大幅提升了单个设备支持的上下文窗口。同时，由于该方案还具有极强的可扩展性，因此在设备资源充足的情况下，我们可以支持无限大的上下文窗口。

接下来是 cost model（成本模型）的介绍。在进行大模型训练时，参数调整是一个非常痛苦的过程，因为模型有大量的参数，并且这些参数之间相互影响，比如 TP、CP、DP 的大小，以及 offload 的比例，还有网络设置中的 CTS。如果对所有参数都进行实际运行测试，将会消耗大量的计算资源。然而，如果不进行实际运行，仅仅通过比例和一些基于 FLOPs 理论算力的简单折算来预测，会导致预测极其不准确。因此，这样的成本模型是不可行的。

为了解决这个问题，我们对 TP、CP、PP 等一系列可能影响性能的因素进行了细致的建模。我们将所需信息分为与模型相关的信息，比如不同组合下单层前向和反向传播的时间；以及与集群相关的信息，比如跨机器的集群通信带宽或者 H2D 的带宽等。整体的测量耗时可以在一个小时内完成，并且这些信息可以多次复用。

在 175b 的案例中，我们建模的预测值和实测值之间的误差控制在 2% 以内。在实际使用过程中，我们的成本模型的误差与实测值的对比也不超过 5%，其中大部分误差来源于网络的不稳定性。下图右边展示了我们的成本模型给出的参数配置表。通常情况下，搜索完成后，我们可以根据 MFU 的前 5 名进行实际测试，最终得到我们的训练配置。这种方法大大提高了参数调整的效率和准确性。

未来展望

未来在训练引擎方面我们会专注于五个主要方向。

1. 万亿参数规模的 MoE 模型：我们期望能够训练具有万亿参数的 MoE 模型，这将推动模型容量和性能的显著提升。
2. 继续扩大序列长度：我们希望能够支持达到百万级别的序列长度，这将极大地扩展模型处理长文本数据的能力。
3. RLHF 框架：目前还没有看到非常高效的 RLHF 框架实现，这将是未来研究的一个重要领域。
4. 低精度训练：随着 Hopper 系列架构的推广以及 FP8、FP6 等多精度配置训练，我们将需要关注低精度训练技术的发展。
5. 异构算力的引入：我们需要考虑引入异构算力来增强训练引擎的灵活性和健壮性。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

第一阶段（10天）：初阶应用

该阶段让大家对大模型 AI有一个最前沿的认识，对大模型 AI 的理解超过 95% 的人，可以在相关讨论时发表高级、不跟风、又接地气的见解，别人只会和 AI 聊天，而你能调教 AI，并能用代码将大模型和业务衔接。

• 大模型 AI 能干什么？
• 大模型是怎样获得「智能」的？
• 用好 AI 的核心心法
• 大模型应用业务架构
• 大模型应用技术架构
• 代码示例：向 GPT-3.5 灌入新知识
• 提示工程的意义和核心思想
• Prompt 典型构成
• 指令调优方法论
• 思维链和思维树
• Prompt 攻击和防范
• …

第二阶段（30天）：高阶应用

该阶段我们正式进入大模型 AI 进阶实战学习，学会构造私有知识库，扩展 AI 的能力。快速开发一个完整的基于 agent 对话机器人。掌握功能最强的大模型开发框架，抓住最新的技术进展，适合 Python 和 JavaScript 程序员。

• 为什么要做 RAG
• 搭建一个简单的 ChatPDF
• 检索的基础概念
• 什么是向量表示（Embeddings）
• 向量数据库与向量检索
• 基于向量检索的 RAG
• 搭建 RAG 系统的扩展知识
• 混合检索与 RAG-Fusion 简介
• 向量模型本地部署
• …

第三阶段（30天）：模型训练

恭喜你，如果学到这里，你基本可以找到一份大模型 AI相关的工作，自己也能训练 GPT 了！通过微调，训练自己的垂直大模型，能独立训练开源多模态大模型，掌握更多技术方案。

到此为止，大概2个月的时间。你已经成为了一名“AI小子”。那么你还想往下探索吗？

• 为什么要做 RAG
• 什么是模型
• 什么是模型训练
• 求解器 & 损失函数简介
• 小实验2：手写一个简单的神经网络并训练它
• 什么是训练/预训练/微调/轻量化微调
• Transformer结构简介
• 轻量化微调
• 实验数据集的构建
• …

第四阶段（20天）：商业闭环

对全球大模型从性能、吞吐量、成本等方面有一定的认知，可以在云端和本地等多种环境下部署大模型，找到适合自己的项目/创业方向，做一名被 AI 武装的产品经理。

• 硬件选型
• 带你了解全球大模型
• 使用国产大模型服务
• 搭建 OpenAI 代理
• 热身：基于阿里云 PAI 部署 Stable Diffusion
• 在本地计算机运行大模型
• 大模型的私有化部署
• 基于 vLLM 部署大模型
• 案例：如何优雅地在阿里云私有部署开源大模型
• 部署一套开源 LLM 项目
• 内容安全
• 互联网信息服务算法备案
• …

学习是一个过程，只要学习就会有挑战。天道酬勤，你越努力，就会成为越优秀的自己。

如果你能在15天内完成所有的任务，那你堪称天才。然而，如果你能完成 60-70% 的内容，你就已经开始具备成为一名大模型 AI 的正确特征了。

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】