在section 3 distributed training这一节给了很详细的背景介绍和分析，包括：
- FSDP和pp的冲突点
- compute/communicate overlap的定义
- Critical batch size的分析，在appendix上还给了详细的推导过程，以及直观的理解
a (mini-)batch is used to approximate the true gradients of the weights with respect to the loss. Increasing the batch size B generally improves this approximation, leading to more efficient steps
However, for large batches the approximation is already accurate and additional samples provide a negligible improvement, leading to a waste of computing power.
- 整体是针对后面的方案，就是low batch size训练来说的。也提到了在low batch size的时候，各种并行策略会遇到的问题。
同时给出了很多不错的reference，都是比较经典的论文值得读一下。
[图片]
[图片]
这里画的深度优先的调度和1F1B耦合起来了。
- 目前感觉调度的顺序，和是否使用1F1B是分开的
- 不过1F1B可以尽早进入backward，和深度优先的思想（快速进入下一个stage，进入backward）是相同的
- 整体目标都是为了减少memory
- 那广度优先的是否可以1F1B呢？
- 好像是不行的，因为能开始进行backward的时候，前面的forward都已经跑完了
[图片]
如何理解这几种调度的区别呢？
看上面的图
- 横向的就是microbatch，纵向的则是stage
- 相同颜色的在相同的GPU上
调度的时候，可以遵循下面的规则
- 同一行中，左边的块优先于后面的块执行。因为是按照microbatch切分，重排序这里没有意义
- 同一列中，上面的块优先于下面的块执行。因为是stage的依赖
- 同时这里我把backward也作为一个stage引入进来
[图片]
先考虑最基本的AFAB（GPipe）
- 根据GPU进行了聚合，每个GPU的执行流都是并行的
- AFAB表示的是执行完所有的forward，再执行backward
- 图中的含义就是执行完上面一整行后，才会执行下一行
[图片]
1F1B
- 因为考虑了warmup的阶段，所以前面几个stage会先横向执行一下。
- 不过在最后一个stage可以看出来，执行的效果是先纵向，再横向的。
[图片]
Interleaved 1F1B，looped pipeline
- 这里主要画了最后一个stage，因为偏差的没有太多。不过可以看出来思路
把依赖关系展开看会更好看一点：
[图片]
- 这里横轴同时代表了时间
- 这里可以看出来一点小的决策点。看黄色的GPU
- 在执行完第一个microbatch的3F之后，可以有两种选择。选择执行第二个microbatch的3F，或者是执行第一个microbatch的3B。对应的就是第二个step的两个黄色块的选择。
- 这里其实就有了深度优先和广度优先的思想了。
- 如果我们选择一直执行新的microbatch的3F，就变成了上面的AFAB
- 如果选择执行下面的黄色块，也就是第一个microbatch的3B，就是这里的1F1B
- 用这个图来理解megatronlm中pp实现的step代码也会比较容易一些。因为这里的横轴就比较接近step的概念了。
- 唯一的区别是代码中一个step是按照1F1B来执行的。所以有一个隐含的依赖，让上面的第二行可以整体往左挪动一格。
- 如果是理解pytorch的代码，就是对应的，因为torch的实现中是1B1F的。
[图片]
然后是interleaved 1F1B，核心思路还是优先执行backward，也就是优先调度更下面的block
[图片]
[图片]
看黄色的GPU：
- 执行4个microbatch之后，就会有选择，可以执行7F，或者是新的microbatch的3F
- 按照深度优先的原则，优先执行7F，然后开始7F/7B的interleaved的执行
- 因为一个GPU上有多个stage，这种looped pipeline的情况下，两个stage之间需要再过N_pp个设备。为了保证pipeline没有bubble，也就需要填充N_pp个microbatch
- 所以这里是按照N_pp个microbatch为一组来执行的。对应的MegatronLM中，也就是microbatch group
- 当microbatch group size小于N_pp的时候，就会引入bubble
- 大于N_pp的时候，会延缓backward进入的速度，导致占用更多的activation memory
[图片]
对于广度优先来说，则是更像是AFAB，把所有microbatch的一个stage都执行完了，再执行下一个stage。
- 优点
- 没有频繁的stage切换，每个stage只会执行一次forward和一次backward。这种情况对FSDP比较友好，不需要频繁的进行reshard/unshard
- Bubble size和深度优先的应该是一样的，都是N_pp / (N_microbatch * N_loop)
- N_microbatch > N_pp的时候，可以做overlap
[图片]
- 当N_microbatch N_pp的时候，看第一个图，会发现深蓝色的块直接依赖上一个step的输出。导致step之间的communication无法被隐藏掉
- 而N_micobatch > N_pp的时候，深蓝色的块的依赖在很久之前就计算完了。所以通信的时间是中间4个microbatch计算的时间。这样就可以做到通信/计算的overlap了。
- 不过深度优先的版本，1F1B也可以相互overlap
- 缺点
- Activation memory比较多
With DP
[图片]
看这张图，广度优先的算法还有一个优势：就是可以更早的完成一整个stage的backward，从而可以更早的触发gradient reduction，让gradient reduction和后面的backward compute overlap到一起。
直观来想的话：
- 深度优先的调度是纵向的调度，尽可能的完成一个microbatch。这样可以做单个microbatch的收尾逻辑（释放内存）
- 广度优先则是横向，尽可能的完成一个stage，这样可以做单个stage的收尾，比如FSDP reshard，gradient reduction
Experiment
[图片]
对比N_loop的数量，固定的batch size
- Batch size小的时候，N_loop变大之后，对于depth-first，网络开销影响比较大
- 论文中提到，根据计算显示，理论的网络开销只有1.6%，但是实际上因为同步开销（应该是kernel launch的开销），影响超过40%
- breadth-first也有影响，不过因为overlap的原因，影响相对较少
[图片]
小batch情况下的对比：
- 小batch的情况下，大家bubble都比较多。
- Depth first的方案引入loop虽然减少了bubble，但是引入了更多的网络开销。在小batch场景效果更差
- 所以小batch的时候depth first的方案效果还没有普通的1F1B好
- 而breadth-first的方案，（我个人认为），因为有batch内的overlap，以及更好的和DP做配合。效果会更好。
逻辑线
- Looping schedule允许小的microbatch，减少bubble，但是需要增加looping的数量
- 增加looping的数量会导致网络开销变大。
- 这里的开销并非带宽，而是collective本身的同步开销
- breadth-first的方案是优先计算完整个stage，再进行下一个stage的计算
- 允许聚合stage的计算，FSDP场景下只需要通信一次
- 更早完成整个stage的backward，允许更早进行gradient reduction，增加gradient reduction和后续backward的overlap
- N_microbatch > N_pp时，允许进行overlap，后续的通信开销可以被N_microbatch - N_pp的这些数据的计算overlap起来。
文章评论