基本信息

训练大型 Transformer 模型是现代人工智能最重要的计算挑战之一。在本文中，我们展示了如何通过减少激活重新计算来显着加速大型 Transformer 模型的训练。激活重新计算通常用于解决内存容量限制。传统上，它们不是存储反向传播的激活，而是重新计算，这节省了内存，但增加了冗余计算。在这项工作中，我们证明了大部分冗余计算是不必要的，因为我们可以在没有它的情况下充分减少内存消耗。我们提出了两种新颖但非常简单的技术：序列并行性和选择性激活重新计算。与张量并行相结合，这些技术几乎消除了重新计算激活的需要。我们在参数规模高达一万亿的语言模型上评估了我们的方法，结果表明我们的方法将激活内存减少了 5 倍，同时将激活重新计算的执行时间开销减少了 90% 以上。例如，在 2240 个 NVIDIA A100 GPU 上训练 530B 参数 GPT-3 样式模型时，我们实现了 54.2% 的模型触发器利用率，这比使用重新计算实现的 42.1% 快了 29%。我们的实施将在 Megatron-LM 和 NeMo-Megatron 中提供 ...