Core Concepts
整数行列乗算ユニットを用いることで、浮動小数点行列乗算ユニットを用いるよりも、精度、メモリ消費量、演算量の面で理論的な利点がある。また、NVIDIA消費者向けGPUでも、浮動小数点行列乗算よりも高速に実行できることを示す。さらに、量子回路シミュレーションへの適用により、精度を維持したまま最大4.33倍の高速化を達成した。
Abstract
本研究は、深層学習ハードウェアで用いられる整数行列乗算ユニット(IMMU)を、高性能計算(HPC)アプリケーションに活用する方法を探っている。
まず、Ozaki方式を用いることで、IMMUを使うことの理論的な利点を示した。具体的には以下の3点が挙げられる:
1スライスあたりの有効ビット数が多く、同じ精度を維持するためのスプリット数を減らせる
作業メモリの使用量を50%~75%削減できる
行列乗算の回数を2乗に減らせる
次に、NVIDIA GPUのINT8 Tensor Coresを用いてOzaki方式を実装し、精度、スループット、電力効率を評価した。入力行列の指数分布幅によってはトレードオフがあるものの、NVIDIA消費者向けGPUでは、cuBLASのDGEMMや既存のFP16 Tensor Coresを用いた実装よりも最大6倍高速に動作することを示した。
さらに、量子回路シミュレーションへの適用を行い、FP64の精度を維持したまま最大4.33倍の高速化を達成した。実用性を高めるため、自動的な精度調整メカニズムも実装した。
Stats
入力行列の指数分布幅が狭い場合(𝜙= 0.1)、INT8x9の誤差はDGEMMよりも小さい
入力行列の指数分布幅が広い場合(𝜙= 4)、INT8x11とINT8x13の誤差はDGEMMと同等かそれ以下
A行列とその近似逆行列A†の乗算では、INT8x𝑋の誤差はDGEMMよりも小さい