Training large transformer models is one of the most important computational
challenges of modern AI. In this paper, we show how to significantly accelerate
training of large transformer models by reducing activation recomputation.
Activation recomputation is commonly used to work around memory capacity
constraints. Rather than storing activations for backpropagation, they are
traditionally recomputed, which saves memory but adds redundant compute. In
this work, we show most of this redundant compute is unnecessary because we can
reduce memory consumption sufficiently without it. We present two novel yet
very simple techniques: sequence parallelism and selective activation
recomputation. In conjunction with tensor parallelism, these techniques almost
eliminate the need to recompute activations. We evaluate our approach on
language models up to one trillion parameters in scale and show that our method
reduces activation memory by 5x, while reducing execution time overhead from
activation recomputation by over 90%. For example, when training a 530B
parameter GPT-3 style model on 2240 NVIDIA A100 GPUs, we achieve a Model Flops
Utilization of 54.2%, which is 29% faster than the 42.1% we achieve using
recomputation. Our implementation will be available in both Megatron-LM and
NeMo-Megatron.

本文通过降低激活记录的方法展示了如何显著加速大型 Transformer 模型的训练，并介绍了两种新颖且非常简单的技术：序列并行和选择性激活重算，这些技术几乎消除了重新计算激活的需要，同时减少了超过 90% 的执行时间开销和 5 倍的激活存储器减少。