Improved Scaling for Fast Mode of Ozaki Scheme II

技術要約：Ozaki Scheme II の Fast Mode におけるスケーリングの改善

問題提起
Ozaki Scheme II は、中国剰余定理（CRT）に基づき、低精度整数ハードウェア（例：INT8 Tensor Core）を用いて高精度行列乗算（例：FP64/FP32）をエミュレートする手法である。このスキームでは、高精度入力行列を整数行列へとスケーリングする必要があるが、その際、整数積を CRT による復元が可能な特定の閾値（$P/2$）を超えないようにしなければならない。もしこの閾値を超えると、CRT による復元が失敗し、壊滅的な精度低下を招く。

本スキームは、スケーリング係数（$\mu$ および $\nu$）を決定するために、2つのモードを提供している：

1. Accurate Mode（高精度モード）: 補助的な INT8 行列乗算（GEMM）を使用してスケーリング係数を推定する。高い精度を保証するが、大幅な計算オーバーヘッドが発生する。
2. Fast Mode（高速モード）: コーシー＝シュワルツの不等式を用いて、補助的な GEMM を行わずにスケーリング係数を推定する。高いスループットを提供するが、歴史的に精度が低いという課題があった。

本論文は、既存の Fast Mode のスケーリング公式における重大な欠陥を特定している。それは、**スケール不変性（scale invariance）**の欠如である。入力行列にスカラー定数 $\alpha$ を乗じた場合、結果として得られるスケーリング済み整数行列の実効ビット幅が予測不能に変化してしまう。

• 大きな $\alpha$ に対しては、実効ビット幅が減少し、精度低下を引き起こす。
• 小さな $\alpha$ に対しては、スケーリングされた整数の積が CRT 閾値（$P/2$）を超え、完全な復元失敗を招く。
  全要素が 1 の行列およびランダム行列を用いた数値実験により、特定のスカラー範囲において、Accurate Mode は安定している一方で、Fast Mode は正しい結果を復元できないことが示された。

手法
著者らは、コーシー＝シュワルツの不等式を介して CRT の一意性の条件から直接導出された、Fast Mode 用の改訂スケーリング公式を提案している。

1. 理論的導出:
   著者らは、スケーリング係数 $\mu_i$ および $\nu_j$ に関する条件を再定式化している。従来のヒューリスティックの代わりに、以下の条件を定義する：
   $$2\mu_i \|a_{i,:}\|_2 \|b_{:,j}\|_2 \nu_j \leq P - 1$$
   GPU 数学 API が提供する log2f 関数の絶対誤差境界を考慮しつつ、底 2 の対数を用いて $\mu_i$ と $\nu_j$ について解くことで、新しい公式を導出する：
   $$\mu_i := 2^{\lfloor P'_{prop} - 0.5 \cdot (\widehat{\log_2}(\sum a_{ih}^2) + 4u_{32}) \rfloor}$$
   ここで、$P'_{prop}$ は事前計算された定数であり、$\widehat{\log_2}$ は有界な対数近似である。

2. スケール不変性の証明:
   著者らは、提案する公式がスケール不変であることを理論的に証明している。入力行列 $A$ が $\alpha = 2^s$ でスケールされた場合、スケーリングされた整数行列の実効ビット幅の計算において $\log_2(\alpha)$ が相殺される。その結果、入力の大きさに依存せずビット幅が一定に保たれ、従来の Fast Mode で観察された精度低下や CRT 失敗が防止される。

3. 実装:
   提案手法（OS II-prop）は、GEMMul8 ライブラリの Fast Mode におけるスケーリング公式を置き換えるものである。これは、元の Fast Mode の $O(mk)$および$O(kn)$ の計算量特性を維持しており、追加のオーバーヘッドを生じさせず、かつ補助的な GEMM を必要とする Accurate Mode よりも高速である。

主な貢献

• 限界の特定: 既存の Fast Mode のスケーリング公式にはスケール不変性が欠けており、特定の入力スケーリングに対して精度低下や CRT 復元失敗を招くことを示した。
• 改訂された公式: スケール不変性を保証し、CRT の一意性の条件が常に満たされる新しいスケーリング公式を提案した。
• ゼロ・オーバーヘッド: 新しい公式は、元の Fast Mode と比較して追加の計算コストを導入しない。
• 包括的な評価: NVIDIA GH200 GPU を用い、倍精度（DGEMM）および単精度（SGEMM）の行列乗算の両方について評価を行った。

結果
ランダム行列を用い、要素の大きさの広がり（$\phi$）および行列次元（$k \in \{1024, 16384\}$）を変化させて実験を行った。

• 精度: 提案手法（OS II-prop）は、すべての $\phi$ の値において Accurate Mode（OS II-accu）と同等の精度を達成した。対照的に、従来の Fast Mode（OS II-fast）は、大きな $\phi$ の値（不良設定の場合）において顕著な精度低下と失敗を示した。
• スループット: OS II-prop は、元の Fast Mode と同等のスループットを維持している。モジュラスの数が少ない特定のシナリオでは、スケーリングされた整数内のゼロ要素が少なくなるため、OS II-prop は元の Fast Mode よりもわずかに遅くなるが、Accurate Mode に対しては大幅に優れたスループットを示す。
• トレードオフ: 精度とスループットのトレードオフ分析において、OS II-prop は OS II-accu を圧倒（同等の精度でより高いスループット）し、かつ OS II-fast の精度限界を克服している。これは、幅広い条件下で cuBLAS を精度とスループットの両面で上回ることができる唯一の手法である。

意義
本論文は、提案手法が性能上の利点を損なうことなく、Ozaki Scheme II の Fast Mode における根本的な精度制限を解決することを主張している。スケール不変性を確保することで、入力のスケーリングに関わらず CRT による復元が成功することを保証し、精度と性能の優れたバランスを提供している。これにより、ユーザーは Fast Mode の速度で、Accurate Mode の信頼性を備えた高精度な結果を低精度ハードウェア上で得ることができ、Accurate Mode のスループット制約と、従来の Fast Mode の精度制約を効果的に克服できる。

著者らは、現在の研究は DGEMM および SGEMM 用の INT8 行列エンジンに焦点を当てているが、このアプローチは FP8 エンジンや複素行列乗算（ZGEMM/CGEMM）にも適用可能であり、これらは今後の課題としている。