三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

【三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的】三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片

三個月、零基礎手搓一塊TPU,能推理能訓練,還是開源的

文章圖片


機器之心報道
編輯:澤南

最近 , 大模型技術的發展 , 讓人們再次重視起 AI 專用芯片 。

對于計算任務負載來說 , 越是專用 , 效率就越高 , 谷歌的 TPU 就是其中的一個典型例子 。 它自 2015 年開始在谷歌數據中心部署后 , 已經發展到了第 7 代 。 目前的最新產品不僅使用了最先進的制程工藝打造 , 也在架構上充分考慮了對于機器學習推理任務的優化 。 TPU 的出現 , 促進了 Gemini 等大模型技術的進展 。

這種芯片為何性能如此強大?或許最直接的了解方法就是嘗試復刻它 。 近日 , 來自加拿大西安大略大學的工程師 Surya Sure 等人宣布已經利用暑假時間構建出了 TinyTPU:一種開源的 ML 推理、訓練芯片 。



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

有趣的是 , 他們并非芯片設計專業的學生 , 打造 TPU 是從理解多層感知機(MLP)這樣的神經網絡基本概念開始的 。 為此 , 他們還手工計算了網絡推理和訓練所需的數學運算 。



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

手搓 TPU 的原因

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

構建一個用于機器學習工作負載的芯片看起來很酷; 之前還沒有一個同時進行推理和訓練的機器學習加速器的完整開源代碼庫 。
我們都沒有真正的硬件設計專業經驗 , 這在某種程度上使得 TPU 更具吸引力 , 因為我們無法準確估計它的難度 。 在項目初期 , 我們確立了嚴格的設計理念:始終嘗試「Hacky Way」(不靠譜的方法) 。 這意味著在咨詢外部資源之前 , 先嘗試那些我們首先想到的「愚蠢」想法 。 這種理念確保了我們沒對 TPU 進行逆向工程 , 而是對其進行重新發明 , 因此我們自己推導出了 TPU 中使用的許多關鍵機制 。

我們也想把這個項目當作一次不依賴人工智能代寫代碼的練習 , 因為我們感覺最近每當遇到小問題時 , 我們的第一反應就是求助于這些 AI 工具 。 我們希望培養一種特定的思維方式 , 以便我們能夠將其付諸實踐 , 并在未來的任何努力中運用它來解決難題 。

在整個項目過程中 , 我們盡可能多地學習深度學習、硬件設計和算法創建方面的基礎知識 。 我們發現 , 學習這些知識的最佳方式是將所有知識都畫出來 , 并將其作為我們的第一反應 。 閱讀這篇文章時 , 你會發現我們的解釋是如何受到這種理念的啟發的 。

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

什么是 TPU?

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

硬件設計快速入門:

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


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

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



在上面的例子中 , 信號 b 在下一個時鐘周期的值被設置為信號 a 的當前值 。 你會發現 , 在大多數情況下 , 信號(變量)是按順序的時鐘周期更新的 , 而不是像軟件設計中那樣立即更新 。

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

那么如何構建一塊 TPU 呢?

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


用于解決 XOR 問題的 2→2→1 多層感知器的架構 。

我們選擇這個特定網絡的原因是 , 我們的目標是針對 XOR 問題(神經網絡的「Hello World」)進行推理和訓練 。 XOR 問題是神經網絡可以解決的最簡單的問題之一 。 所有其他門(AND、OR 等)都可以僅使用一條線性線(一個神經元)來區分哪些輸入對應于 0 , 哪些輸入對應于 1 , 從而根據輸入預測輸出 。 但是 , 要對所有 XOR 進行分類 , 需要一個 MLP , 因為它需要彎曲的決策邊界 , 而這僅靠線性方程無法實現 。

對于幾何和第一性原理的論述 , 免費書籍《理解深度學習》(Understanding Deep Learning)非常值得一讀:https://udlbook.github.io/udlbook/



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

XOR 輸入矩陣和目標輸出:


每一行代表四個可能的異或輸入之一 , 輸出向量顯示預期的異或結果 。

我們對脈動陣列示例進行的另一項簡化是 , 我們將使用 2x2 陣列 , 而不是 TPUv1 中使用的 256x256 陣列 。 然而 , 數學運算仍然忠實于原有的格式 , 因此實際上并沒有簡化 , 而是進行了規模縮小 。

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



其中 X 是我們的輸入矩陣 , W 是我們的權重矩陣 , b 是我們的偏差向量 。

我們如何在硬件中執行矩陣乘法?嗯 , 我們可以使用一個叫做脈動陣列的單元!

TPU 的核心是一個叫做脈動陣列(systolic array)的單元 。 它由稱為處理單元 (PE) 的獨立構建塊組成 , 這些構建塊以網格狀結構連接在一起 。 每個 PE 執行乘法 - 累加運算 , 這意味著它將傳入的輸入 X 與固定權重 W 相乘 , 并將其與傳入的累加和相加 , 所有這些都在同一個時鐘周期內完成 。


處理元件(PE)架構顯示乘法累加運算(無負載權重和啟動標志) 。



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


脈動陣列架構展示了 PE 如何連接以執行矩陣乘法 。

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

示例

現在 , 讓我們來看看異或問題的示例 。

我們的脈動陣列接受兩個輸入:輸入矩陣和權重矩陣 。 對于我們的異或網絡 , 使用以下權重和偏差進行初始化:



輸入和權重調度

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

將 X 矩陣旋轉 90 度


矩陣旋轉 90 度 , 為脈動陣列輸入做準備 。

錯開輸入(將每行延遲 1 個時鐘周期)


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

要輸入我們的權重矩陣 , 我們需要:

交錯排列權重矩陣(與輸入類似)


用于脈動陣列處理的權重矩陣交錯模式

轉置它!


權重矩陣轉置以實現正確的數學對齊 。

請注意 , 旋轉和交錯沒有任何數學意義 —— 它們只是為了使脈動陣列正常工作 。 轉置也只是為了進行數學記賬 —— 它是使矩陣數學運算正常工作所必需的 , 因為我們在神經網絡圖中設置權重指針的方式 。

交錯和先進先出 (FIFO)

為了執行交錯 , 我們為權重和輸入設計了幾乎相同的累加器 , 分別位于脈動陣列的上方和左側 。

由于激活是逐個輸入到脈動陣列中的 , 我們認為先進先出隊列 (FIFO) 是最佳的數據存儲方案 。 然而 , 傳統的 FIFO 和我們構建的累加器之間略有不同 。 我們的累加器有兩個輸入端口 —— 一個用于手動將權重寫入 FIFO , 另一個用于將激活模塊的上一層輸出寫回到輸入 FIFO(上一層的輸出是當前層的輸入) 。

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

脈動陣列矩陣乘法:



偏差和激活

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


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

現在 , 我們的方程看起來很像我們在高中學到的 —— 只不過是多維形式 , 其中從脈動陣列流出的每一列都代表其自身的特征!

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

Leaky ReLU 函數逐元素應用:



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

以我們的異或運算示例為例 , 我們來看看第 1 層如何處理數據 。 首先 , 脈動陣列計算





然后添加偏差:



最后 , LeakyReLU 逐個元素應用:



負值乘以 0.5 , 正值保持不變 。

帶偏置和 LeakyReLU 的脈動陣列:



流水線

現在你可能會問 —— 為什么我們不把偏置項和激活項合并在一個時鐘周期內?這是因為流水線技術!流水線允許多個操作在 TPU 的不同階段同時執行 —— 無需等待一個完整的操作完成后再開始下一個操作 , 而是將工作分解成可以重疊的階段 。

可以把它想象成一條裝配線:當一個工作模塊(激活模塊)處理一個零件時 , 前一個工作模塊(偏置模塊)已經在處理下一個零件了 。 這使得所有模塊都保持忙碌狀態 , 而不是讓它們閑置等待上一個階段完成 。 這也會影響 TPU 的運行速度 —— 如果某個模塊試圖在一個周期內執行多個操作 , 那么該模塊將成為 TPU 的瓶頸 , 因為其他模塊的運行速度最終只能與該模塊相同 。 因此 , 盡可能將操作分解成多個時鐘周期是高效且最佳實踐 。


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

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

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

雙倍緩沖

現在 , 我們知道開始一個新的層意味著我們必須使用新的權重矩陣計算相同的

。 如果我們的脈動陣列是權重平穩的 , 我們該如何做到這一點?我們該如何改變權重?

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

為了實現這一點 , 我們還需要添加一些信號來移動數據 。 首先 , 我們需要一個信號來指示何時切換影子緩沖區和活動緩沖區中的權重 。 我們將這個信號稱為「切換」信號(用藍點表示) , 它將影子緩沖區中的值復制到活動緩沖區 。 它從脈動陣列的左上角傳播到右下角(與移動芯片使能的路徑相同 , 但僅在脈動陣列內傳播) 。 然后 , 我們需要另一個信號來指示何時需要將權重向下移動一行 , 我們將其稱為「接受」標志(用綠點表示) , 因為每一行都表示接受一組新的權重 。 這會將新的權重移動到脈動陣列的頂行 , 并將每一行權重向下移動到脈動陣列的下一行 。 這兩個控制標志協同工作 , 使我們的雙緩沖機制正常工作 。



如果你還沒注意到 , 這讓脈動陣列能夠執行一項強大的功能…… 持續推理!我們可以持續輸入新的權重和輸入 , 并計算任意層級的前向傳播 。 這觸及了脈動陣列的核心設計理念:我們希望最大化 PE 的利用率 。 我們希望始終保持脈動陣列的饋送!

對于第 2 層 , 第 1 層的輸出 (H1)現在成為我們的輸入:



添加偏差并應用激活:



所有值均為正數 , 因此它們會保持不變 。 這些就是我們對異或問題的最終預測!

正向傳遞演練(使用雙倍緩沖)



控制單元和指令集(ISA)

我們推理的最后一步是創建一個控制單元 , 使其使用自定義指令集 (ISA) 來斷言所有控制標志并通過數據總線加載數據 。 包括數據總線在內 , 我們的 ISA 長度為 24 位 , 這使得我們的測試平臺更加優雅 , 因為我們可以每個時鐘周期傳遞一串位 , 而無需單獨設置多個標志 。

然后 , 我們將所有東西整合在一起 , 推理功能完全正常運行!這是一個重要的里程碑 。



反向傳播與訓練

好了 , 我們已經解決了推理問題 —— 那么訓練呢?有趣的是 , 我們可以將用于推理的架構用于訓練!為什么?因為訓練只是矩陣乘法 , 只是多了一些步驟 。

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



因此 , 在完成最后一層的激活函數(我們稱之為 H2)的計算后 , 立即將它們輸入到損失模塊中 , 以計算我們的預測有多糟糕 。 這些損失模塊位于激活模塊的正下方 , 并且只有在到達最后一層時才會使用它們 。 但關鍵在于:訓練時實際上并不需要計算損失值本身 。 你只需要它的導數 。 因為這個導數告訴我們應該向哪個方向調整權重以減小損失 。 這就像擁有一個指向「更好性能」的指南針 。

鏈式法則的魔力

這就是微積分發揮作用的地方 。 為了改進我們的模型 , 我們需要弄清楚每個權重的變化如何影響損失 。 鏈式法則讓我們把這個龐大的計算分解成更小、更易于管理的部分 。

梯度的鏈式法則:



這使我們能夠逐層計算梯度 , 并通過網絡向后傳播它們 。




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

繪制出完整的計算圖后 , 我們發現了一個驚人的現象:反向傳播中的最長鏈與前向傳播非常相似!在前向傳播中 , 我們將激活矩陣與轉置后的權重矩陣相乘 。 在后向傳播中 , 我們將梯度矩陣與未轉置的權重矩陣相乘 。 這就像照鏡子一樣!



將梯度傳播到隱藏層:



并通過第一層的激活:



當 Z1 中正負值混合時 , 梯度為:



一旦我們有了所有這些單獨的導數 , 我們就可以將它們相乘 , 以找到關于損失的任何導數 。





硬件中的 Leaky ReLU 導數實現展示了條件邏輯 。

它的好處在于它只是一個簡單的比較 —— 無需復雜的算法 。 硬件可以在一個時鐘周期內計算出這個導數 , 從而保持我們的流水線順暢運行 。 同樣的原則也適用于其他激活函數:它們的導數通常簡化為硬件可以高效執行的基本操作 。 這一洞見促使我們首先計算長鏈 —— 獲取所有

梯度 , 就像我們在前向傳播中計算激活一樣 。 我們可以緩存這些梯度并重復使用 , 遵循我們已經掌握的高效模式 。

你會注意到一個很棒的模式正在浮現:所有這些位于脈動陣列下方的模塊都會處理逐個輸出的列向量 。 這促使我們想到將它們統一起來 , 形成一個我們稱之為向量處理單元 (VPU) 的東西 —— 因為它們正是這樣做的 , 逐元素地處理向量![6


這不僅使用起來更加優雅 , 而且當我們將 TPU 擴展到 2x2 脈動陣列以上時也非常有用 , 因為我們將擁有 N 個這樣的模塊(N 是脈動陣列的大?。 ?, 每個模塊都必須單獨與它們交互 。 將這些模塊統一到一個父模塊下 , 使我們的設計更具可擴展性和優雅性!


矢量處理單元 (VPU) 架構展示統一的元素級操作 。

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

接下來的幾個導數很有趣 , 因為我們實際上可以使用矩陣乘法(和脈動陣列)借助以下三個恒等式來計算導數:



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

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


統一緩沖區(UB)架構展示雙端口讀取機制 。


顯示讀取操作的統一緩沖區時序波形 。

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


經過修改以執行訓練的新 TPU 架構如下所示:


完整的 TPU 架構 , 展示推理和訓練的所有組件 。

現在我們可以進行反向傳播了!

回到計算圖 , 這里有一個驚人的發現:反向傳播中的最長鏈與正向傳播非常相似!在正向傳播中 , 我們將激活矩陣與轉置的權重矩陣相乘 。 在反向傳播中 , 我們將梯度矩陣與未轉置的權重矩陣相乘 。 這就像照鏡子一樣!


顯示矩陣運算的前向傳遞計算流程 。


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



這時 , 真正神奇的事情發生了:我們可以在計算這些權重梯度的同時 , 將它們直接輸入到梯度下降模塊中!該模塊獲取存儲在內存中的當前權重 , 并使用梯度更新它們 。

梯度下降更新規則:



其中 α 為學習率 , θ 表示任意參數(權重或偏差) 。

計算異或網絡的權重梯度:



無需等待 —— 一切都像水一樣流經我們的管道 。


整合起來

通過不斷重復同樣的過程 —— 前向傳播、后向傳播、權重更新 —— 我們可以訓練網絡 , 直到它的性能達到我們預期 。 之前用于推理的脈動陣列現在也用于訓練 , 只需添加幾個模塊來處理梯度計算 。

最初只是一個簡單的矩陣乘法想法 , 如今已發展成為一個完整的訓練系統 。 每個組件都和諧地協同工作:數據流經管道 , 模塊并行運行 , 而脈動陣列則持續提供有用的工作 。


GTKWave 中的最終波形模擬顯示一個時期后內存中的權重和偏差更新 。

參考內容:
https://www.tinytpu.com/
https://x.com/suryasure05/status/1957518913648095376

    推薦閱讀