?十年來，機器學習軟件開發的格局發生了重大變化。許多框架如雨后春筍般涌現，但大多數都嚴重依賴于英偉達的 CUDA，并在英偉達的 GPU 上才能獲得最佳的性能。然而，隨著 PyTorch 2.0 和 OpenAI Triton 的到來，英偉達在這一領域的主導地位正在被打破。

谷歌早期在機器學習模型架構、訓練、模型優化方面都具有很大優勢，但現在卻難以充分發揮這些優勢。而在硬件方面，其他 AI 硬件公司很難削弱英偉達的統治地位。直到 PyTorch 2.0 和 OpenAI Triton 出現，機器學習模型的默認軟件堆棧將不再是英偉達的閉源 CUDA。

TensorFlow vs. PyTorch

類似的競爭出現在機器學習框架中。幾年前，框架生態系統相當分散，但 TensorFlow 是領跑者。表面看來，谷歌穩操機器學習框架行業，他們憑借 TensorFlow 設計了 AI 應用特定加速器 TPU，從而獲得了先發優勢。

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

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

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

Eager 模式可以說是標準的腳本執行方法，與普通的 Python 代碼沒什么區別。這使得調試和理解代碼更加容易，因為用戶可以看到中間操作的結果，以及模型是如何運行的。

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

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

雖然 TensorFlow 現在也默認使用 Eager 模式，但研究社區和大多數大型科技公司都選擇使用 PyTorch。

機器學習訓練組件

如果將機器學習模型訓練簡化為最簡單的形式，影響機器學習模型訓練的因素主要有兩點：

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

以前，影響機器學習訓練時間的主要因素是計算時間，等待系統執行矩陣乘法。隨著英偉達 GPU 的不斷發展，這很快就不再是主要問題了。

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

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

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

內存墻

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

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

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

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

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

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

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

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

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

從 A100 到 H100，FLOPS 增長到原來的 6 倍以上，但內存帶寬僅增長到 1.65 倍。這導致許多人擔心 H100 的利用率會很低。A100 需要很多技巧才能繞過內存墻，而 H100 還需要實現更多技巧。

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

算子融合

如果將所有時間都花在內存傳輸上（即處于內存帶寬限制狀態），那么增加 GPU 的 FLOPS 將無濟于事。另一方面，如果將所有時間都花在執行大型 matmuls 上，那么即使將模型邏輯重寫為 C++ 來減少開銷也將無濟于事。

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

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

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

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

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

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

英偉達曾是王者

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

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

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

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

PyTorch 2.0

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

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

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

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

PrimTorch

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

TorchDynamo

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

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

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

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

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

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

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

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

TorchInductor

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

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

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

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

OpenAI Triton

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

CUDA 在加速計算領域很受歡迎，但在機器學習研究人員和數據科學家中卻鮮為人知。使用 CUDA 可能會帶來重重挑戰，并且需要深入了解硬件架構，這可能導致開發過程變慢。因此，機器學習專家可能就要依賴 CUDA 專家來修改、優化和并行化他們的代碼。

Triton 彌補了這一缺陷，使高級語言實現了與低級語言相當的性能。Triton 內核本身對典型的 ML 研究者來說非常清晰，這對可用性來說非常重要。Triton 在 SM 中自動執行內存合并、共享內存管理和調度。Triton 對逐元素矩陣乘法不是特別有用，但矩陣乘法已經可以非常高效地完成。Triton 對于成本高昂的逐點運算和減少復雜操作的開銷非常有用。

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

英偉達龐大的軟件體系缺乏遠見，無法利用其在 ML 硬件和軟件方面的巨大優勢，也就沒能成為機器學習的默認編譯器。他們缺乏對可用性的關注，而 OpenAI 和 Meta 也正是得益于此才能夠創建出可移植到其他硬件的軟件堆棧。

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