後藤弘茂のWeekly海外ニュース
NVIDIA Ampereにおけるプルーニング対応の特徴
2020年7月20日 06:50
Ampereでのプルーニング対応の特徴
NVIDIAは新しいGPUアーキテクチャ「Ampere」に、スパース(疎)ネットワークのハードウェア機構を搭載した。Ampereアーキテクチャでは、マトリックス計算用のテンサーコアでスパースネットワークにハードウェアで対応する。Ampereでのスパースサポートは、考え抜いたアーキテクチャとなっており、複雑な実装を避けている。モデル精度を落とさずに、容易にAmpereでアクセラレートできるテクニックを採っている。
Ampereでのスパース対応の重要なポイントは、「Fine-Grained Structured Sparsity」、つまり、細粒度(Fine-Grained)で、かつ構造化(Structured)されたスパースネットワークに対応していること。Ampereのハードウェア対応では、0値化されたパラメータが、ブロックやベクタ単位の粗粒度(Coarse-Grained)でまとまっている必要はない。細かく0値が分布した細粒度のスパースネットワークでも、ハードウェアでアクセラレートできる。その代わり、スパース度は50%と、比較的緩い2分の1の圧縮度に留める。
PDF版は@img|p27.pdf|こちら||t@@
Ampereの方式では、スパースネットワークを作るトレーニングの段階でも、トリッキーな手法は要求されない。また、ハードウェアの対応レベルとしては、メモリ、コンピュート、演算のすべてのレベルで対応する。圧縮されたスパースネットワークを、そのままメモリから読み込むことが可能で、プルーニング(刈り込み)された0値のパラメータは演算されず、空いた演算アレイのスロットでは別なパラメータを演算できる。そのため、Ampereのスパースネットワーク対応では、必要とするメモリ量やメモリ帯域を50%に減らすだけでなく、演算性能を2倍にし、消費電力あたりの性能も2倍にする。
細粒度だが構造化されたスパースネットワーク
Ampereのプルーニングは、ちょっと特殊だ。これまでの研究では、ニューラルネットワークでプルーニングできるパターンを見つけ出し、その上で、プルーンしたネットワークを効率的に実行できるハードウェアの仕組みを作るという方向だった。それに対して、NVIDIAがAmpereで採ったのは、自社のハードウェアに実装しやすいように、プルーニングのアプローチ自体を合わせるという方法だ。
ただし、ハードウェアに合わせるといっても、CNN(Convolutional Neural Network)でフィルタやチャネル単位でプルーニングして粗粒度(Coarse-Grained)のスパースネットワークを作ってハードウェアに合わせるような、強引な手法は取らない。プルーニングの粒度は、あくまでも柔軟な対応が可能な細粒度(Fine-Grained)にできるようにメカニズムを作る。その方がインファレンス(推論)精度を保ちやすく、また、多様なアルゴリズムにも対応がしやすい。結果として、NVIDIAはハードウェアの実装コストを最少にしながら、精度を保ち、かつプルーニングの効果を得ることができた。
Ampere向けのプルーニングの特徴は、細粒度(Fine-Grained)で、かつ構造化(Structured)されたスパースネットワークという点だ。粒度が大きければ、ハードウェアの対応は容易で、インデックス比率も小さくできる。しかし、NVIDIAは細粒度(Fine-Grained)を取った。さまざまなアルゴリズムへの対応を考えると細粒度対応は重要となる。しかし、通常は細粒度(Fine-Grained)のスパース性は、イレギュラーな非構造化(Unstructured)となり、対応が難しい。NVIDIAが採った方法は、細粒度のスパースネットワークを、構造化してしまうというアプローチだ。
インデックスを小さくしてオーバーヘッドを減らす
Ampereでは、ウェイト(重み)マトリックスをプルーニングしてスパースにする。アクティベーション側はデンス(密)のままだ。Ampere向けの処理では、まず、プルーニングの粒度を小ブロックで決め打ちにしてしまう。マトリックスを、4要素ずつの1x4の小さなベクタに区切る。その上で、1x4の4要素のうち2要素をプルーニングでゼロ値にする。
もっとも一般的な、プルーニングは絶対値で一定のしきい値より下の値をゼロ化する。しかし、Ampere対応の場合、1x4の中にゼロ値化できる要素が3個以上あっても、ゼロ値化で刈り込むのは2要素までとする。たとえば、3個の0値が1x4ブロックに存在したら、1個の0値はそのまま残される。逆に、しきい値以上の値が3個以上あっても、2個の値は0値化しなければならない。
この規則性のある「2:4」の仕組みでスパース化したネットワークは、圧縮も容易になる。4要素のうち2要素だけを残すかたちでそれぞれの1x4ブロックを圧縮する。残った値が、元のマトリックスのどこに位置していたかは、残った各要素につき2-bitのインデックスで示される。1x4ブロックの中での位置なので、元位置メタデータは最少で済む。たとえば、左側の要素は、4個のスロットのうち1、2、3のどれかに位置していたことになるので、2-bitで示すことができる。
各要素2-bitと極めて小さいインデックスで済むため、Ampere向けスパースネットワークでは、インデックスオーバーヘッドは小さい。ただし、圧縮率は50%と決め打ちなので、それなりのオーバーヘッドがある。16-bitパラメータの場合はインデックスオーバーヘッドは12.5%、8-bitパラメータの場合は25%だ。逆を言えば、圧縮率が小さくても、オーバーヘッドが小さくなるようにされている。
Ampereのテンサーコアアーキテクチャ
NVIDIAが提案するのは、Ampereのテンサーコアにインプットするウェイトのマトリックスをプルーニングで圧縮する方法だ。大きな特徴は、すでに説明したように、プルーニング時に規則性を持たせて構造化すること。それがどういう意味を持つかは、Ampereテンサーコアのアーキテクチャを知る必要がある。
Ampereのテンサーコアは、スパースネットワークへの対応だけでなく、大幅に強化されている。Ampereテンサーコアは、ハードウェア的には256個の積和算(MAC)ユニットで構成されている。実際には、256個の乗算(Mul)ユニットと、32個の加算(Add)ユニットだが、加算ユニットはカラムのすべての要素を加算するため、実質的に256個の積和算ユニットとなる。
FP16の場合、8x8のマトリックスと、8x4のマトリックスの積和算が1サイクルスループットで可能だ。8x8の64要素のマトリックスと、8x4の32要素のマトリックスで、256積和算/サイクルとなる。
命令発行レベルで見ると、32スレッドオペランドシェアリングの場合、1サイクルの命令発行で、16x16のマトリックスと16x8のマトリックスを掛け合わせる積和算を、テンサーコアに発行することができる。テンサーコアは、ハードウェア的には、16x16の中の8x8のマトリックスと、16x8の中の8x4のマトリックス毎に演算する。つまり、マトリックスのうち1/4ずつ積和算を行なう。16x16と16x8の積和算ではトータルで2048の積和算となり、Ampereのテンサーコアで8サイクルのスループットで処理できる。
PDF版は@img|p54.pdf|こちら||t@@
テンサーコアでのスパース対応の実装
Ampereのスパースネットワークへのハードウェア対応は、このテンサーコアのアーキテクチャに合わせたメカニズムとなっている。テンサーコアでは、インプットされる8x8のマトリックスを、4-wayのロウ毎に区切ってスパースの管理を行なう。
すでに述べたように、トレーニング時には、4値のブロックの中で4値のうち2値までをゼロ値にするプルーニングを行なう。マトリックス上の個々の1x4のリージョンで、4個のウェイトを2個に減らすことで、2分の1にウェイトを減らす。テンサーコアでは、そうして半減したウェイトを2倍詰め込むことで、性能を2倍に引き上げる。
具体的な例として、大きなマトリックスをタイル分けしてテンサーコアに乗せるケースがGTC(GPU Technology Conference)では示された。
下のスライドは、一般的に多用される、大きなマトリックスからタイル切り分けでマトリックス演算GEMM(General Matrix Multiply)を行なう場合だ。Ampereのテンサーコアの場合、通常のデンスネットワークでは、MxKのマトリックスから16x16のタイルを、KxNのマトリックスから16x8のタイルを切り出して、テンサーコアでマトリックス演算を行なう。16x16タイルに16x8タイルを乗算する。そして、MxKxNの中で、タイルを次々に処理して行く。これがAmpereのテンサーコアの標準的なオペレーションとなる。
では、プルーニングによるスパースネットワークではどうなるのか。プルーニングを前提とする場合、Ampereでは、2倍のサイズの16x32のタイルをMxKから、32x8のタイルをKxNから切り出してGEMMする。下のように、ウェイトのタイルが2倍になり、アクティベーションも2倍になる。2倍のサイズのタイルになるので、デンスネットワークの場合は、テンサーコア命令が2発行分となる。
しかし、MxKマトリックスのウェイトタイルがAmpere対応でプルーニングされていると、ウェイトのパラメータは実際には半分になっている。下のスライドのように、細粒度に刈り込まれたウェイトとなっている。スライドで緑がウェイト値が残されている部分で、白が0値化された部分だ。
ちなみに、本当のタイルサイズは16x32だが、スライドで左上に引き出されたウェイトマトリックスはわかりやすくするため8x16となっている。実際には、スパースウェイトは、左上の図の4倍のサイズがある。かなり混乱を招くスライドとなっている。
タイルサイズを半分に圧縮してテンサーコアに割り当てる
Ampere対応のスパースネットワークの場合、実際には、ウェイトタイルは、Ampere対応のフォーマットで圧縮されている。そのため、16x32タイル分のウェイトは、その半分の16x16タイルに押し込められている。圧縮のためのインデックスも付加されるので正確には半分ではないが、ウェイトの値自体は半分になる。インデックスには、それぞれの値が、各アレイのうち元々どの位置にあったかのメタデータが記録される。スライドでは、中央の緑の正方形が圧縮された16x16タイルだ。その右の紫色がインデックスとなる。
圧縮されたスパースネットワークでは、16x32タイルは圧縮され半分の16x16タイルになる。個々のタイルが半分のサイズになるため、ウェイトのMxKマトリックス全体が半分のサイズへと圧縮される。16x32タイルを圧縮した16x16タイルは、1命令でテンサーコアで実行されるため、16x16タイル相当の演算量となる。
ウェイトが圧縮されると、アクティベーションの方はどうなるのか。プルーニングによって削られたウェイトに対応するアクティベーションは、不要となる。KxNのアクティベーションマトリックスの32x8のタイルも、使用するパラメータは半分となる。
ただし、アクティベーションマトリックス側は事前に刈り込んで圧縮することはできない。ウェイト側のタイルが変わると、同じアクティベーションタイルの中で選択されるアクティベーションが変わるからだ。
そのため、テンサーコアがウェイトを読み込む時点で、動的にアクティベーションを選択する。インデックスを読み込んで、アクティベーションの32x8のタイルから、残ったウェイトに対応する必要な値だけが選択される。実際に使うのは16x8分の値となる。
ウェイトが50%に削られ、演算が半分になることでテンサーコアの演算スループットは、スパースネットワークでは2倍となる。実際のマトリックス演算実行時でも、最大で1.78倍のパフォーマンスアップになるとNVIDIAは説明する。また、CNN(Convolutional Neural Network)のコンボリューションタスクでも同様のスピードアップが達成できるという。
再トレーニングが必要なプルーニングのトレーニング
NVIDIAのAmpereでのプルーニングは、0値化できるパラメータがそれなりに多いネットワークで、2:4と決め打ちでスパース化することで、構造化したスパースネットワークを作るというアプローチだ。単純にプルーニングすると、非構造化になってしまうケースを、Ampereの枠に押し込むことで構造化されたスパースネットワークにする。決め打ちの構造化によって、もっと圧縮できるチャンスは失うが、その代わり、ハードウェアでの対応が容易になり、高速化しやすくなる。
NVIDIAのホワイトペーパーには、この仕組みの図式化が掲載されている。通常のFine-Grained(細粒度)スパースネットワークでは、ノード毎に0値化されていない結合の数が異なる規則性のない非構造化状態であるため、データフェッチとコンピュテーションが不揃いとなり、非効率性となる。しかし、プルーニング時に、ノード毎の結合を揃えることでFine-Grainedであっても、規則的な構造化にする。すると、データフェッチとコンピュテーションが揃い、ハードウェアにとって効率的になる。
Ampere対応で50%のプルーニングを行なうと、モデル精度は低下する。そのため、再トレーニングが必要になる。これは、ほとんどのプルーニングで同様だ。Ampere対応のスパースネットワークのトレーニングでもデンストレーニング後にプルーンした後、スパーストレーニングを行なうレシピが示されている。量子化を組み合わせる場合は、デンストレーニング→プルーニング→スパース再トレーニング→量子化→量子化ファインチューニングとなる。
