大規模言語モデルの学習知見

NVIDIA AI Summit 2024 Japan Tokyo
https://www.nvidia.com/ja-jp/events/ai-summit/

横田理央 , 教授, 東京科学大学
藤井一喜, 修士課程, 東京科学大学

GENIAC プロジェクトにおける 172B 規模の LLM の学習をはじめとする大規模モデルの学習には、分散学習に関する知識と、大規模な計算機環境を使いこなす知見が必要です。本セッションでは、Megatron-LM を用いた大規模モデルの実例と知見について紹介します。

Transcript

1. 1 大規模モデルの学習知見 東京科学大学(旧 東京工業大学) 横田理央/藤井一喜
2. 自己紹介: 横田理央 東京科学大学　総合研究院　スーパーコンピューティング研究センター　教授 専門：高性能計算、大規模深層学習、大規模行列演算 2 ImageNet on ResNet50 を2分で学習 フラクタルによる ViTの 事前学習 大規模言語モデルの　 事前学習 Tensor Coreの精度保証 Fugaku-LLM 13B LLM-jp 175B, 172B Swallow 7B,13B,70B Swallow-MS 7B Swallow-MX 8x7B Llama3-Swallow 8B,70B Llama3.1-Swallow 8B,70B [CVPR'19,KDD'20,TPAMI'24] [CVPR'22,CVPR'23,ICCV'23] [COLM'24] FP16 Tensor Coreを用いて FP32の精度を実現 INT8 Tensor Coreを用いて FP64の精度を実現 [IJHPCA'22,IJHPCA'24]
3. 深層学習の大規模化 3
4. 計算資源で殴るだけではない 4 https://situational-awareness.ai/from-gpt-4-to-agi/ MATHで同じスコアを得るために必 要な推論コストは2年で1/600に モデルの足枷を外すことによる大 幅な性能向上が期待できる
5. Unhobbling: The Test-Time Scaling Law 5
6. 分散並列学習 6
7. 7 データ並列
8. 8 ZeRO (FSDP)
9. 9 パイプライン並列
10. 10 テンソル並列
11. 自己紹介: 藤井 一喜 (Kazuki Fujii) • 2023〜 Turing株式会社 リサーチインターン • 2023〜 Swallow Project • 2024〜 SB Intuitionsリサーチインターン • 2024〜 Sakana AI リサーチインターン https://twitter.com/okoge_kaz 11
12. 学習に必要なコスト LLM-jp 172Bモデルを2.1Tトークン学習するには 約 2.3 x 1e24 FLOPs 仮に 380 TFLOP/s で学習する場合、AWSのH100料金換算で 12.3 億円 もし、同じ数値精度でも 500 TFLOP/s で学習できるなら？ → 9.3 億円 (3億円 の節約) 学習効率次第で、必要なコストは 大きく変化する 12
13. 学習を高速化する技術 1. 適切な分散学習設定 a. テンソル並列とパイプライン並列 b. micro-batch-size c. シーケンス並列 2. パイプライン並列 3. 通信と計算の overlap a. データ並列に関する overlap 4. distributed checkpoint 5. 低精度化 (FP8 hybrid training) 13
14. Megatron-LM Swallow LLMの開発にも利用 高速に大規模モデルを学習するためのサンプルや機能が備わっているライブラリ (以降で紹介する高速化技術をすべて網羅) 14
15. 学習を高速化する技術 1. 適切な分散学習設定 a. テンソル並列とパイプライン並列 b. micro-batch-size c. シークエンス並列 2. パイプライン並列 3. 通信と計算のoverlap 4. distributed checkpoint 5. 低精度化 (FP8 hybrid training) 15
16. テンソル並列とパイプライン並列 TP = ノードあたりのGPU数 とするのが最速？ 16 Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM より (8, 8) が最速になっている
17. テンソル並列とパイプライン並列 TP = ノードあたりのGPU数とするのが常に最速ではない 17 TP数を増大させる → 行列積のサイズが小さくなる → GPUの利用率が低下 → 学習において FLOP/s が低下 ← A100(40GB) ABCI 2.0にて計測
18. micro batch size MBS(micro batch size) を適切に設定することで学習効率が上昇 18 ← 8GPUと128GPUで同程度の TFLOP/s/GPU ← A100 (40GB) ABCI 2.0にて計測
19. micro batch size なぜ、MBSを増加させると FLOP/s/GPUが上昇するのか？ 1. GPUに対する処理要求が集中 2. GPU上で実行される算術演算密度 (atithmetic intensity)が上昇 ↑ 3. GPUリソースが効率的に利用できる → パフォーマンスの向上 ↑ では、常にFLOP/s/GPUは上昇するのか？ パイプライン並列と併用しているときは 場合による GPUの効率的利用 ↔ パイプラインバブルの増加 どちらの影響が上回るかにかかっている 19
20. micro batch size MBSをできるだけ増加させるのが良い訳ではない (トレードオフ) 20 パイプライン並列 と併用時は注意 Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM より
21. シークエンス並列 テンソル並列で分割できない LayerNormとDropoutをシークエンス(sequence)方向に分割 LayerNorm, Dropoutをsequence方向に分割することは演算に影響を与えない 21 Reducing Activation Recomputation in Large Language Models より
22. シーケンス並列 backward計算のために必要な Activationを以下の式のように削減することができる • テンソル並列 (Tensor Parallel) のみ Transformer 1層 あたり = • テンソル並列 + シーケンス並列 (Sequence Parallel) Transformer 1層 あたり = 22 10 → 10/t となるので 1GPUあたりの activation memory を削減
23. 学習を高速化する技術 1. 適切な分散学習設定 a. テンソル並列とパイプライン並列 b. micro-batch-size c. シーケンス並列 2. パイプライン並列 3. 通信と計算のoverlap 4. distributed checkpoint 5. 低精度化 (FP8 hybrid training) 23
24. パイプライン並列 : 1F1B vs. interleaved 1F1B P2P通信の増加を無視すれば interleaved 1F1Bの方が速い 24
25. 1F1B vs. interleaved 1F1B micro batch sizeが同じなとき activationが消費するメモリが増加 → OOMの危険性 余分なP2P communication(前述)が発生する → 下図ほど速くはならない 25
26. 学習を高速化する技術 1. 適切な分散学習設定 a. テンソル並列とパイプライン並列 b. micro-batch-size c. シークエンス並列 2. パイプライン並列 3. 通信と計算のoverlap 4. distributed checkpoint 5. 低精度化 (FP8 hybrid training) 26
27. 通信と計算の overlap 学習にかかる時間を短縮するために使用する GPU数を増加させる → 通信(communication) 時間が増え始める → 効率の悪化 (最悪の場合 DPの通信のために20%を消費) 27 forward backward communication communication forward forward backward forward 通信と計算のoverlapの模式図
28. データ並列に関する overlap Megatron-LM の --overlap-grad-reduce, --overlap-param-gather 28 GTC24 Optimizing Parallelization and Overlap to Increase Training Efficiency using Megatron-Core より
29. 学習を高速化する技術 1. 適切な分散学習設定 a. テンソル並列とパイプライン並列 b. micro-batch-size c. シーケンス並列 2. パイプライン並列 3. 通信と計算のoverlap 4. distributed checkpoint 5. 低精度化 (FP8 hybrid training) 29
30. distributed checkpoint 通常の torch checkpoint 特定のGPUにparameter, optimizer statesを集めて保存を行う (通信と保存時間が必要 ) torh-dist checkpoint 複数のGPUで分割してparameter, optimizer statesを保存 → 通信量が減少 & 1つあたりのサイズが減 ↓ → 並列書き込み性能が高いストレージでは保存にかかる時間が減少 ↓ async checkpoint • main stream: 計算 • ckpt stream: checkpoint saveを行う 30
31. 通常のcheckpoint async checkpoint async checkpoint 31 forward backward parameter update checkpoint save forward backward … forward backward forward backward …
32. 学習を高速化する技術 1. 適切な分散学習設定 a. テンソル並列とパイプライン並列 b. micro-batch-size c. シーケンス並列 2. パイプライン並列 3. 通信と計算のoverlap 4. distributed checkpoint 5. 低精度化 (FP8 hybrid training) 32
33. 低精度化 (FP8 hybrid) 33 FP8 E4M3: 精度優先 → forward に向いている FP8 E5M2: 数値幅優先 → backward に向いている forwardはE4M3 backwardはE5M2 を使用すれば、影響を最小化 → FP8 hybrid NVIDIA TransformerEngine Docsより
34. FP8 Scaling 34 NVIDIA TransformerEngine Docsより FP8 hybridでも何も工夫なしでは学習が上手くいかない → overflow, underflow を防ぐために(主にgradient) scaling を行う (左図) FP16 mixed precision trainingでも同様の措置が取られている → FP8 では、activation, gradient全体でなくテンソルごとにscalingを調整することで対応
35. FP8 Delayed Scaling scalingする値をどのように決定するのか？ → NVIDIA Transformer Engine では Delayed Scalingが使用されている ユーザーが決めた期間 (history)に基づき、maximums of absolute values(amaxes)からscalingを決定 35 1. amaxes から機械的にscaling factorを求める 2. 採用された値に基づき FP8 operator は演算を行う 3. FP8 output と 新しい amax が得られる 4. amax historyを更新 5. 1へ(繰り返し)
36. FP8 TransformerEngine ユーザーは、細かな FP8に関することを意識する必要があるか？ → TransformerEngine (ライブラリ) がすべてやってくれる TransformerEngineに用意されているcomponet を利用するだけで良い 36
37. ライブラリの関係性 37 • TransformerEngine → componentの効率的な定義、FP8 support • Megatron-LM (MCore) → 分散学習のための実装、モデル定義など • NeMo → 学習のためのfrontend
38. まとめ • 大規模モデルの学習には多額のコストがかかる → 学習効率は非常に重要 • 学習効率を上げる方法 ◦ 適切な分散学習設定 ◦ 発展的なパイプライン並列の採用 ◦ 通信と計算のoverlap ◦ checkpoint形式の工夫 ◦ 低精度計算 (FP8) • Megatron-LM、NeMo は、上記をすべてサポート ◦ optionを適切に設定することで学習効率を上げられる • 各環境、学習するモデルごとに最適な設定は異なる ◦ optionをイジるのにも、知識と何をしているのかの 理解が必要 38