圖片來源@視覺中國

鈦媒體注：本文來源于微信公眾號半導(dǎo)體行業(yè)觀察（ID：icbank），作者｜semianalysis，鈦媒體經(jīng)授權(quán)發(fā)布。

在過去十年中，機(jī)器學(xué)習(xí)軟件開發(fā)的格局發(fā)生了重大變化。許多框架來來去去，但大多數(shù)都嚴(yán)重依賴于利用 Nvidia 的 CUDA，并且在 Nvidia GPU 上表現(xiàn)最佳。然而，隨著 PyTorch 2.0 和 OpenAI 的 Triton 的到來，英偉達(dá)在該領(lǐng)域主要依靠其軟件護(hù)城河的主導(dǎo)地位正在被打破。

筆者認(rèn)為，機(jī)器學(xué)習(xí)模型的默認(rèn)軟件堆棧將不再是 Nvidia 的閉源 CUDA。球在 Nvidia 的球場上，他們讓 OpenAI 和 Meta 控制了軟件堆棧。由于 Nvidia 的專有工具失敗，該生態(tài)系統(tǒng)構(gòu)建了自己的工具，現(xiàn)在 Nvidia 的護(hù)城河將被永久削弱。

TensorFlow 與 PyTorch

幾年前，框架生態(tài)系統(tǒng)相當(dāng)分散，但 TensorFlow 是領(lǐng)跑者。谷歌看起來準(zhǔn)備控制機(jī)器學(xué)習(xí)行業(yè)。他們憑借最常用的框架 TensorFlow 以及通過設(shè)計(jì)/部署唯一成功的 AI 應(yīng)用特定加速器 TPU 獲得了先發(fā)優(yōu)勢。

相反，PyTorch 贏了。谷歌未能將其先發(fā)優(yōu)勢轉(zhuǎn)化為對新興 ML 行業(yè)的主導(dǎo)地位。如今，谷歌在機(jī)器學(xué)習(xí)社區(qū)中有些孤立，因?yàn)樗皇褂?PyTorch 和 GPU，而是使用自己的軟件堆棧和硬件。按照典型的谷歌方式，他們甚至還有一個(gè)名為 Jax 的第二個(gè)框架，它直接與 TensorFlow 競爭。

由于大型語言模型，尤其是來自 OpenAI 和使用 OpenAI API 或正在構(gòu)建類似基礎(chǔ)模型的各種初創(chuàng)公司的大型語言模型，谷歌在搜索和自然語言處理方面的主導(dǎo)地位甚至沒完沒了。雖然我們認(rèn)為這種厄運(yùn)和陰霾被夸大了，但這是另一個(gè)故事。盡管存在這些挑戰(zhàn)，谷歌仍處于最先進(jìn)機(jī)器學(xué)習(xí)模型的前沿。他們發(fā)明了transformers ，并在許多領(lǐng)域（PaLM、LaMBDA、Chinchilla、MUM、TPU）保持最先進(jìn)的水平。

回到為什么 PyTorch 贏了。雖然谷歌在爭奪控制權(quán)方面存在一些因素，但這主要是由于 PyTorch 與 TensorFlow 相比具有更高的靈活性和可用性。如果我們將其歸結(jié)為第一個(gè)主要級別，PyTorch 與 TensorFlow 的不同之處在于使用“ Eager mode ”而不是“ Graph Mode ”。

Eager 模式可以被認(rèn)為是一種標(biāo)準(zhǔn)的腳本執(zhí)行方法。深度學(xué)習(xí)框架會像任何其他 Python 代碼一樣，逐行立即執(zhí)行每個(gè)操作。這使得調(diào)試和理解您的代碼更容易，因?yàn)槟梢钥吹街虚g操作的結(jié)果并查看您的模型的行為方式。

相反，Graph 模式有兩個(gè)階段。第一階段是表示要執(zhí)行的操作的計(jì)算圖的定義。計(jì)算圖是一系列相互連接的節(jié)點(diǎn)，表示操作或變量，節(jié)點(diǎn)之間的邊表示它們之間的數(shù)據(jù)流。第二階段是延遲執(zhí)行計(jì)算圖的優(yōu)化版本。

這種兩階段方法使理解和調(diào)試代碼更具挑戰(zhàn)性，因?yàn)樵趫D形執(zhí)行結(jié)束之前您無法看到發(fā)生了什么。這類似于“解釋”與“編譯”語言，如 Python 與 C++。調(diào)試 Python 更容易，主要是因?yàn)樗墙忉屝偷摹?/p>

雖然 TensorFlow 現(xiàn)在默認(rèn)具有 Eager 模式，但研究社區(qū)和大多數(shù)大型科技公司已經(jīng)圍繞 PyTorch 安定下來。

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

如果我們將機(jī)器學(xué)習(xí)模型訓(xùn)練簡化為最簡單的形式，那么機(jī)器學(xué)習(xí)模型的訓(xùn)練時(shí)間有兩個(gè)主要的時(shí)間組成部分。

1、計(jì)算 (FLOPS)：在每一層內(nèi)運(yùn)行密集矩陣乘法

2、內(nèi)存（帶寬）：等待數(shù)據(jù)或?qū)訖?quán)重到達(dá)計(jì)算資源。帶寬受限操作的常見示例是各種規(guī)范化、逐點(diǎn)操作、SoftMax和ReLU 。

過去，機(jī)器學(xué)習(xí)訓(xùn)練時(shí)間的主導(dǎo)因素是計(jì)算時(shí)間，等待矩陣乘法。隨著 Nvidia 的 GPU 不斷發(fā)展，這很快就不再是主要問題。

Nvidia 的 FLOPS 通過利用摩爾定律提高了多個(gè)數(shù)量級，但主要是架構(gòu)變化，例如張量核心和較低精度的浮點(diǎn)格式。相比之下，存儲并沒有走同樣的道路。

如果我們回到 2018 年，那時(shí) BERT 模型是最先進(jìn)的，Nvidia V100 是最先進(jìn)的 GPU，我們可以看到矩陣乘法不再是提高模型性能的主要因素。從那時(shí)起，最先進(jìn)的模型在參數(shù)數(shù)量上增長了 3 到 4 個(gè)數(shù)量級，而最快的 GPU 在 FLOPS 上增長了一個(gè)數(shù)量級。

即使在 2018 年，純計(jì)算綁定的工作負(fù)載也占 FLOPS 的 99.8%，但僅占運(yùn)行時(shí)的 61%。與矩陣乘法相比，歸一化和逐點(diǎn)運(yùn)算分別實(shí)現(xiàn)了 250 倍和 700 倍的 FLOPS，但它們消耗了模型運(yùn)行時(shí)間的近 40%。

內(nèi)存墻

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

內(nèi)存遵循從近、快到慢、便宜的層次結(jié)構(gòu)。最近的共享內(nèi)存池在同一芯片上，一般由SRAM構(gòu)成。一些機(jī)器學(xué)習(xí) ASIC 試圖利用巨大的 SRAM 池來保存模型權(quán)重，但這種方法存在問題。即使是 Cerebras 的價(jià)值約 5,000,000 美元的晶圓級芯片也只有 40GB 的 SRAM。內(nèi)存容量不足以容納 100B+ 參數(shù)模型的權(quán)重。

Nvidia 的體系結(jié)構(gòu)在裸片上一直使用的內(nèi)存量要少得多。當(dāng)前一代A100有40MB，下一代H100有50MB。臺積電 5 納米工藝節(jié)點(diǎn)上的 1GB SRAM 需要約 200mm^2 的硅。一旦實(shí)現(xiàn)了相關(guān)的控制邏輯/結(jié)構(gòu)，將需要超過 400mm^2 的硅，或 Nvidia 數(shù)據(jù)中心 GPU 總邏輯面積的大約 50%。鑒于 A100 GPU 的成本為 1 萬美元以上，而 H100 更接近 2 萬美元以上，從經(jīng)濟(jì)角度來看，這是不可行的。即使忽略 Nvidia 在數(shù)據(jù)中心 GPU 上約 75% 的毛利率（約 4 倍加價(jià)），對于完全量產(chǎn)的產(chǎn)品，每 GB SRAM 內(nèi)存的成本仍將在 100 美元左右。

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

內(nèi)存層次結(jié)構(gòu)的下一步是緊密耦合的片外內(nèi)存 DRAM。DRAM 的延遲比 SRAM 高一個(gè)數(shù)量級（~>100 納秒對~10 納秒），但它也便宜得多（$1sa GB 對 $100s GB。）

幾十年來，DRAM 一直遵循著摩爾定律。當(dāng)戈登摩爾創(chuàng)造這個(gè)詞時(shí)，英特爾的主要業(yè)務(wù)是 DRAM。他對晶體管密度和成本的經(jīng)濟(jì)預(yù)測在 2009 年之前對 DRAM 普遍適用。不過自 2012 年以來，DRAM 的成本幾乎沒有改善。

對內(nèi)存的需求只會增加。DRAM 現(xiàn)在占服務(wù)器總成本的 50% 。這就是內(nèi)存墻，它已經(jīng)出現(xiàn)在產(chǎn)品中。將 Nvidia 2016年的P100 GPU 與2022 剛剛開始出貨的H100 GPU 進(jìn)行比較，內(nèi)存容量增加了 5 倍（16GB -> 80GB），但 FP16 性能增加了 46 倍（21.2 TFLOPS -> 989.5 TFLOPS）。

雖然容量是一個(gè)重要的瓶頸，但它與另一個(gè)主要瓶頸帶寬密切相關(guān)。增加的內(nèi)存帶寬通常是通過并行性獲得的。雖然如今標(biāo)準(zhǔn) DRAM 的價(jià)格僅為每 GB 幾美元，但為了獲得機(jī)器學(xué)習(xí)所需的海量帶寬，Nvidia 使用 HBM 內(nèi)存，這是一種由3D 堆疊 DRAM 層組成的設(shè)備，需要更昂貴的封裝。HBM 在每 GB 是10 到 20 美元的范圍內(nèi)，這包括封裝和產(chǎn)量成本。

內(nèi)存帶寬和容量的成本限制不斷出現(xiàn)在 Nvidia 的 A100 GPU 中。如果不進(jìn)行大量優(yōu)化，A100 往往具有非常低的 FLOPS 利用率。FLOPS 利用率衡量訓(xùn)練模型所需的總計(jì)算 FLOPS 與 GPU 在模型訓(xùn)練時(shí)間內(nèi)可以計(jì)算的理論 FLOPS。

即使領(lǐng)先研究人員進(jìn)行了大量優(yōu)化，60% 的 FLOPS 利用率也被認(rèn)為是大型語言模型訓(xùn)練的非常高的利用率。其余時(shí)間是開銷，空閑時(shí)間花在等待來自另一個(gè)計(jì)算/內(nèi)存的數(shù)據(jù)，或者及時(shí)重新計(jì)算結(jié)果以減少內(nèi)存瓶頸。

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

H100 為Hopper 帶來了分布式共享內(nèi)存和 L2 多播(multicast)。這個(gè)想法是不同的 SM（think cores）可以直接寫入另一個(gè) SM 的 SRAM（共享內(nèi)存/L1 緩存）中。這有效地增加了緩存的大小并減少了DRAM 讀/寫所需的帶寬。未來的架構(gòu)將依賴于向內(nèi)存發(fā)送更少的操作，以最大限度地減少內(nèi)存墻的影響。應(yīng)該注意的是，較大的模型往往會實(shí)現(xiàn)更高的利用率，因?yàn)?FLOPS 需要按 2^n 縮放，而內(nèi)存帶寬和容量需求往往會按 2*n 縮放。

算子融合——解決方法

有分析人士認(rèn)為，就像訓(xùn)練 ML 模型一樣，了解您所處的狀態(tài)可以讓您縮小重要的優(yōu)化范圍。例如，如果您將所有時(shí)間都花在內(nèi)存?zhèn)鬏斏希茨幱趦?nèi)存帶寬限制狀態(tài)），那么增加 GPU 的 FLOPS 將無濟(jì)于事。另一方面，如果您將所有時(shí)間都花在執(zhí)行大型 matmuls（如計(jì)算綁定機(jī)制）上，那么將您的模型邏輯重寫為 C++ 以減少開銷將無濟(jì)于事。

而回顧 PyTorch 獲勝的原因，這是由于 Eager 模式提高了靈活性和可用性，但轉(zhuǎn)向 Eager 模式并不全是陽光和彩虹。在 Eager 模式下執(zhí)行時(shí)，每個(gè)操作都從內(nèi)存中讀取、計(jì)算，然后在處理下一個(gè)操作之前發(fā)送到內(nèi)存。如果不進(jìn)行大量優(yōu)化，這會顯著增加內(nèi)存帶寬需求。

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

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

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

我們會說軟件開發(fā)人員很懶惰，但說實(shí)話，幾乎所有人都是懶惰的。如果他們習(xí)慣了 PyTorch 中的一個(gè)新算子，他們將繼續(xù)使用它。開發(fā)人員可能甚至沒有意識到性能的提高，而是使用該算子，因?yàn)檫@意味著編寫更少的代碼。

此外，并非所有算子都可以融合。通?；ㄙM(fèi)大量時(shí)間來決定要融合哪些操作以及將哪些分配給芯片和集群級別的特定計(jì)算資源。哪些算子在哪里融合的策略雖然大體相似，但根據(jù)架構(gòu)的不同會有很大差異。

英偉達(dá)為王

算子的增長和默認(rèn)位置對 Nvidia 有所幫助，因?yàn)槊總€(gè)算子都針對其架構(gòu)進(jìn)行了快速優(yōu)化，但并未針對任何其他硬件進(jìn)行優(yōu)化。如果一家 AI 硬件初創(chuàng)公司想要全面實(shí)施 PyTorch，那就意味著以高性能原生支持不斷增長的 2,000 個(gè)算子列表。

由于提取最大性能所需的所有技巧，在 GPU 上訓(xùn)練具有高 FLOPS 利用率的大型模型所需的人才水平越來越高。Eager mode 執(zhí)行加上operator fusion 意味著開發(fā)的軟件、技術(shù)和模型被推動(dòng)以適應(yīng)當(dāng)前一代 GPU 具有的計(jì)算和內(nèi)存比率。

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

但易用性為王。

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

PyTorch 2.0

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

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

對于當(dāng)前未優(yōu)化的硬件，PyTorch 2.0 的性能改進(jìn)會更大。Meta 和其他公司對 PyTorch 的巨大貢獻(xiàn)源于這樣一個(gè)事實(shí)，即他們希望在他們價(jià)值數(shù)十億美元的 GPU 訓(xùn)練集群上以更少的努力，更容易地實(shí)現(xiàn)更高的 FLOPS 利用率。他們也有動(dòng)力使他們的軟件堆棧更易于移植到其他硬件，以將競爭引入機(jī)器學(xué)習(xí)領(lǐng)域。

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

請輸入圖說

PrimTorch

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

TorchDynamo

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

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

TorchDynamo 已經(jīng)適用于7,000 個(gè)經(jīng)過測試的 PyTorch 模型中的 99% 以上，包括來自 OpenAI、HuggingFace、Meta、Nvidia、Stability.AI 等的模型，而無需對原始代碼進(jìn)行任何更改。測試的 7,000 個(gè)模型是從 GitHub 上使用 PyTorch 的最流行項(xiàng)目中不分青紅皂白地挑選出來的。

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

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

受保護(hù)的圖捕獲檢查捕獲的圖是否對執(zhí)行有效。守衛(wèi)是需要重新編譯的更改。這很重要，因?yàn)槎啻芜\(yùn)行相同的代碼不會多次重新編譯。

如果捕獲的圖對于執(zhí)行無效，則即時(shí)重新捕獲允許重新捕獲圖。

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

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

TorchInductor

TorchInductor 是一個(gè) python 原生深度學(xué)習(xí)編譯器，可以為多個(gè)加速器和后端生成快速代碼。Inductor 將采用具有約 250 個(gè)算子的 FX 圖，并將它們降低到約 50 個(gè)算子。

Inductor 然后進(jìn)入調(diào)度階段，在該階段融合運(yùn)算符，并確定內(nèi)存規(guī)劃。

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

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

分析人士表示，我們不想構(gòu)建只支持 GPU 的編譯器。我們想要一些可以擴(kuò)展以支持各種硬件后端的東西，并且擁有 C++ 和 [OpenAI] Triton 會強(qiáng)制實(shí)現(xiàn)這種通用性。

OpenAI 的Triton

OpenAI 的 Triton 對 Nvidia 的機(jī)器學(xué)習(xí)閉源軟件護(hù)城河具有顛覆性的角度。Triton 直接采用 Python 或通過PyTorch Inductor 堆棧提供數(shù)據(jù)。后者將是最常見的用例。Triton 然后將輸入轉(zhuǎn)換為 LLVM 中間表示，然后生成代碼。對于 Nvidia GPU，它直接生成 PTX 代碼，跳過 Nvidia 的閉源 CUDA 庫（例如 cuBLAS），轉(zhuǎn)而使用開源庫（例如 cutlass）。

CUDA 通常被專門從事加速計(jì)算的人員使用，但在機(jī)器學(xué)習(xí)研究人員和數(shù)據(jù)科學(xué)家中卻鮮為人知。高效使用可能具有挑戰(zhàn)性，并且需要深入了解硬件架構(gòu)，這可能會減慢開發(fā)過程。因此，機(jī)器學(xué)習(xí)專家可能依賴 CUDA 專家來修改、優(yōu)化和并行化他們的代碼。

Triton 彌合了使高級語言能夠?qū)崿F(xiàn)與使用低級語言的語言相當(dāng)?shù)男阅艿牟罹?。Triton 內(nèi)核本身對于典型的 ML 研究人員來說非常清晰，這對于可用性來說非常重要。Triton 在 SM 中自動(dòng)執(zhí)行內(nèi)存合并、共享內(nèi)存管理和調(diào)度。Triton 對逐元素矩陣乘法不是特別有用，這已經(jīng)非常有效地完成了。Triton 對于昂貴的逐點(diǎn)操作和減少更復(fù)雜操作的開銷非常有用，例如Flash Attention涉及矩陣乘法作為較大融合操作的一部分。

OpenAI Triton 目前僅正式支持 Nvidia GPU，但在不久的將來這種情況會發(fā)生變化。未來將支持多個(gè)其他硬件供應(yīng)商，這個(gè)開源項(xiàng)目正在獲得令人難以置信的動(dòng)力。其他硬件加速器直接集成到作為 Triton 一部分的 LLVM IR 的能力大大減少了為新硬件構(gòu)建 AI 編譯器堆棧的時(shí)間。

Nvidia 龐大的軟件組織缺乏遠(yuǎn)見，無法利用其在 ML 硬件和軟件方面的巨大優(yōu)勢，成為機(jī)器學(xué)習(xí)的默認(rèn)編譯器。他們?nèi)狈捎眯缘年P(guān)注，這使得 OpenAI 和 Meta 的外部人員能夠創(chuàng)建可移植到其他硬件的軟件堆棧。

763人已贊賞 >

敬原創(chuàng)，有鈦度，得贊賞

贊賞支持

0

13

0

掃描下載App