Tensor-Parallelism with Partially Synchronized Activations

作者/机构: Itay Lamprecht†‡, Asaf Karnieli†, Yair Hanani†, Niv Giladi◦∗, Daniel Soudry‡ (†Intel, Israel ‡Department of Electrical and Computer Engineering - Technion, Haifa, Israel ◦AWS AI Labs)

A1 主要贡献

大规模语言模型（LLM）的训练和推理日益扩大，通常需要数百甚至数千台设备，因此对高效的分布式训练和推理技术的需求不断增长。张量并行（Tensor-parallelism）是许多著名LLM（如GPT-3、Llama和BLOOM）训练中的关键组成部分，它通过在多个设备上划分权重张量，显著减少了单个设备的内存占用，使得在内存有限的硬件上训练超大模型成为可能。然而，张量并行也带来了巨大的通信开销，主要源于同步激活或梯度张量的all-reduce操作。当张量并行的规模扩大时，每个设备的计算负载减少，而通信负载相对恒定，导致通信成本的相对增加，甚至可能主导整个训练过程。

本文研究的核心问题是如何在保持模型性能的同时，减少张量并行中的通信开销。研究发现，通过对现有实践进行微小调整，LLM可以在不完全同步激活张量的情况下进行训练。基于此，本文提出了“部分通道归约”（partial channel–reduce）操作，即只同步激活张量隐藏维度中的一部分通道，允许不同设备上的激活在通信后存在差异。

为利用这一思想，本文提出了面向张量并行的通信感知架构（Communication-Aware Architecture for Tensor-parallelism, CAAT-Net）。这是一种新的模型架构，通过使用部分通道归约来减少通信开销，同时保持模型参数量和总计算量不变。

本文的主要贡献如下：

  • 证明了部分同步的可行性：展示了在使用张量并行时，LLM可以在不完全同步激活张量的情况下进行训练并收敛。
  • 提出了CAAT-Net架构：提出了一种新颖的架构，通过仅同步部分激活张量，显著减少了张量并行训练中的网络流量。
  • 验证了方法的有效性：通过实验证明，CAAT-Net在各种设置下都能加速训练和推理，并且其准确性与完全同步的训练基本相当。例如，训练一个7B参数的Llama2模型，通信负载减少50%，而在几乎所有评估基准上都没有显著的性能下降。

A2 方法细节

3 CAAT-Net

为了减少LLM训练和推理中的通信带宽，我们用“部分通道归约”（partial channel–reduce）操作替换了旨在同步设备间激活的all-reduce操作，使得只有隐藏维度中的一部分通道在设备间同步。在部分通道归约中，发送到其他设备的通道被称为共享通道（shared channels），而每个设备独有的通道被称为私有通道（private channels）。共享通道在初始化时被选为前 $h \cdot p$ 个通道，其中 $h$ 是隐藏维度的大小，$p$ 是一个同步因子，控制隐藏维度的同步比例。当 $p = 1$ 时，部分通道归约就变成了完整的all-reduce操作（没有私有通道）。该操作在图2(c)中进行了可视化。本节我们提出了一种新颖的通信感知模型架构，它利用部分通道归约来加速LLM的训练和推理。我们方法的理论加速分析见附录C。


图1: CAAT-Net模型架构。左图：我们在单个设备上用一个两层全连接神经网络来举例说明我们的方法。中图：当使用两个设备进行张量并行时，输入激活X在两个设备上是相同的。每个设备使用自己的一组权重来乘以输入，产生中间激活，然后这些激活被归约为两个设备上相同的Z副本。右图：CAAT-Net为设备1和设备2分别接收不同的输入激活X1和X2。这些输入产生的中间激活在设备之间部分同步，从而分别在设备1和设备2上产生Z1和Z2。私有通道标记为P，共享通道标记为S。

3.1 架构

CAAT-Net改变了模型架构。尽管将通信原语从all-reduce更改为partial channel–reduce看起来只是一个微小的改动，但它实际上改变了模型架构。部分通道归约从根本上改变了模型的计算图和信息流。这使得CAAT-Net与那些保留或近似底层数学运算的通信优化方法（如流水线或压缩）有所区别。因此，在考虑部分通道归约时，多层感知器（MLP）和注意力（attention）层必须被重新定义。本节我们使用部分通道归约来定义MLP。注意力操作的定义在附录A.1中描述。

回顾标准MLP中的张量并行。在数学上描述使用部分通道归约的MLP之前，我们先描述在张量并行设置下使用常规all-reduce的MLP。为简单起见，我们考察张量并行维度为2的情况。对于一般情况，请参考附录A.2。给定连续的MLP权重矩阵A和B，输入张量X和激活函数 $\sigma$ ，MLP层的输出Z由下式给出：

$$Y = \sigma ( X A ) \quad ; \quad Z = Y B$$

其中Y是第一个MLP层的输出。Transformer MLP的基本架构如图1（左）所示。当使用张量并行时，我们沿着A的列来划分它，即：

$$A = [A_1, A_2]$$

所以Y的计算是：

$$Y_{1}=\sigma(XA_{1}) \quad ; \quad Y_{2}=\sigma(XA_{2}) \quad ; \quad Y=[Y_{1}, Y_{2}] .$$

我们沿着B的行来划分它，即：

$$\begin{aligned} B=\begin{bmatrix}B_1\\B_2\end{bmatrix} \end{aligned}$$

并且，在一次all-reduce之后，MLP的输出是：

$$Z=Y_{1}B_{1}+Y_{2}B_{2},$$

其中 $Y_1B_1$ 在第一个设备上计算，$Y_2B_2$ 在第二个设备上计算。带有张量并行的all-reduce的实现如图1（中）所示。

CAAT-Net中MLP的实现。我们首先注意到，使用部分通道归约时，每个设备的MLP输入将会不同（因为前一个注意力层在其输出中也将使用部分通道归约）。因此，每个设备的输入被表示为 $X_1$ 和 $X_2$，如图1（右）所示。因此，使用部分通道归约时，第一个线性层的输出是：

$$ Y_{1}=\sigma\left(X_{1} A_{1}\right) \quad ; \quad Y_{2}=\sigma\left(X_{2} A_{2}\right) \quad ; \quad Y=\left[Y_{1}, Y_{2}\right] . $$

接下来，我们需要同时沿着B的列和行来划分它：

$$\begin{aligned} B=\begin{bmatrix} B_{11} & B_{21} \\ B_{12} & B_{22} \end{bmatrix} \end{aligned}$$

然后，使用部分通道归约的MLP输出，表示为 $Z_1$ 和 $Z_2$，在每个设备上是不同的，计算方式如下：

$$Z_{1}=\left[Y_{1} B_{11}+Y_{2} B_{12}, Y_{1} B_{21}\right] \quad ; \quad Z_{2}=\left[Y_{1} B_{11}+Y_{2} B_{12}, Y_{2} B_{22}\right],$$

其中 $B_{11}$ 和 $B_{21}$ 在第一个设备上，$B_{12}$ 和 $B_{22}$ 在第二个设备上。$Y_1B_{11} + Y_2B_{12}$ 中的值位于共享通道，而 $Y_1B_{21}$ 和 $Y_2B_{22}$ 都位于私有通道，这在图1和图2中进行了可视化。

3.2 CAAT-Net 推理

推理过程中的通信减少。使用CAAT-Net训练的模型在推理时也使用部分通道归约。正如我们在5.2节中展示的，这在服务模型时能减少通信并带来加速。

处理与训练时不同的张量并行维度。如果模型服务的张量并行维度与训练时相同，那么服务模型是直接的。但当在推理时使用的张量并行维度与训练时不同时，情况就不同了。举个例子，我们考察一个模型在张量并行维度为2的情况下训练，而在单个设备上进行推理的场景。在训练中，一个transformer层的输出在每个设备上是不同的。在推理中，一个设备需要处理这两个副本。这可以通过使用‘逻辑设备’来实现。在这种情况下，一个物理设备可以模拟多个张量并行等级（rank），并为每个‘逻辑’张量并行设备顺序计算每一层的输出。此外，部分通道归约被替换为单个设备内部的局部求和。也可以通过微调将 $p$ 值增加到1，从而在推理期间恢复为完整的all-reduce操作。这样就得到了一个常规的transformer架构，推理可以在任何张量并行维度上运行。

4 部分同步：实现

实现部分同步所需的调整。为了在没有完全同步激活张量的情况下训练LLM，必须对当前的训练框架【索引1，Megatron-lm: Training multi-billion parameter language models using model parallelism，2020，https://arxiv.org/abs/1909.08053】进行两项主要调整。第一项是调整张量并行的反向传播实现，这对于避免训练中的前向-反向不匹配是必要的。第二项是在反向传播 的all-reduce中以32位（fp32）累积张量。我们凭经验发现，这对于避免训练中的数值问题是必要的。

反向传播中的前向-反向不匹配问题。为了用部分激活同步来训练LLM，我们需要进一步研究transformer模型的反向传播过程。在传统的LLM训练中，反向传播中的all-reduce操作是在MLP或注意力层产生的神经梯度上进行的，如图2a所示。通过显式计算通过网络传播的梯度，我们发现all-reduce操作可以应用于反向传播中的多个位置而不会改变反向传播算法。我们还发现，在部分同步激活的情况下，归约操作只能在一个地方应用。为了分析性地展示这一点，我们考察设备 $m$ 上的MLP层输入，记为 $X_m$，它是设备 $m$ 上前一个注意力层在all-reduce操作之前的输出（记为 $Z_m$）的函数，其中 $R$ 是残差连接，我们有：


图2: 部分同步与部分通道归约。(a) 当前训练框架中的原生transformer。操作f在前向传播中是all-reduce，在反向传播中是恒等操作。操作g在反向传播中是all-reduce，在前向传播中是恒等操作。(b) 使用部分同步时，h表示前向和反向传播中的归约操作，因为两者必须在相同位置进行。归一化函数参数的同步是必要的。(c) 在2个设备上使用参数p的部分通道归约。

$$X_{m}=\operatorname{Norm}\left(\sum_{m} Z_{m}+R\right),$$

其中 Norm 可以是任何归一化函数（通常是LayerNorm【索引20，Layer normalization，2016，https://arxiv.org/abs/1607.06450】或RMSNorm【索引21 ，Root mean square layer normalization，2019，https://arxiv.org/abs/1910.07467】）。我们考察损 失 $L$ 作为所有 $m$ 的 $X_m$ 的函数，以及 $X_m$ 作为 $Z_1$ 的函数，并计算损失函数关于 $Z_1$ 的导数。为简单起见，我们下面忽略残差连接（在附录B中，我们展示了包含残差的推导）。使用链式法则：

$$\frac{\partial \mathcal{L}\left(X_{1}\left(Z_{1}\right), \dots, X_{M}\left(Z_{1}\right)\right)}{\partial Z_{1}}=\sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}} \cdot \frac{\partial X_{m}\left(Z_{1}\right)}{\partial Z_{1}}\right)=\sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}} \cdot \frac{\partial X}{\partial Z_{1}}\right)$$

标准框架中的实现与部分同步的限制。在第二个等式中，我们回顾到 $\partial X_m / \partial Z_1$ 对所有 $m$ 都是相等的（现在记为 $\partial X / \partial Z_1$），因为在具有完整all-reduce操作的张量并行中，$X_1(Z_1) = X_2(Z_1) = ... = X_M(Z_1)$，因此它们关于 $Z_1$ 的导数也相等。这里，对 $m$ 的求和就是反向传播中的all-reduce操作。我们的关键洞见在于，由于矩阵乘法对加法具有分配性，即：

$$\sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}} \cdot \frac{\partial X}{\partial Z_{1}}\right)=\sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}}\right) \cdot \frac{\partial X}{\partial Z_{1}},$$

all-reduce操作可以在反向传播通过归一化函数之前使用（等式右侧），或者在归一化函数之后使用（等式左侧）。在标准的训练框架中，例如Megatron-LM，all-reduce操作是在反向传播通过归一化函数之前使用的（即图2a中的‘g’操作）。当使用部分同步进行训练时，$\partial X_m / \partial Z_m$ 对所有 $m$ 都相等的假设不再成立。因此，反向传播中的all-reduce操作必须在计算归一化函数导数之后进行（即图2b中的‘h’操作），这样它才能与前向传播中求和的位置（即图2a中的‘f’操作）相匹配，以获得正确的梯度。

归一化函数参数更新的同步。类似地，我们考察使用完整all-reduce操作的归一化函数参数更新，记为 $\beta$：

$$\frac{\partial \mathcal{L}\left(X_{1}(\beta), \ldots, X_{M}(\beta)\right)}{\partial \beta}=\sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}} \cdot \frac{\partial X_{m}}{\partial \beta}\right)$$

在梯度下降步骤中，每次迭代后，更新权重之前，更新量 $\partial L / \partial \beta$ 需要通过all-reduce操作进行同步。在流行的训练框架中，再次利用了 $\partial X_m / \partial \beta$ 对所有 $m$ 设备都相同这一事实。如果all-reduce发生在反向传播通过归一化函数之前，则无需在每一步之后进行梯度同步：

$$\sum_{m}\left(\frac{\partial\mathcal{L}(X_{m})}{\partial X_{m}}\right)\cdot\frac{\partial X}{\partial\beta},$$

其中，同样地，$\partial X_m / \partial \beta$ 对所有 $m$ 都是相等的（记为 $\partial X / \partial \beta$）。

总结两种等价方法与部分同步的约束。总结来说，当使用完整all-reduce时，以下两种方式是等价的：
* 在神经梯度反向传播通过归一化函数之前对其进行归约。
* 在梯度反向传播通过归一化函数之后对其进行归约，并在权重更新前对归一化函数参数进行归约。
然而，当使用部分同步（例如部分通道归约）进行训练时，只有第二种方法是可行的。我们修改后的反向传播实现在图2b中有详细说明。需要注意的是，每层归一化函数参数的数量通常是 $\Theta(h)$，其中 $h$ 是隐藏层大小，并且归一化函数参数的同步在每一步之后都会发生。因此，归一化函数参数的额外通信量与张量并行的all-reduce操作相比是微不足道的。

通过32位梯度累积实现数值稳定性。在前一节中，我们提出了两种不同的反向传播实现方式。如果我们在前向传播中完全同步激活，这两种实现方式在数学上是等价的。尽管存在这种等价性，我们在实验中发现，即使在激活完全同步的情况下，使用我们的替代实现方式在训练收敛性上仍可能存在很大差距。这个差距源于实现方式之间的数值差异，通过将归约后的梯度以32位精度累积而不是16位精度，可以完全消除这个差距。虽然LLM通常以16位精度进行训练，但在某些操作中32位精度至关重要，例如矩阵乘法操作中的累积或归一化层。虽然张量并行的梯度all-reduce通常不属于这些操作，但我们发现，当移动反向传播中的all-reduce位置时，32位累积至关重要。需要注意的是，虽然累积需要以32位精度进行，但通信本身可以以16位精度进行，值在通信后、累积前简单地向上转型即可。附录E.1中提供了比较fp32和bfloat16中all-reduce累积的训练实验损失曲线。

私有通道的缩放以匹配方差。部分通道归约导致共享通道和私有通道的统计特性存在差异，这会影响网络中的信号传播。考虑使用2个张量并行设备的部分通道归约情况。假设归约前的MLP输出（$Y_1B_{11}$, $Y_2B_{12}$, $Y_1B_{21}$）中每个元素的均值为零，方差为 $\sigma_A^2$，并且在设备间是独立的。共享通道的激活方差为：

$$\operatorname{Var}(Y_{1}B_{11} + Y_{2}B_{12}) = \operatorname{Var}(Y_{1}B_{11}) + \operatorname{Var}(Y_{2}B_{12}) = 2\sigma_{A}^{2},$$

私有通道的激活方差为：

$$\text{Var}(Y_1 B_{21}) = \sigma_A^2$$

方差计算与校正因子。这种方差不匹配会导致激活和梯度计算中跨通道的信号不均匀。为了纠正这种不匹配，我们将私有通道中的激活乘以一个校正因子 $\sqrt{2}$（如果使用 $r$ 的张量并行维度，则为 $\sqrt{r}$）。这样，在部分通道归约之后，所有通道的激活具有相同的方差。私有通道缩放的消融研究可在附录E.3中找到。虽然我们发现在不同的 $p$ 值下验证损失有轻微改善，但私有通道缩放并非获得足够准确性结果的强制要求。

A4 实验环境

  • 硬件配置: 实验在Intel Gaudi3 HPU加速器上进行，该加速器具有128GB板载内存，每个设备拥有525 GB/s的节点内连接和75 GB/s的节点间连接。部分速度测试还在8x NVIDIA H100-80GB-HBM3和8x NVIDIA A100-SXM4-80GB（均带NVLink）上进行。

  • 模型架构:

    • Llama2-7B: 一个Llama2-7B【索引8，Llama 2: Open foundation and fine-tuned chat models，2023，https://arxiv.org/abs/2307.09288】的变体，用于大规模训练和性能测试 。
    • TinyLlama-1.1B: 使用TinyLlama模型【索引9，Tinyllama: An open-source small language model，2024，https://arxiv.org/abs/2401.02385】进行消融研究 。
    • 130M模型: 一个基于LLaMA架构的自定义小型模型，包含16个注意力头，隐藏维度大小为768，共12个transformer层。
    • GPT3-XL: 用于验证方法对不同架构的鲁棒性。
    • Llama-34B/70B: 用于额外的推理速度测试。

  • 数据集:

    • 训练: 使用RedPajama数据集。
    • 评估: 使用语言模型评估框架（Language Model Evaluation Harness）【索引22，A framework for few-shot language model evaluation，2023，https://zenodo.org/records/10256836】，涵 盖HellaSwag, WikiText103, LAMBADA, WinoGrande, PIQA, OpenBookQA和BoolQ等常识推理任务。

  • 软件配置:

    • 训练框架: Intel的Megatron-LM分支和Optimum-Habana。
    • 推理框架: 使用gpt-fast库在NVIDIA硬件上进行实验。
    • 优化器: AdamW。
    • 分词器: GPTSentencePiece tokenizer。

A4 实验结果

5.1 大规模训练

  • Llama2-7B实验：研究者训练了一个Llama2-7B的变体，使用p=0.5的部分通道归约和8的张量并行维度，在RedPajama数据集上训练了160B个token。如表1所示，与基线模型相比，CAAT-Net在多个零样本评测任务上没有出现显著的性能下降，在Bool-Q上略有统计显著性差异，在WikiText和验证损失上甚至略有改善。这表明在减少50%通信负载的情况下，模型精度得以保持。

    表1: CAAT-Net与基线模型：预训练后零样本准确率。7B参数模型，p = 0.5，张量并行度为8。

  • 可扩展性与鲁棒性研究：为了进一步验证方法，研究者还训练了130M和1.1B（TinyLlama）参数的模型。

    • 同步因子p的影响（图3左）：对于130M和1.1B模型，当p值较大（>0.75）时，性能与基线相似。在中间值（0.25-0.75）时，130M模型验证损失略有改善，而1.1B模型与基线持平或略差。当p值过小时，两个模型性能都下降。实验中的7B模型在p=0.5时也表现出轻微改善，表明该方法具有扩展潜力。
    • 张量并行维度的影响（图3右）：在130M模型上，对于所有测试的张量并行维度（2, 4, 8, 16），随着p从1开始减小，验证损失先是轻微改善，然后恶化。张量并行维度越大，性能恶化的拐点对应的p值也越高。对于所有维度，p=0.5都未导致性能下降或有所改善。
    • 架构鲁棒性：研究者还训练了GPT3-XL模型（p=0.5, TP=8），同样观察到CAAT-Net能够保持准确性，证明该方法不局限于Llama架构。


    图3: 多种场景下的训练准确率。左图：130M和1.1B模型在不同p值下的验证损失，以及7B模型在p=0.5时的验证损失，均归一化至p=1时的损失。右图：130M模型在不同p值和张量并行维度（TP）下的验证损失。

5.2 加速效果

  • 推理加速：

    • Gaudi3硬件（图4下）：在Llama2-7B模型上测量了首个令牌生成时间（TTFT）。当张量并行维度为16且p=0.25时，对于批大小为32的情况，最大可获得26%的加速。在与大规模训练设置相同（TP=8, p=0.5）且不降低精度的条件下，最大可获得14%的加速。
    • NVIDIA硬件（图4右上）：在H100和A100硬件上，使用p=0.25和TP=8，推理TTFT最高可加速13%。
    • 硬件约束：实验还发现，当选择的p值使得通信量是通信缓冲区大小的倍数时，加速效果更显著。

  • 训练加速（图4左上）：

    • 测量了Llama2-7B的训练吞吐量（每秒token数，TPS）。
    • 在TP=16的设置下，当p=0.25时吞吐量可提升14%，当p=0.5时可提升9%。
    • 注意：此处的训练加速实验由于技术限制，反向传播累积使用的是16位精度，因此结果应被视为性能预测。预计使用32位累积（通信仍为16位）不会显著影响报告的加速比。


    图4: Llama 7B模型的训练和推理加速。左上：不同p值和张量并行维度下的训练加速。右上：不同硬件上，使用批大小128的推理首个令牌生成时间（TTFT）加速。左下：使用张量并行8时，推理TTFT加速随批大小变化的函数，对应不同p值。右下：使用张量并行16时，推理TTFT加速随批大小变化的函数，对应不同p值。

5.3 与压缩方法的比较

  • 对比方法：将CAAT-Net与两种压缩激活通信的方法进行比较：随机掩码（相当于在通信前应用dropout）和Top-K掩码（只发送每个token激活中最大的K个值）。
  • 结果（图5）：实验在130M Llama模型上进行。结果显示，即使是轻微的压缩，随机掩码和Top-K掩码都会导致验证损失严重恶化。相比之下，CAAT-Net在p值较低时仍能保持较低的验证损失。
  • 效率分析：此外，如附录D所述，随机掩码和Top-K掩码在减少通信量方面也不如CAAT-Net高效。


图5: 与压缩方法的比较。130M模型使用CAAT-Net、Top-k掩码和随机掩码时的验证损失与p值的关系。

A5 结论

本文证明了通过对当前训练框架进行微小改动，张量并行通信中的激活无需完全同步即可实现训练收敛。基于此，本文提出了CAAT-Net，一个显著减少LLM训练中张量并行相关网络流量的架构，且对模型准确性的影响微乎其微甚至没有影响。通过在小型网络上的实验，我们探索了张量并行维度和同步因子的影响。最后，我们展示了该方法如何加速推理和训练工作负载。

A6 附录

A 使用部分通道归约的MLP和注意力的数学描述

A.1 注意力

标准all-reduce下的注意力层。在本节中，我们描述一个使用部分通道归约的注意力层。与3.1节类似，我们考察张量并行维度为2的情况。此外，为简单起见，我们假设模型仅包含2个注意力头。下面，我们描述使用完整all-reduce的注意力层的输出。第一个和第二个头的注意力头输出分别为 $H_1$ 和 $H_2$：

$$H_{1}=\operatorname{Attention}\left(X W_{1}^{K}, X W_{1}^{Q}, X W_{1}^{V}\right) \quad ; \quad H_{2}=\operatorname{Attention}\left(X W_{2}^{K}, X W_{2}^{Q}, X W_{2}^{V}\right)$$

在张量并行中，注意力层的输出投影 $O$ 在设备间分片。all-reduce后的输出 $Z$ 为：

$$\begin{aligned} Y = [H_1, H_2] \quad ; \quad Z = YO, \quad ; \quad O = \begin{bmatrix} O_1 \\ O_2 \end{bmatrix} \quad ; \quad Z = H_1 O_1 + H_2 O_2 , \end{aligned}$$

其中 $H_1O_1$ 在第一个设备上计算，$H_2O_2$ 在第二个设备上计算。

部分通道归约下的注意力层。与MLP情况一样，当引入部分通道归约时，注意力的输入在每个设备上不一定相同。每个设备上的输入表示为 $X_1$ 和 $X_2$：

$$\begin{aligned} \begin{aligned} H_{1} & =\text { Attention }\left(X_{1} W_{1}^{K}, X_{1} W_{1}^{Q}, X_{1} W_{1}^{V}\right) \\ H_{2} & =\text { Attention }\left(X_{2} W_{2}^{K}, X_{2} W_{2}^{Q}, X_{2} W_{2}^{V}\right) \end{aligned} \end{aligned}$$

使用部分通道归约的注意力层输出在每个设备上是不同的，表示为 $Z_1$ 和 $Z_2$：

$$\begin{aligned} Y = [H_1, H_2] . \quad ; \quad O = \begin{bmatrix} O_{11} & O_{21} \\ O_{12} & O_{22} \end{bmatrix} \end{aligned}$$

$$Z_1 = [H_1 O_{11} + H_2 O_{12}, H_1 O_{21}] \quad ; \quad Z_2 = [H_1 O_{11} + H_2 O_{12}, H_2 O_{22}] .$$

A.2 使用多个设备的部分通道归约

多设备下的通用公式。在本节中，我们放宽在3.1节和附录A.1中作出的假设，以允许超过2个张量并行设备。

MLP。在transformer中，使用完整all-reduce操作的MLP操作可以描述如下。给定权重矩阵 $A \in R^{h \times f}$ 和 $B \in R^{f \times h}$，其中 $f$ 是前馈网络（FFN）的隐藏维度，以及输入激活 $X \in R^{b \times t \times h}$，MLP可以用索引表示法写成：

$$Z_{il} = \sum_{k} \sigma \left( \sum_{j} X_{ij} A_{jk} \right) B_{kl},$$

其中 $t$ 是序列长度，$b$ 是批大小，$\sigma$ 是非线性函数。为了引入部分通道归约操作，我们首先将公式20重写为包含张量并行的表示法：

$$\begin{aligned} \begin{aligned} \tilde{Z}_{i l m} & =\sum_{k} \sigma\left(\sum_{j} X_{i j} \tilde{A}_{j k m}\right) \tilde{B}_{k l m} \\ Z_{i l} & =\sum_{m} \tilde{Z}_{i l m} \end{aligned} \end{aligned}$$

其中 $\tilde{A}$ 是以列并行表示法写的 $A$，最后一个维度对应张量并行等级，类似地 $\tilde{B}$ 是以行并行表示法写的 $B$。对 $m$ 的求和是all-reduce操作。最后，我们将带有部分通道归约操作的MLP写成如下形式：

$$\tilde{Z}_{i l m}=\sum_{k} \sigma\left(\sum_{j} X_{i j m} \tilde{A}_{j k m}\right) \tilde{B}_{k l m}$$

$$\begin{aligned} Z_{i l m}=\begin{cases}\sum_{m} \tilde{Z}_{i l m} & \text { if } l<\text { floor }(h p) \\ \tilde{Z}_{i l m} & \text { if } l \geq \text { floor }(h p)\end{cases} \end{aligned}$$

这里，对 $m$ 的求和是部分通道归约操作。公式21和公式22之间的主要区别是：
* 对 $\tilde{Z}$ 的求和仅限于隐藏维度的前 $p \cdot h$ 个通道。
* 当使用部分通道归约时，每个设备的MLP输入是不同的，因为前一个注意力层在其输出中也使用了部分通道归约。因此，MLP的输入 $X$ 有一个额外的索引 $m$，对应张量并行设备等级。类似地，$Z$ 也有一个额外的索引 $m$，因为在部分通道归约操作之后，每个设备在私有通道中包含不同的参数。

注意力。我们展示带有部分通道归约的注意力操作。为简单起见，我们假设张量并行度等于头的数量，尽管这个假设可以很容易地放宽。使用完整all-reduce操作的注意力操作是：

$$\begin{aligned} \begin{aligned} \mathrm{H}_{i j m} & =\operatorname{Attention}\left(X W_m^K, X W_m^Q, X W_m^V\right)_{i j} \\ \tilde{Z}_{i l m} & =\sum_j \mathrm{H}_{i j m} \tilde{O}_{j l m} \\ Z_{i l} & =\sum_m \tilde{Z}_{i l m}, \end{aligned} \end{aligned}$$

其中 $H_{ijm}$ 是第 $m$ 个注意力头输出中的一个元素，$\tilde{O}_{jlm}$ 是注意力输出投影中的一个元素，在将注意力输出投影 $O$ 重塑为行并行表示法之后。因此，带有部分通道归约操作的注意力操作的输出是：

$$\begin{aligned} \begin{gathered} \mathrm{H}_{i j m}=\operatorname{Attention}\left(X_m W_m^K, X_m W_m^Q, X_m W_m^V\right)_{i j} \\ \tilde{Z}_{i l m}=\sum_j \mathrm{H}_{i j m} O_{j l m} \\ Z_{i l m}= \begin{cases}\sum_m \tilde{Z}_{i l m} & \text { if } l<\text { floor }(h p) \\ \tilde{Z}_{i l m} & \text { if } l \geq \text { floor }(h p)\end{cases} . \end{gathered} \end{aligned}$$

公式23和公式24之间的主要区别是：
* 对 $\tilde{Z}$ 的求和仅限于隐藏维度的前 $p \cdot h$ 个通道。
* 除了第一个注意力层，当使用部分通道归约时，每个设备的注意力输入是不同的，因为前一个MLP层在其输出中也使用了部分通道归约。因此，注意力的输入 $X$ 有一个额外的索引 $m$，对应张量并行等级。类似地，$Z$ 也有一个额外的索引，因为在部分通道归约操作之后，每个设备包含不同的参数。

B 考虑残差连接分析前向-反向不匹配

包含残差连接的推导。对于一个带有all-reduce的transformer，我们考察设备 $m$ 上的MLP层输入，记为 $X_m$，它是设备 $m$ 上前一个注意力层在all-reduce操作之前的输出，记为 $Z_m$，$R$ 是来自注意力块之前的残差连接，我们有：

$$X_{m}=\operatorname{norm}\left(\sum_{m} Z_{m}+R\right).$$

$R$ 是来自注意力块之前的残差连接。我们考察损失作为所有 $m$ 的 $X_m$ 的函数，以及第 $m$ 个设备上跳过下一层的残差连接 $R_m$ 的函数。我们还考察 $X_m$ 和 $R_m$ 作为 $Z_1$ 的函数，并使用链式法则计算损失函数关于 $Z_1$ 的导数：

$$\begin{aligned} \begin{aligned} \frac{\partial \mathcal{L}\left(X_{1}\left(Z_{1}\right), \ldots, X_{M}\left(Z_{1}\right), R_{1}\left(Z_{1}\right), \ldots, R_{M}\left(Z_{1}\right)\right)}{\partial Z_{1}} & = \\ \sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}} \cdot \frac{\partial X_{m}\left(Z_{1}\right)}{\partial Z_{1}}+\frac{\partial \mathcal{L}\left(R_{m}\right)}{\partial R_{m}} \cdot \frac{\partial R_{m}\left(Z_{1}\right)}{\partial Z_{1}}\right) & = \\ \sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}} \cdot \frac{\partial X_{m}\left(Z_{1}\right)}{\partial Z_{1}}+\frac{\partial \mathcal{L}\left(R_{m}\right)}{\partial R_{m}}\right) & = \\ \sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}} \cdot \frac{\partial X}{\partial Z_{1}}+\frac{\partial \mathcal{L}\left(R_{m}\right)}{\partial R_{m}}\right) \end{aligned} \end{aligned}$$

其中在第二个等式中我们使用了以下事实：

$$\frac{\partial R_{m}\left(Z_{1}\right)}{\partial Z_{1}}=\frac{\partial}{\partial Z_{1}} \sum_{m} Z_{m}=1,$$

在最后一个等式中，我们认为所有 $m$ 的 $\partial X_m / \partial Z_1$ 都相等（现在记为 $\partial X / \partial Z_1$），因为在具有完整all-reduce操作的张量并行中，$X_1(Z_1) = X_2(Z_1) = ... = X_M(Z_1)$，因此它们关于 $Z_1$ 的导数也相等。这里，对 $m$ 的求和是反向传播中的all-reduce操作。我们的关键洞见在于，由于矩阵乘法的分配性：

$$\sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}} \cdot \frac{\partial X}{\partial Z_{1}}+\frac{\partial \mathcal{L}\left(R_{m}\right)}{\partial R_{m}}\right)=\sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}}\right) \cdot \frac{\partial X}{\partial Z_{1}}+\sum_{m} \frac{\partial \mathcal{L}\left(R_{m}\right)}{\partial R_{m}} .$$

结论与部分同步的约束。换句话说，all-reduce操作可以在反向传播通过归一化函数之前发生（等式右侧），也可以在归一化函数之后发生（等式左侧）。在等式右侧，残差梯度的求和发生在下一层的all-reduce中，即在残差被累加回激活张量之后的层。在标准的训练框架中，例如Megatron-LM，all-reduce操作发生在反向传播通过归一化函数之前。当使用CAAT-Net进行训练时，$\partial X_m / \partial Z_m$ 对所有 $m$ 都相等的假设不再成立。因此，反向传播中的all-reduce操作必须在计算归一化函数导数之后进行。

归一化函数参数更新。类似地，我们考察使用完整all-reduce操作的归一化函数参数更新，记为 $\beta$。因为残差在应用归一化函数之前从激活中分离，所以归一化函数参数更新不是残差状态的函数：

$$\frac{\partial \mathcal{L}\left(X_{1}(\beta), \ldots, X_{M}(\beta)\right)}{\partial \beta}=\sum_{m}\left(\frac{\partial \mathcal{L}\left(X_{m}\right)}{\partial X_{m}} \cdot \frac{\partial X_{m}}{\partial \beta}\right)$$

在梯度下降步骤中，这意味着在每次迭代之后，更新权重之前，更新量 $\partial \beta$ 需要自己通过all-reduce操作进行同步。在流行的训练框架中，再次利用了 $\partial X_m / \partial \beta$ 对所有 $m$ 都相等的事实。当all-reduce发生在反向传播通过归一化函数之前时，就不需要在每一步之后进行梯度同步了：

$$\sum_{m}\left(\frac{\partial\mathcal{L}(X_{m})}{\partial X_{m}}\right)\cdot\frac{\partial X}{\partial\beta}.$$

再一次，我们认为所有 $m$ 的 $\partial X_m$ 都相等（现在记为 $\partial X / \partial \beta$）。

C 加速分析

理论加速模型。我们提出了一个简单的模型来进一步理解计算和通信时间之间的关系，并分析我们方法的潜在加速效果。我们作出以下假设：
* 计算仅建模为GEMM操作——特别是在MLP和注意力中。通信仅建模为张量并行的all-reduce。我们忽略所有其他计算（嵌入、语言模型头等）和通信（流水线并行、数据并行等）。
* 我们假设计算和通信容量之间存在一个恒定的比率 $C$，单位为 FLOP/s / (GB/s)。
* 我们假设当我们调整张量并行维度时，通信时间大致不变。这个假设在Gaudi处理器中使用多节点连接时是成立的，因为在all-reduce操作期间，节点间和节点内通信会重叠。此外，我们假设计算时间与张量并行维度成反比。这个假设只是近似成立，因为计算时间取决于许多因素，例如张量形状。

计算与通信负载公式。我们考察一个带有简单MLP的transformer的前向传播，其中MLP中的第一个权重矩阵维度为 $h \times 4h$，第二个权重矩阵维度为 $4h \times h$。这种类型的MLP在像GPT和BERT这样的LLM中很常见。我们考虑批大小为1，因为批大小不影响本节的结果。当一个 $m \times p$ 矩阵乘以一个 $p \times n$ 矩阵时，总操作数（乘法和加法）是 $m \cdot p \cdot (2n - 1)$，我们近似为 $2m \cdot p \cdot n$。基于此，一个序列长度为 $s$，有 $n$ 个头和隐藏大小为 $h$ 的transformer层中的总计算操作数是：

$$\begin{aligned} \begin{aligned} \text { MLP ops } & =\underbrace{8 s h^2}_{\mathrm{h} \text { to } 4 \mathrm{~h}}+\underbrace{8 s h^2}_{4 \mathrm{~h} \text { to } \mathrm{h}}=16 s h^2 \\ \text { Attention ops } & =\underbrace{6 s h^2}_{\text {qkv proj }}+\underbrace{4 s^2 h}_{\text {attention }}+\underbrace{2 s h^2}_{\text {out proj }}= \\ & 4 s^2 h+8 s h^2 . \end{aligned} \end{aligned}$$

为了计算总通信负载，我们回想起在每个完整的all-reduce操作中，每个设备发送和接收 $s \cdot h$ 个激活，这在每个transformer块中发生两次。因此，每个设备的总计算操作数 $G$ 和通信负载 $P$ 作为 $p$ 的函数是：

$$G = \frac{24sh^{2} + 4s^{2}h}{r}$$ $$P(p) = 2shp,$$

其中 $r$ 是张量并行维度，$p$ 是CAAT-Net同步因子。因此，前向传播中单个transformer层的计算和通信总时间为：

$$T(p)=G+P(p)$$

加速比与最优p值。我们发现我们的方法带来的加速是：

$$\frac{T(1)-T(p)}{T(1)}=(1-p) \frac{1}{1+\frac{12 h+2 s}{C r}} .$$

我们的方法对于大的 $C$ 和 $r$ 更为显著，但对于非常大的 $h$ 和 $s$ 则效果减弱。这是因为计算在 $h$ 和 $s$ 上是二次的，而通信是线性的。如果通信被所有计算隐藏，那么存在一个最大值 $p$，使得通信和计算时间相等，此时减少 $p$ 以提高加速的动机就减少了。如果 $G \ge P(1)$，那么最优的 $p$ 值是 $p^ = 1$。否则我们可以设置 $G = P(p)$。综合考虑这两种情况，我们发现 $p^$ 的值是：

$$p^* = \min ( \frac{12h + 2s}{Cr}, 1 ).$$

D Top-K和随机掩码的通信减少

掩码方法通信减少效率分析。分析这些实验中的通信量减少情况，我们注意到只有前向传播中的通信被压缩了。这是因为我们使用的是普通的张量并行，而不是我们在第4节中详述的替代实现，所以前向和反向传播中的通信位置不同。因此，与CAAT-Net不同，前向传播中的通信压缩并不意味着反向传播中的通信压缩。此外，all-reduce分两步进行——reduce-scatter和all-gather。在通信前应用掩码相当于压缩了reduce-scatter操作，这只占all-reduce总通信量的一半。因此，使用参数为 $p$ 的Top-K或随机掩码压缩通信，导致张量并行通信量减少了 $(100 \cdot (1-p) / 4)\%$。在参数为 $p$ 的CAAT-Net中，通信在前向和反向传播的reduce-scatter和all-gather中都得到了压缩。CAAT-Net中的总张量并行通信减少了 $(100 \cdot (1-p))\%$。

E 附加结果

E.1 反向传播中All-Reduce位精度的影响

16位累积导致训练发散。本节我们展示了在反向传播all-reduce中使用16位累积会导致损失发散。为了证明这一点，我们训练了TinyLlama模型的变体，如5.1节所述。第一个变体使用Intel的Megatron-LM分支仓库进行训练，没有任何额外改动。第二个变体在应用了第4节中详述的改动后进行训练，反向传播all-reduce累积使用bfloat16精度。第三个变体在应用了第4节中的所有改动（包括在反向传播中使用32位累积）后进行训练。图6中的结果显示，在使用16位累积训练15K步后，出现了严重的性能下降。尽管两种实现在数学上是等价的，但数值差异导致了不稳定性。


图6: 反向传播中All-Reduce累积精度对TinyLlama训练损失的影响。

E.2 不同p值下TinyLlama和GPT-XL的准确率评估

详细准确率数据。我们报告了使用不同p值的TinyLlama的完整零样本准确率评估。评估指标与5.1节中描述的相同。总体而言，p=0.75取得了最好的准确率，在8个指标中有5个得分最高。结果见表3。此外，我们报告了p=0.5时GPT-XL的完整零样本准确率评估。结果见表2。

表2: GPT-XL上的零样本准确率结果

表3: 不同p值下TinyLlama上的零样本准确率结果

E.3 私有通道缩放消融研究

私有通道缩放的有效性。本节我们展示了私有通道缩放的消融研究。我们训练了5.1节中介绍的130M模型，训练了7.8B个token，张量并行维度为8，p=0.5。我们记录了训练结束时有无私有通道缩放的验证损失，发现在大多数p值下，使用私有通道缩放有轻微的改善。结果呈现在图7中。


图7: 部分通道缩放。在130M模型上，有无私有通道缩放的验证损失与p值的关系。

E.4 额外的加速结果

更大模型的推理加速。在本节中，我们报告了34B和70B Llama模型的推理加速。这些模型使用了分组查询注意力（GQA），数量为8，因此我们只能在张量并行维度为8（或2的更低倍数）上评估我们的加速效果。绝对测量值见表4和表5。

表4: Llama 70B推理首个令牌生成时间（TTFT）的绝对测量值，张量并行8。以毫秒为单位，对应不同批处理大小和p值。

表5: Llama 34B推理首个令牌生成时间（TTFT）的绝对测量值，张量并行8。以毫秒为单位，对应不同批处理大小和p值。

E.5 绝对性能指标

原始性能数据。在论文的5.2节中，我们报告了Llama-7b模型在训练和推理中相对于基线的加速。这里我们报告了计算加速所依据的绝对测量值。

表6: 训练吞吐量的绝对测量值。以每秒token数（tps）为单位，对应不同张量并行（TP）维度和p值。

表7: 推理首个令牌生成时间（TTFT）的绝对测量值，张量并行16。以毫秒为单位，对应不同批处理大小和p值。

表8: 推理首个令牌生成时间的绝对测量值，张量并行8。以毫秒为单位，对应不同批处理大小和p值。

E.6 在NVIDIA硬件上的测量

NVIDIA硬件性能数据。本节我们展示了在H100和A100上的完整推理加速测量值，如论文5.2节所述。

表9: H100和A100上的TTFT。以毫秒为单位，对应不同批处理大小和p值。