当模型的大小对单个 GPU 来说过大时,训练大型模型可能会导致内存溢出。为了训练如此大型的模型,可以将各层在不同的 GPU 设备上用管道式并行(Pipeline Parallelism)技术处理。
所谓管道式并行,是按照下图所示,将模型不同的层放在不同的GPU上并行处理。相应的前馈传播和后向传播都需要放在同一个设备中计算。
另外批量的训练数据也划分成更小的批量,称为微批次(micro-batches),在当前设备处理完毕微批次后,把输出传送到下一个设备,同时自身接受下一个微批次,从而实现并行处理。
(图片来自FairScale文档)
FairScale的管道并行实现的源代码在这个位置:
https://github.com/facebookresearch/fairscale/tree/main/fairscale/nn/pipe
有以下源码可以关注:
1. 微批次(micro-batch) 代码: