Context并行

最初是Sequence并行+Tensor并行；

但是，Tensor并行加多GPU parition数目后，每个GPU上的计算量减小，无法和通信overlap，变慢；

CP can better address the issues. With CP, each GPU only computes on a part of the sequence, which reduces both computation and communication by CP times. Therefore, there are no concerns about the overlapping between them. The activation memory footprint per GPU is also CP times smaller, hence no OOM issue any more.

Context并行的动图：

[并行训练]Context Parallelism的原理与代码浅析 - 知乎 (zhihu.com)

3个要点：

1. 使用了RingAttention；（Overlap: Q和当前K、V计算时，下一份K、V在通信过来的路上了）

2. 因为有casual mask，每个token只和前面token的K、V进行计算，不计算后面的；因此为了更好的load-balance，用了将sequqnce划分为2N份的技巧；

3. 复用了最新的OSS和cuDNN FlashAttention

1. RingAttention

即每张卡做1/N的sequence tokens，先计算好自己这些tokens的Q、K、V，然后每次用自己的Q，计算1/N的K和V，GPU之间达成环形来每次传输1/N的K和V；每个GPU和所有K、V计算完毕后，将N个结果向量加和起来，就得到了自己这1/N tokens的attention计算结果；

 2. casual mask引发的Load balance优化

N张卡参与CP，将sequence划分为2N份，GPU-0: 0+5, GPU-1:1+4, GPU-2: 2+3，这样每卡之间的计算量就均匀了；

Context并行，天然就是接上Sequence并行，一起用的；

TP+CP+SP:

[GPU0, GPU1]拿到sequence的前半部分tokens，

[GPU2, GPU3]拿到sequence的后半部分tokens；

GPU0和GPU1将模型做了Tensor并行划分；

GPU2和GPU3将模型做了Tensor并行划分；

Dropout、LN(LayerNorm)，[GPU0, GPU1]只用Sequence并行；（[GPU2, GPU3]同理)

FC1、GeLU、FC2，[GPU0, GPU1]只用Tensor并行；（[GPU2, GPU3]同理)

计算attention那里，特殊，因为每个token的Q要和前面的所有tokens的K、V做计算，但是它卡上只有自己这1/N部分的K、V；因此，采用Context并行来解决，轮转通信K、V，每个GPU计算N次，加和得到结果；即，[GPU0, GPU2]彼此成环，[GPU1, GPU3]彼此成环；

上例中，CP在外层，TP在内层；

Context parallelism overview - NVIDIA Docs

优势：

适合long context的训练；

比较好的overlap了计算和通信；

节约显存，正向时得到的K、V的activation(用于反向计算梯度)，按token划分，缓存在各个GPU上；