當前位置: 妍妍網 > 資訊

萬字長文,Google TPU往事

2024-03-29資訊

【編者按】 盡管 Google 的 AGI 之戰稍顯落後,不敵微軟 X OpenAI、Amazon X Claude 以及 Meta 的 LLaMA,但許多人仍對 Google 懷以希望,畢竟 Google 的 TPU 這張王牌還沒有大顯身手。本文深入回顧了 TPU 的起源以及架構的細節,還提到了 AlexNet 作者 Alex Krizhevsky 是如何讓模型訓練走上了 GPU 之路。故事與技術相交織,可以一窺 AI 的演進。

作者 | Babbage

責編 | 沭七

出品丨AI 科技大本營(ID:rgznai100)

計算歷史中被遺忘或過時的想法仍具有研究價值,並且有時極具實用性。使用神經網路的例子,經歷了幾次「AI 寒冬」與熱潮交替的歷程,幾乎七十年來時冷時熱,無疑是其中最具影響力的案例。

若說對神經網路的興趣歷經起伏,那麽為其加速的專用硬體的發展同樣經歷了多次浮沈。然而,這類硬體面臨著一個典型的「雞生蛋蛋生雞」的問題。CPU 並非執行神經網路的最高效方式。但是,沒有神經網路的實用證明,如何為專用硬體的持續投入提供理由呢?同時,神經網路的有效運作又需要更強大的硬體。隨著炒作周期的變化,一次又一次地,建造專用硬體的計畫被擱置。最終,原本為其他目的而構建的硬體——GPU,打破了這一僵局。

DALL·E 生成

隨後,專用硬體的出現似乎是必然的。我們現在正處於一個「炒作期」,這一炒作導致許多公司,從初創公司到超大型企業,紛紛打造自己的加速器。盡管大多數設計最終可能會失敗,專用神經網路加速器肯定會在未來的計算領域占據一席之地。谷歌的張量處理單元( TPU ),作為首批大規模部署的加速器,無疑將成為其中的幸存者。

為何谷歌的 TPU 值得關註並具有重要性?它們並非可以從谷歌或任何其他方購買的產品。你可以在谷歌雲平台上「租用」它們,但繼續使用 TPU 依賴於谷歌繼續提供支持的承諾。這似乎是一個不穩定的基礎。

但這畢竟是谷歌,僅僅因為谷歌在其服務中使用 TPU,就足以證明它們的重要性。有時人們容易忘記谷歌服務的規模之大。該公司有九項服務( 包括搜尋、Android、Chrome、Gmail、地圖、Play 商店、YouTube、照片 )擁有超過十億使用者。此外,谷歌還擁有世界第三大的「公共雲」平台。

在一個由 Nvidia 在數據中心機器學習硬體領域所主導的時代,谷歌的 TPU 提供了一個已被大規模部署的、證明有效的 Nvidia GPU 的替代方案的例項。谷歌 TPU 計畫背後的核心理念——建立專門用於加速機器學習任務的數據中心硬體,已經促使許多其他公司,從初創企業( 如 Cerebras 和 Tenstorrent )到大型企業( 如 AWS 和 Microsoft ),構建自己的硬體。其中一個風頭無兩的初創公司 Groq,由原 TPU 團隊成員 Jonathan Ross 領導。

第一款 TPU 的開發歷史同樣引人入勝。它展示了即使是像谷歌這樣的公司,也能在迫切需要時迅速且靈活地行動。一個擁有必要資源的公司可以迅速開發出創新的硬體。它還展示了數十年前的想法如何重新浮現,成為現代重要設計的核心。

首先,我們澄清一下,當我們討論「張量處理單元」時,我們討論的究竟是什麽。

張量硬體與軟體

近年來,「張量」一詞開始頻繁出現在一些電腦硬體和軟體的命名中。以下是一個簡短的列表:

1. TensorFlow:谷歌的機器學習框架;

2. 谷歌張量處理單元(Tensor Processing Unit,TPU):谷歌的自訂數據中心加速器;

3. Tensor Core:Nvidia 最新 GPU 中的一個執行單元;

4. 谷歌 Tensor:最新的 Pixel 智慧型手機中的基於 Arm 的系統級芯片。

本文將重點討論第二項,即谷歌的自訂數據中心加速器,涵蓋第一代谷歌 TPU 的歷史和架構。如有必要,我將使用 TPU v1 來區分其後續版本。

這些「張量」的突然出現,當然是因為機器學習,特別是深度學習的重要性迅速增加。張量是深度學習公式化和實作的核心。

那麽,什麽是張量?已經熟悉張量的讀者可以跳過這一節。

維基百科給出了一個不那麽直觀的定義:

在數學中,張量是一種代數物件,它描述了與向量空間相關的一組代數物件之間的多線性關系。張量可以對映不同物件之間的關系,如向量、純量,甚至其他張量。

不過,不要因此而關閉頁面!實際上,我們可以簡化這一定義,幾乎完全依賴於張量的一種表現形式:

張量可以表示為一個(可能是多維的)陣列。

因此,根據張量的性質,它可以被表示為一個 n 維陣列,其中 n 是 0、1、2、3 等等。這些表示形式中的一些有更熟悉的名字:

  • 維度 0 - 純量

  • 維度 1 - 向量

  • 維度 2 - 矩陣

  • 為什麽稱之為張量處理單元( TPU )?因為它設計來加速涉及張量的操作。具體來說,是哪些操作呢?我們原始的維基百科定義中提到的操作,將張量描述為「在不同物件之間的對映( 多線性關系 ),如向量、純量,甚至其他張量」。

    讓我們舉個簡單的例子。一個二維陣列可以描述兩個一維陣列之間的多線性關系。數學傾向的人會認識到,這個過程是透過將一個向量乘以一個矩陣來得到另一個向量。

    這可以推廣到表示更高維陣列之間關系的張量。然而,盡管張量描述了任意高維陣列之間的關系,實際上我們將考慮的 TPU 硬體是設計來執行與一維和二維陣列相關的計算。或者,更具體地說,向量和矩陣操作。

    谷歌 TPU 的起源故事

    要了解 TPU 的起源,我們需要回顧谷歌與機器學習和深度學習長期且不斷發展的歷史。谷歌對機器學習的興趣是顯而易見的。谷歌的使命是「組織世界上的資訊,使之普遍可存取且有用。」機器學習提供了一套強大的工具,幫助谷歌組織那些資訊並使其更有用。從影像和語音辨識到語言轉譯,再到「大型語言模型」,當然,還有在谷歌每年賺取數十億美元的廣告業務中的眾多套用。對谷歌來說,機器學習至關重要。

    因此,谷歌對機器學習的興趣始於 2000 年代初。隨著時間的推移,公司的關註點開始轉向深度學習。一些關鍵的裏程碑展示了谷歌對深度學習的快速增長的關註和承諾,包括:

  • 2011 年:由 Jeff Dean、Greg Corrado 和 Andrew Ng 發起的深度學習研究計畫 Google Brain 作為 Google X 的一部份啟動。

  • 2013 年:在 2012 年底成功的 AlexNet 影像辨識計畫之後,經過與 Microsoft、百度和 DeepMind 的競標過程,谷歌以 4400 萬美元的價格收購了由 Geoffrey Hinton、Alex Krizhevsky 和 Ilya Sutskever 組成的團隊。

  • 2014 年:谷歌以高達 6.5 億美元的價格收購了由 Demis Hassabis、Shane Legg 和 Mustafa Suleyman 創立的 DeepMind。

  • 2012 年是一個關鍵時刻,當時 Google Brain 團隊能夠利用深度學習在影像辨識效能上實作突破性進展。

    使用這個大規模神經網路,我們還顯著提高了標準影像分類測試的狀態——事實上,我們看到了 70% 的相對準確性改進。我們透過利用網路上大量的未標記數據來增強有限的標記數據集來實作這一點。這是我們非常關註的一個方向——如何開發能夠很好地擴充套件的機器學習系統,以便我們可以利用大量的未標記訓練數據。

    早在 2006 年,就有關於開發機器學習客製硬體的討論。當時的結論是,設想的套用所需 的計算能力可以使用谷歌龐大數據中心中已有的伺服器容量以低成本提供。記錄這一點的論文「利用大規模無監督學習構建高級特征」展示了他們使用「一個擁有 1000 台機器(16,000 核心)的集群訓練了三天」來訓練論文中描述的模型。那是 1000 個 CPU,每個 CPU 有 16 個核心。

    Cade Metz 的書【天才制造者】回顧了當 Alex Krizhevsky( AlexNet 作者 )在 2013 年來到谷歌時,他發現他們現有的模型都在 CPU 上執行。Krizhevsky 需要 GPU。所以他親自解決了這個問題!

    在公司的最初幾天,他從當地的電子商店購買了一台 GPU 機器,把它放在離他的辦公桌不遠的走廊的衣櫥裏,將其接入網路,並開始在這台獨立的硬體上訓練他的神經網路。

    最終,Krizhevsky 的新同事們意識到,他們需要 GPU,而且是很多 GPU。2014 年,谷歌決定購買大約 40,000 個 Nvidia GPU,花費約 1.3 億美元。這些 GPU 被用於訓練谷歌業務中使用的深度學習模型。

    順便說一句,像谷歌這樣的公司購買 40,000 個 GPU 並不會被忽視。 如果 Nvidia 需要一個訊號,說明將 GPU 套用於深度學習可能是一個規模可觀的業務,那麽谷歌幾乎可以肯定在 2014 年提供了這個訊號。

    但這些 GPU 並不一定解決了谷歌最大的潛在問題。它們非常適合訓練谷歌正在開發的模型。但是,部署這些模型在谷歌的規模上會如何呢?深度學習的套用,如語音辨識,可能會非常受歡迎。這種潛在的受歡迎程度給谷歌帶來了一個問題。讓我們考慮一下谷歌與之同義的產品,搜尋。

    谷歌搜尋對任何擁有網路瀏覽器和互聯網連線的人來說都是免費的。這種搜尋的品質以及它的免費性,是推動其爆炸性增長和迅速占主導地位的關鍵因素。據估計,到 2023 年,谷歌提供了大約兩萬億次搜尋結果。要能夠以這種規模免費提供搜尋,谷歌需要能夠非常便宜地提供每一組搜尋結果。

    谷歌團隊可以預見到使用這些新的深度學習技術構建的一系列服務將會非常受歡迎。谷歌可以利用其品牌,以及將這些服務與其現有廣受歡迎的產品如搜尋、Gmail 和 Android 整合,再次使這些服務占據主導地位。

    但有一個重大問題。在這種規模下,這些服務將需要投資更多的硬體。2013 年,谷歌團隊設計了一個使用語音辨識的範例,說明了規模。如果使用者每天只對他們的 Android 手機說三分鐘話,並且谷歌使用深度學習將這些語音轉換為文本使用 CPU,那麽公司將需要增加 2 倍或 3 倍的伺服器數量,而這些伺服器已經在使用中。這將是極其昂貴的。

    一個替代方案是擴大使用 GPU,Nvidia GPU 是顯而易見的首選。到 2013 年,Nvidia 的 GPU 及其相關的軟體生態系已經非常成熟,被用作機器學習研究的關鍵工具,用於訓練和推理。Nvidia 的第一款針對通用計算任務的 GPU 出現在 2006 年,Nvidia 的 CUDA 框架,用於在 GPU 上編程通用計算任務,出現在 2007 年。關鍵的是,Nvidia GPU 可以立即購買。

    正如我們所看到的,谷歌確實購買了 Nvidia GPU 並將它們安裝在其數據中心。然而,依賴 Nvidia GPU 並不一定是技術上或戰略上的最佳解決方案。GPU 比 CPU 更適合深度學習所需的計算,但它們並不完全專門用於這些計算。依賴 GPU 會讓潛在的效率損失,而在谷歌的規模下,這些效率損失代表了重大的成本節省損失。此外,依賴單一供應商提供對谷歌戰略至關重要的硬體會增加該戰略的顯著風險。

    有兩個其他選擇。第一個是谷歌可以使用 FPGA,它們可以編程執行深度學習所需的專用計算。

    第二個是谷歌可以設計和構建自己的客製硬體,形式為特定套用積體電路(ASIC)。使用這種專門用於深度學習的客製硬體,將開放更多的效率潛力,並將減少谷歌對單一供應商的依賴。

    與已經安裝的 GPU 相比,FPGA 的效能很快就顯示出來會被擊敗。這些 FPGA 並沒有完全被放棄。

    FPGA 仍然在數據中心中使用,比 ASIC 可用的時間早六個月;它充當了一個「清潔管道」,用於所有支持生產中新加速器的谷歌部署過程。

    但重點將是 ASIC。目標將是開發一個 ASIC,與 GPU 相比,在推理上產生 10 倍的成本效能優勢。

    但開發客製硬體需要時間。具體需要多長時間將是至關重要的。對於該計畫來說,它不能是一個拖延多年的「研究」計畫。它需要快速地向谷歌的數據中心交付硬體,並且要以大規模交付。為此,該計畫將需要利用谷歌能夠快速獲取的任何資源和現有知識。

    TPU 團隊和現有研究

    如何在沒有現有內部團隊的情況下快速設計復雜和創新的硬體。令人驚訝的是,谷歌很快就能組建一個非常有效的團隊。它是如何做到的?

    谷歌已經在為其數據中心建立客製硬體,並且可以召集該團隊的成員來幫助新計畫。然而,該公司以前沒有建立過自己的處理器芯片,所以它開始引入具有相關經驗的新團隊成員。他們是如何做到這一點的?

    其中一位新成員 Norm Jouppi 回憶說,他 是如何在一次會議上與 Jeff Dean 交談後被招募到 TPU 團隊的,以及如何被一個創新的「green field」計畫所吸引。加入團隊的其他知名人物包括 David Patterson( 圖靈獎得主 ),原始 Berkeley RISC 設計的開發者和 RISC-V 指令集架構開發的關鍵人物。

    谷歌很快就擁有了一個才華橫溢且經驗豐富的團隊。但即使如此,如果他們要從頭開始開發新系統的架構,他們可能還是難以滿足為他們設定的目標。幸運的是,對於谷歌來說,有一個現有的架構方法可以很容易地套用。這種方法是在開始工作於 TPU 之前三十五年多的時間裏首次開發的。

    在他們 1978 年的論文 Systolic Arrays (for VLSI)中,卡內基梅隆大學的 H.T Kung 和 Charles E. Leiserson 提出了他們所稱的「systolic system」(脈動系統)的建議。

    systolic system 是一個處理器網路,它有節奏地計算並透過系統傳遞數據....在脈動電腦系統中,處理器的功能類似於心臟。每個處理器定期地輸入和輸出數據,每次進行一些簡短的計算,以便在網路中保持數據的正常流動。

    然後他們確定了 systolic system 的一個套用,矩陣計算。

    許多基本的矩陣計算可以在具有陣列結構的脈動網路上優雅而有效地進行流水線處理。作為一個例子,六邊形連線的處理器可以最優地執行矩陣計算...這些脈動陣列享有簡單和規則的通訊路徑,而且網路中使用的幾乎所有處理器都是相同的。因此,基於脈動陣列的專用硬體裝置可以使用 VLSI 技術廉價構建。

    作為 Kung 和 Leiserson 所處理的技術的年代的一個插圖,他們說明了如何使用「脈動系統」作為 DEC PDP-11 小型電腦的附件。

    該論文描述了脈動陣列的處理元素可能連線的替代方式。

    其中一個提議的安排 - 正交連線 - 與 ILLIAC IV 超級電腦中處理元素的安排產生了共鳴,我們在 ILLIAC IV 超級電腦:DARPA、SIMD、Fairchild 和 Stanley Kubrick 的'2001'中討論了這一點)。

    Kung 和 Leiserson 繼續描述了如何使用六邊形連線的脈動陣列來執行矩陣乘法。

    Kung 和 Lieberson 的這項原始工作在接下來的十年中得到了更多的研究和更多的論文,展示了如何使用脈動陣列解決各種問題。

    脈動陣列是如何工作的?我們將在下一篇貼文中更詳細地檢視 Google 的 TPU 團隊如何有效地實作矩陣運算的概念。基本思想是數據透過陣列的一側或多側輸入系統,並且數據和中間結果隨著每個「脈沖」的流動透過系統。所需的結果在經過足夠的脈沖導致所需計算完成後,從陣列的一側或多側出現。

    到 2013 年,Kung 和 Lieberson 想法背後的一些原始動機,特別是處理 1970 年代制造技術的限制,已經消失。然而,這種方法對於像矩陣乘法這樣的任務的固有效率,特別是在 2013 年相關的,相對較低的功耗,仍然存在。因此,TPU 將使用脈動陣列。

    谷歌擁有工程師和架構方法,但在將想法轉化為工作矽片的能力上仍然存在差距。為此,它將轉向 LSI 公司( 現為 Broadcom 的一部份 )。乍一看,現在的 Broadcom 可能看起來像一個令人驚訝的合作夥伴。它不是一個通常與機器學習硬體相關聯的公司。但它們可以與像 TSMC 這樣的制造合作夥伴合作,使谷歌的設計成為現實,並且以大規模實作。

    最後,讓 TPU 能夠用於實際工作負載意味著的不僅僅是制造大量芯片。這意味著構建軟體,以便谷歌的深度學習工具可以在新設計上執行。新設計必然意味著一個全新的指令集架構,這意味著編譯器必須適應新架構,這本身就是一個不小的壯舉。

    部署和使用

    TPU 團隊的成員總結了他們的目標如下:

    快速構建

    實作 高效能

    ......在大規模......

    ...對新工作負載開箱即用......

    同時具有成本效益

    他們確實快速構建了它。 計畫啟動僅 15 個月後, 第一款 TPU 就於 2015 年初部署在谷歌的數據中心。

    他們是如何做到這一點的?論文【「數據中心效能分析的回顧:張量處理單元」】總結了一些因素:

    我們有時會吹噓 TPU 從計畫啟動到數據中心部署的 15 個月時間,遠遠短於標準的生產芯片。確實,谷歌依賴於生產的後續 TPUs 具有多年的設計周期。快速上市得益於單一的日程焦點,不僅在架構上——其中 700MHz 的時鐘速率使容易實作時間閉環,老化的 28nm 工藝已經完全偵錯——還在我們的數據中心部署團隊的英勇工作上。

    「老化的 28nm 工藝」和相對較低的時鐘速率是有幫助的因素。然而,這忽略了谷歌能夠調動必要資源來支持這樣的計畫的能力。

    TPU 很快就會在谷歌內部用於廣泛的套用。很快迎來了一個有趣且高調的例子:谷歌的 DeepMind 使用 TPU 用於其 AlphaGo Zero 計畫。AlphaGo Zero 是在 2016 年 3 月擊敗了世界圍棋冠軍李世石的模型。

    AlphaGo Zero 是本文所描述的程式。它從自我對弈的強化學習中學習,從隨機的初始權重開始,不使用展開,沒有人類的監督,僅使用原始棋盤歷史作為輸入特征。它僅使用谷歌雲中的一台機器和 4 個 TPU。

    2016 年 AlphaGo Zero 的成功似乎是一個重要時刻。以如此少的硬體實作這一成就是非凡的。

    TPU 的宣布

    像谷歌這樣的公司不需要分享其數據中心幕後發生的細節。因此,TPU 在 2015 年首次引入後一年多保持秘密。

    然後谷歌的 CEO Sundar Pichai 在 2016 年 5 月 18 日谷歌 IO 大會的主題演講中宣布:

    我們在我們的數據中心內執行 TPUs 超過一年了,並且發現它們在機器學習任務中提供了每瓦效能的數量級提升。

    這個釋出伴隨著一個簡短的博文,谷歌使用 TPU 客製芯片為機器學習任務增加動力。除了一些非常簡短的技術細節,該博文還分享了一些 TPU 如何幫助支持谷歌的服務。

    TPU 已經為谷歌的許多套用提供動力,包括用於提高搜尋結果相關性的 RankBrain,以及用於提高我們的地圖和導航的準確性和品質的模組屋檢視。AlphaGo 在與圍棋世界冠軍李世石的比賽中由 TPU 提供動力,使其能夠「更快地思考」並在移動之間更遠地看到。

    深入 TPU 架構

    正如我們所見,TPU v1 團隊采用的方法是 H.T Kung 和 Charles E. Leiserson 在 1978 年的論文【脈動陣列(針對 VLSI)】中首次提出的架構。

    脈動系統是一個由處理器組成的網路,這些處理器有節奏地計算並傳遞數據……在脈動電腦系統中,處理器的功能類似於心臟。每個處理器定期地輸入和輸出數據,每次進行一些簡短的計算,以此保持網路中數據的正常流動。

    那麽,在 TPU v1 中如何使用脈動方法有效執行矩陣乘法呢?讓我們回到我們的 2x2 矩陣乘法範例。

    如果我們有一個 2x2 的乘法單元陣列,這些單元在一個簡單的網格中連線,並且我們以正確的順序將我們要相乘的矩陣的元素輸入到網格中,那麽矩陣乘法的結果自然會從陣列中顯現。

    計算可以用以下圖表示。每個角落的方塊代表一個乘法/累加單元(MAC),它可以執行乘法和加法操作。

    在這張圖中,黃色的值是從頂部和左側輸入矩陣的輸入值。淺藍色的值是儲存的部份和。深藍色的值是最終結果。

    讓我們一步步來分析。

    第 1 步:

    值 a11 和 b11 被載入進左上角的乘法/累加單元(MAC)。它們被相乘,結果被儲存。

    第 2 步:

    值 a12 和 b21 被載入進左上角的 MAC。它們被相乘,並加到之前計算的結果上。這給出了結果矩陣的左上角值。

    與此同時,b11 被傳輸到右上角的 MAC,在那裏它與新載入的值 a21 相乘,結果被儲存。同時,a11 被傳輸到左下角的 MAC,在那裏它與新載入的值 b12 相乘,結果被儲存。

    第 3 步:

    b21 被傳輸到右上角的 MAC,在那裏它與新載入的值 a22 相乘,結果被加到之前儲存的結果上。同時,a12 被傳輸到左下角的 MAC,在那裏它與新載入的值 b22 相乘,結果被加到之前儲存的結果上。在這一步中,我們計算了結果矩陣的右上角和左下角值。

    與此同時,a12 和 b21 被傳輸到右下角的 MAC,在那裏它們被相乘,結果被儲存。

    第 4 步:

    最後,a22 和 b22 被傳輸到右下角的 MAC,在那裏它們被相乘,結果被加到之前儲存的值上,給出了結果矩陣的右下角值。

    因此,矩陣乘法的結果沿著 MAC 矩陣中移動的「對角線」顯現出來。

    在我們的範例中,進行 2x2 矩陣乘法需要 4 個步驟,但這僅是因為在計算開始和結束時一些 MAC 沒有被利用。實際上,一旦 MAC 變為空閑,新的矩陣乘法就會從左上角開始。因此,該單元每兩個周期就能完成一次新的矩陣乘法。

    這是對脈動陣列工作方式的簡化表示,我們簡要地介紹了 TPU v1 中脈動陣列實作的一些細節。我希望這種架構工作的原理已經清楚。

    這是最簡單可能的矩陣乘法,但可以擴充套件到更大的矩陣和更大的乘法單元陣列。

    關鍵點是,如果數據以正確的順序輸入到脈動陣列中,則透過系統的值和結果的流動將確保所需結果隨時間從陣列中出現。沒有必要將中間結果儲存並從「主記憶體儲」區域提取。由於矩陣乘法單元的結構和輸入被輸入到單元的順序,中間結果在需要時自動可用。

    當然,矩陣乘法單元不是孤立存在的,整個系統的最簡單呈現如下:

    首先要註意的是,TPUv1 依賴於透過 PCIe( 高速序列匯流排 )介面與主機電腦的通訊。它還可以直接存取自己的 DDR3 動態 RAM 儲存。

    我們可以將其擴充套件為設計的更詳細展示:

    讓我們從這個設計的展示中 挑選一些關鍵元素,從頂部開始,大致順時針移動:

  • DDR3 DRAM / Weight FIFO: 權重儲存在透過 DDR3-2133 介面連線到 TPU v1 的 DDR3 RAM 芯片中。權重是從主電腦的記憶體透過 PCIe 預載入到這些芯片上的,然後可以轉移到「Weight FIFO」記憶體中,為矩陣乘法單元的使用做好準備。

  • 矩陣乘法單元: 這是一個「脈動」陣列,擁有 256 x 256 的矩陣乘法/累加單元,由頂部的 256 個「權重」值和左側的 256 個數據輸入供給。

  • 累加器: 結果從脈動矩陣單元的底部出現,並儲存在「累加器」記憶體儲存中。

  • 啟用: 在此套用上述神經網路中描述的啟用函式。

  • 統一緩沖區/脈動數據設定: 套用啟用函式的結果儲存在「統一緩沖區」記憶體中,準備作為輸入反饋到矩陣乘法單元以計算下一層所需的值。

  • 乘法/累加單元(MACs)

    到目前為止,我們還沒有具體說明矩陣乘法單元執行的乘法的性質。TPU v1 執行 8 位 x8 位整數乘法,利用量化來避免對更占用更多矽片面積的浮點計算的需求。

    指令集

    TPU v1 采用 CISC(復雜指令集電腦)設計,大約只有 20 條指令。值得註意的是,這些指令是由主電腦透過 PCIe 介面發送給它的,而不是從記憶體中獲取的。

    五個關鍵指令如下:

    Read_Host_Memory

    透過 PCIe 從主電腦的記憶體讀取輸入值到統一緩沖區。

    Read_Weights

    從權重記憶體讀取權重到 Weight FIFO。註意,權重記憶體將已經透過 PCIe 從電腦的主記憶體中讀取了權重。

    Matrix_Multiply / Convolve

    根據論文,這條指令:

    …使矩陣單元執行矩陣乘法或摺積,從統一緩沖區到累加器。矩陣操作接受一個變量大小的 B*256 輸入,將其與一個 256x256 的恒定權重輸入相乘,產生一個 B*256 的輸出,完成 B 流水線周期。

    這是實作脈動陣列矩陣乘法的指令。它還可以執行摺積計算,這對摺積神經網路是必需的。

    Activate

    根據論文,這條指令:

    執行人工神經元的非線性函式,有 ReLU、Sigmoid 等選項。它的輸入是累加器,輸出是統一緩沖區。

    如果我們回到我們簡單的神經網路模型,隱藏層中的值是將「啟用函式」套用於輸入的權重乘以輸入的總和的結果。ReLU 和 Sigmoid 是兩種最受歡迎的啟用函式。在硬體中實作這些將提供套用啟用函式的有用加速。

    Write_Host_Memory

    透過 PCIe 將結果從統一緩沖區寫入主電腦的記憶體。

    值得暫停一刻,反思這五條指令在提供幾乎完整的 TPU v1 推理實作方面的優雅。在虛擬碼中,我們可以大致描述 TPU v1 的操作如下:

    Read_Host_Memory
    Read_Weights
    Loop_Start
    Matrix_Multiply
    Activate
    Loop_End
    Write_Host_Memory

    強調脈動單元在使這成為可能和高效中的重要性也是有用的。正如 TPU v1 團隊所描述:

    ..矩陣單元使用脈動執行以節省能量,透過減少對統一緩沖區的讀寫....它依靠來自不同方向的數據在陣列中的單元以規律間隔到達時被組合。...數據從左側流入,權重從頂部載入。給定的 256 元素乘法-累加操作以對角線波前的形式透過矩陣移動。

    軟體

    沒有軟體支持的 TPU v1 硬體將毫無用處。谷歌開發並使用了 Tensorflow,因此建立「驅動程式」以使 Tensorflow 能夠與 TPU v1 協同工作是主要步驟。

    TPU 軟體棧必須與為 CPU 和 GPU 開發的軟體棧相容,以便應用程式可以快速移植到 TPU。在 TPU 上執行的應用程式的部份通常用 TensorFlow 編寫,並編譯成可以在 GPU 或 TPUs 上執行的 API。

    與 GPUs 一樣,TPU 棧分為使用者空間驅動程式和內核驅動程式。內核驅動程式很輕量級,僅處理記憶體管理和中斷。它設計用於長期穩定。使用者空間驅動程式變化頻繁。它設定和控制 TPU 執行,將數據重新格式化為 TPU 順序,將 API 呼叫轉換為 TPU 指令,並將它們轉換為應用程式二進制檔。

    制造和模具

    TPU v1 是由 TSMC 使用相對「成熟」的 28nm TSMC 工藝制造的。谷歌表示,與谷歌此時在其數據中心使用的更先進工藝制造的 Intel Haswell CPU 和 Nvidia 的 K80 GPU 芯片相比,模具面積不到一半。

    我們已經看到了 TPU v1 的指令集有多簡單,只有 20 條 CISC 指令。ISA 的簡單性導致 TPU v1 的模具中用於解碼和相關活動的「開銷」非常低,僅占模具面積的 2%,標記為「控制」。

    相比之下,24%的模具面積專用於矩陣乘法單元,29%專用於儲存輸入和中間結果的「統一緩沖區」記憶體。

    效能

    此時,回顧一下 TPU v1 旨在使推理——即在谷歌規模的真實世界服務中使用已經訓練好的模型——更高效是有用的。它不旨在提高訓練的速度或效率。盡管推理和訓練有一些共同的特征,但在開發專用硬體時,推理和訓練提供了完全不同的挑戰。

    那麽 TPU v1 的表現如何呢?

    2013 年,TPU v1 的關鍵比較是與 Intel 的 Haswell CPU 和 Nvidia 的 K80 GPU。

  • TPU v1 擁有 K80 GPU 的 25 倍 MAC 和 3.5 倍的片上記憶體。

  • TPU v1 在推理方 面比 K80 GPU 和 Haswell CPU 快大約 15X - 30X。

  • 並且關鍵的是 TPU v1 在能源效率方面遠超 GPU:

  • TPU v1 的相對增量效能/瓦特是 GPU 的 25 到 29 倍。

  • 教訓

    在本文中,我們關註了像谷歌這樣的組織如何能夠快速構建 TPU v1。同時,也看到 了 TPU v1 的客製架構對於使其能夠以比當代 CPU 和 GPU 更低的能耗產生更好的效能至關重要。

    TPU v1 只是故事的開始。TPU v1 設計得很快,唯一的目標是使推理更快和更節能。它有一些明顯的限制,不是為訓練設計的。不僅谷歌內部,外部公司很快就開始考慮如何改進 TPU v1。

    本文使用 GPT-4 生成,原文參見:

    https://thechipletter.substack.com/p/googles-first-tpu-architecture

    https://thechipletter.substack.com/p/googles-first-tensor-processing-unit

    4 月 25 ~ 26 日,由 CSDN 和高端 IT 咨詢和教育平台 Boolan 聯合主辦的「 全球機器學習技術大會 」將在上海環球港凱悅酒店舉行,特邀近 50 位技術領袖和行業套用專家,與 1000+ 來自電商、金融、汽車、智慧制造、通訊、工業互聯網、醫療、教育等眾多行業的精英參會聽眾,共同探討人工智慧領域的前沿發展和行業最佳實踐。 歡迎所有開發者朋友存取官網 http://ml-summit.org、點選「閱讀原文」或掃碼進一步了解詳情。