TPU/GPUの電力効率について考える。
TPUの電力効率はGPU比80倍ではない
何故かTPUはGPUの30～80倍の電力効率があると書かれるのを読むことがある。
これは例えば以下の記事にこうある。
・The performance/Watt of the TPU is 30X – 80X that of contemporary products; the revised TPU with K80 memory would be 70X – 200X better.
(翻訳)TPUの性能/ワットは従来製品の30倍～80倍、K80メモリ搭載の改訂版TPUは70倍～200倍となる。
ここでTPUの従来製品との性能比較はCPUとGPUの結果を含む。
また、論文の図は8bit演算機(性能は92TFLOP)しか乗せてないように見える。これはTPUv1を指している。改良版(v2ではない？)はネックになっているメモリ帯域を増加させて電力効率を2倍にしたものである。理論性能は増加しない。
グラフ図を見れば30～80倍というのはIncremental(増分)のWMのTPU/CPU比が83倍でIncrementalのWMのTPU/GPU比が29倍である。従って80倍という部分はCPU比の値であってGPU比に限れば29倍である。この時点ではメモリ帯域を増やせば68倍と2倍に性能を増やせる。
TPUの性能向上は並列化の方が大きい
よく知られているTPU性能でpodという塊の性能はTPUv2 podが11.5P FLOP、TPUv3 podが100P FLOP、TPUv4 podが1000P(1exa) FLOPと約10倍ずつ大きくなっている。
この数年でこの演算量の上昇はめざましいものがある。ところが1core当たりの性能は各vごとに2倍程度しか増えておらず、残りの4倍は並列化の数を増やした効果だと思われる。
TPUv2 podを8core×64台(512core)と見なせばTPUv3 pod(2048core)、TPUv4 pod(8192core?(4096チップ))と考えられる。(コアとチップとデバイスの関係は8:4:1である。1デバイスに4個チップが載っており、1チップあたり2コア、1デバイス当たり8コアである。1個のTPUデバイスには4つのチップが搭載されており、それぞれのチップに2つのTPUコアが含まれています。)
GPUのTensorコアの影響
上記の比較対象になっているGPUはK80である。ところがこれは相当古いGPU(2014)である。
GPUの世代はK80=>M60=>P100=>V100=>(T4)=>A100=>H100と5世代くらい進んでいる。
K80が300WでFP32で8.74T FLOPだったが、H100は700WでFP16で2P FLOPである。
TensorコアのおかげでGPUの理論演算性能はK80と比べH100では230倍に上がっている。TDPの電力で割っても98倍である。
RTX3090Tiも450WでFP16で320T FLOPであり、RTX1080Tiで250WでFP32で11.3T FLOPや、RTX2080Tiで260Wで114T FLOPと比べ、Tensorコアの影響で理論性能が上がっている。
正直、Tensorコアの影響でどのくらい学習が高速化しているのかよく知らないが、理論値だけなら先ほどのTPUとK80の電力効率である30倍くらいは最新のGPUでも上がっている。
ただ、実際には1080Ti-3090の性能間に理論性能差の30倍も差はなく、ResNet50の学習で4倍程度らしい。
Volta100 vs TPUv3
V100とTPUv3の比較がある。
TPUv3のチップ当たりの理論演算速度はV100とほぼ等しく、また消費電力(TDP)もほぼ同じである。
パフォーマンスもTPUv3とV100共にTPUv2比で2倍前後である。
ResNet50ではTPUv3有利で、TransformerはV100が僅かに有利である。
また、ほかにも以下のことが推定される。
・TPUv2はv1からメモリ帯域が大幅に引き上げられている(34=>700GB/s)。
・チップ当たりのメモリ容量はTPUv2が16GB、TPUv3が32GB。1GPUのメモリ容量は2080Tiが11GB、3090Tiが24GB、V100が32GB、A100が40GBか80GBである。
・TPUv1でcoreあたり8bit演算機が65536個あるがこれは256行×256列という風に並んでいる。
・理論性能はこれにクロック周波数を掛けたものである。700MHz*65536*2=91.75T FLOP。
・pod当たりの消費電力がTPUv2 podが124kW、TPUv3 podが594kWである。
・TPUv4 podはv3の4倍の約2.4MWであると推定され膨大な電力を消費する装置である。300WのGPU換算で8000枚である。こう考えると言うほど省電力な装置には思えない。
A100 vs TPUv4
A100とTPUv4の学習速度比較ではResNet50ではTPUv4がかなり有利で、BERT(Transformer)は差が僅かである。V100とTPUv3の比較から若干TPU有利になっている。
どのようなシステムで比較しているのか明確に書かれてないが、V100とTPUv3の比較から1GPUと1チップ(2core)当たりの比較かと考えられる。ただし、TPUv4では2つ（またはそれ以上）のTPUコアを単一の論理デバイスとして操作するとあるので、コア当たりではなく、チップ当たりでしか操作できないのかもしれない。
ここでV100 vs TPUv3、A100 vs TPUv4で比較すると同等かTPU有利と考えられるが、逆にV100 vs TPUv2、A100 vs TPUv3、H100 vs TPUv4の比較ではGPUの方が有利かもしれない。
| TPUv2 1チップ(2core) | TPUv3 1チップ(2core) | TPUv4 1チップ | |
|---|---|---|---|
| V100 1枚 | GPU有利 | 同等 | TPU有利 |
| A100 1枚 | GPU有利 | GPU有利 | 同等(かTPU有利) |
| H100 1枚 | GPU有利 | GPU有利 | GPU有利 |
H100での新要素(Transformer Engine)
Transformer EngineはH100に搭載された機能で、FP16で計算させる際に一部をFP8の精度で計算することで速度を2倍に向上させる。以前は推論のみFP8で、学習はFP16までしか可能ではなかったが、これは学習も一部FP8で出来る。
| 名称 | 指数部 | 仮数部 | 合計 |
|---|---|---|---|
| FP8 | 4or5bit | 3or2bit | 8bit |
| FP16 | 5bit | 10bit | 16bit |
| FP32 | 8bit | 23bit | 32bit |
| BF16 | 8bit | 7bit | 16bit |
| TF32 | 8bit | 10bit | 19bit |
TPUv4は本当に2倍なのか？
TPUv4のチップ当たり性能は下記に275 TFLOPSとある。TPUv3のチップ当たり性能は123 TFLOPSだからこれで見ると倍になってるように思える。
一方、TPUv4と全くの別物なのか、類似なのか分からないがTPUv4iという論文がある。
そのTPUv4iの性能はチップ当たり138 TFLOPSであり、前述の275 TFLOPSの丁度半分である。TPUv4iではチップ当たりの性能は微増しているだけで2倍にはなっていない。
16nm=>7nmになってTDPはかなり下がっているように見えるのでTPUv3からMXUの数を倍にしても大丈夫そうだが、v4iから再度設計し直して倍にしたのかは疑問である。というのも旧2チップを新1チップとみなせるよう置き換えただけなのではと考えられるからである。もしそうならチップ当たりの性能が向上したのではなく、便宜上のチップの数え方を変えただけである。
また、H100のようにFP8使って理論性能2倍とした可能性もある。
TPUv4は情報が少なく、実際どうなのかは定かではない。
Google ColaboratoryのTPU
INFO:tensorflow:*** Num TPU Cores: 8
Google Colaboratoryの使用できるTPUのコア数は8らしい。
とするならTPUv2 pod(512core)が11.5P FLOPだから8coreなら180T FLOP。
TPUv3 pod(2048core)が100P FLOPだから8coreなら390T FLOP。
ここでV100はFP16で126T FLOP、A100はFP16で312T FLOPS、RTX2080TiはFP16で114T FLOPで、RTX3090TiはFP16で320T FLOPである。V100の1枚とTPUv3の1チップ(2core)の性能が近いから、TPUv3の4チップ(8core)はV100の4枚(または2080の4枚)に相当する。
世代ごとに性能は約2倍違うとすれば、A100(または3090Ti)はGoogle ColaboratoryのTPUv2の8coreと性能的には同等である。また、H100(または4090)はTPUv3の8coreと性能的には同等である。
| TPUv2 8core(180TFLOP) | TPUv3 8core(390TFLOP) | TPUv4 8core | |
|---|---|---|---|
| V100 2枚(252TFLOP) | 同等 | TPU有利 | TPU有利 |
| V100 4枚(504TFLOP) | GPU有利 | 同等 | TPU有利 |
| V100 8枚(1008TFLOP) | GPU有利 | GPU有利 | 同等 |
| A100 1枚(312TFLOP) | 同等 | TPU有利 | TPU有利 |
| A100 2枚(624TFLOP) | GPU有利 | 同等 | TPU有利 |
| A100 4枚(1248TFLOP) | GPU有利 | GPU有利 | 同等 |
| H100 1枚(1000TFLOP) | GPU有利 | 同等 | TPU有利 |
| H100 2枚(2000TFLOP) | GPU有利 | GPU有利 | 同等 |
ただし、FP16を使用してもRTX2080Ti(V100相当)の2GPUがTPUv2に及んでいないという情報もある。ResNetはTPU有利という情報を勘案すればここでのGPUとTPU性能を同等とみなせるのではと思われる。
ColabでTPUv2とv3を見分けるにはメモリを溢れさせた時、v2なら8coreあたり64GB、v3なら8coreあたり128GBの容量で判別つくかと思う。
また、タダでGPU4枚相当の性能の出るTPUを使用できるのは素晴らしいが、Google Colaboratoryで使用できるのはTPUのpod全部ではなくその中の8core分だけである。従ってTPUv2とTPUv3とTPUv4のpod当たりの性能差は10倍あるが、core当たりの性能差は2倍くらいしかない。
Cloud TPUのコスパ
TPUv2の8coreで4.5USD/hである。
仮にこれを一か月使用したとすれば30*24*4.5*128=41万5000円である。
RTX3090Tiの値段が32万で電気代が月1万としても、一か月以上の使用でGPUの方が有利になる。
使用後、GPUを半額で売れるとすれば2週間使用すればGPUの方が有利になる。
とはいえCloud TPUのメリットは使用状況に合わせてもっとたくさんのcoreを借りて学習を短期間で終わらせる事ができることである。8coreを長期間借りても2048coreを短期間借りてもトータルの学習に必要な値段は変わらないとするなら、多く借りて早く終わらせるのが(時間的な)得である。
AWSのGPUインスタンス vs Cloud TPU
p4dというA100のGPU8個のインスタンスが32.77USD/hである。
A100の性能は312T FLOPS(スパース性を使うと624T FLOP)だからx8枚で2.5P FLOP。
一方、Cloud TPUでTPUv3の32coreが1.56P FLOPで近い。これはV100×16枚と同等でこの値段は32USD/hであるから、AWSのGPUインスタンスとCloud TPUの価格差はあまりない。
またAWSではUltraClustersというのでp4dインスタンスを更に並列化できるようである。GPU4000枚分出来るとあるので、その演算量は1248P FLOPでTPUv4の演算量に匹敵する。
電力効率の比較
TDPと理論演算量のプロットを書いてみる。
このグラフの書き方はリンク先のFig.2を参考にした。データが左上にあるほど電力効率が良い。TPUv1のみINT8性能でそれ以外はFP16性能である。
TPUv1の電力あたりの理論演算量は確かに非常に優れているが、実はTPUv2やv3はそこまででもない。
また、電力効率はTPUv1がK80比で29倍であったが、TPUv3とV100比2倍も差がなかった。
TPUv3(16nm)=>v4(7nm)で2倍の性能向上があったとしても、GPU側もV100(12nm)=>A100(7nm)=>H100(5nm)と進化している。(RTX3090はSamsung 8nm)
従ってH100がTPUv4に対して電力効率で負ける理由はない。
今後(2022/4月以降)：TPUv4 vs RTX4090 vs H100
TPUはCloud TPUで今後TPUv4の使用解禁、GPUは次のRTX4090、H100が発売が上げられる。
TPUv4はコア当たりTPUv3の2倍の性能である。RTX4090は発売が22年の9月らしいが消費電力が多いという噂を聞いた。また、22年第3四半期（7～9月）発売のH100の消費電力700Wで、A100の消費電力400Wに比べ2倍近い。
今後しばらくはTPU、GPU共に性能向上を見られそうだが、消費電力の増加が懸念される。
参考