[論文レビュー] Reducing Activation Recomputation in Large Transformer Models

Vijay Anand Korthikanti, Jared Casper+5 more|arXiv (Cornell University)|May 10, 2022

Parallel Computing and Optimization Techniques被引用数 52

ひとこと要約

この論文は、シーケンス並列性と選択的活性化再計算を導入し、活性化メモリと再計算オーバーヘッドを劇的に削減し、少ないメモリで非常に大規模なトランスフォーマーの学習を可能にし、スループットを約30%向上させる。

ABSTRACT

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.

研究の動機と目的

• 大規模トランスフォーマーの学習における活性化メモリのボトルネックを動機づける。
• 大きな計算オーバーヘッドを伴わずに活性化メモリを削減する方法を開発する。
• テンソル並列性とシーケンス並列性を組み合わせて再計算を最小化する。
• 選択的活性化再計算を導入してさらなるメモリとオーバーヘッドを削減する。
• 1兆パラメータ近いモデルまでの実証評価を提供する。

提案手法

• メモリニーズを定量化するために、トランスフォーマーレイヤの活性化メモリ推定を導出する（式1）。
• 注意機構とMLPブロック内の活性化を分割するためにテンソル並列性を適用する（式2）。
• テンソル並列でない領域をシーケンス次元に沿って分割するためにシーケンス並列性を導入する（式4へ）。
• テンソル並列性とシーケンス並列性をパイプライン並列性と組み合わせ、追加通信を最小化するコンバータ g および g を導出する（図5および図6）。
• 高い活性化を生じさせる領域（例：注意関連演算）を対象とする選択的活性化再計算を提案し、FLOPオーバーヘッドを控えめに抑えつつメモリを削減する（式6）。
• Megatron-LM および NeMo-Megatron 実装を用いて、最大1Tパラメータのモデルに対するメモリと実行時間の影響を評価する。

実験結果

リサーチクエスチョン

• RQ1 標準的なメモリレイアウトで大規模トランスフォーマーモデルを学習する際、活性化メモリはどれくらい必要か？
• RQ2 ハイブリッドなテンソル/シーケンス/パイプライン並列性は、大きな計算・通信ペナルティを伴うことなく活性化メモリを削減できるか？
• RQ3 選択的活性化再計算は、FLOPオーバーヘッドを低く保ちながら、実質的にメモリを削減できるか？
• RQ4 これらの手法を適用した場合、 trillion-parameter transformers でどの程度のエンドツーエンドのスループット向上が達成できるか？

主な発見

• シーケンス並列性と選択的活性化再計算を併用すると、活性化メモリを約5倍削減できる。
• 選択的活性化再計算は大規模モデルで再計算FLOPsのオーバーヘッドを3%未満に低下させ、全再計算の30-40%と比較して低くなる。
• 全再計算と比較して、テスト済み構成全体でエンドツーエンドのイテレーションスループットが約29-32%向上。
• GPT-3 530B および MT-NLG では、本アプローチはトランスフォーマーレイヤあたりのメモリを大幅に削減し、重い再計算なしでのトレーニングを可能にする。
• 2240-A100 環境で 530B GPT-3スタイルモデルの場合、Model FLOPs Utilization は 54.2%、ベースラインの 42.1% より29%高速。
• メモリ削減はモデルサイズとともに拡大し、全体で約5xの削減を達成し、ベースラインの活性化がデバイスメモリを超えるトレーニングを可能にする。