NVIDIA RTX PRO 6000 Blackwell Max-Q 上で尾崎スキーム I を動かしてみる

はじめに

こんにちは、Fixstars でアルバイトをしている堀毛晴輝です。

近年の生成 AI、大規模言語モデル需要の高さから、そのようなモデル学習・推論に特化したアクセラレータが多数、開発・販売されています。
例えば、最近の NVIDIA GPU は、従来の科学技術計算向けの高精度演算器よりも、AI 向けの低精度演算器の搭載を重視するようになり、この流れは今後も加速すると考えられています。
Fixstars Techblog にて、これまでに検証を続けてきたワークステーション向け最新 Blackwell GPU: NVIDIA RTX PRO 6000 Blackwell Max-Q (以下 6000 Blackwell Max-Q) も、この流れに乗ったGPUの一つです。

• NVIDIA RTX PRO 6000 Blackwell Max-Q はどのようなGPUなのか？
• NVIDIA RTX PRO 6000 Blackwell Max-Q 上でBLASベンチマークを動かしてみる
• NVIDIA RTX PRO 6000 Blackwell Max-Q 上での LLM fine-tuning の実行速度・メモリ効率調査
• NVIDIA Blackwell 上での LoRA チューニングに向けたプロファイリングと最適化

具体的には、6000 Blackwell Max-Q の INT8 性能 (1755.7 TOPS) はデータセンタ向け NVIDIA H100 GPU の INT8 性能 (3,958 TOPS) の約半分に匹敵し (疎性あり)、ワークステーション向け GPU としては低精度演算において強みを持っています。
一方で、6000 Blackwell Max-Q の倍精度 (FP64) 演算性能は 1.71 TFLOPS と、先述した INT8 性能の約 1000 分の 1 です。
そのため、高い精度を要求するアプリケーションは計算速度に不安が残ります。

このような、高精度演算性能よりも低精度演算性能が重視される昨今の GPU トレンドにおいて、低精度行列積ユニットを用いて高精度行列積をエミュレーションする手法「尾崎スキーム」が注目されています。
この手法は 2012 年に初めて提案され^[1]、昨年には従来とは異なるアプローチによる「尾崎スキーム II」^[2] が提案されました。
本記事では、これらを区別するために前者の手法を「尾崎スキーム I」と呼びます。

尾崎スキームは昨今の低精度/高精度演算性能比の増大に伴い、高速な高精度行列積手法として重要性を増しています。
そうした背景を受け、CUDA Toolkit 13.0 Update 2 以降の cuBLAS では尾崎スキーム I による行列積エミュレーションが正式にサポートされました。
このアップデートは、ハードウェアベンダが行列積性能を強化するために、演算器を増強するのではなくソフトウェアによるエミュレーション技術によって性能を増強することを公式に認めたという、インパクトの大きい内容になっています。

そこで本記事では、倍精度演算性能が限定的な 6000 Blackwell Max-Q において、cuBLAS の尾崎スキーム I による行列積エミュレーションオプションを有効化し、計算速度と精度がどのように変化するかを検証します。
なお、6000 Blackwell Max-Q は、ローカルAI運用向けGPUワークステーション Fixstars AIStation (Fixstars AIStation について)の搭載機を使用しました。

尾崎スキーム I のアルゴリズム

尾崎スキーム I ^[1:1] は、高精度行列積を複数回の低精度行列積でエミュレートする手法です。
2012 年に初めて提案され、近年では低精度行列積に INT8 行列積ユニットを用いる改良 ^[3] などがなされています。

尾崎スキーム I では、高精度浮動小数点数行列 $A, B$ について大まかに以下のステップで行列積 $A \times B$ を計算します。

1. 高精度浮動小数点数行列 $A, B$ を低精度浮動小数点数行列の和:
   

   $$A \risingdotseq A_1 + A_2 + \cdots + A_k, \quad B \risingdotseq B_1 + B_2 + \cdots + B_\ell$$
   
   で近似する。
   ここで、行列 $A$ と $B$ の分割方法は、各スライス同士の積 $A_i \times B_j$ を計算する際に丸め誤差が発生しないビット幅に収まるよう決定する。

2. 行列積 $A \times B$ を、
   

   $$A \times B \risingdotseq \sum_{i=1}^k \sum_{j=1}^\ell A_i \times B_j$$
   
   で近似する。各 $A_i \times B_j$ の行列積計算には、Tensor Core 等の最適化された演算ユニットを用いて高速に計算する。
   ここで、低精度浮動小数点数行列同士の積 $A_i \times B_j$ の結果は高精度浮動小数点数行列として求めることに注意。

実際にはステップ 2 におけるスライス同士の行列積を計算する際に仮数部を切り出し、INT8 行列積を用いて処理しています。

ステップ 1 における行列のスライス数を増やすことにより精度が向上しますが、スライス数の 2 乗に従って計算量が増加します。
行列内の最大値と最小値の比率 (ダイナミックレンジ) が狭い行列であれば少ないスライス数で元の行列を表現することができ、高速に行列積を計算できます。

cuBLAS での尾崎スキーム I の有効化

CUDA Toolkit 13.0 Update 2 以降の cuBLAS では、以下の設定を行うことで尾崎スキーム I を用いた高精度行列積エミュレーションを有効化できます (CUDA Toolkit 13.0 Update 2 - Release Notes)。
用途に応じて、環境変数による一括適用か、API による明示的な制御かを選択できます。

1. 環境変数による設定

   環境変数を設定することで、アプリケーションコードを変更せずにエミュレーションを実行する方法です。

   export CUBLAS_EMULATE_DOUBLE_PRECISION=1
   

   これにより、cublasDgemm() 関数で尾崎スキーム I を用いたエミュレーションが実行されるようになります。

   参照: https://docs.nvidia.com/cuda/cublas/index.html?highlight=CUBLAS_EMULATE_DOUBLE_PRECISION#floating-point-emulation

2. cublasSetMathMode による明示的な指定

   アプリケーション内で cublasSetMathMode 関数を呼び出し、行列積の計算方法をエミュレーションに変更する方法です。

   // エミュレーションを行うように Math Mode を指定
   cublasSetMathMode(handle, CUBLAS_FP64_EMULATED_FIXEDPOINT_MATH); 
   // その後の DGEMM 呼び出しに対して適用される！
   cublasDgemm(handle, transa, transb, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
   

   参照: https://docs.nvidia.com/cuda/cublas/index.html?highlight=cublasSetMathMode#cublassetmathmode

注意点として、尾崎スキーム I による FP64 行列積のエミュレーションは、cuBLAS 内部のエミュレーション戦略に基づき実行するか否か判断されます。
上記の設定でエミュレーションを有効化した場合でも、デフォルトの設定では cuBLAS が演算コストを評価し、パフォーマンス上のメリットがないと判断した場合には標準の DGEMM が用いられます。
エミュレーション可能な場合に、パフォーマンスの評価に関わらずエミュレーションを積極的に採用するためには、以下に示すようにエミュレーション戦略を EAGER に指定します。

export CUBLAS_EMULATION_STRATEGY=eager

もしくは

cublasSetEmulationStrategy(handle, CUBLAS_EMULATION_STRATEGY_EAGER);

とすることで、エミュレーション戦略を指定できます。

参照: https://docs.nvidia.com/cuda/cublas/index.html?highlight=CUBLAS_EMULATION_STRATEGY#cublasemulationstrategy-t

また、先述したとおり尾崎スキーム I を用いた行列積エミュレーションでは、スライス数によって行列積精度や実行時間が変化します。
cuBLAS では、精度とコストのトレードオフを制御するために、Dynamic Mantissa Control と Fixed Mantissa Control の 2 種類の方法が提供されています。

参照: https://docs.nvidia.com/cuda/cublas/index.html?highlight=CUBLAS_EMULATE_DOUBLE_PRECISION#fixed-point

• Dynamic Mantissa Control (デフォルト):

  FP64 と同等以上の精度が出るように、cuBLAS が入力行列を見て自動でスライス数を決定します。
  必要なスライス数がライブラリのデフォルト値 (あるいはユーザー指定値) を超えた場合には、ネイティブの FP64 演算へと動的に切り替えます。

• Fixed Mantissa Control:

  計算に使用するビット数をユーザーが直接指定します。
  指定したビット数 (スライス数) で FP64 と同等以上の精度が出ない場合でもそのまま計算が行われます。

これらは cublasSetFixedPointEmulationMantissaControl() 関数を介して設定できます。

ベンチマーク

今回の検証に使用した計算環境は以下の通りです。

• GPU: NVIDIA RTX PRO 6000 Blackwell Max-Q
• ホスト OS: Ubuntu 24.04.3 LTS
• コンテナイメージ: nvidia/cuda:13.1.0-devel-ubuntu24.04
• CUDA Toolkit: 13.1.80

検証方法

今回のベンチマークでは尾崎スキーム II の実装がなされている GEMMul8 のテストコードをベースに、行列積の演算部分を cuBLAS 標準の倍精度行列積 (DGEMM) と、尾崎スキーム I エミュレーションを有効化した倍精度行列積 (DGEMM-OS1) に書き換えた測定環境を用いています。

ベンチマークに使用する行列の各要素は、正の実数パラメータ $\phi$ を用いて $(\text{rand} - 0.5) \cdot \exp(\phi \cdot \text{randn})$ で生成します。
$\text{rand} \in (0, 1]$ は一様分布から取得された乱数、$\text{randn} \in \mathbb{R}$ は標準正規分布から取得された乱数とします。
$\phi$ が大きくなるほど、ダイナミックレンジが広くなり、精度よく行列積を計算することが難しくなります。
特に尾崎スキーム I では、ダイナミックレンジが広い場合にスライス数を増やす必要があり、スライス数が不足すると、表現しきれなかった小さな値の成分が消えてしまい、精度が悪化してしまいます。
今回のベンチマークでは $\phi = 0.5, 1.0, 2.0, 3.0$ それぞれについて検証します。

比較する 2 つの関数の設定は以下の通りです。

• DGEMM: 標準の倍精度行列積

  CUBLAS_DEFAULT_MATH を指定することで、特殊な高速化をせず通常の倍精度行列積を実行します (DEFAULT 戦略は比較のため記述しています)。

  cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH);
  cublasSetEmulationStrategy(handle, CUBLAS_EMULATION_STRATEGY_DEFAULT);
  cublasDgemm(handle, transa, transb, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
  

• DGEMM-OS1: エミュレーションを有効にした倍精度行列積

  FIXEDPOINT_MATH を指定し尾崎スキーム I を用いたエミュレーションを有効化します。
  エミュレーションを必ず行うために、EAGER 戦略を指定します。

  cublasSetMathMode(handle, CUBLAS_FP64_EMULATED_FIXEDPOINT_MATH);
  cublasSetEmulationStrategy(handle, CUBLAS_EMULATION_STRATEGY_EAGER);
  cublasDgemm(handle, transa, transb, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
  

  なお、今回はデフォルトの Dynamic Mantissa Control を使用します。

実行速度

$n = 1024, 2048, 4096, 8192, 16384$ について、先述した方法で $n \times n$ 行列を 2 つ生成し、これらの行列積にかかった時間を計測します。
計測した時間は、比較のため FLOPS (Floating-point Operations Per Second) に換算して評価します。

ここで、FLOPS とは、コンピュータが 1 秒間に実行できる浮動小数点演算の回数を示す指標です。
サイズ $n$ の正方行列の積に必要な演算は乗算が $n^3$ 回、加算が $n^3 - n^2$ 回で合計 $2 n^3 - n^2$ 回ですが、$n$ が大きいため $2 n^3$ 回の演算と近似します。
行列積の FLOPS は、この演算回数 $2 n^3$ 回を行列積の実行時間(秒)で割ることで計算されます。
また、1 TFLOPS = $10^{12}$ FLOPS です。

計測の結果、標準の DGEMM は行列サイズに関わらず約 1.5 TFLOPS で推移しました。
この値は、FP64 の理論性能である 1.71 TFLOPS に対し約 88 % に相当します。
ハードウェアの仕様上ネイティブな倍精度演算ではこの数値がほぼ限界で、実用的な速度とは言えません。

一方、尾崎スキーム I によるエミュレーションを有効化した DGEMM-OS1 では、$\phi$ の値により性能が大きく変化しました。
ダイナミックレンジが比較的小さい $\phi = 0.5, 1.0$ では、最大で約 8 TFLOPS に達していて、これは標準の DGEMM と比較して約 5 倍の実行速度です。
とくに、この性能は FP64 の理論性能 1.71 TFLOPS を大幅に上回っており、本来倍精度演算が弱い 6000 Blackwell Max-Q においても、INT8 行列積ユニットを転用することでネイティブな計算能力を超えた高速化が実現できています。

ダイナミックレンジが大きい $\phi = 2.0, 3.0$ では、通常の DGEMM とほぼ同等の性能となっています。
これは、ダイナミックレンジが広い場合には精度の維持のためにスライス数が増加してしまい、INT8 行列積を計算する回数が増えてしまうからと考察できます。

計算精度

DGEMM, DGEMM-OS1 それぞれについて、疑似 4 倍精度 (Double-Double) を用いて計算した行列積の結果を真の値として最大相対誤差を評価します。
行列サイズは $1024 \times k$ および $k \times 1024$ とし、$k = 1024, 2048, 4096, 8192, 16384$ の各条件で精度を比較しました。

$\phi = 0.5, 1.0, 2.0$, $k < 16384$ の条件下では、標準の DGEMM の誤差が $10^{-8}$ から $10^{-10}$ 程度であるのに対し、　DGEMM-OS1 は $10^{-13}$ から $10^{-14}$ であり、非常に精度よく行列積を計算できています。
$k = 16384$ では精度が悪化しましたが、これは行列サイズの増加に伴い、INT8 行列積の過程でオーバーフローを防ぐため 1 スライス当たりの情報量が減ってしまったことが要因と考えられます。
しかし、この場合でも標準の DGEMM より高い精度で行列積が計算できています。

また、$\phi = 3.0$ とダイナミックレンジが大きい場合でも、DGEMM-OS1 の誤差は標準の DGEMM と同等程度の精度を確保できています。
以上の結果から、尾崎スキーム I による行列積エミュレーションは計算速度を劇的に向上させつつ、実用上十分な計算精度を備えている手法であるといえるでしょう。

まとめ

本記事では、6000 Blackwell Max-Q において、FP64 行列積を INT8 行列積でエミュレートする尾崎スキーム I の cuBLAS による実装を検証しました。
倍精度演算性能が 1.71 TFLOPS と低い 6000 Blackwell Max-Q であっても、CUDA 13.0 Update 2 からサポートされた尾崎スキーム I によるエミュレーション機能を有効化するだけで、倍精度行列積の性能が引き上げられることを確認できました。
入力行列のダイナミックレンジが適切な範囲内であれば、標準の DGEMM と比較して約 5 倍の最大約 8 TFLOPS が達成でき、これは 6000 Blackwell Max-Q の FP64 理論性能を大幅に超えています。
6000 Blackwell Max-Q が持つ豊富な INT8 演算リソースを倍精度行列積へ活用できており、低精度演算に特化した最新アーキテクチャのポテンシャルを、高精度演算の領域でも引き出せています。

また、本記事で紹介した尾崎スキーム I の実装は cuBLAS によって提供されているため、高度なアルゴリズムを独自に実装する必要はなく、cuBLAS から簡単に呼び出せます。
環境変数や Math Mode のフラグを変更するだけで、既存の倍精度行列積を用いたプログラムを修正なしに高速化することが可能です。

近年の GPU アーキテクチャが低精度演算へ最適化されていく傾向にある中、尾崎スキームのような低精度演算を活用して高精度演算をエミュレーションする手法は、極めて実用性の高いアプローチといえるでしょう。

脚注

1. Generalization of error-free transformation for matrix multiplication and its application, https://www.jstage.jst.go.jp/article/nolta/4/1/4_2/_article

2. Ozaki Scheme II: A GEMM-oriented emulation of floating-point matrix multiplication using an integer modular technique, https://arxiv.org/html/2504.08009v1

3. DGEMM on Integer Matrix Multiplication Unit, https://arxiv.org/html/2306.11975v4