特斯拉在 8 月中旬的 AI 日推出了用于機器學(xué)習(xí)和神經(jīng)網(wǎng)絡(luò)訓(xùn)練的汽車芯片、系統(tǒng)和軟件。他們將共同推進自動駕駛汽車模型的培訓(xùn)。
埃隆馬斯克和他的芯片和系統(tǒng)設(shè)計師團隊在三個多小時的演示中提供了大量技術(shù)細節(jié)。以下是重點。
神經(jīng)網(wǎng)絡(luò)
特斯拉設(shè)計了一種靈活、可擴展的分布式計算機架構(gòu),專為神經(jīng)網(wǎng)絡(luò)訓(xùn)練量身定做。Tesla 的架構(gòu)始于具有 354 個訓(xùn)練節(jié)點的 D1 專用芯片,每個訓(xùn)練節(jié)點都有一個強大的 CPU。這些訓(xùn)練節(jié)點 CPU 專為高性能 NN 和 ML 任務(wù)而設(shè)計,32 位浮點運算的最大性能為 64 GFLOP。
對于具有 354 個 CPU 的 D1 芯片,32 位浮點運算的最大性能為 22.6 TFLOPs。對于 16 位浮點計算,D1 最大性能躍升至 362 TFLOP。
特斯拉推出了兩個用于神經(jīng)網(wǎng)絡(luò)訓(xùn)練的系統(tǒng):Training Tile 和 ExaPOD。一個訓(xùn)練塊在一個多芯片封裝中具有 25 個連接的 D1 芯片。具有 25 個 D1 芯片的訓(xùn)練塊構(gòu)成 8,850 個訓(xùn)練節(jié)點,每個訓(xùn)練節(jié)點都具有上面總結(jié)的高性能 CPU。對于 32 位浮點計算,訓(xùn)練 tile 的最大性能為 565 TFLOP。
ExaPOD 將 120 個訓(xùn)練塊連接到一個系統(tǒng)中,或者將 3,000 個 D1 芯片與 106.2 萬個訓(xùn)練節(jié)點連接起來。對于 32 位浮點計算,ExaPOD 的最大性能為 67.8 PFLOP。
特斯拉神經(jīng)網(wǎng)絡(luò)公告詳情
D1芯片和Dojo神經(jīng)網(wǎng)絡(luò)訓(xùn)練系統(tǒng)的推出,昭示了特斯拉的方向。將這些產(chǎn)品投入生產(chǎn)的研發(fā)投入無疑是非常高的。特斯拉可能會與其他公司分享這項技術(shù)——以創(chuàng)造另一個類似于出售給其他 OEM 的 BEV 積分的收入來源。
下表列出了特斯拉神經(jīng)網(wǎng)絡(luò)產(chǎn)品公告的特點。數(shù)據(jù)是從 8 月 19 日事件的視頻中提取的。我在幾個地方增加了我對芯片和系統(tǒng)架構(gòu)的理解。
特斯拉的設(shè)計目標(biāo)是在其芯片和系統(tǒng)中擴展三個系統(tǒng)特性:計算性能、計算節(jié)點之間的高帶寬和低延遲通信。高帶寬和低延遲一直難以擴展到數(shù)百或數(shù)千個計算節(jié)點。看起來特斯拉已經(jīng)成功地縮放了以連接的 2D 網(wǎng)格格式組織的所有三個參數(shù)。
訓(xùn)練節(jié)點
訓(xùn)練節(jié)點是D1芯片上最小的訓(xùn)練單元。它有一個 64 位處理器,具有 4 寬標(biāo)量和 4 路多線程程序執(zhí)行。CPU 還具有 2 寬矢量數(shù)據(jù)路徑和 8×8 矢量乘法。
CPU 的指令集架構(gòu) (ISA) 專為機器學(xué)習(xí)和神經(jīng)網(wǎng)絡(luò)訓(xùn)練任務(wù)量身定制。CPU 支持多種浮點格式 — 32 位、16 位和 8 位:FP32、BFP16,以及一種新格式:CFP8 或可配置 FP8。
該處理器具有 1.25MB 高速 SRAM 存儲器,用于存儲程序和數(shù)據(jù)。內(nèi)存使用 ECC 或糾錯碼來提高可靠性。
為了獲得訓(xùn)練節(jié)點之間的低延遲,Tesla 選擇了信號在 2GHz+ 時鐘頻率的一個周期內(nèi)可以傳播的最遠距離。這定義了訓(xùn)練節(jié)點的距離以及 CPU 及其支持電子設(shè)備的復(fù)雜程度。這些參數(shù)還允許 CPU 以每秒 512 Gbit 的速度與四個相鄰的訓(xùn)練節(jié)點進行通信。
訓(xùn)練節(jié)點的最大性能因使用的算法而異。浮點性能通常用于比較。最大訓(xùn)練塊 32 位浮點性能 (FP32) 為 64 GFLOP。BFP16 或 CFP8 算術(shù)的最大性能為 1,024 GFLOP。
D1芯片
令人印象深刻的 Tesla D1 芯片是專門用于神經(jīng)網(wǎng)絡(luò)訓(xùn)練的設(shè)計。D1 采用 7 納米工藝制造,在一個 645 平方毫米的芯片中封裝了 500 億個晶體管。該芯片的電線長度超過 11 英里,功耗在 400 瓦范圍內(nèi)。
D1 芯片有一個帶有高速、低功耗 SerDes 的 I/O 環(huán)——芯片周圍共有 576 個通道。每個通道的傳輸速率為 112 Gbps。最大 D1 片上傳輸速率為 10 Tbps(每秒 10 TB)。芯片每側(cè)的最大板外傳輸速率為 4 Tbps。
D1 芯片上的 354 個 CPU 中的每一個都具有 1.25 MB 的 SRAM,這加起來超過 442 MB 的 SRAM。D1芯片的最大性能也是基于354個訓(xùn)練節(jié)點的CPU陣列。
32 位浮點計算的 D1 最大性能達到 22.6 TFLOP。16 位浮點計算的最大性能為 362 TFLOP。
訓(xùn)練瓷磚
Tesla 的 Training Tile 是擴展 AI 訓(xùn)練系統(tǒng)的基石。Training Tile 將 25 個 D1 裸片集成到晶圓上,并封裝為多芯片模塊 (MCM)。特斯拉認(rèn)為這可能是芯片行業(yè)最大的MCM。訓(xùn)練塊被封裝為一個大芯片,可以通過一個保留訓(xùn)練塊帶寬的高帶寬連接器連接到其他訓(xùn)練塊。
訓(xùn)練塊封裝包括多層電源和控制、電流分配、計算平面(25 個 D1 芯片)和冷卻系統(tǒng)。培訓(xùn)板塊用于 IT 中心,而非自動駕駛汽車。
訓(xùn)練塊為 16 位浮點計算提供 25 倍的單個 D1 芯片或高達 9 Peta FLOP 的性能,為 32 位浮點計算提供高達 565 TFLOP 的性能。
12 個 2x3x2 配置的訓(xùn)練圖塊可以裝在一個機柜中,特斯拉將其稱為訓(xùn)練矩陣。
ExaPOD
特斯拉描述的最大系統(tǒng)是 ExaPOD。它由 120 個訓(xùn)練圖塊組成。這增加了 3000 個 D1 芯片和 106.2 萬個訓(xùn)練節(jié)點。它適合 10 個櫥柜。它顯然是供 IT 中心使用的。
ExaPOD 的最大性能是 16 位浮點計算的 1.09 Exa FLOPs 和 32 位浮點計算的 67.8 Peta FLOPs。
Dojo 軟件和 DPU
Dojo 軟件旨在支持大型和小型神經(jīng)網(wǎng)絡(luò)的訓(xùn)練。Tesla 有一個編譯器來創(chuàng)建利用訓(xùn)練節(jié)點、D1 芯片、訓(xùn)練塊和 ExaPOD 系統(tǒng)的結(jié)構(gòu)和功能的軟件代碼。它使用帶有擴展的 PyTorch 開源機器學(xué)習(xí)庫來利用 D1 芯片和 Dojo 系統(tǒng)架構(gòu)。
這些能力允許對大型神經(jīng)網(wǎng)絡(luò)進行分區(qū)和映射,以提取不同的并行度、模型、圖、數(shù)據(jù)并行度,從而加速大型神經(jīng)網(wǎng)絡(luò)的訓(xùn)練。編譯器使用多種技術(shù)來提取并行性。它可以使用數(shù)據(jù)模型圖并行技術(shù)轉(zhuǎn)換網(wǎng)絡(luò)以實現(xiàn)細粒度并行,并且可以優(yōu)化以減少內(nèi)存占用。
Dojo 接口處理器用于與 IT 和數(shù)據(jù)中心的主機進行通信。它通過上面解釋的高帶寬與 PCIe 4.0 連接到主機和基于 D1 的系統(tǒng)。接口處理器還為 D1 系統(tǒng)提供高帶寬 DRAM 共享內(nèi)存。
基于 D1 的系統(tǒng)可以細分并劃分為稱為 Dojo 處理單元的單元。DPU由一個或多個D1芯片、一個接口處理器和一個或多個計算機主機組成。DPU 虛擬系統(tǒng)可以根據(jù)運行在其上的神經(jīng)網(wǎng)絡(luò)的需要進行放大或縮小。
底線
特斯拉神經(jīng)網(wǎng)絡(luò)訓(xùn)練芯片、系統(tǒng)和軟件讓人印象深刻。有很多創(chuàng)新,例如從芯片到系統(tǒng)保持巨大的帶寬和低延遲。Training Tile 的電源和冷卻包裝看起來很創(chuàng)新。
神經(jīng)網(wǎng)絡(luò)訓(xùn)練系統(tǒng)用于數(shù)據(jù)中心,肯定會用于改進特斯拉的 AV 軟件。其他公司很可能也會使用這些特斯拉神經(jīng)網(wǎng)絡(luò)訓(xùn)練系統(tǒng)。
一個關(guān)鍵問題是神經(jīng)網(wǎng)絡(luò)系統(tǒng)將如何用于自動駕駛汽車的推理應(yīng)用?在當(dāng)前版本中,Training Tile 的功耗看起來太高,無法自動使用。演示中的一張圖片有“15 KW 熱排斥”標(biāo)簽用于訓(xùn)練瓷磚。D1 芯片可能在幻燈片中列出的 400 瓦 TDP 范圍內(nèi)。
看起來特斯拉希望和/或依賴這種神經(jīng)網(wǎng)絡(luò)訓(xùn)練創(chuàng)新,將其 Autopilot 變成支持 L3 或 L4 的系統(tǒng)——僅使用基于攝像頭的傳感器。這是一個好賭注嗎?時間會證明一切,但到目前為止,埃隆馬斯克的大部分賭注都是好的——只是有些延遲。