ZenDNN ライブラリによって実現される AMD EPYC™ プロセッサでの最適な推論性能

2022年11月、AMDは第4世代 AMD EPYC™ プロセッサを発売しました。第4世代AMD EPYC™ プロセッサには、AVX-512やVNNI命令セット拡張など、前世代を上回る多数のハードウェアに対する改良が施されており、推論性能の向上に適しています。しかし、ハードウェアはパズルの1ピースに過ぎず、ソフトウェアは基礎となるハードウェアを効果的に活用するための重要なコンポーネントです。

このプラグインは、AMDがフレームワークのアップストリームを通じて、ZenDNN最適化のコミュニティへのアクセス性を向上させるために継続的かつ重点的に取り組んでいることを示すものです。そして、TensorFlow v2.12 以上用のTensorFlow-ZenDNNプラグインの提供を新たに発表することを嬉しく思います。このプラグインは、AMD ZenDNNライブラリを使用した AMD EPYC™ CPU 上でのニューラルネットワーク推論を可能にします。

ZenDNN

GitHubからオープンソースで入手できるZenDNNは、 AMD EPYC™ CPU 向けに最適化された基本的なニューラルネットワークのビルディングブロックを含む、低レベルのAMDディープニューラルネットワークライブラリです。ZenDNNは、深層学習アプリケーションやフレームワークの開発者が、コンピュータ・ビジョン、自然言語処理、推薦システムなどのさまざまなワークロードにおいて、 AMD EPYC™ CPU での推論性能を向上させることを目的として構築されています。

TF-ZenDNN

ZenDNNを高レベルのAIフレームワークに統合し、使いやすさを実現しました。TF-ZenDNNと呼ばれるTensorFlowとの統合のプロトタイプは、特定のバージョンでTensorFlowリポジトリをフォークし、TensorFlowコードを直接修正することで実現されています。TF-ZenDNNは、AMDのZenDNN開発者向けリソースページ（下図1）から直接統合するためのバイナリパッケージとして入手でき、インストール方法は「TensorFlow + ZenDNNユーザーガイド」で確認できます。

図1. ZenDNNの開発者向けリソースページで公開されているZenDNN v4.0のバイナリパッケージを、本ブログではTF-ZenDNN直接統合版と呼びます。

TF-ZenDNNは、ネットワークレベルでグラフを最適化し、ライブラリレベルでコンボリューション、MatMul、エレメントワイズ、プーリング（MaxとAverage）など、調整されたプリミティブ実装を提供しています。グラフ1の下のオレンジの線で描かれた畳み込みニューラルネットワークの幅を含め、さまざまなニューラルネットワークモデルで性能の向上が確認されています。ZenDNN AI推論ライブラリによるTencentのAIアプリケーションの最適化とTinyDefectNetにおけるTF-ZenDNNの影響は、それぞれZenDNNの高い性能とTensorFlowとの統合を実証しています。

グラフ1. TF-vanilla（ZenDNNなし）と比較したTensorFlow-ZenDNNプラグインv0.1およびTF-ZenDNN直接統合v4.0の性能向上率。TensorFlow-ZenDNNプラグインには今後も最適化が追加されるため、性能向上の程度はTF-ZenDNN直接統合と比較することが予想されます。本ブログ末尾の巻末資料ZD-045～ZD-051をご参照ください。

TensorFlow-ZenDNNプラグイン

TF-ZenDNNの直接統合は、上のセクションで説明したバイナリ形式のように、TensorFlowのコードに大きな変更を加える必要があります。このような変更をTensorFlowのリポジトリにアップストリームすることは、面倒で持続不可能です。そこでTensorFlow v2.5では、 AMD EPYC™ CPU 用のTensorFlow-ZenDNNプラグインを実装する際に、デバイス固有のコードをモジュール化してプラグアンドプレイで統合できるPluggableDevice mechanismを提供しています。TensorFlow-ZenDNNプラグインは、カーネルとop登録のC APIを通じて、 AMD EPYC™ CPU 特有のカスタムカーネルの実装と操作をTensorFlowに追加します（下図2）。

図2. TFv2.12にアップストリームされたTensorFlow-ZenDNNプラグインは、 AMD EPYC™ プロセッサでの性能向上のためにAMDが開発したカスタムカーネルや演算を追加することができます。

TensorFlow-ZenDNNプラグインとTF-ZenDNN直接統合の主な違いは、標準のTensorFlowパッケージとの互換性です。TF-ZenDNN direct integrationは、標準的なTensorFlowのパッケージを置き換えるスタンドアロンパッケージです。TensorFlow-ZenDNN plug-inは、TFバージョン2.12以降、TensorFlowの標準パッケージと一緒にインストールする補助的なパッケージとなっています。

TensorFlow開発者の視点から見ると、TF-ZenDNNプラグインアプローチは、TF-ZenDNN直接統合アプローチと比較して、ZenDNN最適化を活用するプロセスを単純化します。TF-ZenDNN直接統合の場合、開発者は基盤となるTensorFlowビルドをダウンロードし、AMD ZenDNN開発者リソースページに移動して、統合用の特定のTF-ZenDNNバイナリをダウンロードする必要があります。一方、TensorFlow-ZenDNNプラグインアプローチでは、次のセクション「AMD EPYC™ プロセッサでZenDNNを使うためのステップバイステップガイド」で詳しく説明するように、ユーザーがZenDNNを活用するために必要なすべてがTensorFlowページ上に存在するようになるという違いがあります。

TensorFlow-ZenDNNプラグインは、最初の反復（v0.1）で、現在、Conv2D、MatMul、BatchMatMul、FusedBatchNorm、AvgPool、MaxPoolなど16種類の共通ZenDNN opsを提供します。カバーされていないその他のオペは、TensorFlowのネイティブカーネルにフォールバックされます。TensorFlow-ZenDNNプラグインは、ResNet、Inception、VGGのようなモデルにおいて、TF-ZenDNN直接統合パッケージと競合する性能を提供します（上のグラフ1では、青い棒がTensorFlow-ZenDNNプラグインの性能、オレンジ線がTF-ZenDNN直接統合性能を示しています）。しかし、MobileNetやEfficientNetなどの他のモデルでは、TF-ZenDNN直接統合の方がプラグインを上回っています。これは、TF-ZenDNN直接統合パッケージで現在サポートされているグラフ最適化をプラグインがまだサポートしていないためです。TensorFlow-ZenDNNプラグインがTF-ZenDNN直接統合と同等の機能を持つようになれば、性能はより近くなると予想されます。

AMD EPYC™ プロセッサでZenDNNを使うためのステップバイステップガイド

TensorFlowでZenDNNの最適化を利用するのは簡単です：
1.TensorFlow Community Supported BuildsのウェブページからZenDNN Plug-in CPU wheelファイルをダウンロードします。

2.次のコマンドを使用してZenDNNプラグインをPipインストールします。
pip install tensorflow-cpu==2.12.0
pip install tensorflow_zendnn_plugin-0.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

3.以下の環境変数を設定することで、推論フローでZenDNNの最適化を有効にします：
export TF_ENABLE_ZENDNN_OPTS=1
export TF_ENABLE_ONEDNN_OPTS=0

推論実行時のZenDNNの最適化を無効にするには、対応するZenDNNの環境変数を0に設定します：
export TF_ENABLE_ZENDNN_OPTS=0

TensorFlow-ZenDNN プラグインは ZenDNN v3.3 でサポートされています。パフォーマンスチューニングのガイドラインについては、TensorFlow-ZenDNNプラグイン・ユーザーガイドの第5章を参照してください。

最適な推論パフォーマンスを得るために、AMDはAMD ZenDNN開発者リソースページで入手できるTF-ZenDNN直接統合バイナリを使用することを推奨しています。

ZenDNNの今後の展開について

TensorFlow v2.12は、TensorFlow-ZenDNNプラグインの最初のリリースとなります。AMDは、今後のTensorFlow-ZenDNNプラグインのリリースにおいて、より多くのZenDNN ops、グラフ最適化、量子化をサポートすることにより、現在および将来の世代の AMD EPYC™ プロセッサにおけるTensorFlow-ZenDNNプラグインの性能向上を継続する意向です。このような機能強化には、第4世代 AMD EPYC™ プロセッサのAVX-512とVNNI機能を活用した最適化を可能にするため、ZenDNN v3.3からZenDNN v4.0へのプラグインバージョン移行が予定されています。

TensorFlow-ZenDNNプラグインをコミュニティのために継続的に改善することを目的として、TensorFlow開発者の皆様には、この新しいTensorFlow-ZenDNNプラグインのテストとZenDNN GitHubページでのコメントや懸念の共有を推奨します。また、テクニカルサポートリソースは、以下のメールアドレスから連絡することができます：zendnnsupport@amd.com.

私たちは、TensorFlowとのコラボレーションを継続し、より広いTensorFlow開発者コミュニティのためにZenDNNエクスペリエンスを向上させることを嬉しく思っています！

謝辞
TensorFlow-ZenDNNプラグインの開発とアップストリームは、AMDとGoogleのTensorFlowチームの多くの人たちの仕事です。
AMDから：Chandra Kumar Ramasamy、Aakar Dwivedi、Savan Anadani、Arun Ramachandran、Avinash-Chandra Pandey、Ratan Prasad、Aditya Chatterjee、Alok Ranjan Srivastava、Prakash Raghavendra、Pradeep Kumar Sinha、Vincent Dee.
グーグルより Penporn Koanantakool、Eugene Zhulenev、Douglas Yarrington。

リーガルエンディングノート
ZD-045からZD-051まで：
2023年2月7日（火）時点でAMD Performance Labsが以下のテストシステムで実施したテストです：
AMD System: AMD Eng Sample of the AMD EPYC™ 9004 96-core processor, dual socket, with hyperthreading on, 2150 MHz CPU frequency (Max 3700 MHz), 768GB RAM, 768MB L3 Cache, NPS1 mode, Ubuntu® 20.04.5 LTS version, kernel version 5.4.0-131-generic, BIOS TQZ1000F, GCC/G++ version 11.1.0, GNU ID 2.31, Python 3.8.15. For no ZenDNN, Tensorflow 2.12. For the ZenDNN plug-in, AOCL BLIS 3.0.6, Tensorflow 2.12, ZenDNN version 3.3; for Direct Integration AOCL BLIS 4.0, Tensorflow Version 2.10, ZenDNN 4.0.
Tests run all from Unified Inference Frontend 1.1 (UIF1.1) model zoo:
FP32 EfficientNet
FP32 Inception v3
FP32 MobileNet v1
FP32 VGG16
FP32 RefineDet
FP32 BERT Base
FP32 ResNet50
Results may vary based on factors such as software versions and BIOS settings. ZD-045 thru ZD-051