LLM并行训练1-流水线并行
并行训练-流水线
简述#
并行训练主要有三种策略:
- 数据并行训练加速比最高,但要求每个设备上都备份一份模型,显存占用比较高,通信量大
- 张量并行,通信量比较高,适合在机器内做模型并行。
- 流水线并行,训练设备容易出现空闲状态,加速效率没有DP高;但能减少通信边界支持更多的层数,适合在机器间使用。
流水线并行#
Micro-batch(Gpipe)#
将网络结构进行纵向拆分, 每张卡训练其中的几层. 下图把网络拆成4层. 如果是按照纯粹的mini-batch训练, 每层之间是纯串行的. 后面的卡会始终等待前面的卡. 所以引入了micro-batch的概念. 把mini-batch进行更细粒度的拆分, 这样在完成batch0的fp之后, 卡0可以进行batch1的fp, 卡1就能开始batch0的fp. 从而提高并行度.
存在的问题:
- 存在bubble_time: 每张卡的空闲时间 = (stage_num - 1) * (fp_time + bp_time)
实际应用中 当mico-batch个数大于stageNum的4倍时, 可以忽略bubble_time
- 显存浪费: 当进行stage3的micro-batch 3时, 还需要保存前面所有mico-batch的fp中间结果用于bp.
- 在每个mini-batch之间无法并行. 因为下一个minibatch需要等当前所有的micro-batch更新完参数
PipeDream(非交错式1F1B DeepSpeed)#

在每个micro-batch fp完成之后立刻优先进行bp. 这样可以把当前batch的中间变量释放掉, bp完成后更新本机参数, 但这种方式存在参数更新冲突, 机器1和机器2使用的参数不一样, 机器1的batch5只用了 batch1反向后更新的参数, 但机器2的使用了batch2的, PipeDream通过多版本参数cache的思想来解决这个问题

为啥worker1需要保存4个版本参数, 而worker4只需要1个呢? 这里的版本数和同一个batch fp和bp的间隔决定的. 如果我跑完fp后, 中间有其他batch更新的bp. 那就需要把这些bp结果给缓存起来, 不然就会导致fp和bp使用的不是同一份参数. 可以看到worker1的batch5 中间间隔了2,3,4 3次bp, 再加上它本身. 就得保存4份...这种方法对显存极度不友好, 所以有了下面的flush方式
1F1B-flush#

对比上面的F-then-B的方式, 1F1B优先bp计算. 每个micro-batch完成后直接释放掉了对应micro-batch的计算中间值.
只需要保存1份w, 在固定micro-batch个数后进行一次flush, 同步所有worker的权重使其保持同一个版本.
另外在stage3中 batch1 fp时, 因为batch0已经算完了. 所以可以直接复用batch0的显存不用重新分配.
[!NOTE]
这里有个疑问..越底层的stage需要缓存的中间值其实越多, 这种造成存储不均匀的问题怎么解决? 通过stage切分不同大小参数的方式么
1F1B-flush(交错式, megatron)#
这个方案有个新的概念, virtual_pipeline, 方案要求一个小批次中的微批次数量是管道并行大小(流水线中的设备数量)的整数倍
按之前非交错式的方法. 一共有8层, worker1如果是1/2层, worker2是3/4层..worker4是7/8层, 每个worker计算连续的层
那么virtual_pipeline如果是2的话, 会把每个worker进一步拆分, worker1变成了计算1/5层, worker2: 2/6层..类推, 相当于通过把每个worker从单一流水线拆成了virtual_pipeline个流水线.
- 在之前的1F1B模式里, 因为每个机器计算是有先后顺序的, worker2的通信接收worker1的fp结果必须等worker1的fp完成.
- 而在交错式设计里, worker2计算的是2/6层, 当他计算2的时候, 可以同步从worker1拿上一个batch的5层结果, 算完2后的理想状态就是直接算5. 能更好的把通信隐藏起来.
总结这个方案的优点:
- 相邻的计算与通信操作无依赖关系, 可以加速并行执行
- 发起通信操作时,通信的对端通常已经准备好了要通信的数据,通信操作不需要额外的等待时间。
作者:sunstrikes
出处:https://www.cnblogs.com/sunstrikes/p/18270017
版权:本作品采用「署名-非商业性使用-相同方式共享 4.0 国际」许可协议进行许可。
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】博客园社区专享云产品让利特惠,阿里云新客6.5折上折
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 清华大学推出第四讲使用 DeepSeek + DeepResearch 让科研像聊天一样简单!
· 推荐几款开源且免费的 .NET MAUI 组件库
· 实操Deepseek接入个人知识库
· 易语言 —— 开山篇
· 【全网最全教程】使用最强DeepSeekR1+联网的火山引擎,没有生成长度限制,DeepSeek本体