MXNorm: MXFPブロックスケールの再利用による効率的なテンソル正規化

要約: 行列乗算の性能は長年、ディープラーニングワークロードをスケールさせる際の主要なボトルネックであり、これはますます低精度の数値形式を用いる新しいアクセラレータの設計を促してきました。しかし、行列乗算の性能の向上は、リダクションと要素ごとの計算の性能向上をはるかに上回っており、それらは依然として高精度で実行されています。本研究では、MXNorm を RMSNorm のドロップイン置換として提案します。MXFP8 キャストの一部として計算されるブロックスケールのみを用いて RMS を推定し、正規化に必要なリダクションの規模を32分の1に縮小します。125M、1B、8B のパラメータを持つ Llama 3 モデルの事前学習の検証を行い、MXFP8 の matmuls を用いた RMSNorm をベースラインとした場合と比較して訓練精度の損失が最小限であることを確認しました。また、torch.compile のみを用いた MXNorm が RMSNorm に対して最大で 2.4 倍のカーネル速度向上を示し、MXFP8 の Llama 3 8B トランスフォーマ層で 1.3% のスピードアップ、NVFP4 で 2.6% のスピードアップに対応します。