三個月、零基礎(chǔ)手搓一塊TPU，能推理能訓(xùn)練，還是開源的

2025-08-25 09:05:00

近日，來自加拿大西安大略大學(xué)的工程師 Surya Sure 等人宣布已經(jīng)利用暑假時間構(gòu)建出了?TinyTPU：一種開源的 ML 推理、訓(xùn)練芯片。

最近，大模型技術(shù)的發(fā)展，讓人們再次重視起 AI 專用芯片。

對于計算任務(wù)負(fù)載來說，越是專用，效率就越高，谷歌的 TPU 就是其中的一個典型例子。它自 2015 年開始在谷歌數(shù)據(jù)中心部署后，已經(jīng)發(fā)展到了第 7 代。目前的最新產(chǎn)品不僅使用了最先進(jìn)的制程工藝打造，也在架構(gòu)上充分考慮了對于機(jī)器學(xué)習(xí)推理任務(wù)的優(yōu)化。TPU 的出現(xiàn)，促進(jìn)了 Gemini 等大模型技術(shù)的進(jìn)展。

這種芯片為何性能如此強(qiáng)大？或許最直接的了解方法就是嘗試復(fù)刻它。近日，來自加拿大西安大略大學(xué)的工程師 Surya Sure 等人宣布已經(jīng)利用暑假時間構(gòu)建出了 TinyTPU：一種開源的 ML 推理、訓(xùn)練芯片。

項目 GitHub 地址：https://github.com/tiny-tpu-v2/tiny-tpu

有趣的是，他們并非芯片設(shè)計專業(yè)的學(xué)生，打造 TPU 是從理解多層感知機(jī)（MLP）這樣的神經(jīng)網(wǎng)絡(luò)基本概念開始的。為此，他們還手工計算了網(wǎng)絡(luò)推理和訓(xùn)練所需的數(shù)學(xué)運(yùn)算。

讓我們看看他是如何做到的。

手搓 TPU 的原因

選擇造一塊 TPU 的原因很簡單：

構(gòu)建一個用于機(jī)器學(xué)習(xí)工作負(fù)載的芯片看起來很酷；
之前還沒有一個同時進(jìn)行推理和訓(xùn)練的機(jī)器學(xué)習(xí)加速器的完整開源代碼庫。

我們都沒有真正的硬件設(shè)計專業(yè)經(jīng)驗(yàn)，這在某種程度上使得 TPU 更具吸引力，因?yàn)槲覀儫o法準(zhǔn)確估計它的難度。在項目初期，我們確立了嚴(yán)格的設(shè)計理念：始終嘗試「Hacky Way」（不靠譜的方法）。這意味著在咨詢外部資源之前，先嘗試那些我們首先想到的「愚蠢」想法。這種理念確保了我們沒對 TPU 進(jìn)行逆向工程，而是對其進(jìn)行重新發(fā)明，因此我們自己推導(dǎo)出了 TPU 中使用的許多關(guān)鍵機(jī)制。

我們也想把這個項目當(dāng)作一次不依賴人工智能代寫代碼的練習(xí)，因?yàn)槲覀兏杏X最近每當(dāng)遇到小問題時，我們的第一反應(yīng)就是求助于這些 AI 工具。我們希望培養(yǎng)一種特定的思維方式，以便我們能夠?qū)⑵涓吨T實(shí)踐，并在未來的任何努力中運(yùn)用它來解決難題。

在整個項目過程中，我們盡可能多地學(xué)習(xí)深度學(xué)習(xí)、硬件設(shè)計和算法創(chuàng)建方面的基礎(chǔ)知識。我們發(fā)現(xiàn)，學(xué)習(xí)這些知識的最佳方式是將所有知識都畫出來，并將其作為我們的第一反應(yīng)。閱讀這篇文章時，你會發(fā)現(xiàn)我們的解釋是如何受到這種理念的啟發(fā)的。

在繼續(xù)之前，我們想先明確一下這篇文章涵蓋的內(nèi)容和不涵蓋的內(nèi)容。請注意，本文并非 TPU 的 1:1 復(fù)制品 —— 而是我們自己重新發(fā)明 TPU 的嘗試。

什么是 TPU？

TPU 是谷歌設(shè)計的一款專用芯片 (ASIC)，旨在專門提高機(jī)器學(xué)習(xí)模型的推理和訓(xùn)練速度。GPU 既可以渲染圖像，又可以運(yùn)行機(jī)器學(xué)習(xí)工作負(fù)載，而 TPU 專用于執(zhí)行數(shù)學(xué)運(yùn)算，這使得它能夠更高效率地完成其設(shè)計初衷。在芯片領(lǐng)域，嘗試掌握單個任務(wù)比嘗試掌握多個任務(wù)更容易，而且效果也更好，而 TPU 正是秉承了這一理念。

硬件設(shè)計快速入門：

在硬件中，我們處理的時間單位稱為時鐘周期。作為開發(fā)者，我們可以根據(jù)需要設(shè)置任意的時間段。通常，單個時鐘周期的范圍從 1 皮秒 (ps) 到 1 納秒 (ns)，我們運(yùn)行的任何操作都將在時鐘周期之間執(zhí)行。

時鐘周期時序圖展示了硬件中操作的同步方式

我們用來描述硬件的語言叫做 Verilog。它是一種硬件描述語言，允許我們描述給定硬件模塊的行為（類似于軟件中的函數(shù)），但它不是以程序的形式執(zhí)行，而是合成布爾邏輯門（AND 與、OR 或、NOT 非等），這些邏輯門可以組合起來，構(gòu)建我們想要的任何芯片的數(shù)字邏輯。以下是 Verilog 中加法運(yùn)算的一個簡單示例：

在上面的例子中，信號 b 在下一個時鐘周期的值被設(shè)置為信號 a 的當(dāng)前值。你會發(fā)現(xiàn)，在大多數(shù)情況下，信號（變量）是按順序的時鐘周期更新的，而不是像軟件設(shè)計中那樣立即更新。

具體來說，TPU 在執(zhí)行矩陣乘法方面非常高效，它占 Transformer 計算操作的 80-90%（在超大型模型中高達(dá) 95%），以及 CNN 計算操作的 70-80%。每個矩陣乘法都代表 MLP（多層感知器）中單個層的計算，而在深度學(xué)習(xí)中，我們擁有多層級的 MLP，這使得 TPU 在處理大型模型時更加高效。

那么如何構(gòu)建一塊 TPU 呢？

當(dāng)我們開始這個項目時，我們只知道方程 y = mx + b 是神經(jīng)網(wǎng)絡(luò)的基礎(chǔ)構(gòu)建塊。然而，我們需要完全理解神經(jīng)網(wǎng)絡(luò)背后的數(shù)學(xué)原理，才能在 TPU 中構(gòu)建其他模塊。因此在開始編寫任何代碼之前，我們每個人都計算出了一個簡單的 2 → 2 → 1 多層感知器 (MLP) 的數(shù)學(xué)原理。

用于解決 XOR 問題的 2→2→1 多層感知器的架構(gòu)

我們選擇這個特定網(wǎng)絡(luò)的原因是，我們的目標(biāo)是針對 XOR 問題（神經(jīng)網(wǎng)絡(luò)的「Hello World」）進(jìn)行推理和訓(xùn)練。XOR 問題是神經(jīng)網(wǎng)絡(luò)可以解決的最簡單的問題之一。所有其他門（AND、OR 等）都可以僅使用一條線性線（一個神經(jīng)元）來區(qū)分哪些輸入對應(yīng)于 0，哪些輸入對應(yīng)于 1，從而根據(jù)輸入預(yù)測輸出。但是，要對所有 XOR 進(jìn)行分類，需要一個 MLP，因?yàn)樗枰獜澢臎Q策邊界，而這僅靠線性方程無法實(shí)現(xiàn)。

對于幾何和第一性原理的論述，免費(fèi)書籍《理解深度學(xué)習(xí)》（Understanding Deep Learning）非常值得一讀：https://udlbook.github.io/udlbook/。

現(xiàn)在，假設(shè)我們要進(jìn)行連續(xù)推理（例如，自動駕駛汽車每秒進(jìn)行多個預(yù)測）。這意味著我們需要同時發(fā)送多條數(shù)據(jù)。由于數(shù)據(jù)本質(zhì)上是多維的，并且具有許多特征，因此我們需要非常大的矩陣維度。然而，XOR 問題簡化了維度，因?yàn)橹挥袃蓚€特征（0 或 1）和 4 種可能的輸入數(shù)據(jù)（0 和 1 的四種二進(jìn)制組合）。這為我們提供了一個 4x2 矩陣，其中 4 是行數(shù)（批處理大小 batch size），2 是列數(shù)（特征大小）。

XOR 輸入矩陣和目標(biāo)輸出：

每一行代表四個可能的異或輸入之一，輸出向量顯示預(yù)期的異或結(jié)果

我們對脈動陣列示例進(jìn)行的另一項簡化是，我們將使用 2x2 陣列，而不是 TPUv1 中使用的 256x256 陣列。然而，數(shù)學(xué)運(yùn)算仍然忠實(shí)于原有的格式，因此實(shí)際上并沒有簡化，而是進(jìn)行了規(guī)模縮小。

等式的第一步是將 m 與 x 相乘，以矩陣形式表示為 XW^T：

其中 X 是我們的輸入矩陣，W 是我們的權(quán)重矩陣，b 是我們的偏差向量。

我們?nèi)绾卧谟布袌?zhí)行矩陣乘法？嗯，我們可以使用一個叫做脈動陣列的單元！

TPU 的核心是一個叫做脈動陣列（systolic array）的單元。它由稱為處理單元 (PE) 的獨(dú)立構(gòu)建塊組成，這些構(gòu)建塊以網(wǎng)格狀結(jié)構(gòu)連接在一起。每個 PE 執(zhí)行乘法 - 累加運(yùn)算，這意味著它將傳入的輸入 X 與固定權(quán)重 W 相乘，并將其與傳入的累加和相加，所有這些都在同一個時鐘周期內(nèi)完成。

處理元件（PE）架構(gòu)顯示乘法累加運(yùn)算（無負(fù)載權(quán)重和啟動標(biāo)志）。

當(dāng)這些 PE 連接在一起時，它們可以以脈動方式執(zhí)行矩陣乘法，這意味著每個時鐘周期都可以計算輸出矩陣的多個元素。輸入從左側(cè)進(jìn)入脈動陣列，然后每個時鐘周期移動到右側(cè)相鄰的 PE。累加和從第一行 PE 的乘法輸出開始，向下移動，并與每個連續(xù) PE 的乘積相加，直到到達(dá)最后一行 PE，成為輸出矩陣的一個元素。

脈動陣列架構(gòu)展示了 PE 如何連接以執(zhí)行矩陣乘法

由于 TPU 只有一個單元（并且矩陣乘法在模型中占據(jù)了主要計算量），因此它可以非常輕松地推理和訓(xùn)練任何模型。

示例

現(xiàn)在，讓我們來看看異或問題的示例。

我們的脈動陣列接受兩個輸入：輸入矩陣和權(quán)重矩陣。對于我們的異或網(wǎng)絡(luò)，使用以下權(quán)重和偏差進(jìn)行初始化：

輸入和權(quán)重調(diào)度

為了將輸入批次輸入到脈動陣列中，我們需要：

將 X 矩陣旋轉(zhuǎn) 90 度。

矩陣旋轉(zhuǎn) 90 度，為脈動陣列輸入做準(zhǔn)備

錯開輸入（將每行延遲 1 個時鐘周期）。

用于脈動陣列處理的輸入矩陣交錯模式

要輸入我們的權(quán)重矩陣，我們需要：

交錯排列權(quán)重矩陣（與輸入類似）

用于脈動陣列處理的權(quán)重矩陣交錯模式

轉(zhuǎn)置它！

權(quán)重矩陣轉(zhuǎn)置以實(shí)現(xiàn)正確的數(shù)學(xué)對齊。

請注意，旋轉(zhuǎn)和交錯沒有任何數(shù)學(xué)意義 —— 它們只是為了使脈動陣列正常工作。轉(zhuǎn)置也只是為了進(jìn)行數(shù)學(xué)記賬 —— 它是使矩陣數(shù)學(xué)運(yùn)算正常工作所必需的，因?yàn)槲覀冊谏窠?jīng)網(wǎng)絡(luò)圖中設(shè)置權(quán)重指針的方式。

交錯和先進(jìn)先出 (FIFO)

為了執(zhí)行交錯，我們?yōu)闄?quán)重和輸入設(shè)計了幾乎相同的累加器，分別位于脈動陣列的上方和左側(cè)。

由于激活是逐個輸入到脈動陣列中的，我們認(rèn)為先進(jìn)先出隊列 (FIFO) 是最佳的數(shù)據(jù)存儲方案。然而，傳統(tǒng)的 FIFO 和我們構(gòu)建的累加器之間略有不同。我們的累加器有兩個輸入端口 —— 一個用于手動將權(quán)重寫入 FIFO，另一個用于將激活模塊的上一層輸出寫回到輸入 FIFO（上一層的輸出是當(dāng)前層的輸入）。

我們還需要以類似的方式為每一層加載權(quán)重，因此我們復(fù)制了權(quán)重 FIFO 的邏輯，但沒有第二個端口。

脈動陣列矩陣乘法：

偏差和激活

等式的下一步是添加偏差。為了在硬件中實(shí)現(xiàn)這一點(diǎn)，我們需要在脈動陣列的每一列下創(chuàng)建一個偏差模塊。我們可以看到，當(dāng)總和移出脈動陣列的最后一行時，我們可以立即將它們輸入到偏差模塊中，以計算預(yù)激活。我們將用變量 Z 表示這些值。

偏差向量 b 會在矩陣的所有行上轉(zhuǎn)發(fā) —— 這意味著它會被添加到每個 Z 行。

現(xiàn)在，我們的方程看起來很像我們在高中學(xué)到的 —— 只不過是多維形式，其中從脈動陣列流出的每一列都代表其自身的特征！

接下來，我們必須應(yīng)用激活函數(shù)，為此我們選擇了 Leaky ReLU。[5] 這也是一個逐元素的操作，類似于偏差函數(shù)，這意味著我們需要在每個偏差模塊下（以及通過代理在脈動陣列的每一列下）都有一個激活模塊，并且我們可以將偏差模塊的輸出立即流式傳輸?shù)郊せ钅K中。我們將用 H 表示這些激活后的值。

Leaky ReLU 函數(shù)逐元素應(yīng)用：

其中 α=0.5 是我們的泄漏因子。對于矩陣，這適用于每個元素。

以我們的異或運(yùn)算示例為例，我們來看看第 1 層如何處理數(shù)據(jù)。首先，脈動陣列計算：

然后添加偏差：

最后，LeakyReLU 逐個元素應(yīng)用：

負(fù)值乘以 0.5，正值保持不變。

帶偏置和 LeakyReLU 的脈動陣列：

流水線

現(xiàn)在你可能會問 —— 為什么我們不把偏置項和激活項合并在一個時鐘周期內(nèi)？這是因?yàn)榱魉€技術(shù)！流水線允許多個操作在 TPU 的不同階段同時執(zhí)行 —— 無需等待一個完整的操作完成后再開始下一個操作，而是將工作分解成可以重疊的階段。

可以把它想象成一條裝配線：當(dāng)一個工作模塊（激活模塊）處理一個零件時，前一個工作模塊（偏置模塊）已經(jīng)在處理下一個零件了。這使得所有模塊都保持忙碌狀態(tài)，而不是讓它們閑置等待上一個階段完成。這也會影響 TPU 的運(yùn)行速度 —— 如果某個模塊試圖在一個周期內(nèi)執(zhí)行多個操作，那么該模塊將成為 TPU 的瓶頸，因?yàn)槠渌K的運(yùn)行速度最終只能與該模塊相同。因此，盡可能將操作分解成多個時鐘周期是高效且最佳實(shí)踐。

流水線階段顯示操作如何跨時鐘周期重疊

為了盡可能高效地運(yùn)行芯片，我們采用了另一種機(jī)制，即傳播「啟動」信號，我們稱之為「移動芯片使能」（用紫色圓點(diǎn)表示）。由于我們設(shè)計中的所有組件都是交錯排列的，我們意識到可以非常優(yōu)雅地在第一個累加器上斷言一個時鐘周期的啟動信號，并在相鄰模塊需要開啟時將其準(zhǔn)確傳播到它們。

這將延伸到脈動陣列，最終延伸到偏置和激活模塊，其中相鄰的 PE 和模塊（從左上角到右下角）在連續(xù)的時鐘周期內(nèi)開啟。這確保了每個模塊僅在需要時執(zhí)行計算，并且不會在后臺浪費(fèi)電量。

雙倍緩沖

現(xiàn)在，我們知道開始一個新的層意味著我們必須使用新的權(quán)重矩陣計算相同的。如果我們的脈動陣列是權(quán)重平穩(wěn)的，我們該如何做到這一點(diǎn)？我們該如何改變權(quán)重？

在思考這個問題時，我們偶然發(fā)現(xiàn)了雙倍緩沖的概念，它源自電子游戲。雙緩沖存在的原因是為了防止顯示器上出現(xiàn)所謂「畫面撕裂」的情況。歸根結(jié)底，像素加載需要時間，我們希望以某種方式「隱藏」這段時間。如果你仔細(xì)觀察，就會發(fā)現(xiàn)這與我們目前在脈動陣列中面臨的問題完全相同。幸運(yùn)的是，游戲設(shè)計師已經(jīng)想出了解決這個問題的方案。通過添加第二個「影子」緩沖區(qū)（在當(dāng)前層計算時保存下一層的權(quán)重），我們可以在計算過程中加載新的權(quán)重，從而將總時鐘周期數(shù)減少一半。

為了實(shí)現(xiàn)這一點(diǎn)，我們還需要添加一些信號來移動數(shù)據(jù)。首先，我們需要一個信號來指示何時切換影子緩沖區(qū)和活動緩沖區(qū)中的權(quán)重。我們將這個信號稱為「切換」信號（用藍(lán)點(diǎn)表示），它將影子緩沖區(qū)中的值復(fù)制到活動緩沖區(qū)。它從脈動陣列的左上角傳播到右下角（與移動芯片使能的路徑相同，但僅在脈動陣列內(nèi)傳播）。然后，我們需要另一個信號來指示何時需要將權(quán)重向下移動一行，我們將其稱為「接受」標(biāo)志（用綠點(diǎn)表示），因?yàn)槊恳恍卸急硎窘邮芤唤M新的權(quán)重。這會將新的權(quán)重移動到脈動陣列的頂行，并將每一行權(quán)重向下移動到脈動陣列的下一行。這兩個控制標(biāo)志協(xié)同工作，使我們的雙緩沖機(jī)制正常工作。

如果你還沒注意到，這讓脈動陣列能夠執(zhí)行一項強(qiáng)大的功能…… 持續(xù)推理！我們可以持續(xù)輸入新的權(quán)重和輸入，并計算任意層級的前向傳播。這觸及了脈動陣列的核心設(shè)計理念：我們希望最大化 PE 的利用率。我們希望始終保持脈動陣列的饋送！

對于第 2 層，第 1 層的輸出 (H1）現(xiàn)在成為我們的輸入：

添加偏差并應(yīng)用激活：

所有值均為正數(shù)，因此它們會保持不變。這些就是我們對異或問題的最終預(yù)測！

正向傳遞演練（使用雙倍緩沖）

控制單元和指令集（ISA）

我們推理的最后一步是創(chuàng)建一個控制單元，使其使用自定義指令集 (ISA) 來斷言所有控制標(biāo)志并通過數(shù)據(jù)總線加載數(shù)據(jù)。包括數(shù)據(jù)總線在內(nèi)，我們的 ISA 長度為 24 位，這使得我們的測試平臺更加優(yōu)雅，因?yàn)槲覀兛梢悦總€時鐘周期傳遞一串位，而無需單獨(dú)設(shè)置多個標(biāo)志。

然后，我們將所有東西整合在一起，推理功能完全正常運(yùn)行！這是一個重要的里程碑。

反向傳播與訓(xùn)練

好了，我們已經(jīng)解決了推理問題 —— 那么訓(xùn)練呢？有趣的是，我們可以將用于推理的架構(gòu)用于訓(xùn)練！為什么？因?yàn)橛?xùn)練只是矩陣乘法，只是多了一些步驟。

接下來才是真正令人興奮的地方。假設(shè)我們剛剛對異或問題進(jìn)行了推理，得到的預(yù)測結(jié)果類似于 [0.8, 0.3, 0.1, 0.9]，而我們實(shí)際想要的是 [1, 0, 0, 1]。我們的模型表現(xiàn)很差，需要改進(jìn)它。這時訓(xùn)練就派上用場了。我們將使用一個叫做損失函數(shù)的東西來準(zhǔn)確地告訴模型它的表現(xiàn)有多差。為了簡單起見，我們選擇了均方誤差 (MSE)—— 可以把它想象成測量預(yù)測結(jié)果與實(shí)際期望結(jié)果之間的「距離」，就像測量籃球投籃偏離目標(biāo)的距離一樣。我們用 L 表示損失。

因此，在完成最后一層的激活函數(shù)（我們稱之為 H2）的計算后，立即將它們輸入到損失模塊中，以計算我們的預(yù)測有多糟糕。這些損失模塊位于激活模塊的正下方，并且只有在到達(dá)最后一層時才會使用它們。但關(guān)鍵在于：訓(xùn)練時實(shí)際上并不需要計算損失值本身。你只需要它的導(dǎo)數(shù)。因?yàn)檫@個導(dǎo)數(shù)告訴我們應(yīng)該向哪個方向調(diào)整權(quán)重以減小損失。這就像擁有一個指向「更好性能」的指南針。

鏈?zhǔn)椒▌t的魔力

這就是微積分發(fā)揮作用的地方。為了改進(jìn)我們的模型，我們需要弄清楚每個權(quán)重的變化如何影響損失。鏈?zhǔn)椒▌t讓我們把這個龐大的計算分解成更小、更易于管理的部分。

梯度的鏈?zhǔn)椒▌t：

這使我們能夠逐層計算梯度，并通過網(wǎng)絡(luò)向后傳播它們。

讓我們一步步回顧一下發(fā)生的事情。

1、計算損失相對于我們的最終激活有多少變化。

2、計算通過對激活函數(shù)（在我們的例子中是 LeakyReLU）求導(dǎo)。

3、計算

由于 z2 的所有元素均為正，因此 LeakyReLU 梯度為：

前向傳播和后向傳播的美妙對稱性

繪制出完整的計算圖后，我們發(fā)現(xiàn)了一個驚人的現(xiàn)象：反向傳播中的最長鏈與前向傳播非常相似！在前向傳播中，我們將激活矩陣與轉(zhuǎn)置后的權(quán)重矩陣相乘。在后向傳播中，我們將梯度矩陣與未轉(zhuǎn)置的權(quán)重矩陣相乘。這就像照鏡子一樣！

將梯度傳播到隱藏層：

并通過第一層的激活：

當(dāng) Z1 中正負(fù)值混合時，梯度為：

一旦我們有了所有這些單獨(dú)的導(dǎo)數(shù)，我們就可以將它們相乘，以找到關(guān)于損失的任何導(dǎo)數(shù)。

之后，我們必須計算激活導(dǎo)數(shù)。

公式為：

這也是一種逐元素計算，這意味著我們可以像損失模塊（以及偏差和激活模塊）一樣構(gòu)建它，但它會執(zhí)行不同的計算。然而，關(guān)于這個模塊，需要注意的是，它需要我們在前向傳播過程中計算的激活值。

現(xiàn)在你可能想知道 —— 我們實(shí)際上是如何在硬件中計算導(dǎo)數(shù)的？讓我們以 Leaky ReLU 為例，它非常簡單，但卻演示了關(guān)鍵原理。請記住，Leaky ReLU 會根據(jù)輸入是正還是負(fù)應(yīng)用不同的操作。導(dǎo)數(shù)遵循相同的模式：對于正輸入，它輸出 1，對于負(fù)輸入，它輸出一個小的常數(shù)（我們使用了 0.01）。

Leaky ReLU 梯度：

硬件中的 Leaky ReLU 導(dǎo)數(shù)實(shí)現(xiàn)展示了條件邏輯。

它的好處在于它只是一個簡單的比較 —— 無需復(fù)雜的算法。硬件可以在一個時鐘周期內(nèi)計算出這個導(dǎo)數(shù)，從而保持我們的流水線順暢運(yùn)行。同樣的原則也適用于其他激活函數(shù)：它們的導(dǎo)數(shù)通常簡化為硬件可以高效執(zhí)行的基本操作。這一洞見促使我們首先計算長鏈 —— 獲取所有梯度，就像我們在前向傳播中計算激活一樣。我們可以緩存這些梯度并重復(fù)使用，遵循我們已經(jīng)掌握的高效模式。

你會注意到一個很棒的模式正在浮現(xiàn)：所有這些位于脈動陣列下方的模塊都會處理逐個輸出的列向量。這促使我們想到將它們統(tǒng)一起來，形成一個我們稱之為向量處理單元 (VPU) 的東西 —— 因?yàn)樗鼈冋沁@樣做的，逐元素地處理向量！[6]

這不僅使用起來更加優(yōu)雅，而且當(dāng)我們將 TPU 擴(kuò)展到 2x2 脈動陣列以上時也非常有用，因?yàn)槲覀儗碛?N 個這樣的模塊（N 是脈動陣列的大小），每個模塊都必須單獨(dú)與它們交互。將這些模塊統(tǒng)一到一個父模塊下，使我們的設(shè)計更具可擴(kuò)展性和優(yōu)雅性！

矢量處理單元 (VPU) 架構(gòu)展示統(tǒng)一的元素級操作

此外，通過為每個模塊整合控制信號（稱之為 VPU 通路位），我們可以選擇性地啟用或跳過特定操作。這使得 VPU 足夠靈活，能夠同時支持推理和訓(xùn)練。例如，在正向傳遞過程中，我們希望應(yīng)用偏差和激活函數(shù)，但跳過計算損失或激活函數(shù)導(dǎo)數(shù)。當(dāng)轉(zhuǎn)換到反向傳遞時，所有模塊都會參與，但在反向鏈中，我們只需要計算激活函數(shù)導(dǎo)數(shù)。由于流水線技術(shù)，所有流經(jīng) VPU 的值都會經(jīng)過四個模塊，任何未使用的模塊都只是充當(dāng)寄存器，將其輸入轉(zhuǎn)發(fā)到輸出，而不執(zhí)行計算。

接下來的幾個導(dǎo)數(shù)很有趣，因?yàn)槲覀儗?shí)際上可以使用矩陣乘法（和脈動陣列）借助以下三個恒等式來計算導(dǎo)數(shù)：

1、如果我們有并對權(quán)重求導(dǎo)，我們得到：

2、如果我們有并對輸入 X 取導(dǎo)數(shù)，我們得到：

3、對于偏差項，導(dǎo)數(shù)就是 1。

這意味著我們可以將前面的與 X、W^T 和 1 相乘，得到和，然后我們可以將所有這些乘以，得到損失函數(shù)關(guān)于所有第二層參數(shù)的梯度。由于所有梯度實(shí)際上都是梯度矩陣，我們可以使用收縮陣列。

現(xiàn)在需要注意的是，激活函數(shù)導(dǎo)數(shù)以及權(quán)重導(dǎo)數(shù)都需要我們在前向傳播過程中計算后激活函數(shù) (H)。這意味著我們需要將每一層的輸出存儲在某種內(nèi)存中才能進(jìn)行訓(xùn)練。為此，我們創(chuàng)建了一個新的暫存器內(nèi)存模塊，我們稱之為統(tǒng)一緩沖區(qū) (UB)。這樣，就可以在前向傳播過程中計算 H 值后立即存儲它們。

我們意識到，通過使用 UB 來存儲輸入和權(quán)重累加器，以及手動將偏差和泄漏因子加載到各自的模塊中，我們也可以省去這些累加器。相比每個時鐘周期使用指令集加載新數(shù)據(jù)，這也是一種更好的做法。由于我們想要同時訪問兩個值（脈動陣列每行 / 列的 2 個輸入或 2 個權(quán)重），我們添加了兩個讀寫端口。我們對每個數(shù)據(jù)原語（輸入、權(quán)重、偏差、泄漏因子、后激活）都進(jìn)行了這樣的操作，以最大限度地減少數(shù)據(jù)爭用，因?yàn)槲覀儞碛性S多不同類型的數(shù)據(jù)。

要讀取值，我們需要提供起始地址和值的數(shù)量，以及我們希望 UB 讀取的起始地址和位置數(shù)量，這樣它每個時鐘周期就會讀取 2 個值。寫入是一種類似的機(jī)制，我們需要指定要寫入兩個輸入端口的值。讀取機(jī)制的優(yōu)點(diǎn)在于，一旦我們提供起始地址，它就會在后臺運(yùn)行，直到讀取給定位置的數(shù)量，這意味著我們只需要每隔幾個時鐘周期提供一條指令。

統(tǒng)一緩沖區(qū)（UB）架構(gòu)展示雙端口讀取機(jī)制

顯示讀取操作的統(tǒng)一緩沖區(qū)時序波形

最終，取消這些機(jī)制不會破壞 TPU 的性能 —— 但它們使我們能夠始終保持脈動陣列的供電，這是我們不容妥協(xié)的核心設(shè)計原則。

在進(jìn)行這項工作的過程中，我們意識到可以對激活導(dǎo)數(shù)模塊進(jìn)行最后一個小小的優(yōu)化 —— 由于我們只使用一次 H2 值（用于計算），我們在 VPU 內(nèi)部創(chuàng)建了一個小型緩存，而不是將它們存儲在 UB 中。其余的 H 值將存儲在 UB 中，因?yàn)樗鼈冃枰糜谟嬎愣鄠€導(dǎo)數(shù)。

用于存儲臨時激活值的 H-cache 優(yōu)化

經(jīng)過修改以執(zhí)行訓(xùn)練的新 TPU 架構(gòu)如下所示：

完整的 TPU 架構(gòu)，展示推理和訓(xùn)練的所有組件

現(xiàn)在我們可以進(jìn)行反向傳播了！

回到計算圖，這里有一個驚人的發(fā)現(xiàn)：反向傳播中的最長鏈與正向傳播非常相似！在正向傳播中，我們將激活矩陣與轉(zhuǎn)置的權(quán)重矩陣相乘。在反向傳播中，我們將梯度矩陣與未轉(zhuǎn)置的權(quán)重矩陣相乘。這就像照鏡子一樣！

顯示矩陣運(yùn)算的前向傳遞計算流程。

這一洞見引導(dǎo)我們首先計算計算圖的長鏈（以黃色突出顯示）—— 獲取所有梯度，就像我們在前向傳播中計算激活值一樣。我們可以緩存這些梯度并重復(fù)使用，遵循已經(jīng)掌握的高效模式。

我們創(chuàng)建一個循環(huán)，其中：

1、從統(tǒng)一緩沖區(qū)獲取橋節(jié)點(diǎn)

2、同樣從統(tǒng)一緩沖區(qū)中獲取相應(yīng)的 H2 矩陣

3、通過我們的脈動陣列來計算權(quán)重梯度

反向傳播通過第二個隱藏層

這時，真正神奇的事情發(fā)生了：我們可以在計算這些權(quán)重梯度的同時，將它們直接輸入到梯度下降模塊中！該模塊獲取存儲在內(nèi)存中的當(dāng)前權(quán)重，并使用梯度更新它們。

梯度下降更新規(guī)則：

其中 α 為學(xué)習(xí)率，θ 表示任意參數(shù)（權(quán)重或偏差）。

計算異或網(wǎng)絡(luò)的權(quán)重梯度：

無需等待 —— 一切都像水一樣流經(jīng)我們的管道。

你可能會想：“我們已經(jīng)用矩陣乘法恒等式計算了長鏈梯度和權(quán)重梯度 —— 那么如何計算偏差梯度呢？” 其實(shí)，我們已經(jīng)完成了大部分工作！由于我們處理的是批量數(shù)據(jù)，所以我們可以簡單地對批量維度上的梯度求和（專業(yè)術(shù)語是「約簡」reduce）。好處就在，我們可以在計算長鏈時直接進(jìn)行約簡 —— 無需額外工作！