和TensorFlow一樣，NVIDIA CUDA的壟斷格局將被打破？

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

TensorFlow vs. PyTorch

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

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

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

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

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

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

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

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

機器學習訓練組件

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

浮點運算 ( FLOPS ) ：在每一層內運行密集矩陣乘法；
記憶體頻寬。

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

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

2018 年，最先進的模型是 BERT，NVIDIA 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+ 參數模型的權重。

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

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

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

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

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

記憶體頻寬和容量的成本限制問題在NVIDIA 的 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 中的一個新運算元，他們就會繼續用它。開發人員甚至可能沒有意識到性能在提高，而是繼續使用該運算元，因為這樣就不用編寫更多的程式碼。

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

NVIDIA 曾是王者

運算元的增長和預設的地位對NVIDIA 來說是優勢，因為每個運算元都針對其架構進行了快速最佳化，但並未針對任何其他硬體進行最佳化。如果一家 AI 硬體初創公司想要全面實施 PyTorch，那就意味著以高性能支援不斷增長的 2000 個運算元列表。

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

每個開發機器學習晶片的人都受制於同一個記憶體牆。ASIC 受制於支援最常用的框架，受制于預設的開發方法、GPU 最佳化的 PyTorch 程式碼以及NVIDIA 和外部庫的混合。在這種情況下，避開 GPU 的各種非運算包袱而支援更多 FLOPS 和更嚴格的程式模型的架構意義不大。

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

PyTorch 2.0

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

PyTorch 2.0 在NVIDIA A100 上的訓練性能提升了 86%，在 CPU 上的推理性能提升了 26%。這大大減少了訓練模型所需的運算時間和成本。這些好處可以擴充到來自 AMD、Intel、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

對於除NVIDIA GPU 之外的每個機器學習 ASIC 來說，為 PyTorch 編寫一個完全支援所有 2000 多個運算元的高性能後端並非易事。PrimTorch 將運算元的數量減少到約 250 個原始運算元，同時還保持 PyTorch 最終使用者的可用性不變。PrimTorch 使 PyTorch 的不同非NVIDIA 後端的實現變得更加簡單和易於存取。定制硬體和系統供應商可以更輕鬆地推出他們的軟體堆疊。

TorchDynamo

轉向圖模式需要可靠的圖定義。為了實現這一轉向，Meta 和 PyTorch 已經嘗試了大約 5 年的時間，但是他們提出的每個解決方案都存在明顯的缺點。最後，他們用 TorchDynamo 破解了這個難題。TorchDynamo 將攝取任何 PyTorch 使用者腳本，包括系統呼叫外部協力廠商庫的腳本，並產生 FX 圖。

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

在測試的 7000 個 PyTorch 模型中，TorchDynamo 已經適用於 99% 以上的模型，包括來自 OpenAI、HuggingFace、Meta、NVIDIA 、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，Intel編譯器生成 C++（也適用於非Intel CPU）。

未來他們將支援更多硬體，但關鍵是 Inductor 大大減少了編譯器團隊在為其 AI 硬體加速器製作編譯器時必須做的工作量。此外，程式碼針對性能進行了更最佳化，記憶體頻寬和容量要求得到了顯著降低。

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

OpenAI Triton

對NVIDIA 的機器學習閉源軟體來說，OpenAI Triton 是一個顛覆性的存在。Triton 直接採用 Python 或通過 PyTorch Inductor 堆疊提供資料，後者是最常見的用法。Triton 負責將輸入轉換為 LLVM 中間表徵，並生成程式碼。NVIDIA GPU 將直接生成 PTX 程式碼，跳過NVIDIA 的閉源 CUDA 庫（如 cuBLAS），轉而使用開源庫（如 cutlass）。

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

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

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

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

資料來源：

How Nvidia’s CUDA Monopoly In Machine Learning Is Breaking - OpenAI Triton And PyTorch 2.0