AI模型的规模在过去4年维持了每年10倍的增长。2021年的产生的大模型已经达到1 trillion parameters。而最近分布式训练号称可以支撑100trillion parameters。似乎这个趋势还能继续维持两年以上。

　　更大的模型以一种暴力美学的形式，进一步提升了模型的效果，同时也展现出一些很有趣的特性，比如解决多种不同类型的任务（zero-shot, one-shot）。训练巨大的模型必然需要底层基础软件和芯片的支撑。然而GPU在过去4年中，无论是显存空间，或者是算力的增长，都在10倍这个数量级，显然跟不上10000倍的模型规模增长。硬件不够，软件来凑。深度学习框架的分布式训练技术强势的支撑起了模型暴涨性的增长。这篇文章总结了过去几年支撑模型规模上涨的关键分布式训练技术。总体来说，核心解决的问题是如何在保障GPU能够高效计算的同时，降低显存的开销。关键可以总结在几个方面的技术上：

　　DP关注的问题是大规模Batch Size下，如何降低显存的开销。我们知道，模型在forward和backward的中间计算过程都会有中间状态，这些中间状态通常占用的空间是和batch size成正比的。通过将大的batch size切分到多个n个GPU上，每个GPU上部分状态的空间开销能降到1/n，而省下的空间可以用来保存模型的parameter，optimizer state等，因此可以提高GPU上容纳的模型规模。DP是一种非常简单易用的分布式训练方法，几乎所有的训练框架都支持这种方法。早期比较流行的开源实现是horovod。现在pytorch ddp和tensorflow mirror strategy都原生的支持了。DP的极限在于两个方面：

　　关于的DP的极限提到，当batch size=1时，DP难以进一步的提高模型规模的上限。另外，传统DP仅仅是从数据的层面切分，降低显存空间压力。而我们知道每个GPU上依然需要一份完整的parameter等副本。MP基本含义在于把模型本身进行切分，使得每个GPU卡只需要存模型的一部分。多个GPU配合起来完成一个完整的minibatch。按这个比较宽泛的定义，MP的使用方式就比较的多样，EBET易博官方网站比如：

　　对一个算子进行拆分，比如FC，把参数和计算的切分到多个GPU上，通过通信完成这个原子计算。

　　单纯的把模型的参数切分到多个GPU上，在使用时通过数据驱动的方式，每个GPU从其他GPU上拉取需要的那部分。比如大embedding参数如果有100GB，可以切分成8份放在8个GPU上。每个minibatch计算时embedding层仅需要gather 100GB中很少的一部分。

　　简单的把模型横向切分几份，分别放在不同GPU上。比如一个1000层的ResNet，可以每个GPU放其中100层的参数。当一个minibatch进来后，依次通过这10个GPU，完成这1000层的计算。

　　TensorFlow可以说是支持MP的典型框架，通过将device placement暴露给用户，开放了几乎所有的玩法。但是弊端就是大部分用户其实并不能合理的使用，达到一个比较好的效果。反而是经常因为配置错误导致性能急剧下降。导致TensorFlow的开发团队不得不投入许多时间帮用户调优。MP的玩法很多，可以在DP的基础上显著降低单个GPU上的显存压力。但是MP的问题也很多，使用起来远比DP复杂：

　　模型切分的可选项太多，次优解会导致较高的通信开销，显著降低硬件的利用率和训练速度。记得有个同事发了篇论文，用1000个GPU，基于RL算法搜索了一个星期，找到了NMT模型上一个超越所有人工的device placement。（仅适用于特定机房，网络，设备，模型等等的组合）

　　MP在一些应用场景（比如上面说的ResNet例子）在计算一个minibatch时，硬件是依次激活的，其他硬件都在等待，硬件的利用率会非常的低。

　　前面提到了MP一个比较大的问题是GPU利用率低。当没有计算到某个GPU上的模型分片时，这个GPU常常是闲着的。PP一定程度上解决了这个问题。PP的思想也比较简单，使用了经典的Pipeline思想。在模型计算流水线上，每个GPU只负责模型的一个分片，计算完就交给下一个GPU完成下一个模型分片的计算。当下个GPU在计算时，上一个GPU开始算下一个minibatch属于它的模型分片。

　　这里需要注意到不同模型分片使用的parameter分片是否同步的问题（防止流水线更新parameter的问题）。从数学上，一个minibatch数据训练使用的所有模型参数都需要是同步的(不完全一定）。因此，在optimize时，整个PP需要进行某种形式的同步。PP依然没有完全解决部分硬件时间空闲的问题。同时MP使用的复杂性和调优依然是个问题。

　　在每个minibatch的计算过程中，forward中的一些中间结果会被保存在显存里，等到backward的时候会被读取出来用于计算gradient。对于一些比较宽和深的模型，这部分forward的空间的开销是非常显著的。同样输入下，每次forward的中间结果都是一样的（dropout is exception)。Checkpoint利用了这个特点，是一种用时间换空间的技术。在forward的过程中，我们只选择记录其中一部分的中间结果，其他中间结果被扔掉，节约显存。当backward是如果需要这些已经被扔掉的中间结果时，它们可以基于没有被扔掉的上一个中间结果（checkpoint），被重新计算出来。有些模型中间有一些“小蛮腰”。比如一些sequeeze、encoder模型结构。或者是模型中间有些激活计算比较轻量，重新计算开销较小。基于这些情况，checkpoint很多时候可以大幅节省空间开销，同时额外的计算耗时增加较小。

　　前面提到DP的显存节省的限制，而许多MP又存在易用性差和性能的问题。Sharding技术比较好的解决了这些问题。一方面它可以基于DP叠加，同时我也把它理解成是一种比较易用的MP技术。Sharding分为3个层面空间开销的Sharding，使得可以通过增加GPU个数，等比降低每个GPU需要的显存。

　　通过这种方式，模型的规模可以极大的提高。而性能的问题，也有非常巧妙的解决办法。由于模型的结构通常是固定的，因此在forward阶段，可以提前并行预拉取parameter，掩盖这部分的通信开销。比较经典的就是Zero论文，已经开源的DeepSpeed实现。顺便提一下，过去几年，GPU之间的通信技术得到了比较突破性的提高，所以Sharding技术导致的通信开销往往不会非常显著，而分布式GPU训练集群的规模也因此可以做的更大。

　　其实到这里，感觉优化的已经非常极致了。然而还有更猛的。现在业界有两种超大规模模型：

　　对于这种模型，把所有参数放在GPU显存里面都显得非常的奢侈。特别是sparse embedding超大的模型，虽然总体embedding parameter规模很大，然而每个minibatch仅使用其中1%不到的一小部分。Offload的思想简单来说就是：

　　cpu计算开销于batch size无关。gpu的计算开销于batch size有关。在足够大的batch size下，CPU的计算显得不比GPU慢。

　　另外还有一些优化可以将当前optimize和下个minibatch gpu计算并行起来。虽然有参数过期的风险，实验发现通过预热的方式可以有效避免模型效果的下降。

　　本文总结了深度学习框架技术是如何在硬件条件限制下，通过多年的技术升级，让AI模型的规模突破天际的。如果我们从更高的视角看，AI模型的增长还有硬件，算法等领域的许多突破。整个生态正在飞速的往前发展。另外，一些瓶颈也显得越来越明显，这里列几点：

　　模型的无效计算。当前大部分DNN在训练或者推理时都需要完成整个模型的“激活”。而我们知道人类的大脑在思考时，往往只需要大脑的一部分参加工作。相比之下，当前AI模型的计算就显得非常低效。

　　模型的可解释性。巨大的模型虽然带来了意想不到的效果，但是也将AI模型不可解释问题推向了一个极致。我们几乎无法有效的判断模型的能力，模型的潜在风险，或是对模型进行有效的干预。

　　训练Pipeline的易用性。虽然很多技术被提出来，并实现。但是在真实场景中，受限于数据中心网络拓扑，分布式数据系统的能力，机型配置和软件栈，深度学习框架本身的调优等许多原因，很难有效的发挥理论的较优结果。