Colossal AI 多维TP

Colossal AI 多维TP

1. 2D TP

1.1. SUMMA 2D 矩阵乘法

数值示例：

条件：每个矩阵都可以均匀的拆分为 p=q^2块（行q块，列q块·）

1.2. Transformers上的应用

b: batch size s: seq_len h: hidden size p: GPUs q: p=q^2
输入shape为{b, s, h}{bs, h}{bs/q, h/q}，实际使用时将b和h进行拆分，如下图所示。

通信量和计算量对比（包含activation checkpointing）

性能对比
1）Weak scaling concerns the speedup for a scaled problem size with respect to the number of processors

2）Strong scaling concerns the speedup for a fixed problem size with respect to the number of processes

3）memory performance

注：测试环境
4 NVIDIA Quadro RTX 5000 GPUs on each node, and nodes are inter-connected with Mellanox InfiniBand

2. 2.5D TP

2.1. Tesseract 矩阵乘法：将一个矩阵乘法拆分为多个SUMMA 2D矩阵乘法

数值示例

2.2. Transformers上的应用

性能对比：
1） Weak scaling

2） Strong scaling

注：4 NVIDIA A-100 GPUs per node

2.3. 对比和结论（C=AB）

1） 当d=1时，即为SUMMA 2D 的矩阵乘法。
2） 2D和2.5D不仅拆分了B，还拆分了A。
3） 2.5D中，p=dq^2。
4） 2.5D相对于2D而言，通过把B复制了d份来减少通信次数。
5） P=dq^2固定的情况下，不同的{d,q,q}配置性能相差甚远。

3. 3D TP

3.1. 算法

3.2. Transformers上的应用

1） 输入在 {b,s,h} 三个维度上都进行拆分。
2） 参数在in_size上按p拆分，out_size按p^2拆分。
3） XA后都有一个reduce-scatter操作。

性能对比：
1） Weak scaling

2）Strong scaling