トップページ

最新記事

AmpereのTensorコアの話

これはrioyokotalab Advent Calendar 2020 2日目の記事です．

TF32について

Ampereに搭載されているTensorコア Gen 3ではTF32と呼ばれる浮動小数点数を用いることができます．
これは指数部8bit，仮数部10bitの計19bitの浮動小数点数です．
何が32bitだ，とお思いでしょうが，CUDAではfloatとして保持するため，メモリ使用量は32bitです．
要するに，普通に使う分にはメモリの節約には全くなりません．
NVIDIAはこれをcuBLASのGEMM等で使えるようにしており，たまに単精度行列積と謳ってその計算性能を表示しています．
しかし，TF32は仮数部が大きい，すなわち表現範囲可能範囲が広いだけで，仮数部はFP16 (IEEE Binary16)と同じであるため，精度は単精度の足元にも及びません．
FP16ではアンダーフローが起きやすい計算は多少は精度がマシになるかもしれませんが．

この記事はそんなTF32を自分のカーネル関数から使う場合の注意点の話です．

TF32のTensorコアでの使い方

以前からTensorコアはWMMA APIと呼ばれるAPIで利用可能なわけですが，TF32もこれを用いることで利用可能です．
ただし注意点があります．
NVIDIAの資料でこのようなものを見たことがある方も多いかと思います．

出典 : NVIDIA A100 Tensor コアGPU アーキテクチャ - NVIDIA

これを見ると，メモリ（本当はfragment/レジスタ）からFP32の行列データを読み，Tensorコアへ送ることで使えるようになる気になります．
しかし，これを信じて使うとアホほど計算精度が低いのです．
上述したとおりTF32はメモリ上ではfloatをして保持され，Tensorコアはこの内MSBから19bitのみを読み込みます．
これは丸めとしてはかなりおそまつなものです．
このため，Tensorコアへ送る前に適切な丸めを自分で行う必要があります．
要するに，先程の図に「TF32への適切な丸めでの変換」という処理が追加で必要なのです（"Format to TF32"はただの19bit切り取りを意味しているので）．

ではどう変換すればいいかと言いますと，PTXのcvt.rna.tf32.f32命令で変換できます．
この命令は零捨一入みたいな丸めを行います．
これはmma.hをincludeしている場合は__float_to_tf32関数がこれをwrapしているため，これを用いて以下のように変換できます． 

const float fp32 = 1.0f;
const float tf32 = __float_to_tf32(fp32);

これを用い，fragmentを作る前か後かで変換してあげれば，理論値通り仮数部10bit分程度の精度がでます．

おわりに

自分で変換を行ってねというのはCUDA Toolkit Documentationにかかれています．
今回の件以外にもNVIDIAの発表で用いられる図は基本的には信用しないほうがいいと個人的には思っています...．

記事作成日:2020-12-02