LLM 预训练加速的新方法：八种模型增长方案总结

一、背景

LLM 的涌现能力依赖于其模型规模的增长，而 Scaling Law 也在推进 LLM 朝着越来越大的方向发展。然而，LLM 预训练的成本非常高，尤其是其与模型规模、数据量成正比，一个千亿参数量的模型往往需要几千个 GPU 训练几个月的时间。加速 LLM 预训练也因此称为一个非常有前景的研究方向。

当前常见的优化方案为优化分布式策略，通信，以及训练稳定性等。与此同时，很多时候大家都会训练各种规模的 LLM，例如 LLaMA 系列模型。也有许多工作在尝试利用已经训练好的较小 LLM 的权重，然后通过渐进式学习的方法加快较大 LLM 模型的训练，比如使用训练好的 LLaMA3 8B 模型来初始化 LLaMA3 30B 模型并继续训练，这种技术可以称为模型增长（Model Growth）。

本文中，我们将总结一系列模型增长的方案，以便对模型增长的发展以及各种方案有一个更清晰的认识。具体来说，我们将分别介绍 Net2Net，StackBert，bert2BERT，LiGO，LEMON、MSG 和 Mango 几种方案。

PS：其实当前很多 MoE 模型的训练正是属于模型增长的范畴。比如说，Mixtral-8x7B（ Mixtral of experts | Mistral AI | Frontier AI in your hands） 是从 Mistral-7B 初始化而来，专家直接拷贝了 Mistral-7B 中的 FFN。此外，也可以使用 Mistral-7B 针对不同任务微调 FFN 后再来作为专家。

相关工作可以参考：

• ​​​Meta 发布 BTX：模型融合与 MoE 的结合​​​
• ​​​MoE 系列论文解读：Gshard、FastMoE、Tutel、MegaBlocks 等​​​
• ​​​7 种 LLM 融合扩展总结：LLaMA-Pro、SOLAR、MoE 等​​​

二、Net2Net

2.1 摘要

Net2Net（[1511.05641] Net2Net: Accelerating Learning via Knowledge Transfer）是模型增长领域的开创性工作，作者是陈天琦和 GAN 的作者 Ian Goodfellow 等。作者提出了基于存在的小模型来初始化大模型，并继续训练的方案。其主要包含宽度（widthwise）扩展和深度（depthwise）扩展两个方面。如下图 Figure 1 所示为 Net2Net 与传统模型训练的差异：

• 传统 Workflow：训练小模型和大模型时分别设计模型，并从 0 开始训练。
• Net2Net Workflow：训练小模型从 0 开始，训练大模型时用已经训练好的小模型来初始化大模型，然后继续训练。​

Function Preserving Initialization（FPI）：目标是给定一个源模型，用它初始化目标模型，能保证给定相同的输入，目标模型和源模型有相同的输出。

2.2 宽度扩展（Net2WiderNet）

宽度扩展，主要是指在模型的层内扩展，增加层的宽度。如下图所示：

• 左图为原始网络，输入包含 x[1] 和 x[2]，输出为 y，其包含两个线性的隐藏单元 h[1] 和 h[2]。
• 右图为扩展之后的网络，新增一个隐藏单元 h[3]，其对应的权重直接拷贝 h[2]，这样 h[3] 和 h[2] 就是相等的。同时 y 中新增一个 h[3] 的权重，同时 h[2] 和 h[3] 的权重都变为 f/2。最终 y = e*h[1] + f*h[2] = e*h[1] + f/2*h[2] + f/2*h[3]。​

2.3 深度扩展（Net2DeeperNet）

深度扩展，主要是指增加模型的层数。如下图所示：其主要是在网络中插入初始化的 Identity Mapping 层，其输入输出是相同的，保证插入后依然是等价的。

三、StackingBert

3.1 StackingBert-v1

在 Efficient Training of BERT by Progressively Stacking 中，作者提出了通过堆叠 Bert 模型的 Transformer Layer 来扩展模型规模的方法。Transformer 模型的 Decoder 和 Encoder 中不同层之间结构完全一样，输入和输出 Shape 相同，这就为通过 Copy 来复制 Layer 提供了很大的遍历。如下图 Figure 3 所示，假设已经训练好 L 层 Encoder 的 Bert 模型，直接通过拷贝即可扩展为 2L 层，具体来说，第 0 层和 第 L 层完全一样，第 i 层和第 L+i 层完全一样。然后继续训练即可以得到 2L 层的 Bert 模型：

3.2 StackingBert-v2(MSLT)

在 [2011.13635] Progressively Stacking 2.0: A Multi-stage Layerwise Training Method for BERT Training Speedup 中，作者进一步对 StackingBert-v1 进行了扩展。具体来说，将一个 N 层 Encoder 的模型分 K+1 次训练，第一次训练一个 N/k 层的 Bert 模型，然后每次扩展 N/k 层并且进行训练。其中绿色为冻结的层，红色为训练的层。也就是每次扩展后只训练新扩展的层，全部扩展完之后再解冻所有层继续训练：

四、bert2BERT

4.1 摘要

在上述 StackingBert 的两个版本中，都是在深度上扩展模型规模。然而，实际上 Bert Large 相比 Bert Base 除了深度更深外，每一层也更宽，如下图所示，其 hidden size 从 768 扩展到 1024，相应的 head 数也有所增加。

在 [2110.07143] bert2BERT: Towards Reusable Pretrained Language Models 中，作者提出 bert2BERT，其同样是为了利用已有的小的预训练模型来加快更大的预训练模型的速度，降低训练成本。作者测试发现，对于 Bert Base 和 GPT base，通过重用一半左右大小的模型，可以节约 45% 和 47% 的计算成本。

4.2 FPI 矩阵扩展

在模型宽度上，其最主要的就是矩阵乘法计算，大部分的模型参数都是一个权重矩阵，对应 h=W*x。因此，最常见的就是矩阵扩展，如下图 Figure 3 所示，一般分为两步：

• 第一步：输入扩展，新增 x3，相应的参数矩阵要增加一列。
• 第二步：输出扩展，新增 h3，相应的参数矩阵要增加一行。​

如下图 Figure 4 所示为一个满足 FPI 的模型宽度扩展示例，其分为三步：

• 第一步：输入扩展，在x1, x2 的基础上新增 x3，不过 x3 是直接 copy 的 x1。因此，直接在权重矩阵上扩展一列，并将第一列的权重和第三列均分。此时 h1 = o/2*x1+p*x2+o/2*x1 = o*x1+p*x2；h2 同理，因此可以保证 FPI。
• 第二步：输出扩展，在h1, h2 的基础上新增 h3，不过 h3 是直接 copy 的 h2。此时，需要在权重矩阵扩展一行，新的第三行直接复制第二行即可保证 h3=h2。
• 第三步：上述扩展了第一层的输出，也就是扩展了第二层的输入，要保持输出结果不变，采用第一步的方法在第二层的权重矩阵扩展一列即可。​

4.3 AKI 矩阵扩展

为了加快训练的收敛速度，作者进一步提出了 Advanced Knowledge Initialization（AKI），其不仅考虑当前层的参数，也考虑下一层的参数。这样做是因为之前其它的工作中发现相邻的 Transformer 层比较相似，因此综合考虑两个相邻层并不会对模型效果有太大影响。当然，这也会打破 FPI 约束。

如下图 Table 2 所示，作者通过实验对比了不同方案的效果，首先 AKI 会比 FPI 收敛更快，计算资源更少。此外，作者进一步加上了两阶段预训练，也就是 bert2BERT，训练成本节约 45.2%，明显优于 StackBERT 和 MSLT：

五、LiGO

5.1 摘要

在 [2303.00980] Learning to Grow Pretrained Models for Efficient Transformer Training 中，作者提出了通过线性投影来使用小模型初始化大模型的方案。具体来说，作者将线性变换分解为线性宽度增长算子和线性深度增长算子的组合，并进一步采用这些增长算子的 Kronecker 分解来嵌入结构知识。在语言模型和视觉 Transformer 模型上的实验表明，采用这些线性增长算子（Linear Growth Operator，LiGO）可以节约高达 50% 的训练成本，优于之前的方法。

5.2 方法

LiGO 的实现方式如下图 Figure 1 所示，作者定义了一个线性映射函数 M，其可以将小模型的参数 ϴ 映射为大模型参数 ϴnew。由于直接学习 M 的代价非常非常高，因此将其分解为宽度增长算子 Rwidth 和深度增长算子 Ldepth 的组合。为了减少可训练参数的数量并嵌入结构知识，进一步通过 Kronecker 积来分解 Rwidth 和 Ldepth，这样每个增长算子都可以表示为较小矩阵的 Kronecker 积。

Kronecker 积：假设矩阵 A 的大小为 m×n，矩阵 B 的大小为 p×q。那么 A 和 B 的 Kronecker 积 A⨂B 的大小将是 (m⋅p)×(n⋅q)。它的元素由以下方式确定：

5.3 实验和结果

其训练过程分为 3 步：

• 为了训练大模型，首先需要学习线性映射 M。作者通过 100 次的梯度迭代来优化 M，这个过程相对于整个训练来说代价很小。
• 然后使用 M 将小模型的参数 ϴ 映射为大模型参数 ϴnew。
• 接着使用常规的训练方式来训练大模型。

作者在 Bert 模型上进行了相关实验，如下图 Figure 2 和 Table 1 所示，LiGO 相比从头训练可以节约 40.7% 的训练成本。（PS：比较奇怪的是，作者测试 bert2BERT 的收益反而比 StackBERT 和 MSLT 更低，这与 bert2BERT 论文不太相符。）

六、LEMON

6.1 摘要

在 [2310.07999] LEMON: Lossless model expansion 中作者提出了 LosslEss MOdel expansioN（LEMON），对之前的模型扩展方案进行了增强，主要聚焦在无损扩展。其在 Vision Transformer 模型上可以减少 56.7% 的训练代价，BERT 上可以减少 33.2%。如下图 Table 1 所示为 LEMON 与其它方法的对比：

6.2 方法

作者提出的 LEMON 方法包含三个基本组件：

• 非均匀的通用无损宽度扩展。
• 针对 LayerNorm 的平均宽度扩展。
• 无损深度扩展。

6.2.1 无损宽度扩展

主要包含两个部分，一个是 MLP 扩展，一个是 MHA 扩展：

• MLP 扩展：如下图 3(a) 所示，其与之前宽度扩展的主要不同就是：之前的扩展中一般 a=β=1/2，其复制出来的神经元和原始神经元提取完全相同的表征，而在这里的 a≠β。
• MHA 扩展：如下图 3(b) 所示，在 Transformer 模型中，模型变宽通常也意味着 MHA 中 Head 的增加，这里作者采用直接拷贝整个 Head 的方案。​

6.2.2 平均宽度扩展

如下图 Figure 4 所示，LayerNorm 的平均宽度扩展就是对于新增的行直接使用之前行的平均，同时对 LayerNorm 中的 μ 增加缩放因子。以此可以保证最终的分布不变，满足 FPI。

具体的证明如下图所示，其要点是通过扩展均值可以保证扩展后均值不变，方差有一个固定的缩放因子，以此可以保证扩展后的输出位置还为 0，扩展前的位置通过缩放因子也可以还原：

如下图所示为 LayerNorm 扩展和 MHA 扩展的结合：

6.2.3 无损深度扩展

深度扩展基本上都是通过堆叠 Layer 的方式实现，要想保证堆叠后的无损，需要保证堆叠的层的输入和输出一样，也就是等效于 Identify Layer。幸运的是，Transformer Layer 中的 MHA 和 MLP 层都有残差连接，如下图 Figure 2 中的红框所示，因此只要保证新增层中 MHA 和 MLP 的输出为 0 就可以保证无损。

具体的实现方式有两种：

• 对应下图 Figure 5(b) ，直接将最后一个全连接层置为全 0 即可。
• 对应下图 Figure 5(c)，同样针对最后一个全连接层处理，将对应同一个神经元的权重设置为相反值，保证累积后和为 0。​

6.3 实验&结果

如下图 Figure 7 所示，作者与之前的模型扩展及模型蒸馏方案进行了对比，本文提出的 LEMON 在 ViT、Bert 模型上都收敛更快：

七、MSG

7.1 摘要

在 [2305.02869] Masked Structural Growth for 2x Faster Language Model Pre-training 中，智源（BAAI）等作者将渐进式增长分为两个方面：

• 确定最优的增长规划（Growth Schedule）：主要是探讨不同维度（比如深度、宽度）对增长效率的影响。
• 设计高效的增长算子（Growth Operator）：当前的方法主要依赖新权重的初始化来继承知识，并且很多只实现了非 FPI 的方案，从而限制了训练的进一步提升。

为了解决以上问题，作者提出了 Masked Structural Growth（MSG），包括：

• 涉及所有可能维度的增长规划。
• 与新权重初始化无关的严格满足 FPI 的增长算子。

如下图 Table 1 所示为 MSG 与其他方案的对比，可见 MSG 支持更多的维度，都满足 FPI，可以获得高达 2.2x 的加速：

7.2 增长算子

如下图 Figure 1 所示，之前的各种方案（比如 Net2Net）都是通过特殊的初始化方式来保证尽量满足 FPI。而本文的 MSG 的核心思路就是不管是什么样的初始化，都通过 Mask 的方式让新增的部分为 0，以满足 FPI。同时，在训练中逐渐增大 Mask，直到其 mask=1，此时就可以删除 Mask。针对 MLP，LayerNorm，MHA 以及残差连接都可以通过 Mask 方式实现。

7.3 增长规划

对于 Transformer 模型，其决定模型规模的超参数主要有 4 个：hedden_dim, ffn_dim, head_num, layer_num。除了 layer_num 为深度扩展外，其它 3 个都是宽度扩展。在模型扩展过程中，可以一次扩展所有参数，也可以逐个扩展，但是又会存在扩展顺序的问题，因此要找到一个最优的扩展方案是一个很有挑战的工作。如下图 Table 2 所示，作者针对 Bert 和 GPT-2 模型制定了不同的扩展规划，并通过实验验证了各自的影响：

7.4 实验

7.4.1 bert2BERT 对比

作者首先与 bert2BERT，以及从头训练进行了效果和速度的对比，可以看出，MSG 基本实现了效果和速度的最优，其最多可以实现 2.2x 加速：

7.4.2 LLM 预训练对比

作者也进一步验证了在 LLM 预训练上的效果，具体来说，作者验证了 LLM 从 16B 扩展到 51B 再扩展到 101B 的方案。其使用 24 DGX-A800（8x80G）机器，共192 A800，先在 16B 规模训练 245.37B Token，然后在 51B 规模训练 39.64B Token，最后在 101B 规模训练 26.54B Token，总共训练了 21.54 天，花费 100K 美元。如下图 Figure 4 所示为其训练的 Loss 曲线，作者与相似数据规模的 GLM-130B 模型对比，其只使用 10% 的训练成本即可以达到 80% 的性能（具体可以参考作者的论文 [2309.03852] FLM-101B: An Open LLM and How to Train It with $100K Budget）：

八、Mango

8.1 摘要

在 [2310.10699] Reusing Pretrained Models by Multi-linear Operators for Efficient Training 中，作者肯定了 bert2BERT 和 LiGO 中通过小模型来初始化大模型的方法，但也提出这些方法可能只考虑了局部相关性，而忽略了整个模型的相关性，这种部分映射的方法可能限制扩展模型的加速能力。因此，本文中，作者提出了一种将目标模型的每个权重与源模型的所有权重线性关联的方案，并利用多线性算子（Multi-Linear Operator，Mango）来降低计算和空间复杂度。实验结果表明，从 DeiT-small 到 DeiT-base，可以节省 76% 的计算成本，比 bert2BERT 和 LiGO 分别高出 12.0% 和 20.7%。

8.2 方案

本文的方案概览如下图 Figure 4 所示：

• 左图：Transformer Layer 中的参数表示。
• 右上：一个 Transformer 模型中的全部参数可以用一个大的 Tensor M 表示，并由 B、I、O、L 这 4 个超参数决定。
• 右下：本文的方案，也就是 Mango 算子，可以学习一个线性映射函数 S，将小模型权重 M1 映射为大模型权重 M2。显然，S 的空间极大，因此作者使用张量环矩阵乘法算子（Tensor Ring Matrix Product Operator，TR-MPO）来降低计算和空间复杂度，将其分为 4 个较小的张量，并通过 Rank 连接，通过训练来学习到 4 个小的张量后就可以用于构建 M2：

• SB：表示在同一层中参数之间相互作用。
• SI和SO：分别表示输入和输出维度上的转换。
• SL：表示层与层之间的关系。
• R：表示 S 的低秩级别。​

8.3 实验和结果

如下图 Table 1 所示，作者将 Mango 与之前的 bert2BERT 和 LiGO 的复杂度、能力进行了对比：

如下图 Figure 7 所示，作者与 StackBERT、bert2BERT 以及 LiGO 的训练效果进行了对比，可以看出，本文提出的 Mango 收敛更快，在 DeiT 上可以节约 76.4% 的成本：

九、参考链接

1. ​​​https://mistral.ai/news/mixtral-of-experts/​​​
2. ​​​https://arxiv.org/abs/1511.05641​​​
3. ​​​https://proceedings.mlr.press/v97/gong19a.html​​​
4. ​​​https://arxiv.org/abs/2011.13635​​​
5. ​​​https://arxiv.org/abs/2110.07143​​​
6. ​​​https://arxiv.org/abs/2303.00980​​​
7. ​​​https://arxiv.org/abs/2310.07999​​​
8. ​​​https://arxiv.org/abs/2305.02869​​​
9. ​​​https://arxiv.org/abs/2309.03852​​​
10. ​​​https://arxiv.org/abs/2310.10699​​​​

本文转载自​​AI闲谈​​，作者： AI闲谈 ​​​​