特斯拉在 8 月中旬的 AI 日推出了用于機(jī)器學(xué)習(xí)和神經(jīng)網(wǎng)絡(luò)訓(xùn)練的汽車(chē)芯片、系統(tǒng)和軟件。他們將共同推進(jìn)自動(dòng)駕駛汽車(chē)模型的培訓(xùn)。

埃隆馬斯克和他的芯片和系統(tǒng)設(shè)計(jì)師團(tuán)隊(duì)在三個(gè)多小時(shí)的演示中提供了大量技術(shù)細(xì)節(jié)。以下是重點(diǎn)。

神經(jīng)網(wǎng)絡(luò)

特斯拉設(shè)計(jì)了一種靈活、可擴(kuò)展的分布式計(jì)算機(jī)架構(gòu)，專(zhuān)為神經(jīng)網(wǎng)絡(luò)訓(xùn)練量身定做。Tesla 的架構(gòu)始于具有 354 個(gè)訓(xùn)練節(jié)點(diǎn)的 D1 專(zhuān)用芯片，每個(gè)訓(xùn)練節(jié)點(diǎn)都有一個(gè)強(qiáng)大的 CPU。這些訓(xùn)練節(jié)點(diǎn) CPU 專(zhuān)為高性能 NN 和 ML 任務(wù)而設(shè)計(jì)，32 位浮點(diǎn)運(yùn)算的最大性能為 64 GFLOP。

對(duì)于具有 354 個(gè) CPU 的 D1 芯片，32 位浮點(diǎn)運(yùn)算的最大性能為 22.6 TFLOPs。對(duì)于 16 位浮點(diǎn)計(jì)算，D1 最大性能躍升至 362 TFLOP。

特斯拉推出了兩個(gè)用于神經(jīng)網(wǎng)絡(luò)訓(xùn)練的系統(tǒng)：Training Tile 和 ExaPOD。一個(gè)訓(xùn)練塊在一個(gè)多芯片封裝中具有 25 個(gè)連接的 D1 芯片。具有 25 個(gè) D1 芯片的訓(xùn)練塊構(gòu)成 8,850 個(gè)訓(xùn)練節(jié)點(diǎn)，每個(gè)訓(xùn)練節(jié)點(diǎn)都具有上面總結(jié)的高性能 CPU。對(duì)于 32 位浮點(diǎn)計(jì)算，訓(xùn)練 tile 的最大性能為 565 TFLOP。

ExaPOD 將 120 個(gè)訓(xùn)練塊連接到一個(gè)系統(tǒng)中，或者將 3,000 個(gè) D1 芯片與 106.2 萬(wàn)個(gè)訓(xùn)練節(jié)點(diǎn)連接起來(lái)。對(duì)于 32 位浮點(diǎn)計(jì)算，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ā)投入無(wú)疑是非常高的。特斯拉可能會(huì)與其他公司分享這項(xiàng)技術(shù)——以創(chuàng)造另一個(gè)類(lèi)似于出售給其他 OEM 的 BEV 積分的收入來(lái)源。

下表列出了特斯拉神經(jīng)網(wǎng)絡(luò)產(chǎn)品公告的特點(diǎn)。數(shù)據(jù)是從 8 月 19 日事件的視頻中提取的。我在幾個(gè)地方增加了我對(duì)芯片和系統(tǒng)架構(gòu)的理解。

特斯拉的設(shè)計(jì)目標(biāo)是在其芯片和系統(tǒng)中擴(kuò)展三個(gè)系統(tǒng)特性：計(jì)算性能、計(jì)算節(jié)點(diǎn)之間的高帶寬和低延遲通信。高帶寬和低延遲一直難以擴(kuò)展到數(shù)百或數(shù)千個(gè)計(jì)算節(jié)點(diǎn)?？雌饋?lái)特斯拉已經(jīng)成功地縮放了以連接的 2D 網(wǎng)格格式組織的所有三個(gè)參數(shù)。

訓(xùn)練節(jié)點(diǎn)

訓(xùn)練節(jié)點(diǎn)是D1芯片上最小的訓(xùn)練單元。它有一個(gè) 64 位處理器，具有 4 寬標(biāo)量和 4 路多線(xiàn)程程序執(zhí)行。CPU 還具有 2 寬矢量數(shù)據(jù)路徑和 8×8 矢量乘法。

CPU 的指令集架構(gòu) (ISA) 專(zhuān)為機(jī)器學(xué)習(xí)和神經(jīng)網(wǎng)絡(luò)訓(xùn)練任務(wù)量身定制。CPU 支持多種浮點(diǎn)格式 — 32 位、16 位和 8 位：FP32、BFP16，以及一種新格式：CFP8 或可配置 FP8。

該處理器具有 1.25MB 高速 SRAM 存儲(chǔ)器，用于存儲(chǔ)程序和數(shù)據(jù)。內(nèi)存使用 ECC 或糾錯(cuò)碼來(lái)提高可靠性。

為了獲得訓(xùn)練節(jié)點(diǎn)之間的低延遲，Tesla 選擇了信號(hào)在 2GHz+ 時(shí)鐘頻率的一個(gè)周期內(nèi)可以傳播的最遠(yuǎn)距離。這定義了訓(xùn)練節(jié)點(diǎn)的距離以及 CPU 及其支持電子設(shè)備的復(fù)雜程度。這些參數(shù)還允許 CPU 以每秒 512 Gbit 的速度與四個(gè)相鄰的訓(xùn)練節(jié)點(diǎn)進(jìn)行通信。

訓(xùn)練節(jié)點(diǎn)的最大性能因使用的算法而異。浮點(diǎn)性能通常用于比較。最大訓(xùn)練塊 32 位浮點(diǎn)性能 (FP32) 為 64 GFLOP。BFP16 或 CFP8 算術(shù)的最大性能為 1,024 GFLOP。

D1芯片

令人印象深刻的 Tesla D1 芯片是專(zhuān)門(mén)用于神經(jīng)網(wǎng)絡(luò)訓(xùn)練的設(shè)計(jì)。D1 采用 7 納米工藝制造，在一個(gè) 645 平方毫米的芯片中封裝了 500 億個(gè)晶體管。該芯片的電線(xiàn)長(zhǎng)度超過(guò) 11 英里，功耗在 400 瓦范圍內(nèi)。

D1 芯片有一個(gè)帶有高速、低功耗 SerDes 的 I/O 環(huán)——芯片周?chē)灿?576 個(gè)通道。每個(gè)通道的傳輸速率為 112 Gbps。最大 D1 片上傳輸速率為 10 Tbps(每秒 10 TB)。芯片每側(cè)的最大板外傳輸速率為 4 Tbps。

D1 芯片上的 354 個(gè) CPU 中的每一個(gè)都具有 1.25 MB 的 SRAM，這加起來(lái)超過(guò) 442 MB 的 SRAM。D1芯片的最大性能也是基于354個(gè)訓(xùn)練節(jié)點(diǎn)的CPU陣列。

32 位浮點(diǎn)計(jì)算的 D1 最大性能達(dá)到 22.6 TFLOP。16 位浮點(diǎn)計(jì)算的最大性能為 362 TFLOP。

訓(xùn)練瓷磚

Tesla 的 Training Tile 是擴(kuò)展 AI 訓(xùn)練系統(tǒng)的基石。Training Tile 將 25 個(gè) D1 裸片集成到晶圓上，并封裝為多芯片模塊 (MCM)。特斯拉認(rèn)為這可能是芯片行業(yè)最大的MCM。訓(xùn)練塊被封裝為一個(gè)大芯片，可以通過(guò)一個(gè)保留訓(xùn)練塊帶寬的高帶寬連接器連接到其他訓(xùn)練塊。

訓(xùn)練塊封裝包括多層電源和控制、電流分配、計(jì)算平面(25 個(gè) D1 芯片)和冷卻系統(tǒng)。培訓(xùn)板塊用于 IT 中心，而非自動(dòng)駕駛汽車(chē)。

訓(xùn)練塊為 16 位浮點(diǎn)計(jì)算提供 25 倍的單個(gè) D1 芯片或高達(dá) 9 Peta FLOP 的性能，為 32 位浮點(diǎn)計(jì)算提供高達(dá) 565 TFLOP 的性能。

12 個(gè) 2x3x2 配置的訓(xùn)練圖塊可以裝在一個(gè)機(jī)柜中，特斯拉將其稱(chēng)為訓(xùn)練矩陣。

ExaPOD

特斯拉描述的最大系統(tǒng)是 ExaPOD。它由 120 個(gè)訓(xùn)練圖塊組成。這增加了 3000 個(gè) D1 芯片和 106.2 萬(wàn)個(gè)訓(xùn)練節(jié)點(diǎn)。它適合 10 個(gè)櫥柜。它顯然是供 IT 中心使用的。

ExaPOD 的最大性能是 16 位浮點(diǎn)計(jì)算的 1.09 Exa FLOPs 和 32 位浮點(diǎn)計(jì)算的 67.8 Peta FLOPs。

Dojo 軟件和 DPU

Dojo 軟件旨在支持大型和小型神經(jīng)網(wǎng)絡(luò)的訓(xùn)練。Tesla 有一個(gè)編譯器來(lái)創(chuàng)建利用訓(xùn)練節(jié)點(diǎn)、D1 芯片、訓(xùn)練塊和 ExaPOD 系統(tǒng)的結(jié)構(gòu)和功能的軟件代碼。它使用帶有擴(kuò)展的 PyTorch 開(kāi)源機(jī)器學(xué)習(xí)庫(kù)來(lái)利用 D1 芯片和 Dojo 系統(tǒng)架構(gòu)。

這些能力允許對(duì)大型神經(jīng)網(wǎng)絡(luò)進(jìn)行分區(qū)和映射，以提取不同的并行度、模型、圖、數(shù)據(jù)并行度，從而加速大型神經(jīng)網(wǎng)絡(luò)的訓(xùn)練。編譯器使用多種技術(shù)來(lái)提取并行性。它可以使用數(shù)據(jù)模型圖并行技術(shù)轉(zhuǎn)換網(wǎng)絡(luò)以實(shí)現(xiàn)細(xì)粒度并行，并且可以?xún)?yōu)化以減少內(nèi)存占用。

Dojo 接口處理器用于與 IT 和數(shù)據(jù)中心的主機(jī)進(jìn)行通信。它通過(guò)上面解釋的高帶寬與 PCIe 4.0 連接到主機(jī)和基于 D1 的系統(tǒng)。接口處理器還為 D1 系統(tǒng)提供高帶寬 DRAM 共享內(nèi)存。

基于 D1 的系統(tǒng)可以細(xì)分并劃分為稱(chēng)為 Dojo 處理單元的單元。DPU由一個(gè)或多個(gè)D1芯片、一個(gè)接口處理器和一個(gè)或多個(gè)計(jì)算機(jī)主機(jī)組成。DPU 虛擬系統(tǒng)可以根據(jù)運(yùn)行在其上的神經(jīng)網(wǎng)絡(luò)的需要進(jìn)行放大或縮小。

底線(xiàn)

特斯拉神經(jīng)網(wǎng)絡(luò)訓(xùn)練芯片、系統(tǒng)和軟件讓人印象深刻。有很多創(chuàng)新，例如從芯片到系統(tǒng)保持巨大的帶寬和低延遲。Training Tile 的電源和冷卻包裝看起來(lái)很創(chuàng)新。

神經(jīng)網(wǎng)絡(luò)訓(xùn)練系統(tǒng)用于數(shù)據(jù)中心，肯定會(huì)用于改進(jìn)特斯拉的 AV 軟件。其他公司很可能也會(huì)使用這些特斯拉神經(jīng)網(wǎng)絡(luò)訓(xùn)練系統(tǒng)。

一個(gè)關(guān)鍵問(wèn)題是神經(jīng)網(wǎng)絡(luò)系統(tǒng)將如何用于自動(dòng)駕駛汽車(chē)的推理應(yīng)用?在當(dāng)前版本中，Training Tile 的功耗看起來(lái)太高，無(wú)法自動(dòng)使用。演示中的一張圖片有“15 KW 熱排斥”標(biāo)簽用于訓(xùn)練瓷磚。D1 芯片可能在幻燈片中列出的 400 瓦 TDP 范圍內(nèi)。

看起來(lái)特斯拉希望和/或依賴(lài)這種神經(jīng)網(wǎng)絡(luò)訓(xùn)練創(chuàng)新，將其 Autopilot 變成支持 L3 或 L4 的系統(tǒng)——僅使用基于攝像頭的傳感器。這是一個(gè)好賭注嗎?時(shí)間會(huì)證明一切，但到目前為止，埃隆馬斯克的大部分賭注都是好的——只是有些延遲。