## Memo ![[Pasted image 20251127233747.png]] ## Memo with LLM ### 論文情報 - 論文のタイトル: Reducing Activation Recomputation in Large Transformer Models - 著者と所属: Vijay Anand Korthikanti (NVIDIA), Jared Casper (NVIDIA), Sangkug Lym (NVIDIA), Lawrence McAfee (NVIDIA), Michael Andersch (NVIDIA), Mohammad Shoeybi (NVIDIA), Bryan Catanzaro (NVIDIA) - カンファレンス/ジャーナル名: MLSys 2023（Memory Optimization ポスターセッション） - 発表年: 2023 ### 論文概要 - 高速化とメモリ節約を両立させるために、活性化の再計算を最小限に抑えつつ大規模トランスフォーマーを訓練する手法を紹介している。 - シーケンス並列化と選択的活性化再計算をテンソル並列化と組み合わせて適用することで、活性化のメモリを5倍削減しながら、従来の再計算に比べて実行時間オーバーヘッドを90%以上削減している。 - 実装はMegatron-LMおよびNeMo-Megatronで公開予定で、5300億パラメータ規模のGPT-3スタイルモデルを2240台のA100で訓練した際にモデルFLOPs利用率が54.2%となり再計算を使うケースより29%高速に動作することを示している。 ### 詳細解説 #### 問題設定 - 入力はトランスフォーマーの多層構造とメモリ容量制限で、出力は大規模モデルをデバイスに収めつつ高いデバイス利用率を維持した訓練スループット。 - 通常はバックプロパゲーションのために活性化を保存せず再計算するため、30～40%の時間オーバーヘッドが発生し、訓練効率を下げる。 - この課題に対して、活性化のメモリを削減しつつ再計算を減らし、通信や計算のバランスを新しい並列化構成で取ることを目指す。 #### 提案手法 - テンソル並列化とシーケンス並列化を併用して、活性化をシーケンス次元に沿って分割し、ほぼ全ての活性化がテンソルグループ内で分散されるように設計。 - シーケンス並列化では層ノルムやドロップアウトも並列に処理し、オールギャザーとリデューススキャッターにより通信帯域は既存のオールリデュースと変わらないようにすることで追加の通信コストを避けた。 - 選択的活性化再計算では、バックパスで必要なテンソルだけを保存し、それ以外は必要に応じて再計算。テンソル並列ランクごとに分配される部分だけを保持し、余分なオールギャザーを計算と重ね合わせて隠蔽することでオーバーヘッドを低減する。 #### 新規性 - データ並列化やオフロードなど既存のメモリ削減技術よりも実装コストや通信オーバーヘッドを抑えつつ、テンソル・シーケンス並列化を混合することで活性化の分布を細かく調整できる。 - 活性化をテンソル並列サイズ分で分割することで、メモリの必要量を線形に削減し、ほぼ全ての再計算を不要にできる点が独自。 - 選択的再計算は層全体の再計算ではなく計算コストの高い操作だけを再計算し、通信と計算を重ねることでオーバーヘッドを大幅に小さくしている。 #### 実験設定 - 22B/175B/530B/1Tパラメータの言語モデルを対象に、主にNVIDIA A100上でMegatron-LM準拠構成を訓練し、テンソル並列・シーケンス並列・再計算の組み合わせを比較。 - 評価指標はモデルFLOPs利用率、各層のフォワード・バックワード・再計算時間、メモリ使用量、エンドツーエンドのイテレーション時間など。 - 2240台のGPUクラスタを用い、再計算の有無とメモリ最適化の影響をスケールアップで測定。 #### 実験結果 - シーケンス並列化と選択的再計算を併用すると活性化メモリがテンソル並列ベースラインの20%未満となり、従来の再計算構成の約10%に匹敵する。 - フル再計算のオーバーヘッドは64%に達したが、選択的再計算で11%に低減し、さらにシーケンス並列化を加えると4%まで落ちた。 - イテレーション全体では29.0～32.1%の高速化が得られ、530Bモデルでは再計算ありの42.1%から54.2%へモデルFLOPs利用率を改善。 - 追加のマイクロバッチレベルの再計算を併用すると、175B/530Bモデルでそれぞれ+0.7/+0.4ポイントのモデルFLOPs利用率向上が得られ、既に小さい再計算コストがさらに圧縮された。 ## Abstract 大規模なトランスフォーマーモデルのトレーニングは、現代のAIにおいて最も重要な計算上の課題の一つである。本論文では、活性化の再計算を削減することで大規模トランスフォーマーモデルのトレーニングを大幅に高速化する方法を示す。活性化の再計算はメモリ容量制約を回避するためによく用いられるが、逆伝播に用いる活性化を保存せずに再計算するため余分な計算が発生する。本研究ではその多くの冗長な計算は不要で、十分なメモリ削減が可能であることを示す。シーケンス並列化と選択的活性化再計算という2つの簡潔で新規の手法を提示し、テンソル並列化と組み合わせることで活性化の再計算をほとんど不要とする。1兆パラメータ規模までの言語モデルで評価した結果、活性化メモリを5倍削減し、活性化再計算による実行時間オーバーヘッドを90%以上削減した。例えば、5300億パラメータの[[2020__NeurIPS__Language Models are Few-Shot Learners|GPT-3]]スタイルモデルを2240台のNVIDIA A100 GPUで訓練すると、モデルFLOPs利用率は54.2%となり、再計算を用いた42.1%より29%高速であった。実装はMegatron-LMおよびNeMo-Megatronで利用可能になる予定である。