Transformer 中锋利不平衡原则及其在加速语言模型预训练中的应用

Transformer由多样化的构件组成，如嵌入层、归一化层、自注意力机制和逐点前馈网络。因此，理解这些构件之间的差异和相互作用至关重要。本文发现，这些构件之间存在明显的锋利不平衡，这在训练初期即显现，并且在整个训练过程中保持不变。针对这一发现，我们提出块级学习率（Blockwise LR）策略，为每个构件的锋利度量身定制学习率，从而加速大型语言模型（LLM）的预训练。通过将Blockwise LR集成到AdamW中，我们一致实现更低的终端损失和近乎$2\times$的加速速度。我们在GPT-2和LLaMA上验证了这一加速，模型规模从0.12B到2B不等，数据集包括OpenWebText、MiniPile和C4。最后，我们将Blockwise LR集成到Adam-mini（Zhang等，2024）中，该方法是Adam的最近提出的一种内存高效变体，实现了$2\times$的加速和$2\times$的内存节省。这些结果凸显了利用锋利不平衡改进LLM训练的潜力。