## Memo ## Memo with LLM ### 論文情報 - **論文のタイトル**: FP8-LM: Training FP8 Large Language Models - **著者と所属**: Houwen Peng, Kan Wu, Yixuan Wei, Guoshuai Zhao, Yuxiang Yang, Ze Liu, Yifan Xiong, Ziyue Yang, Bolin Ni, Jingcheng Hu, Ruihang Li, Miaosen Zhang, Chen Li, Jia Ning, Ruizhe Wang, Zheng Zhang, Shuguang Liu, Joe Chau, Han Hu, Peng Cheng (Microsoft Azure and Microsoft Research) - **カンファレンス/ジャーナル名**: arXiv (2023年10月27日初版、2023年12月19日改訂版) - **発表年**: 2023年 ### 論文概要 本論文では、大規模言語モデル（[[LLM]]）の効率的な学習のためのFP8低ビットデータフォーマットを探求し、勾配やオプティマイザ状態などの大部分の変数が低精度データフォーマットを使用してもモデル精度を損なうことなくハイパーパラメータの変更も不要であることを示している。実験結果では、[[H100]] GPUプラットフォーム上でGPT-175Bモデルの学習において、提案するFP8混合精度学習フレームワークが実メモリ使用量を39%削減し、広く採用されているBF16フレームワーク（Megatron-LM）よりも75%高速化を実現している。 ### 詳細解説 #### 問題設定 **入力**: 大規模言語モデル（125M〜175Bパラメータ）の学習において、従来のFP32全精度またはFP16/BF16混合精度学習 **出力**: FP8混合精度を用いた効率的な学習手法 **必要なデータ**: CommonCrawl、The Pile、C4、OpenWebText、CC-NEWS、CC-Stories、Redpajama、Wikipediaなどのオープンソース言語データセット（最大100Bトークン） #### 提案手法 本論文では3段階のFP8最適化レベル（O1、O2、O3）を提案している： **O1最適化**: FP8勾配とAllReduce通信を実装。自動スケーリング技術により、勾配のアンダーフローとオーバーフローを解決： ``` g'_i = μ · g_i (式3) ``` ここで、μは動的に変更される自動スケーリング係数 **O2最適化**: Adamオプティマイザにおける補助テンソルの低ビット化。精度分離（precision decoupling）により： - 一次勾配モーメント：FP8 - 二次勾配モーメント：FP16 - マスター重み：FP16（テンソルスケーリング付き） - メモリ消費量：パラメータあたり6バイト（従来16バイトから2.6倍削減） **O3最適化**: 分散並列学習にFP8を適用。テンソル並列、パイプライン並列、シーケンス並列、ZeROオプティマイザに対応 #### 新規性 従来のNVIDIA [[Transformer Engine]]（TE）がFP8をlinear層の計算のみに限定していたのに対し、本手法は： 1. FP8勾配とAllReduce通信の実現（自動スケーリングによるアンダーフロー/オーバーフロー解決） 2. 初のFP8オプティマイザの実装（精度分離による安定化） 3. FP8分散並列学習の包括的サポート（テンソル、パイプライン、シーケンス並列、ZeRO） 4. 学習全体のFP8計算、ストレージ、通信への統合 #### 実験設定 **データセット**: CommonCrawl、The Pile、C4、OpenWebText等の組み合わせ（最大100Bトークン） **モデル**: GPTスタイルのTransformerアーキテクチャ（125M〜175Bパラメータ） **評価指標**: - 学習損失の比較（BF16との対比） - 下流タスク性能（HellaSwag、Lambada、BoolQ、PIQA、COPA、Winogrande、Arc、OpenBookQA） - システム性能（メモリ使用量、スループット、通信量） **ハードウェア**: Azure NDv5 H100 GPUプラットフォーム #### 実験結果 **精度保持**: GPT-7B、13B、175Bすべてでの学習損失曲線がBF16とほぼ重複し、下流タスクでの同等性能を確認 **システム性能向上**: - GPT-175B: メモリ使用量39%削減、BF16比75%高速化、TE比37%高速化 - 重み関連通信量63-65%削減 - モデルFLOPs使用率（MFU）34.2%達成 **ファインチューニング適用**: - 指示調整（SFT）: メモリ14%削減、スループット27%向上 - 人間フィードバックからの強化学習（RLHF）: モデル重み32%削減、オプティマイザ状態62%削減 **アブレーション研究**: - 自動スケーリングがpre-scaling、post-scalingより優秀（SNR向上、アンダーフロー/オーバーフロー率削減） - マスター重みはFP16以上の精度が必要（FP8では性能劣化） - 二次勾配モーメントはFP16必須（FP8では学習発散） 論文の提案手法により、大規模言語モデルの学習コストを大幅に削減しながら精度を維持する実用的な解決策を実現している。 ## Abstract 本論文では、大規模言語モデル（LLM）の効率的な学習のためのFP8低ビットデータフォーマットを探求する。我々の主要な知見は、勾配やオプティマイザ状態などのLLM学習における大部分の変数が、モデル精度を損なうことなく、またハイパーパラメータの変更を必要とすることなく、低精度データフォーマットを使用できるということである。具体的に、我々はLLM学習のための新しいFP8自動混合精度フレームワークを提案する。このフレームワークは、LLMの混合精度および[[分散深層学習|分散並列学習]]を合理化するために、FP8利用の3つのレベルを提供する。これは、8ビット勾配、オプティマイザ状態、分散学習を段階的に組み込む。実験結果は、H100 GPUプラットフォーム上でのGPT-175Bモデルの学習において、我々のFP8混合精度学習フレームワークが実メモリ使用量の著しい39%削減を達成しただけでなく、広く採用されているBF16フレームワーク（すなわち、Megatron-LM）よりも75%高速に実行され、Nvidia Transformer Engineを37%上回る速度を実現したことを示している。これにより、大規模基盤モデルの学習コストが大幅に削減される。さらに、我々のFP8混合精度学習手法は汎用的である。これは、LLMの指示調整や人間フィードバックからの強化学習などの他のタスクにシームレスに適用でき、ファインチューニング費用の節約を提供する。我々のFP8低精度学習フレームワークは、aka.ms/MS.AMPにてオープンソース化されている。