> [!abstract] 概要 > 本論文では、大規模言語モデル(LLM)の効率的な訓練を目的として FP8 低ビット数値形式を探索する。著者らの核心的な洞察は、LLM 訓練における大多数の変数——勾配やオプティマイザ状態など——が、モデル精度を損なうことなくかつハイパーパラメータを変更することなく、低精度数値形式を利用できるという点である。具体的には、LLM 訓練のための新しい FP8 自動混合精度フレームワークを提案する。このフレームワークは、LLM 向けの混合精度・分散並列訓練を効率化するために 3 段階の FP8 活用レベルを提供する。8 ビット勾配、オプティマイザ状態、分散学習を段階的・加算的に取り込む設計である。実験結果では、H100 GPU 上での GPT-175B モデル訓練において、提案 FP8 混合精度訓練フレームワークが実メモリ使用量の 39% 削減と、広く採用されている BF16 フレームワーク(Megatron-LM)との比較で 75% 高速化を達成し、NVIDIA Transformer Engine を 37% 上回ることが示された。これにより大規模基盤モデルの訓練コストが大幅に削減される。さらに、提案 FP8 混合精度訓練手法は汎用性が高く、LLM インストラクション・チューニングや人間フィードバックによる強化学習など他のタスクにシームレスに適用でき、ファインチューニングコストの削減に貢献する。FP8 低精度訓練フレームワークはオープンソースとして aka.ms/MS.AMP で公開されている。 ## 論文情報 | 項目 | 内容 | |---|---| | 論文タイトル | 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) | | 掲載媒体 | arXiv:2310.18313v2 | | 公開日 | 2023-12-19(v2) | | URL | https://arxiv.org/abs/2310.18313 | | コード | https://aka.ms/MS.AMP | ## 概要 本論文は LLM 訓練における FP8(8 ビット浮動小数点)の活用を BF16 混合精度の代替として体系化した初の大規模実証研究である。FP8 を GEMM 計算のみに限定していた先行研究(NVIDIA Transformer Engine)を超え、勾配通信・オプティマイザ状態・分散並列の全段階へ FP8 を浸透させることで、端末間での高速化・メモリ削減を実現した。 3 つの最適化レベルを定義する: - **L1**: FP8 線形レイヤ(GEMM 計算のみ TE と同等) - **L2**: + FP8 オプティマイザ(勾配 FP8 + Adam 第一次モーメント FP8 + マスタ重み FP16) - **L3**: + FP8 分散並列訓練(テンソル/シーケンス並列 + FP8 ZeRO) ## 問題設定 LLM の訓練コストは巨大で、GPT-3 175B は数千 petaflop/s-日を消費する。低精度訓練は速度・メモリ・通信コストの三側面で訓練コスト削減の有望な方向であるが、FP8 の狭い動的範囲(S1E4M3 で max 448、S1E5M2 で max 57,344)はアンダーフロー・オーバーフローによる数値不安定性をもたらす。 既存手法の限界: - NVIDIA Transformer Engine: GEMM 計算のみ FP8 化。勾配とマスタ重みは FP16/FP32 のまま。端末間の速度向上・メモリ削減は限定的 - BF16 混合精度(Megatron-LM 等): 安定しているが計算コストが高く、集団通信は FP32 を使用 課題は 2 点: (1) FP8 の動的範囲外のデータがアンダーフロー/オーバーフローを起こす、(2) NCCL の現実装はテンソルごとのスケーリングファクタを考慮した all-reduce に対応していない。 ## 提案手法 ### FP8 勾配と all-reduce 通信 既存の事前スケーリング(pre-scaling)は GPU 数 N で割り込むためアンダーフローを生じやすく、事後スケーリング(post-scaling)は加算前に大きな値が溢れてオーバーフローを生じやすい。提案するオートスケーリングは動的な共有スカラー $\mu$ を導入し、FP8 最大値を超える値の比率が 0.001% を超えたら $\mu \leftarrow 1/2$、継続的に下回れば 1,000 ステップかけて指数的に $\mu \leftarrow 2$ と増大させる。 NCCL の all-reduce でテンソルごとのスケーリングファクタを扱う問題は、GPU ごとのスケーリングファクタの最小値 $s'_g = \min(s'_1, \ldots, s'_N)$ を事前に共有し、全 GPU が同一の $s'_g$ で FP8 量子化することで解決する。これにより通常の NCCL all-reduce がそのまま使える。 ### FP8 オプティマイザ 精度分離の原則: 勾配統計(1 次モーメント)は低精度を許容できるが、マスタ重みは高精度が必要。 - マスタ重み: FP16 + テンソルスケーリング(または FP32) - 勾配: FP8 - 第一次モーメント: FP8 (方向が重要で大きさの精度要件は低い) - 第二次モーメント: FP16 (二乗計算でアンダーフローしやすいため) 結果として従来の Adam(16 バイト/パラメータ)から 6 バイト/パラメータへ 2.6 倍のメモリ削減を達成。 ### FP8 分散並列訓練 テンソル並列とシーケンス並列において、シャード化された重みとアクティベーションを FP8 化して線形レイヤ計算と forward-backward の集団通信を FP8 で実行する(図 2)。シーケンス並列とテンソル並列の境界にある all-gather/reduce-scatter(operator $g$)にも FP8 変換を施し、通信量を半減させる。 ZeRO への FP8 適用: 従来の ZeRO はテンソルを分割して各 GPU に配るため、テンソルごとのスケーリングファクタが分断されて扱いに困る。提案 FP8 ZeRO ではテンソルを分割せず、アルゴリズム 1(貪欲アルゴリズム)でテンソル全体をメモリ残量が多い GPU へ優先配布する。これによりスケーリングファクタと FP8 テンソルを一体として扱え、かつ GPU 間の負荷均衡も得られる。 ## 新規性 1. FP8 計算・ストレージ・通信を訓練全段階へ浸透させた最初のフレームワーク 2. LLM 訓練用初の FP8 オプティマイザ(精度分離原則: 1 次 FP8 / 2 次 FP16 / マスタ重み FP16) 3. FP8 テンソルとスケーリングファクタを一体で扱う FP8 ZeRO 貪欲分散アルゴリズム 4. FP8 all-reduce のオートスケーリング機構(NCCL 互換の単一スカラー共有) 5. SFT / RLHF 等のファインチューニングへの汎用性の実証 ## 実験設定 - ハードウェア: Azure NDv5 H100 GPU プラットフォーム(80 GB HBM) - モデル規模: GPT 125M, 7B, 13B, 175B(decoder-only Transformer、RoPE + Flash Attention) - TP/PP 設定: 175B で TP=8, PP=4, SP 有効 - 訓練データ: CommonCrawl・The Pile・C4・Wikipedia 等の 100B トークン(175B は評価のため 40B トークン) - ベースライン: BF16(Megatron-LM)および NVIDIA Transformer Engine ## 実験結果 ### システム性能(表 5: H100 × 80GB) | モデル | 手法 | GPU メモリ | スループット | 重み関連通信量 | |---|---|---|---|---| | GPT-7B | BF16 | 69.6 GB | 159.2 samples/s | 37.2 GB | | GPT-7B | FP8 (TE) | 77.3 GB | 224.5 samples/s | 37.2 GB | | GPT-7B | FP8 (提案) BS=2 | **49.4 GB (−29%)** | 219.8 (+38%) | **13.9 GB (−63%)** | | GPT-7B | FP8 (提案) BS=4 | 69.3 GB | 230.5 (+45%) | 13.9 GB (−63%) | | GPT-13B | FP8 (提案) BS=2 | 48.9 GB (−28%) | 109.5 (+38%) | 12.4 GB (−64%) | | GPT-175B | FP8 (提案) BS=1 | **40.3 GB (−39%)** | 27.1 (+21%) | **8.2 GB (−65%)** | | GPT-175B | FP8 (提案) BS=4 | 57.7 GB | **39.3 (+75%)** | 8.2 GB (−65%) | - GPT-175B BS=4 でスループット BF16 比 +75%、MFU 34.2%(TE 比 37.3% 改善) - 重み関連通信量を 63〜65% 削減(BF16 の集団通信が FP32 を使うのに対し提案手法は FP8) - TE 比でメモリを GPT-175B で 42% 追加削減 ### モデル精度(表 2: ゼロショット) GPT-7B と GPT-13B の 9 ダウンストリームタスクで FP8 と BF16 の差は 0.6 ポイント以内(7B: FP8 58.0 vs BF16 58.4、13B: FP8 60.4 vs BF16 61.0)。訓練損失曲線は重なり、ハイパーパラメータ変更不要。 ### ファインチューニング(表 3, 4) - SFT(表 3): GPU メモリ −14%、スループット +27%、AlpacaEval 67.20(BF16: 66.15)、MT-Bench 5.70(BF16: 5.75) - RLHF(表 4): モデル重みメモリ −32%、オプティマイザ状態 −62%、AlpacaEval 72.42(BF16: 72.05)、MT-Bench 6.04(BF16: 6.16) ### アブレーション(表 6-8) - 精度分離実験(図 8): FP8 マスタ重みは発散しないが損失が高め。第二次モーメントを FP8 化(設定 #4)すると訓練損失が発散する - アクティベーション関連通信: シーケンス+テンソル並列で通信コストを GPT-13B で 59%、GPT-175B で 65% 削減 - FP8 ZeRO メモリ負荷均衡(表 8): TE 比 GPU 間最大メモリ差を大幅縮小しつつ総メモリも削減 ## 考察 精度分離の原則は実用上の重要な洞察を提供する。Adam の第一次モーメントは方向情報が主であり FP8 の量子化誤差を許容できるが、第二次モーメントは二乗計算でアンダーフロー感度が高くより高精度が必要である。マスタ重みの高精度保持は、更新量が非常に小さいときに精度損失を防ぐために不可欠である。 FP8 ZeRO のテンソル単位分散は、テンソルをサブパーティションに分割せず貪欲アルゴリズムで丸ごと割り当てることで、ZeRO 本来の「テンソルをシャードに分割」という設計と意図的に乖離する。これはスケーリングファクタの整合性保持を最優先とした設計判断であり、負荷均衡はアルゴリズムで補っている。 ## 強み - 勾配・オプティマイザ・分散の全段階への FP8 浸透という包括的アプローチ - ハイパーパラメータ無変更で BF16 同等精度を実証(7B〜175B の広い規模範囲) - SFT/RLHF 等のファインチューニングへの汎用性実証 - OSS 公開(aka.ms/MS.AMP) ## 弱点・課題 - 通信削減率(63〜65%)は FP8 の理論値(75%)より低く、システム転送損失が残る - FP8 ZeRO のテンソル単位分散はサブパーティション分割より実装が複雑で、極端なテンソルサイズ分布でメモリ偏りのリスクがある - GPT 系 decoder-only アーキテクチャのみの実証。エンコーダ系や MoE モデルへの適用性は未検証 - 現論文時点(2023 年末)では H100 専用(FP8 対応 GPU はまだ限定的) - FP8 の 2 フォーマット(E4M3 / E5M2)の選択はアブレーションが限定的