Megatron-LM

一、通讯原语操作

NCCL 英伟达集合通信库，是一个专用于多个 GPU 乃至多个节点间通信的实现。它专为英伟达的计算卡和网络优化，能带来更低的延迟和更高的带宽。

Broadcast

Broadcast代表广播行为，执行Broadcast时，数据从主节点0广播至其他各个指定的节点（0~3）

Scatter

Scatter与Broadcast非常相似，都是一对多的通信方式，不同的是Broadcast的0号节点将相同的信息发送给所有的节点，而Scatter则是将数据的不同部分，按需发送给所有的节点。

Reduce

Reduce称为规约运算，是一系列简单运算操作的统称，细分可以包括：SUM、MIN、MAX、PROD、LOR等类型的规约操作。Reduce意为减少/精简，因为其操作在每个节点上获取一个输入元素数组，通过执行操作后，将得到精简的更少的元素。

AllReduce

Reduce是一系列简单运算操作的统称，All Reduce则是在所有的节点上都应用同样的Reduce操作

Gather

Gather操作将多个sender上的数据收集到单个节点上，Gather可以理解为反向的Scatter。

AllGather

收集所有数据到所有节点上。从最基础的角度来看，All Gather相当于一个Gather操作之后跟着一个Broadcast操作。

ReduceScatter

Reduce Scatter操作会将个节点的输入先进行求和，然后在第0维度按卡数切分，将数据分发到对应的卡上。

AllReduce操作，目标是高效得将不同机器中的数据整合（reduce）之后再把结果分发给各个机器。在深度学习应用中，数据往往是一个向量或者矩阵，通常用的整合则有Sum、Max、Min等。AllReduce具体实现的方法有很多种，最单纯的实现方式就是每个worker将自己的数据发给其他的所有worker，然而这种方式存在大量的浪费。一个略优的实现是利用主从式架构，将一个worker设为master，其余所有worker把数据发送给master之后，由master进行整合元算，完成之后再分发给其余worker。不过这种实现master往往会成为整个网络的瓶颈。AllReduce还有很多种不同的实现，多数实现都是基于某一些对数据或者运算环境的假设，来优化网络带宽的占用或者延迟。如Ring AllReduce：

第一阶段，将N个worker分布在一个环上，并且把每个worker的数据分成N份。
第二阶段，第k个worker会把第k份数据发给下一个worker，同时从前一个worker收到第k-1份数据。
第三阶段，worker会把收到的第k-1份数据和自己的第k-1份数据整合，再将整合的数据发送给下一个worker
此循环N次之后，每一个worker都会包含最终整合结果的一份。

假设每个worker的数据是一个长度为S的向量，那么个Ring AllReduce里，每个worker发送的数据量是O(S)，和worker的数量N无关。这样就避免了主从架构中master需要处理O(S*N)的数据量而成为网络瓶颈的问题。

二、Megatron-LM介绍

Megatron-LM 能够大规模训练大型 Transformer 语言模型。它为基于预训练 Transformer 的语言模型提供高效的张量、流水线和基于序列的模型并行性。

2. 1 数据并行

数据并行模式会在每个worker之上复制一份模型，这样每个worker都有一个完整模型的副本。输入数据集是分片的，一个训练的小批量数据将在多个worker之间分割；worker定期汇总它们的梯度，以确保所有worker看到一个一致的权重版本。对于无法放进单个worker的大型模型，人们可以在模型之中较小的分片上使用数据并行。

数据并行扩展通常效果很好，但有两个限制：

超过某一个点之后，每个GPU的batch size变得太小，这降低了GPU的利用率，增加了通信成本；
可使用的最大设备数就是batch size，限制了可用于训练的加速器数量。

2.2 张量并行(Tensor Parallelism,TP)

在张量并行 (TP) 中，每个 GPU 仅处理张量的一部分，并且仅当某些算子需要完整的张量时才触发聚合操作。( Megatron-LM 论文:Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism)

Transformer 类模型的主要模块为: 一个全连接层 nn.Linear，后面跟一个非线性激活层 GeLU。可以将其点积部分写为Y = GeLU(XA)，其中 X 和 Y 是输入和输出向量， A 是权重矩阵。如果以矩阵形式表示的话，很容易看出矩阵乘法可以如何在多个 GPU 之间拆分:

如果将权重矩阵A按列拆分到N个 GPU 上，然后并行执行矩阵乘法XA₁到XA_n，那么最终将得到 N 个输出向量 Y₁, Y₂, ⋯, Y_n ，它们可以独立输入 GeLU:

[Y₁, Y₂] = [GeLU(XA₁), GeLU(XA₂)]

因为 Y 矩阵是按列拆分的，因此随后的 GEMM(General Matrix Multiplication，通用矩阵乘)可以选择按行拆分方案，这样它就可以直接获取前面层的 GeLU 的输出，而无需任何额外的通信。使用该原理，可以更新任意深度的 MLP，只需在每个 拆列 - 拆行 序列之后同步 GPU。Megatron-LM 论文作者为此提供了一个不错的图示:

f是前向传递中的恒等运算符，后向传递中的all-reduce,g是前向传递中的all-reduce，后向传递中的恒等运算符。f和g是共轭的。

并行化多头注意力层甚至更简单，因为它们本来就是并行的，因为有多个独立的头！

对于自注意力块，我们利用多头注意力操作中固有的并行性，对与𝐾、𝑄和V关联的 GEMM 进行以列分区并行方式，使得与每个注意力头相对应的矩阵乘法在一个 GPU 上本地完成。这能够在 GPU 上分割每个注意力头参数和工作负载，并且不需要任何立即通信来完成自注意力。来自输出线性层（自注意力之后）的后续 GEMM 沿其行并行化，并直接获取并行注意力层的输出，而不需要 GPU 之间的通信。这种 MLP 和自注意力层的方法融合了两个 GEMM 组，消除了之间的同步点，并实现了更好的缩放。这能够在一个简单的transformer层中仅使用前向路径中的两个all-reduce和后向路径中的两个all-reduce来执行所有 GEMM。如下图：

需要特别考虑的是: 由于前向和后向传播中每层都有两个 all reduce，因此 TP 需要设备间有非常快速的互联。因此，除非你有一个非常快的网络，否则不建议跨多个节点进行 TP。实际上，如果节点有 4 个 GPU，则最高 TP 度设为 4 比较好。如果需要 TP 度为 8，则需要使用至少有 8 个 GPU 的节点。

2.3 流水线并行(Pipeline Parallelism,PP)

Naive PP是将模型各层分组分布在多个 GPU 上，并简单地将数据从 GPU 移动到 GPU，就好像它是一个大型复合 GPU 一样。该机制相对简单 - 将所需层用.to() 方法绑到相应设备，现在只要数据进出这些层，这些层就会将数据切换到与该层相同的设备，其余部分保持不变。

这其实就是垂直模型并行，因为在画大多数模型的拓扑图时，其实是垂直切分模型各层的。例如，如果下图显示一个 8 层模型:

===================  ===================
|  0 | 1 | 2 | 3  |  |  4 | 5 | 6 | 7  |
===================  ===================
        GPU0                 GPU1

将它垂直切成 2 部分，将层 0-3 放置在 GPU0 上，将层 4-7 放置在 GPU1 上。现在，当数据从第 0 层传到第 1 层、第 1 层传到第 2 层以及第 2 层传到第 3 层时，这就跟单 GPU 上的普通前向传播一样。但是当数据需要从第 3 层传到第 4 层时，它需要从 GPU0 传输到 GPU1，这会引入通信开销。如果参与的 GPU 位于同一计算节点 (例如同一台物理机器) 上，则传输非常快，但如果 GPU 位于不同的计算节点 (例如多台机器) 上，通信开销可能会大得多。然后第 4 到 5 到 6 到 7 层又像普通模型一样，当第 7 层完成时，通常需要将数据发送回标签所在的第 0 层 (或者将标签发送到最后一层)。现在可以计算损失，然后使用优化器来进行更新参数了。

问题：

该方案在任意给定时刻除了一个 GPU 之外的其他所有 GPU 都是空闲的。因此，如果使用 4 个 GPU，则几乎等同于将单个 GPU 的内存量翻两番，而其他资源 (如计算) 相当于没用上。另外还需要加上在设备之间复制数据的开销。所以 4 张使用Naive PP的 6GB 卡将能够容纳与 1 张 24GB 卡相同大小的模型，而后者训练得更快，因为它没有数据传输开销。但是，比如说，如果你有 40GB 卡，但需要跑 45GB 模型，你可以使用 4x 40GB 卡 (也就刚刚够用，因为还有梯度和优化器状态需要显存)。
共享嵌入可能需要在 GPU 之间来回复制。使用的流水线并行 (PP) 与上述Naive PP 几乎相同，但它解决了 GPU 闲置问题，方法是将传入的 batch 分块为 micros batch 并人工创建流水线，从而允许不同的 GPU 同时参与计算过程。

下图来自于GPipe 论文，图上把网络分成4块，每一块放在一个GPU上，不同的颜色表示不同的GPU，于是就有了 F₀、F₁、F₂、F₃ 这 4 个管级的前向路径，然后是 B3、B2、B1、B0 的逆序后向路径

其b部分表示朴素 PP 方案:为一般的模型并行的运算模式，在每个时间点只有一台设备在处理计算逻辑，完成计算后将结果发送给下一台设备
c部分是 PP 方法 PP 引入了一个新的超参数来调整，称为块(chunks)。它定义了通过同一管级按顺序发送多少数据块。例如，在图(c)的部分，可以看到 chunks = 4。GPU 0 在 chunk 0、1、2 和 3 (F0,0、F0,1、F0,2、F0,3) 上执行相同的前向路径，然后等待，等其他 GPU 完成工作后，GPU 0 会再次开始工作，为块 3、2、1 和 0 (B0,3、B0,2、B0,1、B0,0) 执行后向路径请注意，从概念上讲，这与梯度累积 (gradient accumulation steps，GAS) 的意思相同。PyTorch 叫它块，而 DeepSpeed 叫它 GAS。

因为块，PP 引入了 micro-batches (MBS) 的概念。DP 将全局 batch size 拆分为小 batch size，因此如果 DP 度为 4，则全局 batch size 1024 将拆分为 4 个小 batch size，每个小 batch size 为 256 (1024/4)。而如果块 (或 GAS) 的数量为 32，最终得到的 micro batch size 为 8 (256/32)。每个管级一次处理一个 micro batch 计算 DP + PP 设置的全局批量大小的公式为: mbs * chunks * dp_degree(8 * 32 * 4 = 1024)

这种调度机制被称为全前全后。

GPipe实验发现，通过将mini-batch进一步划分成更小的micro-batch，同时利用pipipline方案，每次处理一个micro-batch的数据，得到结果后，将该micro-batch的结果发送给下游设备，同时开始处理后一个 micro-batch的数据，通过这套方案减小设备中的Bubble(设备空闲的时间称为 Bubble)

使用 chunks=1 你最终得到的是Navie PP，这是非常低效的。而使用非常大的块数，你最终会得到很小的微批量大小，这很可能也不是很有效。因此，必须通过实验来找到能最有效地利用 GPU 的块数。

该图显示存在无法并行化的 “死” 时间气泡，因为最后一个 forward 阶段必须等待 backward 完成流水。那么，找到最佳的块数，从而使所有参与的 GPU 达到高的并发利用率，这一问题其实就转化为最小化气泡数了。

Megatron-LM 和 DeepSpeed 都有自己的 PP 协议实现。

其他一些可选方案:

交错一前一后

一前一后(one-forward-one-backward)

2.4 序列并行 (Sequence Parallelism,SP)

张量并行对训练期间花费最多时间的 Transformer 层部分进行并行化，因此计算效率很高。然而，它保留了层范数以及注意力和 MLP 块之后的dropout，因此，它们在张量并行组中复制。这些元素不需要大量计算，但需要大量激活内存。

在transformer层的非张量并行区域中，操作沿序列维度是独立的。这一特性使得能够沿着序列维度s划分这些区域。沿序列维度进行分区可减少激活所需的内存。这种额外级别的并行性在f之前和f̄之后引入了新的通信集合，它们将充当序列和张量并行区域之间的转换器。例如，在前向传递中，需要在上图中的运算符f之前进行额外的all-gather。这些额外的通信会带来开销，并会减慢训练速度。

为了避免这些额外的通信，将这些操作与f和f̄运算符结合起来，并引入新操作g和ḡ，如下所示图可以看出，g和ḡ是序列并行区域和张量并行区域之间的转换器。

g和ḡ是共轭的。g在前向传播中是all-gather，在后向传播中是reduce-scatter。ḡ在前向传递中是reduce-scatter，在后向传播中是all-gather。

在非并行形式中，transformer中的MLP 块所遵循的层范数可以表述为：

其中 X 是大小为 s × b × h 的层范数的输入，而 A 和 B 分别是大小为 h × 4h 和 4h × h 的线性层的权重矩阵。上述操作的组合张量和序列并行形式如下图所示。下标表示加速器之间的拆分，上标表示拆分的维度。例如，X₁^s 是第一个加速器沿 s 维度（序列维度）拆分的 X 部分，而 Z₂^h 是第二个加速器沿 h 维度（隐藏维度）拆分的 Z 部分。

层范数的输入沿着序列维度 X = [X₁^s, X₂^s] 并行。因此，层范数的输出也将沿着序列维度 Y = [Y₁^s, Y₂^s] 并行。具有 GeLU 非线性的线性层需要整个输入 Y，因此需要执行all-gather。这意味着 g 是前向传递中沿着序列维度的all-gather操作。通过沿其列（A₁^c 和 A₂^c）拆分 A 并沿其行（B₁^r 和 B₂^r）拆分 B，避免了通信并得到 W₁ 和 W₂。这两个张量不再平行，需要将其相加为 W = W₁ + W₂，然后才能将它们输入到 dropout 层。然而，dropout 需要其输入在序列维度 s 上是并行的。不是将这两个操作相加然后在序列维度上并行化，而是将其组合成一个 Reduce-Scatter 操作。因此，ḡ 可以是前向传递中的单个 Reduce-Scatter 操作。综合起来，得到：

如果对反向传递进行类似的分解，会发现 g 和 ḡ 是彼此共轭的。g 在正向传递中是all-gather，在反向传递中是reduce-scatter，而 ḡ 在正向传递中是reduce-scatter，在反向传递中是all-gather。对 layer-norm 进行类似的分解，然后是 Transformer 层的注意部分。

张量并行需要在一次正向和反向传递中进行四次all-reduces，而张量与序列并行需要在一次正向和反向传递中进行四次all-gather和四次reduce-scatter。乍一看，与张量并行相比，具有序列并行的张量似乎需要更多的通信。但是，Ringall-reduces由两个步骤组成：reduce-scatter，然后是all-gather。因此，张量并行和张量与序列并行使用的通信带宽相同。因此，序列并行不会引入任何通信开销。

从公式可知，序列并行和张量并行将反向传递所需的所有激活沿着并行维度划分，但第一个线性运算所需的张量 Y 除外。为了缓解这个问题，不存储反向传递的完整张量 Y。相反，只存储第 i 个张量并行等级上的 Y_i^s 部分，并在反向传递中执行额外的all-gather。为了消除这种额外的all-gather带来的延迟，将此通信与计算 Y 梯度所需的计算重叠，从而减少了开销。

2.5 Context Parallelism (CP)

megatron中的context并行(简称CP)与sequence并行(简称SP)不同点在于，SP只针对Layernorm和Dropout输出的activation在sequence维度上进行切分，CP则是对所有的input输入和所有的输出activation在sequence维度上进行切分，可以看成是增强版的SP。除了Attention模块以外，其他的模块(Layernorm、Dropout)由于没有多token的处理，在CP并行时都不用任何修改。

为什么Attention模块是个例外? 因为Attention计算过程中每个token的Q(query)要跟同一个sequence中其他token的K(key)和V(value)一起进行计算，存在计算上的依赖，所以通过CP并行后，在计算Attention前要通过allgather通信拿到所有token的KV向量，在反向计算时对应需要通过reduce_scatter分发gradient梯度。

为了减少显存占用，在前向时每个gpu只用保存一部分KV块，反向时通过allgather通信拿到所有的KV数据。KV的通信发生在相邻TP通信组相同位置的rank之间。allgather和reduce_scatter在ring拓扑架构实现时，底层会通过send和recv来进行实现。

以上图TP2-CP2的transformer网络为例，在Attention前的是CP的通信算子，其他都是TP的通信算子。AG表示allgather, RS表示reduce_scatter, AG/RS表示前向allgather反向reduce_scatter, RS/AG表示前向reduce_scatter反向allgather。

这里TP2对应为[GPU0, GPU1], [GPU2, GPU3], CP2对应为TP组相同位置的rank号，也就是[GPU0, GPU2], [GPU1, GPU3]。CP并行与Ring Attention类似，但是提供了新的OSS与FlashAttention版本，也去除了low-triangle causal masking的冗余计算。

LLM经常由于sequence长度过长导致显存OOM，这时之前的一种方式是通过重计算的方式保存中间的activation产出，全量重计算的劣势会带来30%的计算代价；另外一种方式是扩大TP(tensor parallel)的大小，扩大TP的劣势在于会对tensor切的更小，从而导致linear fc的计算时间变少，从而与通信很难进行计算的掩盖。

通过CP可以更好解决OOM的问题，每个GPU只用处理一部分的sequence, 同时减少CP倍的通信和计算，但保持TP不变，同时activation也会减少CP倍。total_gpu_count = TP × CP × PP × DP

2.6 可选Activation重计算(Selective Activation Recomputation)

a	注意力头数量	p	PP大小
b	微批量大小	s	序列长度
h	隐层维度	t	TP大小
L	transformer层数	v	词汇量

Attention block:包括自注意力，随后是线性投影和注意力dropout。线性投影以大小 2sbh 存储其输入激活，而注意力dropout需要大小为 sbh 的掩码。自注意力由几个元素组成：

Q、K 和V矩阵相乘：只需要存储大小为 2sbh 的共享输入。
QK^T 矩阵相乘：它需要存储 Q 和 K，总大小为 4sbh。
Softmax：反向传播需要大小为 2as²b 的 Softmax 输出。
Softmax dropout：只需要大小为 as²b 的掩码。
V：需要存储 dropout 输出 (2as²b) 和V (2sbh)，因此需要 2as²b + 2sbh 的存储空间。

注意力块需要 11sbh + 5as²b 字节的存储空间。

MLP: 两个线性层以大小 2sbh 和 8sbh 存储其输入。GeLU 非线性也需要其大小为 8sbh 的输入来进行反向传播。最后，dropout 以大小 sbh 存储其掩码。总的来说，MLP 块需要 19sbh 字节的存储空间。

Layer norm： 每个层范数以大小 2sbh 存储其输入，因此总的来说，我们需要 4sbh 的存储空间。

将注意力、MLP 和层规范所需的内存相加，存储 Transformer 网络单层激活所需的内存为：

张量并行不仅并行化了注意力和 MLP 块内的模型参数和优化器状态，而且还并行化了这些块内的激活。请注意，这些块的输入激活（例如，Q、K 和 V 矩阵乘法的输入或 h → 4h 线性层的输入）未并行化，并且只有每个块内的激活在张量并行组中划分。假设 t 路张量并行，存储激活所需的每层内存从公式减少到：

SP+TP

使用序列并行和张量并行，存储每个 Transformer 层激活所需的内存从公式减少到：

流水线并行性只是将 Transformer 的 L 层划分为 L/p 组层，其中 p 是流水线并行大小。但是，流水线并行性不会将激活所需的总内存均匀地除以 p。这是由于流水线并行调度引入了重叠以减少流水线”气泡”。

1F1B 流水线调度。流水线气泡最小化的调度将最大的内存压力放在流水线的第一阶段（流水线的第一阶段是指第一组 L/p 层，其中还包括输入嵌入）。为了保持流水线压力并避免额外的空闲时间，第一阶段必须存储 p 个微批次的激活（。每个阶段包含 L/p 层，因此无论管道并行大小 p 如何，第一阶段都必须存储 p × L/p = L 层的激活。因此，在第一阶段存储激活所需的总内存为：

对于其他流水线调度，所需的总内存会略有不同。例如，交错调度需要存储层的激活，其中 m 是交错阶段的数量。因此，如果使用交错调度，则总激活内存应按缩放。

Selective Activation Recomputation

与存储所有激活（公式）相比，对所有 Transformer 层进行checkpoint操作可显著减少训练模型所需的内存量。这种减少是以重新计算（额外的前向传递）为代价的，这可能会带来多达 30 − 40% 的计算时间开销。为了平衡内存节省和计算开销，理想的情况是只checkpoint足够的激活，以允许给定的模型并行配置在设备内存受限的情况下进行训练。序列并行提供的内存节省允许更多的配置进行训练而无需重新计算，但大型模型的最佳模型并行配置通常仍需要保存和重新计算一些激活。选择存储与重新计算的激活数量的简单方法是仅checkpoint部分 Transformer 层并存储其他层的所有激活。这种方法不能很好地扩展到大型模型；例如，在训练 MT-NLG 时，每个设备只有三层，这限制了平衡内存与计算的粒度。此外，并非所有激活都需要相同数量的操作来重新计算，因此更明智地选择要存储哪些激活以及要重新计算哪些激活是有益的。

建议不checkpoint并重新计算整个 Transformer 层，而是仅checkpoint并重新计算每个 Transformer 层中占用大量内存但重新计算计算成本不高的部分，或选择性激活重新计算。为此，方程中的项 5as/h 是由于在通过计算 Q、K 和 V 值的线性层增加网络宽度之后的注意操作；即 QK^T 矩阵乘法、softmax、softmax dropout 和对 V 的注意。这些操作通常具有较大的输入大小，因此激活也较大，但是每个输入元素的浮点运算 (FLOP) 数量非常低。Transformer层的其余部分占公式中的第 34 项。因此，对于 5as/h > 34 的大型模型，如果我们checkpoint并重新计算 Transformer 层的这一部分，存储的激活不到一半，并且重新计算未存储的激活的成本很小。

Configuration	Activations Memory Per Transformer Layer
no parallelism
tensor parallel (baseline)
tensor + sequence parallel
tensor parallel + selective activation recomputation
tensor parallel + sequence parallel + selective activation recomputation
full activation recomputation	sbh(2)

表：针对不同技术的每个transformer层的激活内存（以字节为单位）

2.7 DataSets

Megatron-LM 的另一个重要特性是高效的数据加载器。在首次训练启动前，每个数据集中的每个样本都被分成固定序列长度 (BLOOM 为 2048) 的样本，并创建索引以对每个样本进行编号。基于训练超参，我们会确定每个数据集所需要参与的 epoch 数，并基于此创建一个有序的样本索引列表，然后打乱它。举个例子，如果一个数据集中有 10 个样本并应参与 2 个 epoch 的训练，则系统首先按 [0, …, 9, 0, …, 9] 顺序排好样本索引，然后打乱该顺序为数据集创建最终的全局顺序。请注意，这意味着训练不会简单地遍历整个数据集然后重复，你有可能在看到另一个样本之前看到同一个样本两次，但在训练结束时模型将只看到每个样本两次。这有助于确保整个训练过程中的训练曲线平滑。这些索引，包括每个样本在原始数据集中的偏移量，被保存到一个文件中，以避免每次开始训练时都重新计算它们。最后，可以将其中几个数据集以不同的权重混合到训练最终使用的数据中。

2.8 Checkpoint重塑和互操作性

用于将可变张量和流水线并行大小的 Megatron-LM Checkpoint重塑为Transformers 分片Checkpoint的实用程序，因为它对大量工具具有强大的支持，例如Accelerate Big Model Inference、Megatron-DeepSpeed Inference 等。还支持将Transformers 分片Checkpoint转换为可变张量和流水线并行大小的 Megatron-LM Checkpoint，以进行大规模训练。

2.9 CUDA 融合核函数

GPU 主要做两件事。它可以将数据写到显存或从显存读数据，并对这些数据执行计算。当 GPU 忙于读写数据时， GPU 的计算单元就会空闲。如果想有效地利用 GPU，我们希望将空闲时间降至最低。

核函数是一组实现特定 PyTorch 操作的指令。例如，当你调用 torch.add 时，它会通过一个 PyTorch 调度器，它会根据输入张量及其他变量的取值来决定它应该运行哪些代码，最后运行它。CUDA 核函数使用 CUDA 来实现这些代码，因此只能在 NVIDIA GPU 上运行。

现在，当使用 GPU 计算 c = torch.add (a, b); e = torch.max ([c,d]) 时，一般情况下，PyTorch 将执行的操作是启动两个单独的核函数，一个执行 a 和 b 的加法，另一个执行取 c 和 d 两者的最大值。在这种情况下，GPU 从其显存中获取 a 和 b，执行加法运算，然后将结果写回显存。然后它获取 c 和 d 并执行 max 操作，然后再次将结果写回显存。

如果要融合这两个操作，即将它们放入一个 “融合核函数” 中，然后启动那个内核，不会将中间结果 c 写到显存中，而是将其保留在 GPU 寄存器中，并且仅需要获取 d 来完成最后的计算。这节省了大量开销并防止 GPU 空闲，因此整个操作会更加高效。

融合核函数就是这样。它们主要将多个离散的计算和进出显存的数据移动替换为有很少数据移动的融合计算。此外，一些融合核函数会对操作进行数学变换，以便可以更快地执行某些计算组合。

为了快速高效地训练，有必要使用 Megatron-LM 提供的几个自定义 CUDA 融合核函数。特别地，有一个 LayerNorm 的融合核函数以及用于融合缩放、掩码和 softmax 这些操作的各种组合的核函数。Bias Add 也通过 PyTorch 的 JIT 功能与 GeLU 融合。这些操作都是瓶颈在内存的，因此将它们融合在一起以达到最大化每次显存读取后的计算量非常重要。因此，例如，在执行瓶颈在内存的 GeLU 操作时同时执行 Bias Add，运行时间并不会增加。这些核函数都可以在 Megatron-LM repository 代码库中找到。

参考

LLM

#LLM #Megatron-LM

Megatron-LM

https://mztchaoqun.com.cn/posts/D50_Megatron-LM/

作者

mztchaoqun

发布于

2024年12月3日

许可协议

DeepSpeed 上一篇

LLM(十一)——Mamba 下一篇

Megatron-LM

一、 通讯原语操作

二、Megatron-LM介绍

2. 1 数据并行

2.2 张量并行(Tensor Parallelism,TP)

2.3 流水线并行(Pipeline Parallelism,PP)

2.4 序列并行 (Sequence Parallelism,SP)

2.5 Context Parallelism (CP)

2.6 可选Activation重计算(Selective Activation Recomputation)

2.7 DataSets

2.8 Checkpoint重塑和互操作性

2.9 CUDA 融合核函数

参考

一、通讯原语操作