Pipeline Parallel

本篇目标： 为什么要重构 因为要支持vpp（interleave 1f1b） zerobubble dualpipe dualpipev，之前的只支持1f1b，gpipe。 从vpp思索stage的执行序 在gpipe和1f1b里面，一个stage就是一个rank，但如果加上vpp，一个rank上就很多虚拟的stage了。以前设计vpp的初衷就是，一个mb在等通信的时候可以让另一mb去做计算。我们给每个物理rank都绑定了一个executor和一个active queue，为了达成这种目的，我们必须保证两个虚拟stage的active在指令执行顺序没有依赖和申请的资源不会死锁，这样才能让virtual stage A的通信和virtual stage B的计算相互overlap。 那么上面说的这种overlap要怎么排布active才能实现呢？这个必须先定义边界。 从dual（奇美拉）思考stage的执行序 dual V的思想是让一个物理rank不会只绑定模型的一层，而是模型的某两层。这个和vpp的区别是，vpp的优化是不同mb不同时做通信and计算带来的掩盖效果吗，而dual的思想是一个rank可以做几个model layer进而降低了单任务的IDLE，pipeline从两个方向同时发任务，气泡少。可能带来的代价是任务切换带来的开销，但直觉上这部分开销不会很大。（加个TODO先，有空也会验一验） TODO: 验证 folded V / VPP 下更细粒度 action interleave 带来的 kernel launch、stream sync、activation memory 和 comm fragmentation 开销。 mb编码逻辑 随着foldv，dw分离，vpp的引入，mb的编码应该用一个结构体保存所有metadata而不是简单的id。 action之痛 做到vpp的时候，开始发现femtotron这套基于active-mb的扩展性不够好，在做dw分离的时候我把一个backward action拆成一个d action和一个w action，其实这个设计就给后面的复杂埋下了问题。或者说一开始femtotron的forward，backward，sfrb等设计就让这种设计注定变得不可扩展。因为如果一个schedule包含多种优化，一个action的含义必然是复杂的。 所以没办法，先研究一下megatron是怎么实现的。

本篇目标： schedule要不要做成异步 gpipe的非异步调度 就是简单的做sf和rf，sf和rf之间依赖于torch.dist的api做阻塞同步。 在其他博客中也说这种静态调度和下发cpu指令差不多，按顺序一条条执行，执行完了整个程序就跑完了。 异步的问题 那为什么不直接把任务丢到下游（backward是上游）然后直接做下一个mb的计算呢，这样看起来还可以让sm利用率高。 搞成异步之后快的那个stage确实会更快推进，但是慢的那边会更慢。要么就是两遍差不多快，一样没有什么提升。 另外就是action memory的值会比较不确定。可能会非常大。而静态调度actionmemory的大小是可控的。 语义问题 如果搞成异步下发，那说明会有任务的积压，这些任务做了一半以后，checkpoint要按照哪个标准做记录呢？这也是很难做的。 静态调度的状态就是比动态调度状态更少。动态调度存储，加载，更新状态都会更难而且可能bound在控制流。 好处？ 动态调度对慢节点的容忍度好，但是绝对不是pp schedule pipe里面的慢节点。因为llm场景每个mb的时间都差不多。就算是卡慢了也不会这样处理，直接换掉就行了。这种对慢节点的容忍度指的是对于一些异构的流程，比如说rl的几个步骤，以及搜推处理sparse数据等。但是这两个流程我都暂时不特别熟悉，后续还要继续学习。 如果做同步pipe，算子做还是框架做？ 放在框架 那就是调用action之前barrier一下，因为这样涉及到多节点，所以启动开销会比较大，好处是位置浅好定位。 放在算子 那训练框架侧就只启动torch.dist的接口，torch.dist底下再接入通信算子库。算子内部barrer。这种问题是调用栈会很深，以及算子级更难定位，好处是算子层可以做更深度的优化，比如做smfree把单节点的mfu打上去。（不过还是那个问题，单节点mfu可能真高了，全局不好说） 那其实这里再给自己开个新todo，试试sm free的算子实现，用CE做通信，到时候跑训练看下效果。

本篇目标： 问题背景 什么是schedule，这个词含义很广但是在训练框架这里一般考虑的是f和b任务之间的编排。 PP pp开几一般就是把所有layer除以几，然后每个就是一个stage的layer数量。一般按照layer切。 GPipe 有很多mb，每个mb要做很多stage（模型的layer或op，跨卡或跨机，这些都行），GPipe就是要等到所有的mb都做完他们自己的所有前向stage，然后开始反向stage。气泡比较多，此外因为前向和反向的layer是反过来的，所以对于一个mb来说，他做的这些stage里面做前向越早的那个stage，做反向越晚。也就是inflight越多。inflight越多就代表得保存中间的状态，占显存。所以2个肉眼可见的缺点一个是空泡另一个是inflight。 在 Perfetto 中打开 GPipe trace 如何实现一个GPipe呢？（底层组件假设已经分好了，我们只需要考虑怎么把任务排好发出来，底层组件的事情可以见后文如何实现一个调度器） 非常的简单，给每个stage执行的载体（GPU）从mb0下发到mbn就可以了。然后执行那个stage对应的layer的前向传播/反向传播。 1F1B 做F的预取，然后让F和B同时进行。中间的卡交替进行f和b。好处是inflight少，但是空泡不减。 在 Perfetto 中打开 1F1B trace 如何实现一个1F1B呢？也并非很难，假设我们的stage执行载体（GPU，虽然总是括号里写GPU但是某些场景不一定是GPU，目前为了便于理解先这么写） 有m个，那只需要给stage编号（这个编号代表第一次启动任务的顺序）为i的stage提前分配m-i+1个mb就行了，mb的序号是从0到m-i。 然后这样预填充完之后，只需要做简单的配对+交叉即可。因为是1f1b，所以只需要交替下发f和b任务，f和b任务对应的mb编号只需要匹配最近一次任务即可，如果是b则找最旧的未完成mb任务id让fb闭合，如果是f则找最新的未完成mb任务id+1。 interleaved 1F1B 也叫vpp，把一个stage再划分为几个虚拟stage，用interleave的形式排到几张卡上。这个场景为什么能减少bubble在我第一次理解的时候其实不是很直观，因为我思考的是，就算切细了那三角形的空泡依然存在，为什么空泡会少。所以就计算了一下size。只算开始部分的三角形空泡（结束时候是对称的就不管了）不计算的很细的话，我们看三角形空泡里面最长的部分，也就是最底下的那条，长度正比于每个f/b的时间*（pp-1），但这里要注意一个很容易想当然的问题，这里的pp是物理pp数，也就是真实的stage，而不是虚拟的stage。因为我们真实的stage数量一般和gpu数量一样，所以就算很多虚拟stage，一次填充到流水线的阶段也最多只有物理个gpu数。那切完以后f/b的t就变小了，显然空泡就小了。 这么解释不太直观，，最直观的其实是，让最底下的那个rank早启动。假设就是rank0到7，rank7得等好几个阶段才能启动，那就把阶段切细，然后启动的就快了。但是如果切得太细，跨rank（其实是stage）通信不能忽略，那就也不行。 然后写这个还想一个问题就是stage到底跨卡还是跨机还是跨什么东西，问了下ai说具体情况具体分析（等于没说）然后翻了下之前（未发布）的训练框架学习笔记，原则上stage没有跨什么东西的限制，但是在机内有高速互联的情况下一般是跨节点的。因为高速互联要留给tp。优先级tp>dp>pp因为我们假设tp每一层都开一次，那么tp的通信量是 $$ seq\_len \times batch\_size \times layer \times hidden $$dp没有layer这个维度肯定要少点。pp一般都可以overlap了。如果节点内没有高速互联是需要开pp的。 Chimera 最接近dualpipe的办法。 初始流水线：s0f-s1f-s2f-s3f-s3b-s2b-s1b-s0b。 Chimera主要减少了bubble，前面二者有bubble都是因为GPU来任务的时间难免有pipeline式的三角形空泡问题。但是三角形空泡来源一个先入为主的假设就是我们总假设只能gpu0开始做mb0stage0。如果让其他gpu也同时开始一个任务，三角形空洞就能补上很多。（拓展，Chimera只是同时走两段pipe，能不能更多的pipe，收益如何） 其实就是排两个交叉的流水线。依然要vpp把stage加倍。 假设原先4stage，vpp成8个。 s0 0 7 s1 1 6 s2 2 5 s3 3 4 就大概这样的。 如果纯做vpp的话是这样的： s0 0 4 s1 1 5 s2 2 6 s3 3 7 那为什么vpp的效果不如Chimera呢。可以观察一件事情，stage越“在时间上靠前”被下发做f的，在做b收口的时候越晚，占用的显存就越大。所以有一个直观的结论，不同stage的显存开销在时间上是不均匀的。最影响显存开销的就是f的第一stage-b的最后一个stage这一对。这一对fb启动最早释放最晚，所以我们如果多同时启动几个这样的f，就可以让显存开销在时间上更均匀，进而降低了显存需求量的峰值。我们做vpp的话很难让不同stage的显存分配量是均匀的，甚至还有可能让inflight叠加。 ...