Google 訪談實錄:共同設計 TPU、分散式系統與前沿 AI 模型
Google 開發者實驗室(Dev Labs)爐边對談:Jeff Dean 與 Bill Jia 談 AI 系統規模化,2026 年 9 月 20 日
簡介與 Google 的早期歲月
主持人:那麼,大家午安,非常歡迎各位的熱情參與。這是我們的第七場 Dev Labs 活動,而且每一次舉辦,規模都更大、辦得更好。這無疑是目前最棒的一次。我們收到了非常多的關注,謝謝大家的迴響。今天的進行方式是這樣的:我們會探討幾個問題,談談 TPU 的演進。這位是 Jeff,他完全不需要多做介紹。他是 Google 的第 30 號員工。那我們就從這個故事開始。我們今天也請到了 Bill,他等一下也會參與提問,同時我們也希望他能分享一些關於 AI 基礎設施在執行面上的內幕、願景,以及現在乃至未來幾年是如何落實執行的,這樣可以嗎?好的,我們開始吧!那麼 Jeff,跟我們多聊聊吧。你是 Google 的第 30 號員工。
Jeff Dean:對,沒錯。
主持人:那是 1999 年。這裡的大多數人當時可能還沒出生,或者正處於包尿布的年紀。那時候是什麼情況?Jeff Dean:不,那真的非常有趣。我的意思是,我剛加入時,我們全都在巴洛阿圖(Palo Alto)市中心的一個超小辦公室裡,擠在現在 T-Mobile 門市的樓上,差不多就是那種規模的辦公室。當時我們正試圖打造一個高品質的搜尋產品。我在上一份資訊檢索的工作中接觸過這方面技術,特別是關於網頁的部分,利用網頁的圖結構(graph structure)來擴充頁面上現有的資訊等等。因此,Google 似乎是一個很自然的地方,可以讓我們在一個比較不像研究機構、而是更多人直接使用我所做成果的環境中去嘗試,這真的超好玩。那時候充滿了高張力的活力,我們的流量大約每週成長 7%。如果你把 1.07 的 52 次方乘下去,就會知道每年簡直就像在融化一樣,我們得拼命避免在每週二中午系統崩潰。所以那段時間真的非常好玩。隨著公司持續成長,我們也做了越來越多的事情,每次看到都讓人感到興奮。
主持人:那麼,如果我們快轉到 2001 年,發生了什麼事?
Jeff Dean:為了避免系統「融化」,我們當時一直在做的一件事,就是不斷重寫整個搜尋、索引和查詢服務系統,好讓它變得更有效率、使用不同的資料結構以及更精簡的資料表示法。到了 2001 年,我們達到了一個時間點,意識到我們當時擴充索引大小以及查詢服務容量的方式,是為索引建立越來越多的分割區(我們稱之為 shards)來擴大索引規模。然後針對每一個分割區,你又會製作越來越多的複本(copies),好提供同時處理更多查詢的容量。我們最終發現,與其搞 7 個分割區、每個分割區 10 個複本,我們竟然可以在每個資料中心裡做到 60 個分割區、每個分割區 20 個複本。我們最後算了一下帳,意識到與其把索引放在磁碟上,不如把整個全球網頁的索引全部放進記憶體裡——也就是原本那 60 個分割區各 20 個複本、總共 1,200 台機器中。那真的超好玩。我們在短短 3 天內,把整個搜尋系統的速度提升了大約 4 倍,這真的非常過癮,沒錯。
TPU 的誕生:從語音辨識到自研 ASIC
主持人:然後在 2013 年,你做了一些簡單的快速估算。
Jeff Dean:對,沒錯。當時我開始投入深度學習訓練基礎設施,並著手訓練深度學習模型,因為這似乎是正確的抽象層。其實我在 1990 年讀大學時就接觸過類神經網路,我的大學畢業論文就是關於類神經網路的平行訓練,因為我當時就覺得:噢,這才是正確的做法。我們只需要 32 個效能較弱的處理器而不是一個,就能訓練出驚人的模型。事實證明,我們需要的是大約高出 100 萬倍的算力,而不是 32 個。不過最終,透過摩爾定律的進步等等,大約在 2008 年或 2009 年左右,我們開始具備了這樣的條件。接著在 2011 年,我們啟動了一個專案,真正試圖擴大規模,針對語音、視覺和語言進行模型的分散式訓練——當時我們的資料中心裡有 CPU,大量的 CPU。因此我們用了大約 16,000 個 CPU 核心來訓練一個大型視覺模型,並取得了驚人的改善。我們對語音模型也做了同樣的事情,這些模型在準確率上也帶來了驚人的提升。因此,我們在嘗試建立深度聲學語音模型的短短兩個月內,在降低字詞錯誤率(word error rate)方麵所取得的成效,就相當於過去 20 年語音研究的總和。當時的感覺是:噢,天哪,這太棒了!但我們要怎麼提供它上線服務呢?
當時我做了一些計算,我說:好吧,假設有 1,000 萬人開始對著手機講話、每天講幾分鐘,那要怎麼支應這樣的服務?結果發現,光是為了在 Google 某個極小、極冷門的角落推出這個更好的語音辨識模型,我們就得把 Google 擁有的電腦數量翻倍。這聽起來未免太過頭了,更不用說根本不切實際。於是我們決定,專用硬體才是正確的出路。這就是 TPU 家族的起源:我們想要打造專門的加速器,它們基本上非常、非常擅長低精確度線性代數運算,除此之外別無其他。這就是為什麼你能以高出許多的效率來提供語音模型、視覺模型等服務。結果證明,在 2015 年我們推出該晶片時,其每瓦效能比當時同期的 CPU 和 GPU 高出 30 到 80 倍。這就是最初的迭代。
Bill Jia:對,TPU v1 真的主要是針對推論(inference),對吧?隨後的世代則主要是針對訓練和推論。而在最近,我們甚至開始把這兩條路線再次分開,因為訓練和推論在本質上還是有些許不同。
主持人:當你們在 2017 年發表這篇論文時,它被接受了。這是 50 年歷史中被引用次數最多的——
Jeff Dean:對,是 ISCA 歷史上被引用最多的論文。關於 TPU v1 的那篇論文有非常多的作者,大概有 35 位左右吧。考慮到這在計算機架構 50 年的歷史中相對是近期的事情,這確實很不錯。
TPU 架構、軟體堆疊與光學交換的演進
Bill Jia:Google 一向被視為頂尖的軟體公司,卻開始積極推進硬體藍圖,這真的非常了不起。而且做硬體是一回事,因為要讓 TPU 變得強大,背後需要進行大量的主機板設計、晶片組設計以及 TPU 本身的設計。但同時,我們也必須確保編譯器(compiler)和框架(framework)運作得極為順暢。早期你實際上推動了 XLA(這是我們的 TPU 編譯器),同時也推動了 TensorFlow。能否跟我們介紹一下,當初推動 TPU 軟體堆疊的早期歷程是什麼樣的?
Jeff Dean:好。我的意思是,我認為身為機器學習開發者或研究人員,你想要的其實很簡單:你希望能夠表達自己的高階想法,然後讓它像是會魔法一樣在大規模系統上執行,而不需要去過多思考為了獲得真正高效能、可能需要在底層執行的所有效能調校細節。這意味著我們需要一個好的編譯器、非常優秀的互連技術(interconnects),以及能夠表達這種抽象概念的框架:我是在和本地機器的 4 個晶片對話,還是在和大型分散式架構中的 1 万個晶片對話?你最好不需要花太多心思去想這些,也不用為了這兩種情境去做截然不同的事情。
主持人:聊聊不同代 TPU 的迭代,你認為這一切未來會走向何方?
Jeff Dean:好的。所以繼 TPU v1 之後——它基本上只是一張可以插進插槽的 PCIe 卡,我們當時甚至在還不知道怎麼用它的情況下就買了一大堆。我當時把我們的財務長(CFO)堵在角落說:我們需要買很多這個,因為我們會用到它;雖然還不確定具體用途,但拜託讓我們買大量好嗎?於是我們買了。然後 TPU v2 是我們第一款將「訓練」納入設計考量的系統,它不只是單一晶片,而是一個類似超級電腦的完整系統,由許多透過 2D 環面(2D torus)互連的晶片組成。歷經多代之後,從 TPU v3 開始,我們導入了液冷技術。所以當你的系統晶片表面接滿了管線時,總是讓人感到興奮——當然,漏水可不是什麼好事。
接著是 TPU v4:隨著我們擴大 Pod 的規模(TPU v2 和 v3 在 256 個晶片與 1,024 個晶片之間有著固定的配線),當系統變得越來越大時,個別晶片、主機板或系統開始會出現故障。你必須具備更高的容錯能力,因為即使某些零件壞了,你仍然希望能夠使用更大規模 Pod 的完整拓樸。因此,TPU v4 在機器機架(racks)之間導入了光學可重組網路(optical reconfigurable network)的概念。
主持人:所以你們基本上可以像拼樂高積木一樣,從資料中心的各個角落把算力組裝起來,對吧?
Jeff Dean:沒錯,透過那些配有微調微型鏡片的奇妙小開關,讓這 8 個機架看起來彷彿就緊挨彼此,儘管它們其實分散在資料中心的許多走道中。這真的非常有用。TPU v5、v6 和 v7 都導入了越來越低的精確度運算格式,這極具價值。你能從中獲得更高的效能,儘管拿 16 位元、8 位元和 4 位元的 FLOPS 來互相比較總是有些奇妙。如果你能在機器學習演算法中善用它們,低精確度格式確實是榨出更高效能的絕佳途徑。
Pathways:協調不規則性與大規模分散式訓練
主持人:第一代 TPU 是在 2015 年,2017 年發表論文,現在它已經成為傳奇。接著在 2018 年,Pathways 問世了。我們來聊聊 Pathways,以及為什麼它如此重要。
Jeff Dean:好。我們當初著手打造 Pathways,主要是為了解決處理更加不規則、更加稀疏(sparse)的模型的訓練問題;同時,作為底層的系統基礎設施,它能讓我們擁有單一的程式設計模型、單一的行程(process),來驅動大型分散式訓練系統。對程式設計師來說,如果你的單純 Python 行程看起來就像是連接了 8,000 個裝置,那會舒心很多。你會覺得:太棒了,我只需要跨這 8,000 個裝置執行一個 AllReduce,系統就會自動搞定,完全不需要你去做任何特別的處理。
XLA 編譯器本身的底層軟體負責了一部分這類工作,而在稍微高一點的層級上,Pathways 則負責協調大量不同晶片之間的資料傳輸。因此,如果你在同一個 TPU Pod 內部有晶片,需要進行通訊時,它就會使用 TPU Pod 內部的 ICI 高速連結。但如果你需要跨多個 Pod 進行通訊,或者需要從這個 Pod 裡的晶片跟另一個 Pod 裡的晶片對話,Pathways 就會利用資料中心網路(或甚至廣域網路,如果你把奧克拉荷馬州、德州和愛荷華州等地的 Pod 拼湊成一個訓練工作時)來協調這項傳輸。這基本上賦予了你一個很棒的抽象概念:我手頭上有一大堆算力,我只需要讓系統自己去找出最有效的利用方式。
將硬體、基礎設施與 Gemini 模型進行共同設計
Bill Jia:太棒了。Jeff,你很早就深度參與了 TPU 的網路設計、TPU 本身的設計、編譯器與框架,當然還有負責協調整體資料傳輸、網路流量重組以及所有事務的 Pathways。這些全部都是 AI 基礎設施堆疊。現在,你又深度參與了 Gemini 模型(當前的 Gemini 模型以及未來的 Gemini 模型)的設計。因此,Gemini 模型能夠與整個基礎設施堆疊進行真正的共同設計(co-design)。從你的角度來看,我們已經收穫了哪些好處?對於如何推動模型的前沿邊界,你又有什麼樣的未來想法?
Jeff Dean:我認為,整個 TPU 世代之所以能受惠良多,很大程度上歸功於我們自己就是機器學習算力的大用戶,而且我們擁有一群積極推動全新研究邊界的人員,他們會以現有工作負載意想不到的方式去對硬體施加壓力。因此,藉由將所有東西都納在 Alphabet 的同一個屋簷下,未來的 TPU 世代設計者、編譯器與基礎設施的軟體設計者,以及機器學習研究人員之間,實際上能夠產生非常密集的互動。研究人員可以說:「嘿,這種做法在小規模下似乎行得通,我們認為這對未來一兩年的大規模訓練至關重要,我們希望確保我們的硬體能夠支援它。」
因為身為計算機架構師,如果你閉門造車,試圖去猜測快速變遷的 AI 領域在你當前設計的晶片預計上線的 2 到 6 年時間窗口內會走向何方,這會變得無比困難。你能得到的洞察越多——「對,我們認為這行得通」——而且迭代次數越多,情況就會越好。這通常不只是「我們覺得我們需要這麼做」,硬體團隊往往會回應:「嗯,這很難,但我們可以這樣做,這樣有用嗎?」於是你們之間就會展開這種來回對話:「對,這實際上能讓我們做出非常接近這個目標、甚至更好的成果。」這對於實現有效的共同設計至關重要。
大規模可靠性:橫跨 10萬個 TPU 的 Goodput 與容錯能力
Bill Jia:昨天我在大會上發表開幕致詞時提到,隨著我們開始擴大 Gemini 模型的規模——大家要知道,我記得我們每個月都會坐在同一個會議室裡討論可靠性。我是負責基礎設施(infra)這一塊,而 Jeff 代表的是 Google DeepMind。我們當時會說:「嘿,你知道嗎,Gemini 模型變得越來越大,接著我們需要動用大量的 TPU 來訓練模型。最終,我們使用了 10萬個 TPU 來訓練模型。」這就是我們共同關注的一個指標,叫做 Goodput(有效產出率)。Goodput。
起初,我們的 Goodput 表現並不好。我記得我們當時共同討論過:Badput(無效產出率)壓倒了 Goodput。一開始大概只有 60%。隨著我們投入越來越多的 TPU 來進行模型預訓練,情況變得非常複雜。我們需要更大的規模、更多的資料、更多的平行化。Goodput 一開始大概是 50%、60%、70%,但這遠遠不夠,對吧?不過簡單來說,現在我們使用大規模的訓練基礎設施來訓練複雜度高得多的模型,而 Goodput 可以高達 95%,甚至 98%。
Jeff Dean:對,我認為這其實是許多因素的結合:更好的維運實務、晶片部署時更嚴格的品管測試(QA testing),以及能夠更妥善處理故障並確保持續向前推進的軟體機制——即使系統的一部分已經當機。所有這些因素結合在一起至關重要,並帶來了巨大的改變。
Bill Jia:對,我認為如果我們同時擁有模型、資料、整套硬體、網路、所有軟體堆疊以及維運團隊,這就是另一大優勢。
Jeff Dean:你可以說:嘿,這超級重要,因為如果你把單一一盤(tray)TPU 視為獨立個體,不把它當成更大型系統的一部分,你可能會覺得「好吧,沒問題,我們下週順道經過時再修理它就好」。但如果它是正在運作、且其中包含故障機盤的 Pod 的一部分,那你就會想要立刻去修好它,因為你需要修復的目標所帶來的影響範圍(blast radius),實際上比看上去要大得多。
主持人:所以,當你談到 10萬個晶片時,可靠性其實已經不再只是一個軟體修補程式(patch)的問題了。你會說這更多是系統設計的問題嗎?或者它更偏向物理學範疇?
Jeff Dean:我認為可靠性本質上是整個系統的屬性,而你想讓很多不同的面向都變得強固。你能做到的一件事,就是用不可靠的零件拼湊出強固的整體。這其實是 Google 從創立之初就有的作法。我們購買便宜的消費級 PC 來支應搜尋流量,並在這些機器之上打造出強固的軟體系統,讓我們能夠應對單機故障,同時依然提供系統理應具備的功能。比方說分散式檔案系統,你把資料複製到多台機器上,即使某個資料複本掛了,總還是有其他複本可以用。
我認為在大規模 AI 訓練系統中也可以做到這一點。如果你有一個 20 個 Pod 的預設配置,但其中一個壞了,你完全可以在維修第 20 個 Pod 的同時,利用其餘 19 個 Pod 繼續推進工作。
維運協調與緩解隱性資料損壞
主持人:Bill,你昨天在主題演講中也稍微談到了 OCS 和 Jupiter。你想不想順便聊聊這個,以及它是如何融入可靠性機制的?
Bill Jia:好,這是一段漫長的歷程。首先,我們觀察每 1萬個晶片、每一個零件的中斷(interrupts)頻率,看看每天有多少次中斷。我們必須把這個數字降到最低,因為如果中斷次數太多,隨著訓練叢集規模擴大,過多的中斷會帶來嚴重負面影響,對吧?所以我們做的一件事就是提升硬體可靠性。不僅如此,正如我昨天提到的,甚至在預訓練開始之前,軟體就會掃描整個集群、每一個元件:你的生命徵象(vital sign)如何?如果生命徵象有問題,我們會在訓練開始前就進行修復或直接將其排除。這是第一點。
第二點:訓練開始了。但在訓練啟動之後,即便硬體展現出非常強健的生命徵象,訓練往往會持續數週之久,過程中依然可能出問題。我們該如何識別到底是哪一個特定的元件和伺服器出了問題,然後進行修復、排除、替換並修復它?這是第二點。
此外,正如 Jeff 所提到的,這裡面牽涉到大量的維運協調。如果資料中心正在維護整個資料中心機架列(row)的電力,而剛好這一列正在執行訓練,那當然非常糟糕。因此,我們必須與資料中心建置團隊、網站可靠性工程師(SRE)、以及所有 ML 研究人員和工程師進行大量的協調。這需要大量的協調工作,我們必須保持高度同步。
我認為 Jeff 在 Gemini 模型設計和軟體設計方面也主導了許多工作。有時候事情會出狀況,但訓練本身結合了資料平行化(data parallelism)與模型平行化(model parallelism)。如果問題發生在單一資料複本(replica)內,也許其餘的複本會繼續運作,對吧?系統會自動對權重進行平均,所以你不需要為單一複本的故障而煩惱。這是一個非常全面性的改善。還有什麼要補充的嗎?
Jeff Dean:沒有。我的意思是,各種狀況都有可能發生,要預測所有可能的故障方式是非常困難的。因此,建立能夠偵測故障的強固機制至關重要——即便你未必完全清楚故障的原因是什麼,但你能偵測到它。偶爾你會碰到這樣的晶片:當溫度上升時,它們開始出現不可靠的問題,比如把 2 加 2 算成 5 之類的。
Bill Jia:那是最棘手的情況,也就是隱性資料損壞(silent data corruption)錯誤。那簡直是場惡夢,太可怕了。
Jeff Dean:對,這可能發生在晶片本身、可能發生在不穩定的網路連線上,也可能發生在許多地方。我們正在運用許多與 Google 早期相同的原則:因為我們當時買的是便宜的消費級 PC,它們不僅沒有 ECC(錯誤檢查與校正),連記憶體的同位檢查(parity)都沒有。因此,如果你用大量電腦來做事而它們都沒有同位檢查,你就會在資料中遇到隨機的位元翻轉(bit flips)。我們所做的許多運算,都必須對位元翻轉具備強固性。其中一種做法是:沒關係,我正在處理 10 億個網頁,如果漏掉其中一個,大概也不是世界末日。因此,你可以在硬體層之上的軟體層透過檢查碼(checksumming)來忽略這一點,讓系統保持強固,即使特定的機器或網路連線不可靠也無所謂。
開源策略:JAX、StableHLO 與 TPU 上的 PyTorch
主持人:現在,Google 是開源(open source)的堅定信仰者,所以我稍微轉換一下話題來談談開源。我知道台下有很多學者,也有很多我們的合作夥伴。讓我們來談談 Google 的 AI 基礎設施策略,確保它不會變成一個被「設限封閉」的資源。我們討論的所有這些技術,我們都希望能確保與全世界分享。讓我們聊聊 JAX、StableHLO 以及整體開源的演進,這對於確保它不會變成 gated resource(封閉受限資源),在今天究竟有多重要?
Jeff Dean:對我們來說,與更廣泛的生態系互動至關重要。將技術開源,讓大家能夠看到原始碼、進行修改、協助我們改進,並讓它成為一個跨多個組織的協同合作專案,而不是我們自己緊抓不放的東西,這極為重要。這就是為什麼我們開源了 TensorFlow、JAX,以及 XLA 編譯器的 StableHLO 表示法。未來我們還會開源更多東西。多年來,我們一直是開源社群(如 Linux 核心等)的大力貢獻者。通常,我們是集體開源貢獻中最主要的幾個大型組織之一,因為我們深信這一點,並且認為當大家齊心協力時,整個生態系都會受益。
主持人:這簡直是完美的轉場,讓你來談談 PyTorch 以及你在 Meta 的經歷。
Bill Jia:對,我剛才在爐邊對談開始前還和 Jeff 聊過。Jeff 一直在主導許多 AI 基礎設施的開源策略。正如 Jeff 所提到的,過去 Google 開源了 TensorFlow 和 JAX,Kubernetes 也是由 Google 創造並開源的,還有 Android 以及許多極棒的成果。現在,隨著我們將 TPU 推向社群並強調 Google Cloud 上的 TPU,我們正在加倍押注這項開源策略。
事實上,我們不僅開源了 JAX 核心(core),還建構了許多高階函式庫,例如我們如何進行強化學習(RL)——我們稱之為 TuneX,它其實就是建立在 JAX 核心之上。我們將其開源。如何進行檢查點儲存(checkpointing)、如何處理推論事項:我們在 JAX 之上建構了大量的高階函式庫與框架,並且正在將其開源,因為我們希望人們使用它們,並利用這項開源策略直接使用 TPU。這是第一點。
當然,我們希望滿足客戶的需求,對吧?所以我們也希望擁抱開源社群中現有的任何成熟產品。PyTorch 非常成熟,在社群中被廣泛使用。你知道嗎?讓我們在 TPU 上擁抱 PyTorch。我們實際上正在進行一個名為 Torch-TPU 的專案。目前它處於私人預覽階段(private preview)。我們希望在下個季度(也就是短短幾個月後)推出公開預覽(public preview),然後在第四季於 GitHub 上推出公開版本,讓每個人都能使用。正如我昨天提到的,如果使用 Torch-TPU,只需修改幾行簡單的程式碼:將後端裝置設為 TPU,接著希望所有的訓練與推論都能在 TPU 上執行。我們實際上同時支援「即時模式」(eager mode)與「編譯模式」(compiled mode)。這是框架層面的部分。
我們同時也積極與社群互動——我確信台下有很多 vLLM 和 SGLang 的使用者。我們也與 vLLM 和 SGLang 進行交流,確保這類高階的開源推論框架也能在 TPU 上使用。所以我們目前有大量的工作在進行:我們開源自己的東西,同時也擁抱社群中成熟的開源產品。
以機器學習自動化硬體設計
主持人:我們的聽眾中確實也有很多大一和大二的博士生。你們有多少人是博士生?他們應該聚焦在哪些事情上?順便告訴大家,其中有些人還是大學生呢。
Jeff Dean:好的。處在你職業生涯的這個階段總是令人興奮的,因為我認為你可以找到自己真正喜愛且認為重要的事情,並找出推動某個領域前進的方法,在世界上產生影響力。我覺得那是一段非常好玩的時光。
主持人:單從基礎設施的角度來看,有沒有哪一個物理瓶頸是你在思考的?
Jeff Dean:有太多了。我非常看好越來越多的專用硬體(specialized hardware),因為我認為這是打造更高效率系統的唯一途徑。我們現在擁有的工作負載中,有一小部分工作負載將佔據世界算力的很大一部分,對吧?因此,如果你思考一下,這正是對專用化(specialization)的強烈呼喚。而專用化的問題在於,如果你未來想做的事情改變了,那麼你花費兩年心血精心打造的硬體可能就不再那麼合用了。
因此,要讓專用化真正發揮作用,你必須自動化更多的硬體設計流程。目前,設計硬體的方式是集結一大群人。其中一些人拿著高階規格書並從中建構低階的 RTL(暫存器傳輸級)。接著,因為這是人工手動翻譯的,你還需要另一群人來驗證第一群人做的事是否正確。然後你還有另一組人來精心佈局(lay out)晶片。這聽起來像是:如果你能建立更多可透過強化學習或其他演進式技術進行搜尋的自動化迴圈(automated loops),並且能讓這些迴圈跑得夠快——目前的 EDA 工具通常不具備快速執行的設計——那麼你就有機會在設計流程中擁有更自動化的探索迴圈,並或許能以驚人的幅度縮短設計週期。
比方說,如果你能用 10 個人在 3 個月內設計出一款新晶片,而不是用 150 個人花 2 年時間,那麼世界上的專用硬體就會多得多。而且你對未來的押注也不需要那麼冒險:你究竟想要 2 到 6 年後的算力,會變成更像是 3 到 6 個月到 4 年後的算力,這是一個更容易下注的方向。
推論工作負載、智慧代理系統與工具瓶頸
Bill Jia:如果我們回顧一兩年前,當時整個產業(包括 Google 在內以及許多前沿實驗室)的主要焦點都是如何讓模型變得無比強大。大家把大量精力放在預訓練和後訓練(post-training)上,因此許多硬體策略都偏向訓練端。但現在,隨著大模型走向成熟,大量的流量開始湧向代理(agent)世界以及推論端。因此,Jeff,如果你能從硬體設計的角度分享一些想法:我們該如何設計硬體來聚焦並強化推論流量?
Jeff Dean:推論和訓練有些許不同。對於推論來說,你基本上擁有一個模型,而你的目標是為它服務大量的請求。因此,你需要盡可能少移動那些不會改變的資訊。不會改變的東西是模型權重(model weights),會改變的東西則是請求(requests)和 KV 快取(KV cache)等等。所以你需要設計出一個非常擅長將資料搬移降到最低的系統。
此外我也認為,當你看到的不再只是單純的「提示詞(prompt)接回應」,而是更獨立的互動——代理執行某些動作、決定呼叫工具、工具執行、工具結果回傳,然後你把這些整合到模型的上下文中,讓它決定下一步該做什麼——我們會意識到,我們所有的工具都慢得離譜,因為它們本來就是為人類速度的迭代設計的。如果你在工具執行過程中需要進行程式碼編譯,你會發現編譯器太慢了會令人十分沮喪。如果你把推論硬體做得超級快,它生成程式碼的速度可能甚至比你編譯它的速度還要快,更不用說比你執行它的速度快得多了。
所以我們一直在做的一件事,就是讓我們的內部工具變得更快,好讓它們變得更好。事實證明,你其實可以非常有效率地將一種程式設計語言翻譯成另一種程式設計語言,因為你手中握有的正是你想要的完整、充分指定的規格書。你擁有用像 Python 這種直譯式語言寫成的整個程式,而你只需要在 Go、Rust、C++ 或其他語言中獲得完全等效的程式。代理其實可以把這件事做得相當好,這與我們平時和代理進行的編碼互動有很大的不同——平常我們是說:「請幫我做一個網頁伺服器」,然後代理必須填入各種充滿假設的細節,而這些假設可能根本不是你想要的。但有了充分指定的工具,它就能把這件事做得非常好:它能夠執行所有單元測試(unit tests)、翻譯所有單元測試、在新的系統中執行它們,並進行並排(side-by-side)行為驗證,以確保其行為完全一致。
Bill Jia:我們內部其實有一個活生生的例子,我們稱之為 Project Tern(燕鷗專案)。因為許多模型原本是建構在 TensorFlow 上的,但我們正在遷移到 JAX,所以我們利用這個專案自動將所有 TensorFlow 模型翻譯並遷移到 JAX,同時自動執行所有的單元測試與程式碼測試。
主持人:說到這個,你們之前向 NeurIPS 提交了一篇被拒絕的論文,內容提到效果甚微。那是一篇關於知識蒸餾(distillation)的論文。
Jeff Dean:噢,對。這是一個激勵人心的故事,本意是要給人啟發的。事實證明,蒸餾非常重要。這是我和我的同事 Geoff Hinton、Oriol Vinyals 一起提交的論文,探討如何取得一個大型模型,並透過蒸餾(distillation)把它濃縮到一個學生模型(student model)中。
Bill Jia:那是哪一年的事?
Jeff Dean:那是 2015 年或 2016 年。我們最初的構想是將其應用於訓練一個由各種專用視覺模型組成的大型整合模型(ensemble)。那是我們在論文中進行的一系列實驗之一:你擁有 2,000 個視覺類別,但你為動物訓練一個專用模型、為汽車訓練另一個、為其他東西再訓練一個,然後把這些全部蒸餾成一個單一模型,讓它擅長所有這些事情。結果它被拒絕了,不過沒關係;我們把它發到了 arXiv 上。大家反正還是讀了。
主持人:所以千萬別氣餒。
Jeff Dean:完全正確。
Q&A:通用硬體 vs. 專用超級電腦
主持人:好的,我們還剩下幾分鐘。我們來回答幾個問題。現場有幾位工作人員正拿著麥克風。找到了,這邊有一位。我們從這裡開始好嗎?請站起來自我介紹一下,並簡單簡述你的背景。
聽眾(Ying):大家好,我叫 Ying。我曾在 Google Brain 工作了十年。很高興再次見到你們。我想請教 Bill 和 Jeff:我們剛才談到 Google 早期的通用軟體(commodity software)與通用算力(commodity compute);這正是 Google 在 2000 年代初期實現規模化的方式。如今,AI 超級電腦變得越來越專門化,更像是原本的超級運算框架。那麼,你們如何看待硬體與機器學習系統這兩種不同的設計模式?你們認為未來我們應該更多地轉向通用硬體,還是應該繼續打造超級運算機器?謝謝。
Jeff Dean:我們當初能夠靠通用硬體來擴展搜尋規模的原因在於,搜尋是一個非常美妙的問題:如果你把它拆解開來,你會發現機器之間的通訊幾乎為零,而且大部分工作都在單機上完成。因此,你在互連方面不需要任何奇特的東西。我們當初在機器上使用的是 100M 乙太網路(100-megabit Ethernet),然後每 40 台機器共享一個 1Gbps 的機架上行鏈路(uplink),所以實際上大約有 4 比 1 的過度認購率(oversubscription)。但這沒關係,因為你傳給每台機器的只是像「巴洛阿圖餐廳」這樣的查詢字串,然後回傳的也只是一小段包含這 10 個搜尋結果的摘要。然而,訓練則真正推動了極限:你需要大量的連通性。
你當然希望能夠在單一晶片上完成訓練——那會是最理想的狀態——但這實在太花時間,或者根本塞不下。因此,你最終不得不把問題分割到許多晶片上。無論你怎麼切分——模型平行化或資料平行化——你通常都會面臨相當大量的通訊。這就是為什麼機器之間需要這些更奇特的互連技術。你希望每個晶片的效能達到最大化,這樣你就不必把任務分攤到過多的晶片上,但你依然必須把它分攤到很多晶片上。這就是為什麼我們正在使用液冷技術,來確保每個晶片在這種情境下的效能都能發揮到極致。
我認為推論的走向會更偏向:你可以擁有專用硬體,但它沒那麼奇特,特別是對小型模型而言。但隨著模型規模越來越大,你開始面臨通訊需求,相較於乙太網路通用硬體與通用網路,這些需求開始顯得相當奇特,不過它絕對比訓練更加主流和通用。這樣講合理嗎?
Q&A:開源貢獻、CUDA 對等性與面向開發者的 AI 代理
主持人:我想這裡還有另一個問題。
聽眾(Andra):Jeff,你好,我是來自 Uber 的 Andra。我有一個關於 JAX 和 OpenXLA 生態系的問題,特別是考慮到人們已經圍繞這個建立起開源社群。我看到這幾乎是一個類似 Android 和 iOS 的 AI 基礎設施時代。就年輕工程師如何為 OpenXLA 的函式庫功能做出貢獻而言,你有什麼看法?因為對大多數使用者來說,他們都是重度 CUDA 使用者,或者他們是基於 CUDA 函式庫來建構函式庫或功能。對於 OpenXLA,你認為年輕工程師可以做出什麼樣的貢獻?
Bill Jia:首先,OpenXLA 和 JAX,以及 JAX 之上的 JAX 堆疊已經完全開源了。這是第一大塊工作。另一大塊工作是 Torch-TPU。我們目前在內部正與 Meta 合作開發,但最終我們會把它放到 GitHub 上,讓所有社群的研究人員和工程師都能參與貢獻。這是我們的兩大並行工作。社群中任何想要為開源儲存庫做出貢獻的年輕工程師,我們都非常歡迎與 Google 合作。我們可以討論如何共同開發這個開源儲存庫。
此外,我認為最終許多使用者都會在 TPU 或 GPU 上使用 JAX 堆疊或 PyTorch 堆疊。目前,如果社群中的人們在使用時遇到任何問題,很多人會打電話給 Google 或 Nvidia,因為他們有經驗豐富的工程師,或者他們會把問題發到社群討論板上,希望有人看得懂問題並擁有類似經驗來回答。但在我看來,這實在太慢了。社群能在開源策略上做出貢獻的地方在於:想像一下,如果我們能有一個在 GPU 或 TPU 上執行的 JAX 代理和 PyTorch 代理——這都無所謂。如果社群中的每個人都貢獻資料來源,來協助訓練出一個非常強大的代理,那麼想像一下,如果我擁有一個 TPU 同事代理。當我執行訓練、後訓練或推論時,我不需要去問 Google 或逛討論板,我可以直接使用這個代理來協助我。如果這個代理能夠回答很大一部分問題,那就太棒了。但这需要全社群共同貢獻,因為我們需要大量的資料點和經驗。
Jeff Dean:或許我可以就開源貢獻補充一個統合的觀點。大家有很多種方式可以做出貢獻。主動與不同儲存庫的擁有者互動,試著釐清:「我正在考慮做這個,這會有用嗎?或者你們對我有什麼建議可以著手進行的工作嗎?」這是一個很好的習慣。在你不分青紅皂白地丟給別人 5,000 行程式碼並說「嘿,拿去」之前,先獲得一些認同(buy-in)或關於什麼才是有用之物的指導,我認為是非常明智的做法。如果你渴望貢獻,但心裡未必有特定想做的事,大家通常都有一長串認為很有價值、正在等待有人去做的候選清單。
聽眾(Andra):謝謝你的回答。我所遇到的其中一個遲疑之處在於,許多功能或模組核心(module kernels)仍然極度依賴 CUDA。你是否看到使用 JAX、OpenXLA 或 Torch-TPU 來實作某些等效功能的趨勢,好讓我們能擁有某種相較於 CUDA 功能的泛用相容性?這會是一個好的方向嗎?
Bill Jia:是的,我認為最終會朝這個方向發展。在 Google 內部,當我們審視內部的 JAX 堆疊以及 TPU 上的 PyTorch 堆疊時,我們會檢視所有的 CUDA 功能,同時也會檢視我們自己的功能。我們正在確保我們至少能做到不相上下,甚至更好。所以這是我們必須攻克的必要任務。而在那個層級上,我們實際上也希望開源低階的 TPU SDK 工作。
Jeff Dean:我也認為人們應該在更高的抽象層級上思考。以 JAX 或 PyTorch 表達式的方式來思考才是正道,而不一定是去想:「我該如何用某種核心程式碼(無論是 TPU 核心語言還是 GPU 核心語言)來平行化這個運算,以獲得最高效能?」希望編譯器和底層系統能夠為你搞定這些,而你只需要專注於像矩陣乘法這樣美妙的抽象概念即可。
Q&A:擴展網路、3D 環面與同步訓練
主持人:我想我們還有時間問最後一個問題。就最後一個問題,請說。
聽眾(John):嗨,我是來自紐約大學(NYU)的 John。我多年來一直從事網路研究。Bill Jia 先生以及另一位講者先前關於擴展性的演講讓我印象深刻。我知道這些 TPU 是以 3D 環面(3D torus)相連的。相較於超乙太網路聯盟(Ultra Ethernet Consortium)和其他正在研究脂肪樹(fat-tree)互連技術的團隊,從擴展性的角度來看,當你們把 TPU 或 GPU 擴展到數萬個甚至高達 100 萬個時,這個 3D 環面還能夠在如此龐大的規模下處理全對全(all-to-all)通訊嗎?
Jeff Dean:我們針對 TPU v5e 或 TPU v4 的最大 Pod,據我所知大約是 9,600 多個晶片。這就是我們採用 3D 環面架構的規模極限。在此之上,我們以 Pathways 作為軟體抽象層,疊加在許多由 3D 環面連接的 Pod 之上,而這會利用資料中心網路——也就是你資料中心裡擁有的任何網路織網(fabric),甚至包括多都會區訓練設置(例如你在奧克拉荷馬州的某個大樓裡有 5 個 Pod,在愛荷華州的另一棟大樓裡有 8 個 Pod,它們之間透過高速 WAN 連結)。這種方式在我們的擴展經驗中運作得相當不錯。
有時候你確實會希望對正在執行的運算進行對應配置,以便擁有資料平行(data-parallel)與模型平行(model-parallel)維度。你通常會希望模型平行的面向集中在單一 Pod 或 Pod 的某個切片(slice)內,然後在這些 Pod 或 Pod 切片之間擁有資料平行的複本。這對我們來說運作得相當順暢。基於環面(torus)的網路的好處在於,它們在本地連接時極為簡單,而不像在資料中心地板上需要極其複雜的配線。
主持人:除了那些有趣的鏡面光學開關例外。
Jeff Dean:Bill,有什麼要補充的嗎?
Bill Jia:對,這完全屬實。我們在機架內部、在立方體內部擁有 ICI。然後我們使用 OCS 將規模擴大到 9,600 個。再往上一層,我們擁有由 Pathways 軟體協調的資料中心織網。要超越 100,000 個晶片,我們會利用雲端資料中心網路將所有東西連接在一起。這種擴展性對我們來說運作得相當良好。
Jeff Dean:我想補充的是,即使在這樣的規模下,我們依然能夠做到完全同步訓練(fully synchronous training),這對機器學習的可重現性(reproducibility)和可解釋性(interpretation)來說是一大優勢。或許在某些時間點非同步訓練會捲土重來,但到目前為止,我們已經能夠把同步訓練推向相當遠的境界,我們拭目以待接下來會如何發展。
主持人:這真是個完美的結尾。我們的時間到了。非常謝謝 Jeff。謝謝 Bill 百忙之中抽空參與。我們在外面露台為大家準備了一些非常酷的紀念品,其中包括印有知名 Jeff Dean 迷因的原子筆以及一件 Dev Labs T 恤。非常謝謝大家!