LLM推理，显存都去哪儿了

Mastering LLM Techniques: Inference Optimization | NVIDIA Technical Blog

Pipeline并行：

缓解"气泡"(Bubble)的办法：把大batch拆成多个小batch；

Megatron的Tensor并行：

1. MLP: 连续乘以2个矩阵；第一个W矩阵，竖着切分为左右2个；第二个W矩阵，横着切分为上下2个；g表示2个结果矩阵加和到一起；这样可以避免中间结果Y的AllGather通信；

2. Self-Attention: 一张卡负责前一半heads的Q、K、V（KV cache也是放前一半heads的）；后一张卡负责后一半heads的Q、K、V（KV cache也是放后一半heads的）；转换矩阵B横着切分为上下2个，各自乘完之后，2个结果矩阵相加；可以避免中间结果Y的AllGather；

最后一层FN的W矩阵（也就是复用的Embedding组成的矩阵，几千行几万列，因为vocabulary size是几万）很大；把W按列切分到多个GPU，模型并行来计算；跑完的结果，不需要all-gather，只需要每张卡element-wise得计算e^x，然后加和自己那部分的e^x，然后all-reduce所有GPU上的加和值做分母，再将本卡上的那部分e^x除以该分母得到概率p；

Sequence并行

和Megatron Tensor并行，一起用；

把LayerNorm、Dropout，进行多卡并行，每张卡处理一部分token的完整激活；

用AllGather和Reduce-Scatter，代替之前的AllReduce，总通信量没有增加；

如上，Linear完后，所有卡上的矩阵，要进行加和到一起；但现在改成Reduce-Scatter，每张卡加和自己那部分token的子矩阵，也就是大矩阵横着切；加完后，每张卡计算Dropout和LayerNorm都只在自己那部分token的子矩阵上计算；输入到后面的Linear层之前，所有GPU卡上的加和后的子矩阵，进行AllGather；

上面，forward阶段，从左往右，依次是
g: AllGather

g~: Reduce-Scatter

g: AllGather

g~: Reduce-Scatter

注意：同一位置的通信操作，在forward和backward时，是不一样的通信pattern;

FlashAttention:

1. 巧妙分片，减少读取显存次数；

2. Fused kernel；

读写显存：

设M为shared-memory大小，N为sequence length，d为单个head的hidden size;

老方法：O(N*N);

FlashAttention: 外层循环一共做N*d/M次(K、V每个的总size为N*d，shared-memory大小为M）；每次总共读Q的size为N*d；所以是O(N*N*d*d/M)；

因为d*d/M远小于1，所以优于老方法；