别让显卡闲着！聊聊多卡训练模型那些事儿，从入门到放弃（不是）

哎,说到训练AI模型，现在这玩意儿真是越来越“吃”硬件了，早些年一张高端卡还能玩得转，现在动不动就几百亿参数，感觉显卡就像个无底洞，多少算力都能吞下去，所以啊，身边不少朋友和同行都开始琢磨：单卡不够，多卡来凑呗！但这事儿吧，听着挺美，真动起手来，里头门道可多了，一不小心就从“性能飞跃”变成“麻烦集合”。

我自己也是折腾了好久,踩了无数坑，才慢慢摸出点眉目，今天就跟大伙儿唠唠，多卡训练到底是怎么一回事，有哪些路子可以走，以及——那些没人轻易告诉你的头疼瞬间。

咱得搞清楚为啥要多卡,很简单，模型太大，或者数据太多，一张卡装不下、算得慢，多卡的核心思想就是把计算任务拆开，分给不同的显卡同时干，干完再汇总一下，这就像搬砖，一个人搬太慢，找几个兄弟一起搬，效率理论上就上去了，但怎么个“拆”法，学问就大了。

目前主流的方案,大概可以分成两大派别：数据并行 和 模型并行。

数据并行，这是最常见、相对也最容易上手的路子，想象一下，你有一大锅数据（比如100万张图片），数据并行的做法是，把这一锅数据分成好几份（比如4份，每份25万张），然后复制完全相同的模型到每一张显卡上，每张卡就用自己分到的那份数据，独立地做一次前向传播和反向传播，计算出模型参数该怎么调整（梯度），关键来了：这些计算出来的“调整意见”（梯度）不能各改各的，得把所有卡的结果收集起来，求个平均，得到一个统一的“调整方案”，再用这个方案去更新每一张卡上的模型，这样，下一轮训练开始，所有卡上的模型又是一模一样的新版本了。

这个过程里,那个负责“收集和平均梯度”的角色至关重要，通常由CPU或者其中一张显卡来担任，我们叫它“参数服务器”，现在更流行的是All-Reduce这种集体通信操作，效率更高一些，PyTorch的DistributedDataParallel和TensorFlow的MirroredStrategy，本质上都是在帮你做数据并行的脏活累活，它的好处是实现起来相对直观，很多框架支持得好，只要你模型能完整塞进一张卡的内存，用这个方案扩增卡数来加速训练，效果通常比较线性。

如果你的模型巨大，大到一张显卡的显存放不下它哪怕一个副本，那数据并行就傻眼了，这时候就得请出模型并行了，这招更硬核，思路是把模型这个“大家伙”本身切成几段，不同的段放到不同的卡上，比如一个超深的神经网络，前10层放在卡1，中间10层放卡2，最后10层放卡3，数据则像流水线一样，先流过卡1算完，把结果传给卡2，卡2算完再传给卡3，得到最终输出，反向传播的时候，梯度再沿着反方向传回去。

模型并行听起来很酷,能解决大模型训练的根本问题，但它对通信的要求极高，因为卡和卡之间要频繁传递中间计算结果（激活值），如果切分不好或者通信太慢，大部分时间显卡可能都在“等”数据传来传去，反而造成了闲置，它通常不是提速的首选，而是不得已而为之的解决方案，多见于训练那种千亿参数级别的巨型模型。

除了这两种“纯种”方案，还有各种混合并行的玩法，比如把数据并行和模型并行结合起来，在多个机器、多个卡之间进行更复杂的切分，这就属于高手领域了，得对模型结构和分布式计算有很深的理解才能玩得转。

聊完了路子,再说点实在的“坑”，多卡训练，绝不是插上几张卡、改两行代码就能享受速度飙升的。

第一,通信开销是隐形杀手，卡和卡之间（尤其是不同PCIe插槽甚至不同机器之间）传输数据，是有延迟和带宽限制的，数据并行里梯度同步，模型并行里激活值传递，都会占用大量时间，如果模型本身计算量不大（不够“重”），通信开销可能直接把多卡带来的收益全吃掉，甚至变得更慢，这就好比，你雇了四个搬运工，但他们每次搬砖前都要开个半小时的会协调工作，那还不如一个人默默搬呢。

第二,显存利用不是1+1=2，你以为8张24G显存的卡，就能训练192G显存的模型？太天真了，每张卡上除了要放模型参数，还要放计算过程中的中间变量、优化器状态（比如Adam优化器要存的动量和方差，这玩意儿大小经常是参数的两倍），在数据并行下，这些状态是每张卡都存一份的，所以显存开销几乎是线性增长的，并没有凭空变出更多空间来放更大的模型，模型并行能解决模型放不下的问题，但又回到了通信开销的难题上。

第三,软件配置和调试的噩梦，设置多卡训练的环境，特别是跨机器的，简直就是和驱动程序、CUDA版本、深度学习框架的分布式模块、网络权限搏斗的过程，一个版本不对，一个端口没开，错误信息可能看得你云里雾里，调试也比单卡复杂得多，日志散落在各个进程，出错了可能只是其中一张卡有问题。

我的个人建议是：别盲目上多卡，先好好优化你的单卡代码，看看是不是数据加载有瓶颈，或者模型结构本身有优化空间，如果单卡确实跑满了，而且计算任务足够重，先从数据并行试起，这是最友好的起点，框架提供的工具已经很成熟了，能帮你屏蔽不少底层细节。

至于模型并行和更复杂的混合并行,除非你明确知道自己在做什么，或者要挑战的模型规模实在离谱，否则可以先放一放，那是一片更深的江湖，需要更多的耐心和专业知识去驾驭。

多卡训练是个强大的工具,但它不是魔法，它放大了算力，同时也放大了系统的复杂度和可能出问题的地方，抱着“折腾”和“学习”的心态去尝试，收获会比单纯追求速度更大，如果最后被折腾得够呛，暂时“放弃”，回去用单卡跑个小点的模型，也完全不丢人，毕竟，能让项目跑起来的方案，才是好方案，先动起来，比卡在硬件门槛上更重要，对吧？

（免费申请加入）AI工具导航网

相关标签： # ai模型训练多显卡方案