對于一家自動駕駛電動汽車制造商來說,花費數(shù)億美元從頭開始創(chuàng)建自己的人工智能超級計算機,超大規(guī)模人工智能訓練的成本和難度有多大?公司創(chuàng)始人必須多么自負和肯定才能組建一支能夠做到這一點的團隊?
像許多問題一樣,當您準確地提出這些問題時,他們往往會自己回答。很明顯,SpaceX 和特斯拉的創(chuàng)始人兼 OpenAI 聯(lián)盟的聯(lián)合創(chuàng)始人埃隆馬斯克沒有時間或金錢浪費在科學項目上。
就像世界上的超級大國低估了完全模擬核導彈及其爆炸所需的計算能力一樣,也許自動駕駛汽車的制造商開始意識到,在復雜的世界中教汽車自動駕駛這種情況總是在變化,這將需要更多的超級計算。一旦你接受了這一點,你就可以從頭開始,建造合適的機器來完成這項特定的工作。
簡而言之,這就是特斯拉的 Project Dojo 芯片、互連和超級計算機工作的全部內(nèi)容。
在Hot Chips 34大會上,曾在Dojo超級計算機上工作的芯片、系統(tǒng)和軟件工程師首次公開了該機器的許多架構(gòu)特性,并承諾將在特斯拉AI日上談論Dojo系統(tǒng)的性能。
Emil Talpes 在 AMD 工作了近 17 年,研究各種 Opteron 處理器以及命運多舛的“K12”Arm 服務器芯片,他介紹了他的團隊創(chuàng)建的 Dojo 處理器。Debjit Das Sarma 則同期在 AMD 擔任 CPU 架構(gòu)師,他在演講中受到贊譽,目前是特斯拉的自動駕駛硬件架構(gòu)師,Douglas Williams 也是如此,我們對他一無所知。這家汽車制造商的首席系統(tǒng)工程師 Bill Chang 在 IBM Microelectronics 工作了 15 年,設計 IP 模塊并致力于制造工藝,然后幫助 Apple 將 X86 處理器轉(zhuǎn)移到自己的 Arm 芯片上,而Rajiv Kurian一開始在特斯拉工作,然后在Waymo工作。據(jù)我們所知,在去年 8 月的 Tesla AI Day 1 上發(fā)言特斯拉自動駕駛硬件高級總監(jiān)Ganesh Venkataramanan負責Dojo項目。過去十五年里,Venkataramanan 還是 AMD CPU 設計團隊的領導者。
所以以一種奇怪的方式,Dojo 代表了一個可替代的人工智能未來,如果特斯拉來幫助從頭開始設計定制的人工智能超級計算機,從全新核心內(nèi)核中的向量和整數(shù)單元一直到一個完整的 exascale 系統(tǒng),專為 AI 訓練用例的規(guī)?;鸵子诰幊潭O計。
與來自 AI 初創(chuàng)公司的許多其他相對較新的平臺一樣,Dojo 設計優(yōu)雅而徹底。最引人注目的是特斯拉工程師在關注規(guī)模時拋出的東西。
“我們應用定義的目標是可擴展性,”Talpes 在演講結(jié)束時解釋道?!拔覀儾辉購娬{(diào)典型 CPU 中的幾種機制,例如一致性、虛擬內(nèi)存和全局查找目錄,因為當我們擴展到非常大的系統(tǒng)時,這些機制不能很好地擴展。相反,我們依賴于整個網(wǎng)格中非??焖偾曳浅7植际降?SRAM 存儲。與典型分布式系統(tǒng)相比,互連速度高出一個數(shù)量級?!?/p>
基于此,讓我們深入了解 Dojo 架構(gòu)。
根據(jù) Talpes 的說法,Dojo 核心有一個整數(shù)單元,它從 RISC-V 架構(gòu)中借用了一些指令,并且有一大堆特斯拉自己創(chuàng)建的附加指令。矢量數(shù)學單元“主要由特斯拉從頭開始實現(xiàn)”,Talpes 沒有詳細說明這意味著什么。他確實補充說,這個自定義指令集針對運行機器學習內(nèi)核進行了優(yōu)化,我們認為這意味著它不會很好地運行孤島危機。
Dojo 指令集支持 64 位標量指令和 64 B SIMD 指令,它包括處理從本地內(nèi)存到遠程內(nèi)存?zhèn)鬏敂?shù)據(jù)的原語(primitives),并支持信號量(semaphore)和屏障約束( barrier constraints),這是使內(nèi)存操作符合指令不僅在 D1 內(nèi)核中運行,而且在 D1 內(nèi)核的集合中運行。至于特定于 ML 的指令,有一組通常在軟件中完成的 shuffle、transpose 和 convert 指令,現(xiàn)在蝕刻在晶體管中,核心還進行隨機舍入( stochastic rounding ),可以進行隱式 2D 填充(implicit 2D padding),即通常通過在一條數(shù)據(jù)的兩側(cè)添加零來調(diào)整張量來完成。
Talpes 明確表示,D1 處理器是我們推測的 Dojo 芯片和系統(tǒng)系列中的第一個,是“高吞吐量、通用 CPU”,它本身并不是加速器。或者更準確地說,Dojo 的架構(gòu)旨在加速自身,而不需要一些外部設備來完成。
每個 Dojo 節(jié)點都有一個內(nèi)核,是一臺具有 CPU 專用內(nèi)存和 I/O 接口的成熟計算機。這是一個重要的區(qū)別,因為每個內(nèi)核都可以做自己的事情,而不依賴于共享緩存或寄存器文件或任何東西。
D1 是一個超標量(superscalar)內(nèi)核,這意味著它在其內(nèi)核中支持指令級并行性,就像當今大多數(shù)芯片一樣,它甚至具有多線程設計來驅(qū)動更多指令通過該內(nèi)核。但是多線程更多的是每時鐘做更多的工作,而不是擁有可以將獨立的 Linux 實例作為虛擬機運行的獨立線程,因此同步多線程 (SMT) 的 Dojo 實現(xiàn)沒有虛擬內(nèi)存,保護機制有限,并且Dojo 軟件堆棧和應用程序管理芯片資源的分配。
D1 內(nèi)核是一個 64 位處理器,具有 32 B 的取指窗口( fetch window),最多可容納 8 條指令;一個八寬解碼器每個周期可以處理兩個線程。這個前端饋入( front end feeds into)一個四寬標量調(diào)度器(four-wide scalar schedule),該調(diào)度器具有四路 SMT,它有兩個整數(shù)單元、兩個地址單元和一個用于每個線程的寄存器文件。還有一個帶有四路 SMT 的兩側(cè)向量調(diào)度器,它饋送到一個 64 B 寬的 SIMD 單元或四個 8x8x4 矩陣乘法單元。
每個 D1 內(nèi)核都有一個 1.25 MB 的 SRAM,這是它的主存儲器。它不是緩存,如果有的話,掛在更大的 Dojo 網(wǎng)絡上的 DDR4 內(nèi)存被視為比其他任何東西都更像大容量存儲。該 SRAM 可以以 400 GB/秒的速度加載并以 270 GB/秒的速度存儲,并且該芯片具有明確的指令,可以將數(shù)據(jù)移入或移出 Dojo 機器中其他內(nèi)核的外部 SRAM 存儲器。嵌入在該 SRAM 中的是一個列表解析器引擎(list parser engine),該引擎饋入解碼器對和一個收集引擎(gather engine),饋入向量寄存器文件,它們一起可以將信息發(fā)送到其他節(jié)點或從其他節(jié)點獲取信息,而無需像與其他 CPU 架構(gòu)。
這個列表解析功能是 Dojo 芯片設計獨有的關鍵特性之一:
這本質(zhì)上是一種封裝不同數(shù)據(jù)位的方法,以便可以在系統(tǒng)中的 D1 內(nèi)核之間高效傳輸。
D1 內(nèi)核支持多種數(shù)據(jù)格式。標量單元支持 8、16、32 或 64 位的整數(shù),而向量單元及其關聯(lián)的矩陣單元支持多種數(shù)據(jù)格式,具有精度和數(shù)值范圍的混合,其中不少是動態(tài)的可由 Dojo 編譯器組合。
Talpes 指出,F(xiàn)P32 格式比 AI 訓練應用的許多部分所需的精度和范圍更廣,IEEE 指定的 FP16 格式?jīng)]有足夠的范圍覆蓋神經(jīng)網(wǎng)絡中的所有處理層;相反,Google Brain 團隊創(chuàng)建的 Bfloat 格式范圍更廣,但精度更低。因此,Tesla 不僅提出了用于較低精度和更高吞吐量矢量處理的 8 位 FP8 格式,而且還提出了一組可配置的 8 位和 16 位格式,Dojo 編譯器可以在尾數(shù)的精度附近滑動和上圖所示的指數(shù),以涵蓋更廣泛的范圍和精度。在任何給定時間,最多可以使用 16 種不同的矢量格式,但每個 64 B 數(shù)據(jù)包必須屬于同一類型。
在圖的右上角,您將看到片上網(wǎng)絡路由器( network on chip router),它將多個核心連接在一起形成一個 2D 網(wǎng)格。NOC 可以處理跨節(jié)點邊界的 8 個數(shù)據(jù)包(boundary),每個方向 64 B,每個時鐘周期,即在所有四個方向上一個數(shù)據(jù)包輸入和一個數(shù)據(jù)包輸出到網(wǎng)格中每個核心最近的鄰居。該路由器還可以在每個周期對本地 SRAM 進行一次 64 B 讀取和一次 64 B 寫入,因此可以在內(nèi)核之間移動數(shù)據(jù)。
這些都在在D1 核心上完成所有蝕刻,是由其代工合作伙伴臺積電以 7 納米工藝完成的。之后,特斯拉開始復制 D1 核心并將它們互連網(wǎng)格,像這樣:
十幾個 D1 核心排列成一個本地塊,創(chuàng)建了一個 18 核心乘 20 核心的二維陣列,但由于某種原因,只有 354 個 D1 核心可用。D1 芯片以 2 GHz 運行,在這些內(nèi)核上共有 440 MB 的 SRAM,在 BF16 或 CFP8 上提供 376 teraflops,在 FP32 上提供 22 teraflops。向量單元中沒有 FP64 支持。如此多的 HPC 工作負載無法在此 D1 芯片上運行,一些使用 64 位矢量數(shù)學的 AI 應用也不會。特斯拉不需要關心——它只需要運行自己的人工智能應用程序,如果它想在 D2 或 D3 芯片上添加 FP64 支持來運行其 HPC 模擬和建模工作負載,以便馬斯克的公司可以設計宇宙飛船和汽車,好吧,在完成所有這些工作之后,這相當容易。
D1 die 有 576 個雙向 SerDes 通道,圍繞在 die 周圍以鏈接到其他 D1 die,并且 D1 die 的所有四個邊緣的帶寬為 8 TB/秒。這些芯片重達 645 平方毫米,可通過這些 SerDes 無縫連接到特斯拉所謂的 Dojo 訓練模塊中。像這樣:
訓練tile采用 25 個已知良好的 D1 裸片,并將它們打包成一個 5×5 陣列,相互連接。訓練tile的外部邊緣在 40 個 I/O 芯片上實現(xiàn)了 36 TB/秒的聚合帶寬;這是 2D 網(wǎng)格的網(wǎng)格的二分帶寬的一半,它跨越了 tile 內(nèi)的 D1 芯片。該tile具有 10 TB/秒的塊上二分帶寬,以及跨內(nèi)核的 11 GB SRAM 內(nèi)存。每個 tile 提供 9 petaflops 的 BF16/CFP8 魅力。
那些Dojo訓練tile消耗15千瓦,顯然是水冷的,它們的設計使得多個訓練tile可以與相鄰tile互連。目前尚不清楚這是如何發(fā)生的,但很明顯,您需要一排相互連接的tile,水平或垂直方向,而不是帶有幾個設備托盤的單獨機架,然后需要某種巨大比例的光纜或電纜,圍繞在tile之間承載數(shù)據(jù)。垂直工作,如下所示:
您將在上圖中注意到,在 D1 網(wǎng)格的邊緣有所謂的 Dojo 接口處理器或 DIP,它們連接到 D1 網(wǎng)格以及為 DIP 供電并執(zhí)行各種系統(tǒng)管理的主機系統(tǒng)功能。每個訓練tile總共有 11 GB 的私有 SRAM 主內(nèi)存,但系統(tǒng)需要某種更大的內(nèi)存,該內(nèi)存合理地靠近網(wǎng)格。在這種情況下,Tesla 選擇創(chuàng)建一個 DIP 內(nèi)存和 I/O 協(xié)處理器,其中包含 32 GB 共享 HBM 內(nèi)存——我們還不知道是哪種,但它是 HBM2e 或 HBM3——以及以太網(wǎng)接口到外部世界以及在tile和核心之間進行比通過這個巨大的網(wǎng)格更直接的跳躍。圖片顯示一對主機安裝了十個這樣的 DIP,每組三個 Dojo 訓練圖塊總共有 320 GB 的 HBM 內(nèi)存。但圖表上的措辭表明,每個 tile 分配了 160 GB,這意味著每個 tile 一個主機,而不是此處顯示的三個 tile 兩個主機。
該 DIP 卡有兩個 I/O 處理器,每個處理器帶有兩個 HBM 內(nèi)存組,該卡提供 32 GB 的 HBM 內(nèi)存和 800 GB/秒的帶寬。對我們來說,這看起來像是稍微降低了 HBM2e 內(nèi)存。該卡通過 PCI-Express 實現(xiàn)了 Tesla 傳輸協(xié)議 (TTP:Tesla Transport Protocol ),這是一種專有互連,對我們來說這有點像 CXL 或 OpenCAPI,以將完整的 DRAM 內(nèi)存帶寬提供給 Dojo 訓練塊。在卡的另一端,有一個 50 GB/秒的 TTP 協(xié)議鏈路在以太網(wǎng) NIC 上運行,它連接到現(xiàn)有的以太網(wǎng)交換機,該交換機可以是單個 400 Gb/秒端口或成對的 200 Gb/秒端口。DIP 插入 PCI-Express 4.0 x16 插槽,每張卡提供 32 GB/秒的帶寬。每個磁貼邊緣有 5 個卡,有 160 GB/秒的帶寬進入主機服務器和 4 個。
正如我們已經(jīng)指出的那樣,DIP 不僅將 DRAM 實現(xiàn)為fat本地存儲,而且還提供了另一種網(wǎng)絡維度,可用于繞過 2D 網(wǎng)格,而需要大量躍點才能跨越所有這些核心和tile。像這樣:
Chang 表示,在整個系統(tǒng)中通過 2D 網(wǎng)格實現(xiàn)端到端可能需要 30 hops,但使用 TTO over Ethernet 協(xié)議和fat tree Ethernet 交換機網(wǎng)絡,只需要 4 hops。顯然,帶寬要低得多,但在網(wǎng)絡的第三維(因此是 Z 平面)上的延遲要低得多。
Dojo V1 訓練矩陣是 Tesla 正在構(gòu)建的基礎系統(tǒng),它有 6 個訓練tile、4 個主機服務器上的 20 個 DIP,以及一組連接到以太網(wǎng)交換結(jié)構(gòu)的輔助服務器,如下所示:
基礎 Dojo V1 系統(tǒng)有 53,100 個 D1 內(nèi)核,在 BF16 和 CFP8 格式下的額定速度為 1 exaflops,在 Tile 上有 1.3 TB 的 SRAM 內(nèi)存,在 DIP 上有 13 TB 的 HBM2e 內(nèi)存。完整的 Dojo ExaPod 系統(tǒng)共有 120 個tile,將有 1,062,000 個可用的 D1 內(nèi)核,重量為 20 exaflops。
順便說一句,這個東西運行 PyTorch。沒有像 C 或 C++ 那樣低級的東西,也沒有像 CUDA 這樣的遠程。Dojo 機器的另一個巧妙之處在于 SRAM 將自身呈現(xiàn)為單個地址空間。它是一個平坦的內(nèi)存區(qū)域,本地計算懸掛在其內(nèi)存塊上。
我們期待看到 Dojo 在 AI 基準測試中的表現(xiàn)。
更多信息可以來這里獲取==>>電子技術(shù)應用-AET<<