715
Tesla在上月的AI Day推出了用于機(jī)器學(xué)習(xí)和神經(jīng)網(wǎng)絡(luò)訓(xùn)練的芯片、系統(tǒng)和軟件。它們將共同推進(jìn)用于自動(dòng)駕駛汽車模型的訓(xùn)練。
Elon Musk和他的芯片與系統(tǒng)設(shè)計(jì)團(tuán)隊(duì)在超過三小時(shí)的演講中提供了大量的技術(shù)細(xì)節(jié)，以下是其中的一些亮點(diǎn)。
神經(jīng)網(wǎng)絡(luò)
Tesla設(shè)計(jì)了一個(gè)靈活、可擴(kuò)展的分布式計(jì)算機(jī)架構(gòu)，專門用于神經(jīng)網(wǎng)絡(luò)訓(xùn)練。Tesla的架構(gòu)從D1芯片開始，該芯片有354個(gè)訓(xùn)練節(jié)點(diǎn)，每個(gè)節(jié)點(diǎn)有一個(gè)強(qiáng)大的CPU。這些訓(xùn)練節(jié)點(diǎn)CPU是為高性能的NN和ML任務(wù)設(shè)計(jì)的，32位浮點(diǎn)運(yùn)算的最大性能為64GFLOPs。
對(duì)于擁有354個(gè)CPU的D1芯片，32位浮點(diǎn)運(yùn)算的最大性能是22.6TFLOPs。對(duì)于16位浮點(diǎn)計(jì)算，D1的最大性能躍升至362TFLOPs。
Tesla推出了兩個(gè)用于神經(jīng)網(wǎng)絡(luò)訓(xùn)練的系統(tǒng)：Training Tile和ExaPOD。一個(gè)Training Tile在一個(gè)封裝中擁有25個(gè)連接的D1芯片。一個(gè)由25個(gè)D1芯片構(gòu)成的Training Tile有8850個(gè)訓(xùn)練節(jié)點(diǎn)，每個(gè)節(jié)點(diǎn)都有上面提及到的高性能CPU。一個(gè)Training Tile的32位浮點(diǎn)運(yùn)算最大性能是565TFLOPs。
ExaPOD將120個(gè)Training Tile連接成一個(gè)系統(tǒng)，即3000個(gè)D1芯片、106.2萬個(gè)訓(xùn)練節(jié)點(diǎn)。一個(gè)ExaPOD的32位浮點(diǎn)運(yùn)算最大性能是67.8PFLOPs。
Tesla對(duì)神經(jīng)網(wǎng)絡(luò)介紹的細(xì)節(jié)D1芯片和Dojo神經(jīng)網(wǎng)絡(luò)訓(xùn)練系統(tǒng)的推出顯示出了Tesla的發(fā)展方向。這些產(chǎn)品投入使用的研發(fā)投資無疑是非常高的。Tesla很可能與其他公司分享這項(xiàng)技術(shù)，創(chuàng)造另一個(gè)收入來源，類似于賣給其他OEM的電動(dòng)車積分。
下表列出了Tesla的神經(jīng)網(wǎng)絡(luò)的特點(diǎn)。這些數(shù)據(jù)是從當(dāng)天視頻中提取的。筆者在一些地方加入了對(duì)芯片和系統(tǒng)架構(gòu)的理解。
Tesla的設(shè)計(jì)目標(biāo)是在其芯片和系統(tǒng)中擴(kuò)展三個(gè)系統(tǒng)特性：計(jì)算性能、高帶寬和計(jì)算節(jié)點(diǎn)之間的低延遲通信。高帶寬和低延遲一直以來都很難擴(kuò)展到數(shù)百或數(shù)千的計(jì)算節(jié)點(diǎn)。看起來Tesla已經(jīng)成功地將所有三個(gè)參數(shù)組織在一個(gè)連接的二維網(wǎng)格格式中。
訓(xùn)練節(jié)點(diǎn)
訓(xùn)練節(jié)點(diǎn)是D1芯片上最小的訓(xùn)練單元。它有一個(gè)64位處理器，具有4寬標(biāo)量和4路多線程的程序執(zhí)行。該CPU還具有8*8向量乘法的2寬向量數(shù)據(jù)路徑。
該CPU的指令集架構(gòu)（ISA）是為機(jī)器學(xué)習(xí)和神經(jīng)網(wǎng)絡(luò)訓(xùn)練任務(wù)量身定制的。該CPU支持多種浮點(diǎn)格式，32位、16位和8位。FP32、BFP16，以及一種新的格式CFP8或可配置的FP8。
該處理器有1.25MB的高速SRAM用于程序和數(shù)據(jù)存儲(chǔ)。該存儲(chǔ)器使用ECC或糾錯(cuò)代碼以提高可靠性。
為了降低訓(xùn)練節(jié)點(diǎn)之間的延遲，Tesla挑選了信號(hào)在2GHz+時(shí)鐘頻率的一個(gè)周期內(nèi)可傳播的最遠(yuǎn)距離。這定義了訓(xùn)練節(jié)點(diǎn)的距離，以及CPU及其支持電子設(shè)備的復(fù)雜程度。這些參數(shù)還允許一個(gè)CPU以每秒512G比特的速度與四個(gè)相鄰的訓(xùn)練節(jié)點(diǎn)通信。
訓(xùn)練節(jié)點(diǎn)的最大性能取決于使用的算法，常以浮點(diǎn)性能來比較。Training Tile的32位浮點(diǎn)性能（FP32）的最大性能是64GFLOPs。BFP16或CFP8算術(shù)的最大性能是1,024GFLOPs。
D1芯片
令人印象深刻的D1芯片是一種專用于神經(jīng)網(wǎng)絡(luò)訓(xùn)練的設(shè)計(jì)。D1采用7nm工藝制造，在一個(gè)面積為645平方毫米的裸片中封裝了500億個(gè)晶體管。該芯片電線長(zhǎng)度超過11英里，功耗為400W左右。
D1芯片有一個(gè)帶有高速、低功耗SerDes的I/O環(huán)，總共有576條通道環(huán)繞著芯片。每條通道的傳輸速率為112Gbps。D1的最大片上傳輸速率為10Tbps（每秒10兆位）。芯片每側(cè)的最大片外傳輸速率為4Tbps。
由于D1芯片上的354個(gè)CPU中的每一個(gè)都有1.25MB的SRAM，加起來就有超過442MB的SRAM。D1芯片的最大性能也是基于354個(gè)訓(xùn)練節(jié)點(diǎn)的CPU陣列。
D1的32位浮點(diǎn)計(jì)算的最大性能達(dá)到22.6TFLOPs。16位浮點(diǎn)計(jì)算的最大性能為362TFLOPs。
Training Tile
Tesla的Training Tile是擴(kuò)展AI訓(xùn)練系統(tǒng)的基石。一個(gè)Training Tile將25個(gè)D1芯片集成到一個(gè)晶圓上，并被封裝成多芯片模塊（MCM）。Tesla認(rèn)為這可能是芯片行業(yè)中最大的MCM。Training Tile被封裝成一個(gè)大芯片，可以通過一個(gè)高帶寬連接器與其他Training Tile連接，保留了Training Tile的帶寬。
Training Tile的封裝包括多層電源和控制、電流分配、計(jì)算平面（25個(gè)D1芯片）和冷卻系統(tǒng)。Training Tile用于IT中心，而不是自動(dòng)駕駛汽車。
Training Tile提供單個(gè)D1芯片的25倍性能，或16位浮點(diǎn)計(jì)算的9Peta FLOPs，32位浮點(diǎn)計(jì)算則高達(dá)565TFLOPs。
12塊232配置的Training Tile可以裝在一個(gè)柜子里，Tesla稱它為Training Matrix。
ExaPOD
Tesla描述的最大系統(tǒng)是ExaPOD。它是由120個(gè)Training Tile構(gòu)成的。加起來有3000個(gè)D1芯片和106.2萬個(gè)訓(xùn)練節(jié)點(diǎn)。它裝在10個(gè)機(jī)柜中，顯然是為IT中心使用的。
ExaPOD的最大性能是16位浮點(diǎn)計(jì)算的1.09Exa FLOPs，32位浮點(diǎn)計(jì)算的67.8Peta FLOPs。
Dojo軟件和DPU
Dojo軟件旨在支持大型和小型神經(jīng)網(wǎng)絡(luò)的訓(xùn)練。特斯拉有一個(gè)編譯器來創(chuàng)建軟件代碼，利用訓(xùn)練節(jié)點(diǎn)、D1芯片、Training Tile和ExaPOD系統(tǒng)的結(jié)構(gòu)和能力。它使用的是PyTorch開源機(jī)器學(xué)習(xí)庫，并進(jìn)行了擴(kuò)展來利用D1芯片和Dojo系統(tǒng)架構(gòu)。
這些能力允許大型神經(jīng)網(wǎng)絡(luò)被分割和映射，以提取不同的并行性、模型、圖形、數(shù)據(jù)的并行度，從而加速大型神經(jīng)網(wǎng)絡(luò)的訓(xùn)練。編譯器使用多種技術(shù)來提取并行性。它可以利用數(shù)據(jù)模型圖并行技術(shù)對(duì)網(wǎng)絡(luò)進(jìn)行轉(zhuǎn)換以實(shí)現(xiàn)細(xì)粒度并行，并可以優(yōu)化以減少內(nèi)存占用。
Dojo接口處理器用于與IT和數(shù)據(jù)中心的主機(jī)通信。它用PCIe 4.0連接到主機(jī)，并通過上述的高帶寬連接到基于D1的系統(tǒng)。接口處理器還為D1系統(tǒng)提供高帶寬的DRAM共享內(nèi)存。
基于D1的系統(tǒng)可以被細(xì)分和劃分為稱為Dojo Processing Unit的單元。DPU由一個(gè)或多個(gè)D1芯片、一個(gè)接口處理器和一個(gè)或多個(gè)計(jì)算機(jī)主機(jī)組成。DPU虛擬系統(tǒng)可以根據(jù)在其上運(yùn)行的神經(jīng)網(wǎng)絡(luò)的需要，擴(kuò)大或縮小規(guī)模。
總結(jié)
Tesla的神經(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ù)中心的，肯定會(huì)被用于改進(jìn)Tesla的AV軟件。其他公司很可能也會(huì)使用這些Tesla神經(jīng)網(wǎng)絡(luò)訓(xùn)練系統(tǒng)。
看起來Tesla希望或依賴這種神經(jīng)網(wǎng)絡(luò)訓(xùn)練創(chuàng)新，繼續(xù)使其基于純視覺的自動(dòng)駕駛系統(tǒng)持續(xù)改進(jìn)。這是正確的方向嗎？時(shí)間會(huì)告訴我們，到目前為止，Elon Musk的大部分賭注盡管都有一些時(shí)間上的推遲，但都是正確的。
[參考文章]
Tesla AI Day Perspectives — Egil Juliussen
