手把手教你用多张显卡并行训练大模型

为什么我们需要多卡训练？

哎呀，现在的大模型动不动就几百亿参数，一张显卡根本装不下啊！这就好比你想搬个大衣柜上楼，一个人搬不动，只能多找几个帮手一起抬。多卡训练就是这个道理，把模型或者数据拆分到多张显卡上，大家一起干活，效率自然就上去了。

我记得刚开始接触深度学习的时候，用一张GTX 1080就能跑很多模型了。但现在不行了，随便一个语言模型都是几十个G，单卡训练简直是痴人说梦。所以啊，学会多卡训练已经成为每个深度学习工程师的必备技能了。

多卡训练的两种主要方式

多卡训练主要有两种路子，一种是数据并行，另一种是模型并行。咱们先来搞清楚它们到底有什么区别。

• 数据并行：这个最简单理解，就像复印试卷一样。每张显卡上都放一份完整的模型，然后把训练数据分成几份，每张卡处理一份数据，最后把梯度汇总一下更新模型。
• 模型并行：这个就复杂点了，是把模型本身拆开，不同的层放在不同的显卡上。好比造汽车，你在A车间造发动机，在B车间造底盘，最后组装起来。

这两种方法各有各的用武之地，具体用哪个得看你的模型到底有多大。

数据并行的具体实现方法

数据并行是目前最常用的方法，PyTorch里面用DataParallel或者DistributedDataParallel就能搞定。不过我要提醒你，DataParallel虽然用起来简单，但性能不如DistributedDataParallel，特别是在多机多卡的情况下。

用DistributedDataParallel的基本步骤是这样的：

首先初始化进程组，然后给每个进程分配对应的GPU，接着把模型包装一下，最后用DistributedSampler来分配数据。这套流程虽然稍微复杂点，但效果确实好。

我建议你直接从DistributedDataParallel开始学，虽然门槛高一点，但学会了就能应对各种复杂场景了。

模型并行的适用场景和实现

模型并行听起来很美好，但实际上用得比较少，主要是因为它对模型结构有要求，而且通信开销比较大。不过当你遇到特别大的模型，单张卡连一层都放不下的时候，就只能用模型并行了。

比如说，有一个100层的Transformer模型，你可以把前50层放在第一张卡上，后50层放在第二张卡上。前向传播的时候，数据要从第一张卡传到第二张卡，反向传播的时候梯度又要传回来。

混合并行策略：强强联合

现在最厉害的玩法是把数据并行和模型并行结合起来用，这就是混合并行。比如说你有8张卡，可以先把模型分成4份，每2张卡用模型并行组成一个小组，然后4个小组之间用数据并行。

这种玩法在训练千亿级别参数的大模型时特别有用。不过说实话，配置起来确实麻烦，需要对模型结构和分布式训练都有很深的理解。

实际操作中遇到的坑

多卡训练说起来简单，做起来可是坑不少。我把自己踩过的坑都告诉你，希望能帮你少走点弯路。

• 显存不均衡：有时候某张卡的显存用得特别多，其他卡却很空闲，这时候就要调整模型拆分策略了。
• 通信瓶颈：如果卡之间的通信成了瓶颈，训练速度反而会变慢，这时候要考虑用更快的互联方式，比如NVLink。
• 调试困难：多进程的调试比单卡麻烦多了，经常是这个问题解决了，那个问题又冒出来了。

我最开始搞多卡训练的时候，光是解决各种莫名其妙的问题就花了一个多星期，那段时间真是头发都掉了一大把。

性能优化的小技巧

好不容易把多卡训练跑起来了，接下来就是要让它跑得更快。这里有几个实用的优化技巧：

梯度累积是个好东西，当显存不够的时候，可以多累积几个batch再更新，虽然训练会慢点，但总比跑不起来强。

激活检查点也能省很多显存，原理就是在前向传播的时候不保存中间结果，等到反向传播的时候重新计算。用时间换空间，很划算。

还有就是选择合适的batch size，不是越大越好，也不是越小越好，要找到那个甜点区域。这个需要你多试几次，观察loss曲线和显存使用情况。

总之啊，多卡训练是个实践出真知的活儿，光看教程是不够的，一定要亲手去试，去踩坑，这样才能真正掌握。希望我的这些经验对你有帮助！