Hot Chips 31はマシンラーニングが花盛り(6) Intelの機械学習用トレーニングチップ「Spring Crest」の性能を読み解く

Spring Crestは16bitのBFLOAT16で乗算を行い、積項の和は32bitのFP32で計算している。次の表は、いろいろな数値表現の得失を纏めたものである。チップ面積や消費電力の点ではbit数が少ないほうが効率が良いが、精度が低いので収束性が悪い。その中で、BFLOAT16で乗算しFP32で加算を行うSpring Crestの方式は、最高の精度を維持しつつ電力効率が高く、学習時間を短くできる選択であるという。

FP32は精度が高く、表現できる数の範囲も広いがチップ面積や電力効率は良くない。一方、FP16は面積や電力効率は良いが計算の収束性では中程度の評価である。BFLOAT16は面積や電力効率が良く、収束性も良い。Spring Crestでは、積の計算にはBFLOAT16、その和の計算にはFP32を使用し、最高の精度で、効率の高い計算ができる

行列積の計算はBFLOAT16で行われるが、行列以外の計算のためにベクトル演算器をもっている。行列積以外の計算では、BFLOAT16とFP32の選択ができるようになっている。

また、ベクトル演算器は、アクティベーション関数の計算や乱数発生、リダクションや累積加算などの計算を行うことができる

Spring CrestではBFLOAT16で積を計算する。行列積以外のユニットはBF16かFP32かの選択ができる。ベクタ演算器は、アクティベーション関数や乱数の発生などの機能も持っている

Spring Crestは4個のHBM2-2400スタックをメモリとして搭載している。総メモリ容量はは32GBである。また、TPCごとに2.5MBのスクラッチパッドメモリを持ち、チップ全体では60MBのメモリを持っている。

このメモリはテンソルの転置を行う機能を持っている。また、それぞれのメモリブロックはReadとWriteを同時に行うことができる。1つのTPCのメモリのR/Wバンド幅は1.4Tbpsである。そして、HBMとMRB(説明が無いが、スクラッチパッドのバッファメモリか)はメモリ間転送ができるようになっている。

Spring Crestのメモリシステムは、4個のHBM2-2400が中心となっている。1.22Tbpsのバンド幅で、8GBの容量を持つ。さらにTPCごとに2.5MBのスクラッチパッドメモリを持つ

Spring Crestのチップのフロアプランに見られるように、TPCは4行5列のブロックの両側に4個のTPCが張り出した配置になっているので、次の図に示すメッシュネットワークの接続も、多少、不規則になっている。

クロスセクションのバンド幅は、両方向の合計で2.6Tbpsで、HBMやSerDesへのアクセスもこのメッシュネットワークを共用している。ただし、別種のタイプの通信には別のメッシュが使われている。

TPCは2次元メッシュのネットワークで接続され、4個のHBMやインターチップのクロスバもメッシュに接続されている。ネットワークのクロスセクションバンド幅は2.6Tbpsである

Spring Crestでは、1つのチップに収まらない大きなモデルも複数チップ、あるいは複数シャシーを繋いで走らせることができる。チップ間の通信には、最大112Gbpsのリンクを16本使用できるので、両方向の合計では3.58Tbpsの通信が行なえる。ルーターは完全プログラマブルなので、複数のトポロジをグルーレスで作ることができる。

このネットワークで低レーテンシのローカルメモリ転送が行なえ、ネットワークは最大1024ノードまで拡張ができる。

Spring Crestでは複数のチップ、あるいは複数のシャシーを使って大きなモデルを実行することができる。チップ間は、16チャネルの112Gbpsのリンクで繋ぐことができる。このやり方で、最大1024ノードのネットワークを作ることができる

ディープラーニングのワークロードでは各種のサイズの行列積が出てくる。Spring Crestでは1024×700×512から5124×9124×2048と大きくサイズが異なる行列積でも演算器の利用率は31.1%から57.4%の範囲に収まっている。この結果は、公表された競合製品の結果と比べると～2倍良い結果となっている。結果として、Spring Crestは、遊んでいるリソースを減らし、結果としてトレーニング時間を短縮することができる。

Spring CrestチップのGEMM(行列積)実行性能。1024×700×512から5124×9124×2048と大きくサイズが異なる行列積でも31.1%～57.1%と高い演算器使用率を保っている

次の棒グラフに見られるように、テンソルの配置によって計算効率は59%～87%まで変化する。Spring Crestでは複数のテンソルレイアウトをサポートしており、適当なものを選ぶことにより、計算効率を高めることができる。

Spring Crestの畳み込み性能。7つの異なるサイズの計算が書かれているが、演算器の利用率は59%～87%と高い利用率となっている

全体をリング接続した場合の測定であるが、2枚のカード間の送受の場合のバンド幅は161GB/s、Allreduceのバンド幅は151GB/s、Broadcast のバンド幅は147GB/sが得られている。

Spring Crestの通信性能。リング接続の場合の測定で、2枚のカード間の送受では161GB/s、Allreduceでは151GB/s、Broadcastでは147GB/sの通信速度が得られている

直接のメモリ転送ができオーバヘッドが小さいので、低レーテンシとなっており、2枚のカードでのAllreduceのレーテンシは3μs、4カードでは8μsで、32カードでは36μsとなった。クロスシャシーではケーブルの遅延の増加がレーテンシ増加の原因と思われる。

また、メッセージサイズを1MBから128MBまで変えたAllreduceでは、1MBで8チップの場合は68.7GB/sに対して1MB、32チップでは39.9GB/sとなった。これはチップ数が多くなると通信回数が増えることによると思われる。一方、データサイズが128MBの場合は8チップでは147.1GB/s、32チップの場合は147.4GB/sとなり、データが大きくなるとチップ数の違いはほとんど見えないという結果になっている。

左の表は2kBと小さなデータのAllreduceのレーテンシの測定結果。8カードまでは9μsと高速である。クロスシャシーとなる16カード、32カードではそれぞれ30μs、36μsとレーテンシが伸びているのはシャシー間を接続するケーブルの遅延が効いていると思われる。右の表はデータサイズを変えてAllreduceのデータ伝送速度を測定したもので、1MBのデータで8チップの場合は68.7GB/sに対して32チップでは39.9GB/sに低下している。これはチップ数が増えて通信回数が増えたことによると思われる。128GBの場合は、データ量が大きいので、チップ数の違いは目立たなくなっていると思われる

結論であるが、DLのトレーニングではドメインスペシフィックなアクセラレーションは効果的であった。それは、数値表現や計算方法もDL向きのものを使った。そして、このプラットフォームでは、レガシーなワークロードは実行しないのでDL処理だけに最適化ができたからである。

そして、アーキテクトは、データ移動を減らし、コンピュートユニットを遊ばせないアーキテクチャを考え、また、マイクロベンチマークで高い利用率や効率が得られるようにすることで、速くて、より電力効率の高い学習ができるようになったからである。