Transformer模型分布式并行通信量浅析

1.数据并行DP（朴素数据并行，Zero数据并行之后补充）

$O(h^2\ast l)$
每台机器做完自己的梯度后需要做一次All reduce操作来累积梯度，故一个batch计算发送的数据量为每层梯度大小$h^2$乘以层数$l$

优点：运用简单，效率高，计算和通讯之间可以做异步。

缺点：有时候单个GPU无法容纳下整个大模型做训练。

流水线并行PP（Gpipe）

$O(b\ast s\ast h\ast l/k)$
每个pp层之间发送中间变量大小（和输入x相同大小）为$b\ast s\ast h$，乘以正向和反向一共要传递的次数$2\ast (l/k-1)$，得到上述大致的数量级

优点：通讯量小，数据无关。

缺点：要保证模型能均匀切分，否则会影响性能。需要用到重计算，来支持更大的批量$b$，从而保证流水能发挥作用。反传之前需要等待所以微批量计算。

张量模型并行TMP

$O(b\ast s\ast h\ast l)$
MLP层第一个W1纵向切，第二个W2横向切能保证一个MLP只做一次All reduce操作通讯。同理attn根据头数来切分投影层。两者每次的发送数据大小和PP一样都是$O(b\ast s\ast h)$，不同是，TMP发送量和模型Transformer层数成线性关系。

优点：能切分很大的Transformer模型。

缺点：通讯量大，且通信一般只在节点内部的多卡通讯，不做跨节点通讯。头数需要被GPU整除。

*一般而言$b\ast s$略大于$h$，$k$是$l$的1/10，故一般而言，通讯量TMP>DP>PP.在实际的工程中，一般TMP和PP都在节点内通讯，只有数据并行会做跨节点通讯。