- Google TPU v1は、ディープラーニングサービスの推論コストとスケールに対応するため、GPU比で10倍のコスト性能を目標に、2013年末から15か月で開発されたASIC
- 中核は、1978年にH.T. KungとCharles E. Leisersonが提案したシストリックアレイを256×256 MAC構成で実装し、行列積の中間結果がメモリを往復する回数を減らした点
- TPU v1はPCIeでホストと通信し、DDR3-2133に重みを保存し、
Read_Host_Memory、Read_Weights、Matrix_Multiply / Convolve、Activate、Write_Host_Memoryのような少数の命令で推論フローを構成
- 8-bit×8-bit整数乗算と量子化により、浮動小数点演算が必要とするダイ面積の負担を回避し、TensorFlowと連携するUser Space Driver・Kernel Driverが実行を制御
- Intel Haswell CPUとNvidia K80 GPUが当時の比較対象で、TPU v1は推論で約15〜30倍高速、GPU比の相対的な増分性能/Wattは25〜29倍高かったが、学習用デバイスではなかった
推論コストを下げるためのTPU v1の出発点
- TPU v1プロジェクトの目標は、ディープラーニングベースのサービスが必要とする大規模なハードウェアコストを下げるため、推論用ASICを迅速に作ることだった
- 目標値は、GPU比で推論における10倍のコスト性能優位だった
- 迅速な開発、高い性能、大規模適用、新しいワークロードへの即時対応、コスト効率が同時に求められた
- TPUという名前は、テンソル演算を高速化する装置であることに由来する
- 実際にTPU v1ハードウェアが扱う中核計算は、主にベクトルと行列演算
- ニューラルネットワークの隠れ層と出力層は、入力値ベクトルと重み行列の積に活性化関数を適用した結果として表現できる
- 複数の入力データが入ると、行列積の各項目に活性化関数を適用する形になる
シストリックアレイで行列積を処理
- TPU v1は、H.T. KungとCharles E. Leisersonの1978年の論文 Systolic Arrays (for VLSI) にあるシストリックシステムの概念を利用している
- 複数のプロセッサが規則的なリズムでデータを計算し、受け渡す構造
- 各プロセッサは短い計算を行いながら、データを継続的に内外へ移動させる
- 単純な2×2行列積では、入力値を上側と左側から正しい順序で入れると、結果が2×2のMACアレイから自然に得られる
- 各MACは乗算と累算を実行する
- 部分和はアレイ内に保存され、最終結果は移動する対角線の形で現れる
- 2×2の例では4段階必要だが、実際には左上のMACが空き次第、次の行列積を開始できるため、2サイクルごとに新しい行列積が可能
- 重要なのは、データを正しい順序でシストリックアレイに供給すれば、値と結果の流れそのものが必要な計算順序を作るという点
- 中間結果をメインメモリに保存して再度取り出す必要がない
- 行列積ユニットの構造と入力順序のおかげで、中間結果は必要な時点で自動的に利用できる
TPU v1のシステム構成
- TPU v1はホストコンピュータとPCIe高速シリアルバスで通信し、自身のDDR3 DRAMに直接アクセスする
- 主な構成要素は次のとおり
- DDR3 DRAM / Weight FIFO
- 重みはDDR3-2133インターフェースで接続されたDDR3 RAMチップに保存される
- ホストメモリからPCIe経由で事前にロードされた後、Matrix Multiply Unitで使うためにWeight FIFOへ移動する
- Matrix Multiply Unit
- 256×256 MACで構成されたシストリックアレイ
- 上側から256個の重み、左側から256個のデータ入力を受け取る
- Accumulators
- シストリック行列ユニット下部から出た結果を保存する
- Activation
- Unified Buffer / Systolic Data Setup
- 活性化関数適用後の結果を保存し、次の層の計算のためにMatrix Multiply Unitの入力として再投入する準備を行う
計算形式と命令セット
- TPU v1のMatrix Multiply Unitは8-bit×8-bit整数乗算を実行する
- 量子化を使うことで、より大きなダイ面積を必要とする浮動小数点演算を避けた
- 命令セットは約20命令を持つCISC設計
- 命令はメモリから取得するのではなく、ホストコンピュータがPCIe経由で送信する
- 推論フローの大部分は5つの主要命令で構成される
Read_Host_Memory
- 入力値をホストメモリからPCIe経由でUnified Bufferへ読み込む
Read_Weights
- 重みメモリからWeight FIFOへ重みを読み込む
Matrix_Multiply / Convolve
- Unified Bufferの入力をAccumulatorsへ送りながら、行列積または畳み込みを実行する
- B×256入力と256×256の定数重み入力を掛け合わせてB×256出力を作り、B個のパイプラインサイクルを要する
Activate
- Accumulatorsの入力にReLU、Sigmoidなど人工ニューロンの非線形関数を適用し、結果をUnified Bufferへ出力する
Write_Host_Memory
- Unified Bufferの結果をPCIe経由でホストメモリへ書き込む
- このフローはおおむね次の形になる
Read_Host_Memory
Read_Weights
Loop_Start
Matrix_Multiply
Activate
Loop_End
Write_Host_Memory
- 行列ユニットはUnified Bufferの読み書きを減らしてエネルギーを節約するため、シストリック実行を用いる
- データは左側から入り、重みは上側からロードされる
- 256要素のMAC演算は、対角線状の波面のように行列を通過する
TensorFlowとドライバスタック
- TPU v1ハードウェアを実際のサービスで使うには、これを支えるソフトウェアスタックが必要だった
- GoogleはTensorFlowを開発・利用していたため、TensorFlowがTPU v1と動作するようドライバを作ることが中核的な段階だった
- TPUソフトウェアスタックはCPUおよびGPU向けスタックと互換である必要があった
- アプリケーションをTPUへ迅速に移植できなければならなかった
- TPU上で実行されるアプリケーション部分は通常TensorFlowで書かれ、GPUまたはTPUで実行可能なAPIへコンパイルされる
- GPUと同様に、TPUスタックもUser Space DriverとKernel Driverに分かれる
- Kernel Driverは軽量に保たれ、メモリ管理と割り込みのみを処理し、長期安定性を目指す
- User Space Driverは頻繁に変化し、TPU実行設定と制御、TPUの順序に合わせたデータ再整形、API呼び出しのTPU命令への変換、アプリケーションバイナリ生成を担う
28nmプロセスとダイ配置
- TPU v1はTSMCの比較的成熟した28nmプロセスで製造された
- 当時Googleデータセンターで使われていたIntel Haswell CPUとNvidia K80 GPUチップは、より先進的なプロセスで作られていた
- Googleによれば、TPU v1のダイ面積はこれらチップのダイ面積の半分未満
- 単純なISAはデコードと関連処理に必要なダイオーバーヘッドを低減した
control領域はダイ面積の**2%**しか占めない- Matrix Multiply Unitは24%、Unified Bufferは**29%**を占める
性能比較と明確な限界
- TPU v1は、すでに学習済みのモデルをGoogle規模の実サービスでより効率的に使うための推論用デバイス
- 学習速度や効率を改善するために設計された装置ではない
- 推論と学習は、特化型ハードウェア開発において異なる課題を生む
- 2013年時点の主要な比較対象はIntel Haswell CPUとNvidia K80 GPUだった
- TPU v1はK80 GPUよりMAC数が25倍多い
- TPU v1はK80 GPUよりオンチップメモリが3.5倍多い
- TPU v1はK80 GPUおよびHaswell CPUより推論で約15〜30倍高速
- GPU比の相対的な増分性能/Wattは25〜29倍
- カスタムアーキテクチャにより、TPU v1は当時のCPUやGPUより高い推論性能と低いエネルギー消費を達成した
- 高速推論と電力効率という単一目標に集中した第1世代設計だったため、学習用には設計されていないという限界が残る
1件のコメント
Hacker News の意見
Groq CEO の Jonathan Ross が最近のポッドキャストインタビューで、Google で初期の TPU を作っていた頃の話をしていた。もともとは推論速度の問題に悩んでいたチームの近くに座っていて、20% ルールの時間で作った FPGA だったという。
動くものを作ったあと、Jeff Dean が計算して ASIC に進むことにしたとのこと。
今なら Google は TPU チームを別会社としてスピンアウトすべきだと思う。Nvidia に対抗できる唯一の信頼できる競争相手で、ソフトウェアサポートも Nvidia に次ぐ水準だ。
https://open.spotify.com/episode/0V9kRgNS7Ds6zh3GjdXUAQ?si=q...
ただし、それぞれに弱点がある。ファウンドリ供給枠は逼迫しているが、Nvidia はより高価な AI チップを売れるなら、コンシューマー向け GPU 市場を犠牲にできる。競合が数年前から大きく賭けていたり、Intel のように大きな生産能力を持つ会社が優先順位を変えたりすれば、この強みは消える。
独自ソフトウェアが業界標準なら便利だが、実際の重要度はユースケースに大きく左右される。TPU 向けのハードウェア設計は GPU より本質的にはるかに単純に見えるし、レイトレーシング、テクスチャサンプラー、ラスタライズは不要で、ほとんどは行列乗算と大量のメモリがあればよい。
顧客関係は会話に入り続けるうえでは役に立つが、少しでも優位性を探す市場では、ドルあたり FLOPS が最も高いハードウェア供給元が、生産能力を埋めるだけの顧客は確保するはずだ。だから数年後には、競争はかなり早く現実のものになりそうだ。
ただし TPU は Google 内部でも比較的希少なコンピューティング資源で、社内需要を満たすのも難しい可能性が高い。
会社をここまで尊敬していると言うことはほとんどないが、Groq は本当に尊敬している。
Google が TPU を発明し、Google Research が LLM 論文まで出したのに、なぜ NVDA と AI スタートアップが価値のほぼ 100% を持っていったのか分からない。
Google の場合なら、ChatGPT のようなものを広く提供しようと提案すると、検索エンジンの有料掲載や、人々が訪問しなくてもよくなるサイトの広告収益を侵食しかねない、という比喩が成り立つ。だから競争上必要なときだけ、衝撃が少ない形で慎重に導入しようという判断になったのかもしれない。
実際にはそんなに単純ではないだろうが、それが理由ならかなり笑える。
https://www.linkedin.com/posts/eolver_googles-defense-agains...
Google 社員だが、しばらく TPU を見ていないなら v5 を確認してみることを勧める。今は PyTorch/JAX をサポートしていて、TensorFlow 専用だった頃よりずっと使いやすくなっている。
この記事は、抽象的に散らばっていたいくつもの断片が、実際にシリコンの中でどう流れるのかをうまくつなげてくれた。
単純な CISC 命令が LLM の推論段階にほぼそのまま対応しているのを見るのが特によかった。
無知をさらす馬鹿な質問かもしれないが、コンシューマー方面では M1〜M4 チップ が一部の AI 作業に向いているという話をずっと聞いている。
最近の自分にとって最も重要なのは Photoshop や Resolve のようなツールだが、Apple の新しい独自チップでは、古い自分のマシンよりはるかに速く動くのを見た。
これがこのチップや H100 にできることとうまくつながらない可能性もあるが、多少なりともつながるのか気になっている。もちろん Apple は独自チップを単体では売っていないので、実用化するには外付けサーバーのような形で GPU と AI チップを大量に載せた製品を出す必要があるだろう。
M* チップはユニファイドメモリを使っており、特に Pro/Max/Ultra は 1080 のような GPU と比べてもメモリ帯域幅が非常に高い。M1 Ultra のメモリ帯域幅は 2080 と 3090 の間くらいだ。
小さなバッチサイズ、特にローカル作業の大半のようなバッチ 1 では、推論は演算能力よりメモリ帯域幅に制約される。だから M* チップは機械学習に向いているという話になる。
しかし H100 は主に巨大なバッチサイズでの学習に使われ、大きなモデルを学習するには多くの相互接続が必要になる。その規模では算術強度が非常に高いため、M* チップはネットワークで束ねられるとしても競争力は大きくない。H100 のように電力を大量に消費するチップとは、電力/効率のパレート曲線上で別の点を選んでいるということだ。
Googleが本当にやるべきことは、2nm EUVの領域に入り、2nm未満へ進むことだ。
電子リソグラフィであれ、ASMLがチップに焼き付ける技術であれ、そういうものを手にしたら本当に危険な存在になる。Google X式のハードコアなムーンショット・プロジェクトが必要だと思う。
あるいは、装置1台を買うための5億ドルくらいは持っているのかもしれない。TPUが本当にそれほど優れているなら、自社技術と自社ファブまで垂直統合する良い事業になり得る。
現在の先端ファブが利益を出している理由は、多様な顧客と用途向けに汎用コンピューティングデバイスを作っており、数十年にわたって蓄積された人材とエンジニアリングがあるからだ。加えて、顧客がMicronのチップオンチップHDI歩留まり改善や、Xilinxのダイ間通信ファブリックとマルチチップ基板設計のような重要領域で、独立してイノベーションを後押ししている。
TPUは必要な数量を絶対に作れないし、収益性のある規模の経済を生み出す顧客も引き寄せられない。Googleは競合他社に対して魅力的な価格も提示しなければならない。
十分に説得力のある事業性があれば、既存のファブは喜んで生産枠を割り当てるはずだ。TPUはまだまったくそのレベルで説得力がない。
TensTorrentのJim Kellerが、AIコアを作る別のアプローチを説明した発表を聞いた。RISC-Vコアを5個置き、1つはデータ読み込み、1つはデータアップロード、残りは行列演算専用に使う方式だ。
彼はGoogle TPUにも触れていて、プログラミングはVLIWを扱うような感じで、コンパイラに約500人が付いていたと言っていた。
原文に「TPU v1は約20個の命令しか持たないCISC設計」とあるが、CISC/RISCは鋭い観察から研究プログラム、革命的技術、マーケティングの流行語を経て、ついには完全に意味のない言葉にまで来たようで笑える。
用語のライフサイクルと言うべきかもしれない。
だからTPUは命令数が少なくても、各命令がかなり複雑ならCISCと言えるかもしれない。ただ、最後にコンピュータアーキテクチャを履修したのは15年前の大学院の授業なので、記憶はあいまいだ。その学期の大半を、今では役に立たなくなったItanium関連の作業に費やしたこともある。
ファウンドリの生産枠需要はものすごく大きそうだが、MicrosoftやGoogleが独自チップを作って生産が必要になったとき、どうやって列の最前列に行くのか気になる。
十分に単純で、「古くて需要が少ない」ファブでも作れるものなのだろうか？ AppleとNvidiaはすでに多くのファウンドリ生産枠を押さえていると理解している。
https://en.wikipedia.org/wiki/Tensor_Processing_Unit#Product...
HBMのような領域にもかなりの存在感と支出があり、SemiAnalysisに関連する良い記事がある。
LLMが本当に**-1, 0, 1 量子化**で大きく定着したら、ハードウェアがどう変わるのか気になる。