大模型分布式训练的第四种境界

本文约8500字，建议阅读10分钟
本文将分享大模型分布式训练面临的挑战和相关技术体系。

主要内容包括：

1. 历史背景

2. 分布式训练挑战

3. 分布式训练技术体系

4. 未来挑战

5. Q&A

01 历史背景

自 2019 年以来，大语言模型发展迅猛，不断有新的研究成果涌现，包括各类预训练模型及其应用，如 LLM Infra 等相关技术工作，这些成果令人振奋。然而，对于从事相关基础设施建设的人员来说，这样的发展趋势也带来了诸多挑战。

人们将当前大语言模型领域的机遇与淘金热时期相类比，众所周知其中蕴藏着巨大的价值，但如何有效发掘则是一个问题。在过去淘金需要铲子，而现在在大语言模型时代，则需要更高级别的工具——挖掘机。这一比喻实际上揭示了对基础设施建设提出的更高要求。这意味着不仅要有掌握操作挖掘机技术的专业人才，即具备处理和利用大语言模型技能的专业人员，同时还需要配套完善的基础设施，例如，确保挖掘机能够从一个地方转移到另一个地方的运输链路，以及足够的燃料储备。将当前大模型所需的计算能力建设比作淘金热时期所需工具的升级是非常生动贴切的。

02 分布式训练所面临的挑战

1. 大语言模型对算力的巨大需求

第一个关键挑战在于大语言模型对算力的巨大需求。尽管业界普遍认为大模型需要庞大的算力支持，但具体需求量级往往没有被清晰理解。这里提供一个简单的估算公式：模型所需算力大致等于模型参数数量乘以训练时使用的 token 数量的六倍。

例如，GPT-3 拥有约 1750 亿个参数，在通用预训练阶段使用了 3000 亿个 token，通过计算可知其所需的算力达到了 10 的 23 次方的量级。同样地，我们可以用这个公式估算其他大型模型如 LLM-65B 以及谷歌的相关模型所需要的算力，这些模型所需的算力甚至能达到 10 的 24 次方级别，这是一个极为巨大的数字。达到这样的算力规模后，我们可能会观察到大模型的涌现能力，虽然这一概念仍存在一定的争议。

当模型算力提升至这种水平时，通常会显著提高模型的各项指标性能。对于如此庞大的算力要求，可以直观地联想到其高昂的成本。参考北京资源大会上获取的信息，以LLM-65B 模型为例，单次有效训练或者说找到一个较为优秀的模型可能需要花费四千万到五千万人民币以上。这足以说明大语言模型对算力的需求之大，即使当前算力成本相对降低，但仍处于高位。

2. 大模型规模与显存之间的冲突

第二个挑战是大模型规模与显存之间的冲突问题。目前硬件设备尚未完全突破工业制造上的限制，导致在现有统计数据显示下，即便算力不断提升，但受限于显存容量，我们仍然难以充分利用这些算力资源进行高效的大模型训练。在引用的相关数据中，大约每 3.9 个月模型的大小就会翻一番，这对应着显存消耗的巨大需求。

当前，一个规模达到 5400 亿参数的大模型可能需要接近 2T的显存，其成本极其高昂。同时，尽管单设备算力经常有所提升，但整体算力需求增长速度更快，在过去 18 个月内，算力需求甚至增加了约十倍，从 10 的 23 次方到 10 的 24 次方量级。

然而，GPU 和 TPU 等计算设备的算力发展远低于这一需求增速，随着时间推移，供需差距将进一步扩大。对于技术人员而言，单卡算力难以满足如此大规模的训练需求。以目前最大的算力需求为参照，假设使用单卡 A100 峰值算力 312tFLOPs，若要充分利用所有算力，理论上需要约 254 年才能完成训练。

3. 构建分布式系统

构建分布式系统可以显著缩短这一时间，例如拥有 2048 个 GPU 时，根据计算效率（r）的不同，训练所需时间可降至若干天。

上图右侧表格显示了模型参数规模与单设备硬件利用率之间的关系。通常，模型参数规模越大，硬件利用率越高，因为需要进行更多计算，从而促使利用率得到提升。但在 3000 亿参数规模下，整体利用率仅约为 53%，像 Meta Lab 的 LLM-65B 模型估计利用率可能在 42% 左右。

面对这些挑战，我们需要关注模型在整个集群中的定位、计算流程及其复杂性。深度学习过程主要包括前向传播、反向传播和优化器计算，在大规模集群环境下，网络拓扑结构、硬件间的通信带宽以及不同层级的访问速度都成为设计的关键考量因素。此外，如何优化 CPU 与 GPU 协同工作，提高数据读取预处理和 off-load 任务的效率，以及解决多层交换机下的节点间通讯复杂性，都是分布式大规模训练的重大技术难题。

03 分布式训练技术体系

接下来介绍分布式训练的技术体系，这里分享的内容主要基于公开文献和团队过往的学习积累。

1. 分布式机器学习的基本概念

分布式机器学习的发展历史并不长，只有十来年的时间，但早期的一些关键点如AlexNet 对现今研究仍具有重要启示意义。

实际上，从计算机视觉（CV）的角度看，深度学习首次被引入并取得突破是在 ImageNet 图像识别挑战赛中，显著提升了图像识别性能。在这个领域和方向上，更值得关注的是早期使用两张 GPU 显卡来处理计算密集型任务，这在当时是开创性的尝试。同时，在工业界，百度凤巢系统先前就已开始采用大规模深度神经网络（DNN）进行相关工作。

随着深度学习框架如 TensorFlow 的流行以及分布式训练框架的发展，解决大模型问题的方法也在不断演变。初期可能主要聚焦于稀疏场景下的分布式训练优化，之后百度美颜项目将高性能计算（HPC）中的 All-Reduce 技术引入分布式机器学习，用于加速数据并行处理。再后来，诸如 Word2Vec、DDP 等经典操作的出现推动了这一领域的进步。

到了 2018 年，谷歌提出了将计算图与数据处理平台（如 Spark）中的数据流图概念相结合，通过合理分配资源到各个设备上实现大规模模型训练任务的拆分。Transformer 等高效且易于并行化的模型结构的提出，进一步推动了监督学习算法在自然语言处理（NLP）等领域内的显著提升，同时分布式训练技术也同步快速发展。

在此过程中，英伟达等公司提出了 MegaPipe 等相关工作以实现大规模流水线式训练，并关注张量并行等方面的研究，这些技术近年来愈发成熟并逐步从科研走向实际应用。

2. 数据并行

接下来的重点为数据并行层面，比如微软推出的 DeepSpeed 框架，其核心技术旨在解决大规模分布式训练时的显存占用问题。研究人员发现，显存瓶颈主要包括两大部分：一是参数服务器存储的模型参数、梯度及优化器状态；二是零冗余优化（Zero Redundancy Optimization, ZeRO）相关的显存占用，包括激活函数计算过程中的临时显存分配。

由于 GPU 或 CPU 在进行数据处理和显存分片时必须按照一定的规则快速分配，这就导致了大量临时存储碎片的产生。DeepSpeed 通过科学分析，以及针对性地优化这两部分显存占用，有效提升了计算效率。

在 Zero-DP（数据并行零冗余优化）方面，它将显存优化划分为三个阶段。首先，在 Stage 1 中，研究发现优化器状态无需每个计算节点都存储完整的备份，而是仅需保存属于自身部分的 n 分之一。例如，在拥有 8 张 GPU 卡的情况下，每个设备仅需存储其对应的 0 到 1/8 的数据部分，以此类推。通过这种策略，在前向和反向传播计算过程中，只需在通信同步时进行特定操作，即使只保存 n 分之一的备份，也能确保逻辑上的数据并行正确性。

其次，在 Stage 2 中，梯度同样可以仅保留 n 分之一，计算过程中实时与其他节点交换不属于自身保存范围内的梯度部分，采用类似 All-Reduce 或 Ring-Reduce 的方式来完成梯度聚合。

进一步深入，在 Stage 3 中，模型参数也可以沿用同样的思路，即在前向传播时临时获取所需的部分，以通信时间换取显存空间。

DeepSpeed 框架中的 Zero-DP 实现展现了这一系列思想，尽管实现复杂度较高，但已成为深度学习社区中的一项重要技术突破，并广泛应用于大规模模型训练，如语言模型等场景。其中，Recompute 策略是解决激活函数显存占用问题的关键手段，即在反向传播时重新计算必要的激活值，而非一直保存它们，从而提高显存利用率。此外，随着对 optimizer、梯度和参数切分方法的深入探讨，业界开始考虑是否能将这些数据从 GPU 显存迁移到内存甚至硬盘中存储，这是一种更为激进的方法，旨在进一步释放 GPU 显存资源，优化大规模模型训练过程。然而，这种方法会带来额外的时间成本，目前行业普遍认为前向计算与反向计算的时间比例约为 1:2，因此完全重复前向计算可能会导致大约 20% 至 30% 的时间开销增加。

在资源有限的情况下，可以通过 off-load 技术将计算图的部分放到内存或硬盘中，以减少损耗并支持更大规模的模型训练。这一方法虽然会增加训练时间（例如，从百毫秒到数秒，增幅至少十倍以上），但对于拥有较充裕时间但硬件资源有限的研究者来说，可以尝试采用这种技术来训练模型。

3. 流水线并行

另一个关键领域是流水线并行（pipeline parallelism），这一概念近年来在分布式训练的研究中得到了广泛应用。实际上，在 CPU 操作系统的任务调度层面，流水线并行是一个经典的工作原理，而这里主要聚焦于其在分布式训练中的应用。

流水线并行可以大致划分为两个部分：同步流水线和异步流水线。同步流水线的核心思想是保持与单机训练行为一致，无论最终并行度如何设置，即每个计算阶段的行为都如同在一单一设备上进行训练时那样，包括前向传播、反向传播及通信等操作。

这是一个基本的同步流水线概念，通过将计算图划分成多个 stage，比如这张图展示了四个 stage，并将每个 stage 分配给不同的设备执行，例如 device3 负责某个 stage。

按照从输入到输出的顺序，首先在 device0 上完成计算，然后将产生的激活值（activation）传递给下一个计算节点，以此类推，每个节点在收到前一个节点的激活值后进行相应的计算，并生成新的激活值继续向下传递。在反向传播过程中同样遵循这一逻辑，梯度信息会沿着相反的方向流动。

然而，当我们将这样的流程以图形方式展示给上级或老板时，他们可能会对资源利用效率表示不满，因为在这种同步流水线模式下，所有设备只有在接收到前一阶段的数据之后才能开始工作，而在等待数据传输的过程中，部分设备可能处于空闲状态，无法实现充分利用。

通过将计算图划分为多个 stage，并分配到不同的设备上执行，每完成一个 stage 后就将激活值传递给下一个 stage 进行计算。当设备等待接收上一阶段计算结果（activation 或 gradients）时，会出现空转，这意味着即使拥有大量资源的集群，其实际并行工作比例可能只有四分之一或八分之一，降低了整体利用率，这显然是不可接受的。

为了解决这一问题，谷歌早期的一项研究提出了一种名为“pipeline”的优化思路。该思路通过将每次计算任务切分成更小的部分，并确保计算完成后能迅速传递给下一个设备进行处理，同时当前设备又能立即开始处理来自上一个设备的新任务。

具体而言，这种方法是将 mini-batch 样本进一步细分为 micro-batch 维度，例如原本可能是 64 个 patch，现在可以按 16 个 patch 细分，即每个样本的 micro-batch 数量变为 16。这样操作后，在某些维度上，设备的利用率会有所提高，尽管仍然存在空洞时间（idle time），但在很多情况下，这种时间会被显著缩短。

在讨论分布式训练中的显存管理问题时，阿里巴巴经过一系列科学研究提出了一项名为"1F1B"的工作。尽管随着模型并行度（micro-batch 数）的增加，中间临时保存的副本数量也会增多，导致 GPU 显存峰值持续增长，但这项工作指出通过科学可控的调度策略可以改善这一情况。具体来说，采用“来一个前向计算一个后向”的方式，使得整体显存占用量与设备数量和单个 micro-batch 所占用的激活值大小相关，而非 micro-batch 的数量。这样一来，即使增大 micro-batch 大小，也不会增加显存峰值。

此外，英伟达提出了改进型的工作，在"1F1B"的基础上提出了精细化切分 stage 的方法。不同于先前按连续几层划分 stage 的做法，新的思路是采用 Round Robin 的方式更为细致地分配任务，例如将 8 层网络分为两部分，第一部分从 L1 到 L5，并且在完成一层计算后立即推进至下一设备进行后续计算，从而提高了设备利用率，这种技术被称作交错式(1F1B)流水线方法。

国防科技大学的一项研究则结合了 recompute 技术，发现在流水线并行中的最后一个 stage 无需保存激活值，可以直接进行反向传播计算，通过提前调度相关节点可缩短关键路径，有效提高训练效率。尽管这项工作可能具有一定的局限性，但其带来的性能提升显著，值得深入研究。

在同步训练方面，还有更为激进的方法，例如某位华裔学生的研究，尝试解决一个与传统思路不同的问题，即在考虑时序图时，并不局限于从零到一或从一到零的单向顺序执行。例如，在处理 micro-batch 时，同时利用两个方向进行数据传输和计算，就像在一个双向通道上操作一样。具体来说，设想在设备间构建一个双向通信链路，以 micro-batch 为例子，可以实现数据既从 device 这一端流入流出，也能从另一个方向传递。这样，在对时序进行合理安排、优化 stage 编排时，比如将原本的阶段细分为四个部分，就能够进一步提升整个系统的利用率，包括压缩 pipeline 空间占用的时间，从而实现更高效地运行。

关于流水线并行中的异步变形部分，这里不作详细展开。早期在处理 P 问题时，我们同样面临抉择：是在所有节点的梯度都准备好之后进行一轮整体参数更新（同步），还是只要有部分节点完成就立刻更新参数（异步）。此外，还有诸如 ASP、BSP 和 SSP 等不同版本的同步/异步策略。

在流水线变形中，有一种工作模式是当数据到达后立即执行反向计算，这就要求保存多个副本以应对版本差异问题，确保在异步计算场景下，即使显存使用增多，也能保证充分利用设备算力，并尽量减少显存开销。理论上，这种逻辑能够充分利用所有资源，但在实际应用中，由于中间计算过程存在的参数版本差异，可能导致与单卡训练效果存在一定的差距。

Tensor 运算优化的一项经典工作——m in，强调在处理大规模矩阵计算时（例如数千万个值的矩阵），为了充分挖掘计算能力，可以将矩阵切分并在不同的设备上并行计算。虽然文章没有明确提到这一点，但实际在 m 的实现中，会对 embedding 部分也进行相应的切分处理，以便于后续的 MLP（多层感知机）和 Attention 机制等大矩阵乘法操作。通过合理地在不同设备上分配 tensor 并适时引入通信逻辑，在恰当的时间点进行通讯与结果复用，从而实现对 tensor 变形问题的有效解决。

实际上，在实际工作中，很多人并未充分关注到如何更合理地在不同设备上分配和切分矩阵以优化计算的问题。这里分享一项研究工作，其目标是通过一个创新的通信策略来解决这一问题。该策略主张从二维维度对矩阵进行切分，并基于此对样本进行区分处理。

具体来说，研究者提出了一种类似于 GR=2 的方法，将矩阵在两维空间内划分，并且对权重采用二乘二的方式进行切分，每次仅传输四分之一的数据，在所有切片都完成计算后进行一次整合通信。尽管这个基本概念看似直观且合理，但在实践中，需要在较早阶段就实现同步通信以确保各个设备能拿到相应结果。

以一个假设为例，若在一个设备上将任务划分为八层，那么是否可以在不进行额外通信的情况下，在这八层内部连续进行计算呢？研究指出这是可能的。设想这样一个场景：当完成某一层计算后，当前设备只持有数据的一部分（即一半），然后继续处理下一部分数据，直到最后与其它显卡进行通信时才增加通信逻辑。这样一来，原本输入数据就形成一个斜角，只需要通过对权重进行转置操作，理论上便能在单个设备上连续进行计算，直至必须与其它设备通讯为止。这种技术训练方法旨在通过精细化的任务切分与通信策略优化，提高分布式计算环境下的资源利用率和效率。

04 未来挑战

最后分享我们在实际工作中关注到的几个关键点。

首先，面对当前工业界中复杂且难以通过单一程序、语言或图像清晰表达的问题，我们意识到需要开发一种高度直观且功能丰富的工具。类似于腾讯报道中所展示的那种可视化工具，在深度学习框架推出后起到了重要作用。在现阶段，我们也急需开发出训练过程可视化以及优化调度等方面的辅助工具，这将极大助力于我们的科研团队和同行们的工作进展。

第二点是关于自动变形这一领域的研究，尤其是智能并行相关的技术。值得注意的是，现有的学术研究成果往往更多关注于图（graph）结构本身或设备层面，而忽视了如何处理大量动态信息等实际问题。在前文中提到为了确保算力的有效利用，我们需要考虑如何动态地调整喂给模型的样本数量以及 max-size 等参数，并且要充分应对计算设备内部的异构性问题，包括但不限于计算设备本身的异构性和网卡的异构性。这些因素都可能带来诸多挑战。

此外，在模型结构演进方面，从 Transformer 到当前一些新型模型结构，我们观察到存在变革的需求，例如出现了 WW-KV 这类新型结构。过去的分布式框架主要针对 Transformer 结构进行优化，但在面对其它类型的模型或硬件切换时，原有的策略可能无法有效解决所有问题，甚至于模型结构本身也可能需要适应性改变。

对于 Transformer 结构优化是否转向其它结构方向的问题，确实已有相关研究，但我们所做的工作还远远不够。目前了解到的一些研究不仅局限于优化 Transformer 结构本身，还在探索结合其它结构的可能性，比如参考 CN 等早期和最新的研究成果。

05 Q&A

Q1： 针对于 transformer 结构的优化，有没有想往其它结构的方向优化？

A1：有一个双向工作的概念，这是一个非常有趣的工作思路，尤其在处理大规模数据和不同设备场景时，其关注点和选型考量都有所区别。例如，当模型规模达到千亿级别时，可能需要采用更细致、更具针对性的方法来优化资源使用和计算流程。由于政策原因，GPU 相关的问题在此不便详细回答，请大家谅解。

Q2： 在较小规模的场景下，如何高效利用显存和计算时间？

A2：可能会通过一些微调策略来平衡二者，但这通常会导致一定程度上的显存额外损耗。不存在既能最大化利用资源又不违背基本物理限制的方法。例如，我们可以运用像recompute 这样的经典优化技术，在某些层面上进行更细致的优化，虽然这会增加一定的计算耗时（如 10% 以内），但其适用范围有限，往往仅适用于特定的产品或场景。

Q3： 关于 Machin 与 DeepSpeed 这两套框架哪个使用更为广泛的问题。

A3：Machin 主要专注于 3D 并行，尚未实现 z20 和 z23 那样的特性；而 DeepSpeed则主要采用 zero+off-load 的技术。从我的角度来看，这两者的定位并不特别清晰。举例来说，DeepSpeed 自身并未强调独特的并行机制，可能借鉴或复用了 Machin 等其他技术的一些工作成果。至于哪个框架使用得更多，则需要根据具体应用场景来判断，我个人认为它们各有其适用范围和优缺点。

在大模型训练过程中，我们可以将其大致分为预训练（pre-tuning）和微调（alignment）两个阶段。对于预训练阶段，当模型参数规模庞大，计算设备上载入和训练复杂度较高时，往往会选择使用像 Machin-Deps Pet 等工具进行优化，同时科大讯飞以及其他类似平台如 GPT-3 也采用类似的解决方案来应对这一挑战。

而在微调或具体应用任务中，比如运用 70 亿参数级别的模型进行指令生成任务时，DeepSpeed 可能更为常用。原因在于 DeepSpeed 目前集成了多种高效技术，例如 Hugging Face 的集成模块 night，以及国内一些类似于 Hugging Face 的平台功能，并且其对模型和平台的集成程度相当高。

然而，从代码质量和架构角度来看，Machin-Deps Pet 与 DeepSpeed 之间存在相互交叉和借鉴的现象，可以说是你中有我、我中有你，这并非贬低它们的意思，而是反映了这两种工具在发展过程中由于需求和技术迭代所形成的必然联系。

Q4： 自动并行策略的搜索如何进行？

A4：这是一个颇具挑战性的问题。在自动并行领域中，搜索有效的并行方案虽是难题之一，但相较于模型结构搜索（如 NAS）而言，其难度尚不至那么突出。我认为关键在于如何合理地表示和处理 IR（中间表示），比如 DPT-PPP 这样的并行模式，以及如何描述硬件设备状态与动态信息，例如根据计算资源利用率、同步样本数量等动态调整参数。

随着计算任务需求的增长，尤其是当面临越来越多的计算资源时，资源复用和共享成为必然趋势，这会带来一系列新的问题。对于高性能 AI 容器的相关思路，我目前并不清楚提问的具体含义。

关于 Zero Stage 3 和 Tensor Parallelism 的区别：Zero Stage 3 优化涉及到将optimizer 的状态分片，但它并非作用于操作符级别。举例来说，在一个包含两层网络结构的例子中，第一层为 w0 乘以 x0，第二层为 w1 乘以 y0。在 Stage 3 中，时间步长的切分不会深入到权重 w0 内部。而 Tensor Parallelism 则是针对大规模权重矩阵（如千万级或亿级别的 w）进行切分，并分配到不同的硬件上进行并行计算。

Q5：Softmax 运算与其他操作之间是否存在并行的可能性？

A5：在 Tensor 变形时，M 模型的实现涉及到 Softmax 操作。在进行分布式计算时，Softmax 层会被划分为多个设备处理，这会引入两次通信过程。首先，需要通过 All-Reduce Max 通信来确保每个设备知道全局的最大指数值，以便后续对所有设备上的指数进行平均化处理，这样可以防止因大规模密集计算导致溢出问题。例如，如果有 1 亿个数值都在 1998 或 999 附近，直接相加 1000 次确实容易溢出，但通过减去最大值并求和后再平均，最终结果将是等价的。

第一个通信步骤就是传递这个最大指数值，第二个通信则是整合各个设备上完成的部分计算结果，以实现最终输出的一致性。这里需要注意的是，我们不是按照 Stage 维度或者 Line 维度切分，而是从参数维度出发，将模型参数分成 n 份，每个设备负责其中的 1/n 部分，并在该设备上进行完整的前向和反向计算。这种情况下，即使是在流水线并行中，每个设备仅处理整个模型的 n 分之一，但仍然能保证数据并行，即模型分布在不同设备上，每个设备独立处理不同的数据样本，而在最后阶段通过通信整合结果。

这样的设计具有很高的工程挑战性，尤其是在处理只有 R、只有 1 和只有 2 这类细粒度差异时。对于对此感兴趣的人，建议查看 Deep Speed 的相关实现代码，它提供了一个巧妙的数据并行解决方案。

分享嘉宾

段石石，上海壁仞智能科技有限公司技术总监。过去从事过推荐算法、搜广推分布式训练系统开发、机器学习平台，目前主要负责国产芯片大规模分布式的相关工作，支持包括经典场景以及大模型场景下的生态适配。

编辑：王菁