情報処理学会第 86 回全国大会

5J-02

大規模言語モデルの分散並列学習
藤井 一喜 1 横田 理央 2
1 東京工業大学 情報理工学院
2 東京工業大学

はじめに

1

2.2

学術情報国際センター

学習ライブラリ

近年、様々な研究機関や企業で大規模言語モデル学習には Megatron-LM 3）を利用した。3D 並列化(LLM) の開発が行われている。これらのモデルは、(3D Parallelism) が利用できる点、分散学習設定のセ人間に近い言語理解能力と生成能力、さまざまな分クションで触れる工夫が実装されている点、ライブ野への適用可能性を示し、大きな注目を集めていラリとしての成熟度を加味し選定を行った。

る。大規模言語モデルの効率的な学習には、分散並2.3列学習が不可欠であり、学習効率に大きく影響する重要な要素である。本論文では、Llama 2 から継続

分散学習設定本 研 究 で は Llama 2 7B、13B、70B か ら 継 続 事

事前学習を行った日本語モデル「Swallow 1）」の開前 学 習 を 行 っ た。効 率 的 な 学 習 を 行 う た め に、発に使用した分散並列学習手法について報告する。

デ ー タ 並 列（Data Parallelism）、テ ン ソ ル 並 列 化

また学習に使用した実装を公開する。 2）継続事前（Tensor Parallelism）、パイプライン並列化（Pipeline

学習の有効性に関する議論は、別の論文を参照されParallelism）を統合した 3D 並列化（3D Parallelism）

たい [1]。

を採用し、高い計算効率と効率的なメモリ利用を目指した。表 1 に各モデルサイズにおける分散学習設

分散学習

2定を示す。他にも、以下に挙げる工夫を取り入れ大規模言語モデルを 1 つの GPU で学習することは、GPU メモリの制約と学習に要する時間の両方から困難である。GPU メモリの点では、最新の H100
80GB を利用しても、今回学習した 7B のモデルを学習することは困難である。また、仮にモデルパラメータと勾配、オプティマイザの状態 (OptimizerState) が 1 枚の GPU に収まったとしても、1 枚のた。 4）。
表 1 モデルパラメータごとの分散学習設定: DP、TP、PP、SP はそれぞれデータ並列 (Data Parallelism)、テンソル 並 列 (Tensor Parallelism)、パ イ プ ラ イ ン 並 列 (PipelineParallelism)、シ ー ク エ ン ス 並 列 (Sequence Parallelism) を表す。

Params DPTPPPSPDistributed Optimizer

GPU で は 学 習 を 完 了 す る た め に 非 現 実 的 な 時 間

7B

1622

✓

✓

を要することになる。そのため、今回の学習では、

13B

824✓

✓

データ並列とモデル並列を併用した分散並列学習を70B488✓✓採用した。

2.1

効率的なメモリ消費 Megatron-LM の Distributed Optimizer を用いて、オプティマイザの状態（opti-学習環境

学習には、産総研の AI 橋渡しクラウド (ABCI) をmizer state）をデータ並列プロセス間に分散配置利用した。混合精度（bfloat16）を採用し NVIDIAし、冗長性を排除することで、必要なメモリ使A100 ノードを複数台使用し分散並列学習を行っ用量を削減した。Distributed Optimizer は Reduce

た。各ノードは NVIDIA A100 40GB GPU を 8 基をScatter と All Gather を利用し、効率的に通信を搭載し、ノード間は InﬁniBand HDR にて接続されて3） https://github.com/NVIDIA/Megatron-LM4） いくつかの実験では、実験途中で使用ノード数の変更が行われたため、データ並列数が 2 倍もしくは 1/2 倍されている期間が存在する。これは本論文で説明する実験以外にも同時に進めた実験があるためで、限られた計算資源・期間で有望そうな設定を優先してモデル構築を進めたからである。

いる。

1） https://tokyotech-llm.github.io/swallow-llama2） https://github.com/rioyokotalab/Megatron-Llama2

1-61

Copyright 2024 Information Processing Society of Japan.
All Rights Reserved.

情報処理学会第 86 回全国大会行うため、通常のデータ並列と同じ通信コスト

表2にも関わらずメモリ削減が行える。
トポロジーを考慮した 3D マッピング

Params ノード数 TFLOPS/GPU

3D 並列化において、Transformer ブロックはパイプライン並列により複数の GPU に分散配置され、さらにテンソル並列により層内のパラメーターは分散

7B

4134

43.0 %

13B

8143

45.8 %

70B

32158

50.6%

であり、本研究においてもその有効性が示された。

れているように、通信を多く必要とする分散手

しかしながら、3D 並列化は実装が難しく、あらゆ法のワーカー（テンソル並列ワーカー）はノーる言語モデルの学習に向いているとは言えない。

ド内に配置した。これは、ノード内の通信は今後は、3D 並列化に加え FSDP(Fully Sharded DataParallel) のようにデータ並列だけで 10B 前後のモデためである。また、データ並列の勾配平均化のルを負担が少なく学習できるように、ライブラリ開ための通信を考慮して、データ並列ワーカーも発を進め、LLM の開発に貢献したいと考えている。

可能な限りノード内に配置した。パイプライン
並列は他の並列化手法と比較して通信量が少な

謝辞

い P2P(Point-to-Point) 通信であるため、パイプラ

LLM の継続事前学習の実験では、国立研究開発法インステージはノード間で配置した。

人産業技術総合研究所が構築・運用する AI 橋渡し
クラウド（ABCI: AI Bridging Cloud Infrastructure）のメモリ効率化のための 1F1B の採用 1F1B(one forward pass followed by one backward pass) のパイプ「大規模言語モデル構築支援プログラム」の支援をライン並列である PipeDream-Flush[3] を利用す受けました。本研究は、JST、CREST、JPMJCR2112

ることで、パイプラインステージ数以下のマイの支援を受けたものである。

クロバッチ数しか活性 (activation) を必要としないようにし、GPipe[4] よりもメモリ効率を上昇させた。
Sequence Parallelism による並列化

テンソル並列は Self-Attention や MLP ブロックを並列化するが、レイヤーノーム (Layer-Norms) とドロップアウト (Dropouts) は並列化しないため、テンソル並列プロセス間でこれらが冗長にメモリ上に存在する。これを効率化するために、シークエンス並列 (Sequence Parallelism) [5] を利用した。
シークエンス並列は、テンソル並列と同時に用いることで、通信コストのオーバーヘッドなしにメモリの効率化を行える。

計算効率

実際に学習を行った際のモデルパラメータごとのTFLOPS を表 2 に示す。実行効率として、70B のモデルにおいて 50%以上を記録しており、効率的に学習を行えたと考える。

3

実行効率

配置される。この際、Megatron-LM[2] で提案さNVLink によりノード間通信よりも高速である

2.4

モデルパラメータごとの TFLOPS/GPU

おわりに

本研究では、最大で 70B のモデルの効率的な学習を行った。3D 並列化は大規模言語モデルの事前学習、継続事前学習を効率的に行うために重要な技術

1-62

参考文献
[1] 藤井一喜, 中村泰士, Mengsay Loem, 服部翔, 飯田大貴,
水木栄, 平井翔太, 大井聖也, 横田理央, 岡崎直観. 継続事前学習による日本語に強い大規模言語モデルの構
築. 言語処理学会第 30 回年次大会 (NLP2024), 2024.
[2] Deepak Narayanan, Mohammad Shoeybi, Jared Casper,Patrick LeGresley, Mostofa Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie
Bernauer, Bryan Catanzaro, Amar Phanishayee, and MateiZaharia. Eﬃcient Large-Scale Language Model Training
on GPU Clusters Using Megatron-LM. In Proceedingsof the International Conference for High Performance
Computing, Networking, Storage and Analysis, pp. 1–15, 2021.
[3] Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie
Chen, and Matei Zaharia. Memory-Eﬃcient PipelineParallel DNN Training. In International Conference onMachine Learning, pp. 7937–7947. PMLR, 2021.
[4] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan
Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan
Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Eﬃcient
Training of Giant Neural Networks using Pipeline Parallelism. Advances in neural information processing
systems, Vol. 32, , 2019.
[5] Vijay Anand Korthikanti, Jared Casper, Sangkug Lym,
Lawrence McAfee, Michael Andersch, Mohammad
Shoeybi, and Bryan Catanzaro. Reducing Activation Recomputation in Large Transformer Models. Proceedingsof Machine Learning and Systems, Vol. 5, , 2023.

Copyright 2024 Information Processing Society of Japan.
All Rights Reserved.