TPU的系統(tǒng)架構(gòu)

簡單解釋：專門用于機(jī)器學(xué)習(xí)的高性能芯片，圍繞128x128 16 位乘法累加脈動陣列矩陣單元（“MXU”）設(shè)計(jì)的加速器。如果這句話能為你解釋清楚，那就太好了！如果沒有，那么請繼續(xù)閱讀......

介紹

您可能聽說過 Google 有一個用于機(jī)器學(xué)習(xí)的特殊芯片，它被稱為 TPU（“Tensor Processing Unit，張量處理單元”），它構(gòu)成了 Google 努力將盡可能多的機(jī)器學(xué)習(xí)能力放入單個芯片中。谷歌云為開發(fā)者提供了使用這種能力的機(jī)會，但芯片本身感覺是一個黑匣子……今天，讓我們剝離這些隱藏的“面紗”，看看我們是否能看到魔法里面的東西。

在本文中，我將嘗試講解 TPU 的系統(tǒng)架構(gòu)，同時保持足夠簡單，以使硬件經(jīng)驗(yàn)最少的軟件開發(fā)人員也能看懂。

高性能推理

訓(xùn)練和運(yùn)行神經(jīng)網(wǎng)絡(luò)需要大量的計(jì)算能力。早在 2013 年，Google 進(jìn)行了一些簡單的計(jì)算，以了解他們需要什么來運(yùn)行語音搜索，結(jié)果令人驚訝：

如果我們考慮人們每天只使用三分鐘谷歌語音搜索的場景，并且我們在我們使用的處理單元上為我們的語音識別系統(tǒng)運(yùn)行深度神經(jīng)網(wǎng)絡(luò)，我們將不得不將谷歌數(shù)據(jù)中心的數(shù)量增加一倍！ — Norm Jouppi( Google 的杰出工程師。他以在計(jì)算機(jī)內(nèi)存系統(tǒng)方面的創(chuàng)新而聞名)

ASIC

神經(jīng)網(wǎng)絡(luò)是強(qiáng)大的工具，但要在標(biāo)準(zhǔn)計(jì)算機(jī)上隨處運(yùn)行，它們的成本太高（即使對谷歌來說也是如此）。

值得慶幸的是，這不需要標(biāo)準(zhǔn)計(jì)算機(jī)來完成繁重的工作。在某些時候，設(shè)計(jì)一個定制芯片來承載這個任務(wù)變得具有成本效益。此定制芯片是專用集成電路(ASIC)。

通常，ASIC 帶來的麻煩多于其價值。他們需要很長時間來設(shè)計(jì)：Google 花了15 個月的時間來開發(fā) TPUv1，這個速度快得驚人。它們最初很昂貴，需要專門的工程師和大約一百萬美元的制造成本。而且它們不靈活：一旦完成，就無法更換芯片設(shè)計(jì)。

但是，如果ASIC的數(shù)量足夠，那么經(jīng)濟(jì)性的好處可以彌補(bǔ)最初的缺點(diǎn)。ASIC 通常是完成任務(wù)的最快、最節(jié)能的方式。谷歌希望這種性能能夠運(yùn)行神經(jīng)網(wǎng)絡(luò)，而 TPU 就是結(jié)果。

標(biāo)量、向量、矩陣

神經(jīng)網(wǎng)絡(luò)需要大量的數(shù)學(xué)運(yùn)算，但大多數(shù)數(shù)學(xué)運(yùn)算都非常簡單：將一堆數(shù)字相乘，然后將結(jié)果相加。可以在一個稱為乘法累加(MAC) 的操作中將這兩者連接在一起。如果我們不需要做任何其他事情，我們可以非常非常快地進(jìn)行乘法累加。

如果沒有新芯片，我們將使用 CPU 或 GPU 來實(shí)現(xiàn)。CPU 是一臺標(biāo)量機(jī)器，這意味著它一次處理一個指令。這非常適合通用應(yīng)用程序，例如您的筆記本電腦或服務(wù)器，但我們可以通過專業(yè)化來擠出更多性能。

數(shù)據(jù)維度

一個簡化的向量架構(gòu)

GPU 是一個向量機(jī)(vector machine)。你可以給它一個很長的數(shù)據(jù)列表——一個一維向量——并同時在整個列表上運(yùn)行計(jì)算。這樣，我們每秒可以執(zhí)行更多的計(jì)算，但我們必須對數(shù)據(jù)向量并行執(zhí)行相同的計(jì)算。這種計(jì)算非常適合圖形或加密挖掘，其中一項(xiàng)工作必須執(zhí)行多次。

但我們?nèi)匀豢梢宰龅酶?。神?jīng)網(wǎng)絡(luò)的數(shù)據(jù)以矩陣形式排列，即二維向量。因此，我們將構(gòu)建一個矩陣機(jī)(matrix machine)。而且我們真的只關(guān)心乘法累加，所以我們會優(yōu)先考慮處理器通常支持的其他指令。我們將把大部分芯片用于執(zhí)行矩陣乘法的 MAC，而忽略大多其他操作。

重復(fù)這個N 次，你就得到了圖片

脈動陣列Enter the Systolic Array

提升矩陣計(jì)算性能的方法是通過一種稱為脈動陣列的架構(gòu)。這是有趣的一點(diǎn)，這也是 TPU 具有高性能的原因。脈動陣列是一種硬件算法，它描述了計(jì)算矩陣乘法的芯片上的單元模式。“Systolic”描述了數(shù)據(jù)如何在芯片中以波浪的形式移動，就像人類心臟的跳動。

TPU 中實(shí)現(xiàn)脈動陣列版本設(shè)計(jì)有一些變化。

考慮一個矩陣乘法運(yùn)算：

2x2 矩陣相乘

對于 2x2 輸入，輸出中的每一項(xiàng)都是兩個乘積的總和。沒有元素被重復(fù)使用，但個別元素被重復(fù)使用。

我們將通過構(gòu)建一個 2x2 網(wǎng)格來實(shí)現(xiàn)這一點(diǎn)。（它實(shí)際上是一個網(wǎng)格，而不僅僅是一個抽象——硬件就是這樣有趣）。請注意，2x2 是一個玩具示例，而全尺寸 MXU 是一個巨大的 128x128。

假設(shè) AB/CD 代表我們的激活值，EF/GH代表我們的權(quán)重。對于我們的數(shù)組，我們首先 像這樣加載權(quán)重：

稍后我將討論我們?nèi)绾巫龅竭@一點(diǎn)

接下來激活進(jìn)入輸入隊(duì)列，在我們的示例中，該隊(duì)列位于每一行的左側(cè)。

注意用零填充：這可以確保數(shù)據(jù)在正確的時刻進(jìn)入數(shù)組

每個時鐘周期 ，每個單元都會并行執(zhí)行以下步驟：

乘以我們的權(quán)重和來自左側(cè)的激活。如果左側(cè)沒有單元格，則從輸入隊(duì)列中取出。

將該產(chǎn)品添加到從上面?zhèn)魅氲牟糠挚偤椭?。如果上面沒有單元格，則上面的部分總和為零。

將激活傳遞到右側(cè)的單元格。如果右側(cè)沒有單元格，則丟棄激活。

將部分總和傳遞到底部的單元格。如果底部沒有單元格，則收集部分總和作為輸出。

上面的Python偽代碼：

?

https://github.com/antonpaquin/SystolicArrayDemo/blob/master/systolic.py

通過這些規(guī)則，可以看到激活將從左側(cè)開始，每個周期向右移動一個單元格，部分和將從頂部開始，每個周期向下移動一個單元格。

數(shù)據(jù)流將如下所示：

這就是上面構(gòu)建的陣列！讓我們來看看第一個輸出的執(zhí)行：

周期 1

左上角從輸入隊(duì)列中讀取 A，乘以權(quán)重 E 以產(chǎn)生產(chǎn)品 AE。

AE 從上面添加到部分和 0，產(chǎn)生部分和 AE。

激活 A 傳遞到右上角的單元格。

部分和 AE 傳遞到左下角的單元格。

周期 2

左下角從輸入隊(duì)列中讀取 B，乘以權(quán)重 G 生成產(chǎn)品 BG

BG 從上面添加到部分和 AE，產(chǎn)生部分和 AE + BG

激活 B 傳遞到右下角的單元格

部分和 AE + BG 是一個輸出。

可以看到我們已經(jīng)正確計(jì)算了輸出矩陣的第一項(xiàng)。同時，在第 2 周期中，我們還在左上角計(jì)算 CE，在右上角計(jì)算 AF。這些將通過單元格傳播，到第 4 周期，將產(chǎn)生整個 2x2 輸出。

這里有一些來自谷歌的圖表，它們提供了更多的視覺畫面。

流經(jīng)數(shù)組的數(shù)據(jù)動畫。相對于本文中的其他圖像，這個 gif 是轉(zhuǎn)置的——順時針旋轉(zhuǎn) 90 度，然后水平翻轉(zhuǎn)。

會看到輸入激活與零交錯，以確保它們在正確的時刻進(jìn)入數(shù)組，并且離開數(shù)組的輸出也同樣交錯。完全計(jì)算結(jié)果矩陣需要 3n-2 個周期，而標(biāo)準(zhǔn)的順序解決方案是 n3。這是一個不錯的結(jié)果，將計(jì)算量大大降低。

【科普】Xilinx 3D IC技術(shù)簡介

谷歌的 MXU 圖。

因?yàn)槲覀冋诓⑿羞\(yùn)行 128x128 MAC 操作。在硬件中實(shí)現(xiàn)乘法器通常很大且成本很高，但脈動陣列的高密度讓 Google 可以將其中的 16,384 個裝入 MXU。這直接轉(zhuǎn)化為速度訓(xùn)練和運(yùn)行網(wǎng)絡(luò)。

權(quán)重的加載方式與激活方式大致相同——通過輸入隊(duì)列。我們只是發(fā)送一個特殊的控制信號（上圖中的紅色）來告訴數(shù)組在權(quán)重經(jīng)過時存儲權(quán)重，而不是運(yùn)行 MAC 操作。權(quán)重保留在相同的處理元素中，因此我們可以在加載新集合之前發(fā)送整個批次，從而減少開銷。

就是這樣！芯片的其余部分很重要，值得一試，但 TPU 的核心優(yōu)勢在于它的 MXU——一個脈動陣列矩陣乘法單元。

TPU的其余部分

上面設(shè)計(jì)了出色的脈動陣列，但仍有大量工作需要構(gòu)建支持和基礎(chǔ)部分以使其運(yùn)行。首先，我們需要一種將數(shù)據(jù)輸入和輸出芯片本身的方法。然后我們需要在正確的時間將它進(jìn)出數(shù)組。最后，我們需要一些方法來處理神經(jīng)網(wǎng)絡(luò)中不是矩陣乘法的內(nèi)容。讓我們看看這一切是如何在硬件中發(fā)生的。

完整的系統(tǒng)

下面是舊 TPUv1 的系統(tǒng)圖和布局模型。這些不一定適用于較新的 TPUv2 和 TPUv3 — 自本文檔發(fā)布以來，Google 已經(jīng)進(jìn)行了一些架構(gòu)更改。

在最高級別，TPU 被設(shè)計(jì)為加速器。這意味著它將插入主機(jī)系統(tǒng)，主機(jī)將加載要在加速器上計(jì)算的數(shù)據(jù)和指令。結(jié)果通過相同的接口返回給主機(jī)。通過這種模型，加速器（TPU）可以加速耗時且昂貴的矩陣運(yùn)算，而主機(jī)可以處理其他所有事情。

讓我們用一些框圖來檢查加速器內(nèi)部的內(nèi)容。我們將逐步介紹這些。

TPUv1 的系統(tǒng)圖和布局模型

主機(jī)接口將通過 PCIe 連接到加速器（TPU）。我們可以看到通過這個接口傳輸 3 種形式的數(shù)據(jù)：權(quán)重（到DDR3）、激活activations（到統(tǒng)一緩沖區(qū)）和控制指令（到紅色控制路徑）。

批次中所有輸入的權(quán)重相同，因此每批次將加載一次。與激活相比，這種情況很少見，因此可以將權(quán)重存儲在速度較慢的片外DDR3 DRAM中。我們可以連接主機(jī)接口來寫入這個 DDR3，當(dāng)我們加載模型時我們會把所有的權(quán)重放在那里。在計(jì)算之前，權(quán)重從 DDR3 讀取到權(quán)重 FIFO中，這意味著我們可以在計(jì)算當(dāng)前批次時預(yù)取下一組權(quán)重。

統(tǒng)一緩沖區(qū)保存我們的激活。在操作期間，主機(jī)需要快速訪問此緩沖區(qū)，以讀取結(jié)果并寫入新輸入。統(tǒng)一緩沖器直接連接到MXU，這兩個組件占據(jù)了芯片空間的最大份額（53%）。該緩沖區(qū)最多可容納 384 個尺寸為 256x256 的激活矩陣，這將是芯片支持的最大批次。這些激活被非常頻繁地讀取和寫入，這就是為什么我們將 30% 的布局專用于片上內(nèi)存，以及 MXU 和統(tǒng)一緩沖區(qū)路徑周圍的 167 GiB/s 總線。

MXU 通過Accumulators寫回統(tǒng)一緩沖區(qū)。首先，累加器從 MXU 中收集數(shù)據(jù)。然后，激活管道（Activation Pipeline）應(yīng)用標(biāo)準(zhǔn)的神經(jīng)網(wǎng)絡(luò)函數(shù)（如 ReLU 和 Maxpool），這些函數(shù)的計(jì)算量不如矩陣乘法。完成后，我們可以將輸出放回統(tǒng)一緩沖區(qū)，為下一階段的處理做好準(zhǔn)備。

然后是控制流：從主機(jī)獲取指令的紅色路徑。該路徑將執(zhí)行諸如控制 MXU 何時乘法、權(quán)重 FIFO 將加載哪些權(quán)重或激活管道將執(zhí)行哪些操作等操作。把它想象成告訴芯片的其余部分做什么的leader。

完整的流程如下：

1、芯片啟動，緩沖區(qū)和 DDR3 為空

2、用戶加載 TPU 編譯的模型，將權(quán)重放入 DDR3 內(nèi)存

3、主機(jī)用輸入值填充激活緩沖區(qū)

4、發(fā)送控制信號將一層權(quán)重加載到 MXU（通過權(quán)重 FIFO）

5、主機(jī)觸發(fā)執(zhí)行，激活通過 MXU 傳播到累加器

6、當(dāng)輸出出來時，它們通過激活管道運(yùn)行，新層替換緩沖區(qū)中的舊層

7、重復(fù) 4 到 6 直到我們到達(dá)最后一層

8、最后一層的激活被發(fā)送回主機(jī)

TPUv1 推理的全貌。在較新的 TPUv2 中經(jīng)歷了類似的事情……

新一代 TPU 允許訓(xùn)練（即更新權(quán)重），因此需要有一條從 MXU 到重量存儲。在 TPUv1 框圖中，情況并非如此。

但是，只要知道 TPUv2 能做什么，我們就可以猜到一些不同之處：

TPUv1 中的 MXU 是一個 8 位整數(shù) 256x256 數(shù)組，比 TPUv2 中的 16 位 bfloat16 128x128 MXU 更大但精度更低。

TPUv1 中的激活管道被 TPUv2 中的全向量和標(biāo)量單元所取代。與 TPUv1 中有限的“激活”和“標(biāo)準(zhǔn)化/池”相比，這些功能為其提供了更廣泛的可用功能。

統(tǒng)一緩沖區(qū)似乎被高帶寬內(nèi)存所取代。這將為芯片的其余部分騰出一些空間，但代價是延遲。

這些差異主要是因?yàn)?TPUv1 是為推理而非訓(xùn)練而設(shè)計(jì)的，因此可以接受低精度算術(shù)和較少的操作。升級意味著新一代 TPU 更加靈活——足以讓谷歌輕松地將它們布置在他們的云上。

其他概念

bfloat16

大多數(shù) CPU/GPU 機(jī)器學(xué)習(xí)計(jì)算都是使用 32 位浮點(diǎn)數(shù)完成的。當(dāng)我們降到 16 位時，ML 工程師往往更擔(dān)心數(shù)字的范圍而不是精度。舍入小數(shù)點(diǎn)的幾個分?jǐn)?shù)是可以的，但是超出數(shù)字表示的最大值或最小值是一件令人頭疼的事情。傳統(tǒng)的 float16 精度很高，但范圍不夠。谷歌對此的回答是 bfloat16 格式，它有更多的位用于指數(shù)，而更少的位用于有效位。

為了比較，浮點(diǎn)數(shù)中的位是：

float32：1 位符號，8 位指數(shù)，23 位有效數(shù)

float16：1 位符號，5 位指數(shù)，10 位有效數(shù)

bfloat16：1 位符號，8 位指數(shù)，7 位有效數(shù)

bfloat16.

因?yàn)樗哂邢嗤瑪?shù)量的指數(shù)位，所以 bfloat16 只是 float32 的兩個高位字節(jié)。這意味著它的范圍與 float32 大致相同，但精度較低。在實(shí)踐中，這種策略效果很好。在 TPU 上，大部分?jǐn)?shù)據(jù)仍以 float32 格式存儲。但是，MXU 具有 bfloat16 乘法器和 float32 累加器。較小的乘法器降低了芯片的功率和面積要求，為更多以更快時鐘運(yùn)行的乘法器留出了空間。這提高了芯片的性能，而不會降低精度。

XLA

XLA 是一個用于 Tensorflow 后端的實(shí)驗(yàn)性 JIT 編譯器。它將您的 TF 圖轉(zhuǎn)換為線性代數(shù)，并且它有自己的后端可以在 CPU、GPU 或 TPU 上運(yùn)行。

Pods

Google云中的 TPU 存在于“pod”中，它們是具有大量計(jì)算能力的大型機(jī)架。每個 pod 可容納 64 個 TPUv2 板，用于 11.5 petaflops的工作性能。

單個 TPU 通常不足以以所需的速度訓(xùn)練大型模型，但訓(xùn)練涉及頻繁的權(quán)重更新，需要在所有相關(guān)芯片之間分配。Pod 使 TPU 保持緊密的網(wǎng)絡(luò)連接，因此分布在 Pod 上的所有芯片之間的訓(xùn)練幾乎可以線性擴(kuò)展。

這種線性縮放是一個重要的原則。如果數(shù)據(jù)科學(xué)家想要更多的權(quán)力，他們會添加更多的芯片，沒有任何問題。

TPU發(fā)展歷史

結(jié)論

這是我能找到有關(guān)TPU工作原理的所有信息，可能她并不完整，但是我希望你明白了TPU的工作原理。

TPU 是一個非常好的硬件，但它可能在 v1 出現(xiàn)之前就已經(jīng)存在多年了。這主要是因?yàn)楣雀枳隽撕芏嗑€性代數(shù)（由機(jī)器學(xué)習(xí)驅(qū)動），這種芯片在商業(yè)上變得可行。其他公司也抓住了這一趨勢：Groq是 TPU 項(xiàng)目的衍生產(chǎn)品，英特爾希望盡快開始銷售Nervana芯片，還有 更多的公司 希望加入競爭。專用硬件有望降低訓(xùn)練和運(yùn)行模型的成本；希望這將加快我們創(chuàng)新的步伐。

審核編輯 ：李倩