圖片來源 @視覺中國
钛媒體注：本文來源于微信公衆号半導體行業觀察（ID：icbank），作者｜semianalysis，钛媒體經授權發布。
在過去十年中，機器學習軟件開發的格局發生了重大變化。許多框架來來去去，但大多數都嚴重依賴于利用 Nvidia 的 CUDA，并且在 Nvidia GPU 上表現最佳。然而，随着 PyTorch 2.0 和 OpenAI 的 Triton 的到來，英偉達在該領域主要依靠其軟件護城河的主導地位正在被打破。
筆者認為，機器學習模型的默認軟件堆棧将不再是 Nvidia 的閉源 CUDA。球在 Nvidia 的球場上，他們讓 OpenAI 和 Meta 控制了軟件堆棧。由于 Nvidia 的專有工具失敗，該生态系統構建了自己的工具，現在 Nvidia 的護城河将被永久削弱。
TensorFlow 與 PyTorch
幾年前，框架生态系統相當分散，但 TensorFlow 是領跑者。谷歌看起來準備控制機器學習行業。他們憑借最常用的框架 TensorFlow 以及通過設計 / 部署唯一成功的 AI 應用特定加速器 TPU 獲得了先發優勢。
相反，PyTorch 赢了。谷歌未能将其先發優勢轉化為對新興 ML 行業的主導地位。如今，谷歌在機器學習社區中有些孤立，因為它不使用 PyTorch 和 GPU，而是使用自己的軟件堆棧和硬件。按照典型的谷歌方式，他們甚至還有一個名為 Jax 的第二個框架，它直接與 TensorFlow 競争。
由于大型語言模型，尤其是來自 OpenAI 和使用 OpenAI API 或正在構建類似基礎模型的各種初創公司的大型語言模型，谷歌在搜索和自然語言處理方面的主導地位甚至沒完沒了。雖然我們認為這種厄運和陰霾被誇大了，但這是另一個故事。盡管存在這些挑戰，谷歌仍處于最先進機器學習模型的前沿。他們發明了 transformers ，并在許多領域（PaLM、LaMBDA、Chinchilla、MUM、TPU）保持最先進的水平。
回到為什麼 PyTorch 赢了。雖然谷歌在争奪控制權方面存在一些因素，但這主要是由于 PyTorch 與 TensorFlow 相比具有更高的靈活性和可用性。如果我們将其歸結為第一個主要級别，PyTorch 與 TensorFlow 的不同之處在于使用 " Eager mode " 而不是 " Graph Mode "。
Eager 模式可以被認為是一種标準的腳本執行方法。深度學習框架會像任何其他 Python 代碼一樣，逐行立即執行每個操作。這使得調試和理解您的代碼更容易，因為您可以看到中間操作的結果并查看您的模型的行為方式。
相反，Graph 模式有兩個階段。第一階段是表示要執行的操作的計算圖的定義。計算圖是一系列相互連接的節點，表示操作或變量，節點之間的邊表示它們之間的數據流。第二階段是延遲執行計算圖的優化版本。
這種兩階段方法使理解和調試代碼更具挑戰性，因為在圖形執行結束之前您無法看到發生了什麼。這類似于 " 解釋 " 與 " 編譯 " 語言，如 Python 與 C++。調試 Python 更容易，主要是因為它是解釋型的。
雖然 TensorFlow 現在默認具有 Eager 模式，但研究社區和大多數大型科技公司已經圍繞 PyTorch 安定下來。
機器學習訓練組件
如果我們将機器學習模型訓練簡化為最簡單的形式，那麼機器學習模型的訓練時間有兩個主要的時間組成部分。
1、計算 ( FLOPS ) ：在每一層内運行密集矩陣乘法
2、内存（帶寬）：等待數據或層權重到達計算資源。帶寬受限操作的常見示例是各種規範化、逐點操作、SoftMax 和 ReLU 。
過去，機器學習訓練時間的主導因素是計算時間，等待矩陣乘法。随着 Nvidia 的 GPU 不斷發展，這很快就不再是主要問題。
Nvidia 的 FLOPS 通過利用摩爾定律提高了多個數量級，但主要是架構變化，例如張量核心和較低精度的浮點格式。相比之下，存儲并沒有走同樣的道路。
如果我們回到 2018 年，那時 BERT 模型是最先進的，Nvidia V100 是最先進的 GPU，我們可以看到矩陣乘法不再是提高模型性能的主要因素。從那時起，最先進的模型在參數數量上增長了 3 到 4 個數量級，而最快的 GPU 在 FLOPS 上增長了一個數量級。
即使在 2018 年，純計算綁定的工作負載也占 FLOPS 的 99.8%，但僅占運行時的 61%。與矩陣乘法相比，歸一化和逐點運算分别實現了 250 倍和 700 倍的 FLOPS，但它們消耗了模型運行時間的近 40%。
内存牆
随着模型規模的不斷飙升，大型語言模型僅用于模型權重就需要 100 GB（如果不是 TB）。百度和 Meta 部署的生産推薦網絡需要數十 TB 的内存來存儲其海量嵌入表。大型模型訓練 / 推理中的大部分時間都沒有花在計算矩陣乘法上，而是在等待數據到達計算資源。顯而易見的問題是，為什麼架構師不将更多内存放在更靠近計算的位置。答案是顯而易見的——成本。
内存遵循從近、快到慢、便宜的層次結構。最近的共享内存池在同一芯片上，一般由 SRAM 構成。一些機器學習 ASIC 試圖利用巨大的 SRAM 池來保存模型權重，但這種方法存在問題。即使是 Cerebras 的價值約 5,000,000 美元的晶圓級芯片也隻有 40GB 的 SRAM。内存容量不足以容納 100B+ 參數模型的權重。
Nvidia 的體系結構在裸片上一直使用的内存量要少得多。當前一代 A100 有 40MB，下一代 H100 有 50MB。台積電 5 納米工藝節點上的 1GB SRAM 需要約 200mm^2 的矽。一旦實現了相關的控制邏輯 / 結構，将需要超過 400mm^2 的矽，或 Nvidia 數據中心 GPU 總邏輯面積的大約 50%。鑒于 A100 GPU 的成本為 1 萬美元以上，而 H100 更接近 2 萬美元以上，從經濟角度來看，這是不可行的。即使忽略 Nvidia 在數據中心 GPU 上約 75% 的毛利率（約 4 倍加價），對于完全量産的産品，每 GB SRAM 内存的成本仍将在 100 美元左右。
此外，片上 SRAM 存儲器的成本不會随着傳統摩爾定律工藝技術的縮小而降低太多。同樣的 1GB 内存，采用台積電下一代 3nm 制程工藝，成本反而更高。雖然 3D SRAM 将在一定程度上幫助降低 SRAM 成本，但這隻是曲線的暫時彎曲。
内存層次結構的下一步是緊密耦合的片外内存 DRAM。DRAM 的延遲比 SRAM 高一個數量級（~>100 納秒對 ~10 納秒），但它也便宜得多（$1sa GB 對 $100s GB。）
幾十年來，DRAM 一直遵循着摩爾定律。當戈登摩爾創造這個詞時，英特爾的主要業務是 DRAM。他對晶體管密度和成本的經濟預測在 2009 年之前對 DRAM 普遍适用。不過自 2012 年以來，DRAM 的成本幾乎沒有改善。
對内存的需求隻會增加。DRAM 現在占服務器總成本的 50% 。這就是内存牆，它已經出現在産品中。将 Nvidia 2016 年的 P100 GPU 與 2022 剛剛開始出貨的 H100 GPU 進行比較，内存容量增加了 5 倍（16GB -> 80GB），但 FP16 性能增加了 46 倍（21.2 TFLOPS -> 989.5 TFLOPS）。
雖然容量是一個重要的瓶頸，但它與另一個主要瓶頸帶寬密切相關。增加的内存帶寬通常是通過并行性獲得的。雖然如今标準 DRAM 的價格僅為每 GB 幾美元，但為了獲得機器學習所需的海量帶寬，Nvidia 使用 HBM 内存，這是一種由 3D 堆疊 DRAM 層組成的設備，需要更昂貴的封裝。HBM 在每 GB 是 10 到 20 美元的範圍内，這包括封裝和産量成本。
内存帶寬和容量的成本限制不斷出現在 Nvidia 的 A100 GPU 中。如果不進行大量優化，A100 往往具有非常低的 FLOPS 利用率。FLOPS 利用率衡量訓練模型所需的總計算 FLOPS 與 GPU 在模型訓練時間内可以計算的理論 FLOPS。
即使領先研究人員進行了大量優化，60% 的 FLOPS 利用率也被認為是大型語言模型訓練的非常高的利用率。其餘時間是開銷，空閑時間花在等待來自另一個計算 / 内存的數據，或者及時重新計算結果以減少内存瓶頸。
從當前一代的 A100 到下一代 H100，FLOPS 增長了 6 倍以上，但内存帶寬僅增長了 1.65 倍。這導緻許多人擔心 H100 的利用率低。A100 需要很多技巧才能繞過内存牆，H100 還需要實現更多技巧。
H100 為 Hopper 帶來了分布式共享内存和 L2 多播 ( multicast ) 。這個想法是不同的 SM（think cores）可以直接寫入另一個 SM 的 SRAM（共享内存 /L1 緩存）中。這有效地增加了緩存的大小并減少了 DRAM 讀 / 寫所需的帶寬。未來的架構将依賴于向内存發送更少的操作，以最大限度地減少内存牆的影響。應該注意的是，較大的模型往往會實現更高的利用率，因為 FLOPS 需要按 2^n 縮放，而内存帶寬和容量需求往往會按 2*n 縮放。
算子融合——解決方法
有分析人士認為，就像訓練 ML 模型一樣，了解您所處的狀态可以讓您縮小重要的優化範圍。例如，如果您将所有時間都花在内存傳輸上（即您處于内存帶寬限制狀态），那麼增加 GPU 的 FLOPS 将無濟于事。另一方面，如果您将所有時間都花在執行大型 matmuls（如計算綁定機制）上，那麼将您的模型邏輯重寫為 C++ 以減少開銷将無濟于事。
而回顧 PyTorch 獲勝的原因，這是由于 Eager 模式提高了靈活性和可用性，但轉向 Eager 模式并不全是陽光和彩虹。在 Eager 模式下執行時，每個操作都從内存中讀取、計算，然後在處理下一個操作之前發送到内存。如果不進行大量優化，這會顯着增加内存帶寬需求。
因此，在 Eager 模式下執行的模型的主要優化方法之一稱為算子融合（operator fusion）。操作被融合，而不是将每個中間結果寫入内存，因此在一次傳遞中計算多個函數以最小化内存讀 / 寫。算子融合改善了運算符調度、内存帶寬和内存大小成本。
這種優化通常涉及編寫自定義 CUDA 内核，但這比使用簡單的 python 腳本要困難得多。作為一種内置的妥協，随着時間的推移，PyTorch 在 PyTorch 中穩定地實現了越來越多的算子（operator ) 。其中許多算子隻是簡單地将多個常用運算融合到一個更複雜的函數中。
算子的增加使得在 PyTorch 中創建模型變得更容易，并且由于内存讀 / 寫更少，Eager 模式的性能更快。缺點是 PyTorch 在幾年内激增到 2,000 多個算子。
我們會說軟件開發人員很懶惰，但說實話，幾乎所有人都是懶惰的。如果他們習慣了 PyTorch 中的一個新算子，他們将繼續使用它。開發人員可能甚至沒有意識到性能的提高，而是使用該算子，因為這意味着編寫更少的代碼。
此外，并非所有算子都可以融合。通常花費大量時間來決定要融合哪些操作以及将哪些分配給芯片和集群級别的特定計算資源。哪些算子在哪裡融合的策略雖然大體相似，但根據架構的不同會有很大差異。
英偉達為王
算子的增長和默認位置對 Nvidia 有所幫助，因為每個算子都針對其架構進行了快速優化，但并未針對任何其他硬件進行優化。如果一家 AI 硬件初創公司想要全面實施 PyTorch，那就意味着以高性能原生支持不斷增長的 2,000 個算子列表。
由于提取最大性能所需的所有技巧，在 GPU 上訓練具有高 FLOPS 利用率的大型模型所需的人才水平越來越高。Eager mode 執行加上 operator fusion 意味着開發的軟件、技術和模型被推動以适應當前一代 GPU 具有的計算和内存比率。
每個開發機器學習芯片的人都依賴于同一個内存牆。ASIC 有責任支持最常用的框架。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 、英特爾、Tenstorrent 、Luminous Computing、特斯拉、谷歌、亞馬遜、微軟、Marvell 、Meta 、Graphcore 、Cerebras 、SambaNova 等的其他 GPU 和加速器。
對于當前未優化的硬件，PyTorch 2.0 的性能改進會更大。Meta 和其他公司對 PyTorch 的巨大貢獻源于這樣一個事實，即他們希望在他們價值數十億美元的 GPU 訓練集群上以更少的努力，更容易地實現更高的 FLOPS 利用率。他們也有動力使他們的軟件堆棧更易于移植到其他硬件，以将競争引入機器學習領域。
PyTorch 2.0 還通過更好的 API 支持數據并行、分片、管道并行和張量并行，為分布式訓練帶來了進步。此外，它在整個堆棧中原生支持動态形狀，在許多其他示例中，這使得 LLM 的不同序列長度更容易支持。這是主要編譯器首次支持從訓練到推理的 Dynamic Shapes。
請輸入圖說
PrimTorch
為 PyTorch 編寫一個完全支持所有 2,000 多個算子的高性能後端對于除 Nvidia GPU 之外的每個機器學習 ASIC 來說都是困難的。PrimTorch 将算子的數量減少到約 250 個原始算子，同時還保持 PyTorch 最終用戶的可用性不變。PrimTorch 使 PyTorch 的不同非 Nvidia 後端的實現變得更加簡單和易于訪問。定制硬件和系統供應商可以更輕松地推出他們的軟件堆棧。
TorchDynamo
轉向圖形模式需要一個可靠的圖形定義。Meta 和 PyTorch 已經嘗試了大約 5 年的時間來實現這一點，但是他們提出的每個解決方案都有明顯的缺點。他們終于用 TorchDynamo 破解了這個難題。TorchDynamo 将攝取任何 PyTorch 用戶腳本，包括調用外部 3rd 方庫的腳本，并生成 FX 圖形。
Dynamo 将所有複雜算子減少到 PrimTorch 中的約 250 個原始算子。一旦圖形成，未使用的算子将被丢棄，圖确定哪些中間算子需要存儲或寫入内存，哪些可能被融合。這極大地減少了模型内的開銷，同時對用戶來說也是無縫的。
TorchDynamo 已經适用于 7,000 個經過測試的 PyTorch 模型中的 99% 以上，包括來自 OpenAI、HuggingFace、Meta、Nvidia、Stability.AI 等的模型，而無需對原始代碼進行任何更改。測試的 7,000 個模型是從 GitHub 上使用 PyTorch 的最流行項目中不分青紅皂白地挑選出來的。
谷歌的 TensorFlow/Jax 和其他圖形模式執行管道通常要求用戶确保他們的模型适合編譯器架構，以便可以捕獲圖形。Dynamo 通過啟用部分圖捕獲、受保護的圖捕獲和即時重新捕獲來改變這一點。
部分圖形捕獲允許模型包含不受支持的 / 非 python 構造。當無法為模型的那部分生成圖時，将插入圖中斷，并且将在部分圖之間以急切模式執行不支持的構造。
受保護的圖捕獲檢查捕獲的圖是否對執行有效。守衛是需要重新編譯的更改。這很重要，因為多次運行相同的代碼不會多次重新編譯。
如果捕獲的圖對于執行無效，則即時重新捕獲允許重新捕獲圖。
PyTorch 的目标是創建一個具有流暢 UX 的統一前端，該前端利用 Dynamo 生成圖形。該解決方案的用戶體驗不會發生變化，但性能可以得到顯着提升。捕獲圖形意味着可以在大量計算資源上更有效地并行執行。
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 的編譯器。我們想要一些可以擴展以支持各種硬件後端的東西，并且擁有 C++ 和 [ OpenAI ] Triton 會強制實現這種通用性。
OpenAI 的 Triton
OpenAI 的 Triton 對 Nvidia 的機器學習閉源軟件護城河具有颠覆性的角度。Triton 直接采用 Python 或通過 PyTorch Inductor 堆棧提供數據。後者将是最常見的用例。Triton 然後将輸入轉換為 LLVM 中間表示，然後生成代碼。對于 Nvidia GPU，它直接生成 PTX 代碼，跳過 Nvidia 的閉源 CUDA 庫（例如 cuBLAS），轉而使用開源庫（例如 cutlass）。
CUDA 通常被專門從事加速計算的人員使用，但在機器學習研究人員和數據科學家中卻鮮為人知。高效使用可能具有挑戰性，并且需要深入了解硬件架構，這可能會減慢開發過程。因此，機器學習專家可能依賴 CUDA 專家來修改、優化和并行化他們的代碼。
Triton 彌合了使高級語言能夠實現與使用低級語言的語言相當的性能的差距。Triton 内核本身對于典型的 ML 研究人員來說非常清晰，這對于可用性來說非常重要。Triton 在 SM 中自動執行内存合并、共享内存管理和調度。Triton 對逐元素矩陣乘法不是特别有用，這已經非常有效地完成了。Triton 對于昂貴的逐點操作和減少更複雜操作的開銷非常有用，例如 Flash Attention 涉及矩陣乘法作為較大融合操作的一部分。
OpenAI Triton 目前僅正式支持 Nvidia GPU，但在不久的将來這種情況會發生變化。未來将支持多個其他硬件供應商，這個開源項目正在獲得令人難以置信的動力。其他硬件加速器直接集成到作為 Triton 一部分的 LLVM IR 的能力大大減少了為新硬件構建 AI 編譯器堆棧的時間。
Nvidia 龐大的軟件組織缺乏遠見，無法利用其在 ML 硬件和軟件方面的巨大優勢，成為機器學習的默認編譯器。他們缺乏對可用性的關注，這使得 OpenAI 和 Meta 的外部人員能夠創建可移植到其他硬件的軟件堆棧。
更多精彩内容，關注钛媒體微信号（ID：taimeiti），或者下載钛媒體 App
