?十年來(lái)，機(jī)器學(xué)習(xí)軟件開發(fā)的格局發(fā)生了重大變化。許多框架如雨后春筍般涌現(xiàn)，但大多數(shù)都嚴(yán)重依賴于英偉達(dá)的 CUDA，并在英偉達(dá)的 GPU 上才能獲得最佳的性能。然而，隨著 PyTorch 2.0 和 OpenAI Triton 的到來(lái)，英偉達(dá)在這一領(lǐng)域的主導(dǎo)地位正在被打破。

谷歌早期在機(jī)器學(xué)習(xí)模型架構(gòu)、訓(xùn)練、模型優(yōu)化方面都具有很大優(yōu)勢(shì)，但現(xiàn)在卻難以充分發(fā)揮這些優(yōu)勢(shì)。而在硬件方面，其他 AI 硬件公司很難削弱英偉達(dá)的統(tǒng)治地位。直到 PyTorch 2.0 和 OpenAI Triton 出現(xiàn)，機(jī)器學(xué)習(xí)模型的默認(rèn)軟件堆棧將不再是英偉達(dá)的閉源 CUDA。

TensorFlow vs. PyTorch

類似的競(jìng)爭(zhēng)出現(xiàn)在機(jī)器學(xué)習(xí)框架中。幾年前，框架生態(tài)系統(tǒng)相當(dāng)分散，但 TensorFlow 是領(lǐng)跑者。表面看來(lái)，谷歌穩(wěn)操機(jī)器學(xué)習(xí)框架行業(yè)，他們憑借 TensorFlow 設(shè)計(jì)了 AI 應(yīng)用特定加速器 TPU，從而獲得了先發(fā)優(yōu)勢(shì)。

然而現(xiàn)在看來(lái)，PyTorch 贏了，谷歌未能將其先發(fā)優(yōu)勢(shì)轉(zhuǎn)化為新興 ML 行業(yè)的主導(dǎo)地位。如今，谷歌在機(jī)器學(xué)習(xí)社區(qū)中似乎有些被孤立了，因?yàn)樗鼪](méi)有使用 PyTorch 和 GPU，而是使用自己的軟件堆棧和硬件。甚至，谷歌研發(fā)了第二個(gè)機(jī)器學(xué)習(xí)框架 ——JAX，直接與 TensorFlow 競(jìng)爭(zhēng)，這是典型的「谷歌行為」。

一些人認(rèn)為，由于大型語(yǔ)言模型的興起，特別是 OpenAI 的大型語(yǔ)言模型和各種利用 OpenAI API 構(gòu)建的語(yǔ)言模型，谷歌在搜索和自然語(yǔ)言處理方面的主導(dǎo)地位正在減弱。也許這種觀點(diǎn)過(guò)于悲觀，畢竟當(dāng)前大多數(shù)模型的基礎(chǔ)架構(gòu)仍是谷歌開發(fā)的 transformer。

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

Eager 模式可以說(shuō)是標(biāo)準(zhǔn)的腳本執(zhí)行方法，與普通的 Python 代碼沒(méi)什么區(qū)別。這使得調(diào)試和理解代碼更加容易，因?yàn)橛脩艨梢钥吹街虚g操作的結(jié)果，以及模型是如何運(yùn)行的。

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

這種兩階段的方法使得理解和調(diào)試代碼更具挑戰(zhàn)性，因?yàn)樵趫D執(zhí)行結(jié)束之前用戶無(wú)法看到發(fā)生了什么。這類似于「解釋型」與「編譯型」語(yǔ)言，如 Python 與 C++，調(diào)試 Python 更容易，因?yàn)樗墙忉屝驼Z(yǔ)言。

雖然 TensorFlow 現(xiàn)在也默認(rèn)使用 Eager 模式，但研究社區(qū)和大多數(shù)大型科技公司都選擇使用 PyTorch。

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

如果將機(jī)器學(xué)習(xí)模型訓(xùn)練簡(jiǎn)化為最簡(jiǎn)單的形式，影響機(jī)器學(xué)習(xí)模型訓(xùn)練的因素主要有兩點(diǎn)：

• 計(jì)算 ( FLOPS ) ：在每一層內(nèi)運(yùn)行密集矩陣乘法；
• 內(nèi)存帶寬。

以前，影響機(jī)器學(xué)習(xí)訓(xùn)練時(shí)間的主要因素是計(jì)算時(shí)間，等待系統(tǒng)執(zhí)行矩陣乘法。隨著英偉達(dá) GPU 的不斷發(fā)展，這很快就不再是主要問(wèn)題了。

英偉達(dá)利用摩爾定律將 FLOPS 提高了多個(gè)數(shù)量級(jí)，但主要是架構(gòu)變化 —— 張量核（tensor core）和更低精度的浮點(diǎn)數(shù)格式。相比之下，存儲(chǔ)方面沒(méi)有太大的變化。

2018 年，最先進(jìn)的模型是 BERT，英偉達(dá) V100 是最先進(jìn)的 GPU，那時(shí)矩陣乘法已經(jīng)不再是提高模型性能的主要因素。之后，模型在參數(shù)數(shù)量上增長(zhǎng)了 3 到 4 個(gè)數(shù)量級(jí)，而最快的 GPU 在 FLOPS 上增長(zhǎng)了 1 個(gè)數(shù)量級(jí)。

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

內(nèi)存墻

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

最近的共享內(nèi)存池一般是同一芯片上的 SRAM。一些機(jī)器學(xué)習(xí) ASIC 試圖利用巨大的 SRAM 池來(lái)保存模型權(quán)重。但即使是 Cerebras 的價(jià)值約 5000000 美元的晶圓級(jí)芯片也只有 40GB 的 SRAM。內(nèi)存容量不足以容納 100B+ 參數(shù)模型的權(quán)重。

英偉達(dá)設(shè)計(jì)的芯片片上內(nèi)存容量要少得多 ——A100 有 40MB，H100 有 50MB。臺(tái)積電 5 納米芯片上的 1GB SRAM 需要約 200 平方毫米的硅，要實(shí)現(xiàn)相關(guān)的控制邏輯 / 結(jié)構(gòu)，將需要超過(guò) 400 平方毫米的硅。鑒于 A100 GPU 的成本在 1 萬(wàn)美元以上，而 H100 更是接近 2 萬(wàn)美元，從經(jīng)濟(jì)角度看，這種方法是不可行的。即使忽略英偉達(dá)在數(shù)據(jù)中心 GPU 上約 75% 的利潤(rùn)率，對(duì)于完全量產(chǎn)的產(chǎn)品，SRAM 內(nèi)存的成本仍在 100 美元 / GB 左右。

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

內(nèi)存層次結(jié)構(gòu)的下一步是緊密耦合的片外內(nèi)存 DRAM。DRAM 的延遲比 SRAM 高一個(gè)數(shù)量級(jí)（約 100nsVS10ns），但它也便宜得多。幾十年來(lái)，DRAM 一直遵循著摩爾定律。戈登摩爾創(chuàng)造這個(gè)詞時(shí)，英特爾的主要業(yè)務(wù)就是 DRAM。他對(duì)晶體管密度和成本的預(yù)測(cè)在 2009 年之前對(duì) DRAM 普遍適用。不過(guò)自 2012 年以來(lái)，DRAM 的成本幾乎沒(méi)有改善。

然而，人們對(duì)內(nèi)存的需求只增不減。DRAM 現(xiàn)在占服務(wù)器總成本的 50%，逐漸形成所謂的「內(nèi)存墻」。將英偉達(dá) 2016 年的 P100 GPU 與最新的 H100 GPU 進(jìn)行比較，我們會(huì)發(fā)現(xiàn)內(nèi)存容量增加到原來(lái)的 5 倍（16GB → 80GB），F(xiàn)P16 性能增加到 46 倍（21.2 TFLOPS → 989.5 TFLOPS）。

雖然內(nèi)存容量是一個(gè)重要瓶頸，但另一個(gè)瓶頸 —— 內(nèi)存帶寬也非常關(guān)鍵。內(nèi)存帶寬的增加通常是通過(guò)并行性獲得的。雖然如今標(biāo)準(zhǔn) DRAM 的價(jià)格僅為幾美元 / GB，但為了獲得機(jī)器學(xué)習(xí)所需的海量帶寬，英偉達(dá)使用 HBM 內(nèi)存 —— 一種由 3D 堆疊 DRAM 層組成的設(shè)備，需要更昂貴的封裝。HBM 的成本大概是 10-20 美元 / GB，包括封裝和產(chǎn)量成本。

內(nèi)存帶寬和容量的成本限制問(wèn)題在英偉達(dá)的 A100 GPU 中尤為明顯。如果不進(jìn)行大量?jī)?yōu)化，A100 只能具有非常低的 FLOPS 利用率。

即使研究人員進(jìn)行了大量?jī)?yōu)化，大型語(yǔ)言模型的 FLOPS 利用率也只能達(dá)到 60% 左右 。很大一部分時(shí)間花費(fèi)在等待來(lái)自另一個(gè)計(jì)算 / 內(nèi)存的數(shù)據(jù)，或者及時(shí)重新計(jì)算結(jié)果以減少內(nèi)存瓶頸。

從 A100 到 H100，F(xiàn)LOPS 增長(zhǎng)到原來(lái)的 6 倍以上，但內(nèi)存帶寬僅增長(zhǎng)到 1.65 倍。這導(dǎo)致許多人擔(dān)心 H100 的利用率會(huì)很低。A100 需要很多技巧才能繞過(guò)內(nèi)存墻，而 H100 還需要實(shí)現(xiàn)更多技巧。

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

算子融合

如果將所有時(shí)間都花在內(nèi)存?zhèn)鬏斏希刺幱趦?nèi)存帶寬限制狀態(tài)），那么增加 GPU 的 FLOPS 將無(wú)濟(jì)于事。另一方面，如果將所有時(shí)間都花在執(zhí)行大型 matmuls 上，那么即使將模型邏輯重寫為 C++ 來(lái)減少開銷也將無(wú)濟(jì)于事。

PyTorch 之所以能勝過(guò) TensorFlow，就是因?yàn)?Eager 模式提高了靈活性和可用性，但轉(zhuǎn)向 Eager 模式并不是只有好處。在 Eager 模式下運(yùn)行時(shí)，每次運(yùn)算都要從內(nèi)存中讀取、計(jì)算，然后在處理下一次運(yùn)算之前發(fā)送到內(nèi)存。如果不進(jìn)行大量?jī)?yōu)化，這會(huì)顯著增加內(nèi)存帶寬需求。

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

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

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

我們可以說(shuō)軟件開發(fā)人員太懶了，但說(shuō)實(shí)話，又有誰(shuí)沒(méi)懶惰過(guò)呢。一旦習(xí)慣了 PyTorch 中的一個(gè)新算子，他們就會(huì)繼續(xù)用它。開發(fā)人員甚至可能沒(méi)有意識(shí)到性能在提高，而是繼續(xù)使用該算子，因?yàn)檫@樣就不用編寫更多的代碼。

此外，并非所有算子都可以融合。決定要融合哪些運(yùn)算，將哪些運(yùn)算分配給芯片和集群級(jí)別的特定計(jì)算資源都需要花費(fèi)大量的時(shí)間。算子在何處融合的策略雖大體相似，但因?yàn)榧軜?gòu)的不同也會(huì)有很大差異。

英偉達(dá)曾是王者

算子的增長(zhǎng)和默認(rèn)的地位對(duì)英偉達(dá)來(lái)說(shuō)是優(yōu)勢(shì)，因?yàn)槊總€(gè)算子都針對(duì)其架構(gòu)進(jìn)行了快速優(yōu)化，但并未針對(duì)任何其他硬件進(jìn)行優(yōu)化。如果一家 AI 硬件初創(chuàng)公司想要全面實(shí)施 PyTorch，那就意味著以高性能支持不斷增長(zhǎng)的 2000 個(gè)算子列表。

因?yàn)樘崛〉阶畲笮阅苄枰芏嗉记桑?GPU 上訓(xùn)練具有高 FLOPS 利用率的大型模型所需的人才水平越來(lái)越高。Eager 模式執(zhí)行加算子融合意味著開發(fā)的軟件、技術(shù)和模型都在不斷地被推動(dòng)，以適應(yīng)當(dāng)前一代 GPU 具有的計(jì)算和內(nèi)存比率。

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

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

PyTorch 2.0

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

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

對(duì)于當(dāng)前未優(yōu)化的硬件，PyTorch 2.0 具有更大的性能改進(jìn)空間。Meta 和其他公司對(duì) PyTorch 做出如此巨大的貢獻(xiàn)，是因?yàn)樗麄兿Ｍ谧约簝r(jià)值數(shù)十億美元的 GPU 訓(xùn)練集群上以更少的努力實(shí)現(xiàn)更高的 FLOPS 利用率。這樣他們也有動(dòng)力使軟件堆棧更易于移植到其他硬件，將競(jìng)爭(zhēng)引入機(jī)器學(xué)習(xí)領(lǐng)域。

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

PrimTorch

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

TorchDynamo

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

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

在測(cè)試的 7000 個(gè) PyTorch 模型中，TorchDynamo 已經(jīng)適用于 99% 以上的模型，包括來(lái)自 OpenAI、HuggingFace、Meta、英偉達(dá)、Stability.AI 等的模型，而無(wú)需對(duì)原始代碼進(jìn)行任何更改。測(cè)試的 7000 個(gè)模型是從 GitHub 上使用 PyTorch 的最受歡迎項(xiàng)目中隨機(jī)挑選出來(lái)的。

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

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

受保護(hù)的圖捕獲會(huì)檢查捕獲的圖是否對(duì)執(zhí)行有效。「保護(hù)」的意思是一種需要重新編譯的更改。這很重要，因?yàn)槎啻芜\(yùn)行相同的代碼不會(huì)多次重新編譯。如果捕獲的圖對(duì)于執(zhí)行無(wú)效，則即時(shí)重新捕獲允許重新捕獲圖。

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

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

TorchInductor

TorchInductor 是 Python 原生深度學(xué)習(xí)編譯器，可以為多個(gè)加速器和后端生成快速代碼。Inductor 將采用具有約 250 個(gè)算子的 FX 圖，并將它們降低到約 50 個(gè)算子。接著，Inductor 進(jìn)入調(diào)度階段，在該階段融合算子，并確定內(nèi)存規(guī)劃。

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

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

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

OpenAI Triton

對(duì)英偉達(dá)的機(jī)器學(xué)習(xí)閉源軟件來(lái)說(shuō)，OpenAI Triton 是一個(gè)顛覆性的存在。Triton 直接采用 Python 或通過(guò) PyTorch Inductor 堆棧提供數(shù)據(jù)，后者是最常見的用法。Triton 負(fù)責(zé)將輸入轉(zhuǎn)換為 LLVM 中間表征，并生成代碼。英偉達(dá) GPU 將直接生成 PTX 代碼，跳過(guò)英偉達(dá)的閉源 CUDA 庫(kù)（如 cuBLAS），轉(zhuǎn)而使用開源庫(kù)（如 cutlass）。

CUDA 在加速計(jì)算領(lǐng)域很受歡迎，但在機(jī)器學(xué)習(xí)研究人員和數(shù)據(jù)科學(xué)家中卻鮮為人知。使用 CUDA 可能會(huì)帶來(lái)重重挑戰(zhàn)，并且需要深入了解硬件架構(gòu)，這可能導(dǎo)致開發(fā)過(guò)程變慢。因此，機(jī)器學(xué)習(xí)專家可能就要依賴 CUDA 專家來(lái)修改、優(yōu)化和并行化他們的代碼。

Triton 彌補(bǔ)了這一缺陷，使高級(jí)語(yǔ)言實(shí)現(xiàn)了與低級(jí)語(yǔ)言相當(dāng)?shù)男阅堋riton 內(nèi)核本身對(duì)典型的 ML 研究者來(lái)說(shuō)非常清晰，這對(duì)可用性來(lái)說(shuō)非常重要。Triton 在 SM 中自動(dòng)執(zhí)行內(nèi)存合并、共享內(nèi)存管理和調(diào)度。Triton 對(duì)逐元素矩陣乘法不是特別有用，但矩陣乘法已經(jīng)可以非常高效地完成。Triton 對(duì)于成本高昂的逐點(diǎn)運(yùn)算和減少?gòu)?fù)雜操作的開銷非常有用。

OpenAI Triton 目前僅正式支持英偉達(dá)的 GPU，但在不久的將來(lái)會(huì)發(fā)生變化，將支持多個(gè)其他硬件供應(yīng)商。其他硬件加速器可以直接集成到 Triton 的 LLVM IR 中，這大大減少了為新硬件構(gòu)建 AI 編譯器堆棧的時(shí)間。

英偉達(dá)龐大的軟件體系缺乏遠(yuǎn)見，無(wú)法利用其在 ML 硬件和軟件方面的巨大優(yōu)勢(shì)，也就沒(méi)能成為機(jī)器學(xué)習(xí)的默認(rèn)編譯器。他們?nèi)狈?duì)可用性的關(guān)注，而 OpenAI 和 Meta 也正是得益于此才能夠創(chuàng)建出可移植到其他硬件的軟件堆棧。

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