情報処理学会-全国大会2024-大規模言語モデルの分散並列学習

Transcript

1. 大規模言語モデルの分散並列学習 東京工業大学 藤井一喜　　横田理央 5J-02
2. 2 概要 • Llama 2をベースに日本語コーパスで継続事前学習 • 7B, 13B, 70B のモデル規模にて分散並列学習 • メモリ効率化のための工夫を行った3D Parallelism
3. 3 継続事前学習 Meta Llama 2 Swallow Swallow Corpus
4. 4 学習時に必要なメモリ (backward時) FP16/FP32 Mixed Precison p = parameter数 parameters gradients optimizer states 2p + 2p + 12p = 16p 必要 注意: activation、中間層の出力、バッチデータ、memory fragmentation などあるため、これだけではない
5. 5 学習時に必要なメモリ (backward時) FP16/FP32 Mixed Precison p = parameter数 parameters gradients optimizer states 2p + 2p + 12p = 16p + α 必要 → 1つのGPUで学習するのは不可能 → 分散学習
6. 6 分散学習 大規模言語モデルの学習には様々な分散並列化手法がある Data Parallel(=DP) Tensor Parallel(=TP) Pipeline Parallel(=PP)
7. 7 分散学習設定 3D Parallelism (DP, TP, PP)を利用 メモリ効率化のために以下を採用 • SP: Sequence Parallelism • Distributed Optimizer (DeepSpeed ZeRO Stage1相当)
8. 8 なぜ3D Parallelismなのか 3D ParallelismとFSDPの比較 3D Parallelism FSDP Llama 2 7B 134 TFLOPS/GPU 134 TFLOPS/GPU Llama 2 13B 143 TFLOPS/GPU 135 TFLOPS/GPU Llama 2 70B 158 TFLOPS/GPU 87 TFLOPS/GPU
9. 9 効率的なメモリ消費 (Distributed Optimizer) データ並列 Distributed Optimizer GPU: 1 GPU: 2 GPU: 3 GPU: 1 GPU: 1 GPU: 2 GPU: 3 3D Parallelism + Distributed Optimizer optimizer states optimizer states optimizer states Gradinets optimizer states Gradinets Gradinets Parameters Parameters Parameters
10. 10 トポロジーを考慮した3Dマッピング TP → DP → PP の順に配置 必要な通信量を考慮 TPは大量のAll Reduce →Tensor Parallel をノード内 Pipeline ParallelはP2P通信 → ノード間 Microsoft Research Blogより
11. 11 学習曲線
12. 12 TFLOPS Swallow Project における TFLOPS Megatron-LM GPT-3 (175B) 51.4 % (=160TFLOPS)
13. 補足資料
14. 14 言語モデルの評価結果 日本語タスク平均スコア
15. 15 学習トークンに対するスケール性
16. 16 学習ライブラリ 1. 3D Parallelism をサポート 2. Llamaアーキテクチャに対応 3. HF → Megatron 変換コードあり 詳細 ↓
17. 17 ABCI環境 1 Interconnect InfiniBand HDR 200Gbps x 4 Intranode NVLink
18. 18 ABCI環境 2 A100 SXM4 PLX PCIe IB HDR PLX A100 SXM4 IB HDR IB Switch PCIe
19. 19 ABCI環境 3 学習に使用したAノード(A100)は フルバイセクションバンド幅のFat Tree → 通信帯域幅のボトルネックは解消されている FSDP < 3D Parallelism であることは変わりないがABCIの環境では FSDPでも致命的に遅くはならない