TVMからTensorCoreを使ってみる
こんにちは、NTTの森永です。
前回の記事では、TVMの自動チューニングや量子化など基本的な機能を紹介しました。今回は、よりTVMの性能を引き出す手段としてTensor Coreを活用した畳み込み実装を使う方法を紹介します。
Tensor Coreとは？
Tensor Coreとは、NVIDIA GPUのVolta以降の世代に搭載される行列演算ユニットです。Tensor Coreは、専用のAPIを使って、4x4の行列A, B, Cに対してAB+Cを1サイクルで高速に計算します。TuringやAmpereなどの最新寄りのアーキテクチャでは、FP32やFP16に加え、INT8, INT4やさらにはINT1といった低精度の演算にTensor Coreを使うことで高い演算性能を発揮でき、Tensor Coreは学習のみならず推論でも有効です。
NVIDIA GPUにおけるTVMの性能については、少し前のアーキテクチャでは、AutoTVMを使ってcuDNNやTensorRTを上回る性能を達成👊といったことも可能だったようです。しかし、前回の記事の結果のように最近のアーキテクチャでは、TVMの性能がNVIDIA謹製のツールたちに再び勝てなくなっているように見えました。そんな中、TVMでもCuDNNなしでTensor Coreを使う畳み込み演算の実装のパッチがマージされていたため、これを活用しTensorRTやcuDNNと比べてみようというのが、本記事の趣旨です。なお、評価環境にはTesla T4を搭載したAWSのg4dnインスタンスを使用します。
TVMからTensor Coreを使ってみる
TVMにおけるTensor Coreの対応状況は、ざっくりと述べると以下のようになっています。本記事で注目しているのは、①の実装です。
① CUDAバックエンド：TVM独自の実装が使われ、TVMの自動チューニング（AutoTVM）が適用可能です。後述する設定を行うと、FP32とFP16の時にTensor Coreが使われます。FP32のTensor Coreには、（AB+C）のAとBをFP16に下げて結果をFP32で取るため僅かに精度が落ちます。
② cuDNNバックエンド：前回紹介した-libs=cudnnを指定する方法です。FP16の時にTensor Coreが使われます。各計算の最適化はcuDNNに委ねるため、AutoTVMを使うことはできません。
CUDAバックエンドには、現在、畳み込み層と密結合層にTensorCoreを使う実装が存在します。畳み込みの実装には通常のconv2d_nhwc_tensorcoreとWinograd版のconv2d_nhwc_winograd_tensorcoreの２種類が存在します。いずれにせよ畳み込みにTensorCoreを使うには、畳み込みのレイアウトをNHWCに変えてあげるだけでOKです。
ただし、各実装についてshapeやバッチサイズに制約があり、例えば通常の畳み込み実装（conv2d_nhwc_tensorcore）は、(バッチサイズ, 入力チャンネル数, 出力チャンネル数) が(16, 16, 16), (8, 16, 32) , (32, 16, 8)の倍数の時しか有効にできません。
Tensor Core実装の効果
早速AutoTVMを適用して、Tensor Core実装の効果を細かくみてみましょう。前回紹介したTVMのデバッガ機能を活用し、最も多くの層でTensor Coreが有効になるバッチサイズ32のResNet-50に対して、Tensor Core実装のCUDAバックエンドとcuDNNバックエンドで各畳み込みレイヤーごとの処理時間を比較してみました。
FP32では、全てのレイヤーでTVMのTensor Core実装がcuDNN を上回り、高速化率では、cuDNNの1.1~3.3倍が達成できました。全畳み込み層の合計実行時間でもcuDNNの1.69倍となっていました。
前章で述べた通り、TVMのcuDNNバックエンドは、FP32の時にはTensor Coreを使用しません。FP32でTensor Coreを使う場合、精度劣化が少し懸念されるので、TVMがcuDNNの設定でそれを回避しているためです。CUDAバックエンドの方だけTensor Coreを使う上記の比較は少しフェアではないので、TVM内でのcuDNNの設定に手を加え、FP32の時もcuDNN経由からTensor Coreが使えるようにし、再測定してみました。
黄色のfp32_cudnn_tensorcoreが、文字通りFP32でTensor Coreを有効にしたcuDNNバックエンドです。後半の層では少し速くなっていますが、全体での実行時間もfp32_cudnnとほぼ同程度でした。FP32では、cuDNNバックエンドよりもCUDAバックエンドの方が上手にTensor Coreを使えていそうです。
FP16でも結果を見てみます。FP16の方は、精度劣化の懸念がなく、特に何もせずともcuDNNもTensor Coreを活用する設定になるので、安心です。
FP16では、合計実行時間こそTensor Core実装が1.4倍ほど速かったものの、いくつかのレイヤーでcuDNNに少し劣っています。特に、チャンネル数が増える後半のレイヤーには現在の実装よりも良いスケジュール実装（AutoTVMの自動最適化のためのテンプレート）が存在しそうで、この部分は今後の改善に期待といったところでしょうか。
畳み込みTensor Core実装の性能評価 （対TensorRT）
Tensor Core実装のおかげで大幅に性能が向上していることがわかったところで、NVIDIA製の最適化コンパイラであるTensorRTと性能を比べてみます。引き続き、バッチサイズ32でResNet-50を用い、全体での実行時間をTensorRTと比較してみます。
FP32では、Tensor Core実装のおかげでTVMがTensorRTを上回る性能を発揮できました。但し、FP32のTensorRTではTensor Coreは使われないので、精度劣化を許容すればの結果となります。一方、FP16では、TensorRTには一回り劣っており、前章同様にスケジュールに改善の余地があることを示唆しています。また、一番右のINT8でのTensorRTの結果のように、Turing以降の世代のTensor CoreではINT8以下の精度での演算性能の強力さが特徴の一つであるため、INT8等の量子化対応は今後のTensor Core実装における重要課題と言えるでしょう。
推論ではリクエストを即時処理したいケースは少なくないため、もちろんバッチサイズ1も重要です。しかし、Tesla T4において、Tensor Coreが有効にできないバッチサイズ1でのTVMの最も速い使い方は現状、cuDNNバックエンドとなります。これでは、以下のようにTensorRTの性能を大きく下回ります。
各Tensor Core実装のバッチサイズやshapeの対応範囲を広げるのも、TVMがTensorRTに対抗していくには今後必須そうです。
終わりに
TVMのTensor Core実装は、多くのレイヤーでcuDNNに優っており、ポテンシャルが高いと感じました。ただし、使えるshapeや精度が限られており、実用的に対応していると言うには、もう少し見守る必要がある段階のようです。実際、Tensor Core対応の今後は、RFCで以下のような計画を挙げています。
- 使用できるバッチサイズやチャンネル数の制約をなくす
- Winograd_TensorCoreの実装 (FP16, INT8)
- 畳み込み層のINT8対応
- 密結合層のINT8対応
このうち、2番のfp16のWinograd_tensorcoreについては既に実装済みであるため、本記事で課題として挙げたshapeの制約の解除や各種演算のINT8対応が行われ、より実用的になっていくと予想されます。性能の改善の余地についてもTensorRT等を意識し活発に議論されているため、TVMのTensor Core対応は暫く追いかけていきたいところです。
ところで、AutoTVMについても、改善版のAutoTVM v 2.0（通称Ansor）の提案がなされています。従来のAutoTVMでは各ハードウェアに詳しい実装者が書いたスケジュールが必要だったのに対し、Ansorではスケジュールの生成まで自動で行うようです。また、膨大な時間がかかっていたチューニングも効率が改善され、大幅に時間短縮されるとのことで、とても楽しみです。ただし、制約の多い専用APIを使うTensor Core実装について、効率の良いスケジュールを自動で組むのは難しいようにも思えるので、AnsorとTensor Coreの関わり方については個人的に今後注目しています。
NTTでは、共にOSSコミュニティ活動を行う仲間を募集しています。興味のある方は、新卒中途問わず、ぜひ採用情報をチェックしてください！
