我扒出了Gemini 2.0超實(shí)時(shí)多模態(tài)幕后的黑科技，第六代TPU芯片Trillium！

最近真是 AI 圈的大年三十初一初二初啊。。。。。各個(gè)大廠都爭先恐后的整花活！OpenAI 放出了 o1 Pro、GPT-4o 高級語音、GPTCanavas，就跟孔雀開屏一樣 ～谷歌最近的大動(dòng)作是發(fā)布了 Gemini 2.0 嘛！2.0 比 1.5 版本快一倍，而且是原生的多模態(tài)大模型，能輸入和生成語言、聲音、圖片、視頻等。而且，最受大家震撼的是超級超級低的實(shí)時(shí)，無卡頓的多模態(tài)交互！只能說非常牛，小瑤灰常好奇究竟咋做到的，然后就去扒了下訓(xùn)練背后的故事，結(jié)果！小瑤發(fā)現(xiàn)谷歌真正牛的是用來訓(xùn)練 Gemini 2.0 的芯片 Trillium！即第 6 代 TPU，比上一代 TPU 5e 性能提升可以說非常非常顯著了！訓(xùn)練效率提高 4 倍以上推理吞吐量提高 3 倍能源效率提高 67%每塊芯片的峰值計(jì)算性能顯著提升 4.7 倍高帶寬內(nèi)存 (HBM) 容量加倍-芯片間互連 (ICI) 帶寬加倍單個(gè) Jupiter 網(wǎng)絡(luò)結(jié)構(gòu)中有 10 萬個(gè) Trillium 芯片每美元可將訓(xùn)練性能提高 2.5 倍，每美元可將推理性能提高 1.4 倍這些硬件上的進(jìn)步都給 AI 訓(xùn)練帶來了實(shí)打?qū)嵉暮锰帲“耍簲U(kuò)展 AI 訓(xùn)練工作負(fù)載訓(xùn)練 LLM，包括密集模型和混合專家 (MoE) 模型推理性能和收集調(diào)度嵌入密集型模型提供訓(xùn)練和推理性價(jià)比要讓像 Gemini 2.0 這樣的多模態(tài)大語言模型加速訓(xùn)練，首先需要更大規(guī)模的數(shù)據(jù)和更強(qiáng)大的計(jì)算資源。Trillium 的技術(shù)可以視為一名超級助手：它將龐大而復(fù)雜的計(jì)算任務(wù)分散到眾多主機(jī)中，并通過高速 Jupiter 數(shù)據(jù)中心網(wǎng)絡(luò) 緊密相連。這種協(xié)同工作就好像將 256 個(gè)芯片聚合成一個(gè)緊密合作的“大家族”。在加速訓(xùn)練的過程中，Trillium 借助名為 “TPU 多切片（TPU multi-slicing）” 的技術(shù)，使得大規(guī)模訓(xùn)練變得更加高效。與此同時(shí)，還有一個(gè)名為 “Titanium” 的系統(tǒng)，它能從主機(jī)適配器到網(wǎng)絡(luò)架構(gòu)層面全面支持動(dòng)態(tài)任務(wù)卸載，確保整個(gè)數(shù)據(jù)中心的運(yùn)行更加順暢。在實(shí)際應(yīng)用中，Trillium 在一個(gè)由 12 個(gè)這樣的“大家族”構(gòu)成、共計(jì) 3072 個(gè)芯片的環(huán)境中，實(shí)現(xiàn)了 99% 的擴(kuò)展效率；在一個(gè)更大的 24 個(gè)“家族”、共 6144 個(gè)芯片的部署中，也依然達(dá)到 94% 的擴(kuò)展效率。這意味著，即便是在訓(xùn)練如 GPT3-175b 這樣超大規(guī)模的模型時(shí)，無論是在單一數(shù)據(jù)中心還是跨數(shù)據(jù)中心環(huán)境中，Trillium 都能以極高的效率推動(dòng)模型訓(xùn)練進(jìn)程。官方鏈接：https://cloud.google.com/blog/products/compute/trillium-tpu-is-ga快來和小瑤一起來看看具體的提升！擴(kuò)展效率大提升！如下圖所示，即使使用 1 片 Trillium-256 芯片的 Trillium-256 芯片艙作為基線，在擴(kuò)展到 24 個(gè)芯片艙時(shí)，仍然能夠達(dá)到超過 90% 的擴(kuò)展效率。在訓(xùn)練 Llama-2-70B 模型時(shí)，谷歌測試表明，Trillium 從 4 片 Trillium-256 芯片艙擴(kuò)展到 36 片 Trillium-256 芯片艙時(shí)，幾乎實(shí)現(xiàn)了線性擴(kuò)展，擴(kuò)展效率達(dá)到了 99%還有一個(gè)更直觀的看法就是，將 Trillium TPU 與 Google 前代產(chǎn)品（v5p）對比一下！在下面的圖表所示，在 12-pod 規(guī)模下，Trillium 的擴(kuò)展效率達(dá)到了 99%，與同等規(guī)模的 Cloud TPU v5p 集群（總峰值浮點(diǎn)運(yùn)算次數(shù)）相比遙遙領(lǐng)先。擅長訓(xùn)練密集型和專家混合（MoE）模型！像 Gemini 這樣的大語言模型（LLM）本身就具有數(shù)十億個(gè)參數(shù)，先天強(qiáng)大且復(fù)雜。要高效訓(xùn)練這些密集型的大模型，不僅需要雄厚的計(jì)算能力，還離不開軟硬件協(xié)同優(yōu)化的解決方案。Trillium 在訓(xùn)練諸如 Llama-2-70b 和 gpt3-175b 這類極其龐大的 LLM 時(shí)，相較于上一代 Cloud TPU v5e 能實(shí)現(xiàn)高達(dá) 4 倍 的加速性能。如下圖所示：此外，除了傳統(tǒng)的“大模型”之外，采用 專家混合（MoE） 架構(gòu)的大語言模型正日趨流行。這類架構(gòu)由多個(gè)“專家”神經(jīng)網(wǎng)絡(luò)組成，每個(gè)專家擅長處理特定的任務(wù)領(lǐng)域。然而，與訓(xùn)練一個(gè)單一、龐大的模型相比，對這些專家進(jìn)行有效協(xié)同與管理，無疑提高了訓(xùn)練的復(fù)雜性。即便如此，Trillium 在訓(xùn)練 MoE 模型 時(shí)也依舊展現(xiàn)出強(qiáng)大的性能提升，如下圖所示，相比上一代 Cloud TPU v5e 能快上 3.8 倍。推理性能與調(diào)度優(yōu)化大升級！面來越重要的多步推理（multi-step inference）需求，具有更高處理效率的加速器變得至關(guān)重要。Trillium 加速并優(yōu)化了 AI 模型的部署，因此它在圖像擴(kuò)散模型和密集型大語言模型的推理性能上均表現(xiàn)出色。如下圖所示，相較 Cloud TPU v5e，使用 Trillium 對 Stable Diffusion XL （SDXL）進(jìn)行推理時(shí)，每秒圖像吞吐量提升超過 3 倍；在 Llama2-70B 的推理中，每秒 Token 吞吐量提升近 2 倍。嵌入密集型模型隨著第三代 SparseCore 的引入 Trillium ，其在處理嵌入密集型模型（embedding-heavy models）時(shí)的性能提升了 2 倍，在 DLRM DCNv2 模型上的性能更是提升了 5 倍。SparseCore 是為嵌入密集型工作負(fù)載打造的數(shù)據(jù)流處理器架構(gòu)，能夠靈活適應(yīng)多樣化的計(jì)算需求。它擅長處理動(dòng)態(tài)且數(shù)據(jù)相關(guān)的操作，例如散布-收集（scatter-gather）、稀疏段求和（sparse segment sum）以及分區(qū)（partitioning），從而在動(dòng)態(tài)數(shù)據(jù)環(huán)境中保持高效運(yùn)轉(zhuǎn)。因此，Trillium 現(xiàn)在不僅在規(guī)模和性能上滿足訓(xùn)練當(dāng)今最大 AI 工作負(fù)載的需求，還注重優(yōu)化每美元性能。在實(shí)際訓(xùn)練中，相較于 Cloud TPU v5e，Trillium 在訓(xùn)練密集型 大語言模型（如 Llama2-70b 和 Llama3.1-405b）時(shí)，可實(shí)現(xiàn)高達(dá) 2.1 倍 的性能提升；而與 Cloud TPU v5p 相比，則可獲得 2.5 倍 的性能提升。在 Trillium 上生成 1000 張圖像的成本比 Cloud TPU v5e 離線推理降低 27%，比 Cloud TPU v5e 在 SDXL 上的服務(wù)器推理成本降低 22%。這意味著用戶可以在保證高性能的同時(shí)，以更合理的投入獲得更高的產(chǎn)出。結(jié)語在目前人工智能創(chuàng)新不斷攀升的時(shí)代，Trillium 再次證明了谷歌云在 AI 基礎(chǔ)設(shè)施領(lǐng)域的獨(dú)特優(yōu)勢。相比之下，OpenAI 數(shù)次嘗試自研芯片卻碰壁收場，至今進(jìn)度成謎。。。而谷歌不僅能擴(kuò)展至數(shù)十萬顆芯片的龐大規(guī)模，更能借助軟硬件協(xié)同優(yōu)化，讓 Trillium 成為行業(yè)新標(biāo)桿。不得不說。。底蘊(yùn)深厚的老牌大廠自有其過人之處，谷歌仍是那個(gè)讓人心服口服的科技巨頭！