在流水线并行、数据并行（DP、DDP与ZeRO）之后，今天我们来深入探讨Transformer预训练中最核心的并行范式——来自NVIDIA的张量模型并行（TP）。其核心思想非常直观：将模型参数沿纵向切开，分散到不同GPU上独立计算，最后再聚合结果。

撰写本文时，我仔细翻阅了Megatron的源码。整体框架并不难理解，但一旦涉及混合并行的细节，需要考量的因素就变得复杂起来。因此，本文配合了大量图例，许多设计思路与实现细节均源自源码阅读。既然已经深入源码，下一章干脆带大家一起逐行解读Megatron的代码——手动并行毕竟不像调用API那般简单，必须根据实际模型设计框架，还要清楚在哪里打日志、做checkpoint。

即便你对Transformer设计不熟悉也没关系，本文会通过图解详细说明。不过建议阅读前先了解：Transformer学习笔记二（self-attention）、图解大模型系列数据并行上篇（DP与DDP，其中包含AllReduce通信的详细讲解）。

全文结构如下：

一、切分权重的两种方法

二、MLP层

三、Self-Attention层

四、Embedding层

五、Cross-Entropy层

六、经典并行：TP + DP（Megatron + ZeRO）

七、实验效果与GPU利用率

八、参考

最后说个题外话，Megatron是变形金刚里的霸天虎领袖，如今在环球影城当脱口秀演员（开个玩笑）。没有它，就显不出擎天柱的战绩；没有它，就没有变形金刚电影，也没有我曾经的暑假快乐。而现在，没有它就没有大Transformer的诞生。所以，把“引燃人类AGI的热情”当作Megatron的阴谋，拍进下一部电影里，是不是挺带感的？

一、切分权重

设输入数据为X，参数为W。X的维度为(b, s, h)，W的维度为(h, h')。其中：b是batch_size，s是sequence_length，h是hidden_size（每个token向量的维度），h'是参数W的hidden_size。每次forward的过程如下：

为方便画图，图中假设b=1。现在W太大，单卡装不下，需要切开放到不同卡上。这时面临三个问题：怎么切分W？切完怎么做forward？做完forward怎么做backward，进而求梯度、更新权重？一般来说，我们可以沿着W的行（h维度）或列（h'维度）切分。下面分别介绍这两种方法。

1.1 按行切分权重

（1）forward

用N表示GPU数量。有几块GPU，就把W按行维度切成几份。下图展示了N=2的情况：

W按行切开后，X的维度与它对不齐了。怎么办？很简单，把X也“按列切开”：

（2）backward

做完forward取得预测值Y，计算出损失L，就可以做backward了。重新画一下forward过程，并加入backward部分：

f和g分别表示两个算子，每个算子包含一组forward+backward操作。图中的每一行表示一块GPU上的计算过程。g的backward：假定要求梯度，只要把同时广播到两块GPU上，就能独立计算各自权重的梯度了。f的backward：当模型有多层时，梯度要从上一层向下一层传播，比如传播到X后才能继续往下传，这就是f的backward的作用。

1.2 按列切分权重

（1）forward

按列切分后的forward计算图如下：

（2）backward

g的backward不难推出。f的backward：因为损失L对X的依赖来自XW1和XW2，所以需要将两块GPU上的梯度相加。

现在，“按行”和“按列”切分权重的方法都已介绍。在Megatron-LM中，权重切分由这两个基础算子组合而成。接下来，针对Transformer的不同部分，看看Megatron-LM具体怎么切分。

二、MLP层

2.1 MLP层的张量模型并行计算方法

MLP层结构最简单，咱们先看它。其计算过程如下图：

其中GELU是激活函数，A和B分别为两个线性层。在Transformer中，一般设h'=4h。假设有N块GPU，要把MLP层的权重拆到上面计算，Megatron的做法如下：

在MLP层中，对A采用“列切割”，对B采用“行切割”。f的forward：把输入X拷贝到两块GPU上，各自独立做forward。g的forward：计算完Z1和Z2后，GPU间做一次AllReduce，相加得到Z。g的backward：把拷贝到两块GPU上，各自独立算梯度。f的backward：当需要把梯度传到下一层时，两块GPU做一次AllReduce，把各自的梯度相加。

为什么对A列切割、对B行切割？目的是尽量保证各GPU计算独立，减少通讯量。GELU是非线性函数，如果对A行切割，必须在GELU前做AllReduce，增加通讯。而列切割则能让每块GPU继续独立计算。A确定做列切割后，B自然就做行切割了。

2.2 MLP层的通讯量分析

MLP层forward一次AllReduce，backward一次AllReduce。一次AllReduce包含Reduce-Scatter和All-Gather两个阶段，每个阶段通讯量相等。设每个阶段通讯量为，则总通讯量为。根据计算图也容易得出具体值。

三、Self-Attention层

3.1 Multihead-Attention的计算方法

先快速过一下multi-head attention层的参数构造。当head数量为1时，self-attention层的计算方法如下：

seq_len和d_model分别对应本文的s和h。attention层需要训练的三块权重。k_dim和v_dim满足特定关系。单头清楚后，来看多头的情况。下图是num_heads=2时attention层的计算方法：对每块权重，沿列方向（k_dim维度）切割。每个head上的维度变为(d_model, k_dim//2)，单独计算后concat起来。

attention的多头计算简直是为张量模型并行量身定做——每个头可以独立计算，最后concat。也就是说，可以把每个头的参数放到一块GPU上。整体流程如下：

对Q、K、V三个参数矩阵进行列切割，每个头放一块GPU上并行计算。对线性层B进行行切割，方式和MLP层基本一致。实际应用中，不一定一个头占一块GPU，也可以多个头占一块GPU，但最好保证head总数能被GPU个数整除。

3.2 Self-Attention层的通讯量分析

和MLP层类似，self-attention层也做两次AllReduce，总通讯量相同。结合来看整体计算逻辑：

四、Embedding层

4.1 输入层Embedding

Embedding层由两部分组成：word embedding（维度(v, h)，v是词表大小）和positional embedding（维度(max_s, h)，max_s是最大序列长度）。positional embedding的max_s通常不长，每个GPU拷贝一份没问题。但word embedding的词表大小很可观，需要拆分到各GPU。具体做法：

输入X过word embedding，就是用token序号查找对应的词向量。比如输入[0, 212, 7, 9]，就要去word embedding中找第0、212、7、9行的词向量。假设词表有300个词，将word embedding拆分到两块GPU上，第一块维护[0, 150)，第二块维护[150, 299)。输入X查找时，能找到的就正常返回，找不到的就把向量全置0。查找完毕后，每块GPU上的数据做一次AllReduce，得到最终输入。

4.2 输出层Embedding

输出层也有一个word embedding，把输入映射回词表。通常输入层和输出层共用一个word embedding。计算过程如下：

必须时刻保证输入层和输出层共用一套word embedding。在backward过程中，输出层会对word embedding算一次梯度，输入层也会算一次。更新权重时，要用两次梯度的总和。如果输入层和输出层在同一块GPU上（流水线并行深度=1），这点不用操心（实践中大部分用Megatron的项目都这么干）。但若在不同GPU上，就要在权重更新前，让两块GPU上的word embedding梯度做一次AllReduce。

五、Cross-Entropy层

最后是计算损失函数的一层。回顾4.2，输出层过完embedding后的样子：

通常需要对Y1和Y2做All-Gather，concat成Y，然后对Y每一行做softmax，得到每个词出现的概率，再与真值做cross-entropy。但All-Gather有额外通讯量，当词表v很大时，开销不容忽视。可以优化：

每块GPU上先按行求和，得到各自的GPU_sum(e)。然后做AllReduce，得到每行最终的sum(e)（softmax的分母）。此时通讯量为。然后每块GPU可以计算各自部分的e/sum(e)，与真值做cross-entropy，得到每行loss，按行加总得到scalar Loss。最后将各GPU上的scalar Loss做AllReduce，得到总Loss，通讯量为N。这样，通讯量从大大降至。

六、张量模型并行 + 数据并行

至此，张量模型并行的计算架构基本说完了。实际应用中，对Transformer类模型，最经典的方法是张量模型并行 + 数据并行，并在数据并行中引入ZeRO做显存优化。架构如下：

node表示一台机器。一般在同一台机器的GPU间做张量模型并行，在不同机器上做数据并行。图中颜色相同的部分为一个数据并行组。这样设计大概率与两种并行的通讯量有关，具体来说，与TP和DP模式下每一层的通讯量以及backward计算方式有关。

6.1 TP与DP通讯量

先分析通讯量。关注Transformer每一层：在TP中，每层做2次AllReduce，通讯总量为。其中，，通讯总量为。在DP中，每层做1次AllReduce（先不考虑ZeRO拆分权重），通讯总量为。通讯量主要来自梯度，总通讯量为。比较和，忽略常数后，最终比较和。实际应用中，前者可能略大，但量级在左右。通讯量大的尽量放一台机器内（带宽大，通讯时间小），通讯量相对小的可以跨机器。

6.2 TP与DP计算backward的方式

回顾TP在backward过程中，从上一层向下一层传播时，所有GPU间需要做一次AllReduce。而DP中，本层算完梯度后，可以正常发给同DP组的GPU做AllReduce，同时继续往下一层做backward。下一层不依赖上一层的梯度聚合结果，因此对带宽要求没那么高，可以放到机器间做DP。

七、实验效果

7.1 实验配置与核心数据

原始论文采用32台DGX-2H（每台16张Tesla V100 SXM3 32GB），共512块GPU，训练不同大小的GPT2模型。核心实验数据如下：

每一行表示模型不同大小。Model parallel列是只用TP时消耗的GPU数，对应蓝色柱状图；model + data parallel列是TP+DP时消耗的GPU数，对应绿色柱状图。

7.2 纯TP实验效果

第一行是整个实验的基线：一个能装进单卡的GPT模型，不涉及卡间通信，GPU计算效率100%。随着模型增大，需要的GPU数变多，通讯量增大，单卡计算效率逐步下降。从1卡到8卡，模型大小和GPU数量成正比。但论文中最大只做到8卡——一台机器明明有16张卡，为什么不做16B左右的模型？回想ZeRO对显存消耗的分析，模型增大时，不仅是参数变多，activation这样的中间结果也占大头，不引入显存优化，一台机器装不下16B的模型。

7.3 TP + DP实验效果

再看TP+DP。第一行基线中，模型大小可塞进单卡，每台机器选2块GPU，共64块，每块上都放完整模型，做纯数据并行，单卡效率96%。第二行模型变大，单卡装不下，每台机器选4块GPU，每2块组成一个TP组，共2个TP组。所有机器共64个TP组（128卡），组间做数据并行。以此类推。蓝色和绿色柱状图一一对应，分别表示单个模型在1、2、4、8卡上的单卡效率。引入数据并行（跨机）后，绿色单卡效率并未大幅下降，原因在于DP组内做backward时，下一层计算不依赖上一层的梯度AllReduce结果——你算你的，我发我的，计算通讯比影响不大。

最后，一台DGX-2H售价约40万刀，硬件成本40×32=1280万刀。只能说，Megatron注定生于NVIDIA家。

7.4 GPU效率计算

说说柱状图中weak scaling指标怎么算。单卡装下模型时，单GPU每秒计算量39TeraFlops（可近似理解为矩阵乘法和加法的次数），不涉及通讯，效率100%。上到512卡TP+DP后，计算总量为15100TeraFlops。假设512卡无通讯影响，理论算力应为39×512=19968TeraFlops。算力利用率=15100/19968=76%。这就是贯穿全文的76%指标——即使做分布式训练，GPU算力也不会被通信浪费太多，这正是Megatron的卖点之一。

关于Megatron的理论介绍就到这里。在大模型训练的下个篇章里，一起读Megatron源码，学习经典的大模型预训练方法，顺便看看Megatron的第二个卖点：“我们的代码用起来很简单！”是不是事实。