SparseなDNNモデルの分散学習

大規模言語モデル (LLM) をはじめとする巨大なDNNモデルでは、計算量の削減のためsparseなアーキテクチャが使われることがあります。本講演では、このようなsparseなDNNモデルに対する分散学習の手法について解説します。12/20のW&Bミートアップで発表しました。

Transcript

1. sparseなDNNモデル の分散学習 株式会社 Preferred Networks リサーチャー 三上 裕明
2. 2 • 三上 裕明 (みかみ ひろあき) • (主な) 業務: DNN学習等の高速化/分散処理 ◦ LLMの学習高速化 ▪ PLaMo-13Bを公開しました ◦ MN-Core 向けコンパイラの開発 ◦ その他GPUを用いた分散処理の最適化 • 経歴 ◦ 株式会社 Preferred Networks (2019/9 〜) ◦ ソニー株式会社 (2017/4 〜 2019/8) ◦ 東京大学大学院 (2015/4 〜 2017/3) 自己紹介
3. 3 • sparse DNNモデルとは ◦ sparse DNNモデル ◦ LLM の影響 • sparse DNNモデルの分散学習 手法紹介 ◦ 手法1: Mixture of Experts (MoE) ◦ 手法2: Branch-Train-Merge (BTM) • さいごに 目次
4. 4 sparse DNNモデル (LLM以外)
5. 5 sparsityによってDNNの性能を改善したモデル sparse DNNモデルとは dense DNN sparseな重みを持つ 入力によって使用する重みを変える hello こんにちは
6. 6 sparsityによってDNNの性能を改善したモデル sparse DNNモデルとは dense DNN sparseな重みを持つ 入力によって使用する重みを変える hello こんにちは メモリ消費量 を削減できる
7. 7 sparsityによってDNNの性能を改善したモデル sparse DNNモデルとは dense DNN sparseな重みを持つ 入力によって使用する重みを変える hello こんにちは メモリ消費量 を削減できる 大きなモデルを小さい計算量で 動かせる
8. 8 sparse DNNモデルの利用分野は限られている sparse DNNモデル: 現状 (LLM注目以前) 問題1: 多くの場合GPUで効率的に扱えない処理が必要 [Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks] GPUの高速化のためにsparseモデルを使うのは難しい 問題2: 学習が不安定になりやすい [ST-MoE: Designing Stable and Transferable Sparse Expert Models] 推論の効率化に使われることが多かった Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity loss
9. 9 LLM とsparse DNNモデル
10. 10 LLM LLMとsparse DNNモデル 巨大なTransformer Attention Is All You Need
11. 11 LLM LLMとsparse DNNモデル 巨大なTransformer 非常に計算量が大きい Attention Is All You Need
12. 12 LLM LLMとsparse DNNモデル 巨大なTransformer 非常に計算量が大きい Attention Is All You Need 実効効率よりも 計算量の削減が大事
13. 13 LLM LLMとsparse DNNモデル 巨大なTransformer 非常に計算量が大きい Attention Is All You Need 実効効率よりも 計算量の削減が大事 学習に時間がかかる
14. 14 LLM LLMとsparse DNNモデル 巨大なTransformer 非常に計算量が大きい Attention Is All You Need 実効効率よりも 計算量の削減が大事 学習に時間がかかる 学習の高速化が重要
15. 15 LLM LLMとsparse DNNモデル 巨大なTransformer 非常に計算量が大きい Attention Is All You Need 実効効率よりも 計算量の削減が大事 学習に時間がかかる 学習の高速化が重要 LLMではsparseなモデルの学習にメリットがある
16. 16 LLMとsparse DNNモデル: 難しい点 問題1: 学習が不安定になりやすい [ST-MoE: Designing Stable and Transferable Sparse Expert Models] 問題2: 大規模学習における知見/best practiceが乏しい loss
17. 17 LLMとsparse DNNモデル: 難しい点 問題1: 学習が不安定になりやすい [ST-MoE: Designing Stable and Transferable Sparse Expert Models] 問題2: 大規模学習における知見/best practiceが乏しい 主に論文に基づいて手法を紹介 loss
18. 18 手法1: Mixture of Experts (MoE)
19. 19 手法1: Mixture of Experts (MoE) Router DNN0 DNN1 DNN2 DNN3 hello tokenごとに使用するDNNを変える ⇒ sparseなモデルとなる
20. 20 手法1: Mixture of Experts (MoE) Router DNN0 DNN1 DNN2 DNN3 hello tokenごとに使用するDNNを変える ⇒ sparseなモデルとなる 使用するDNNを決める 小さなモデル それぞれのDNNをexpertと呼ぶ
21. 21 手法1: Mixture of Experts (MoE) Router DNN0 DNN1 DNN2 DNN3 hello tokenごとに使用するDNNを変える ⇒ sparseなモデルとなる このtokenは DNN0を使う
22. 22 手法1: Mixture of Experts (MoE) Router DNN0 DNN1 DNN2 DNN3 hello tokenごとに使用するDNNを変える ⇒ sparseなモデルとなる このtokenは DNN0を使う
23. 23 手法1: Mixture of Experts (MoE) Router DNN0 DNN1 DNN2 DNN3 hello こんにちは Router DNN0 DNN1 DNN2 DNN3 tokenごとに使用するDNNを変える ⇒ sparseなモデルとなる このtokenは DNN0を使う このtokenは DNN2を使う
24. 24 手法1: Mixture of Experts (MoE) Router DNN0 DNN1 DNN2 DNN3 hello こんにちは Router DNN0 DNN1 DNN2 DNN3 tokenごとに使用するDNNを変える ⇒ sparseなモデルとなる パラメータ数はexpert数 に比例して増加する
25. 25 手法1: Mixture of Experts (MoE) Router DNN0 DNN1 DNN2 DNN3 hello こんにちは Router DNN0 DNN1 DNN2 DNN3 tokenごとに使用するDNNを変える ⇒ sparseなモデルとなる 計算量はexpert数に よらず固定 パラメータ数はexpert数 に比例して増加する
26. 26 手法1: Mixture of Experts (MoE) 性能例 LLMに適用した例としてSwitch Transformer の性能を引用 モデル モデルサイズ Negative Log Perplexity 学習速度 (example/sec) T5-Base (denseモデル) 0.2B -1.731 1600 T5-Large (denseモデル) 0.7B -1.550 470 Switch-Base (MoE) 7B -1.534 1000 Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
27. 27 手法1: Mixture of Experts (MoE) 難しさ Router DNN0 DNN1 DNN2 DNN3 hello 問題: 通信コストが大きく、分散学習の効率を出しづらい MoEの学習 通常のDNN学習 DNN hello
28. 28 手法1: Mixture of Experts (MoE) 難しさ Router DNN0 DNN1 DNN2 DNN3 hello 問題: 通信コストが大きく、分散学習の効率を出しづらい MoEの学習 通常のDNN学習 DNN hello 計算量はほぼ同じ パラメータは expertの個数倍
29. 29 手法1: Mixture of Experts (MoE) 難しさ Router DNN0 DNN1 DNN2 DNN3 hello 問題: 通信コストが大きく、分散学習の効率を出しづらい MoEの学習 通常のDNN学習 DNN hello 計算量はほぼ同じ パラメータは expertの個数倍 データ並列の通信がネックとなりやすい
30. 30 手法1: MoE: Expert Parallelism Router DNN0 hello ワーカーがexpertの一部のみを担当することで通信コストを削減する ワーカー0 ワーカー1 こんにちは Router DNN1
31. 31 手法1: MoE: Expert Parallelism Router DNN0 hello ワーカーがexpertの一部のみを担当することで通信コストを削減する ワーカー0 ワーカー1 こんにちは Router DNN1 ワーカーが持つパラメータは小さくなる データ並列の通信コスト減
32. 32 手法1: MoE: Expert Parallelism Router DNN0 hello ワーカーがexpertの一部のみを担当することで通信コストを削減する ワーカー0 ワーカー1 こんにちは Router DNN1 このtokenは DNN1を使う
33. 33 手法1: MoE: Expert Parallelism Router DNN0 hello ワーカーがexpertの一部のみを担当することで通信コストを削減する ワーカー0 ワーカー1 こんにちは Router DNN1 担当expertの workerへ送信
34. 34 手法1: MoE: Expert Parallelism Router DNN0 hello ワーカーがexpertの一部のみを担当することで通信コストを削減する ワーカー0 ワーカー1 こんにちは Router DNN1 担当expertの workerへ送信 もとの workerへ送信
35. 35 手法1: MoE: Expert Parallelism Router DNN0 hello ワーカーがexpertの一部のみを担当することで通信コストを削減する ワーカー0 ワーカー1 こんにちは Router DNN1 担当expertの workerへ送信 もとの workerへ送信
36. 36 手法1: MoE: Expert Parallelism Router DNN0 hello ワーカーがexpertの一部のみを担当することで通信コストを削減する ワーカー0 ワーカー1 こんにちは Router DNN1 担当expertの workerへ送信 もとの workerへ送信 担当するワーカーへのtokenの通信が必要
37. 37 手法1: MoE: Expert Parallelism Router DNN0 hello ワーカーがexpertの一部のみを担当することで通信コストを削減する ワーカー0 ワーカー1 こんにちは Router DNN1 担当expertの workerへ送信 もとの workerへ送信 担当するワーカーへのtokenの通信が必要 この通信の総量はexpertの個数によらず一定
38. 38 手法1: MoE: Expert Parallelism Router DNN0 hello ワーカーがexpertの一部のみを担当することで通信コストを削減する ワーカー0 ワーカー1 こんにちは Router DNN1 担当expertの workerへ送信 もとの workerへ送信 担当するワーカーへのtokenの通信が必要 この通信の総量はexpertの個数によらず一定 データ並列よりもexpertの個数を増やせる
39. 39 手法1: MoE: Expert Parallelism 実用上のポイント • データ並列とexpert parallelの使い分けはモデル・計算機クラスタ次第 [ST-MoE: Designing Stable and Transferable Sparse Expert Models] • tokenの偏りによって効率低下する Router DNN0 hello ワーカー0 ワーカー1 hello Router DNN1
40. 40 手法1: Mixture of Experts (MoE): まとめ • Mixture of Experts (MoE): tokenごとに使用するDNNを切り替える ◦ 高速に高い性能のモデルを学習できる • 分散学習で速度を出すのが難しい ◦ expert parallelismによる分散学習を考慮する必要がある ◦ 分散学習の手法、モデルアーキテクチャ、計算機クラスタの構成が 強く関連する 分散学習手法 モデルアーキテクチャ 計算機クラスタの構成 expert parallelの並列数 expertの個数 モデル並列の設定 データ並列かexpert parallelか
41. 41 手法2: Branch-Train-Merge (BTM)
42. 42 手法2: Branch-Train-Merge (BTM) Router DNN0 DNN1 DNN2 DNN3 hello こんにちは Router DNN0 DNN1 DNN2 DNN3 MoEではtokenの割り当て方が問題となる
43. 43 手法2: Branch-Train-Merge (BTM) Router DNN0 DNN1 DNN2 DNN3 hello こんにちは Router DNN0 DNN1 DNN2 DNN3 MoEではtokenの割り当て方が問題となる データセットの情報を割り当て に使えないか？
44. 44 手法2: Branch-Train-Merge (BTM) 英語 Transformer hello こんにちは 日本語 Transformer BTM: 学習データセットごとにTransformerのexpertを用意 英語 データセット 日本語 データセット … Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models
45. 45 手法2: Branch-Train-Merge (BTM) 英語 Transformer hello こんにちは 日本語 Transformer 英語 データセット 日本語 データセット … 利点: 分散学習に自明並列性が使える ワーカー0 ワーカー1
46. 46 手法2: Branch-Train-Merge (BTM) 英語 Transformer hello こんにちは 日本語 Transformer 英語 データセット 日本語 データセット … 利点: 分散学習に自明並列性が使える それぞれのTransformerは完全に独立したモデル 学習中にワーカー間の通信が必要ない ワーカー0 ワーカー1
47. 47 手法2: Branch-Train-Merge (BTM) モデル モデルサイズ Perplexity Transformer LM 125M 19.9 BTM (8 expert) 1B (125M x8) 17.2 Transformer LM 750M 14.7 BTM (8 expert) 6B (750M x8) 13.4 Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models
48. 48 手法2: Branch-Train-Merge (BTM) 難しさ: データセットの分割が性能に大きく影響する Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models
49. 49 手法2: Branch-Train-Merge (BTM) 難しさ: データセットの分割が性能に大きく影響する ランダムなデータセットの分割だと もとのTransformerより性能が劣化する Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models
50. 50 手法2: Branch-Train-Merge (BTM) 難しさ: データセットの分割が性能に大きく影響する ランダムなデータセットの分割だと もとのTransformerより性能が劣化する Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models 適切なデータセット分割が不可欠
51. 51 手法2: Branch-Train-Merge (BTM): まとめ • BTM: データセットごとに使用するDNNを切り替える ◦ 自明並列性を使った簡単で効率の良い分散学習ができる ◦ 少ない計算資源で大きいモデルを学習できる • 学習したモデルの性能はデータセットの分割方法に大きく影響する ◦ 分散学習の手法・モデルアーキテクチャが データセットの前処理に依存する
52. 52 さいごに
53. 53 • LLMのような超巨大DNNモデルではsparseモデルの大規模な分散学習が 重要になりつつある ◦ 手法1: Mixture of Experts (MoE) ◦ 手法2: Branch-Train-Merge (BTM) • sparseモデルの分散学習は様々な分野を考慮する必要がある点が特徴的 ◦ モデルアーキテクチャを踏まえた分散学習手法 ◦ 計算機クラスタに合わせたモデルアーキテクチャ ◦ データセット構成を踏まえたモデルアーキテクチャと分散学習手法 . . . sparseなDNNモデルの分散学習 さいごに
54. Making the real world computable