arutema47's blog

書いたり書かなかったり。

ハードウェアの速度をどう評価するか考える(2) ~メモリ、メモリ律速~

前回のあらすじとこの記事の目的

前編: ハードウェアの速度をどう評価するか考える(1) ~クロック、OPS~

現代ハードウェアの計算性能を評価する尺度であるメモリ律速の概念とルーフラインモデルについて理解を深めることです。 本記事を通し、あるアルゴリズムが速度が十分に出ない時、それがハードウェアのどの性能(メモリか演算)に律速されてるかイメージできるようになるのが目標です。

しかし前編は子供が起きたので前編は演算速度だけで終わってしまった! 今回は起きる前にメモリ律速まで書くぞ!

https://venturebeat.com/wp-content/uploads/2020/05/Jensen-cooking.jpg?resize=1024%2C614&strip=all

  • アチアチのGPUお待ち!

メモリ律速

脱線したが本線のメモリ律速に戻る。前回では演算速度の話をし、演算速度(OPS)はハードウェアの持つ計算能力で決まると書いた。それでは最強・最速のCPU/GPUは無数に演算機を詰め込めば実現できるのだろうか?もちろん違う。

f:id:aru47:20201101173608p:plain
AMD Zen2 chip photo from ISSCC2020

チップの殆どの面積が内蔵メモリ(L3キャッシュ)で占められることがわかる。またCPUコア自体もSRAMエリアが四割ほどを占める(畑みたいなエリアはSRAMセルである)

f:id:aru47:20201101173512p:plain
Apple A12
f:id:aru47:20201101173539p:plain
Apple A12 CPU

Appleの最新A12チップの中身も見てみよう。

CPUエリアの半分以上もSRAM(L2とシステムキャッシュ(L3?))で占められていることがわかる。(GPUはいい画像が見つからなかったけどこれもまたチップの殆どがキャッシュです。。

先端チップの大きな面積はメモリ(L3キャッシュ)であった。要はメモリである。なんでこんなにメモリを敷き詰めているかというとハードウェア性能にメモリはとても影響するからである。

メモリが重要なワケ

当たり前だが、計算機が

out = a + b

という計算をするためにはabの2つのデータが必要であり、outをメモリに格納する必要がある。そのため1演算(OPS)には2データの読み込みと1データの書き込みが必要だ。

ここでもしハードウェアが100GOPSの計算速度と200Gdata/秒の読み込み速度を持つが、1Mdata/秒しか書き込み速度がないとする中で上記演算を繰り返すとしよう。すると1Mデータを出力したところで書き込みリミットに差し掛かってしまい、以後の計算は停止してしまう。CPUのパイプラインストールと同様のことが起きてしまい、例えいくら演算速度が早くても無意味である。 また上記のout = a + bの例でも同様に読み込みデータが極端に遅ければ演算にデータ読み込みが間に合わず、演算はストールしてしまう。

このようにメモリ読み込み/書き込みが原因で演算速度が律速してしまうことをメモリ律速と呼ぶ。

ここまで読むとわかるが、データの読み込み速度というのは演算速度と同等にハードウェア性能を評価するためには重要であることがわかる。 演算性能だけ高めたとしても、実際の計算は早くならないのだ。

GPUから読み解くメモリバンド幅

メモリ読み書きの速度はメモリ帯域やメモリバンド幅とも呼ばれ、ハイエンドGPUでは非常に重視されている。

f:id:aru47:20201101174151p:plain

NVIDIA A100 | NVIDIA

先端nvidia GPU A100は40GBもの専用メモリを備え、それらはHBM2(high bandwidth memory2)という接続IOでGPUーメモリ間を結ぶ。またそのメモリ帯域幅は1.5TB/s!であり、そのデータ量までならばデータを読み書きできることができることを示している。

先程のout = a + bの例で各データが1BのINTであったとすると、1演算で3Bのデータ伝送が必要になる。そのため1秒間で行える演算量は概ね1.5TB/s / 3 = 500GOPSであると試算できる。これはGPUメモリ内にすべての入力が格納されている例を考えた。例えばDNNを走らせる際はDNNの重み、gradient、中間値はすべてGPUメモリ内に格納するようになっており、概ねこのような概算が可能である。ただDNNの入力はCPUからPCIeバスを通じて伝わるため(64GB/s)、そこは律速しないように注意する必要はある。

また実際のGPUはキャッシュと呼ばれる大量の内部メモリ(メモリ帯域は10TBくらいと一桁高速)を備えるため、a,bが固定値であれば一度GPU内に読み込めば自動で再利用してくれる。またCNNのweightのように入力をスライドさせ繰り返し適応するものでは重みをキャッシュ内に格納し外部メモリアクセスを最小化するように賢く再利用している。そのようなトリックがcuDNN等のライブラリにかかれており(暗号化されているが)、頑張って読んでいくと3x3convや5x5convに特化したカーネルなどを見つけることができる。キャッシュの仕組みなどはここでは解説しないが、興味あるならばパターソン&ヘネシーなど読むと良いです。

amzn.to

amzn.to

amzn.to

ルーフラインモデル

f:id:aru47:20201005102653p:plain
TPUのルーフラインモデル

とうとう本題のルーフラインモデル*1まで来ました!今までの話を追えていれば難しい話ではありません。

out = a + b から演算量を求める例え話の時におや、と思った方もいるかと思います。 興味深いことにハードウェアの演算量、メモリ帯域が与えられ、演算したいアルゴリズムのデータ入出力量が既知であればそのハードウェアの演算速度(OPS)が求められるのです!

演算したいアルゴリズムのデータ入出力量をルーフラインモデルでは

ある演算に必要なデータ量=MAC OPS/data[byte]

で表現し、x軸にプロットします。またその時のアルゴリズムでの実測OPSをy軸にプロットします。

注意したいのはCNNのようにデータを計算間で再利用できる文は(キャッシュ利用できる分は)差し引いて考える点です。Alpha GOのCNNはMAC OPS/data=1000とデータ再利用効率が非常に高く、グラフの右上にプロットされます。ここの点ではメモリ帯域で律速されず、ハードウェアのOPSでリミットされます。

一方でRNNやMLPはデータ再利用効率は優れず、MAC OPS/data=100程度です。TPUv1はメモリ帯域が遅く、演算はメモリ律速になってしまいます。このようにルーフラインモデルは複数の演算、メモリヘビーなアルゴリズムを比べることでハードウェアの演算速度とメモリ帯域両方を加味したバランスの良いベンチマークが可能となります。

ちなみにDNNでバッチ数を増やすと速度がちょっとずつ上がる、TPUではバッチ数を出来るだけ上げたほうがよい、と言われるのはバッチ間でweightデータを再利用できるためMAC OPS/dataが上がり、メモリ律速が緩和され高速化できるためです。

Further comments

TPU vs GPU

TPUやGPUの最大速度はCNNといえどもメモリ律速であり、大体メモリバンド幅で決まります。 そしてTPUとGPUは同じIO(HBM)を使用しているので最大速度もとんとんになってしまいます。

https://www.extremetech.com/wp-content/uploads/2018/04/ResNet-50-640x396.png TPUv2の結果ですが、だいたい速度は同じです。

FP16, BF16, TF32などの低精度Mixed Precision学習

qiita.com

https://qiita-user-contents.imgix.net/https%3A%2F%2Fqiita-image-store.s3.amazonaws.com%2F0%2F171915%2F891e685c-48f6-4afb-9839-1074706d8f73.png?ixlib=rb-1.2.2&auto=format&gif-q=60&q=75&s=f8c18922b1c0b7d21b3f11d95d5e2bba

メモリ律速はデータの読み書きする量で決まるのでFP16といったより少ないデータ量で学習を行うと、理想的には二倍ほど同じGPUでも速度向上するポテンシャルがあります。なのでハードウェア性能を引き出す上でMixed Precision学習は非常に強力なテクニックです。例えばTPUv3ではBF16というフォーマットで学習が行われ、指数表現がFP32フォマットと同等の精度であるため学習精度も落ちないことを売りにしています。

また最近nvidiaが導入した低精度表現TF32(実効的にFP19)の学習がpytorchでデフォルトになりました。 これで更に高速化が達成できるかと思います。

pytorch.org

もしpytorchを使っているならば公式のMixed Precision学習導入説明がとてもわかりやすいので読んでみて下さい。

*1:Roofline: An Insightful Visual Performance Model for Floating-Point Programs and Multicore Architectures https://people.eecs.berkeley.edu/~kubitron/cs252/handouts/papers/RooflineVyNoYellow.pdf