?十年來，機器學(xué)習(xí)軟件開發(fā)的格局發(fā)生了重大變化。許多框架如雨后春筍般涌現(xiàn)，但大多數(shù)都嚴(yán)重依賴于英偉達的 CUDA，并在英偉達的 GPU 上才能獲得最佳的性能。然而，隨著 PyTorch 2.0 和 OpenAI Triton 的到來，英偉達在這一領(lǐng)域的主導(dǎo)地位正在被打破。

谷歌早期在機器學(xué)習(xí)模型架構(gòu)、訓(xùn)練、模型優(yōu)化方面都具有很大優(yōu)勢，但現(xiàn)在卻難以充分發(fā)揮這些優(yōu)勢。而在硬件方面，其他 AI 硬件公司很難削弱英偉達的統(tǒng)治地位。直到 PyTorch 2.0 和 OpenAI Triton 出現(xiàn)，機器學(xué)習(xí)模型的默認軟件堆棧將不再是英偉達的閉源 CUDA。

TensorFlow vs. PyTorch

類似的競爭出現(xiàn)在機器學(xué)習(xí)框架中。幾年前，框架生態(tài)系統(tǒng)相當(dāng)分散，但 TensorFlow 是領(lǐng)跑者。表面看來，谷歌穩(wěn)操機器學(xué)習(xí)框架行業(yè)，他們憑借 TensorFlow 設(shè)計了 AI 應(yīng)用特定加速器 TPU，從而獲得了先發(fā)優(yōu)勢。

然而現(xiàn)在看來，PyTorch 贏了，谷歌未能將其先發(fā)優(yōu)勢轉(zhuǎn)化為新興 ML 行業(yè)的主導(dǎo)地位。如今，谷歌在機器學(xué)習(xí)社區(qū)中似乎有些被孤立了，因為它沒有使用 PyTorch 和 GPU，而是使用自己的軟件堆棧和硬件。甚至，谷歌研發(fā)了第二個機器學(xué)習(xí)框架 ——JAX，直接與 TensorFlow 競爭，這是典型的「谷歌行為」。

一些人認為，由于大型語言模型的興起，特別是 OpenAI 的大型語言模型和各種利用 OpenAI API 構(gòu)建的語言模型，谷歌在搜索和自然語言處理方面的主導(dǎo)地位正在減弱。也許這種觀點過于悲觀，畢竟當(dāng)前大多數(shù)模型的基礎(chǔ)架構(gòu)仍是谷歌開發(fā)的 transformer。

那么，為什么 PyTorch 能大獲全勝呢？主要原因是相對于 TensorFlow，PyTorch 具有更高的靈活性和可用性。PyTorch 與 TensorFlow 主要的不同之處在于使用 Eager 模式而非 Graph 模式。

Eager 模式可以說是標(biāo)準(zhǔn)的腳本執(zhí)行方法，與普通的 Python 代碼沒什么區(qū)別。這使得調(diào)試和理解代碼更加容易，因為用戶可以看到中間操作的結(jié)果，以及模型是如何運行的。

相反，Graph 模式分為兩個階段。第一階段代表要執(zhí)行操作的計算圖，其中的節(jié)點代表操作或變量，而節(jié)點之間的邊表示它們之間的數(shù)據(jù)流。第二階段是延遲執(zhí)行計算圖的優(yōu)化版本。

這種兩階段的方法使得理解和調(diào)試代碼更具挑戰(zhàn)性，因為在圖執(zhí)行結(jié)束之前用戶無法看到發(fā)生了什么。這類似于「解釋型」與「編譯型」語言，如 Python 與 C++，調(diào)試 Python 更容易，因為它是解釋型語言。

雖然 TensorFlow 現(xiàn)在也默認使用 Eager 模式，但研究社區(qū)和大多數(shù)大型科技公司都選擇使用 PyTorch。

機器學(xué)習(xí)訓(xùn)練組件

如果將機器學(xué)習(xí)模型訓(xùn)練簡化為最簡單的形式，影響機器學(xué)習(xí)模型訓(xùn)練的因素主要有兩點：

• 計算 ( FLOPS ) ：在每一層內(nèi)運行密集矩陣乘法；
• 內(nèi)存帶寬。

以前，影響機器學(xué)習(xí)訓(xùn)練時間的主要因素是計算時間，等待系統(tǒng)執(zhí)行矩陣乘法。隨著英偉達 GPU 的不斷發(fā)展，這很快就不再是主要問題了。

英偉達利用摩爾定律將 FLOPS 提高了多個數(shù)量級，但主要是架構(gòu)變化 —— 張量核（tensor core）和更低精度的浮點數(shù)格式。相比之下，存儲方面沒有太大的變化。

2018 年，最先進的模型是 BERT，英偉達 V100 是最先進的 GPU，那時矩陣乘法已經(jīng)不再是提高模型性能的主要因素。之后，模型在參數(shù)數(shù)量上增長了 3 到 4 個數(shù)量級，而最快的 GPU 在 FLOPS 上增長了 1 個數(shù)量級。

即使在 2018 年，純計算綁定的工作負載也占 FLOPS 的 99.8%，但僅占運行時的 61%。與矩陣乘法相比，歸一化和逐點運算（pointwise ops）使用的 FLOPS 僅為矩陣乘法的 1/250 和 1/700，但它們消耗了近 40% 的模型運行時間。

內(nèi)存墻

隨著模型規(guī)模的不斷飆升，大型語言模型（LLM）僅用于模型權(quán)重的內(nèi)存就需要 100 GB 以上。百度和 Meta 部署的產(chǎn)品推薦網(wǎng)絡(luò)需要數(shù)十 TB 的內(nèi)存來存儲其海量嵌入表。大型模型訓(xùn)練 / 推理中的大部分時間都沒有花在計算矩陣乘法上，而是在等待數(shù)據(jù)傳輸。顯然，問題在于為什么架構(gòu)師不將更多內(nèi)存放在更靠近計算的位置，問題的答案也是可想而知的 —— 成本。

最近的共享內(nèi)存池一般是同一芯片上的 SRAM。一些機器學(xué)習(xí) ASIC 試圖利用巨大的 SRAM 池來保存模型權(quán)重。但即使是 Cerebras 的價值約 5000000 美元的晶圓級芯片也只有 40GB 的 SRAM。內(nèi)存容量不足以容納 100B+ 參數(shù)模型的權(quán)重。

英偉達設(shè)計的芯片片上內(nèi)存容量要少得多 ——A100 有 40MB，H100 有 50MB。臺積電 5 納米芯片上的 1GB SRAM 需要約 200 平方毫米的硅，要實現(xiàn)相關(guān)的控制邏輯 / 結(jié)構(gòu)，將需要超過 400 平方毫米的硅。鑒于 A100 GPU 的成本在 1 萬美元以上，而 H100 更是接近 2 萬美元，從經(jīng)濟角度看，這種方法是不可行的。即使忽略英偉達在數(shù)據(jù)中心 GPU 上約 75% 的利潤率，對于完全量產(chǎn)的產(chǎn)品，SRAM 內(nèi)存的成本仍在 100 美元 / GB 左右。

此外，片上 SRAM 存儲器的成本不會隨著傳統(tǒng)摩爾定律工藝技術(shù)的縮小而降低太多。同樣的 1GB 內(nèi)存，采用臺積電下一代 3nm 制程工藝，成本反而更高。雖然 3D SRAM 將在一定程度上幫助降低 SRAM 成本，但這只是暫時的。

內(nèi)存層次結(jié)構(gòu)的下一步是緊密耦合的片外內(nèi)存 DRAM。DRAM 的延遲比 SRAM 高一個數(shù)量級（約 100nsVS10ns），但它也便宜得多。幾十年來，DRAM 一直遵循著摩爾定律。戈登摩爾創(chuàng)造這個詞時，英特爾的主要業(yè)務(wù)就是 DRAM。他對晶體管密度和成本的預(yù)測在 2009 年之前對 DRAM 普遍適用。不過自 2012 年以來，DRAM 的成本幾乎沒有改善。

然而，人們對內(nèi)存的需求只增不減。DRAM 現(xiàn)在占服務(wù)器總成本的 50%，逐漸形成所謂的「內(nèi)存墻」。將英偉達 2016 年的 P100 GPU 與最新的 H100 GPU 進行比較，我們會發(fā)現(xiàn)內(nèi)存容量增加到原來的 5 倍（16GB → 80GB），F(xiàn)P16 性能增加到 46 倍（21.2 TFLOPS → 989.5 TFLOPS）。

雖然內(nèi)存容量是一個重要瓶頸，但另一個瓶頸 —— 內(nèi)存帶寬也非常關(guān)鍵。內(nèi)存帶寬的增加通常是通過并行性獲得的。雖然如今標(biāo)準(zhǔn) DRAM 的價格僅為幾美元 / GB，但為了獲得機器學(xué)習(xí)所需的海量帶寬，英偉達使用 HBM 內(nèi)存 —— 一種由 3D 堆疊 DRAM 層組成的設(shè)備，需要更昂貴的封裝。HBM 的成本大概是 10-20 美元 / GB，包括封裝和產(chǎn)量成本。

內(nèi)存帶寬和容量的成本限制問題在英偉達的 A100 GPU 中尤為明顯。如果不進行大量優(yōu)化，A100 只能具有非常低的 FLOPS 利用率。

即使研究人員進行了大量優(yōu)化，大型語言模型的 FLOPS 利用率也只能達到 60% 左右 。很大一部分時間花費在等待來自另一個計算 / 內(nèi)存的數(shù)據(jù)，或者及時重新計算結(jié)果以減少內(nèi)存瓶頸。

從 A100 到 H100，F(xiàn)LOPS 增長到原來的 6 倍以上，但內(nèi)存帶寬僅增長到 1.65 倍。這導(dǎo)致許多人擔(dān)心 H100 的利用率會很低。A100 需要很多技巧才能繞過內(nèi)存墻，而 H100 還需要實現(xiàn)更多技巧。

H100 為 Hopper 架構(gòu)帶來了分布式共享內(nèi)存和 L2 多播 ( multicast ) 。這個想法讓一個 SM 中的數(shù)據(jù)可以直接寫入另一個 SM 的 SRAM（共享內(nèi)存 /L1 Cache）中。這有效地增加了緩存的大小并減少了 DRAM 讀 / 寫所需的帶寬。未來的架構(gòu)將減少向內(nèi)存發(fā)送的操作數(shù)量，以最大限度地減少內(nèi)存墻的影響。值得注意的是，較大的模型往往會實現(xiàn)更高的利用率，因為 FLOPS 需要按參數(shù)量的三次方成比例擴展，而內(nèi)存帶寬和容量需求往往按二次方擴展。

算子融合

如果將所有時間都花在內(nèi)存?zhèn)鬏斏希刺幱趦?nèi)存帶寬限制狀態(tài)），那么增加 GPU 的 FLOPS 將無濟于事。另一方面，如果將所有時間都花在執(zhí)行大型 matmuls 上，那么即使將模型邏輯重寫為 C++ 來減少開銷也將無濟于事。

PyTorch 之所以能勝過 TensorFlow，就是因為 Eager 模式提高了靈活性和可用性，但轉(zhuǎn)向 Eager 模式并不是只有好處。在 Eager 模式下運行時，每次運算都要從內(nèi)存中讀取、計算，然后在處理下一次運算之前發(fā)送到內(nèi)存。如果不進行大量優(yōu)化，這會顯著增加內(nèi)存帶寬需求。

因此對于在 Eager 模式下執(zhí)行的模型，有一種主要的優(yōu)化方法是算子融合。融合運算在一次傳遞中計算多個函數(shù)，以最小化內(nèi)存讀 / 寫，而不是將每個中間結(jié)果寫入內(nèi)存。算子融合改善了運算符調(diào)度、內(nèi)存帶寬和內(nèi)存大小成本。

這種優(yōu)化通常涉及編寫自定義 CUDA 內(nèi)核，但這比使用簡單的 Python 腳本要難得多。隨著時間的推移，PyTorch 中穩(wěn)定地實現(xiàn)了越來越多的算子，其中許多算子只是簡單地將多次常用運算融合到一個更復(fù)雜的函數(shù)中。

算子的增加讓在 PyTorch 中創(chuàng)建模型變得更容易，并且由于內(nèi)存讀 / 寫更少，Eager 模式的性能更快。缺點是 PyTorch 在幾年內(nèi)激增到了 2000 多個算子。

我們可以說軟件開發(fā)人員太懶了，但說實話，又有誰沒懶惰過呢。一旦習(xí)慣了 PyTorch 中的一個新算子，他們就會繼續(xù)用它。開發(fā)人員甚至可能沒有意識到性能在提高，而是繼續(xù)使用該算子，因為這樣就不用編寫更多的代碼。

此外，并非所有算子都可以融合。決定要融合哪些運算，將哪些運算分配給芯片和集群級別的特定計算資源都需要花費大量的時間。算子在何處融合的策略雖大體相似，但因為架構(gòu)的不同也會有很大差異。

英偉達曾是王者

算子的增長和默認的地位對英偉達來說是優(yōu)勢，因為每個算子都針對其架構(gòu)進行了快速優(yōu)化，但并未針對任何其他硬件進行優(yōu)化。如果一家 AI 硬件初創(chuàng)公司想要全面實施 PyTorch，那就意味著以高性能支持不斷增長的 2000 個算子列表。

因為提取到最大性能需要很多技巧，在 GPU 上訓(xùn)練具有高 FLOPS 利用率的大型模型所需的人才水平越來越高。Eager 模式執(zhí)行加算子融合意味著開發(fā)的軟件、技術(shù)和模型都在不斷地被推動，以適應(yīng)當(dāng)前一代 GPU 具有的計算和內(nèi)存比率。

每個開發(fā)機器學(xué)習(xí)芯片的人都受制于同一個內(nèi)存墻。ASIC 受制于支持最常用的框架，受制于默認的開發(fā)方法、GPU 優(yōu)化的 PyTorch 代碼以及英偉達和外部庫的混合。在這種情況下，避開 GPU 的各種非計算包袱而支持更多 FLOPS 和更嚴(yán)格的編程模型的架構(gòu)意義不大。

然而，易用性第一。打破惡性循環(huán)的唯一方法是讓在英偉達的 GPU 上運行模型的軟件盡可能輕松無縫轉(zhuǎn)移到其他硬件。隨著模型架構(gòu)的穩(wěn)定和來自 PyTorch 2.0、OpenAI Triton 和 MLOps 公司（如 MosaicML）的抽象成為默認，芯片解決方案的架構(gòu)和經(jīng)濟性開始成為購買的最大驅(qū)動力，而不是英偉達高級軟件提供的易用性。

PyTorch 2.0

幾個月前，PyTorch 基金會成立，并脫離了 Meta 。除了對開放式開發(fā)和治理模型的更改外，2.0 還發(fā)布了早期測試版本，并于 3 月全面上市。PyTorch 2.0 帶來了許多變化，但主要區(qū)別在于它添加了一個支持圖形執(zhí)行模型的編譯解決方案。這種轉(zhuǎn)變將使正確利用各種硬件資源變得更加容易。

PyTorch 2.0 在英偉達 A100 上的訓(xùn)練性能提升了 86%，在 CPU 上的推理性能提升了 26%。這大大減少了訓(xùn)練模型所需的計算時間和成本。這些好處可以擴展到來自 AMD、英特爾、Tenstorrent、Luminous Computing、特斯拉、谷歌、亞馬遜、微軟、Marvell、Meta、Graphcore、Cerebras、SambaNova 等的其他 GPU 和加速器。

對于當(dāng)前未優(yōu)化的硬件，PyTorch 2.0 具有更大的性能改進空間。Meta 和其他公司對 PyTorch 做出如此巨大的貢獻，是因為他們希望在自己價值數(shù)十億美元的 GPU 訓(xùn)練集群上以更少的努力實現(xiàn)更高的 FLOPS 利用率。這樣他們也有動力使軟件堆棧更易于移植到其他硬件，將競爭引入機器學(xué)習(xí)領(lǐng)域。

在更好的 API 的幫扶下，PyTorch 2.0 還可以支持?jǐn)?shù)據(jù)并行、分片、pipeline 并行和張量并行，為分布式訓(xùn)練帶來了進步。此外，它在整個堆棧中原生支持動態(tài)形狀，在許多其他示例中，這更容易支持 LLM 的不同序列長度。下圖是主要編譯器首次支持從訓(xùn)練到推理的 Dynamic Shapes：

PrimTorch

對于除英偉達 GPU 之外的每個機器學(xué)習(xí) ASIC 來說，為 PyTorch 編寫一個完全支持所有 2000 多個算子的高性能后端并非易事。PrimTorch 將算子的數(shù)量減少到約 250 個原始算子，同時還保持 PyTorch 最終用戶的可用性不變。PrimTorch 使 PyTorch 的不同非英偉達后端的實現(xiàn)變得更加簡單和易于訪問。定制硬件和系統(tǒng)供應(yīng)商可以更輕松地推出他們的軟件堆棧。

TorchDynamo

轉(zhuǎn)向圖模式需要可靠的圖定義。為了實現(xiàn)這一轉(zhuǎn)向，Meta 和 PyTorch 已經(jīng)嘗試了大約 5 年的時間，但是他們提出的每個解決方案都存在明顯的缺點。最后，他們用 TorchDynamo 破解了這個難題。TorchDynamo 將攝取任何 PyTorch 用戶腳本，包括調(diào)用外部第三方庫的腳本，并生成 FX 圖。

Dynamo 將所有復(fù)雜算子減少到 PrimTorch 中的約 250 個原始算子。一旦圖形成，未使用的算子將被丟棄，圖會決定哪些中間算子需要存儲或?qū)懭雰?nèi)存、哪些可能被融合。這極大地減少了模型內(nèi)的開銷，同時對用戶來說也是「無縫」的。

在測試的 7000 個 PyTorch 模型中，TorchDynamo 已經(jīng)適用于 99% 以上的模型，包括來自 OpenAI、HuggingFace、Meta、英偉達、Stability.AI 等的模型，而無需對原始代碼進行任何更改。測試的 7000 個模型是從 GitHub 上使用 PyTorch 的最受歡迎項目中隨機挑選出來的。

谷歌的 TensorFlow/Jax 和其他圖模式執(zhí)行 pipeline 通常要求用戶確保他們的模型適合編譯器架構(gòu)，以便可以捕獲圖。Dynamo 通過啟用部分圖捕獲、受保護的圖捕獲和即時重新捕獲來改變這一點。

部分圖捕獲允許模型包含不受支持的 / 非 python 構(gòu)造。當(dāng)無法為模型部分生成圖時，將插入圖中斷，并且將在部分圖之間以 eager 模式執(zhí)行不支持的構(gòu)造。

受保護的圖捕獲會檢查捕獲的圖是否對執(zhí)行有效。「保護」的意思是一種需要重新編譯的更改。這很重要，因為多次運行相同的代碼不會多次重新編譯。如果捕獲的圖對于執(zhí)行無效，則即時重新捕獲允許重新捕獲圖。

PyTorch 的目標(biāo)是創(chuàng)建一個具有流暢 UX 的統(tǒng)一前端，該前端利用 Dynamo 生成 graph。該解決方案的用戶體驗不會發(fā)生變化，但性能可以得到顯著提升。捕獲圖可以在大量計算資源上更有效地并行執(zhí)行。

隨后，Dynamo 和 AOT Autograd 將優(yōu)化的 FX 圖傳遞給 PyTorch 本機編譯器級別 TorchInductor。硬件公司也可以將此圖輸入到他們自己的后端編譯器中。

TorchInductor

TorchInductor 是 Python 原生深度學(xué)習(xí)編譯器，可以為多個加速器和后端生成快速代碼。Inductor 將采用具有約 250 個算子的 FX 圖，并將它們降低到約 50 個算子。接著，Inductor 進入調(diào)度階段，在該階段融合算子，并確定內(nèi)存規(guī)劃。

隨后，Inductor 進入「Wrapper Codegen」，它生成在 CPU、GPU 或其他 AI 加速器上運行的代碼。封裝器 Codegen 取代了編譯器堆棧的解釋器部分，可以調(diào)用內(nèi)核和分配內(nèi)存。后端代碼生成部分利用適用于 GPU 的 OpenAI Triton 并輸出 PTX 代碼。對于 CPU，英特爾編譯器生成 C++（也適用于非英特爾 CPU）。

未來他們將支持更多硬件，但關(guān)鍵是 Inductor 大大減少了編譯器團隊在為其 AI 硬件加速器制作編譯器時必須做的工作量。此外，代碼針對性能進行了更優(yōu)化，內(nèi)存帶寬和容量要求得到了顯著降低。

研究人員們需要的不是只支持 GPU 的編譯器，而是想要支持各種硬件后端。

OpenAI Triton

對英偉達的機器學(xué)習(xí)閉源軟件來說，OpenAI Triton 是一個顛覆性的存在。Triton 直接采用 Python 或通過 PyTorch Inductor 堆棧提供數(shù)據(jù)，后者是最常見的用法。Triton 負責(zé)將輸入轉(zhuǎn)換為 LLVM 中間表征，并生成代碼。英偉達 GPU 將直接生成 PTX 代碼，跳過英偉達的閉源 CUDA 庫（如 cuBLAS），轉(zhuǎn)而使用開源庫（如 cutlass）。

CUDA 在加速計算領(lǐng)域很受歡迎，但在機器學(xué)習(xí)研究人員和數(shù)據(jù)科學(xué)家中卻鮮為人知。使用 CUDA 可能會帶來重重挑戰(zhàn)，并且需要深入了解硬件架構(gòu)，這可能導(dǎo)致開發(fā)過程變慢。因此，機器學(xué)習(xí)專家可能就要依賴 CUDA 專家來修改、優(yōu)化和并行化他們的代碼。

Triton 彌補了這一缺陷，使高級語言實現(xiàn)了與低級語言相當(dāng)?shù)男阅堋riton 內(nèi)核本身對典型的 ML 研究者來說非常清晰，這對可用性來說非常重要。Triton 在 SM 中自動執(zhí)行內(nèi)存合并、共享內(nèi)存管理和調(diào)度。Triton 對逐元素矩陣乘法不是特別有用，但矩陣乘法已經(jīng)可以非常高效地完成。Triton 對于成本高昂的逐點運算和減少復(fù)雜操作的開銷非常有用。

OpenAI Triton 目前僅正式支持英偉達的 GPU，但在不久的將來會發(fā)生變化，將支持多個其他硬件供應(yīng)商。其他硬件加速器可以直接集成到 Triton 的 LLVM IR 中，這大大減少了為新硬件構(gòu)建 AI 編譯器堆棧的時間。

英偉達龐大的軟件體系缺乏遠見，無法利用其在 ML 硬件和軟件方面的巨大優(yōu)勢，也就沒能成為機器學(xué)習(xí)的默認編譯器。他們?nèi)狈捎眯缘年P(guān)注，而 OpenAI 和 Meta 也正是得益于此才能夠創(chuàng)建出可移植到其他硬件的軟件堆棧。

原文鏈接：https://www.semianalysis.com/p/nvidiaopenaitritonpytorch?