CUDA 閉源庫將和 TensorFlow 一樣逐漸式微。
十年來，機器學習軟體開發的格局發生了重大變化。許多框架如雨後春筍般湧現，但大多數都嚴重依賴於英偉達的 CUDA，並在英偉達的 GPU 上才能獲得最佳的性能。然而，隨著 PyTorch 2.0 和 OpenAI Triton 的到來，英偉達在這一領域的主導地位正在被打破。
Google早期在機器學習模型架構、訓練、模型最佳化方面都具有很大優勢，但現在卻難以充分發揮這些優勢。而在硬體方面，其他 AI 硬體公司很難削弱英偉達的統治地位。直到 PyTorch 2.0 和 OpenAI Triton 出現，機器學習模型的默認軟體堆疊將不再是英偉達的閉源 CUDA。
TensorFlow vs. PyTorch
類似的競爭出現在機器學習框架中。幾年前，框架生態系統相當分散，但 TensorFlow 是領跑者。表面看來，Google穩操機器學習框架行業，他們憑藉 TensorFlow 設計了 AI 應用特定加速器 TPU，從而獲得了先發優勢。
然而現在看來，PyTorch 贏了，Google未能將其先發優勢轉化為新興 ML 行業的主導地位。如今，Google在機器學習社區中似乎有些被孤立了，因為它沒有使用 PyTorch 和 GPU，而是使用自己的軟體堆疊和硬體。甚至，Google研發了第二個機器學習框架 ——JAX，直接與 TensorFlow 競爭，這是典型的「Google行為」。
一些人認為，由於大型語言模型的興起，特別是 OpenAI 的大型語言模型和各種利用 OpenAI API 構建的語言模型，Google在搜尋和自然語言處理方面的主導地位正在減弱。也許這種觀點過於悲觀，畢竟當前大多數模型的基礎架構仍是Google開發的 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、特斯拉、Google、亞馬遜、微軟、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 的最受歡迎項目中隨機挑選出來的。
Google的 TensorFlow/Jax 和其他圖模式執行 pipeline 通常要求使用者確保他們的模型適合編譯器架構，以便可以捕獲圖。Dynamo 通過啟用部分圖捕獲、受保護的圖捕獲和即時重新捕獲來改變這一點。
部分圖捕獲允許模型包含不受支持的 / 非 python 構造。當無法為模型部分生成圖時，將插入圖中斷，並且將在部分圖之間以 eager 模式執行不支持的構造。
受保護的圖捕獲會檢查捕獲的圖是否對執行有效。「保護」的意思是一種需要重新編譯的更改。這很重要，因為多次運行相同的程式碼不會多次重新編譯。如果捕獲的圖對於執行無效，則即時重新捕獲允許重新捕獲圖。
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
