## Memo ## Memo with LLM ### 論文情報 - **論文タイトル**: Reducing Activation Recomputation in Large Transformer Models - **著者と所属**: Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, Bryan Catanzaro (NVIDIA) - **カンファレンス名**: MLSys 2023 (Proceedings of the 6th Conference on Machine Learning and Systems) - **発表年**: 2023年 ### 論文概要 大規模Transformerモデルの訓練において、メモリ制約を解決するための従来のactivation recomputationの必要性を大幅に削減する新しい手法を提案した論文。sequence parallelismとselective activation recomputationという2つのシンプルな技術により、activation memoryを5倍削減し、recomputationによる実行時間オーバーヘッドを90%以上削減することを示した。 ### 詳細解説 #### 問題設定 Transformerモデルが数兆パラメータに拡大する中で、モデルパラメータ、activations、optimizer stateをデバイスメモリに収めることが困難になっている。特にactivationsの保存が重要なボトルネックとなっており、従来はactivation recomputation（gradient checkpointing）によってメモリを節約していたが、これは30-40%の計算時間オーバーヘッドを伴う。入力は各種サイズのTransformerモデル（22B～1Tパラメータ）で、出力は効率的な学習の実現である。必要なデータとして、バッチサイズb、シーケンス長s、隠れ次元h、レイヤー数L、アテンションヘッド数a、語彙サイズvなどがある。 #### 提案手法 **1. Activation Memory解析**: 単一Transformerレイヤーのactivation memoryを数式化し、sbh(34 + 5as/h)バイトと近似した。これにより各コンポーネントのメモリ使用量を定量的に把握できる。 **2. Sequence Parallelism**: tensor parallelismと組み合わせ、layer normやdropoutなどの非並列化領域をシーケンス次元で並列化する。通信演算gとḡを導入し、gはall-gather（順方向）/reduce-scatter（逆方向）、ḡはreduce-scatter（順方向）/all-gather（逆方向）として機能する。これによりactivation memoryをテンソル並列サイズtで均等分割可能となる。 **3. Selective Activation Recomputation**: アテンションブロック内の特定操作（QK^T行列積、softmax、softmax dropout、attention over V）のみ再計算対象とする。これらは大量のメモリを消費するが計算コストは低い。GPT-3では70%のメモリ削減を2.7%のFLOPsオーバーヘッドで実現する。 数式では、メモリ使用量が34sbhL/tに削減され、シーケンス長に対して線形、アテンションヘッド数に対して独立となる。 #### 新規性 従来のtensor parallelismは一部の領域（layer norm、dropout）を並列化できず、10sbhのメモリオーバーヘッドが残っていた。本研究のsequence parallelismは追加通信コストなしにこの問題を解決する。Ring all-reduceがreduce-scatter + all-gatherに分解できることを利用し、通信帯域幅を増加させない。また、selective activation recomputationは従来の全レイヤー再計算ではなく、メモリ使用量と計算コストの観点から最適化された部分再計算を実現している。 #### 実験設定 NVIDIA Seleneスーパーコンピュータ上でA100 80GB GPUを使用し、22B、175B（GPT-3）、530B（MT-NLG）、1Tパラメータのモデルで評価した。tensor parallel size=8、シーケンス長s=2048、語彙サイズv=51200で統一した。混合精度訓練を採用し、175Bと530Bモデルではinterleaved scheduleを使用した。メモリ使用量はPyTorchのGPU active memoryとtotal memoryで検証し、解析モデルとの一致を確認した。 #### 実験結果 **メモリ削減**: sequence parallelismとselective activation recomputationをそれぞれ単独で使用すると約50%のメモリ削減、組み合わせると5倍（20%に削減）の効果を得た。これはfull activation recomputation（10%）の約2倍のメモリ使用量だが、計算オーバーヘッドは大幅に小さい。 **実行時間**: 22Bモデルの単層で、sequence parallelismのみで6%の高速化、selective recomputationで11%のオーバーヘッド（full recomputationは64%）、両者併用で4%のオーバーヘッドを実現した。大規模モデルほど効果が顕著で、530Bや1Tモデルでは2%のオーバーヘッドに留まった。 **エンドツーエンド性能**: 全構成で29.0%～32.1%のスループット向上を達成した。Model FLOPs Utilization（MFU）は1Tモデルで56.3%に達し、530B GPT-3スタイルモデルでは54.2%のMFU（従来の42.1%から29%向上）を達成した。 ## Abstract 大規模なTransformerモデルの訓練は、現代のAIにおける最も重要な計算上の課題の一つです。本論文では、activation recomputationを削減することで、大規模Transformerモデルの訓練を大幅に高速化する方法を示します。Activation recomputationは、メモリ容量の制約を回避するために一般的に使用されています。逆伝播のためにactivationを保存する代わりに、従来はそれらを再計算していましたが、これによりメモリは節約されるものの冗長な計算が追加されます。本研究では、メモリ消費を十分に削減できるため、この冗長な計算の大部分は不要であることを示します。我々は、sequence parallelismとselective activation recomputationという2つの新しくシンプルな技術を提示します。tensor parallelismと組み合わせることで、これらの技術はactivationの再計算の必要性をほぼ完全に排除します。1兆パラメータ規模の言語モデルまでの我々のアプローチを評価し、我々の手法がactivation memoryを5倍削減し、activation recomputationによる実行時間オーバーヘッドを90%以上削減することを示します。例えば、2240台のNVIDIA A100 GPUで530BパラメータのGPT-3スタイルモデルを訓練する場合、Model Flops Utilizationは54.2%を達成し、recomputationを使用した場合の42.1%より29%高速化されます。