为什么说大模型训练很难？ 第1页

给大家推荐一个基于大规模预训练模型的智能应用开发大赛，总奖金有120万元。

比赛任务是在超大规模预训练模型“悟道2.0”的基础上，结合社会热点主题——健康医疗、教育学习、社交生活、效率工具、环境自然等开发新颖实用的智能应用。

大赛开放了包括：文本生成、图片生成、写诗、新闻生成等9个“悟道”大模型API 接口，参赛者不仅能够免费使用这些功能，还能够直接下载“悟道”模型用于应用开发。

具体可以点击下方卡片了解：

由于上面 @ZOMI酱 讲的很细了，我这边说说我的感受，主要来自两方面：策略组合爆炸+工程量大。我用显存优化和分布式两方面来说明这个问题。

显存优化复杂

显存消耗估计

模型训练时消耗的显存，往往是参数量的几倍还多。不说特别大的，仅以100B模型为例，我们大致估算一下 不做任何显存优化 时，使用AdamOptimizer训练的固定消耗：

即使没有列出Activations的消耗，上面的部分应该也比较吓人了。这得堆多少GPU卡，多少台机器？

• 16GB显卡：75张GPU卡，约等于单机八卡服务器10台；
• 32GB显卡：38张GPU显卡，约等于单机八卡服务器，5台。

也就是说，即使抠抠索索的使用，也得消耗这么多资源，显存优化不做不行。于是，大量的显存优化工作就必须得上了。

显存优化方法

1. 合理的计算调度：如DAPPLE中提到的early backward；
2. Gradient Checkpointing：习惯上称之为重算（Re-computation）
3. CPU-offload：ZeRO中提到的借用Host memory做辅助变量的offload
4. Optimization-state优化：切分辅助变量分片存储

上面的有些优化，在目前流行的一些框架上，是需要手动做的。目前也有一些框架在自动化托管这些工作，但是上面这些优化组合，全自动化方案能否做到性能最优，并不容易。

这是第一个难点。

分布式复杂

即使做了显存优化，单卡也绝对不够，分布式方案是必然选择。但这里问题在于：Dense巨无霸模型很难用Naive Data Parallelism训练（单卡放不下一份完整的replica），迫使必须叠加各种各样的并行模式，所以只能是所谓的Hybrid Parallelism。例如可以是如下这种组合方案：

数据并行 + 算子拆分 + 流水并行

大家对数据并行比较熟悉，所以下面只展开后面两项、

算子拆分

算子拆分并不是什么新鲜想法。我用矩阵乘法举个最简单的例子。例如我们有如下变量和矩阵运算：

• variable矩阵：
• batch input矩阵 ：
• 输出： 。

我们可以将 拆分到两个device上进行分块计算，因为我们知道矩阵可以使用分块乘法，即：

所以，单个矩阵乘法可以分到两个device上计算。我们在工程上要做的就是：将 切分到两个device上，将 复制到两个device上，然后两个device分别做矩阵乘法即可。有的时候，切分会带来额外的通信，比如矩阵乘法切到了reduction维度上，为了保持语义正确，就必须再紧跟一个AllReduce通信。

在Megatron-LM中，对Attention Block就做了这样的算子拆分（拆K，Q，V的head维度），像下面这样。

这里复杂之处在于，你不能无脑地将所有算子全部做拆分，因为拆分可能会引入一些额外通信，降低总体吞吐。所以你得做些分析，决定哪些算子被拆分，这得考虑以下因素：

• Workload的特点：计算量、参数量、Tensor shape等；
• 硬件拓扑：GPU算力，通信带宽等。

写个程序自动化行不行？能不能让程序自动搜索出来最佳的拆分方案？这取决于你在什么层次上做搜索。以TensorFlow为例，如果在GraphDef上的Operator这个粒度上做，搜索空间肯定爆炸。如果在Layer这个粗粒度上搜索，则是有可能的。但自动化方案也有一定的工作量：

1. 拆分规则定制：为每种算子都需要定制一个完备的推导方案；
2. 推导传播算法：推导影响链以及通信的引入（比如引入Reshard）；
3. Cost model：快速评估每种方案的性能；
4. 高效地搜索算法：在巨大的空间里快速搜索出较好的方案。

现在大部分框架都不支持这种全自动化策略，要么是半自动或纯手工，要么是针对某种模型把它的拆分方案写死。所以只能造轮子解决这个事。

流水并行

不切算子，而是将不同的Layer切分到不同的Device上，就可以形成Pipeline方案。GPipe就是这样一种方案，提出了将一个batch拆分成若干个micro-batch，依次推入到Pipeline系统中，即可消除Bubble time（打时间差：Device 1计算micro-batch 1时，Device 0为空闲状态，正好推入下一个micro-batch）。下面这个图应该能看明白，就不过多解释了。

然而，这种方案有显存压力。上图(a)中看到，B0的计算依赖于B1和F0，也就是说F0计算完成后，其Activations不能立即释放，直到B0算完才能丢掉。所以针对F0来说，推N个micro-batch时，我们要保存N份这样的Activations，直到他们的反向计算完成才能释放，显存峰值是很大的。

因此我们又要做显存优化了，我们在DAPPLE中提出了Early Backward方案，通过合理调度计算顺序，及早释放Activations，如下图所示。

和算子拆分类似，全自动化方案工作量不小，比如Pipeline怎么切，才能让通信量小，计算还能均匀，这需要一定的算法和工程量。

而且，在TensorFlow这样的框架上实现Early backward方案，除了让用户在编写模型时强行加上control flow严格控制，没有别的办法。大家自己可以试一试，用TF的Graph mode写这样的调度顺序，是有多么的噩梦，最好的解决方案可能还真就是造个轮子。

混合并行

如果把上述数据并行、算子拆分和流水并行混合起来，复杂度就更高了。

所以分布式，是第二个难点。

后记

以上是我总结的一些难点，除此之外还有一些其他系统方面的工程问题。训练大模型，尤其是在已有的框架中从0开始支持，上面列出的问题都是必须要解决的。否则就算给你足够的资源，你也不知道怎么把Workload跑起来。

所以这不仅仅是资源的事，是技术的事，也是轮子的事。