Ascend NPU上での言語モデル事前学習に向けた HiFloat4 形式

要旨: 大規模基盤モデルは現代の機械学習の中核となっており、性能はモデルサイズとデータ量に応じて概ね予測可能な形でスケールします。しかし、このようなモデルの学習とデプロイには大きな計算コストとメモリコストがかかるため、低精度学習手法の開発が動機づけられています。近年の研究では、4-bit 浮動小数点（FP4）形式――MXFP4 や NVFP4 など――を大規模言語モデル（LLM）の線形 GEMM 演算に適用できることが示されており、高精度のベースラインと比べて計算スループットとメモリ効率が最大で 4 倍向上することが報告されています。本研究では、Huawei Ascend NPU 向けに最近提案された HiFloat4 FP4 形式を調査し、大規模学習設定において MXFP4 と体系的に比較します。すべての実験は Ascend NPU クラスター上で実施し、線形およびエキスパート GEMM 演算はすべて FP4 精度で実行します。密（dense）なアーキテクチャ（例: Pangu や LLaMA 系モデル）と、混合エキスパート（MoE）モデルの両方を評価します。ここでは、標準の線形層とエキスパート固有の GEMM の両方が FP4 で動作します。さらに、FP4 学習向けに最適化された安定化手法を検討し、数値的な劣化を大幅に低減します。その結果、4-bit 計算の効率的な利点を維持しつつ、相対誤差をフル精度ベースラインに対して 1% 以内に保つことができます。本研究の結果は、NPUs における FP4 学習について包括的な実証的調査を提供し、大規模な密モデルおよび MoE モデルにおける FP4 形式間の実用上のトレードオフを明らかにします。