論文では語られないLLM開発において重要なこと Swallow Projectを通して

Transcript

1. 論文では語られないLLM開発において重要なこと - Swallow Projectを通して Institute of Science Tokyo, Swallow Project 藤井 一喜
2. 自己紹介 藤井 一喜 / Kazuki Fujii 東京科学大学 情報理工学院 修士課程2年 Preferred Networks PLaMo インターン • Swallow Project 事前学習、データ高品質化を担当 • 研究興味 ◦ 大規模モデルの分散学習、低精度計算による高速化 ◦ データ品質改善によるLLMの性能改善 2
3. Swallow Project オープンモデルを利用して 日本語に強い 大規模言語モデル (LLM) を研究開発する • 東京科学大学、産総研の共同研究 ◦ 岡崎研究室 (NLP) ◦ 横田研究室 (HPC, ML) • 数多くの日本語LLMをリリース ◦ これまでに12モデルシリーズ を公開 ▪ Llama 3, Gemma 2ベースが最新 ◦ 産業応用 にも利用 3
4. 継続事前学習 (Continual Pre-Training) 4 Llama-3, Gemma-2 ... Open LLMs 日本語 + 英語 + コード Llama-3-Swallow Gemma-2-Swallow • 利点 o Open LLMの力を利用できる o 比較的低コストで学習可能 • 欠点 o アーキテクチャの制約 o 元モデルのLicenseに縛られる 課題 • 破滅的忘却 • 英語スコアの低下
5. 目次 1. Swallow Projectについて • SwallowProjectと継続事前学習 2. LLM開発において重要なこと (論文ではあまり言及されない ) • 学習環境整備 ◦ 環境整備 ◦ ライブラリ整備 • モデル評価 ◦ 評価フレームワーク実装 ◦ in domain評価 • 学習速度の向上と安定性 ◦ 論文で提案されている手法のうち採用できるのは数少ない • 学習データの品質改善 • スケジュール管理 5
6. 学習環境整備 (環境整備) • ストレージ ◦ GPUの性能は上昇傾向 (Ampere → Hopper → Blackwell) → forward, backwardが高速に → 学習データのpre-fetchが間に合わないと FLOP/s の低下を招く (下図) ◦ 例) AWS FSx for Lustre, TSUBAME SSD Lustre ◦ on the flyでtokenizeする場合はそれも考慮した実装を検討 • ネットワーク ◦ 通信速度が出るか nccl-testsで確認 ▪ All Gather, Reduce Scatter, Send, Recv ◦ 物理的に使用ノードが近接しているか ▪ 例) AWS Placement Group 6 ストレージ負荷により FLOP/sが一時的に低下する様子 (ABCI 2.0)
7. 学習環境整備 (学習ライブラリ ) 学習ライブラリはどこかが壊れているのが 普通 • OSSライブラリ: NeMo, Megatron-LM, transformersなど • 自作ライブラリ: llm-recipes, moe-recipesなど → unit testと小規模モデルを学習する CIを作成しメンテナンス (それでも壊れる) Swallow Projectで費やした時間の 1/3はライブラリメンテナンス 例) TransformerEngine v2.0, v2.1 にて tensor parallel size > 1 + tp comm overlapを利用すると、モデルの収束性が悪化する • 小規模なモデルは差が見れない問題 • 依存ライブラリのAPI更新による問題 (PyTorch, TE) → 大規模な検証していない commitは信頼できない 7 TransformerEngineに起因する問題により lossの推移が変化する例 → https://github.com/NVIDIA/TransformerEngine/issues/1616
8. 学習環境整備 (学習ライブラリ ) mcore v0.8.0 mcore v0.9.0 CUDA Toolkit 12.1, cuDNN 8.9.7, nccl 2.20.5 torch 2.3.1, TE v1.9 日本語コーパスの改善 ablation実験 ablation実験は同一の環境で実施 前versionとのloss, grad normの差を検証 FLOP/s, forward, backward時間も検証 時間軸 CUDA Toolkit 12.4, cuDNN 9.1.0, nccl 2.21.5 torch 2.5.1, TE v1.11 コードコーパスの改善 ablation実験 実験環境ごとに tag 管理(Github) Swallow Project管理下 総計 42 tags (Pre-training)
9. 学習環境整備 (依存ライブラリ ) • PyTorchをpip installすることは、ほぼない ◦ Swallow Projectのリリースモデルの学習には、 独自build PyTorchかNGC PyTorchを利用 ◦ 配布されているPyTorchはdistributed backendにmpiを指定できない ▪ Tensor Parallel Communication Overlapの実装方法によっては問題になる ◦ NGC PyTorchの最新版が常に使えるとは限らない ▪ BugのあるTEでversion lock ▪ PyTorchが新しすぎてAPIの変更に他の依存ライブラリが追いついていない ◦ 高速化のためには、 PyPIでbuildされているversionとは異なる依存関係が必要 ▪ 特定のCUDA Toolkit, cuDNN, nccl以上のversionを要求するなど ▪ 例) Blockwise Quantization (CUDA Toolkit >= 12.9) 9 FYI: • PyTorch build: https://zenn.dev/turing_motors/articles/3a434d046bbf48 • NGC PyTorch version lock: https://zenn.dev/turing_motors/articles/fdea210f6ed6e9
10. モデル評価 (評価ライブラリ実装 ) • 評価ライブラリのメンテナンスは重要 ◦ 他組織のモデルを正確に評価するのは難しい ▪ HF Fast Tokenizer (Sarashina) ▪ HumanEval(コード評価)における改行 (PLaMo-2) • 評価スコアを見るだけでは不十分 どのように間違っているのかの観察(目視)は必須 (1) Llama-3.1-Swallow-v0.3は、HumanEval(コード評価)において実行すら 出来ないコードまたは、存在しない関数を読み出すことが多数確認された → SwalloCode-v1へ活かされた (2) MMLUなどの選択肢タスクにおけるスコアは、 1. Standard, 2. Continuationの両方を行うべし → 知識が足らないのか (学習データが低品 質 or データ量が足らない ) 、選択肢を選ぶという能力の問題なのか 10 An Empirical Study of Mamba-based Language Models: https://arxiv.org/abs/2406.07887
11. モデル評価 (in-domain評価) GSM8K等のいくつかのベンチマークにはtestセットの他にtrainセットが存在する trainセットを学習の終盤に利用するまたは、trainセットから似た形式の合成データを作成し学習に使用 → GSM8Kや似た形式であるMGSMのスコアのみが上昇 (他の数学タスクであるBBH, MATHは改善しない ) 新たに採用したデータが”in-domain”評価によりスコア改善を引き起こしたのか、それとも”本質的”な改善なのかは 評価スコアを眺めるだけでは分からない → 他組織のモデルとの比較はより困難 (データセットが非公開な場合 ) 対象ベンチマークでの判断を放棄する場合は、”in-domain”評価にあたる可能性があるデータを採用できるが... (評価タスクは貴重であるため、簡単に使い潰すわけにはいかない) (例) llm-jp jasterを学習に利用した後、llm-jp-evalで評価することで高スコアを得る → リリースモデルの性能を極限まで良くしたいという意味では良いが... (他のモデルとの比較は不可) 11
12. 学習速度向上と安定性 (低精度) • 低精度化 (FP8) 手法の多くはモデル構築には導入できない FP8 Delayed Scaling[1] は、175Bモデルで200Bトークン学習する限りでは問題ない しかし、Llama-3からの継続事前学習では、100Bトークン未満でも学習の不安定化、下流タスクの性能低下を招く[2] FP8をMLPに導入することにより高速化できるスループットは実用的には高々1.3倍 (Hopper) → 複雑な提案法の多くはbaselineのチューニング不足を解消するとFP8を導入した恩恵を相殺 してしまう → NVIDIAのBlockwise FP8, MXFP8[3]を導入することが最も実用的 (学会にAcceptされる手法ではなく、NVIDIAがGTCで発表する手法の方が実用に耐える) 12 [1] FP8 Format for Deep Learning https://arxiv.org/abs/2209.05433 [2] Balancing Speed and Stability: The Trade-offs of FP8 vs. BF16 Training in LLMs https://arxiv.org/abs/2411.08719 [3] Microscaling Data Formats for Deep Learning https://arxiv.org/abs/2310.10537 Llama-3-70Bからの継続事前学習 (BF16/FP8 DelayedScaling)
13. 学習速度向上と安定性 (分散学習) 13 A G A G A G 0 1 2 forward 2 1 0 R S R S R S backward A G 0 1 2 A G A G 2 R S 1 0 R S R S time save 通信と計算の Overlapは、下流タスクへの副作用なしで学習を高速化する ↔ ライブラリの複雑性は悪化 (例: TEにおけるTensor Parallel Comm overlapバグ) データ並列における通信と計算の Overlap (GTC24で発表された技術 )
14. 学習データの品質改善 継続事前学習による 下流タスク性能比較 により学習データの品質判定を実施 (タスクLossなど学習データの品質判定手法は存在するが...) 改行の有無など些細な点で下流タスク性能が大きく上下することがあるため、できるだけモデル開発に近い設 定でablationを実施 コードコーパスの改善実験だけ 13実験も実施 → モデル開発に必要な計算資源量の3倍以上の資源を消費 = 上手くいかない実験を含めるとablation実験だけで5倍以上 実験の流れ (概要) 1. 実験立案 2. データ準備 (コーパス構築) 3. 実験設定の決定 4. 実験実施 5. checkpoint 評価 6. 検証 → 実験実施以外にも計算資源は消費される 14 コードコーパスablation実験におけるデータ比率 [4] [4] Rewriting Pre-Training Data Boosts LLM Performance in Math and Code https://arxiv.org/abs/2505.02881
15. 学習データの品質改善 (実例: コード品質改善 ) Llama-3.1-8Bからの継続学習でStack-v2等のOpen Codeデータを利用してもコード生成タスクは改善しない (Task Aligned CodeQAデータの採用により post trainingの前借りを行う、または、 GPT-4等の生成結果を使った Distill QAデータを利用しない場合 ) → コードデータ(Pre-Training)の品質が課題 学習データを観察していると (1) Syntax Error (2) コード規約違反 等が散見された → ルールベースのフィルターではなく、compile, linter によるファイルターを実施しては？ ↓ 15 Llama-3.1-Swallow-8B-v0.1 (2024/10) のスコア推移 HumanEval, JHumanEval: コード生成タスク 左: HumanEval, 右: HumanEval+ フィルタリング手法比較
16. 学習データの品質改善 (実例: コード品質改善 ) さらにコード品質を改善するためにLLMを利用した書き換え(Rewriting)を実施 SwallowCode-v1では、2段階(SGCR→SCOR)のRewritingを実施し、さらに大きくスコアを改善 (OpenCodeデータセット中最良) → Rewriting時のpromptは、評価結果を目視で観察することで決定 → Filter、Rewritingともに、”目視”での観察により改善案を起草 (= 闇雲に合成データを作ることは計算資源の無駄？ ) 高い下流タスク性能 ≠ 高品質コーパス post trainingの前借りになっていないか (QA形式に合成しないで検証 ) リークは存在しないか (HumanEval, HumanEval+のprompt, solutionと一致度検証) 下流タスクの形式に過度に alignさせていないか 16 SwallowCode Rewriting手法の比較 [5] [5] Rewriting Pre-Training Data Boosts LLM Performance in Math and Code https://arxiv.org/abs/2505.02881
17. スケジュール管理 • LLM開発はスケジュールの 管理の戦い ◦ 計算資源はまとめて調達 (ABCI 50ノード占有、AWS HyperPodなど) → 学習開始時期までに (1)学習ライブラリ (2)学習環境整備 (3)学習データ の準備 を完了している必要あり ◦ 大規模実験 (30B, 70B modelを数100 B token以上)は四半期に一度が限界 → 学習scriptで1.0と0.1を間違えるなどは許されない (余裕をもった準備が必要) • 開発寄りの作業と研究寄りの作業を両立する ◦ 次の大規模開発を見据えて (1)低精度 (2)分散学習 (3)学習データ に関して論文化できそうなアイデアを試す 　 ただし、モデル性能の向上に直接貢献しない、または開発に導入するには煩雑すぎるテーマは抱えられない • 継続事前学習に関する研究実験は認可制 ◦ 計算資源を要するため、横田研究室、岡崎研究室、産総研の全体meetingで議論後、実験を投入 ◦ ライブラリの変更、動作検証(実際の学習も含む)、評価チームへのcheckpointの受け渡しなども事前相談 17
18. まとめ 1. Swallow Projectについて • SwallowProjectと継続事前学習 2. LLM開発において重要なこと (論文ではあまり言及されない ) • 学習環境整備 ◦ 環境整備 ◦ ライブラリ整備 • モデル評価 ◦ 評価フレームワーク実装 ◦ in domain評価 • 学習速度の向上と安定性 ◦ 論文で提案されている手法のうち採用できるのは数少ない • 学習データの品質改善 • スケジュール管理 18