マイクロソフトリサーチが初のFP4精度大規模モデル訓練フレームワークを発表、効率が大幅に向上

1月29日、マイクロソフトリサーチは初のFP4精度による大規模モデル訓練フレームワークを発表しました。同じハイパーパラメータ設定で、FP8やBF16に匹敵する訓練効果を得ることができます。これにより、必要なストレージと計算リソースが減少します。
モデルサイズと訓練効果
この方法で訓練されたモデルの最大規模は130億パラメータに達し、訓練トークン数は数百億に達します。現在はFP8で模擬されたFP4を使用していますが、実際のFP4を使用すれば、効果がさらに向上します。
FP4のシミュレーションと訓練方法
研究当初、FP4をネイティブにサポートするハードウェアは存在しなかったため、研究チームはFP8のTensorCoreを使用してFP4をシミュレートしました。具体的には、1.3B、7B、13BのLLaMAモデルを用いて、0から1000万トークンの訓練過程で、FP4訓練はBF16の損失曲線とほぼ一致しました。
FP4精度訓練を実現するため、研究チームはカスタマイズされたFP4行列乗算（GeMM）CUDAカーネルを採用しました。カーネルでは、最初にFP16でFP4のAとB行列を共有メモリに読み込んで変換し、その後FP4で分割行列乗算を行い、最終的にFP16で中間結果を還元して、FP16形式の出力行列を得ます。
量子化戦略とハードウェア適合
このフレームワークはE2M1のFP4フォーマットを採用しており、指数に2ビット、尾数に1ビット、符号ビットに1ビットを使用し、合計4ビットです。このフォーマットは、現在主流のMLアクセラレーションチップの量子計算ユニット設計に適合するように選ばれました。重み行列Wと活性化行列Aには異なる粒度の量子化戦略が採用され、FP4の行列乗算での加速効果を最大限に引き出します。
順伝播と逆伝播の革新
モデルの順伝播が始まると、フレームワークは各線形層の重み行列Wと入力活性化行列Aを同時にFP4量子化します。量子化の際、最初に行列内の数値をスケーリングとオフセットを行い、FP4が表現できる範囲にマッピングした後、テーブル検索を使って最寄りのFP4離散値に丸めます。
逆伝播では、研究チームは計算効率と勾配の正確性を保証する新しい微分可能な勾配推定方法を提案しました。また、活性化行列内の「外れ値」問題に対しては、「外れ値のピーク削減と補償」戦略を提案しました。
混合精度設計と応用
一部の段階で、研究チームは混合精度設計を採用し、例えば勾配通信にはFP8を、オプティマイザの状態保存にはFP16を使用しています。システムの他の部分、例えば非行列乗算操作やLoss ScalingなどではFP16を使用しています。これらの設計は、訓練数値の安定性を保ちながら、計算とストレージのオーバーヘッドを可能な限り減らします。
研究チームと背景
このフレームワークは、マイクロソフトアジア研究院とSIGMAチームによって開発され、全ての研究者は中国系です。第一著者のRuizhe Wangは中科大の博士課程の学生で、現在MSRAでインターン中、研究分野は低精度量子化です。中科大研究部門の部長、查正軍教授もこのプロジェクトに参加しています。通信著者はMSRAのシニアリサーチマネージャーである程鵬氏と首席研究マネージャーのYeyun Gong氏です。MSRAの卓越した科学者である郭百寧氏も本プロジェクトに参加しています。
まとめ
マイクロソフトリサーチが発表したFP4精度の大規模モデル訓練フレームワークは、革新的な量子化戦略と混合精度設計により、より効率的なモデル訓練を実現しました。このフレームワークは訓練効率を高めるだけでなく、計算とストレージリソースの需要を大幅に削減し、大規模モデル訓練の新たな道を切り開きました。
論文リンク：arXiv
参考リンク：Twitter