この資料は、東京工業大学横田研究室の藤井一喜さんがW&Bマンスリーミートアップのために準備してくれた資料です。 「大規模言語モデル開発を支える分散学習技術」 大規模言語モデル（LLM）を学習する過程において、分散学習は避けて通れない重要な技術の一つです。本講演では、分散学習の基本的な概念とそのメカニズムをわかりやすく解説します。さらに、実例やノウハウについてもご紹介します。

1. 大規模言語モデル開発
を支える分散学習技術
1

2. 自己紹介: 藤井 一喜
● 2023〜 東京工業大学 横田理央 研究室
● 2023〜 Turing株式会社 リサーチインターン
● 2023〜 Kotoba Tech リサーチインターン
https://twitter.com/okoge_kaz 2

3. 横田研究室について
どのような研究室？
スパコンを利用した大規模深層学習、大規模行列計算の研究を行なう研究室
https://www.rio.gsic.titech.ac.jp/jp/index.html
最近のプロジェクト
● ABCI GrandChallenge 13B (2023年5月)
● LLM-jp 13B, 175B
● 産総研, 東工大 岡崎研究室 LLaMA-2 継続学習 (7B, 13B, 70B)
● その他進行中のLLM関連プロジェクト多数
3

4. LLMの学習で分散学習は必要？
数B パラメーターのモデルを学習するのであれば、GPUのメモリ制約から1GPUでは学習できない
→ 分散学習が必要
LLM-jp 13Bの学習やLLaMA-2 70Bの継続学習を行なうには？
→ データ並列、テンソル並列、パイプライン並列 を組み合わせる必要あり (後述）
工夫次第で学習速度が大きく変化するので、概念を理解することは重要
4

5. データ並列だけで十分？
純粋なデータ並列だけでは、大きなモデルは学習できない (GPUメモリ制約のため）
→ テンソル並列、パイプライン並列などは大規模なLLMの学習において欠かせない
DeepSpeed ZeRO, PyTorch FSDPなどで代替可能？
→ ある程度の規模までは、テンソル並列、パイプライン並列を併用しないでも学習できる
いま出てきた概念を理解するのが今日の目的
5

6. 目次
1. 分散並列学習の基礎
● データ並列(DP)
● テンソル並列(TP)
● パイプライン並列(PP)
● DeepSpeed ZeRO
1. 3D Parallelism と ZeRO
2. 学習事例
● LLaMA-2 7B, 13B, 70B 継続事前学習
1. W&Bの活用例
6

7. 分散並列学習の種類
データ並列とモデル並列の違いを説明できますか？
● データ並列 (Data Parallel = DP)
● モデル並列 (Model Parallel = MP)
○ パイプライン並列 (Pipeline Parallel = PP)
○ テンソル並列 (Tensor Parallel = TP) *
* 文献によっては、こちらをモデル並列と呼んでいる例もあります 7

8. データ並列
データ並列のポイント
1. データセットを分割
2. GPUごとにモデルをもつ
3. backward後に同期
画像: Colossal-AI より
8

9. データ並列のポイント1
データセットの分割
● 各GPUは割り当てられた分のデータセットだけを処理
● 学習時間は理論的*には 1/N (Nはデータ並列数)
● Global Batch Sizeは、batch size x N x A
(N: データ並列数、A: gradient accumulation step数)
* ノード間、ノード内通信の性能に依存する
9

10. データ並列のポイント2
GPUごとにモデルをもつ
● forward, backward 処理はそれぞれ別々に行う
● 学習データは別、モデルは一緒
→ gradient(勾配) は異なる
● 学習に必要なモデル重み, 勾配, Optimizer stateをそれぞれ持つ*
* 右図の Model Copy はそれぞれのGPUが有しているmodel weight, gradients, optimizer statesを指す 10

11. データ並列のポイント3
backward後に同期
● 別々のデータで学習しbackwardを行ったので勾配は異なる
→ All Reduce で同期を行なう
● 勾配の平均でモデル parameter を更新 *
● 次のstepへ
NVIDIA Developer Technical Blog より
* GPUごとに有しているModel Copyを用いて更新 11

12. データ並列 のまとめ
データ並列ができること、できないこと
● データ並列を使う意味は？
→ 学習時間の短縮 👍*
● データ並列は万能？
→ ❌ ただのデータ並列ではModel Copyをそれぞれで有している
→ 1GPUに載らないサイズのモデルは学習できない
→ モデル並列
* Large Batch Problemなど問題はある
12

13. モデル並列のコンセプト
何かしらの方法でモデルを分割する
● Transformer Layer内 or Layerごと
→ モデルを分割できるとGPUあたりのモデルパラメータは減少
→ 1GPUに収まらなかったモデルが学習可能 👍
keyword: モデル分割
図: Colossal AI より
右上図: Tensor Parallel(テンソル並列), 右下図: Pipeline Parallel (パイプライン並列)
13
↑ Tensor Parallel
↑ Pipeline Parallel

14. モデル並列の種類
1. テンソル並列 (Tensor Parallel: TP)
行列積(テンソル積)を分解する (右上図)
1. パイプライン並列 (Pipeline Parallel: PP)
モデルをある層(layer)でぶつ切りにする(右下図)
1. シークエンス並列(Sequence Parallel: SP)*
14
↑ Tensor Parallel
↑ Pipeline Parallel

15. テンソル並列の概要
テンソル並列のポイント
● 行列x行列 演算を並列化する
● テンソルをN個に分割することで、それぞれの
GPUは 1/N のテンソルを有するだけになる
→ GPUメモリ制約から学習できなかったモデル
も学習可能に 👍
● ただ分割するだけでは、もちろんダメ
→ 重要な点は、計算グラフに影響を与えな
い形で分割を行う *
NVIDIA nemo documentation より
* 詳細についてはMegatron-LMの論文にある行列積の分解式を”手”で解いてみると理解が深まります 15

16. 行列の分解
行列積の分解例
について分解を行う際
とすると
と表せます
これを図示したのが右図です*
* column dimension partitionの例
16
1
2
1
2

17. パイプライン並列の概要
モデルをlayerの束でN個に分割
→ 右の例ではモデルを4つに分割している
仮に32 layerあるモデルであれば、
1GPUに8 layer割り当てられることになる
→ GPUメモリ制約から解放 👍
Colossal AI より 17

18. パイプライン並列におけるBubble
効率的なパイプラインを目指して
● ナイーブな実装だと右上
→ GPUが働いていない白い部分が多い
→ GPUの使用率を上げたい
→ バッチをmicro-batchに分割して
パイプライン化すれば効率化できる 👍
(GPUが働いていない時間をbubbleという)
Google Research より 18
forward backward
forward backward

19. DeepSpeed ZeRO 概要
DeepSpeed ZeRO
気軽に分散学習することができるZeROです
が、何をしているのか説明できますか？
Stage 1: Optimizer States
Stage 2: Stage1 + Gradinets
Stage 3: Stage2 + Parameters
ZeRO: Memory Optimizations Toward Training Trillion Parameter Models より
19

20. ZeROのコンセプト
パラメーターの分散
データ並列のときに登場した右図ですが
よく考えると無駄(冗長性)がないですか？
→ 学習中いつでもモデルのすべてのparameterを利用する
わけではない。
特に、Optimizer State, Gradientsはbackward時しか利用しない
→ 必要なときに集めるだけで良い
→ GPU間で分散して保有しよう
20

21. ZeRO Stage1におけるメモリ消費量
DeepSpeed ZeRO Stage1
Optimizer StatesをN個に分割
→ GPU:0〜GPU:N-1 がそれぞれの担当を
受け持つ
→ Optimizer Stateに関するメモリ消費量
は 1/N になる 👍
21

22. 目次
1. 分散並列学習の基礎
● データ並列(DP)
● テンソル並列(TP)
● パイプライン並列(PP)
● DeepSpeed ZeRO
1. 3D Parallelism と ZeRO
2. 学習事例
● LLaMA-2 7B, 13B, 70B 継続事前学習
1. W&Bの活用例
22

23. 3D Parallelism とは
3D Parallelism
Data Parallel + Tensor Parallel + Pipeline Parallel を組み合わせた分散学習手法
ただ組み合わせるだけで良い？
→ ❌ ダメ
学習GPU環境(トポロジー）を考慮したProccess Groupの配置を行なう必要がある
Megatron-LM, Megatron-DeepSpeedはそこを自動で行ってくれる
23

24. Tensor Parallel + Pipeline Parallel
3D Parallelismを理解するために
Tensor Parallel + Pipeline Parallel (DP=1のとき)
Megatron-LM 論文より
24

25. 3D Parallelismの例
3D Parallelism (Data Parallel + Tensor Parallel + Pipeline Parallel)
下図: ZeRO DP + TP + PP
Data Parallel Process Group = 0 に着目
→ 32層のTransformer Blockを4つに分割
(Pipeline Parallel 次元 = 4)
DP Process Group = 0 & PP stage = 0に着目
→ Layer: 0〜7 が4つに分かれている
→ Tensor Parallel 次元 = 4
この右図は ZeRO DP=2, PP=4, TP=4 の
3D Parallelismを示している
25

26. 3D Parallelism と DeepSpeed ZeROの使い分け
3D Parallelism
3D Parallelismは万能か？
→ 使用できるのであれば、使用したほうが良い (FLOPs面で有利)
Megatron-LM, Megatron-DeepSpeedなど
(使用方法はZennの記事にて紹介しています)
しかし自前で実装するとなるとハードルが高い
加えて変更可能なパラメーターが多数あり、ある程度の知識が必要
→ 継続学習で30B以下のサイズのモデルを数B Token学習させるだけならZeRO, FSDPでも十分
26

27. 3D Parallelismを使うべきタイミングまとめ
結局どれを使うべき？
学習速度を気にしない(短期学習) & モデルサイズが30B以下
→ ZeRO or FSDP
学習速度が気になる(長期学習) or 事前学習 or モデルサイズが大きい
→ 3D Parallelism (Megatron-LM, Megatron-DeepSpeed …etc)
27

28. 目次
1. 分散並列学習の基礎
● データ並列(DP)
● テンソル並列(TP)
● パイプライン並列(PP)
● DeepSpeed ZeRO
1. 3D Parallelism と ZeRO
2. 学習事例
● LLaMA-2 7B, 13B, 70B 継続事前学習
1. W&Bの活用例
28

29. 事前学習 vs 継続事前学習
事前学習: 言語モデルを最初(from scratch)から学習させる (例: LLM-jp 13B)
継続事前学習: すでに学習済みの言語モデルをさらに学習させる (今回の我々）
なにが嬉しいの？
→ LLaMA-2であれば高い英語能力と、ある程度の日本語能力をすでに有している
ここに、自分たちが追加したい能力に寄与するデータセットを学習させる (比較的コスト低)
29
事前学習(from scratch) 継続事前学習
必要データセットサイズ 数100B〜数T Token 数B Token〜 数100 BToken
Tokenizer自由度 高い 制約あり
計算コスト 大 小

30. LLaMA継続学習概要
プロジェクト概要
LLaMA-2 7B, 13B, 70Bに日本語語彙を追加して日本語データ支配のデータセットで学習を行う
継続事前学習による日本語能力の獲得がメインテーマ
昨日(2023/12/19)にモデルをリリース
良質な指示チューニングデータで学習させた改良モデルを準備中
30

31. LLaMA-2 継続学習で行ったこと
本プロジェクトの特徴
1. 13B, 70Bにて語彙拡張版モデルをリリース
a. token化効率の面で優れている
2. 実験網羅性
a. 20B, 40B, 60B, 80B, 100B Token学習させるごとに日本語、英語両方の評価を行っている
b. 語彙拡張の有無の実験を7B, 70Bにて検証
c. LLaMA-2 7B, 13B, 70Bにて同一データセットによる実験
d. LLaMA-2 13Bにおけるデータセットの違いが性能に与える影響の検証
e. 対訳データセットが下流タスクに与える影響の検証
31
分 散 学 習 分散学習
Token化効率上昇例

32. 学習モデルの評価
評価
32
モデル名 日本語スコア
Meta Llama-2-7b-hf 0.3201
Swallow-7B 0.3940
Meta Llama-2-13b-hf 0.3963
Swallow-13B 0.4625
Meta Llama-2-70b-hf 0.4830
Swallow-70B 0.5528

33. LLaMA-2 70B 語彙拡張あり Loss推移
33

34. LLaMA-2 継続学習 学習詳細
学習詳細
Megatron-LMを利用
3D Parallelism + Sequence Parallel + Distributed Optimizer(zero 1相当)の利用
→ 継続学習においても高いFLOPsを実現 👍
LLaMA-2 70B におけるハイパラ(抜粋)
34
global batch size 1024
日本語: 英語(データセット) 9: 1 (= 90B Token, 10B Token)
Learning Rate min 1.6e-6, max: 5.0e-5

35. 目次
● 分散学習の基礎
○ データ並列(DP), テンソル並列(TP), パイプライン並列(PP)
○ DeepSpeed ZeRO
● 3D Parallelism と ZeRO, FSDP
● 学習事例
○ LLM-jp 13B
○ LLM-jp 175B
○ LLaMA-2 継続学習
● W&Bの活用例
35

36. W&Bの活用例: logging範囲
wandb.log()
train loss, valid lossだけでなく
optimizer state, learning rate,
TFLOPs, tokens per secなどloggingでき
るものはすべて記録
→ configの設定ミスに気づきやすい
(数値よりもグラフ化されていると👍)
研究室内で実験を準備していない人が
確認することでミスが格段に減少
LLM-jp 175B optimizer statesのlogging例
36

37. W&Bの活用例: グラフの描画
学習後のvisualize
右図では、複数のrunに分かれ
ている実験をworld_sizeごとに
まとめている
→ 同じ色で
表示すると分かりやすい
学習時はstep数でlogging
→ あとでtoken数をx軸に変更
→ global batch sizeごとの比較
チーム内の議論で👍
37

38. 本日のまとめ
● 分散学習の基礎
○ データ並列(DP), テンソル並列(TP), パイプライン並列(PP)
○ DeepSpeed ZeRO
● 3D Parallelism と ZeRO, FSDP
● 学習事例
○ LLM-jp 13B
○ LLM-jp 175B
○ LLaMA-2 継続学習
● W&Bの活用例
38

39. Appendix: テンソル並列の注意点
テンソル並列(TP)を学習に使用する際の注意
1. hidden size, attention headsを割り切るTP数
→ 右図のようにAttention headsごとに分割
1. TP数 <= 1 node GPU数
→ 頻繁にAll Reduceする必要がある
1. 1 node GPU数 % TP数 == 0
→ Tensor Parallel Process Groupの配置上
の都合*
Efficient Large-Scale Language Model Training on
GPU Clusters Using Megatron-LM より
* Megatron-LM, Megatron-DeepSpeedにおける Process GroupとGPUを紐付ける配置コードより
39

40. Appendix: Tensor Parallelにおける通信
Tensor Parallel における All Reduce
forward
Self-Attention → All-Reduce(g) → MLP → All-Reduce(g)
backward
MLP → All-Reduce(f) → Self-Attention → All-Reduce(f)
結果を別々のGPU間で集める必要がある
forward: Dropoutの前
backward: MLPブロック、Self-Attentionブロックの終わり
40

41. Appendix: メモリ効率が良いパイプライン
1F1B
様々な種類のPipeline Parallelが存在
右上: GPipe(All-Forward All-Backward)
右下: 1F1B(1Forward-1Backward)
上の２つはBubbleの量は同じだが
backward時に必要となる中間Activation
に必要となるメモリ量については1F1Bの方が少ない 👍
Megatron-LM 論文より 41

42. Appendix: ZeROの通信コスト
DeepSpeed ZeRO Stage1
Data ParallelでGradientsをbackward終了後に平均していた
このときの通信コストはAll-Reduceを利用していたのでモデルパラメータ数を とすると *
ZeRO Stage1, 2では、それぞれのGPUの担当領域分の
Gradientsの平均が求まっていれば良い → scatter-reduce
更新を行った後にupdated parametersをすべてのプロセスに
配布する→ all-gather すなわち、合計で (DPと同じ)
* All-Reduceはreduce-scatterとall-gatherのプロセスに分けられるため
ReduceScatter
AllGather
42

43. Appendix: ZeRO stage1,2 と 3の違い
DeepSpeed ZeRO
ZeRO 1, 2では通信量(Communication Volume)自体は変化していない
→ DPから冗長性だけを取り除き、メモリ使用効率を上昇させた 👍
ZeRO Stage3では通信量は1.5倍になっている。(今回は説明省略(参考: ZeRO論文))
今回紹介したのは、ZeRO-DPの機能について
ZeROとはZeRO-DPとZeRO-Rを合わせた総称であるため、これだけではない (説明略)
43

44. Appendix: ZeROだけを利用する場合
DeepSpeed ZeRO
HuggingFace TrainerとのIntegrationがあり簡単に使うには便利
→ 新しいモデルを試してみる用途にマッチしている (wandbのintegarationもアリ)
しかし、長期学習には向かない？(FLOPsが出にくいのでコスト増)
実装コスト vs Computation Cost
お金で解決できる場合は、Trainerクラス + ZeROでも問題ない
transformersを利用するだけで良いので楽
44

45. Appendix: FSDPを利用する場合
PyTorch FSDP
詳細な実装についてはDeepSpeed ZeROとは異なる点も多いが、使用感は似ている
llama-recipesなどFSDPを利用したライブラリが最近増えている
HuggingFace Transformersの.from_pretrained() を利用しつつも、Trainerクラスのような大部分を隠蔽し
た構造ではない学習コードになりがち
→ 利点: カスタマイズ性向上 (wandbのloggingを自分で調節できる 👍)
欠点: 実装コスト、FSDPには不安定な機能も多い
45

46. Appendix: LLM-jp 13B 学習概要
Megatron-DeepSpeedを利用
● ZeRO DP=12 (stage1)
● DeepSpeed PP=8 (1F1B)
● TP=1
東大mdx 12node(96基 A100 40GB) 使用
当時2023年9月頃は、mdxにネットワーク上の問題が
あり145TFLOPs程度しか出ていない
(現在は改善)
46

47. Appendix: LLM-jp 13Bでの分散学習
Megatron-DeepSpeedで効率的に分散学習を行なうには？
1. 学習環境のGPUにおいてDPが最大になる & CUDA out of memoryが発生しないTP, PPを発見
2. micro-batch数やactivation-checkpointingをselective, fullにするなどして最もFLOPsが出る組み合わ
せを探す
a. PP数を増加させてmicro-batch数を増やしてみる
b. activation-checkpointを有効化して、TP, PPが減らせないか確かめる
3. 学習環境における最適な組み合わせを発見する
→ GPUノード間通信環境(inter-node connect)やGPU間通信の充実度によって最適なconfigは変化する
47

48. Appendix: LLM-jp 13Bでのconfig探索
mdxにて最速のconfig探索をしている例
比較軸
● PP次元数
● TP次元数
● micro-batch数
● global batch数
num-micro-batchesが変化するため
48

49. Appendix: LLM-jp 13B での Tips
Tips
1. Megatron-LMやMegatron-DeepSpeedにはTensorboard用のlogging機能がdefaultでついている
→ 必要なメトリクスだけwandbに送信するようにする (オススメはできるだけ全部 loggingす
る)
1. 12ノード程度であれば2週間程度の学習でもノード不良は、ほぼ発生しない
2. FP16でもLoss Spikeなく学習できている
a. データの質や多様性による？
b. LLM-jp v1コーパスの質はそれほど良くないが、問題なく学習できた
49

50. Appendix: LLM-jp 175B 学習概要
ABCI LLM構築支援プロジェクトの一環
49ノード(392基) A100 40GBを利用
Megatron-DeepSpeedを使用
● ZeRO DP=2 stage1
● DeepSpeed PP=49(1F1B)
● TP=2
学習FLOPsは156TFLOPs
GPT-3 175Bの学習を行った
→ 72B Tokenほど学習した時点でLoss Spikeが発生
前のcheckpointからLRを下げて再スタートなどを行な
うも改善されず（時間的制約から断念）
50

51. Appendix: LLM-jp 175B 技術詳細
導入技術
1. FlashAttention v2
2. RoPE(Rotary Positional Embedding)
3. BF16
4. embedding層, lm-headをTransformer Blockとみなしメモリの偏在消費を解消する技術
未導入技術
1. Sequence Length Warmup
2. skip-batch
51

52. Appendix: LLM-jp 175B ノード不良
ノード不良について
49ノードを使用していた際に遭遇
1ノードでも調子が悪くなると全
体が停止(右図)
ただしjob自体はfailedしないので
気づきにくい
→ 自動検知scriptを作成
52

53. Appendix: LLM-jp 175B ノード不良の監視
ノード不良について
wandbのrun statusと
最終update時刻から停止を検知
右図のように実験管理者のslackに
通知を飛ばしていた
53

54. Appendix: LLM-jp 175B 直面した問題達
Megatron-DeepSpeedやDeepSpeedの問題
1. checkpoint load後のLoss乖離問題
checkpointをloadするとLoss curveが以前と乖離する
問題 (次スライド)
1. BF16 checkpoint save & load 問題
DeepSpeed BF16 Optimizerを利用するとcheckpoint
のsave, loadができない問題
→ DeepSpeed側で条件分岐にミスがあったので修正
54

55. Appendix: LLM-jp 175B Loss乖離問題
Megatron-DeepSpeedやDeepSpeedの問題
1. checkpoint load後のLoss乖離問題
checkpointをloadするとLoss curveが以前と乖離する
問題 (checkpoint周りにバグ？ 右図)
Microsoft DeepSpeedチームと議論しつつ、原因究明
のための実験を多数行なう
→ evaluation時にFlashAttentionのdropout確率を0で
完全に上書きしてしまっていた。(Megatron-DeepSpeed)
現在は修正済み ✅
55

56. Appendix: Loss Spikeに関する考察
LLM-jp 175B Loss Spike原因
PaLMの論文に
右のような記述あり
要点は、
1. skip batchは有効
2. 200-500 batch skip
3. データの質が原因と
は限らない
56

57. Appendix: Loss Spikeに関する考察
LLM-jp 175B Loss Spike原因
そのため、skip batch を行
なうと途中のcheckpointか
ら復活可能？
ただし、継続学習できる余
地を残してLearning Rate
decayをあえてしていない
のが最も致命的と思われる
57

58. Appendix: checkpoint convertの困難さ
checkpoint convert script作成
original HF checkpoint
Megatron TP=N, PP=M
Megatron TP=1, PP=1
converted HF checkpoint
完全一致
Megatron-LLaMA, Megatron-LLM
などにconvert scriptは存在するが
問題あり
Megatron-LM公式
HF → Megatron converterで変換した
モデルをHFに再変換すると問題発生
→ 自作する必要に迫られる
他のタスクを行いながら、3週間ほど
でTensorレベルで完全一致する
convert scriptを開発
58