特斯拉Dojo架構大解析——道場微架構 - Hot Chips 34
Hot Chips 34 - 特斯拉的道場微架構
我是斜杠青年,一個熱愛前沿科技的“雜食性”學者!
說特斯拉只對機器學習感興趣是輕描淡寫的。這家電動汽車制造商制造了一臺名為Dojo的內部超級計算機,專為訓練其機器學習模型進行了優(yōu)化。與許多其他超級計算機不同,Dojo沒有使用現(xiàn)成的CPU和GPU,例如來自AMD、Intel或Nvidia。相反,特斯拉根據(jù)他們的需求設計了自己的微架構,讓他們做出更一般的架構無法做出的權衡。在本文中,我們將根據(jù)特斯拉在Hot Chips的演示文稿來了解該架構。該架構沒有單獨的名稱,因此為了簡單起見,每當我們進一步提到Dojo時,我們都會談論架構。
方框圖
在高水平上,Dojo是一個八寬核心,具有四向SMT,以保守2 GHz運行。它具有CPU風格的管道,使其比GPU更能容忍不同的算法和分支代碼。Dojo的指令集在標量方面類似于RISC-V,但特斯拉的工程師添加了一套專注于加速機器學習的自定義矢量指令。

特斯拉將道場描述為“高吞吐量、通用CPU”。從績效的角度來看,這當然是有道理的。但為了增加計算密度,特斯拉做出了犧牲,與我們在臺式機、筆記本電腦和智能手機中熟悉的CPU相比,Dojo內核將非常難以使用。在某些方面,Dojo核心處理的更像IBM Cell中的SPE,而不是傳統(tǒng)的通用CPU核心。
像細胞的SPE?
在21世紀初中期,IBM的Cell處理器具有八個“辛能處理元件”(SPE),由功能齊全的CPU核心(“電源處理元件”或PPE)控制。乍一看,道場與SPE有很多相似之處。Dojo和SPE都針對矢量處理進行了優(yōu)化,并依賴于單獨的主機處理器進行工作分配。在Dojo或SPE上運行的代碼無法直接訪問系統(tǒng)內存。相反,應用程序預計將主要在一小塊本地SRAM中工作。此本地SRAM由軟件管理,不能用作緩存。如果需要來自主內存的數(shù)據(jù),則必須使用DMA操作將其帶入。最后,Dojo和Cell的SPE都缺乏虛擬內存支持。我們稍后會更多地討論這意味著什么,但簡而言之,它使多任務處理變得非常困難。

Dojo在幾個重要方面有所不同。由于Dojo的設計沒有考慮到小規(guī)模部署,因此主機處理器駐留在單獨的主機系統(tǒng)上。這些主機系統(tǒng)具有帶有接口處理器的PCIe卡,然后通過高速網(wǎng)絡鏈路連接到Dojo芯片。相比之下,Cell的主機處理器位于同一個模具上。這使得單獨部署單個Cell芯片成為可能——這在Dojo中是不可能的。Dojo的1.25 MB本地塊是SRAM比Cell SPE大得多,帶寬也比Cell SPE高得多,Cell SPE只有256 KB的SRAM。Cell的256 KB SRAM只有一個端口,每個周期可以交付128B。Dojo的SRAM有五個64B端口。當然,設計目標非常不同。Dojo是寬而低時鐘的,而Cell SPE具有為高時鐘設計的狹窄、深的管道。但有趣的是,Cell的高級設計概念今天仍然具有現(xiàn)實意義。
前端:CPU舒適度等
讓我們從前端開始,簡單地瀏覽一下道場的管道。有一種分支預測器,因為特斯拉的圖表顯示了BTB(分支目標緩沖區(qū))。它的預測能力可能無法接近我們在AMD、ARM和英特爾高性能內核上看到的,因為Dojo需要優(yōu)先將模具區(qū)域用于矢量執(zhí)行。但即使是基本的分支預測器,也是從根本沒有預測器邁出的一大步,Dojo的分支預測能力在處理分支代碼或更大的指令足跡時,它應該比GPU提供更好的性能。
一旦分支預測器生成了下一個指令獲取指針,Dojo可以將每個周期32字節(jié)從“小”指令緩存拉入每個線程獲取緩沖區(qū)。此指令緩存可能有助于減少本地SRAM的指令帶寬壓力,確保數(shù)據(jù)端可以以盡可能少的爭用訪問SRAM。此外,指令緩存不連貫。如果新代碼加載到本地SRAM中,則在分支到該新代碼之前必須刷新指令緩存。
從獲取緩沖區(qū),Dojo的解碼器每個周期可以處理來自兩個線程的八條指令。我對“每個周期兩個線程”的含義有點困惑,因為具有SMT的CPU通常每個周期獲取和解碼處理一個線程,并在循環(huán)邊界上的線程之間切換。也許Dojo將解碼器分成兩個集群,并選擇兩個線程在每個周期喂養(yǎng)它們。這可能會減少從被提取的分支中解碼吞吐量損失。
在解碼時,某些指令,如分支、謂詞操作和即時加載(“列表解析”)可以在前端執(zhí)行并從管道中刪除。這有點像較新的x86 CPU,消除了重命名器中的寄存器到注冊副本。道場沒有通過管道跟蹤“消除”指令,以保持有序的退休。其他處理器將所有內容跟蹤到退休,以便他們可以在任何指令邊界停止,并保持恢復執(zhí)行所需的所有狀態(tài)。這種功能被稱為“精確異?!?,被現(xiàn)代操作系統(tǒng)用于交付各種好東西,例如分頁到磁盤,或確切地告訴你代碼在哪里搞砸了。
特斯拉不在乎提供精確的例外情況。Dojo確實具有調試模式,通過管道中有更多的指令來提供“更精確”的異常,但沒有像正常的無序CPU那樣的有序退出邏輯。
Dojo的執(zhí)行引擎
在看到寬大的前端后,看到一個只有兩個ALU和兩個AGU的4寬整數(shù)執(zhí)行引擎可能會有點奇怪。但這個漏斗形狀的管道是有道理的,因為一些指令被執(zhí)行并丟棄在前端。Dojo也沒有進入客戶端系統(tǒng),在客戶端系統(tǒng)中,標量整數(shù)性能很重要。因此,整數(shù)邊提供了足夠的吞吐量來壓縮控制流和地址生成,以保持矢量和矩陣單元的進給。
Dojo的矢量和矩陣執(zhí)行側位于核心管道中的標量執(zhí)行引擎之后,并有兩個執(zhí)行管道。兩根管子聽起來可能不太像,但道場在這些管道后面有非常寬的執(zhí)行單元。一個管道可以執(zhí)行512位矢量執(zhí)行,而另一個管道可以執(zhí)行8x8x4矩陣乘法。因此,只要指令暴露足夠的顯式并行性,Dojo就能夠實現(xiàn)非常高的吞吐量——特別是在使用矩陣單元的情況下。特斯拉聲稱,一個擁有354個道場核心的模具可以在2 GHz時達到362 BF16 TFLOPS,這表明每個核心每個周期可以達到512 BF16 FLOPS。
我們不確定Dojo是否可以進行完全的無序執(zhí)行。但特斯拉確實表示,整數(shù)側可以遠遠領先于矢量端,這表明它可以執(zhí)行過去停滯不前的指令,直到其中一個調度器填滿。缺乏有序的退休也表明執(zhí)行能力不合序。
通常,實現(xiàn)無序執(zhí)行會帶來很多復雜性。這是因為CPU必須看起來按順序執(zhí)行指令。AMD、ARM和英特爾的高性能CPU使用大型重新排序緩沖區(qū)(和其他結構)跟蹤指令,以便其結果可以按程序順序提交。這意味著,如果一個程序做了一些愚蠢的事情,比如除以零,這些核心可以準確顯示哪個指令是錯誤的。而且,他們可以顯示CPU狀態(tài),該狀態(tài)反映了故障前所有指令中的操作,但之后沒有。這意味著你可以修復任何使指令出現(xiàn)故障的東西,并恢復執(zhí)行。Dojo放棄了這項能力。作為交換,Dojo可以避免與通過管道跟蹤每條指令相關的電源和區(qū)域開銷,以確保結果按程序順序提交。
SRAM訪問
通常,我們會在這里討論緩存。但Dojo無法直接訪問系統(tǒng)內存,因此我們將討論1.25 MB的SRAM塊。它每個周期可以處理兩個512位負載,使其每個周期的帶寬與支持AVX-512的英特爾CPU相匹配。特斯拉表示,SRAM有五個512位端口(2個負載,1個存儲端口,兩個到網(wǎng)格停止)。但標量端只有兩個AGU,這可能意味著核心無法在每個周期都承受兩個512位負載和一個512位存儲。
由于Dojo的本地SRAM塊不是緩存,因此它不需要與數(shù)據(jù)一起存儲的標簽和狀態(tài)位。SRAM前面也沒有L1D緩存,因此它必須足夠快,才能處理所有加載和存儲指令,而不會造成瓶頸,即使其大小更接近L2緩存。不將SRAM實現(xiàn)為緩存可能是特斯拉保持低延遲的方式。如果我們回顧很久以前的AMD錘子架構,我們可以看到,在檢測到L1D丟失后,訪問1 MB L2需要8個周期(對于12個總延遲周期)。如果直接解決1 MB的SRAM而不是用作緩存,則可能會刪除三個階段,使延遲減少到5個周期:

考慮到幾十年來的過程節(jié)點改進,甚至比Athlon更低的時鐘速度目標,很容易看到特斯拉如何訪問具有類似L1延遲的L2大小的SRAM塊。當然,跳過一個級別的緩存可以節(jié)省面積和電力。
為了進一步降低延遲、面積和核心復雜性,Dojo沒有虛擬內存支持。因此,它沒有TLB或頁面行走機制?,F(xiàn)代操作系統(tǒng)利用虛擬內存為每個進程提供自己的內存視圖。程序使用的內存地址不是直接訪問物理內存地址,而是由CPU使用操作系統(tǒng)設置的分頁結構將它轉換為物理地址。這就是現(xiàn)代操作系統(tǒng)如何將程序彼此隔離,并防止一個行為不端的應用程序導致整個系統(tǒng)崩潰的方法。
虛擬內存也是你運行比物理內存更多的程序的方式。當你用完真實內存時,操作系統(tǒng)會解對映頁面,將其寫入磁盤,并向你的程序提供它所需的內存。當另一個糟糕的程序試圖訪問該內存時,CPU試圖將虛擬地址轉換為物理地址,但發(fā)現(xiàn)翻譯不存在。CPU拋出頁面故障異常,操作系統(tǒng)通過將被驅逐的頁面讀回物理內存并填寫頁表條目來處理該異常。
這些在道場上都不可能做到。核心的4向SMT功能更適合單個應用程序公開顯式并行性,而不是提高多任務處理性能。例如,一個線程可以執(zhí)行矢量計算,而另一個線程將數(shù)據(jù)從系統(tǒng)內存異步加載到SRAM(通過DMA)。
為了進一步簡化設計,Dojo只需21個地址位即可尋址SRAM,這可以簡化Dojo的AGU和尋址總線。這些權衡可能允許特斯拉以足夠低的延遲訪問該SRAM,以避免在其面前實現(xiàn)單獨的L1數(shù)據(jù)緩存。
DRAM訪問
說到系統(tǒng)內存,Dojo芯片不會直接連接到內存。相反,它們連接到配備HBM的接口處理器。這些接口處理器還負責與主機系統(tǒng)通信。

帶有25個獨立芯片的Dojo瓷磚可以訪問160 GB的HBM內存。特斯拉表示,他們可以跨越瓷磚邊界從每個模具邊緣傳輸900 GB/s,這意味著可以使用4.5 TB/s的鏈路帶寬訪問接口處理器及其HBM。由于訪問HBM需要通過單獨的芯片,因此訪問延遲可能非常高。
Dojo它很小。你覺得如何?
Dojo是一個8寬的核心,具有至少一些OoO執(zhí)行能力、體面的矢量吞吐量和矩陣乘法單元。但即使有1.25 MB的本地SRAM,它最終也是一個非常小的核心。相比之下,富士通的A64FX占據(jù)了同一工藝節(jié)點的兩倍多區(qū)域。
處理器設計就是做出正確的權衡。特斯拉希望通過將大量核心包裝到模具上來最大限度地提高機器學習的吞吐量,因此單個內核必須很小。為了實現(xiàn)其區(qū)域效率,道場使用了一些熟悉的技術。它以保守的2 GHz運行。較低的時鐘電路往往占用更少的面積。它可能有一個基本的分支預測器和一個小的指令緩存。如果程序具有較大的代碼足跡或許多分支,則會犧牲一些性能。
但特斯拉還通過削減運行內部工作負載不需要的功能,在盡量減少電力和面積使用方面也走得更遠。他們不進行數(shù)據(jù)端緩存,不支持虛擬內存,也不支持精確的異常。

其結果是一個處理器核心,提供了現(xiàn)代CPU核心的性能靈活性,同時在許多方面不如英特爾8086對用戶和程序員友好。除了核心之外,特斯拉還通過專門為大規(guī)模部署設計Dojo芯片來節(jié)省模具區(qū)域。
物理實現(xiàn)
縮小后,道場內核在一個非常大的645毫米2模具上實現(xiàn),稱為D1。與我們熟悉的其他芯片不同,單個道場模具是不自給自足的。它沒有DDR或PCIe控制器。模具邊緣周圍有IO接口,允許模具與相鄰的模具通信,延遲約為100ns。


要獲得系統(tǒng)內存,Dojo D1芯片必須與具有板載HBM的接口處理器通信。該接口處理器反過來通過PCIe連接到主機系統(tǒng)(接口處理器安裝在PCIe卡上)。從理論上講,最小的功能Dojo部署將涉及Dojo模具、接口處理器卡和主機系統(tǒng)。但特斯拉在包含25個模具的模塊中部署了道場模具,以給人一種規(guī)模感。Dojo D1模具是專門為超級計算機設計的積木,僅此而已。

這種專業(yè)化可以節(jié)省更多的模具面積。Dojo D1不會將區(qū)域用于DDR和PCIe控制器。除了外部設計用于與相鄰模具接口的定制IO連接器外,大多數(shù)模具都被道場巖心占據(jù)。

相比之下,設計時考慮到了更靈活的部署靈活性的芯片在IO上花費了大量區(qū)域。AMD的Zen 1“Zeppelin”模具是一個很棒的展示。Zeppelin可以直接連接到DDR4內存、PCIe設備、SATA驅動器和USB設備——非常適合滿足客戶要求。在服務器中,IFOP接口允許它與相鄰的模具通信。放置在IFOP附近的大塊SRAM可能是窺探過濾器,有助于在高核心計數(shù)設置中高效保持緩存一致性。Dojo不會試圖保持跨核心的緩存一致性,也不使用任何SRAM進行窺探過濾器。

AMD通過將大約44%的“Zeppelin”模具面積花在內核和緩存以外的邏輯上來實現(xiàn)這種靈活性。Dojo只將28.9%的模具面積用于SRAM和內核以外的其他東西。
最后
特斯拉Dojo超級計算機背后的微架構展示了如何實現(xiàn)非常高的計算密度,同時仍然保持CPU使用分支代碼良好運行的能力。為了實現(xiàn)這一點,你放棄了定義我們現(xiàn)代計算體驗的大部分舒適感。如果你可以假設圍繞Dojo核心構建桌面,那么任何使用過MS-DOS的人可能會熟悉這種體驗。你無法同時運行多個應用程序。一個行為不端的應用程序可能會迫使你重新啟動系統(tǒng)。如果你沒有足夠的RAM來運行某個程序,你可以完全忘記運行它(無需分頁到磁盤)。
但這些權衡在特斯拉的超級計算機中是完全有意義的。特斯拉不需要Dojo內核來同時處理幾個正在運行的應用程序,Dojo只需要運行內部的可信代碼。因此,特斯拉并不關心虛擬內存支持。同樣,在Dojo上運行的機器學習程序在編寫時將考慮到該特定系統(tǒng)。你不會有一批任意程序,可能要求比可用內存更多。這意味著你不需要精確的異常(和虛擬內存)來允許過度提交內存、內存壓縮或將內存頁面切換到磁盤和從磁盤中切換等技術。精確的異常對調試也很有用,但特斯拉通過單獨的調試模式使調試變得更便宜。這樣,在編寫和調試代碼后,Dojo可以專注于運行它,而無需進行按順序提交指令結果所需的簿記。

當然,特斯拉為獲得高計算密度而采用的權衡在消費者或服務器CPU中是不可能的。但它們在行動中非常有趣,我們必須感謝特斯拉抽出時間在Hot Chips上展示。在過去的二十年里,流程節(jié)點的改進一直在放緩,導致單線程性能的進步放緩。在過去的五年里,電源和冷卻限制一直在減緩多線程性能。但對更多計算能力的胃口并沒有放緩,因此公司正在轉向更專業(yè)的硬件來跟上步伐。特斯拉Dojo超級計算機的架構是一個很好的例子,說明如何進行權衡來提高計算密度,以及當前的趨勢如何有利于為吞吐量綁定應用程序引入專用硬件。
了解更多前沿科技,關注我就是你最好的選擇!