[论文解读] TorchScale: Transformers at Scale

Shuming Ma, Hongyu Wang+9 more|arXiv (Cornell University)|Nov 23, 2022

Topic Modeling被引用 6

一句话总结

TorchScale 是一个开源的 PyTorch 工具包，通过集成先进的建模技术（如 Magneto 架构、DeepNorm 初始化和 X-MoE），并结合一种新颖的梯度裁剪方法（SparseClip），实现了 Transformer 模型在大规模下的高效且稳定的扩展。该工具包在语言建模和神经机器翻译任务中均实现了稳定的训练和性能提升，即使在 96 层深度和 256 专家 MoE 配置下也能保持稳定。

ABSTRACT

Large Transformers have achieved state-of-the-art performance across many tasks. Most open-source libraries on scaling Transformers focus on improving training or inference with better parallelization. In this work, we present TorchScale, an open-source toolkit that allows researchers and developers to scale up Transformers efficiently and effectively. TorchScale has the implementation of several modeling techniques, which can improve modeling generality and capability, as well as training stability and efficiency. Experimental results on language modeling and neural machine translation demonstrate that TorchScale can successfully scale Transformers to different sizes without tears. The library is available at https://aka.ms/torchscale.

研究动机与目标

• 解决将 Transformer 模型扩展至标准尺寸以上时面临的不稳定性和低效性挑战。
• 提升在深层和宽幅 Transformer 模型扩展过程中的训练稳定性，特别是针对稀疏 MoE 架构。
• 在不损害收敛性或性能的前提下，实现模型深度、宽度和专家数量的有效扩展。
• 提供一个统一的、通用的框架，支持包括语言建模、机器翻译和多模态预训练在内的多样化任务。
• 提出一种专为稀疏 MoE 模型设计的新颖梯度裁剪方法（SparseClip），以缓解梯度主导问题并提升训练稳定性。

提出的方法

• 采用 Magneto 架构作为默认主干网络，该架构使用 Sub-LN 初始化并具备统一的编码器-解码器结构，适用于通用建模。
• 基于 DeepNet 理论实现 Sub-LN 初始化，以增强深层 Transformer 中的优化稳定性。
• 将 X-MoE（eXtended Mixture of Experts）与 top-1 和 top-2 路由机制结合，通过稀疏性实现参数高效的扩展。
• 提出 SparseClip，一种重新加权的梯度裁剪方法，通过 $ \kappa = 1/\sqrt{E} $ 对 MoE 梯度进行缩放，以平衡密集层与 MoE 组件之间的梯度范数。
• 采用基于 L2 范数的梯度裁剪，其中范数计算方式经修改为：$ ||g||_2 = ||(g_d, \kappa g_s)||_2 $，其中 $ E $ 为专家数量。
• 在编码器、解码器及编码器-解码器架构中统一支持密集模型和稀疏 MoE 模型，具备一致的训练与推理流水线。

实验结果

研究问题

• RQ1 在不出现训练发散的情况下，Transformer 能否稳定扩展至 96 层和 256 个专家？
• RQ2 与原始梯度裁剪方法相比，所提出的 SparseClip 方法在稀疏 MoE 模型中如何提升训练稳定性和收敛性？
• RQ3 在实现高效扩展的同时，采用 top-2 路由的 X-MoE 架构在多大程度上保持了模型性能？
• RQ4 结合 Sub-LN 初始化的 Magneto 架构是否能在语言建模和机器翻译等多样化任务中实现更好的泛化能力和训练稳定性？
• RQ5 TorchScale 中先进建模技术的组合是否能在不同模型尺寸和任务中实现一致的性能提升？

主要发现

• TorchScale 在语言建模任务上成功训练了深度达 96 层的解码器模型，实现了稳定的收敛，且验证困惑度（PPL）随深度增加而持续改善。
• 采用 top-2 路由的 256 专家 X-MoE 模型尽管参数量超过 100 倍于较小模型，仍实现了更低的 PPL，证明了其高效的扩展能力。
• SparseClip 显著提升了稀疏 MoE 模型的训练稳定性和收敛性，与原始梯度裁剪相比，PPL 显著降低。
• 使用 SparseClip 后，梯度裁剪触发次数随训练时间推移而减少，表明梯度更新更加一致且准确。
• 在 SparseClip 下，路由熵更低，表明专家利用更充分，训练过程中路由行为更加稳定。
• 随着专家数量的增加，训练 FLOPs 几乎保持不变，证实了 TorchScale 中稀疏 MoE 方法的高效性。