以Transformer为例，介绍大模型训练过程中常用的张量并行(TP)和序列并行(SP)技术，推导在不同并行策略下的显存大小和通信量。

Transformer结构

单层Transformer由一个Attention模块和一个MLP模块组成

假设数据类型为fp16，输入shape为(s, b, h)，b为head数量，分析激活内存：

• LayerNorm：一共2层，每层的输入占$2sbh$，共$4sbh$
• MLP：第一层的输入$2sbh$，第二层的输入$8sbh$；$GeLU$的输出$8sbh$,$Dropout$的mask为$sbh$，一共是$19sbh$
• SelfAttention：初始输入$2sbh$，$q,k,v$占用$6sbh$，$Softmax$的输出$2a{s}^{2}b$，$Dropout$的mask和输出$3a{s}^{2}b$，输出$Linear$层的输入$2sbh$，最后一个$Dropout$的mask为$sbh$。一共是$11sbh+5a{s}^{2}b$

$$Activationsmemoryperlayer=sbh(34+5\frac{as}{h})$$

张量和序列并行

Tensor Parallelism

当单卡放不下单个模型时，可以使用张量并行。megatron论文中Transformer中张量并行的流程图：

forward过程中，$f$不做任何操作，$g$表示AllReduce；backward过程中，$g$不做任何操作，$f$表示AllReduce

**MLP：**先沿着列切分第一个线性层的权重，再沿着行切分第二个线性层的权重

**Attention：**先沿着列切分$q,k,v$投影线性层，各自计算selfattention；再沿着行切分输出线性层

通信量分析：前向2次AllReduce，反向2次AllReduce

---

张量并行同时减少了模型大小和激活值，开启后激活值显存占用为
$$Activationsmemoryperlayer=sbh(10+\frac{24}{tp}+5\frac{as}{h})$$

Tensor and Sequence Parallelism

序列并行可用于减少LayerNorm和Dropout模块的激活值，通常与张量并行结合使用

沿着(s, b, h)中s方向切分过的数据，输入到LayerNorm和Dropout层计算

在输入Attention和MLP前，进行AllGather，沿着sp聚合；输入LayerNorm和Dropout前，进行ReduceScatter，沿着sp切分

通信量分析：forward2次AllGather，2次ReduceScatter；反向2次AllGather，2次ReduceScatter

SP并行减少了LayerNorm和Dropout层的激活值，激活值分析：
$$Activationsmemoryperlayer=sbh(\frac{10}{sp}+\frac{24}{tp}+5\frac{as}{h})$$

Ulysses-SP

DeepSpeed-Ulysses 是一种高效的长序列并行技术，通过All2All重排将序列维度并行转为多头并行，无需切分模型权重

通信量分析：前向和反向分别做4次All2All；单次All2All的通信量为$Nd/P$

Ulysses vs. TP-SP

从Attention模块来看：

• Megatron通过TP，显式地把$Wq,Wk,Wv$切分开，每张卡上计算所有seq的部分head的结果。
• Ulysses通过all2All，在每张卡完整保存$Wq,Wk,Wv$的前提下，让每张卡上计算所有seq的部分head的结果。

---

• Megatron Tp-Sp：4 AllGather + 4 ReduceScatter，总通讯量为8Nd
• DeepSpeed Ulysses：8 All2All，总通讯量为8Nd/P，可以增加卡数来降低通信量

参考资料

1. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
2. Reducing Activation Recomputation in Large Transformer Models
3. DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models