本文介绍如何应对大模型训练时的内存挑战，通过管道并行技术将模型层分配到不同GPU处理。内容包括微批次处理、模型分割和管线管理，重点讨论了FairScale库中的实现细节，特别是利用'栅栏同步'确保数据在计算和传输间的正确同步。
当模型的大小对单个 GPU 来说过大时，训练大型模型可能会导致内存溢出。为了训练如此大型的模型，可以将各层在不同的 GPU 设备上用管道式并行(Pipeline Parallelism)技术处理。
所谓管道式并行，是按照下图所示，将模型不同的层放在不同的GPU上并行处理。相应的前馈传播和后向传播都需要放在同一个设备中计算。
另外批量的训练数据也划分成更小的批量，称为微批次（micro-batches），在当前设备处理完毕微批次后，把输出传送到下一个设备，同时自身接受下一个微批次，从而实现并行处理。
（图片来自FairScale文档）
FairScale的管道并行实现的源代码在这个位置：
https://github.com/facebookresearch/fairscale/tree/main/fairscale/nn/pipe
有以下源码可以关注：
1. 微批次（micro-batch) 代码：
最低0.47元/天 解锁文章
扫一扫
1298
到【灌水乐园】发言
