Profile and benchmark

backward相比于forward慢了2倍。
（update：我这里好像backward也把forward算上了，所以可能算出来应该是1:1的，但是经验值看，应该两者是有2～3倍的差距的。所以不清楚这里具体是什么原因）

相比之下标准差并不大，说明benchmark脚本没什么问题

还对比了一下用perf counter和cuda.Event，区别也不大，因为也没有什么排队延迟。为了方便控制，就都用cuda Event了

关掉warmup之后，结果明显不同，同时标准差也升高了：

warmup主要是去掉一些一次性成本以及不稳定因素：
* CUDA相关的初始化
* GPU功耗的爬升。开始执行的时候，可能会以低频率执行。
* cuBLAS等算法库会做算法选择和缓存（甚至探测不同的实现），直观想就是可能前几次还在尝试一些更高效率的代码路径（可能对不同workload影响不同）。后续则会变得稳定
* 内存分配，申请显存池
* 一些代码链路，指令的cache，cpu cache

mix precision:

效果最差的是用float16存+算
最好的是都用float32
用float16作为中间结果，存到float32中效果也还可以

autocast的实验结果：
* grad/param是float32
* ffn的输出是float16
* LN的输出是float32

layer norm有一个计算方差和mean的操作，会用sum。所以用了float32

bf16的结果，看起来ln还是用的float32。
因为bf16只是扩大了精度，但是算方差的时候，计算两个很相近的数值再算平方，比较依赖小数位的精度。

然后跑transformer：
forward：
float32:168ms
float16:104ms
bf16:216ms

backward:
float32:350ms
float16:223ms
bf16:445ms

我的gpu上好像对bf16支持不太好。RTX2080只支持float16，bf16是在Ampere系列上才有。
然后对于这种不支持的计算，torch/cuda（不确定是谁）默认的行为是转化成float32，算完再转化回来。所以反而更慢

因为显存比较少，medium size的我就已经无法执行backward pass了
改小了点context length：


可以看到模型越大这个混合精度的效果是越来越明显的。

memory

直接dump最后的快照就可以看出这个趋势。这个是large model，context length是128

最下面这块是模型占用的内存4G。然后第一个爬坡就是forward，用了3G去保存这些activation

然后到backward阶段，forward的一些activation一边释放，一边增加新的grad的占用。到最后面activation为空，然后都是grad的内存，也差不多3G～4G。因为需要保存grad的就是model size这么大。

然后是optimizer，差不多占用了model size的2倍，因为每一个参数需要保存那两个动量。

跑多个epoch会发现，有一个峰值是前一个epoch的grad还没释放，新的activation已经计算出来了，会有一个小的峰值，随着前一个epoch的grad释放就好了。

用了混合精度好像影响也不大。猜测可能是：
* grad/optimzier因为精度原因，本身就还是float32的
* activation，部分层（比如占大头的attention需要softmax/sum）也是float32的。所以整体能够转化成fp16的比较少

最多的内存就一直是optimizer state和model size了。

然后有关context length的对比，这里我跑不了2.7B的模型，就用small的对比了一下
128长度：

512:

1024:

可以看到到512的时候，内存分配的峰值就受到activation主导了，应该就是那个N方的attention

放大了看，有一些会保留下来的尖刺，看栈也可以看出来就是attention。

至于为什么需要保留这个N方的激活值，是因为softmax的梯度计算依赖原始的输入，所以需要把计算出来的logits保存下来

然后问题中的residual stream的内存占用。就是context_length 乘 d_model。

flash attention

时间上，基本backward都是forward的2倍以上
内存占用上，peak memory都和N方正比了，就是用来保存attention的score。

按照上面分析的，需要把attention的logits保存下来，所以至少有一个16K x 16K x 4byte = 1G的内存占用保存attention的分数

不过具体为什么是这个值也不太清楚。和上面分析的也不太一样。

加了compile影响也不大。内存占用也没区别。可能memory bound的也说不定。因为小的模型配置下forward还是快了一些的

transformer的，backward快了不少。可能FFN那些地方矩阵乘法比较多，优化的比较好一些。

flash atten impl

这块跟着assignment里给的伪代码来就行，一步一步实现也比较简单。
* 记得不要跟着flash atten的原文，里面没有scaled by sqrt(d)，所以容易出错

然后在写triton的时候，可以先把启动TRITON_INTERPRET的版本写完，能过了之后再去写cuda的版本。其实整体代码和pytorch的都差不多，很多地方把torch.xxx换成tl.xxx就可以了。

然后我这里还遇到一个比较离谱的，在这里发现解决方法了：https://github.com/triton-lang/triton/issues/4813
在使用float32 dot float32的时候，因为一些20系列的显卡没有fp32的tensor core的实现，而triton有一些bug，导致会报出奇怪的报错，完全看不出来是哪里出错了。
* 看里面的解释，是因为在找tensor core的实现没找到，然后fallback到普通版本的时候失效了，导致的报错。
* 所以关掉tensor core的实现，在所有的dot中，都加上allow_tf32=False就可以了。
* 或者换一个显卡应该也行。。。

backward这块没搞，所以benchmark主要是针对forward来的：

naive:

flash:

可以看到内存占用大幅度减少了。不过执行时间略有上升，主要是因为我这里的block size设置太小了。导致没有利用好全部的内存。把block size设置大一些估计可以比naive的速度快。

但是我发现在我的GPU上提高block size会爆掉，因为share mem不够用了。

研究了一下，可以把pipeline的stage调小点，改成1，然后block设大点。此时速度会有一定的上升。

然后还有一个优化方法，就是用半精度的QKV，这样load进缓存会快一些，同时还可以用上更快的tensor core实现（毕竟我这个GPU用不了tf32）

DDP

gloo, cpu benchmark all reduce，4 process

nccl, cuda benchmark all reduce，4process

快了5倍，而且我这个机器上是没有nvlink的，可能是nccl有什么特殊的优化。

gloo, 2 process:

cuda, 4 process:

发现2 process到4process并不是线性的增长，可能是用了什么特殊的all reduce算法？比如log什么的

查了一下可能是因为常用的ring-allreduce中，在带宽瓶颈的情况下，延迟是和 2(P - 1) / P成正比的，所以在这个case中，P从2变成4，延迟上升了1.5倍

为了验证这个猜测，再跑一个gloo的 6proc（因为我没有6个gpu了），延迟是83ms

从43 -> 67 -> 83。按照上面公式算的话，就是1倍，1.5倍和1.66。（并没有完全对上

可能还有其他的因素干扰，因为这里是内存带宽，可能还受到NUMA等因素干扰

NaiveDDP

跑的还是small model，大概是600ms一个epoch。reduce的占比相对来说比较小

batch size变大了我的GPU就跑不起来了。变小了就还是10ms，感觉瓶颈主要在一些启动的开销上。

加上flat的优化之后反而更慢了，因为瓶颈不是这里的通信，加上了还需要一些额外的拷贝把grad摊平

重新跑一下large的：


1个step就爆显存了，不过确实快一些。
* 为什么能跑一个step，按照上面perf的，第一个step后optimizer state才会吃内存，所以此时占用的内存会更多。然后因为用flat的方式需要重新拷贝一份grad，（相当于拷贝了一份param）

带上overlap之后明显速度快了不少，一个step速度也快了100多ms

ddp bucket那个profile我这里看不出效果，原因同上，还是不是带宽瓶颈

4D parallelism

XXL模型：

126层，每一层两个矩阵，2 x d_model x d_ff
总共是204B

FP32的话，模型占用的内存就是816G。然后gradient和optimizer state还需要占3倍，总共就是3200G+

对于backward来说，需要保存的就是所有的activation，应该是sequence x (d_model + d_ff) x num_blocks。seq不长的情况下占比不大。

后面还有一些小题，包括optimizer的，看上去没有特别关键，就暂时先跳一下。