TPU 是什麼?專為 AI 設計的專用晶片

TPU,TPU 的全名是 Tensor Processing Unit,中文叫「張量處理器」,它是 Google 專門為了人工智慧設計的晶片。

圖片來源:Ironwood TPUs: Built for the age of inference

簡單來說

一般電腦用的 CPU 什麼都能做,但做 AI 計算不夠快,顯示卡用的 GPU 雖然比較快,但本來是設計來玩遊戲畫圖的,拿來做 AI 還是有點不太合身。

TPU 不一樣,它從頭到尾只為 AI 而活,它最擅長的就是一直做「矩陣乘法」,這剛好是神經網路最需要的計算。

TPU 的設計哲學

根據 Google 發表的研究論文,TPU 的設計有幾個特點:

  • 只專心做一件事:放棄那些用不到的功能(像是快取、亂序執行),把所有力氣放在矩陣計算上
  • 資料盡量留在晶片上:減少搬來搬去的次數,省時又省電
  • 聽話就好:不需要太聰明的猜測功能,一切聽從軟體指揮,這樣時間才抓得準

這種「犧牲彈性,換取效率」的做法,就是 TPU 最核心的設計理念。

CPU、GPU、TPU 到底差在哪?

我用一個廚房的比喻來說明:

晶片類型比喻說明
CPU五星級主廚什麼菜都會做,但一次只能做幾道,客人一多就忙不過來
GPU一整排切菜工只會切菜,但一次可以切幾千顆馬鈴薯,效率超高
TPU專門烤麵包的烤箱只會烤麵包,但一次可以烤幾萬片,又快又省電

簡單說:CPU 是通才,GPU 是大量工讀生,TPU 是專門為 AI 打造的終極工具

Google 為什麼要自己做晶片?

用量太大,買的跟不上

2015 年的時候,Google 的服務已經大量使用 AI:

  • 語音搜尋
  • 圖片辨識
  • Google 翻譯
  • 搜尋結果排名

每天有幾十億次請求,如果用 CPU,要買好幾萬顆才夠用,電費會嚇死人。如果用 GPU,雖然好一點,但 GPU 本來是設計來打電動的,拿來做 AI 還是有很多用不到的功能在空轉,而且 NVIDIA 的晶片不便宜。

摩爾定律變慢了

以前每 18 個月晶片效能就翻一倍,現在這個速度已經慢下來了。但 AI 模型長大的速度卻越來越快,像 ChatGPT 這種大型模型,參數數量每幾個月就翻一倍,這個差距,只能靠做專用晶片來補。

不想被別人掐脖子

如果全世界只有 NVIDIA 能供應 AI 晶片,萬一人家漲價或缺貨,Google 就慘了,自己做晶片,長期來看反而更省錢,也不用擔心被綁架。

TPU 的演進過程和效能有多猛?

每一代有什麼不一樣?

根據 Google 發表的論文,各代 TPU 的差異如下:

世代什麼時候出的主打功能重點
TPU v12015專門做推理(模型訓練好後的預測)6 萬 5 千個計算單元,92 兆次運算/秒
TPU v22017可以訓練模型了晶片、系統、軟體一起設計
TPU v32018大規模訓練,用水冷散熱可以串聯幾千顆晶片一起工作
TPU v42021用光傳訊號,加入稀疏計算加速器比 v3 快 2.1 倍,省電 2.7 倍

到底有多快?

第一代 TPU v1

  • 比同期的伺服器 GPU(NVIDIA K80)快 15 到 30 倍
  • 比同期的 Intel CPU 快 30 到 50 倍
  • 省電程度更是驚人:每瓦特效能比 GPU 好 30 倍,比 CPU 好 80 倍

第四代 TPU v4

  • 比 NVIDIA 當時最強的 A100 GPU 快 1.2 到 1.7 倍
  • 比 A100 省電 1.3 到 1.9 倍
  • 比另一家 AI 晶片公司 Graphcore 的產品快 4.3 到 4.5 倍

不是所有情況都贏

有研究指出,TPU 和 GPU 各有擅長的地方,要看訓練的是什麼樣的模型。有些結構比較特別的神經網路,TPU 不見得跑得比 GPU 快,需要靠軟體調整才能發揮實力。

TPU 有哪些黑科技?

脈動陣列:像流水席一樣的計算方式

這是 TPU 最厲害的設計,一般的晶片算東西,要不斷去記憶體拿資料、算完、存回去,光是搬來搬去就浪費很多時間。

TPU 的脈動陣列就像流水席:資料像菜一樣一直流,經過一個計算單元就處理一點,再流到下一個。這樣計算和搬運同時進行,效率當然高。

光學互連:用光速傳資料

當好幾千顆晶片要一起工作時,互相溝通就變成大問題。傳統用電線傳,又慢又熱。

TPU v4 開始用光訊號代替電訊號,速度快、發熱低,幾千顆晶片可以像一個團隊一樣順暢合作。

液冷散熱:幫晶片泡冷水

從 TPU v3 開始,晶片已經熱到用風扇吹不動了。Google 直接讓晶片泡在特殊的冷卻液裡,這樣才能一直全力運轉不會過熱。

SparseCore:專門對付稀疏資料

現代的 AI 模型裡有一種叫「嵌入層」的東西,它像一個巨大的表格,大部分格子是空的,只有少數格子有資料。這種「稀疏」的資料,傳統晶片處理起來很沒效率。

TPU v4 加入了一個叫 SparseCore 的專用加速器,專門處理這種狀況,速度可以快 5 到 7 倍。Google 的搜尋廣告、推薦系統都很需要這個功能。

怎麼寫程式給 TPU 跑?

用什麼工具?

TPU 主要用 TensorFlow 或 JAX 這兩個框架來寫程式,然後透過 XLA 編譯器把程式轉成 TPU 聽得懂的指令。

怎麼讓幾千顆晶片一起工作?

要用好幾千顆 TPU 訓練一個超大模型,需要一些技巧:

  • 把模型切成很多塊,分給不同晶片算
  • 把權重資料也分開存放
  • 安排好計算的先後順序
  • 用編譯器幫忙最佳化

除了 AI 還能做什麼?

你可能覺得 TPU 只能跑神經網路,但學術界發現,它也能用在其他科學計算:

  • 超大矩陣乘法:用 2048 顆 TPU 算 100 萬 × 100 萬的矩陣,2 分鐘就搞定,傳統超級電腦要幾小時甚至幾天
  • 模擬雲層:比傳統 CPU 快 10 倍以上,以前要算一週的氣象模擬,現在一天就能算完
  • 醫學影像:MRI 影像重建的速度可以跟 NVIDIA V100 GPU 比

TPU 有什麼缺點?

不是所有模型都適合

TPU 擅長的是「整齊」的大規模矩陣運算,如果你的模型有很多 if-else 判斷、很多迴圈,或者有很多小而不規則的矩陣,TPU 就跑不快。

寫程式有門檻

要讓幾千顆 TPU 好好工作,需要對 TensorFlow/JAX 很熟悉,還要懂得怎麼切分模型,Google 有提供編譯器幫忙,但還是有學習成本。

記憶體不夠用

有時候晶片上的記憶體會不夠,或者某些層的計算會卡住,導致 TPU 無法跑到最快的速度。

生態系還不如 NVIDIA

NVIDIA 的 GPU 搭配 CUDA 軟體,已經發展很多年了,幾乎所有 AI 工具都對它做過最佳化。TPU 雖然支援 TensorFlow 和 JAX,但一些比較新的模型或冷門的功能,支援度可能沒那麼快。

未來 TPU 會怎麼發展?

硬體還會更進步

根據研究,未來可能有這些方向:

方向簡單說明效果
可重構資料流晶片可以自己調整計算方式模擬顯示可快 2.75 倍
近閾值計算電壓壓到極限來省電可快 2.5 倍,準確度損失小
更多專用加速器除了 SparseCore 再加其他特化功能特定任務更快

軟體也會更好用

未來要繼續改善編譯器、讓異质晶片更好搭配、減少晶片之間的通訊成本,這樣才能真正把理論速度變成實際速度。

與其做一把什麼都能切的萬用刀,不如做一把專門切某種東西的超級刀

Google TPU 的故事告訴我們一件事:與其做一把什麼都能切的萬用刀,不如做一把專門切某種東西的超級刀。

從第一代專做推理,到第四代變成超大規模訓練的超級電腦,TPU 證明了「領域專用」的價值。它不只用在 Google 自己的服務(搜尋、翻譯、Gmail、YouTube),也被學術界拿去研究氣象、物理、醫學影像。

如果把 CPU 比作燭光、GPU 比作電燈,那 TPU 就像正在升起的太陽,專門為 AI 而生,光芒萬丈。

未來,會有更多公司推出自己的 AI 專用晶片(像亞馬遜的 Inferentia、特斯拉的 Dojo),我們正進入一個算力百花齊放的時代。而 TPU 作為先行者,已經證明了:

在特定領域,把簡單的事做到極致,就是最強大的力量。

參考文獻

  1. Jouppi, N. P., Young, C., Patil, N., et al. (2017). In-Datacenter Performance Analysis of a Tensor Processing Unit. Proceedings of the 44th Annual International Symposium on Computer Architecture. DOI: 10.1145/3079856.3080246
  2. Norrie, T., Patil, N., Yoon, D. H., et al. (2021). The Design Process for Google’s Training Chips: TPUv2 and TPUv3. IEEE Micro, 41(2), 56-63. DOI: 10.1109/MM.2021.3058217
  3. Kumar, S., Bitorff, V., Chen, D., et al. (2019). Scale MLPerf-0.6 Models on Google TPU-v3 Pods. arXiv Preprint, arXiv:1909.09756.
  4. Kumar, S., Bradbury, J., Young, C., et al. (2020). Exploring the Limits of Concurrency in ML Training on Google TPUs. arXiv Preprint, arXiv:2011.03641.
  5. Wang, Y. E., Wei, G. Y., & Brooks, D. (2019). Benchmarking TPU, GPU, and CPU Platforms for Deep Learning. arXiv Preprint, arXiv:1907.10701.
  6. Jouppi, N. P., Kurian, G., Li, S., et al. (2023). TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings. Proceedings of the 50th Annual International Symposium on Computer Architecture. DOI: 10.1145/3579371.3589350 / arXiv:2304.01433.
  7. Kravets, N., Khovrat, A., & Saichyshyna, N. (2020). Analysis Of Tensor Processors Application In Machine Learning Tasks On The Example Of GOOGLE TPU. Bionics of Intelligence, 1(94). DOI: 10.30837/bi.2020.1(94).13
  8. Bhattarai, B. (2022). A TensorFlow simulation framework for scientific computing of fluid flows on tensor processing units. Computer Physics Communications, 274, 108292. DOI: 10.1016/j.cpc.2022.108292
  9. Large-scale distributed linear algebra with tensor processing units. (2022). Proceedings of the National Academy of Sciences, 119(33), e2122762119. DOI: 10.1073/pnas.2122762119
  10. Accelerating large-eddy simulations of clouds with Tensor Processing Units. (2023). arXiv Preprint, arXiv:2301.04698.
返回頂端