大模型训练技术架构、并行策略与优化方案详解

如图所示，PS 架构将计算节点分为 server 与 worker，其中，worker 用于执行网络模型的前向与反向计算。而 server 则对各个 worker 发回的梯度进行合并并更新模型参数，对深度学习模型参数中心化管理的方式，非常易于存储超大规模模型参数。

但是随着模型网络越来越复杂，对算力要求越来越高，在数据量不变的情况下，单个 GPU 的计算时间是有差异的，并且网络带宽之间并不平衡，会存在部分 GPU 计算得比较快，部分 GPU 计算得比较慢。这个时候如果使用异步更新网络模型的参数，会导致优化器相关的参数更新出现错乱。而使用同步更新则会出现阻塞等待网络参数同步的问题。

GPU 强大的算力毋庸置疑可以提升集群的计算性能，但随之而来的是，不仅模型规模会受到机器显存和内存的制约，而且通信带宽也会由于集群网卡数量降低而成为瓶颈。

这个时候百度基于 PS 架构之上提出了 Ring-All-Reduce 新的通讯架构方式。

通过异步流水线执行机制，隐蔽了 IO 带来的额外性能开销，在保证训练速度的同时，使训练的模型大小不再受制于显存和内存，极大提升模型的规模。而 RPC&NCCL 混合通信策略可以将部分稀疏参数采用 RPC 协议跨节点通信，其余参数采用卡间 NCCL 方式完成通信，充分利用带宽资源。

集合通讯模式

集合通讯模式（Collective Communication）摒弃了中心化的参数服务器，转而采用点对点或全互联的通信原语（如 All-Reduce, All-Gather, Reduce-Scatter）。这种模式在现代高性能计算集群中更为常见，特别是在使用 NVLink 和 InfiniBand 的高速网络环境下。

在集合通讯模式下，所有 GPU 节点地位平等。以 Ring-All-Reduce 为例，它将梯度分片，在环状拓扑中进行传递和累加。相比传统的树形结构或广播，Ring-All-Reduce 在网络带宽利用率上更高，能够显著降低通信延迟。NVIDIA 的 NCCL (NVIDIA Collective Communications Library) 库实现了这些高效的原语，支持多机多卡环境下的自动拓扑感知和最优路径选择。此外，还有像 Hierarchy-All-Reduce 这样的变体，结合了 Ring 和 Tree 的优势，以适应更大规模的集群。

04 大模型训练相关论文

2022 年学习大模型、分布式深度学习，不可能错过的 AI 论文，你都读过了吗？根据上面的介绍，我们将分为分布式并行策略相关的论文、分布式框架相关的论文、通讯带宽优化相关的论文等不同的维度对论文进行整理。并给出一个简单的解读，希望大家可以一起去分享好的思想。

分布式并行策略相关

• 数据并行（Data Parallel, DP）：数据并行训练加速比最高，但要求每个设备上都备份一份模型，显存占用比较高。
• 模型并行（Model Parallel, MP)：模型并行，通信占比高，适合在机器内做模型并行且支持的模型类型有限。
• 流水线并行（Pipeline Parallel, PP)：流水线并行，训练设备容易出现空闲状态，加速效率没有数据并行高；但能减少通信边界支持更多的层数，适合在机器间使用。
• 混合并行（Hybrid parallel, HP）：混合并行策略的思想，集三种策略的优势于一身，实现取长补短。具体来说，先在单机内使用模型并行和分组参数切片组合的策略，这么选择的原因是这两个策略通信量较大，适合使用机器内的卡间通信。接着，为了承载千亿规模大模型，叠加流水线并行策略，使用多台机器共同分担计算。最后，为了计算和通讯高效，在外层又叠加了数据并行来增加并发数量，提升整体训练速度。这就是我们目前在 AI 框架中添加的并行策略，业界基本上都是使用这种方式。

并行相关的经典论文

下面就是并行相关的经典推荐论文，首先就是 Jeff Dean 在 2012 年的开创文章，然后介绍 Facebook PyTorch 里面使用到的数据并行中 DDP、FSDP 的策略。然而这并不够，因为有多重并行策略，于是 NVIDIA 推出了基于 GPU 的数据、模型、流水线并行的比较综述文章。实际上流水线并行会引入大量的服务器空载 buffer，于是 Google 和微软分别针对流水线并行优化推出了 GPipe 和 PipeDream。最后便是 NVIDIA 针对自家的大模型 Megatron，推出的模型并行涉及到的相关策略。

1. Large Scale Distributed Deep Networks 2012 年的神作，要知道那个时候神经网络都不多，这是出自于 Google 大神 Jeff Dean 的文章。主要是神经网络进行模型划分，因为推出得比较早，所以会稍微 Native 一点，但是作为分布式并行的开创之作，稍微推荐一下。

2. Getting Started with Distributed Data Parallel PyTorch Distributed: Experiences on Accelerating Data Parallel Training. Facebook 为 PyTorch 打造的分布式数据并行策略算法 Distributed Data Parallel (DDP)。与 Data Parallel 的单进程控制多 GPU 不同，在 distributed 的帮助下，只需要编写一份代码，torch 就会自动将其分配给 n 个进程，分别在 n 个 GPU 上运行。不再有主 GPU，每个 GPU 执行相同的任务。对每个 GPU 的训练都是在自己的过程中进行的。每个进程都从磁盘加载其自己的数据。分布式数据采样器可确保加载的数据在各个进程之间不重叠。损失函数的前向传播和计算在每个 GPU 上独立执行。因此，不需要收集网络输出。在反向传播期间，梯度下降在所有 GPU 上均被执行，从而确保每个 GPU 在反向传播结束时最终得到平均梯度的相同副本。

3. Fully Sharded Data Parallel: faster AI training with fewer GPUs Facebook 发布的 FSDP（Fully Sharded Data Parallel），对标微软在 DeepSpeed 中提出的 ZeRO，FSDP 可以看成 PyTorch 中的 DDP 优化版本，本身也是数据并行，但是和 DDP 不同的是，FSDP 采用了 parameter sharding，所谓的 parameter sharding 就是将模型参数也切分到各个 GPUs 上，而 DDP 每个 GPU 都要保存一份 parameter，FSDP 可以实现更好的训练效率（速度和显存使用）。

4. Efficient Large-Scale Language Model Training on GPU Clusters 很好的一篇综述出品与 NVIDIA，论文中，NVIDIA 介绍了分布式训练超大规模模型的三种必须的并行技术：数据并行（Data Parallelism）、模型并行（Tensor Model Parallelism）和流水并行（Pipeline Model Parallelism）。

5. Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training 在传统的数据并行中，模型参数被复制并在每次训练循环结束后被优化器更新。然而，当每个核的批量数不够大的时候，计算或许会变成一个瓶颈。例如，以 MLPerf 的 BERT 训练为例，在 512 个第三代 TPU 芯片上，LAMB 优化器的参数更新时间可以占到整个循环时间的 18%。Xu 等人在 2020 年提出了参数更新划分技术，这种分布式计算技术首先执行一个 reduce-scatter 操作，然后使得每个加速器有整合梯度的一部分。这样每个加速器就可以算出相应的被更新的局部参数。在下一步，每个被更新的局部参数被全局广播到各个加速器，这样使得每个加速器上都有被更新的全局参数。为了获得更高的加速比，同时用数据并行和模型并行去处理参数更新划分。在图像分割模型中，参数是被复制的，这种情况下参数更新划分类似于数据并行。然后，当参数被分布后到不同的核之后，就执行多个并发的参数更新划分。

6. PipeDream: Fast and Efficient Pipeline Parallel DNN Training 微软研究院宣布了 Fiddle 项目的创立，其包括了一系列的旨在简化分布式深度学习的研究项目。PipeDreams 是 Fiddle 发布的第一个侧重于深度学习模型并行训练的项目之一。其主要采用'流水线并行'的技术来扩展深度学习模型的训练。在 PipeDream 中主要克服流水线并行化训练的挑战，算法流程主要如下。首先，PipeDream 必须在不同的输入数据间，协调双向流水线的工作。然后，PipeDream 必须管理后向通道里的权重版本，从而在数值上能够正确计算梯度，并且在后向通道里使用的权重版本必须和前向通道里使用的相同。最后，PipeDream 需要流水线里的所有 stage 都花费大致相同的计算时间，这是为了使流水线得到最大的通量。

7. GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism GPipe 是 Google 发明的论文，专注于通过流水线并行扩展深度学习应用程序的训练负载。GPipe 把一个 L 层的网络，切分成 K 个 composite layers。每个 composite layer 运行在单独的 TPU core 上。这 K 个 core composite layers 只能顺序执行，但是 GPipe 引入了流水并行策略来缓解这个顺序执行的性能问题，把 mini-batch 细分为多个更小的 macro-batch，提高并行程度。GPipe 还用 recomputation 这个简单有效的技巧来降低内存，进一步允许训练更大的模型。

8. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism 出自 NVIDIA，虽然这两篇文章都是在讲 Megatron 网络模型，实际上里面展开的都是模型并行等多维度并行的相关的技术点。其中第一篇论文共有两个主要的结论：1，利用数据和模型并行的分布式技术训练了具有 3.9B 参数的 BERT-large 模型，在 GLUE 的很多数据集上都取得了 SOTA 成绩。同时，还训练了具有 8.3B 参数的 GPT-2 语言模型，并在数据集 Wikitext103，LAMBADA，RACE 都上取得 SOTA 成绩。这篇论文，一方面体现了算力的重要性，另一方面体现了模型并行和数据并行技术关键性。这两项优化技术在加速模型训练和推断过程中至关重要。

05 大模型算法相关

必须了解的基础大模型结构

基础大模型结构基本上都是由 Google 贡献的，首先要看 17 年只需要 Attention 替代 RNN 序列结构，于是出现了第四种深度学习的架构 Transformer。有了 Transformer 的基础架构后，在 18 年推出了 BERT 预训练模型，之后的所有大模型都是基于 Transformer 结构和 BERT 的预训练机制。后面比较有意思的就是使用 Transformer 机制的视觉大模型 ViT 和引入专家决策机制的 MoE。

1. Attention is all you need Google 首创的 Transformer 大模型，是现在所有大模型最基础的架构，现在 Transformer 已经成为除了 MLP、CNN、RNN 以外第四种最重要的深度学习算法架构。谷歌在 arxiv 发了一篇论文名字教 Attention Is All You Need，提出了一个只基于 attention 的结构来处理序列模型相关的问题，比如机器翻译。传统的神经机器翻译大都是利用 RNN 或者 CNN 来作为 encoder-decoder 的模型基础，而谷歌最新的只基于 Attention 的 Transformer 模型摒弃了固有的定式，并没有用任何 CNN 或者 RNN 的结构。该模型可以高度并行地工作，所以在提升翻译性能的同时训练速度也特别快。

2. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding Google 发布的首个预训练大模型 BERT，从而引爆了预训练大模型的潮流和趋势，这个不用介绍大家肯定有所听闻啦。BERT 的全称为 Bidirectional Encoder Representation from Transformers，是一个预训练的语言表征模型。它强调了不再像以往一样采用传统的单向语言模型或者把两个单向语言模型进行浅层拼接的方法进行预训练，而是采用新的 masked language model（MLM），以致能生成深度的双向语言表征。BERT 论文发表时提及在 11 个 NLP（Natural Language Processing，自然语言处理）任务中获得了新的 state-of-the-art 的结果，令人目瞪口呆。

3. An Image is Worth 16x16 Words: transformer for Image Recognition at Scale ViT Google 提出的首个使用 Transformer 的视觉大模型，基本上大模型的创新算法都是出自于 Google，不得不服。ViT 作为视觉转换器的使用，而不是 CNN 或混合方法来执行图像任务。结果是有希望的但并不完整，因为因为除了分类之外的基于视觉的任务：如检测和分割，还没有表现出来。此外，与 Vaswani 等人（2017 年）不同，与 CNN 相比，transformer 性能的提升受到的限制要大得多。作者假设进一步的预训练可以提高性能，因为与其他现有技术模型相比，ViT 具有相对可扩展性。

4. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding 好像 G 开头的模型都是 Google 的了一样魔性。在 ICLR 2021 上，Google 的进一步将 MoE 应用到了基于 Transformer 的神经机器翻译的任务上。GShard 将 Transformer 中的 Feedforward Network（FFN）层替换成了 MoE 层，并且将 MoE 层和数据并行巧妙地结合起来。在数据并行训练时，模型在训练集群中已经被复制了若干份。GShard 通过将每路数据并行的 FFN 看成 MoE 中的一个专家来实现 MoE 层，这样的设计通过在多路数据并行中引入 All-to-All 通信来实现 MoE 的功能。

具有里程碑意义性的大模型

1. GPT-3: Language Models are Few-Shot Learners OpenAI 发布的首个百亿规模的大模型，应该非常具有开创性意义，现在的大模型都是对标 GPT-3。GPT-3 依旧延续自己的单向语言模型训练方式，只不过这次把模型尺寸增大到了 1750 亿，并且使用 45TB 数据进行训练。同时，GPT-3 主要聚焦于更通用的 NLP 模型，解决当前 BERT 类模型的两个缺点：对领域内有标签数据的过分依赖：虽然有了预训练 + 精调的两段式框架，但还是少不了一定量的领域标注数据，否则很难取得不错的效果，而标注数据的成本又是很高的。对于领域数据分布的过拟合：在精调阶段，因为领域数据有限，模型只能拟合训练数据分布，如果数据较少的话就可能造成过拟合，致使模型的泛化能力下降，更加无法应用到其他领域。

2. T5: Text-To-Text Transfer Transformer Google 把 T5 简单的说就是将所有 NLP 任务都转化成 Text-to-Text（文本到文本）任务。对于 T5 这篇论文，很 Google 的一篇文章啦，让我也很无力，毕竟财大气粗之外，还有想法，这就是高富帅。回到论文本身，T5 意义不在烧了多少钱，也不在屠了多少榜，其中 idea 创新也不大，它最重要作用是给整个 NLP 预训练模型领域提供了一个通用框架，把所有任务都转化成一种形式。

3. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows 微软亚研提出的 Swin Transformer 的新型视觉 Transformer，它可以用作计算机视觉的通用 backbone。在两个领域之间的差异，例如视觉实体尺度的巨大差异以及与文字中的单词相比，图像中像素的高分辨率，带来了使 Transformer 从语言适应视觉方面的挑战。

超过万亿规模的稀疏大模型

1. Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts Google 发布的多任务 MoE。多任务学习的目的在于用一个模型来同时学习多个目标和任务，但常用的任务模型的预测质量通常对任务之间的关系很敏感（数据分布不同，ESMM 解决的也是这个问题），因此，google 提出多门混合专家算法（Multi-gate Mixture-of-Experts）旨在学习如何从数据中权衡任务目标（task-specific objectives）和任务之间（inter-task relationships）的关系。所有任务之间共享混合专家结构（MoE）的子模型来适应多任务学习，同时还拥有可训练的门控网路（Gating Network）以优化每一个任务。

2. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity Google 重磅推出首个万亿参数的超大规模稀疏语言模型 Switch Transformer。声称他们能够训练包含超过一万亿个参数的语言模型的技术。直接将参数量从 GPT-3 的 1750 亿拉高到 1.6 万亿，其速度是 Google 以前开发的语言模型 T5-XXL 的 4 倍。

06 内存和计算优化

最后就是优化方面的，其中主要是并行优化器、模型压缩量化、内存复用优化、混合精度训练等方面的优化，下面各列了几个最经典的文章。

优化器相关

1. Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour 一篇 17 年关于优化器的老文章，文章的一个重要的结论很简单，就是一个线性缩放原则，但里面分析的不错，讲到了深度学习中很多基本知识的一个理解。本文从实验的角度进行细致的分析。虽然文章分析的是如何在更大的 batch 上进行训练，但同样的道理本文也可以用在像我一样的贫民党，当我们没有足够的 GPU 或者显存不足的时候到底该怎么调节一些参数。

2. Adafactor: Adaptive learning rates with sublinear memory cost AdaFactor，一个由 Google 提出来的新型优化器，AdaFactor 具有自适应学习率的特性，但比 RMSProp 还要省显存，并且还针对性地解决了 Adam 的一些缺陷。说实话，AdaFactor 针对 Adam 所做的分析相当经典，值得我们认真琢磨体味，对有兴趣研究优化问题的读者来说，更是一个不可多得的分析案例。

3. ZeRO: Memory Optimization Towards Training A Trillion Parameter Models 微软提出很经典很经典的一个算法了，为了这个算法还基于 PyTorch 开发了一个分布式并行 DeepSpeed 框架。现有普遍的数据并行模式下的深度学习训练，每一台机器都需要消耗固定大小的全量内存，这部分内存和并不会随着数据的并行而减小，因而，数据并行模式下机器的内存通常会成为训练的瓶颈。这篇论文开发了一个 Zero Redundancy Optimizer (ZeRO)，主要用于解决数据并行状态下内存不足的问题，使得模型的内存可以平均分配到每个 gpu 上，每个 gpu 上的内存消耗与数据并行度成反比，而又基本不影响通信效率。 ZeRO 分为三个阶段：

   • Stage 1: 仅分片优化器状态。
   • Stage 2: 分片优化器状态和梯度。
   • Stage 3: 分片优化器状态、梯度和模型参数。这是最激进的优化，允许在有限的显存下训练更大的模型。

内存优化相关论文

1. Training Deep Nets with Sublinear Memory Cost 陈天奇这个名字可能圈内人都会比较熟悉了，在 2016 年的时候提出的，主要是对神经网络做内存复用。这篇文章提出了一种减少深度神经网络训练时内存消耗的系统性方法。主要关注于减少存储中间结果（特征映射）和梯度的内存成本，因为在许多常见深度架构中，与中间特征映射的大小相比，参数的大小相对较小。使用计算图分析来执行自动原地操作和内存共享优化。更重要的是，还提出了一种新的以计算交换内存的方法。

2. Gist: Efficient data encoding for deep neural network training Gist 是 ISCA'18 的一篇顶会文章，不算是新文章了，但是引用量在系统文章中算是非常高的，看完之后发现实验果然扎实，值得学习。主要思想还是围绕如何降低神经网络训练时候的显存使用量。Gist 面向数据压缩，发掘训练模式以及各个层数据的特征，对特定数据进行不同方案的压缩，从而达到节省空间的目的。

底层系统架构相关

1. Parameter Server for Distributed Machine Learning 亚马逊首席科学家李沐在读书时期发表的文章。工业界需要训练大型的机器学习模型，一些广泛使用的特定的模型在规模上的两个特点：1. 深度学习模型参数很大，超过单个机器的容纳能力有限；2. 训练数据巨大，需要分布式并行提速。这种需求下，当前类似 Map Reduce 的框架并不能很好适合。于是李沐大神在 OSDI 和 NIPS 上都发过文章，其中 OSDI 版本偏向于系统设计，而 NIPS 版本偏向于算法层面。关于深度学习分布式训练架构来说是一个奠基性的文章。

2. More Effective Distributed ML via a Stale Synchronous Parallel Parameter Server GeePS: Scalable deep learning on distributed GPUs with a GPU-specialized parameter server. 分布式深度学习可以采用 BSP 和 SSP 两种模式。1 为 SSP 通过允许 faster worker 使用 staled 参数，从而达到平衡计算和网络通信开销时间的效果。SSP 每次迭代收敛变慢，但是每次迭代时间更短，在 CPU 集群上，SSP 总体收敛速度比 BSP 更快，但是在 GPU 集群上训练，2 为 BSP 总体收敛速度比 SSP 反而快很多。

3. Bandwidth Optimal All-reduce Algorithms for Clusters of Workstations Bringing HPC Techniques to Deep Learning. 百度在 17 年的时候联合 NVIDIA，提出了 ring-all-reduce 通讯方式，现在已经成为了业界通讯标准方式或者是大模型通讯的方式。过去几年中，神经网络规模不断扩大，而训练可能需要大量的数据和计算资源。为了提供所需的计算能力，我们使用高性能计算（HPC）常用的技术将模型缩放到数十个 GPU，但在深度学习中却没有充分使用。这种 ring allreduce 技术减少了在不同 GPU 之间进行通信所花费的时间，从而使他们可以将更多的时间花费在进行有用的计算上。在百度的硅谷 AI 实验室（SVAIL）中，我们成功地使用了这些技术来训练最先进的语音识别模型。我们很高兴将 Ring Allreduce 的实现发布为 TensorFlow 的库和补丁程序，并希望通过发布这些库，我们可以使深度学习社区更有效地扩展其模型。

混合精度训练

混合精度训练（Mixed Precision Training）通过使用 FP16（半精度浮点数）代替 FP32（单精度浮点数）来加速计算并减少显存占用，同时保持模型精度。NVIDIA 的 Tensor Cores 专门针对混合精度进行了硬件优化。关键点在于使用动态 Loss Scaling 来防止梯度下溢，并通过保留关键参数（如优化器状态）的 FP32 副本以确保数值稳定性。这一技术已成为现代大模型训练的标配。

总结

大模型训练是一项复杂的系统工程，涉及硬件、算法、框架和优化等多个层面的协同。从早期的参数服务器架构到如今的高效集合通讯模式，从单一的数据并行到多维度的混合并行策略，技术的演进始终围绕着提升训练效率和降低资源消耗展开。Transformer 及其变体（如 MoE）奠定了算法基础，而 ZeRO、FSDP 等优化技术则突破了显存瓶颈。未来，随着硬件算力的持续提升和通信协议的优化，更大规模、更高效能的模型训练将成为常态。开发者需要深入理解这些底层原理，才能在实践中做出正确的技术选型和架构设计。