# 選択的活性化再計算 ## 定義 選択的アクティベーション再計算(selective activation recomputation)は、逆伝播に必要なアクティベーションをすべて保存する「完全保持」でも、すべてを再計算する「完全再計算(gradient checkpointing)」でもなく、「再計算コストが低くメモリ占有が大きい演算の出力のみを選択的に再計算する」手法である。 訓練中のアクティベーションメモリ管理には 3 つの選択肢がある: | 方式 | メモリ | 計算オーバーヘッド | |---|---|---| | 完全保持 | 最大(モデル規模に比例) | なし | | 完全再計算 | 最小(入力のみ保持) | +30〜40% | | **選択的再計算** | **中間(5 倍削減、完全再計算比 2×)** | **+2〜7%** | Korthikanti+ MLSys2023([[@2023__MLSys__Reducing Activation Recomputation in Large Transformer Models]])は [[Megatron-LM]] のテンソル並列化と [[シーケンス並列化]] と組み合わせて選択的再計算を定式化し、トランスフォーマーモデルで実用化した。 ### 選択する演算の基準 トランスフォーマーのアテンションブロックでは、Q/K/V 線形層を経てシーケンス次元の幅が `a`(ヘッド数)倍に拡大した後の演算——QKᵀ 行列積・Softmax・Softmax Dropout・V へのアテンション——が以下の特性を持つ: - **メモリ占有が大**: シーケンス長の 2 乗に比例する `2as²b` バイトのオーダー - **FLOPs/要素比が低い**: 入力要素あたりの浮動小数点演算が少なく、再計算のコストが軽微 Eq.5(テンソル+シーケンス並列化後)のアクティベーション式 `(sbh/t)(34 + 5as/h)` のうち `5as/h` の項がこれに対応する。 GPT-3(a=96, s=2048, h=12288)では `5as/h = 80` であり、定数項 34 を大きく上回る。選択的再計算でこの部分を省くと: - アクティベーション削減率: 70%(GPT-3)、65%(MT-NLG) - 再計算 FLOPs オーバーヘッド: 2.7%(GPT-3)、1.6%(MT-NLG) 選択的再計算後のメモリ式: $\text{Eq.6:} \quad \text{Total} = \frac{34 \cdot sbhL}{t}$ これはシーケンス長 s に対して線形であり、アテンションヘッド数 a に非依存となる。 ## 横断的知見 - **「全か無か」から「選択的」へ**: Korthikanti+ 2023 以前の実践では、トランスフォーマー層境界でのチェックポイント(完全再計算)が主流であった。選択的再計算の最大の貢献は「どこを選ぶか」の分析的基準——FLOPs/要素比とメモリ占有の積——を与えたことにある。この基準は任意アーキテクチャに汎化可能であり、他のモデル設計にも適用できる。 - **シーケンス並列化との相乗効果**: [[シーケンス並列化]] が LayerNorm・Dropout の冗長複製を排除して定数項の `10sbh` を `10sbh/t` に削減し、選択的再計算が `5as/h` 項を削減する。2 手法の組み合わせにより、それぞれが約 50% 削減するのに対し合計で 5 倍削減を実現する。 - **モデル規模とシーケンス長が大きいほど効果的**: `as/h` 比率はモデルが大規模・長コンテキスト化するにつれて増大するため、選択的再計算の効果も拡大する。530B・1T モデルでは再計算オーバーヘッドが 2% にとどまるが、22B モデルでは 7% である。 ## 未解決の問い - 長コンテキスト訓練(s/h >> 1)において選択的再計算の FLOPs オーバーヘッドがどこまで増大するか。論文は定性的に「依然として最良の戦略」と主張するが定量的限界は示されていない。 - Flash Attention 等の I/O 効率化手法と選択的再計算は、どのように相互作用するか。Flash Attention はアテンション内部の中間値を保存しないことで選択的再計算と類似の目標を持つが、実装上の相互依存は別途調査が必要。 - MoE(Mixture-of-Experts)モデルでは Expert Parallelism が加わり活性化パターンが動的になる。MoE の選択的再計算における「再計算すべき演算の選択」はどう変わるか。 - 自動化された最適選択手法(Feng & Huang CVPR2021 の任意計算グラフへの最適チェックポイント探索)との組み合わせは有効か。 ## 関連 - 親概念: [[並列化戦略]] / [[LLM分散学習]] - 姉妹概念: [[シーケンス並列化]] / [[チェックポイント]] - ソース: [[@2023__MLSys__Reducing Activation Recomputation in Large Transformer Models]] - エンティティ: [[Megatron-LM]] / [[NVIDIA]] / [[Vijay Korthikanti]] / [[Bryan Catanzaro]] - 関連 MOC: [[分散深層学習 - MOC]] ## 出典 - [[@2023__MLSys__Reducing Activation Recomputation in Large Transformer Models]](Korthikanti+, MLSys 2023: 選択的再計算の定式化・解析式・530B での 30% 高速化実証)