特斯拉Dojo超算架構(gòu)細(xì)節(jié)首次公開！為自動(dòng)駕駛「操碎了芯」

新智元

+ 關(guān)注

2022-08-30 15:17

583次閱讀

特斯拉Dojo超算架構(gòu)細(xì)節(jié)首次公開！為自動(dòng)駕駛「操碎了芯」

新智元報(bào)道

編輯：Aeneas 好困

【新智元導(dǎo)讀】在剛剛舉辦的硅谷芯片技術(shù)研討會(huì)Hot Chips 34會(huì)議上，備受關(guān)注的特斯拉Dojo超算指令集結(jié)構(gòu)細(xì)節(jié)史上首次被公開。

為了滿足對(duì)人工智能和機(jī)器學(xué)習(xí)模型越來越大的需求，特斯拉創(chuàng)建了自己的人工智能技術(shù)，來教特斯拉的汽車自動(dòng)駕駛。

最近，特斯拉在Hot Chips 34會(huì)議上，披露了大量關(guān)于Dojo（道場(chǎng)）超級(jí)計(jì)算架構(gòu)的細(xì)節(jié)。

本質(zhì)上，Dojo是一個(gè)巨大的可組合的超級(jí)計(jì)算機(jī)，它由一個(gè)完全定制的架構(gòu)構(gòu)建，涵蓋了計(jì)算、網(wǎng)絡(luò)、輸入/輸出（I/O）芯片到指令集架構(gòu)（ISA）、電源傳輸、包裝和冷卻。所有這些都是為了大規(guī)模地運(yùn)行定制的、特定的機(jī)器學(xué)習(xí)訓(xùn)練算法。

Ganesh Venkataramanan是Tesla自動(dòng)駕駛硬件高級(jí)總監(jiān)，負(fù)責(zé)Dojo項(xiàng)目，以及AMD的CPU設(shè)計(jì)團(tuán)隊(duì)。Hot Chips 34會(huì)議上，他和眾位芯片、系統(tǒng)和軟件工程師首次公開了該機(jī)器的許多架構(gòu)特性。

數(shù)據(jù)中心「三明治」

「一般來說，我們制造芯片的過程，是把它們放在包裝上，把包裝放在印刷電路板上，然后進(jìn)入系統(tǒng)。系統(tǒng)進(jìn)入機(jī)架。」Venkataramanan說。

但是這個(gè)過程中存在一個(gè)問題：每次數(shù)據(jù)從芯片移動(dòng)到封裝上并離開封裝時(shí)，都會(huì)產(chǎn)生延遲和帶寬損失。

為了繞過這些限制，Venkataramanan和他的團(tuán)隊(duì)決定從頭開始。

由此，Dojo的訓(xùn)練瓦片誕生了。

這是一個(gè)獨(dú)立的計(jì)算集群，占地半立方英尺，在15千瓦的液冷封裝中能夠達(dá)到556TFLOPS的FP32性能。

每個(gè)瓦片都配備了11GB的SRAM，并在整個(gè)堆棧中使用定制的傳輸協(xié)議，通過9TB/s結(jié)構(gòu)連接。

Venkataramanan說：「這塊訓(xùn)練板代表了從計(jì)算機(jī)到存儲(chǔ)器、到電源傳輸、到通信的無與倫比的集成度，不需要任何額外的開關(guān)。」

訓(xùn)練瓦片的核心是特斯拉的D1，這是一個(gè)500億個(gè)晶體管芯片，基于臺(tái)積電的7納米工藝。特斯拉表示，每個(gè)D1能夠在400W的TDP下實(shí)現(xiàn)22TFLOPS的FP32性能。

特斯拉然后用25個(gè)D1，把它們分到已知的好模具上，然后用臺(tái)積電的晶圓上系統(tǒng)技術(shù)把它們包裝起來，以極低的延遲和極高的帶寬實(shí)現(xiàn)大量的計(jì)算集成。

然而，晶片上的系統(tǒng)設(shè)計(jì)和垂直堆疊架構(gòu)，給電力輸送帶來了挑戰(zhàn)。

據(jù)Venkataramanan說，目前大多數(shù)加速器將電源直接放在硅片旁邊。他解釋說，這種方法雖然行之有效，但這就意味著加速器的很大一部分區(qū)域必須專門用于這些組件，這對(duì)Dojo來說是不切實(shí)際的。于是，特斯拉選擇直接通過芯片底部直接提供電源。

此外，特斯拉還開發(fā)了Dojo接口處理器（DIP），它是主機(jī)CPU和訓(xùn)練處理器之間的橋梁。

每個(gè)DIP都有32GB的HBM，最多可以將五個(gè)這樣的卡以900GB/s的速度連接到一個(gè)訓(xùn)練瓦片上，以達(dá)到4.5TB/s的總量，每個(gè)瓦片共有160GB的HBM。

特斯拉的V1配置成對(duì)的這些瓦片——或150個(gè)D1模具——在陣列中支持四個(gè)主機(jī)CPU，每個(gè)主機(jī)CPU配備五個(gè)DIP卡，以實(shí)現(xiàn)聲稱的BF16或CFP8性能的exaflop。

軟件

這樣一個(gè)專門的計(jì)算架構(gòu)，就需要一個(gè)專門的軟件棧。然而，Venkataramanan和他的團(tuán)隊(duì)認(rèn)識(shí)到，可編程性將決定Dojo的成敗。

「當(dāng)我們?cè)O(shè)計(jì)這些系統(tǒng)時(shí)，軟件同行的易編程性是最重要的。研究人員不會(huì)等待你的軟件人員為適應(yīng)我們想要運(yùn)行的新算法而寫一個(gè)手寫的內(nèi)核。」

為了做到這一點(diǎn)，特斯拉放棄了使用內(nèi)核的想法，圍繞編譯器設(shè)計(jì)了Dojo的架構(gòu)。

「我們的做法是使用PiTorch。我們創(chuàng)建了一個(gè)中間層，它幫助我們并行化，以擴(kuò)展其下面的硬件。所有東西下面都是編譯過的代碼。」為了創(chuàng)建可適應(yīng)任何未來工作負(fù)載的軟件堆棧，這是唯一的方法。

盡管強(qiáng)調(diào)了軟件的靈活性，Venkataramanan指出，目前在他們的實(shí)驗(yàn)室中運(yùn)行的平臺(tái)，暫時(shí)僅限于特斯拉使用。

Dojo架構(gòu)一覽

看完了以上這些，讓我們深入了解一下Dojo的架構(gòu)。

特斯拉擁有用于機(jī)器學(xué)習(xí)的百億億次人工智能級(jí)系統(tǒng)。特斯拉有足夠的資金規(guī)模來雇傭員工，并專門為其應(yīng)用構(gòu)建芯片和系統(tǒng)，就像特斯拉的車載系統(tǒng)一樣。

特斯拉不僅在構(gòu)建自己的AI芯片，還在構(gòu)建超級(jí)計(jì)算機(jī)。

分布式系統(tǒng)分析

Dojo的每個(gè)節(jié)點(diǎn)都有自己的CPU、內(nèi)存和通信接口。

Dojo節(jié)點(diǎn)

這是Dojo處理器的處理管線。

處理管道

每個(gè)節(jié)點(diǎn)有1.25MB的SRAM。在AI訓(xùn)練和推理芯片中，一種常見的技術(shù)是將內(nèi)存與計(jì)算共置，以最大限度地減少數(shù)據(jù)傳輸，因?yàn)閺墓β屎托阅艿慕嵌葋砜矗瑪?shù)據(jù)傳輸非常昂貴。

節(jié)點(diǎn)內(nèi)存

然后每個(gè)節(jié)點(diǎn)都連接到一個(gè)2D網(wǎng)格。

網(wǎng)絡(luò)接口

這是數(shù)據(jù)路徑概述。

數(shù)據(jù)路徑

下面是一個(gè)例子，說明芯片可以做的列表解析。

列表解析

這里有更多關(guān)于指令集的內(nèi)容，屬于特斯拉原創(chuàng)，而不是典型的Intel、Arm、NVIDIA或AMD CPU/GPU的指令集。

指令集

在人工智能中，算術(shù)格式很重要，尤其是芯片支持哪些格式。利用DOJO，特斯拉就可以研究常用格式，例如FP32、FP16和BFP16。這些是常見的行業(yè)格式。

算術(shù)格式

特斯拉也在研究可配置的FP8或CFP8。它有4/3和5/2的范圍選項(xiàng)。這類似于 NVIDIA H100 Hopper配置的FP8。我們還看到Untether.AI Boqueria 1458 RISC-V核心AI加速器專注于不同的FP8類型。

算術(shù)格式 2

Dojo還具有不同的CFP16格式，以實(shí)現(xiàn)更高的精度，并支持FP32、BFP16、CFP8和CFP16。

算術(shù)格式 3

然后將這些核心集成到制造的模具中。特斯拉的D1芯片由臺(tái)積電以7nm工藝制造。每個(gè)芯片有354個(gè)Dojo處理節(jié)點(diǎn)和440MB的SRAM。

First Integration Box D1 模具

這些D1芯片被封裝在一個(gè)道場(chǎng)訓(xùn)練瓦片上。D1芯片經(jīng)過測(cè)試，然后被組裝成一個(gè)5×5的瓦片。這些瓦片每個(gè)邊緣有4.5TB/s的帶寬。它們還具有每個(gè)模塊15kW的功率傳輸包絡(luò)，或者可以說，每個(gè)D1芯片去掉40個(gè)I/O裸片所使用的功率后，大約還有600W。

通過對(duì)比可以看出，如果一家公司不想設(shè)計(jì)這種東西，為什么像Lightmatter Passage會(huì)更有吸引力。

二次集成箱Dojo訓(xùn)練瓦片

Dojo的接口處理器位于2D網(wǎng)格的邊緣。每個(gè)訓(xùn)練塊有11GB的SRAM和160GB的共享DRAM。

Dojo系統(tǒng)拓?fù)?/span>

以下是連接處理節(jié)點(diǎn)的2D網(wǎng)格的帶寬數(shù)據(jù)。

Dojo系統(tǒng)通信邏輯二維網(wǎng)格

每個(gè)DIP和主機(jī)系統(tǒng)提供32GB/s的鏈接。

Dojo系統(tǒng)通信 PCIe鏈接DIP和主機(jī)

特斯拉還具有用于更長(zhǎng)路線的Z平面鏈接。在接下來的演講中，特斯拉談到了系統(tǒng)級(jí)的創(chuàng)新。

通信機(jī)制

這里是die和tiles的延遲邊界，這就是為什么在Dojo中對(duì)它們進(jìn)行不同處理的原因。需要Z平面鏈路的原因是，長(zhǎng)路徑很昂貴。

Dojo系統(tǒng)通信機(jī)制

任何處理節(jié)點(diǎn)都可以跨系統(tǒng)訪問數(shù)據(jù)。每個(gè)節(jié)點(diǎn)都可以將數(shù)據(jù)推送或拉取到SRAM或DRAM。

Dojo系統(tǒng)批量通信

Dojo使用平面尋址方案進(jìn)行通信。

系統(tǒng)網(wǎng)絡(luò)1

這些芯片可以在軟件中繞過錯(cuò)誤的處理節(jié)點(diǎn)。

系統(tǒng)網(wǎng)絡(luò)2

這意味著軟件必須了解系統(tǒng)拓?fù)洹?/span>

系統(tǒng)網(wǎng)絡(luò)3

Dojo不保證端到端的流量排序，因此需要在目的地對(duì)數(shù)據(jù)包進(jìn)行計(jì)數(shù)。

系統(tǒng)網(wǎng)絡(luò)4

以下是數(shù)據(jù)包如何計(jì)入系統(tǒng)同步的一部分。

系統(tǒng)同步

編譯器需要定義一個(gè)帶有節(jié)點(diǎn)的樹。

系統(tǒng)同步2

特斯拉表示，一個(gè)exa-pod擁有超過100萬個(gè)CPU（或計(jì)算節(jié)點(diǎn)）。這些都是大型系統(tǒng)。

總結(jié)

特斯拉專門為大規(guī)模工作而建造了Dojo。通常，初創(chuàng)公司都希望為每個(gè)系統(tǒng)構(gòu)建一個(gè)或幾個(gè)芯片的AI芯片。顯然，特斯拉專注于更大的規(guī)模。

在許多方面，特斯拉擁有一個(gè)巨大的人工智能訓(xùn)練場(chǎng)是合理的。更令人興奮的是，它不僅使用商業(yè)上可用的系統(tǒng)，而且還在構(gòu)建自己的芯片和系統(tǒng)。標(biāo)量方面的一些ISA是借用RISC-V的，但矢量方面和很多架構(gòu)特斯拉都是定制的，所以這需要大量的工作。

參考資料：

https://www.theregister.com/2022/08/24/tesla_supercomputer_dojo/

https://www.servethehome.com/tesla-dojo-ai-system-microarchitecture/