在近期舉辦的第七屆DevLabs大會上，前Google首席科學家Jeff Dean與Google雲AI基礎設施負責人Bill Jia展開了一場深度爐邊對話，圍繞GoogleTPU近十餘年的演進歷程展開系統回顧，話題涵蓋Google搜尋基礎設施的早期規模化、TPU的誕生契機、歷代TPU架構迭代、軟硬體協同設計、十萬晶片級叢集的可靠性工程、開源生態戰略、PyTorch生態適配、專用硬體自動化設計、推理時代的基礎設施重構等，並在最後的觀眾問答環節回應了商用硬體與專用超算路線之爭、3D環形拓撲的擴展極限，以及開發者如何參與開源生態擺脫對CUDA依賴等問題。
Jeff Dean提出，當前晶片設計流程仍高度依賴人工編寫暫存器傳輸級程式碼，若能借助強化學習等自動化搜尋循環實現設計流程自動化，晶片研發周期有望從150人兩年壓縮到10人三個月，這將從根本上改變專用硬體的經濟性，使團隊不必再為兩到六年後的算力需求押注。他還指出，在十萬晶片規模下Google依然堅持完全同步訓練以保證可復現性，此外他認為，Google能夠持續迭代TPU的核心優勢並非單點技術突破，而是硬體設計師、編譯器工程師與機器學習研究人員之間高頻且雙向的內部反饋循環，這種協同設計能力本身就是難以複製的護城河。
01
Google搜尋基礎設施的早期規模化
請回顧一下你剛加入Google時的經歷，以及公司早期在搜尋業務上的發展狀況。到2001年，當時Google在搜尋系統架構上做了那些關鍵調整？
Jeff Dean：剛加入Google時，我們正試圖打造一款真正高品質的搜尋產品。在之前的工作中我已經接觸過一些資訊檢索的業務，尤其是在網路上利用網路的圖結構來增強網頁上的現有資訊。Google似乎是一個很自然的地方，可以在一個不那麼偏向學術研究的環境中嘗試這些想法，更加專注於讓人們直接使用我的工作成果。大家充滿幹勁，我們的流量每周都在以7%的速度增長。如果你算一下1.07的52次方，就知道我們每年的規模都在極速膨脹，我們每天都在拚命避免系統在高峰期崩潰。隨著公司的不斷發展，我們做了越來越多的事情，看著這些進步總是令人興奮。
（關於2001年的架構調整）為了避免系統崩潰，我們當時正在做的一件事就是不斷重寫整個搜尋、索引和查詢服務系統，使其更高效，並使用不同的資料結構以及更緊湊的表示形式。到了2001年我們意識到，擴展索引規模和查詢服務容量的方式是把索引劃分成更多的分區，我們稱之為分片，從而建構更大的索引。對於每一個分片，我們會製作更多的副本來提供並行處理更多查詢的容量。最終我們在每個資料中心不再是7個分片且每個保留10個副本，而是擴大到了60個分片且每個保留20個副本。
我們算了一筆帳，意識到與其把索引放在磁碟上，不如把整個網路的索引直接放進那1200台機器的記憶體中，而這些機器原本就是用來裝那60個分區的20個副本的。我們在三天內將整個搜尋系統的速度提升了4倍，相當過癮。
02
深度學習算力缺口催生TPU
談談2013年你所做的深度學習算力估算，這些工作是如何催生TPU項目的？TPU的後續幾代產品在訓練和推理上的分工有何變化？2017年你們發表了相關論文並被接收，這也是ISCA五十年歷史上被引用次數最多的論文，是這樣嗎？
Jeff Dean：我開始從事深度學習訓練基礎設施和訓練深度學習模型的工作，因為這似乎是正確的抽象層。我在1990年讀本科時就接觸過神經網路，並寫了一篇關於神經網路平行訓練的本科論文，因為我覺得這是正確的發展方向，我們只需要32個性能較弱的處理器協同工作而不是一個單核處理器，這樣就能訓練出令人驚嘆的模型。結果證明我們需要增加100萬倍的算力，而不是僅僅32個處理器。後來得益於摩爾定律的發展，大概在2008年或2009年左右我們開始擁有這種規模的算力。2011年我們啟動了一個項目，旨在擴大規模並進行語音、視覺和語言模型的分佈式訓練。
當時我們的資料中心裡有大量的CPU。我們用了大約16000個CPU核心來訓練一個大型視覺模型並獲得了驚人的提升。我們在語音模型上也做了同樣的嘗試。模型在質量上展現出了驚人的飛躍，僅僅在我們嘗試建構深度聲學語音模型的兩個月時間裡，我們在詞錯誤率上的改進就相當於過去20年語音研究的成果總和。這讓我們意識到這項技術潛力巨大，但問題是我們該如何部署和提供服務。我做了一些估算，假設有1億人每天開始對著手機說話幾分鐘，我們要如何處理這些請求。結果顯示，僅僅為了在Google一個不起眼的邊緣業務中推出更好的語音識別模型，我們就需要將Google現有的電腦數量翻一番。這聽起來有點誇張了。
更不用說這很不切實際。因此我們決定研發專用硬體才是出路。這就是TPU系列的起源。我們希望建構專門的加速器，這些晶片非常擅長低精度線性代數計算，僅此而已。這樣你就能以高得多的效率提供語音和視覺模型服務。結果證明當我們在2015年拿到晶片時，它的每瓦性能比當時的CPU和GPU高出30到80倍。
（關於產品迭代分工）那是TPU v1，實際上是專門針對推理任務設計的。隨後的幾代產品則主要兼顧了訓練和推理。而最近我們又開始將這兩條產品線分開，因為訓練和推理的側重點確實有所不同。
（關於ISCA論文的引用情況）關於TPU v1的那篇ISCA論文是多人合著的，大概有35位作者。考慮到它在電腦體系結構五十年的歷史中發表時間相對較晚，能有這樣的成績確實很不錯。
03
TPU軟體棧的搭建邏輯
Google過去一直被視為頂尖的軟體公司，後來卻大力推進硬體路線圖。做硬體是一回事，但這背後遠不止晶片設計本身，還需要編譯器和框架的完美配合。能否介紹一下Google早期在TPU軟體棧上，也就是XLA編譯器與TensorFlow上的探索歷程？
Jeff Dean：作為一名機器學習開發者或研究人員，你想要的是能夠表達你高維度的想法，然後讓它神奇地在大規模系統上運行，而不需要你去過多考慮為了獲得良好性能而在底層需要做的所有性能最佳化。這意味著需要有一個優秀的編譯器，極其出色的互連技術，以及能夠抽象表達這種邏輯的框架，無論我是與本地機器上的4個晶片通訊，還是在一個大型分佈式環境中的10000個晶片通訊。你希望自己不必過多操心這些底層細節，也不用為了這兩種不同的環境去寫不同的程式碼。
談談不同代次TPU的迭代情況，以及你認為未來的發展方向是什麼？
Jeff Dean：在TPU v1時它實際上只是一塊可以插在插槽裡的PCIe卡，我們其實採購了一大批。我們向當時的首席財務官爭取大量採購這些卡，因為我們肯定會用到它們。TPU v2是我們設計的第一款將訓練需求考慮在內的系統，它不再僅僅是單塊晶片，而是一個類似超級電腦的系統，包含許多通過二維環面拓撲互連的晶片。在後續的幾代產品中我們從TPU v3開始引入了液冷技術。當你看到系統裡的冷卻管道直接連接到晶片表面時總是讓人感到非常興奮。
當然漏水可不是什麼好事。到了TPU v4隨著計算叢集規模的擴大，要知道TPU v2和v3在256塊晶片到1024塊晶片之間採用的是固定線路連接。而當系統規模變得更大時，個別晶片、主機板或系統出現故障的機率就會增加。你需要更強的容錯能力，因為即便有部分元件損壞，你仍然希望能夠使用一個大得多計算叢集的完整拓撲結構。因此TPU v4引入了機架間光學可重構網路的概念。你可以用帶有微調反射鏡的交換機將散佈在資料中心各個機房過道的模組重新組合起來，讓這8個機架的機器感覺就像緊挨在一起一樣，即便它們在物理位置上相隔很遠。
這非常實用。TPU v5、v6和v7均引入了低精度計算格式，這對於大幅提升性能極為有用。將16位、8位和4位的浮點運算次數放在一起比較總是有些奇怪，但如果你能在你的機器學習演算法中善用它們，低精度計算絕對是搾取更高性能的絕佳途徑。
04
單一程式設計模型如何驅動萬卡級分佈式訓練
第一代TPU是2015年推出的，2017年發表了那篇現在已成為傳奇的論文，然後到了2018年你們推出了Pathways。談談Pathways以及為什麼它如此重要？
Jeff Dean：我們當初建構Pathways主要是為了處理更加不規則和更稀疏的模型訓練，同時也將其作為底層系統基礎設施，使我們能夠擁有單一的程式設計模型和單一的處理程序來驅動大規模分佈式訓練系統。對於程式設計師來說，如果你只需寫一個單一的Python處理程序，看起來就像它連接了8000個裝置，而你只想在所有這8000個裝置上執行一次全歸約操作，系統就能自動幫你完成，無需你編寫任何特殊程式碼，這無疑是一種極佳的體驗。
底層的XLA編譯器軟體處理了部分工作，而在稍高一點的層面上，Pathways負責編排海量不同晶片之間的資料移動。同一TPU叢集內的晶片之間需要通訊時，它會使用TPU叢集內部的ICI高速鏈路。如果通訊跨越多個叢集，或者這個叢集裡的晶片需要和那個叢集裡的晶片對話，Pathways就會利用資料中心網路來調度資料傳輸。如果你正在跨越多個不同地區的資料中心拼湊一個訓練任務，它甚至會利用廣域網來完成。這就為你提供了一個非常優雅的抽象概念，你擁有一大堆算力資源，而系統會自己想辦法以最優的方式去利用它們。
05
軟硬體協同設計的護城河：Gemini模型與TPU基礎設施如何共同進化
Jeff你早些時候大量參與了TPU的網路設計、TPU自身的設計，然後是編譯器和框架，接著又主導了Pathways的資料移動與架構重構編排。這些構成了完整的AI基礎設施棧，而你現在又深度參與到Gemini模型的設計中。從你的角度來看，模型與基礎設施棧的協同設計讓我們已經獲得了那些收益？對於未來，你有什麼想法可以讓我們繼續推高模型的性能邊界？
Jeff Dean：歷代TPU的迭代其實都深深受益於一個事實，那就是我們自己就是機器學習算力的超級大客戶。我們內部有很多人在不斷嘗試突破新技術研究的邊界，這些新想法往往會以前所未有的方式給硬體帶來極大的壓力。由於所有這些團隊都在同一家公司的保護傘下，未來的TPU硬體設計師、編譯器與基礎設施的軟體架構師，以及機器學習研究人員之間能夠進行極為頻繁的互動。研究人員可以反饋說，某種方法在小規模下似乎很有效，並且很可能在一年後的大規模模型訓練中變得至關重要。
我們要確保我們的硬體能夠支援這些需求。作為一名電腦架構師，如果處於孤立狀態想要在某款晶片服役的兩到六年生命周期內精準預測瞬息萬變的AI領域的走向，那是極其困難的。你能獲得越多關於技術可行性的洞察，你們之間的迭代循環越快，這不僅僅是研究人員在提需求說我們需要這樣做，通常硬體設計師也會反過來說那太難實現了，但我們可以換種方式做，這樣對你們有幫助嗎。通過這種雙向的溝通，研究人員會發現這能讓我們實現非常接近的目標，甚至可能效果更好。這種互動對於實現高效的軟硬體協同設計超級重要。
Bill Jia：昨天我在台上做開場演講時就談到了，當我們開始擴大Gemini模型的規模時，我記得每個月我們都會坐在同一個會議室裡討論可靠性問題。我代表基礎設施團隊，而Jeff代表模型研發團隊。我們看到Gemini模型的體量變得越來越大。過去我們用幾十塊TPU來訓練模型，而現在我們用數十萬塊TPU來訓練。我們共同關注的一個核心指標叫作有效吞吐量。起初我們的有效吞吐量並不理想。我記得大家坐在一起看著資料，當時無效吞吐量甚至超過了有效吞吐量。
一開始隨著我們投入更多TPU去預訓練這個模型，模型變得非常複雜。我們需要更大的規模、更多的資料以及更高程度的平行化。初期的有效吞吐量只有50%、60%或者70%，這顯然是不夠的。長話短說，現在我們使用規模極其龐大的訓練基礎設施來訓練複雜度極高的模型，而有效吞吐量已經可以高達95%，甚至98%。
Jeff Dean：我認為這是許多舉措共同作用的結果，包括更精細的維運實踐，晶片部署時更嚴格的測試，以及能夠更好地應對故障並在系統部分當機時仍能繼續向前推進的軟體架構。所有這些因素加在一起起到了決定性的作用，帶來了巨大的改觀。
Bill Jia：我認為這是我們獨有的另一個優勢。如果我們自己掌握了模型、資料、完整的硬體網路以及軟體棧，這就給了我們的維運團隊極大的便利。
Jeff Dean：因為如果你孤立地看待某個TPU托盤，如果你沒意識到它是更大系統的一部分，你可能會覺得等下一周例行巡檢時再修也無妨。但如果它是某個正在運行任務的叢集中的一部分，並且叢集中已經有多個損壞的托盤，你肯定會希望立刻去修復它，因為這個待修復部件的故障影響範圍實際上比它看起來要大得多。
06
十萬晶片級可靠性工程
當談到十萬塊晶片規模的可靠性時，這已經不是打個軟體補丁就能解決的問題了。這更多是系統設計問題，還是跟物理極限有關？能詳細談談OCS和Virgo架構，以及它們在提升可靠性方面發揮了什麼作用？
Jeff Dean：可靠性實際上是整個系統的一種屬性。你需要在許多不同的層面上建構穩健性。其中一種方法就是用不可靠的部件拼湊出一個穩健的整體。這源於Google成立之初的理念，我們當時購買了最廉價的消費級PC來處理搜尋流量，然後我們在其之上建構了強大的穩健系統，使我們能夠應對單台機器的故障並依然穩定提供所需的服務。也許這是一個分佈式檔案系統，所以你會跨多台機器複製資料，即使某些副本當機你總能訪問到可用的資料。在大型AI訓練系統中你也可以採用同樣的思路。如果你默認的配置是20個叢集協同工作但其中1個當機了，你可以在修復那第20個叢集的同時，讓剩下的19個叢集繼續推進任務。
（關於OCS與Virgo架構）Bill Jia：這是一段漫長的旅程。首先我們關注了中斷率，比如每1000塊晶片、每個計算叢集每天會發生多少次中斷。我們必須盡力降低這個數值，因為隨著訓練叢集規模的不斷擴大，過高的中斷率是非常致命的。我們所做的一項工作就是提升硬體本身的可靠性。除此之外我昨天也提到，甚至在預訓練開始之前我們就會使用軟體掃描整個算力艦隊的每一個部件，檢查其生命體徵。如果發現某個部件的體徵有潛在問題，我們會在訓練開始前就將其修復或隔離。這是第一步。
其次即使硬體在訓練前展現出了極佳的狀態，長達數周的訓練過程中依然會不可避免地出現問題。當問題發生時我們如何精準定位是那些特定元件或伺服器出了故障並迅速解決，我們需要隔離它、替換它、修復它。這是第二步。同時正如Jeff提到的，這背後有大量的維運協同工作。如果資料中心團隊正在維護一整排機櫃，而這排機櫃偏偏正在執行訓練任務，那情況就很糟糕了。我們在資料中心團隊、站點維運、站點可靠性工程師以及機器學習研究員和工程師之間建立了極其緊密的協同機制。保持步調一致需要大量的協調工作。Jeff也在Gemini模型設計和軟體設計方面發揮了領導作用。有時當故障發生時由於訓練同時採用了資料平行和模型平行，如果故障只發生在一個資料副本內，其他的副本可能就會直接接管繼續運行。系統會自動對權重進行平均，這樣你就不會因為某一個副本當機而導致整個任務受到干擾。這是一種全方位的綜合提升。
Jeff Dean：在運行過程中各種問題都可能發生，你很難預測系統會以什麼樣的方式崩潰。關鍵在於建立一種穩健的機制，即使你不知道是什麼原因導致的故障也能敏銳地捕捉到它。有時候你會遇到這樣的晶片，當溫度升高時它就開始出現可靠性問題，比如計算二加二等於五。你必須要能檢測出這種異常，因為如果你正在建構大規模的AI這絕對是個災難。
Bill Jia：那是最棘手的問題，也就是靜默資料錯誤，那真是太可怕了。
Jeff Dean：這種情況可能發生在晶片內部，可能發生在不穩定的網路鏈路中，也可能發生在其它的很多環節。在這方面我們運用了很多Google早期秉持的原則。因為我們當初採購的是廉價消費級PC，它們不僅沒有糾錯碼記憶體校驗甚至連基礎的奇偶校驗都沒有。如果你用成千上萬台這樣的機器去做算力密集型工作肯定會遇到隨機的位元翻轉問題。因此我們的很多計算任務都必須具備抵抗位元翻轉的穩健性。一種實現方式是達成共識，如果你在處理十億個網頁時丟掉其中一個並不是世界末日。你可以通過在硬體層之上的軟體層加入總和檢查碼機制來忽略這些底層硬體錯誤，從而確保整體業務的穩健，即便某台特定的機器或網路鏈路並不可靠。
07
開源是為了讓技術成為跨組織的協作項目
談談Google在AI基礎設施戰略中，如何確保它不會成為一個被圍牆封閉的特權資源？請談談JAX和StableHLO以及開源生態整體的演進，探討這種開源理念在當今是如何確保這項技術不會走向封閉的。
Jeff Dean：對我們來說，與更廣泛的生態系統進行互動是至關重要的。將我們的技術開源，讓人們能夠查看原始碼、修改它、幫助我們改進它，使其成為一個跨越眾多組織的協作項目，而不是僅僅我們自己內部獨享的秘密，這是極其重要的。這也是為什麼我們開源了TensorFlow、JAX以及用於XLA編譯器的StableHLO中間表示，並且我們未來還會開源更多的技術。多年來我們一直是整個開源社區的重要貢獻者。通常情況下我們都是那些大型開源合作項目中貢獻最大的組織之一，因為我們堅信開放原始碼的力量，我們也深信當大家攜手共進時整個生態系統都會從中受益。
Bill Jia：Jeff一直在領導著眾多AI基礎設施的開源戰略。正如Jeff所提，Google過去開源了TensorFlow和JAX，Kubernetes也是由Google孵化並開放原始碼的，還有Android等諸多優秀項目。現在隨著我們將TPU推向社區，並強調在Google雲平台上使用TPU，我們正在進一步加碼這項開源戰略。我們不僅開源了JAX的核心，還建構了上層庫。例如在強化學習領域，我們推出了基於JAX核心建構的trlX並將其全面開源。
對於檢查點和推理，我們在JAX之上建構了上層的庫與框架並將其開源，因為我們希望大家能直接利用這項開源戰略來使用TPU。另一方面，我們也希望貼近開發者的真實需求，所以我們要積極擁抱開源社區中已經成熟的產品。PyTorch在社區中生態成熟且被廣泛使用。既然如此，我們就選擇PyTorch，積極推進PyTorch在TPU上的運行。我們目前有一個名為Torch TPU的項目，現在正處於內測階段，我們計畫在下個季度推出公開預覽版。然後到第四季度，我們計畫在GitHub上公開發佈，讓所有人都能使用。
如果在TPU上使用PyTorch，只需修改幾行簡單的程式碼將後端裝置設定為TPU即可。如此一來，所有的訓練和推理任務就有望原生運行在TPU上。我們同時支援動態圖模式和編譯模式。這是在框架層面的工作。此外，我們也在積極與vLLM和SGLang的使用者以及團隊進行溝通，確保這些上層的開源推理框架同樣能在TPU上流暢使用。我們正在做大量的工作，不僅開源自家的軟體，同時也積極擁抱社區中成熟的開源產品。
08
如何將晶片研發周期從兩年壓縮到三個月
純粹從基礎設施的角度來看，你目前在思考的物理瓶頸是什麼？
Jeff Dean：瓶頸有很多。我個人非常看好專用硬體的發展，因為這是實現系統極致高效的關鍵途徑。如今，少數幾種工作負載已經佔據了全球絕大部分的計算量，這極其需要硬體的專業化定製。但專用硬體面臨的難題是，如果未來計算需求發生變化，你花費兩年時間精心打造的硬體可能就不再適用了。為了讓硬體專用化真正發揮作用，我們需要實現硬體設計流程的高度自動化。目前設計硬體的方式是組建一個龐大的團隊，一部分人根據高層規範編寫底層的暫存器傳輸級程式碼。
由於這是人工轉換的過程，你需要另一個團隊來驗證第一個團隊的程式碼是否正確無誤。接著又需要另一批人投入大量心血進行晶片的物理佈局。如果我們能建立出更多可以通過強化學習或其他演化演算法進行搜尋的自動化循環，並讓這些循環快速運行，目前的EDA工具通常不具備這種速度。如果能讓它們高速運行，我們就有機會在設計流程中引入高度自動化的探索循環，從而極大地壓縮設計周期。如果以前需要150人耗時兩年，而未來只需10個人在三個月內就能設計出一款新晶片，那麼我們將看到大量專用硬體的湧現。這樣一來，你就不必再去押注兩到六年後的計算需求，你的預判窗口縮短到了未來三個月到四年，這種預測顯然要容易得多。
09
推理時代的基礎設施：為AI Agent場景重新設計硬體
包括Google在內的許多公司和前沿實驗室此前主要精力都在如何打造出色的大語言模型上，硬體策略也大多圍繞訓練側展開。如今隨著大語言模型的逐漸成熟，巨大的流量開始向AI Agent與推理側轉移，應該如何設計硬體來重點應對激增的推理流量？
Jeff Dean：推理和訓練有很大的區別。在推理時模型是固定的，你需要處理海量的並行請求，因此你需要儘可能減少不可變資料的移動。不可變的是模型權重，而變化的是請求和KV快取。你需要設計一個能在這個環節極致最小化資料移動的高效系統。此外我認為隨著互動模式從簡單的提示響應演變為更加獨立的AI Agent互動，AI Agent執行某些操作，決定呼叫一個工具，運行工具獲取結果，然後將其放回模型上下文以決定下一步動作。
我們將會意識到現有的工具都太慢了，因為它們原本是為人類的互動迭代速度而設計的。如果你的工具呼叫中包含程式碼編譯步驟，你會因為編譯器太慢而感到痛苦。如果你將推理硬體做得極其快速，它生成程式碼的速度將遠超編譯速度，更不用說運行速度了。我們一直在做的一項工作就是提升內部工具的速度，使其更適合AI Agent使用。事實證明模型在不同程式語言之間的程式碼翻譯上表現得非常出色，因為你提供了一個意圖完全明確的規範，比如你有一個用Python等解釋型語言編寫的完整程序，而你需要一個功能完全等效的Go、Rust或C++程序。
AI Agent實際上能非常好地完成這項翻譯工作。這與我們平時讓AI Agent幫我們寫個Web伺服器這種常規編碼互動不同，常規互動中模型需要自行假設並填補很多細節，可能並不符合你的預期。但在程序邏輯完全明確的情況下，它的翻譯工作非常出色。它甚至能把所有的單元測試一併翻譯過去，在新系統中運行這些測試，並排驗證程序行為，確保結果完全一致。這是一種非常有效的方法，也許我們可以用一種語言編寫程式碼，然後再通過AI Agent翻譯成另一種。
Bill Jia：Google內部有一個現成的項目叫Project Turn，這就是一個很好的例子。因為我們有很多模型仍在使用TensorFlow，但在向JAX遷移的過程中，我們將所有的TensorFlow模型翻譯成JAX，並自動運行所有的單元和質量測試。
你們團隊曾向NeurIPS提交過一篇關於蒸餾的論文卻被拒了，當時的評審意見說它影響力甚微，能聊聊這件事嗎？
Jeff Dean：事實證明蒸餾非常重要。那是我和同事Geoff Hinton以及Oriol Vinyals共同提交的一篇論文，探討如何將一個模型整合作為教師模型，將其知識蒸餾到一個學生模型中。
那是那一年的事？
Jeff Dean：應該是2014年。當時我們最初的設想是用於視覺領域，訓練一個龐大的專家模型集合。我們在論文中做了一組實驗，總共有兩萬個視覺類別，我們專門為動物訓練了一個模型，為汽車訓練了另一個模型等等，然後將它們蒸餾成一個單一模型，讓它在所有這些類別上都表現出色。論文當時被拒了，但這無所謂。我們把它傳到了arXiv上，大家照樣會讀。所以從事科研不要氣餒。
10
現場問答：商用硬體與專用超算的路線之爭
Google早期通過商用計算裝置進行橫向擴展，而如今AI超級電腦變得越來越專用化，更像是早期的超級計算架構。如何看待硬體與機器學習系統在這兩種設計模式上的差異？未來會回歸通用商用硬體，還是繼續建構專用的超級電腦？
Jeff Dean：我認為我們在搜尋引擎上能依靠商用硬體進行擴展，是因為搜尋任務非常特殊。如果任務拆解得當，機器之間幾乎不需要通訊，而單台機器上則有大量的計算工作。因此你不需要任何特殊的網路互連技術。當時我們的機器上只有百兆乙太網路卡，整個機架的40台機器共享一個千兆的上行鏈路。雖然這是40比1的超額配置，但這毫無影響，因為你只需向每台機器傳送諸如“帕洛阿爾托餐廳”的檢索請求，然後接收一個包含十條結果的簡短程式碼片段。
但模型訓練徹底打破了這種邊界，因為它需要極高的資料連通性。最理想的情況是能在一塊晶片上完成所有訓練，但那要麼耗時太長，要麼記憶體根本放不下。所以你不得不將問題拆分到多塊晶片上，而無論你是採用模型平行還是資料平行，通常都會產生海量的通訊開銷。這就是為什麼機器之間需要採用特殊的高速互連技術。你需要極致提升單晶片的性能，從而減少所需的晶片總數。即便如此，你仍然需要跨越大量的晶片進行分散式運算，這也是為什麼我們要在這種架構中引入液冷技術，以確保每塊晶片都能發揮出極致性能。
我認為推理端會採用一些專用硬體，但其網路架構不會像訓練端那樣特殊和昂貴，特別是對於規模較小的模型而言。不過隨著模型規模越來越大，推理端的通訊需求也會逐漸超出標準乙太網路和傳統商用網路的能力範疇。但總體而言，推理硬體肯定會比訓練硬體更加主流和通用化。
11
現場問答：3D環形拓撲能否撐起十萬乃至百萬晶片規模的AI超算
注意到TPU採用了3D環形拓撲的連接方式，與UEC等其他聯盟推崇的Clos交換架構相比，從擴展性的角度來看，當計算規模擴展到數萬乃至上百萬晶片時，3D環形拓撲還能否支撐如此大規模的全對全通訊？
Bill Jia：我可以透露的是，我們目前最大的TPU pod大約包含9600多塊晶片。那就是我們採用3D環形拓撲的規模上限。在這個規模之上，我們擁有Pathways作為軟體抽象層，統管眾多採用3D環形拓撲的pod。它通過資料中心網路進行互連，這就取決於你資料中心內採用了什麼網路架構了。我們甚至支援跨地域的多資料中心分佈式訓練，比如在不同地區的資料中心分別部署多個pod叢集，中間通過高速廣域網鏈路連接。這種擴展方式對我們來說效果非常好。有時你確實需要對計算任務進行對應規劃，切分出資料平行和模型平行的維度。通常我們會把模型平行限制在單個pod內部或pod的局部切片中，然後在不同的pod之間進行資料平行複製。事實證明這種策略非常有效。環形拓撲網路的優勢在於它的本地互連非常簡單，不需要在機房鋪設極其複雜的線纜。
Jeff Dean：除了光路交叉連接開關。
Bill Jia：對，那項技術已經非常成熟了。我們在機架內和叢集內使用ICI，然後使用OCS將其擴展到9600塊晶片。再往上，我們有Virgo網路架構，由Pathways軟體進行統一編排調度。當規模突破十萬塊晶片時，我們就會利用跨資料中心網路將所有資源連接起來。這種擴展能力對我們來說運行得非常順暢。
Jeff Dean：即使在這種龐大的規模下，我們依然能夠實現完全的同步訓練，這對於機器學習的可復現性和模型可解釋性來說是非常棒的。也許未來的某個階段非同步訓練會再次興起，但就目前而言，我們在同步訓練的規模化上已經走得很遠了。讓我們拭目以待。
12
現場問答：開發者如何參與OpenXLA生態，擺脫對CUDA的依賴
現在AI基礎設施領域感覺就像Android和iOS並存的時代，大多數使用者依然嚴重依賴CUDA來編寫庫或函數，對於年輕工程師來說，參與OpenXLA庫函數開發有那些切入點？（追問）很多函數或底層核心模組仍然與CUDA深度繫結，是否會出現用JAX、OpenXLA或Torch TPU實現等效底層功能、從而擺脫對CUDA依賴的趨勢？
Bill Jia：首先，OpenXLA、JAX以及建構在JAX之上的整套AI算力棧已經完全開源，這是一個重要的工作方向。另一個工作方向是Torch TPU項目。在內部我們正與Meta合作開發，但最終我們會把它發佈到GitHub上，讓廣大研究人員和工程師社區共同參與貢獻。這是我們齊頭並進的兩條戰線。當然，對於社區中任何有志於為開放原始碼庫做貢獻的年輕工程師，我們都極其歡迎與Google展開合作，我們可以一起探討如何共同開發這些開放原始碼專案。
此外，許多使用者正在使用基於TPU或GPU的JAX和PyTorch框架。目前在社區中，如果有人在使用時遇到困難，很多人最終只能尋求Google或輝達的官方技術支援，或者把問題發到社區論壇，期盼著能有遇到過類似問題的人來解答。在我看來，這種解決問題的路徑太慢了。其實大量的社區互助都可以轉化為對這兩大開源戰略的貢獻。試想一下，如果有一個精通JAX的AI Agent和一個精通PyTorch的AI Agent，當我在開發中遇到任何問題時，如果社區裡的每個人都貢獻出自己的資料和經驗來訓練這個AI Agent，最終在進行模型訓練、後訓練或推理時，我就相當於擁有了一位隨叫隨到的AI Agent助手。我不再需要去求助Google或翻看社區論壇，直接通過這個AI Agent就能解決問題。但這需要整個社區的共同參與，我們需要匯聚海量的資料點和實戰經驗。
Jeff Dean：關於參與開源貢獻，我想提供一個更高層面的建議。參與開源有很多種方式。一個很好的習慣是先主動與程式碼倉庫的維護者溝通，說明你想開發的功能是否有用，或者直接詢問他們目前最需要什麼。在你一股腦提交大量程式碼之前，先獲得一些方向性的認可和指導是非常必要的。如果你滿腔熱情想要做貢獻但又不知道從何下手，他們手裡通常會有一份大家期待已久的待辦願望清單。
（關於CUDA繫結核心的替代趨勢）Bill Jia：在Google內部，我們在評估TPU上的內部JAX和PyTorch技術堆疊時，會全面對標所有的CUDA函數和我們自有的函數。我們的目標是必須確保性能至少打平甚至超越對方。這是我們必須跨越的關鍵一步。在這個層面上，我們也在評估那些底層架構同樣可以開源，比如底層的TPU SDK開發工作。
Jeff Dean：我也認為開發者更希望在更高的抽象層級上進行思考。你應該把精力放在如何建構JAX或PyTorch的數學表示式上，而不是去糾結如何用核心程式碼來實現平行化，不管它是TPU還是GPU的核心語言。關於極致性能的搾取，最好的狀態是把這一切都交給編譯器和底層系統去處理，這樣你就可以專注於像矩陣乘法這樣優美的數學抽象。(數字開物)