Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
TL;DR 精炼摘要
本文提出了一种新的交错流水线并行调度,结合张量、流水线和数据并行方法,有效提高了大规模语言模型在GPU集群上的训练效率。在3072个GPU上可达到502 petaFLOP/s的训练速度,吞吐量提高超过10%。
摘要
Large language models have led to state-of-the-art accuracies across a range of tasks. However, training these models efficiently is challenging for two reasons: a) GPU memory capacity is limited, making it impossible to fit large models on even a multi-GPU server, and b) the number of compute operations required to train these models can result in unrealistically long training times. Consequently, new methods of model parallelism such as tensor and pipeline parallelism have been proposed. Unfortunately, naive usage of these methods leads to fundamental scaling issues at thousands of GPUs, e.g., due to expensive cross-node communication or devices spending significant time waiting on other devices to make progress. In this paper, we show how different types of parallelism methods (tensor, pipeline, and data parallelism) can be composed to scale to thousands of GPUs and models with trillions of parameters. We survey techniques for pipeline parallelism and propose a novel interleaved pipeline parallelism schedule that can improve throughput by 10+% with memory footprint comparable to existing approaches. We quantitatively study the trade-offs between tensor, pipeline, and data parallelism, and provide intuition as to how to configure distributed training of a large model. Our approach allows us to perform training iterations on a model with 1 trillion parameters at 502 petaFLOP/s on 3072 GPUs with achieved per-GPU throughput of 52% of theoretical peak. Our code is open sourced at this https URL.
思维导图
论文精读
中文精读
1. 论文基本信息
1.1. 标题
高效地在 GPU 集群上使用 Megatron-LM 训练大规模语言模型 (Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM)
1.2. 作者
- Deepak Narayanan (微软研究院)
- Mohammad Shoeybi (英伟达)
- Jared Casper (英伟达)
- Patrick LeGresley (英伟达)
- Mostofa Patwary (英伟达)
- Vijay Korthikanti (英伟达)
- Dmitri Vainbrand (英伟达)
- Prethvi Kashinkunti (英伟达)
- Julie Bernauer (英伟达)
- Bryan Catanzaro (英伟达)
- Amar Phanishayee (微软研究院)
- Matei Zaharia (斯坦福大学)
1.3. 发表期刊/会议
预印本 (arXiv),但论文内容已在国际会议上广泛引用和讨论。
1.4. 发表年份
2021年
1.5. 摘要
大规模语言模型 (Large Language Models, LLMs) 在多种任务上取得了最先进的 (state-of-the-art) 准确率。然而,高效训练这些模型面临两大挑战:a) GPU 显存容量有限,即使是多 GPU 服务器也无法容纳大型模型;b) 训练这些模型所需的计算操作数量巨大,可能导致不切实际的训练时间。因此,学界提出了新的模型并行 (model parallelism) 方法,如张量并行 (tensor parallelism) 和流水线并行 (pipeline parallelism)。不幸的是,这些方法的简单应用在数千个 GPU 规模下会导致根本性的扩展问题,例如,由于昂贵的跨节点通信或设备花费大量时间等待其他设备推进。
本文展示了如何组合不同类型的并行方法(张量并行、流水线并行和数据并行 (data parallelism)),以扩展到数千个 GPU 和拥有万亿参数的模型。我们调研了流水线并行技术,并提出了一种新颖的交错流水线并行调度 (interleaved pipeline parallelism schedule),该调度可以在内存占用与现有方法相当的情况下,将吞吐量 (throughput) 提高 10% 以上。我们定量研究了张量并行、流水线并行和数据并行之间的权衡,并提供了如何配置大型模型分布式训练的直觉。我们的方法使我们能够在 3072 个 GPU 上以 502 petaFLOP/s 的速度对万亿参数模型执行训练迭代,每个 GPU 的吞吐量达到理论峰值的 52%。我们的代码已在 GitHub 开源。
1.6. 原文链接
https://arxiv.org/abs/2104.04473 该论文于 2021-04-09T00:00:00.000Z 发布在 arXiv。
2. 整体概括
2.1. 研究背景与动机
近年来,基于 Transformer 的语言模型 (Transformer-based language models) 在自然语言处理 (Natural Language Processing, NLP) 领域取得了飞速发展,这得益于计算资源的普及和数据集规模的扩大。大型语言模型 (Large Language Models, LLMs) 已被证明在零样本 (zero-shot) 或少样本 (few-shot) 学习中表现出色,并在许多 NLP 任务和数据集上实现了高准确率。这些模型在客户端反馈摘要、自动对话生成、语义搜索和代码自动补全等应用中具有广阔前景。
然而,随着模型参数数量呈指数级增长(如图 1 所示),训练这些模型面临着两大严峻挑战:
-
显存限制: 即使是拥有最大显存的 GPU(例如 NVIDIA 80GB-A100),也无法容纳这些模型的全部参数和中间激活值。将参数在主机内存和设备内存之间交换虽然可行,但会显著增加训练时间。
-
计算量巨大: 训练这些模型所需的浮点运算 (FLOPs) 数量庞大,可能导致不切实际的训练时间。例如,训练一个 1750 亿参数的
GPT-3模型在单个V100GPU 上需要大约 288 年。传统的数据并行 (data parallelism) 虽然可以扩展,但也存在局限性:
-
GPU 利用率下降: 当每个 GPU 的批次大小 (per-GPU batch size) 过小时,GPU 利用率会降低,通信成本会增加。
-
扩展性受限: 数据并行所能使用的最大设备数量受限于批次大小,限制了可用于训练的加速器数量。
为了应对这些挑战,研究人员提出了多种模型并行 (model parallelism) 技术,如张量模型并行 (tensor model parallelism) 和流水线模型并行 (pipeline model parallelism)。张量模型并行将单个
Transformer层内的矩阵乘法 (GEMMs) 拆分到多个 GPU 上,这对于高达 200 亿参数的模型在单个多 GPU 服务器内(如 NVIDIADGX A100服务器)效果良好。但对于更大的模型,需要跨多个多 GPU 服务器进行拆分,这会导致两个问题:a) 张量并行所需的all-reduce通信需要通过较慢的服务器间链路,而非服务器内部高速的NVLink;b) 过高的模型并行度可能导致小规模的矩阵乘法,降低 GPU 利用率。
流水线模型并行则将模型的层 (layers) 划分到不同的 GPU 上,并通过将批次 (batch) 拆分为更小的微批次 (microbatches) 来实现流水线式执行。然而,无论是何种调度策略,为了保持严格的优化器语义 (optimizer semantics),优化器步骤需要在设备间同步,导致每个批次结束时出现流水线冲洗 (pipeline flush),即所有微批次必须完成执行且不注入新的微批次。这可能导致高达 50% 的时间浪费在冲洗上,形成流水线气泡 (pipeline bubble)。虽然存在异步和有界陈旧性 (bounded-staleness) 的流水线方法,但它们放松了权重更新语义。
在实践中,用户可以使用多种并行技术训练大型模型,每种技术都有不同的权衡,并且这些技术可以组合使用。然而,组合这些技术会产生复杂的相互作用,需要仔细权衡以实现良好性能。本文旨在解决的核心问题是:在给定批次大小的情况下,如何组合并行技术以最大化大型模型的训练吞吐量,同时保持严格的优化器语义?
2.2. 核心贡献/主要发现
本文的主要贡献总结如下:
- 提出
PTD-P并行策略: 提出了一种名为PTD-P的新颖并行化策略,它有效地结合了流水线并行、张量并行和数据并行。该策略利用流水线并行跨多 GPU 服务器划分模型,张量并行在多 GPU 服务器内部划分模型,并结合数据并行进行扩展,从而在数千个 GPU 上实现了万亿参数模型的实用训练,并实现了优异的计算性能(达到峰值设备吞吐量的 52%)。 - 引入交错流水线并行调度: 提出了一种新颖的交错流水线并行调度 (interleaved pipeline parallelism schedule),与现有调度(如
GPipe和PipeDream-Flush)相比,在内存占用相当的情况下,可将吞吐量提高 10% 以上。 - 定量研究并行性权衡: 对张量并行、流水线并行和数据并行之间的性能交互进行了深入的定量研究,并提供了分布式训练配置的实践指导原则。
- 达到创纪录的训练吞吐量: 使用 3072 个
A100GPU 成功以 502 petaFLOP/s 的聚合吞吐量训练万亿参数的GPT模型,每个 GPU 的吞吐量达到理论峰值的 52%。这是目前针对此类规模模型最快的训练吞吐量,使得万亿参数模型的端到端训练时间估计为约 3 个月,具有实际可行性。 - 性能优化和工程实现: 实现了多轴创新和精心工程,包括高效的核函数 (kernel) 实现(使大部分计算为计算密集型而非内存密集型)、智能的计算图分区以减少网络通信量和设备空闲时间、领域特定的通信优化(如
scatter/gather优化),以及利用高速硬件(A100GPU 和高带宽互连)。 - 开源软件: 将
Megatron-LM代码库开源 (https://github.com/nvidia/megatron-lm),以便其他研究团队能够高效地训练大规模 NLP 模型。
3. 预备知识与相关工作
3.1. 基础概念
为了更好地理解本文,以下是一些关键的基础概念:
3.1.1. 大规模语言模型 (Large Language Models, LLMs)
- 概念定义:
LLMs是指参数量巨大(通常从数十亿到万亿级别)的深度学习模型,它们通过在海量文本数据上进行自监督学习(如预测下一个词元)来学习语言的复杂模式和结构。这些模型通常基于Transformer架构。 - 在本文中的重要性: 本文的核心目标就是高效训练这些参数规模不断增长的
LLMs,特别是Transformer架构的GPT系列模型。
3.1.2. Transformer 架构
- 概念定义:
Transformer[42] 是一种深度学习架构,最初用于机器翻译,后来成为LLMs的基石。其核心是自注意力机制 (self-attention mechanism),允许模型在处理序列数据时权衡不同位置的重要性。 - 在本文中的重要性:
Megatron-LM和本文提出的并行策略都是围绕Transformer层的结构进行优化的。理解Transformer层(由自注意力块和多层感知机 (MLP) 组成)如何被分割是理解张量并行和流水线并行的关键。
3.1.3. GPU 显存容量与计算量
- 概念定义:
- GPU 显存容量 (GPU Memory Capacity): 指 GPU 上用于存储模型参数、激活值、梯度和优化器状态的板载高速内存。
- 浮点运算 (Floating-point Operations, FLOPs): 是衡量计算复杂度的单位,表示每秒进行的浮点运算次数。通常使用
teraFLOP/s(每秒万亿次浮点运算) 或petaFLOP/s(每秒千万亿次浮点运算) 来衡量 GPU 或系统的性能。
- 在本文中的重要性: 这两者是训练
LLMs的两大主要瓶颈。显存限制了单个 GPU 能容纳的模型大小,而巨大的FLOPs需求导致训练时间过长。本文的并行策略旨在克服这些限制。
3.1.4. 分布式训练 (Distributed Training)
- 概念定义: 当单个设备无法处理模型或数据集的规模时,将训练任务分散到多个计算设备(如 GPU)上进行处理的技术。
- 在本文中的重要性: 本文的所有工作都是关于如何高效地进行分布式训练。
3.1.5. 数据并行 (Data Parallelism)
- 概念定义: 最常见的分布式训练方法。每个计算设备(例如一个 GPU)都拥有模型的完整副本。输入数据集被分成小批次 (mini-batches) 并分发给不同的设备。每个设备独立地计算其分配到的数据的梯度。然后,这些梯度通过通信操作(通常是
all-reduce)进行聚合,以更新所有设备上的模型参数,确保模型副本的一致性。 - 在本文中的重要性: 数据并行用于在模型参数和中间激活值能够适应单个模型并行组 (
model-parallel group) 的内存限制时,进一步扩展训练规模。
3.1.6. 模型并行 (Model Parallelism)
- 概念定义: 当单个设备无法存储模型的全部参数或激活值时,将模型本身的不同部分分配到不同的设备上。它允许多个设备协同工作来处理单个训练样本的计算。
- 在本文中的重要性: 这是本文解决超大模型训练问题的核心策略,包括张量并行和流水线并行。
3.1.7. 张量模型并行 (Tensor Model Parallelism)
- 概念定义: 一种细粒度的模型并行技术,它将单个模型层内部的计算(例如,矩阵乘法中的大矩阵)分解到多个设备上。例如,一个大的矩阵乘法 可以将矩阵 拆分到多个 GPU 上。
- 在本文中的重要性: 本文主要采用
Megatron-LM[40] 的张量并行策略,在单个多 GPU 服务器内部使用,以利用高速NVLink互连。
3.1.8. 流水线模型并行 (Pipeline Model Parallelism)
- 概念定义: 一种粗粒度的模型并行技术,它将模型的不同层分配给不同的设备(或设备组)。每个设备只负责模型的一部分层。训练批次被分割成更小的微批次 (microbatches),这些微批次像流水线一样依次通过各个设备。当前设备完成其层的计算后,将中间激活值传递给下一个设备。
- 微批次 (Microbatch): 原始训练批次 (
batch) 的一个更小的子集。流水线并行通过处理多个微批次来提高设备利用率。 - 流水线冲洗 (Pipeline Flush): 在流水线并行的某些调度中,为了保证优化器语义的严格性(即所有设备在更新参数时使用相同版本的权重),在每个完整批次的结束时,必须等待所有正在进行中的微批次完成其前向和反向传播,而不能注入新的微批次。
- 流水线气泡 (Pipeline Bubble): 由于流水线冲洗或流水线启动/结束时的设备空闲,导致设备未能充分利用其计算资源的时间。这个空闲时间是流水线并行效率损失的主要来源。
- 在本文中的重要性: 流水线并行用于跨多个多 GPU 服务器划分模型,以克服显存限制,并与张量并行和数据并行结合。本文还提出了改进的交错调度来减少流水线气泡。
3.1.9. 优化器语义 (Optimizer Semantics)
- 概念定义: 指优化器(如
SGD、Adam)在训练过程中更新模型参数的行为规范。严格的优化器语义要求所有设备在计算梯度和更新权重时,都基于相同版本的模型参数。如果不同设备看到不同版本的权重,可能导致收敛性问题或最终模型性能下降。 - 在本文中的重要性: 本文强调保持严格的优化器语义,这意味着必须处理好流水线并行带来的权重更新不一致问题(通过流水线冲洗),而不是采用放松语义的异步方法。
3.1.10. All-reduce 和点对点通信 (All-reduce and Point-to-point Communication)
- 概念定义:
- All-reduce: 一种集体通信操作,所有参与的进程(或 GPU)都贡献一个值,然后所有进程都收到所有贡献值的聚合结果(例如,求和、平均)。它通常用于数据并行中聚合梯度。
- 点对点通信 (Point-to-point Communication): 两个设备之间直接发送和接收数据的通信方式。
- 在本文中的重要性: 张量并行大量使用
all-reduce,而流水线并行主要使用点对点通信。不同通信模式的开销是决定并行策略组合的关键因素,尤其是在跨节点通信时。高速的NVLink用于节点内通信,而InfiniBand用于跨节点通信。
3.1.11. 激活重计算 (Activation Recomputation)
- 概念定义: 一种内存优化技术,通过在反向传播 (backward pass) 过程中重新计算部分前向传播 (forward pass) 的激活值,而不是将其全部存储在内存中。这可以显著减少显存占用,但会增加计算量(因为需要再次执行前向传播)。
- 在本文中的重要性: 对于训练非常大的模型,激活重计算是必不可少的,以将显存占用保持在可接受的水平。
3.1.12. 混合精度训练 (Mixed Precision Training)
- 概念定义: 在训练深度学习模型时,同时使用 16 位浮点数(如
FP16或BF16)和 32 位浮点数 (FP32)。FP16可以节省显存并加速计算,而FP32则用于需要更高数值精度的操作,以保持训练的稳定性和模型准确性。 - 在本文中的重要性: 本文的所有实验都使用了混合精度,这是提高训练效率和模型规模的常用技术。
3.2. 前人工作
本文建立在现有分布式训练和 Transformer 模型研究的基础上,并在此基础上进行了创新:
- Transformer 模型及应用 [11, 13, 27, 33-35, 42, 46]:
Transformer架构 [42] 及其变体(如BERT[13],GPT[33, 34],RoBERTa[27],T5[35])是LLMs的基石。GPT-3[11] 等大型模型展示了强大的零样本/少样本学习能力。本文的目标是高效训练这类模型。- Attention 机制:
Attention机制是Transformer的核心。其计算公式如下: 其中:- (Query), (Key), (Value) 分别是查询、键和值矩阵。它们通过将输入嵌入向量与不同的线性投影矩阵相乘得到。
- (其中 是序列长度, 是键/查询向量的维度)。
- 计算查询和键之间的相似度分数。
- 是一个缩放因子,用于防止点积结果过大,导致
softmax函数的梯度过小。 - 函数将分数转换为概率分布,确保注意力权重和为 1。
- 矩阵被加权求和,得到注意力机制的输出。
- Attention 机制:
- Megatron-LM [39, 40]:
Megatron率先展示了如何在GPU上使用张量模型并行训练大型Transformer模型。本文的实现是基于Megatron-LM代码库的扩展。Megatron的张量并行策略(如在MLP块和自注意力块中的矩阵分割)是本文张量并行部分的基础。 - 流水线并行调度 [14, 20, 23, 29, 30, 45]:
GPipe[20] 提出了一种默认调度,先执行所有微批次的前向传播,再执行所有微批次的反向传播。其缺点是需要存储所有微批次的激活值,内存占用高。PipeDream-Flush[30] 改进了GPipe,通过 1F1B (1 Forward, 1 Backward) 模式,将内存占用限制在流水线深度相关的微批次数量,而非总微批次数量。本文的非交错调度基于此。- 其他工作如
PipeDream[29]、PipeMare[45]、PipeDream-2BW[30, 45] 和Kosson et al.[23] 探索了异步或有界陈旧性 (bounded-staleness) 的流水线并行,这些方法通过放松严格的优化器语义来减少流水线气泡,但可能影响收敛性。本文侧重于严格优化器语义。
- 分片数据并行 (Sharded Data Parallelism) 和 ZeRO [24, 28, 36, 37, 44]:
ZeRO(Zero Redundancy Optimizer) [36, 37] 及其扩展ZeRO-Infinity[37] 通过分片优化器状态、梯度和模型参数来大幅减少每个GPU的内存占用。ZeRO-3是ZeRO最完整的阶段,它将所有模型状态(参数、梯度、优化器状态)分片到数据并行工作器中。本文与ZeRO-3进行了对比,并指出PTD-P在某些情况下表现更优。 - 自动分区 [14, 22, 29, 41]:
FlexFlow[22]、PipeDream[29]、DAPPLE[14] 和Tarnawski et al.[41] 等工作尝试自动分区模型训练图到多个设备上。本文采取的是启发式方法,而非自动搜索。
3.3. 技术演进
LLM 训练的并行化技术经历了从简单到复杂的演进:
-
数据并行 (Data Parallelism): 最早也是最直接的并行方法,通过复制模型并在不同数据批次上训练来加速,主要受限于单 GPU 显存和通信开销。
-
模型并行 (Model Parallelism) 出现: 随着模型规模超出单 GPU 显存,模型并行成为必然。
- 张量并行 (Tensor Parallelism):
Megatron-LM[40] 等工作展示了在层内进行张量分割的有效性,尤其在高速互连(如NVLink)的节点内部表现良好。 - 流水线并行 (Pipeline Parallelism):
GPipe[20]、PipeDream[29] 等工作引入了层间并行,通过微批次流水线化来减少空闲时间。
- 张量并行 (Tensor Parallelism):
-
分片数据并行 (Sharded Data Parallelism) 和优化器优化:
ZeRO[36] 等技术通过分片优化器状态、梯度乃至参数来进一步降低单GPU内存占用,使得更大模型的数据并行成为可能。 -
混合并行策略: 随着模型和集群规模的进一步扩大,单一并行策略的局限性显现。例如,张量并行在跨节点时通信开销巨大,而流水线并行存在流水线气泡。因此,将多种并行策略(如张量、流水线、数据并行)组合成为趋势。
本文的工作正处于这一演进的最新阶段,旨在通过巧妙组合和优化多种并行技术,解决当前
LLMs训练中面临的显存和计算效率瓶颈,尤其是在大规模GPU集群上的扩展问题。
3.4. 差异化分析
本文 Megatron-LM 的 PTD-P 方法与相关工作的核心区别和创新点在于:
- 多维度并行策略的组合与优化:
- 不同于单一模型并行: 像早期
Megatron[40] 专注于张量并行,或PipeDream[29] 专注于流水线并行,PTD-P明确提出并优化了张量并行 (节点内)、流水线并行 (跨节点) 和数据并行的组合使用。这种分层、智能的并行组合策略是其能够扩展到万亿参数模型的关键。 - 优于简单组合: 本文不仅简单组合,更深入研究了这些并行维度之间的非平凡交互,并提供了配置分布式训练的指导原则(例如,张量并行应在节点内部使用,流水线并行应跨节点使用)。
- 不同于单一模型并行: 像早期
- 新颖的交错流水线并行调度:
- 本文提出了交错流水线并行调度 (interleaved pipeline parallelism schedule),相比于传统的
GPipe[20] 或PipeDream-Flush[30] 调度,它通过让每个设备处理多个模型块(model chunk),有效减少了流水线气泡 (pipeline bubble),从而在内存占用相当的情况下,将吞吐量提高了 10% 以上。
- 本文提出了交错流水线并行调度 (interleaved pipeline parallelism schedule),相比于传统的
- 通信和计算的深度优化:
Scatter/Gather通信优化: 针对张量并行和流水线并行结合时的跨节点重复通信问题,提出了scatter/gather优化,将冗余数据分散传输,然后在接收端通过高速NVLink进行all-gather。这显著减少了跨节点InfiniBand链路的通信量,使交错调度更具可行性。- 计算核函数优化: 实现了数据布局优化(避免转置,启用步进批次
GEMM)、元素级操作融合 (PyTorch JIT自动融合) 和自定义核函数融合 (scale, mask, softmax),确保大部分计算是计算密集型 (compute-bound) 而非内存密集型 (memory-bound),从而提高GPU利用率。
- 在超大规模上的实际性能优势:
- 高于
DeepSpeed: 尽管DeepSpeed[2] 也结合了流水线、张量和数据并行来训练万亿参数模型,但本文的Megatron-LM实现了更高的吞吐量(峰值效率的 52% vs. 36%)。这得益于本文的核函数融合、更高效的流水线调度、更快的硬件(A100vs.V100)以及更大的GPU规模。 - 高于
ZeRO-3: 在与没有模型并行的ZeRO-3[36, 37] 进行比较时,PTD-P在增加GPU数量时展现出更优异的扩展性,对于 1750 亿和 5300 亿参数模型,其性能可高出 70%,主要原因是减少了跨节点通信。
- 高于
- 详细的性能分析和指导原则: 本文提供了对并行化维度之间权衡的深入分析,包括流水线气泡、通信量、计算效率和微批次大小的影响,为如何配置分布式训练提供了实用的指导原则。
4. 方法论
4.1. 方法原理
本文的核心思想是,为了高效地训练超大规模语言模型(万亿参数级别),单一的并行策略不足以应对 GPU 显存限制和巨大的计算需求。因此,需要智能地组合多种并行技术:
-
流水线并行 (Pipeline Parallelism) 负责在跨多 GPU 服务器的层级上划分模型,以克服单个服务器的显存限制。
-
张量并行 (Tensor Parallelism) 则用于在单个多 GPU 服务器内部划分模型层,以利用服务器内部高速的
NVLink互连,提高计算效率。 -
数据并行 (Data Parallelism) 用于在模型参数和中间激活值能够适应单个模型并行组 (
model-parallel group) 的内存限制时,进一步扩展训练规模,提高整体吞吐量。这种分层组合的策略被称为
PTD-P。此外,本文还通过提出新颖的交错流水线调度和**scatter/gather通信优化**,以及一系列计算优化,来解决组合并行带来的通信开销和效率问题,从而实现接近线性的大规模扩展。
4.2. 核心方法详解 (逐层深入)
4.2.1. PTD-P 的并行模式组合
本文提出结合流水线模型并行、张量模型并行和数据并行,称之为 PTD-P。这种组合方式如图 2 所示。
该图像是示意图,展示了两层变换器(Transformer layer)中的张量并行(Tensor MP)和流水线并行(Pipeline MP)分区方案。图中清晰地标示了两个变换器层及其对应的张量和流水线分区设计,帮助理解并行处理在模型训练中的应用。
图 2 (原文 Figure 2) 展示了 PTD-P 中张量并行和流水线并行的组合。
-
数据并行 (Data Parallelism):
- 每个工作器(或工作器组,对应于一个模型并行组)都拥有模型的一个完整副本(或模型的一个分片)。
- 输入数据集被分片并分发给不同的数据并行组。
- 每个数据并行组独立计算其分配到的数据的梯度。
- 周期性地,所有数据并行组通过
all-reduce操作聚合梯度,以确保所有工作器上的模型权重版本一致。 - 限制: 对于无法在单个工作器上存储的超大模型,数据并行需要与其他模型并行技术结合使用。
-
流水线模型并行 (Pipeline Model Parallelism):
- 模型的层 (
layers) 被分片(sharded)到多个设备上。如果模型由重复的Transformer块组成,则每个设备可以被分配相等数量的Transformer层。 - 一个批次 (
batch) 被拆分成更小的微批次 (microbatches),然后这些微批次在设备之间以流水线方式执行。 - 为了保持严格的优化器语义(确保输入在前向和反向传播中看到一致的权重版本),本文引入了周期性的流水线冲洗 (pipeline flushes)。这意味着在每个批次的开始和结束时,设备会处于空闲状态,形成流水线气泡 (pipeline bubble)。本文的目标是尽可能减小这个气泡。
- 模型的层 (
4.2.1.1. 流水线调度策略
* 默认调度 (Default Schedule) - PipeDream-Flush 变体:
* `GPipe` [20] 提出了一种调度,其中一个批次的所有微批次的前向传播首先执行,然后是所有微批次的反向传播(如图 3 所示)。这种调度的问题是内存占用高,因为它需要为所有 个微批次存储中间激活值。
* 本文采用的是 `PipeDream-Flush` [30] 调度。在这种调度中,首先进入一个预热阶段 (warm-up phase),工作器执行不同数量的前向传播(如图 4 顶部所示)。预热后进入稳态 (steady state),每个工作器执行一个前向传播 () 紧接着一个反向传播 ()(简称 1F1B)。最后,在一个批次结束时,完成所有剩余的、正在进行中的微批次的反向传播。
* 这种调度将飞行中的微批次(即反向传播尚未完成且需要维护激活值的微批次)数量限制在流水线深度 (),而非批次中的微批次总数 ()。因此,当 时,`PipeDream-Flush` 比 `GPipe` 内存效率更高。
* 流水线气泡大小 (Pipeline Bubble Size) 公式:
令 为批次中的微批次数量, 为流水线阶段数(用于流水线并行的设备数量), 为每个迭代的理想时间, 和 为单个微批次前向和反向传播的时间。
流水线气泡由开始时的 次前向传播和结束时的 次反向传播组成。
总的流水线气泡时间为: 。
理想的批次处理时间为: 。
因此,流水线气泡时间分数(即流水线气泡大小)为:
为了使气泡时间分数较小,需要 。

*该图像是一个示意图,展示了在四个设备上进行模型训练的前向传播和反向传播过程。图中时间轴上标记了每个设备在不同时间步进行前向(蓝色)和反向(绿色)传播的状态,并显示了管道冲洗期间设备的闲置状态。*
图 3 (原文 Figure 3) 展示了 `GPipe` 的调度方式,先所有前向,再所有反向,导致了大的流水线气泡。
* 交错阶段调度 (Schedule with Interleaved Stages) - 本文提出的新调度:
* 为了进一步减小流水线气泡,每个设备可以为多个层子集(称为模型块 (model chunk))执行计算,而不是单个连续的层集合。例如,如果之前每个设备有 4 层,现在可以将其分配为两个模型块(每个 2 层),即设备 1 负责层 1、2、9、10,设备 2 负责层 3、4、11、12,以此类推。
* 通过这种方式,流水线中的每个设备被分配了多个流水线阶段,每个阶段的计算量更少。
* 本文开发了一个适应内存高效 1F1B 调度的交错调度,如图 4 所示。此调度要求批次中的微批次数量是流水线并行度(流水线中的设备数量)的整数倍。
* **流水线气泡大小公式:**
如果每个设备有 个阶段(或模型块),则每个阶段或块的微批次前向和反向时间将变为 和 。
流水线气泡时间因此减少为: 。
气泡时间分数变为:
这意味着新调度将气泡时间减少了 倍。
* **权衡:** 气泡时间的减少不是免费的,此调度需要**额外的通信量**,通信量也增加了 倍。本文利用多 `InfiniBand` 网卡和 `scatter/gather` 优化来减轻此影响。

*该图像是一个示意图,展示了如何在多个设备上分配不同阶段的前向传递和反向传递计算。上方为传统的单阶段分配,下方为多阶段分配,提高了设备的利用率。该方法通过将多个阶段分配给每个设备来优化并行处理。*
图 4 (原文 Figure 4) 对比了 `PipeDream-Flush` (上图) 和本文提出的交错调度 (下图)。交错调度通过将每个设备分配多个模型块,使得流水线冲洗更快发生,从而减少了流水线气泡。
-
张量模型并行 (Tensor Model Parallelism):
-
模型中的单个层被分片到多个设备上。本文沿用
Megatron[40] 对Transformer层的特定分区策略,如图 5 所示。 -
MLP块中的并行化:MLP块由两个GEMM(矩阵乘法)和一个GeLU非线性激活函数组成:- 将第一个权重矩阵 沿其列进行分割 。这允许
GeLU非线性激活函数独立应用于每个分区的GEMM输出: 这种分割方式的优点在于无需同步,因为GeLU是非线性的。 - 第二个权重矩阵 的行可以沿其行进行分割,以消除
GEMM之间的通信需求(如图 5a 所示): 第二个GEMM的输出在经过dropout层之前,会在GPU之间进行all-reduce操作。
-
自注意力块中的并行化:
- 多头注意力操作的内在并行性被利用来分区自注意力块(如图 5b 所示)。
- 键 ()、查询 () 和值 () 矩阵可以进行列并行分区。
- 输出线性层可以直接作用于注意力操作的并行输出(权重矩阵按行分区)。
-
通信: 这种方法在
MLP和自注意力块中分割GEMM,在前向传播中只需要两次all-reduce操作( 运算符),在反向传播中也只需要两次all-reduce操作( 运算符)。 和 运算符在代码中实现。
Figure 5: Blocks of transformer model partitioned with tensor model parallelism (figures borrowed from Megatron [40]). and are conjugate. is the identity operator in the forward pass and allreduce in the backward pass, while is the reverse.
图 5 (原文 Figure 5) 展示了
Transformer模型的块如何通过张量模型并行进行分区。 和 是共轭的,其中 在前向传播中是恒等操作,在反向传播中是all-reduce; 则相反。 -
4.2.2. 并行化配置的性能分析
本文分析了将流水线并行、张量模型并行与数据并行结合时的性能影响。在给定 GPU 预算和批次大小的情况下,不同的并行度选择会导致内存占用、设备利用率和通信量之间的权衡。
- 符号约定:
- : 分别代表流水线模型并行大小、张量模型并行大小和数据并行大小。
- :
GPU总数。要求 。 - : 全局批次大小 (Global Batch Size)。
- : 微批次大小 (Microbatch Size)。
- : 每个流水线的微批次数量。
4.2.2.1. 张量并行与流水线模型并行
-
流水线气泡: 假设 (数据并行大小),则 。流水线气泡大小为 。
- 随着 增加(张量并行度提高),流水线气泡会减小(在固定
B, b, d的情况下, 也固定)。
- 随着 增加(张量并行度提高),流水线气泡会减小(在固定
-
通信量:
- 流水线模型并行: 采用廉价的点对点通信。每个连续设备对(前向或反向传播)的通信量为
bsh,其中 为序列长度, 为隐藏层大小。 - 张量模型并行: 采用
all-reduce通信。每个微批次,每个层的前向和反向传播都需要两次all-reduce操作,涉及大小为bsh的张量。每个设备通常有多个层,总的张量并行通信量为 ,其中 是流水线阶段中的层数。 - 结论: 张量模型并行增加了设备之间的通信量。当 大于单个节点中的
GPU数量时,跨较慢的节点间链路执行张量模型并行会导致性能下降。
- 流水线模型并行: 采用廉价的点对点通信。每个连续设备对(前向或反向传播)的通信量为
-
经验法则 #1: 在考虑不同形式的模型并行时,张量模型并行通常应在 个
GPU服务器内部使用,并行度最高为 ;然后流水线模型并行可以用于跨服务器扩展到更大的模型。
4.2.2.2. 数据并行与模型并行
-
数据并行与流水线模型并行:
-
假设 (张量模型并行大小)。每个流水线的微批次数量为 ,其中 。
-
GPU总数为 ,流水线阶段数为 。 -
流水线气泡大小为:
-
随着 变大(数据并行度提高),
n-d变小,因此流水线气泡变小(如图 6 所示)。 -
结论: 如果数据并行所需的
all-reduce通信不随 急剧增加(环形实现的通信时间按 缩放),总吞吐量将会提高。增加批次大小 也会增加吞吐量,因为 增加,气泡减小,且数据并行all-reduce通信频率降低。
Figure 6: Fraction of time spent idling due to pipeline flush (pipeline bubble size) versus data-parallel size ( d ), for different numbers of GPUs( n )and ratio of batch size to microbatch size .
图 6 (原文 Figure 6) 展示了流水线冲洗导致的空闲时间(流水线气泡大小)与数据并行大小 之间的关系,针对不同数量的
GPU和批次大小与微批次大小的比率 。 -
-
数据并行与张量模型并行:
- 张量模型并行需要在每个微批次中执行
all-reduce通信,这在跨多GPU服务器时可能非常昂贵。 - 数据并行只需在每个批次中执行一次昂贵的
all-reduce通信。 - 此外,随着张量模型并行度增加,每个模型并行等级 (
model-parallel rank) 执行的计算量减少,对于不够大的层,现代GPU可能无法以峰值效率执行这些子矩阵计算。
- 张量模型并行需要在每个微批次中执行
-
经验法则 #2: 在使用数据并行和模型并行时,应使用总模型并行大小 ,以确保模型的参数和中间元数据能够适应
GPU显存;然后数据并行可以用于扩展训练到更多GPU。
4.2.2.3. 微批次大小 (Microbatch Size)
-
影响: 微批次大小 的选择会影响模型训练吞吐量。它同时影响操作的算术强度 (arithmetic intensity) 和流水线气泡大小。
-
权衡: 增加微批次大小会减少流水线中的微批次数量 ,导致更大的流水线气泡;但同时也可以通过增加执行核函数的算术强度来提高
GPU利用率(如图 7 所示)。 -
分析: 忽略通信成本,总批次计算时间为 。
-
结论: 最佳微批次大小是模型相关的(如图 8 所示),需要权衡上述两个因素。
Figure 7: Per-GPU throughput versus microbatch size for a GPT model with a billion parameters (128 attention heads, hidden size of 4096, 4 transformer layers).图 7 (原文 Figure 7) 展示了
GPT模型(10 亿参数)的每个GPU吞吐量与微批次大小的关系。
Figure 8: Behavior of normalized estimated throughput (time computed as with respect to the microbatch size for the same GPT model from Figure 7.图 8 (原文 Figure 8) 展示了对于与 Figure 7 相同的
GPT模型,归一化估计吞吐量(时间计算公式为 )随微批次大小 的变化行为。 -
经验法则 #3: 最佳微批次大小 取决于模型的吞吐量和内存占用特性,以及流水线深度 、数据并行大小 和批次大小 。
4.2.2.4. 激活重计算 (Activation Recomputation)
- 原理: 一种可选技术,通过在反向传播之前再次运行前向传播来增加计算操作,从而减少内存占用(只存储给定流水线阶段的输入激活值,而不是整个中间激活值集)。
- 必要性: 对于训练合理规模的大型模型,激活重计算是必需的,以将内存占用保持在可接受的低水平。
- 检查点数量: 激活检查点 (
activation checkpoints) 的数量不影响吞吐量,但影响内存占用。如果一个模型阶段有 层, 是检查点数量,总内存占用约为 ,其中 是层的输入激活值大小, 是每层的中间激活值。理论最佳是在 时获得。实践中,每 1 或 2 个Transformer层设置一个检查点是最佳的。 - 其他技术: 激活分区 (
activation partitioning) [36] 也可以与张量模型并行结合使用,以进一步减少激活值导致的内存占用。
4.2.3. 实现细节
本文在 Megatron-LM 代码库的基础上,使用 PyTorch [32] 实现了 PTD-P,并使用 NCCL [7] 进行设备间通信。为了获得良好性能,实现了以下通信和计算优化:
4.2.3.1. 通信优化
-
Scatter/Gather通信优化:-
问题: 在使用流水线并行时,希望并行发送和接收前向和反向张量。一个
DGX A100配备 8 块InfiniBand(IB) 网卡,但点对点通信(在两个服务器上的GPU对之间进行)难以利用所有 8 块网卡。 -
冗余: 在流水线并行和张量并行结合时,相邻流水线阶段中执行张量并行的
rank会发送和接收完全相同的张量集(如图 9a 所示),导致大量冗余通信。 -
优化: 对于足够大的模型(例如,张量模型并行大小为 8),我们将发送端的张量分割成等大小的块 (
chunks),然后每个rank只将一个块发送到下一个节点的对应rank,使用各自的InfiniBand网卡(例如,rank 1发送到rank 3,rank 2发送到rank 4,如图 9b 所示)。在接收端,通过NVLink执行all-gather操作,重新构建完整的张量。 -
效果: 这种优化更好地利用了
DGX A100服务器上的多块IB网卡,并使交错调度等通信密集型调度变得可行。定量上,有了scatter/gather优化,每对连续阶段之间所需的总通信量减少到 ,其中 是张量模型并行大小。
该图像是示意图,展示了Infiniband连接和数据分散、合并机制的工作原理。左侧部分为Infiniband连接的配置,右侧展示了Scatter和All-gather通信方式。
图 9 (原文 Figure 9) 展示了
scatter/gather通信优化。左图 (a) 未优化,相同张量通过跨节点InfiniBand链接冗余发送。右图 (b) 优化后,发送端将张量分散成小块,每块只发送一次,然后在接收端通过all-gather重新组合,减少了跨节点通信量。 -
4.2.3.2. 计算优化
为了实现高性能,本文对计算图进行了三项模型特定优化:
- 数据布局修改: 将
Transformer层中的数据布局从[b, s, a, h](批次、序列、注意力头、隐藏层大小)更改为[s, b, a, h]。这避免了内存密集型的转置操作,并使得使用步进批次GEMM核函数成为可能。 - 融合核函数 (Fused Kernels) for 元素级操作: 使用
PyTorch JIT[10] 为一系列元素级操作(如bias + GeLU和bias + dropout + add)生成了融合核函数。 - 自定义融合核函数 for
Attention操作: 创建了两个自定义核函数,用于融合scale、mask和softmax(归约) 操作:一个支持通用掩码(如BERT),另一个支持隐式因果掩码(如自回归模型GPT)。
5. 实验设置
5.1. 数据集
实验使用了不同规模的 GPT 模型进行评估,模型参数从 10 亿到 1 万亿不等。
- 所有模型均使用词汇表大小 (1024 的倍数) 和序列长度 。
- 通过调整隐藏层大小 、注意力头数量和层数 来改变模型规模。
- 模型参数 的计算公式为:
- 符号解释:
- : 模型参数总数。
- :
Transformer层数。 - : 隐藏层大小。
- : 词汇表大小。
- : 序列长度。
- 符号解释:
5.2. 评估指标
本文使用了以下评估指标来衡量模型的训练效率和性能:
5.2.1. 每 GPU 吞吐量 (Achieved teraFLOP/s per GPU)
- 概念定义: 每 GPU 吞吐量衡量每个
GPU每秒执行的万亿次浮点运算 (teraFLOP/s)。这是一个衡量单个GPU计算效率和利用率的关键指标,越高越好。它反映了GPU在实际训练任务中的有效计算能力,包括了计算、通信、数据处理和优化器步骤等所有开销。 - 数学公式:
- 符号解释:
- : 每个
GPU的吞吐量 (teraFLOP/s)。 - : 单次迭代中模型执行的总浮点运算次数 (FLOPs)。
- : 完成单次训练迭代所需的时间。
- : 使用的
GPU总数。
- : 每个
5.2.2. 理论峰值 FLOP/s 的百分比 (Percentage of theoretical peak FLOP/s)
- 概念定义: 该指标表示实际达到的每
GPU吞吐量占GPU理论最大浮点运算能力(即峰值FLOP/s)的百分比。它衡量了GPU硬件被算法和实现利用的效率,百分比越高说明利用率越高,系统开销越小。 - 数学公式:
- 符号解释:
- : 达到理论峰值
FLOP/s的百分比。 - : 实际达到的每个
GPU吞吐量 (teraFLOP/s)。 - : 单个
GPU的理论峰值浮点运算能力 (teraFLOP/s)。对于A100GPU,16 位精度下的理论峰值为 312 teraFLOP/s。
- : 达到理论峰值
5.2.3. 聚合 petaFLOP/s (Achieved aggregate petaFLOP/s)
- 概念定义: 聚合
petaFLOP/s是所有GPU共同达到的总浮点运算能力,以每秒千万亿次浮点运算表示。它反映了整个GPU集群的总体计算吞吐量,是衡量大规模分布式系统性能的指标。 - 数学公式:
- 符号解释:
- : 聚合
petaFLOP/s。 - : 每个
GPU的吞吐量 (teraFLOP/s)。 - : 使用的
GPU总数。
- : 聚合
5.2.4. 训练时间估算 (Training Time Estimates)
- 概念定义: 基于模型大小、总训练词元数和实际达到的吞吐量,估算出完成整个模型训练所需的时间。这是衡量训练实践可行性的一个重要指标。
- 数学公式:
结合模型参数 和
FLOPs计算公式 ,以及模型参数 ,并考虑到文中提到的近似条件 (, , ),可以推导出: - 符号解释:
- : 端到端训练所需总时间。
- : 总训练词元数 (tokens)。
- : 模型参数总数。
- : 使用的
GPU总数。 - : 每个
GPU的吞吐量 (teraFLOP/s)。
5.2.5. FLOPs 计算 (Floating-Point Operations Calculation)
根据论文附录,Transformer 模型中 FLOPs 的计算方式如下:
- 矩阵乘法 : 需要 次
FLOPs(乘法和加法各算一次)。 Transformer层: 由一个注意力块和一个 2 层前馈网络 (feed-forward network) 组成。- 注意力块:
- 键、查询、值变换:
FLOPs。 - 注意力矩阵计算:
FLOPs。 - 注意力 over values:
FLOPs。 - 后注意力线性投影:
FLOPs。
- 键、查询、值变换:
- 前馈网络: 隐藏层大小从 增加到
4h,再减少到 。需要FLOPs。 - 单次前向传播总和:
FLOPs。 - 考虑反向传播和激活重计算: 反向传播需要两倍的前向传播
FLOPs。激活重计算额外需要一次前向传播。因此,每个Transformer层总FLOPs为 。
- 注意力块:
Logit层:Logit层将特征维度 转换为词汇维度 。- 前向传播:
2 B s h VFLOPs。 - 反向传播:
4 B s h VFLOPs。 - 总计:
6 B s h VFLOPs。
- 前向传播:
- 总
FLOPs(Total Floating-Point Operations): 对于一个包含 个Transformer层的模型,总的FLOPs为:- 符号解释:
- : 单次迭代中的总浮点运算次数。
- : 训练批次大小。
- : 序列长度。
- :
Transformer层数。 - : 隐藏层大小。
- : 词汇表大小。
- 符号解释:
5.3. 对比基线
论文将 PTD-P 方法与以下基线进行了比较:
-
ZeRO-3 [36, 37]:
DeepSpeed库 [3] 中实现的ZeRO-3。ZeRO-3是一种分片数据并行方法,它将优化器状态、梯度和模型参数分片到数据并行工作器中。在比较中,ZeRO-3在没有使用任何模型并行的情况下进行,以展示PTD-P组合模型并行的优势。 -
Megatron [40] (仅张量并行):
Megatron-LM之前的版本主要侧重于张量模型并行。论文通过实验对比了仅使用张量并行和结合使用张量与流水线并行的性能差异。 -
PipeDream [30] (仅流水线并行): 之前的流水线并行工作(如
PipeDream)在不结合张量并行时,在处理超大模型时可能面临挑战。论文通过实验对比了仅使用流水线并行和结合使用张量与流水线并行的性能差异。这些基线具有代表性,因为它们是当前分布式训练领域最先进或最流行的并行化策略。与
ZeRO-3的比较旨在突显PTD-P在减少跨节点通信方面的优势;与仅使用张量或流水线并行的方法的比较,则旨在证明PTD-P组合策略的必要性和优越性。
5.4. 硬件平台
所有实验均在 Selene 超级计算机 [8] 上使用混合精度进行。
- 节点配置: 每个集群节点有 8 块 NVIDIA 80-GB
A100GPU[6]。 - 节点内互连: 节点内的
GPU之间通过NVLink和NVSwitch[9] 连接,提供高带宽的节点内通信。 - 节点间互连: 每个节点有 8 块 NVIDIA Mellanox 200Gbps HDR
InfiniBand HCA用于应用通信,另有 2 块HCA用于专用存储。节点之间以三级(叶、脊、核心)胖树拓扑连接,包含 850 个交换机,这种拓扑结构支持高效的all-reduce通信。 - 存储: 集群使用全
NVME共享并行文件系统,提供高性能数据访问和存储。 - GPU 峰值吞吐量:
A100 GPU在 16 位精度下的理论峰值设备吞吐量为 312 teraFLOP/s。
6. 实验结果与分析
6.1. 核心结果分析
6.1.1. 端到端性能 (End-to-End Performance)
下表展示了 GPT 模型从 10 亿参数到 1 万亿参数的弱扩展吞吐量 (weak-scaling throughput)。弱扩展是指随着 GPU 数量的增加,模型规模也相应增加,旨在评估方法在处理更大问题时的可扩展性。
以下是原文 Table 1 的结果:
| Number of parameters (billion) | Attention heads | Hidden size | Number of layers | Tensor model-parallel size | Pipeline model-parallel size | Number of GPUs | Batch size | Achieved teraFLOP/s per GPU | Percentage of theoretical peak FLOP/s | Achieved aggregate petaFLOP/s |
| 1.7 | 24 | 2304 | 24 | 1 | 1 | 32 | 512 | 137 | 44% | 4.4 |
| 3.6 | 32 | 3072 | 30 | 2 | 1 | 64 | 512 | 138 | 44% | 8.8 |
| 7.5 | 32 | 4096 | 36 | 4 | 1 | 128 | 512 | 142 | 46% | 18.2 |
| 18.4 | 48 | 6144 | 40 | 8 | 1 | 256 | 1024 | 135 | 43% | 34.6 |
| 39.1 | 64 | 8192 | 48 | 8 | 2 | 512 | 1536 | 138 | 44% | 70.8 |
| 76.1 | 80 | 10240 | 60 | 8 | 4 | 1024 | 1792 | 140 | 45% | 143.8 |
| 145.6 | 96 | 12288 | 80 | 8 | 8 | 1536 | 2304 | 148 | 47% | 227.1 |
| 310.1 | 128 | 16384 | 96 | 8 | 16 | 1920 | 2160 | 155 | 50% | 297.4 |
| 529.6 | 128 | 20480 | 105 | 8 | 35 | 2520 | 2520 | 163 | 52% | 410.2 |
| 1008.0 | 160 | 25600 | 128 | 8 | 64 | 3072 | 3072 | 163 | 52% | 502.0 |
分析:
- 超线性扩展: 随着模型规模和
GPU数量的增加,本文的方法展示出接近线性的甚至在某些情况下是超线性的扩展。例如,从 1.7 亿参数(32GPU)到 1 万亿参数(3072GPU),每个GPU的吞吐量从 137 teraFLOP/s 提高到 163 teraFLOP/s。这表明GPU利用率随着模型规模的增大而提高(更大的矩阵乘法),而通信时间相对于计算时间的增加并不显著。 - 高
GPU利用率: 对于最小的模型(17 亿参数),实现了 44% 的理论峰值设备吞吐量;对于最大的万亿参数模型,实现了 52% 的理论峰值设备吞吐量。这表明PTD-P能够有效地利用A100 GPU的计算能力。 - 聚合吞吐量: 对于万亿参数模型,在 3072 个
A100 GPU上实现了惊人的 502.0 petaFLOP/s 聚合吞吐量。 - 实际训练时间:
- 训练 1750 亿参数的
GPT-3模型(3000 亿词元)在 1024 个A100 GPU上,估计需要 34 天。 - 训练 1 万亿参数模型(假设需要 4500 亿词元)在 3072 个
A100 GPU上,估计需要 84 天(约 3 个月)。 这些时间表明,本文的方法使得训练超大规模模型在合理的时间范围内成为可能。
- 训练 1750 亿参数的
6.1.2. 与 ZeRO-3 的比较
下表和图 10 展示了 PTD-P 与 ZeRO-3 在 1750 亿和 5300 亿参数 GPT 模型上的吞吐量比较。
以下是原文 Table 2 的结果:
| Scheme | Number of parameters (billion) | Model-parallel size | Batch size | Number of GPUs | Microbatch size | Achieved teraFLOP/s per GPU | Training time for 300B tokens (days) |
| ZeRO-3 without Model Parallelism | 174.6 | 1 | 1536 | 384 | 4 | 144 | 90 |
| 768 | 2 | 88 | 74 | ||||
| 1536 | 1 | 44 | 74 | ||||
| 529.6 | 1 | 2560* | 640 | 4 | 138 | 169 | |
| 2240 | 1120 | 2 | 98 | 137 | |||
| 2240 | 1 | 48 | 140 | ||||
| PTD Parallelism | 174.6 | 96 | 1536 | 384 | 1 | 153 | 84 |
| 768 | 1 | 149 | 43 | ||||
| 1536 | 1 | 141 | 23 | ||||
| 529.6 | 280 | 2240 | 560 | 1 | 171 | 156 | |
| 1120 | 1 | 167 | 80 | ||||
| 2240 | 1 | 159 | 42 |
Figure 10: Throughput per GPU of PTD-P and ZeRO-3 for two different GPT models (the 175B GPT-3 model is shown with dotted lines, and the 530B model is shown with solid lines). Global batch sizes are fixed and ZeRO-3 is used without any model parallelism.
图 10 (原文 Figure 10) 展示了 PTD-P 和 ZeRO-3 在两个不同 GPT 模型(1750 亿参数 GPT-3 模型为虚线,5300 亿参数模型为实线)上的每个 GPU 吞吐量。全局批次大小固定,ZeRO-3 在没有模型并行的情况下使用。
分析:
- PTD-P 性能优势: 在
GPU数量较少且微批次大小为 4 时,PTD-P对于 1750 亿参数模型实现了 6% 的更高吞吐量,对于 5300 亿参数模型实现了 24% 的更高吞吐量。 - 扩展性更优: 随着
GPU数量的增加(例如,GPU数量翻倍,而全局批次大小保持不变),PTD-P比单独使用ZeRO-3的扩展性更好。对于两种模型,PTD-P的性能比ZeRO-3高出 70%,这主要归因于PTD-P减少了跨节点通信。 - 跨节点通信是瓶颈:
ZeRO-3需要在数据并行工作器之间分片参数和梯度,这导致了额外的通信。当GPU数量增加时,这种跨节点通信成为主要瓶颈,而PTD-P的分层并行策略(张量并行在节点内,流水线并行跨节点)更好地管理了通信开销。
6.1.3. 流水线并行 (Pipeline Parallelism)
该图像是一个图表,展示了在不同的管道并行大小下,每个GPU所实现的teraflop/s数值。蓝色线条代表批量大小为8的情况,橙色线条则代表批量大小为128,在管道并行大小从1增加到8时,性能相对下降趋势明显。
图 11 (原文 Figure 11) 展示了在弱扩展实验设置下(模型大小随流水线并行大小增加),使用两种不同批次大小的流水线并行每个 GPU 吞吐量。
分析:
-
弱扩展: 在弱扩展设置中,随着流水线并行大小的增加,模型大小也按比例增加。图 11 显示,较高批次大小(128)的吞吐量比低批次大小(8)的吞吐量更高,并且随着流水线并行大小的增加,其性能下降相对更小。
-
流水线气泡影响: 较高的批次大小能更好地分摊流水线气泡(如公式 所示,当 增大时,气泡影响减小),从而实现更好的扩展性。
Figure 11: Throughput per GPU of pipeline parallelism using two different batch sizes in a weak-scaling experiment setup (model size increases with the pipeline-parallel size). Figure 12: Throughput per GPU of interleaved and non-interleaved schedules for a GPT model (175 billion parameters) on 96 GPUs.
图 12 (原文 Figure 12) 展示了在 96 个 GPU 上,对于 1750 亿参数的 GPT 模型,交错调度和非交错调度的每个 GPU 吞吐量。
分析:
- 交错调度的优势: 交错调度结合
scatter/gather通信优化,比非交错(默认)调度具有更高的计算性能。对于通信密集型场景,交错调度可以提高高达 10% 的吞吐量。 - 批次大小影响: 随着批次大小的增加,两种调度之间的性能差距缩小。原因有二:a) 批次大小增加时,默认调度的气泡大小减小;b) 流水线内部的点对点通信量与批次大小成正比,通信量增加时,非交错调度也能赶上。
Scatter/Gather的重要性: 如果不使用scatter/gather优化,在较大批次大小下,默认调度甚至会优于交错调度,这表明scatter/gather优化对于实现交错调度的性能优势至关重要。
6.1.4. 并行配置比较
6.1.4.1. 张量并行与流水线并行
Figure 13: Throughput per GPU of various parallel configurations that combine pipeline and tensor model parallelism using a GPT model with 162.2 billion parameters and 64 A100 GPUs.
图 13 (原文 Figure 13) 展示了在 64 个 A100 GPU 上,对于 1612 亿参数 GPT 模型,结合流水线并行和张量模型并行的各种并行配置的每个 GPU 吞吐量。
分析:
- 组合并行的重要性: 图 13 强烈表明,单独使用张量模型并行或流水线模型并行都无法达到最佳性能。当张量并行度和流水线并行度结合时,可以实现峰值性能。
- 张量并行在节点内: 当张量并行度等于单个节点内的
GPU数量(DGX A100节点为 8)时,性能达到峰值。这验证了经验法则 #1:张量并行最适合在节点内部利用高速NVLink进行,因为它涉及昂贵的all-reduce通信。 - 流水线并行跨节点: 流水线并行使用更便宜的点对点通信,因此更适合跨节点扩展。但流水线阶段总数应限制在合理范围内,以避免过大的流水线气泡。
- 结论: 结合两种技术可以有效降低通信开销并提高计算资源利用率,超越任何单一模型并行策略。
6.1.4.2. 流水线并行与数据并行
Figure 14: Throughput per GPU of various parallel configurations that combine data and pipeline model parallelism using a GPT model with 5.9 billion parameters, three different batch sizes, microbatch size of 1, and 64 A100 GPUs.
图 14 (原文 Figure 14) 展示了在 64 个 A100 GPU 上,对于 59 亿参数 GPT 模型,结合数据并行和流水线模型并行的各种并行配置的每个 GPU 吞吐量,使用三种不同批次大小和微批次大小为 1。
分析:
- 流水线并行度的影响: 对于每个批次大小,随着流水线并行度增加,每个
GPU的吞吐量都会下降。这与分析模型(公式 )相符,即流水线气泡随流水线并行度()的增加而增大。 - 数据并行的作用: 流水线模型并行主要用于支持无法在单个工作器上容纳的大型模型训练;一旦模型能够适应,数据并行应被优先用于扩展训练规模,因为其通信效率更高。
6.1.4.3. 张量并行与数据并行
Figure 15: Throughput per GPU of various parallel configurations that combine data and tensor model parallelism using a GPT model with 5.9 billion parameters, three different batch sizes, microbatch size of 1, and 64 A100 GPUs.
图 15 (原文 Figure 15) 展示了在 64 个 A100 GPU 上,对于 59 亿参数 GPT 模型,结合数据并行和张量模型并行的各种并行配置的每个 GPU 吞吐量,使用三种不同批次大小和微批次大小为 1。
分析:
- 张量并行的通信开销: 随着批次大小的增加(在微批次大小为 1 的情况下),数据并行通信变得不频繁。然而,张量模型并行中所需的
all-to-all通信在每个微批次中都需要执行,这会主导端到端训练时间,尤其是在跨多GPU节点进行通信时。 - GPU 利用率下降: 随着张量模型并行度增加,每个
GPU执行的矩阵乘法规模变小,可能导致GPU利用率下降。 - 数据并行的优势: 数据并行虽然扩展性有限(受限于批次大小),但其通信效率较高。对于非常大的模型,单一数据并行受限于显存容量和批次大小。
6.1.5. 微批次大小 (Microbatch Size)
Figure 16: Throughput per GPU of a parallel configuration for different microbatch sizes on a GPT model with 91 billion parameters, for two different batch sizes using 64 A100 GPUs.
图 16 (原文 Figure 16) 展示了在 64 个 A100 GPU 上,对于 910 亿参数 GPT 模型,并行配置 在两种不同批次大小下,每个 GPU 吞吐量与微批次大小的关系。
分析:
- 最佳微批次大小: 对于 910 亿参数模型,最佳微批次大小为 2。这表明微批次大小并非越大越好,存在一个最佳值。
- 权衡: 增加微批次大小会减少流水线中的微批次数量 ,导致更大的流水线气泡;但同时也能通过增加执行核函数的算术强度来提高
GPU利用率。这两个因素相互冲突,使得最佳微批次大小的选择具有挑战性。 - 指导意义: 本文的分析模型(第 3.3 节)可以作为确定不同训练配置和模型最佳微批次大小的代理。
6.1.6. 激活重计算 (Activation Recomputation)
Figure 17: Throughput (in sequences per second) with and without activation recomputation for a GPT model with 145 billion parameters using 128 A100 GPUs .
图 17 (原文 Figure 17) 展示了在 128 个 A100 GPU 上,对于 1450 亿参数 GPT 模型(并行配置 ),带激活重计算和不带激活重计算的吞吐量(以每秒序列数计)。
分析:
- 小批次下的性能下降: 对于小批次大小,激活重计算导致吞吐量降低高达 33%,因为在反向传播过程中需要额外执行一次前向传播。
- 大批次下的必要性与性能提升: 然而,激活重计算对于支持更大的批次大小是必需的,因为它显著减少了内存占用。在内存限制下,启用激活重计算能够使用更大的批次,从而减小流水线气泡。因此,在大批次下,带激活重计算的吞吐量比不带激活重计算(但使用较小批次)的最高吞吐量高出 2 倍。这表明,虽然激活重计算增加了计算量,但通过允许更大的批次,最终可以提高整体训练效率。
6.1.7. Scatter/Gather 优化
Figure 18: Throughput per GPU with and without the scatter/gather optimization for a GPT model with 175 billion parameters using 96 A100 GPUs and the interleaved schedule.
图 18 (原文 Figure 18) 展示了在 96 个 A100 GPU 上,对于 1750 亿参数 GPT-3 模型,使用交错调度时,带和不带 scatter/gather 优化的每个 GPU 吞吐量。
分析:
Scatter/gather优化:对于通信密集型调度(特别是大批次下的交错调度),scatter/gather通信优化带来了高达 11% 的吞吐量提升。这是通过减少跨节点链路上的通信量实现的,验证了其在复杂并行策略中的重要性。
6.1.8. 融合算子 (Fused Operators)
- 效果: 计算优化中描述的算子融合显著提升了性能。
- 对于 1750 亿参数
GPT-3模型,吞吐量提升了 19%(从 113 teraFLOP/s/GPU 到 135 teraFLOP/s/GPU)。 - 对于 5300 亿参数
GPT模型,吞吐量提升了 11%(从 133 teraFLOP/s/GPU 到 148 teraFLOP/s/GPU)。 这表明确保大部分计算是计算密集型而非内存密集型是提高效率的关键。
- 对于 1750 亿参数
6.1.9. 节点间通信带宽 (Inter-Node Communication Bandwidth)
- 重要性: 论文的强大结果离不开优化的软硬件栈。特别是高带宽的
GPU之间以及服务器之间的通信链路至关重要。 - 实际带宽: 在 3072 个
GPU上训练万亿参数模型时,观察到流水线阶段之间点对点通信的有效二分带宽 (bisection bandwidth) 为 892 GB/s;数据并行副本之间all-reduce操作的有效二分带宽为 12.9 TB/s。 - 结论: 任何优化较差的算子分区策略都会导致更多的节点间通信,从而阻碍扩展性能。
6.1.10. 检查点加载和保存 (Checkpoint Loading and Saving)
- 挑战: 大型模型的检查点(
checkpoints)尺寸巨大。例如,万亿参数模型的检查点大小为 13.8 TB。 - 性能:
- 万亿参数模型首次加载检查点时,所有 384 个节点(3072
GPU)达到了 1 TB/s 的峰值读取带宽,这是并行文件系统可达到的最大读取吞吐量。 - 检查点保存达到了峰值写入带宽的 40%(273 GB/s)。 这些结果表明,在实践中,检查点管理对于大规模模型训练同样至关重要,需要高性能存储系统支持。
- 万亿参数模型首次加载检查点时,所有 384 个节点(3072
6.2. 消融实验/参数分析
本文没有明确划分“消融实验”章节,但其对不同并行策略组合、微批次大小、激活重计算和 scatter/gather 优化的分析,实际上起到了消融实验的作用。
-
并行策略组合分析(图 13, 14, 15): 通过比较不同
(t, p, d)组合下的吞吐量,展示了各个并行维度如何相互作用以及各自的最佳应用场景。例如,图 13 有效地“消融”了单一张量并行或单一流水线并行的方案,证明了其组合的优越性。 -
微批次大小分析(图 16): 通过改变微批次大小,评估其对流水线气泡和
GPU利用率的权衡影响,找到了最佳配置。 -
激活重计算分析(图 17): 对比了开启和关闭激活重计算时的吞吐量,揭示了其在内存和计算之间的权衡,以及在大批次训练中的必要性。
-
Scatter/Gather优化分析(图 18): 对比了带和不带scatter/gather优化时的吞吐量,直接量化了该通信优化的性能贡献。 -
融合算子分析: 论文指出,算子融合使 1750 亿参数模型吞吐量提高了 19%,5300 亿参数模型提高了 11%。这表明融合算子是提升计算效率的关键组件。
这些分析共同验证了
PTD-P各个组件和优化措施的有效性,并提供了如何配置这些超参数以获得最佳性能的指导。
7. 总结与思考
7.1. 结论总结
本文成功展示了 PTD-P(节点间流水线并行、节点内张量并行和数据并行)如何有效地组合,以在 GPU 集群上实现万亿参数语言模型的高效训练。核心发现和贡献包括:
- 分层并行策略: 通过智能地将张量并行应用于节点内部(利用高速
NVLink)和流水线并行应用于节点之间(利用点对点通信),有效克服了GPU显存和通信带宽的限制。数据并行则在模型并行组之上进一步扩展了训练规模。 - 新颖的交错流水线调度: 提出的交错流水线并行调度结合
scatter/gather通信优化,显著减少了流水线气泡,并将吞吐量提高了 10% 以上,同时保持了内存效率。 - 创纪录的性能: 在 3072 个
A100 GPU上实现了 502 petaFLOP/s 的聚合吞吐量,每个GPU达到理论峰值的 52%。这使得训练万亿参数模型所需的 3 个月时间在实践中变得可行。 - 全面的优化: 结合了计算(如融合核函数、数据布局优化)和通信优化(
scatter/gather),以及对微批次大小、激活重计算等关键超参数的深入分析,共同促成了卓越的扩展性能。 - 开源贡献: 开源了
Megatron-LM代码库,为社区提供了高效训练大规模LLMs的工具。
7.2. 局限性与未来工作
论文作者指出了以下局限性及可能的未来研究方向:
- 非自动化的并行策略搜索: 本文不自动探索并行策略的搜索空间(如
FlexFlow[22]、PipeDream[29] 等),而是基于实践经验提供启发式配置。未来可以研究自动化的并行策略优化。 - 严格的优化器语义: 本文坚持使用流水线冲洗来确保严格的优化器语义。未来工作可以探索放松语义的异步或有界陈旧性 (
bounded-staleness) 方法(如PipeDream-2BW[30]、PipeMare[45]),这些方法可能会进一步提高吞吐量,但可能以收敛速度或最终准确性为代价。 - 对称模型架构的假设: 在流水线并行中,本文主要考虑了具有重复
Transformer块的对称模型架构,使得层到流水线阶段的分配相对容易。对于更不对称的模型架构,层分配会更复杂。 - 与
ZeRO-3结合: 本文的ZeRO-3比较是其单独使用的情况。未来可以探索将ZeRO-3与模型并行结合,以潜在地改进其扩展行为。 - 硬件依赖: 尽管论文强调了优化的软硬件栈,但其高性能很大程度上依赖于顶级的
A100 GPU和高带宽InfiniBand互连。在通用或低配集群上的表现可能有所不同。
7.3. 个人启发与批判
7.3.1. 个人启发
- 分层并行设计的必要性: 这篇论文给我最大的启发是,在面对超大规模的计算挑战时,简单地应用单一并行策略是不够的。必须深入理解不同硬件层级(节点内
GPU互连、节点间网络)的通信特性,并设计分层的并行策略(如节点内张量并行,节点间流水线并行)来充分利用各层级的优势,规避其劣势。这种“软硬件协同设计”的理念是实现极致性能的关键。 - 工程优化与算法创新的结合: 论文不仅提出了新的并行调度,还深入到核函数优化、数据布局调整、通信模式改进等工程细节。这表明在深度学习系统领域,理论创新与精细的系统工程优化同样重要,两者缺一不可。很多时候,10% 的性能提升可能并非来自新的理论突破,而是来自于对现有系统的深度理解和巧妙优化。
- 微观与宏观权衡的艺术: 论文详细分析了微批次大小、激活重计算等超参数对吞吐量和内存占用的影响。这展示了在分布式训练中,存在许多相互竞争的因素,需要仔细权衡。没有一劳永逸的最佳配置,而是需要针对特定模型、硬件和批次大小进行细致的分析和调整。
- 超大规模训练的实践可行性: 论文通过在 3072 个
GPU上训练万亿参数模型,并给出实际的训练时间估算(约 3 个月),有力地证明了这种规模的LLM训练在工程上是可行的。这为未来更大规模模型的探索奠定了基础,并鼓舞了研究人员向更复杂的模型架构和应用迈进。
7.3.2. 批判
- 硬件依赖性强: 论文的卓越性能是在
NVIDIA Selene超级计算机这一顶级硬件栈上实现的,包括 80GBA100 GPU、高速NVLink节点内互连和 200GbpsInfiniBand胖树拓扑网络。这样的硬件配置并非所有研究机构或企业都能拥有。在更普遍的、可能存在网络瓶颈或GPU显存较小的集群上,PTD-P的性能优势可能会打折扣。论文虽然讨论了通信带宽的重要性,但并未在更受限的硬件环境下进行对比实验。 - 启发式而非自动化配置: 论文承认其并行策略配置是基于启发式方法,而非像
FlexFlow或PipeDream那样的自动化搜索。虽然启发式方法在实践中可能表现良好,但对于新模型架构或新的硬件配置,寻找最佳并行策略仍然需要大量的人工经验和试错。缺乏自动化的配置工具可能会增加其普适性应用的门槛。 - 严格优化器语义的取舍: 论文为了保持严格的优化器语义而引入了流水线冲洗,并为此设计了交错调度来减少气泡。然而,放松优化器语义的异步或有界陈旧性方法在一些工作中被证明可以提供更高的吞吐量(尽管可能影响收敛性或最终准确性)。论文并未深入探讨在何种情况下这种权衡是值得的,以及
PTD-P在放松语义下的表现如何。对于某些对最终准确性要求不那么极致但对训练速度要求极高的场景,放松语义可能更具吸引力。 - 模型架构的限制: 论文假设模型由重复的
Transformer块组成,这使得流水线阶段的划分相对均匀。对于更复杂的、异构的模型架构,如何高效地进行层间分区(流水线并行)仍是一个挑战。 - 比较基线的选择: 尽管与
ZeRO-3进行了比较,但ZeRO-3是在没有模型并行的情况下使用的。ZeRO系列本身也可以与模型并行结合。如果将PTD-P与结合了模型并行的ZeRO变体进行比较,结果可能会有所不同,这可能会提供更全面的性能图景。
相似论文推荐
基于向量语义检索推荐的相关论文。