過(guò)去十年機(jī)器學(xué)習(xí)軟件開(kāi)發(fā)行業(yè)概覽：英偉達(dá) CUDA 壟斷地位下降，PyTorch 超越谷歌 Tensorflow

作者 | Dylan Patel

譯者 | 馬可薇

策劃 | Tina

在過(guò)去的十年間，機(jī)器學(xué)習(xí)軟件開(kāi)發(fā)的格局翻天覆地。流水的框架鐵打的英偉達(dá)，這些框架大多都極度依賴(lài)英偉達(dá)的 CUDA（統(tǒng)一計(jì)算架構(gòu)），且在英偉達(dá) GPU 上性能最好。不過(guò)，隨著 PyTorch 2.0 和 OpenAI 的 Triton 的到來(lái)，英偉達(dá)因其軟件護(hù)城河而在這一領(lǐng)域的制霸地位恐將不保。

本篇報(bào)告中將涵蓋以下主題：為何谷歌的 TensorFlow 輸給了 PyTorch，谷歌為何沒(méi)能利用其在人工智能領(lǐng)域的早期領(lǐng)導(dǎo)地位，機(jī)器學(xué)習(xí)模型訓(xùn)練時(shí)間的主要構(gòu)成成分，內(nèi)存容量、帶寬、成本墻，模型優(yōu)化，為何其他人工智能硬件公司至今無(wú)法撼動(dòng)英偉達(dá)的統(tǒng)治地位，為何硬件地位逐漸重要，英偉達(dá)在 CUDA 上的競(jìng)爭(zhēng)優(yōu)勢(shì)是如何消失的，以及英偉達(dá)的競(jìng)爭(zhēng)對(duì)手在大型云的芯片訓(xùn)練上所取得的重大勝利。

問(wèn)題簡(jiǎn)述是，英偉達(dá)閉源 CUDA 將不再涵蓋機(jī)器學(xué)習(xí)模型的默認(rèn)軟件棧。英偉達(dá)雖然搶占先機(jī)，但卻讓 OpenAI 和 Meta 后來(lái)居上掌控了軟件棧，英偉達(dá)專(zhuān)有工具的失敗致使后者在生態(tài)系統(tǒng)中建立了自己的工具，英偉達(dá)的護(hù)城河也遭到了永久地削弱。

TensorFlow vs. PyTorch

僅僅幾年前，在相當(dāng)分散的框架生態(tài)系統(tǒng)中，掌握最常用框架 TensorFlow 的谷歌作為領(lǐng)跑者，設(shè)計(jì)并部署了唯一成功的人工智能指定應(yīng)用加速器 TPU，他們看似已經(jīng)搶占先機(jī)、準(zhǔn)備好制霸機(jī)器學(xué)習(xí)行業(yè)了。

2019 年機(jī)器學(xué)習(xí)框架狀態(tài)，PyTorch 掌控研究領(lǐng)域，Tensorflow 掌控行業(yè)領(lǐng)域

可事實(shí)卻是 PyTorch 贏了，而谷歌沒(méi)能將先手優(yōu)勢(shì)轉(zhuǎn)換為對(duì)新生機(jī)器學(xué)習(xí)行業(yè)的主導(dǎo)權(quán)。如今的谷歌使用自研軟硬件棧而非 PyTorch 和對(duì)應(yīng) GPU，導(dǎo)致其在機(jī)器學(xué)習(xí)社區(qū)的地位頗為尷尬，谷歌甚至還有另一套名為 Jax 的框架，直接和 TensorFlow 相競(jìng)爭(zhēng)。

不僅如此，關(guān)于谷歌在搜索及自然語(yǔ)言處理方面的主導(dǎo)地位會(huì)因大型語(yǔ)言模型而衰退之類(lèi)的言論也是甚囂塵上，尤其是來(lái)自那些 OpenAI 及各種利用 OpenAI 的 API 的、或基于類(lèi)似基礎(chǔ)模型的初創(chuàng)公司。雖說(shuō)可能是有些杞人憂(yōu)天，但我們今天不談這些。目前來(lái)說(shuō)，雖然挑戰(zhàn)不斷，但谷歌仍是處于機(jī)器學(xué)習(xí)模型的最前沿。谷歌所發(fā)明的 Transformer 在 PaLM、LaMBDA、Chinchilla、MUM，以及 TPU 等諸多領(lǐng)域依舊是最先進(jìn)的。

讓我們回到 PyTorch 贏得一籌的話(huà)題中去。雖然也有從谷歌手中奪取掌控權(quán)的成分在，但 PyTorch 主要還是贏在了其相對(duì) TensorFlow 而言更高的靈活性和可用性，從原則為先的角度來(lái)看，PyTorch 與 TensorFlow 的不同在于前者使用的是“動(dòng)態(tài)圖模式（Eager Mode）”而非“圖模式（Graph Mode）”。

動(dòng)態(tài)圖模式可被看作是一種標(biāo)準(zhǔn)的腳本執(zhí)行方式。深度學(xué)習(xí)框架會(huì)隨調(diào)用立即逐行執(zhí)行所有操作，這點(diǎn)和任何的 Python 代碼執(zhí)行都一樣。因此，代碼的調(diào)試和理解也都更加容易，操作間的結(jié)果和模型表現(xiàn)也更直觀。

相較之下，圖模式則有兩個(gè)階段，第一階段定義需要執(zhí)行操作的計(jì)算圖，其中計(jì)算圖是一系列代表操作、變量的相交節(jié)點(diǎn)，節(jié)點(diǎn)相連的邊則代表其間的數(shù)據(jù)流。第二階段定義延遲執(zhí)行計(jì)算圖的優(yōu)化版本。因?yàn)閳D的執(zhí)行過(guò)程中我們無(wú)從得知到底發(fā)生著什么，所以這種分階段的方式讓代碼調(diào)試更具挑戰(zhàn)，也更難理解。這與“解釋性”和“編譯性”編程語(yǔ)言類(lèi)似，可解釋的 Python 調(diào)試要比 C 更容易。

盡管現(xiàn)在 TensorFlow 也默認(rèn)擁有了動(dòng)態(tài)圖模式，但研究社區(qū)和多數(shù)大型技術(shù)公司都已經(jīng)習(xí)慣了 PyTorch 的解決方案。至于 PyTorch 更勝一籌的深層解釋?zhuān)?qǐng)見(jiàn)這里?？傊?，主流 AI 大會(huì) NeurIPS 上最優(yōu)秀的那些非谷歌的生成性人工智能都是用的 PyTorch。

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

追根究底，影響機(jī)器學(xué)習(xí)的模型訓(xùn)練耗時(shí)主要有兩個(gè)因素:1. 計(jì)算（FLOPS），即每層內(nèi)運(yùn)行密集的矩陣乘法 2. 內(nèi)存（帶寬），即等待數(shù)據(jù)或?qū)訖?quán)重獲取計(jì)算資源。常見(jiàn)受帶寬限制的操作有歸一化、點(diǎn)式操作、SoftMax、ReLU。

曾經(jīng)主要影響機(jī)器學(xué)習(xí)訓(xùn)練時(shí)長(zhǎng)的計(jì)算時(shí)間、等待矩陣乘法等因素，隨著英偉達(dá) GPU 的不斷發(fā)展都已不再重要。英偉達(dá)的 FLOPS 在摩爾定律下英偉達(dá)的 FLOPS 在摩爾定律下提升了多個(gè)數(shù)量級(jí)，但主要架構(gòu)變化還是集中在張量核心及低精度浮點(diǎn)格式上，內(nèi)存方面則沒(méi)有太多變化。

英偉達(dá) GPU 增長(zhǎng)趨勢(shì)

2018 年的 BERT 模型和英偉達(dá)的 GPU V100 均是時(shí)代頂尖產(chǎn)品，我們不難發(fā)現(xiàn)，矩陣乘法已經(jīng)不再是提高模型性能的主要因素。自此之后，最為先進(jìn)的模型在參數(shù)數(shù)量上有了 3 至四個(gè)數(shù)量級(jí)的增長(zhǎng)，而最快的 GPU 也有了一個(gè)數(shù)量級(jí)的增長(zhǎng)。

PyTorch中操作類(lèi)占比

即使是在 2018 年，純粹的計(jì)算性質(zhì)工作負(fù)載在占據(jù) 99.8%的 FLOPS 同時(shí)，僅占用了 61%的運(yùn)行時(shí)間。歸一化和逐點(diǎn)操作相較矩陣乘法而言，分別擁有 250 倍和 700 倍 FLOPS 的減少，但同時(shí)也消耗了模型近 40%的運(yùn)行時(shí)間。

內(nèi)存墻

隨著模型規(guī)模的不斷擴(kuò)張，大型語(yǔ)言模型的權(quán)重本身便占據(jù)了千億、乃至太字節(jié)。百度和 Meta 所部署的生產(chǎn)型推薦網(wǎng)絡(luò)其中的大規(guī)模嵌入表就可占用幾十兆字節(jié)內(nèi)存，大型模型訓(xùn)練或推理中的大部分時(shí)間并沒(méi)有花在矩陣乘法的計(jì)算，而是在等待數(shù)據(jù)到達(dá)計(jì)算資源?；蛟S你會(huì)問(wèn)，為什么架構(gòu)師不把更多內(nèi)存放在靠近計(jì)算模塊的地方？答案只有一個(gè)字，貴。

存儲(chǔ)層級(jí)

存儲(chǔ)層級(jí)所遵循的規(guī)律是從近且快，到慢但廉價(jià)的。共享內(nèi)存池最近可以在同一塊芯片上，通常這種情況是由 SRAM（靜態(tài)隨機(jī)存取存儲(chǔ)器）組成。有的機(jī)器學(xué)習(xí) ASIC 嘗試?yán)么笮?SRAM 池保存模型權(quán)重，但即使是Cerebras價(jià)值約5百萬(wàn)美元的晶圓規(guī)模芯片上也只有 40G 的 SRAM。我們沒(méi)有足夠的內(nèi)存來(lái)容納 100B 以上的參數(shù)模型權(quán)重。

英偉達(dá)的架構(gòu)常常會(huì)在芯片上使用更為少量的內(nèi)存，當(dāng)前一代 A100 包括了 40MB 內(nèi)存，而下一代的 H100 上也只有 50MB。臺(tái)積電 5 納米工藝節(jié)點(diǎn)上 1GB 的 SRAM 需要大約 200 平方毫米的硅，加上相關(guān)的控制邏輯和結(jié)構(gòu)的實(shí)現(xiàn)后就需要超過(guò) 400 平方毫米的硅，或使用英偉達(dá)數(shù)據(jù)中心 GPU 總邏輯面積的 50%左右?？紤]到 A100 GPU 超過(guò) 1 萬(wàn)美元的價(jià)格，H100 的價(jià)格很可能也會(huì)兩萬(wàn)美元起步，經(jīng)濟(jì)層面上這條路行不通。就算我們忽略英偉達(dá)數(shù)據(jù)中心 GPU 的 75%毛利率（約為四倍加價(jià)），實(shí)現(xiàn)產(chǎn)品的完全產(chǎn)出所需要每 GB 的 SRAM 內(nèi)存，成本仍在一百美元上下。

此外，芯片上 SRAM 內(nèi)存的花銷(xiāo)并不會(huì)隨著傳統(tǒng)摩爾定律所帶來(lái)的工藝技術(shù)縮減而下降太多，在下一代臺(tái)積電 3 納米工藝技術(shù)下，同樣 1GB 的內(nèi)存成本實(shí)際是在增長(zhǎng)的。3D SRAM 雖然會(huì)在一定程度上降低 SRAM 成本，但這也只是價(jià)格曲線(xiàn)的暫時(shí)下跌。

存儲(chǔ)層次中的下一層是緊密耦合的片外內(nèi)存 DRAM。DRAM 相較于 SRAM，延遲要高上一個(gè)數(shù)量級(jí)（約 100 納秒和 10 納秒的區(qū)別），但 DRAM 也要便宜許多（DRAM 每 GB 一美元，SRAM 每 GB 一百美元）

數(shù)十年來(lái) DRAM 一直遵循摩爾定律，事實(shí)上，在戈登·摩爾創(chuàng)造“摩爾定律”這個(gè)詞時(shí)，英特爾的主要業(yè)務(wù)就是 DRAM。摩爾關(guān)于晶體管密度與成本的經(jīng)濟(jì)預(yù)測(cè)對(duì) 2009 年之前的 DRAM 通常都是準(zhǔn)確的，但自 2012 年來(lái)，DRAM 的成本幾乎沒(méi)有提升。

DRAM 每 GB 價(jià)格

我們對(duì)內(nèi)存的需求只增不減，目前 DRAM 已經(jīng)占據(jù)了服務(wù)器總成本的50%。內(nèi)存墻的存在已經(jīng)開(kāi)始在產(chǎn)品中顯露出來(lái)了。相比英偉達(dá) 2016 年的 GPU P100，2022 年剛剛發(fā)售的 GPU H100 在內(nèi)存容量上提升了五倍（16GB 到 80GB 的提升），而 FP16 性能卻提升了足足 46 倍（21.1 TFLOPS 至 989.5 TFLOPS）。

容量瓶頸與另一同樣重要的帶寬瓶頸息息相關(guān)。并行化是增加內(nèi)存帶寬的主要手段，如今的 DRAM 每 GB 價(jià)格區(qū)區(qū)幾美元，而英偉達(dá)為了達(dá)到機(jī)器學(xué)習(xí)所需要的巨大帶寬而用上了 HBM 內(nèi)存，這是一種由3D堆疊的DRAM層所組成的設(shè)備，且需要更為昂貴的包裝。HBM 的價(jià)格區(qū)間為 10 至 20 美元每 GB，其中包含有包裝與產(chǎn)量成本。

內(nèi)存帶寬與容量限制在英偉達(dá) A100 GPU 中被反復(fù)提及。沒(méi)有大量?jī)?yōu)化過(guò)的 A100 常常會(huì)有極低的 FLOPS 利用率，F(xiàn)LOPS 利用率是通過(guò)計(jì)算（模型訓(xùn)練時(shí)總計(jì)算 FLOPS）/(GPU 在模型訓(xùn)練時(shí)間內(nèi)理論可計(jì)算 FLOPS）而得出。

即使是在頂尖研究者所做的大量?jī)?yōu)化下，60%的 FLOPS 利用率對(duì)于大型語(yǔ)言模型訓(xùn)練而言也算是非常高的了。剩下的時(shí)間都是開(kāi)銷(xiāo)，包括等待其他計(jì)算或內(nèi)存數(shù)據(jù)的空閑期，或?yàn)闇p少內(nèi)存瓶頸而進(jìn)行即時(shí)重新計(jì)算結(jié)果。

FLOPS 在 A100 至 H100 兩代間增長(zhǎng)了 6 倍有余，但內(nèi)存帶寬卻只有 1.65 倍的增長(zhǎng)，從而導(dǎo)致了許多對(duì) H100 低利用率問(wèn)題的擔(dān)憂(yōu)。人們?yōu)樽?A100 繞過(guò)內(nèi)存墻搞出了許多變通方案，而這種努力在 H100 上恐怕會(huì)只多不少。

H100為Hopper架構(gòu)帶來(lái)了分布式共享內(nèi)存和二級(jí)組播，其中不同 SM（可看作內(nèi)核）可直接寫(xiě)入其他 SM 的 SRAM（共享內(nèi)存/L1 緩存）。此舉在有效增加緩存大小的同時(shí)，縮減了DRAM讀寫(xiě)所需的帶寬。后續(xù)架構(gòu)也將通過(guò)減少向內(nèi)存?zhèn)鬏數(shù)牟僮骶徑鈨?nèi)存墻的影響。值得注意的是，因?yàn)閷?duì) FLOPS 的需求會(huì)隨參數(shù)數(shù)量增加而立方擴(kuò)展，對(duì)內(nèi)存帶寬及容量的需求則常呈二次曲線(xiàn)發(fā)展，所以較大型的模型也更傾向于實(shí)現(xiàn)更高的利用率。

算子融合 – 治標(biāo)不治本

同機(jī)器學(xué)習(xí)模型訓(xùn)練一樣，明白自己所處狀態(tài)才能更精確地進(jìn)行重要的優(yōu)化。舉例來(lái)說(shuō)，如果我們處于內(nèi)存帶寬約束的狀態(tài)，時(shí)間全花在了內(nèi)存?zhèn)鬏斏希敲丛黾?GPU 的 FLOPS 并不能解決問(wèn)題。而如果是處于計(jì)算約束的狀態(tài)，笨重的矩陣乘法非常耗時(shí)的話(huà)，那么試圖通過(guò)將模型邏輯改寫(xiě)成 C 來(lái)削減開(kāi)銷(xiāo)也是沒(méi)效果的。

雖然 PyTorch 是通過(guò)動(dòng)態(tài)圖模式增加了靈活性和可用性而贏得的比賽，但動(dòng)態(tài)圖模式也不是完美的。在動(dòng)態(tài)圖模式中執(zhí)行的每個(gè)操作都需要從內(nèi)存中讀取、計(jì)算，再發(fā)送到內(nèi)存之后，才能處理下一個(gè)操作。在缺乏大量?jī)?yōu)化的情況下，這種模式將大大增加對(duì)內(nèi)存帶寬的需求。

因此，動(dòng)態(tài)圖模式中執(zhí)行模型的主要優(yōu)化手段之一是算子融合。用融合操作符的方式取代將每次的中間結(jié)果寫(xiě)入內(nèi)存，在一次計(jì)算中計(jì)算多個(gè)函數(shù)，從而盡可能減少對(duì)內(nèi)存的讀寫(xiě)。算子融合改善了操作符的調(diào)度，也削減了內(nèi)存帶寬和容量的成本。

算子融合簡(jiǎn)圖

這種優(yōu)化方式常常需要編寫(xiě)自定義 CUDA 內(nèi)核，這可比簡(jiǎn)單使用 Python 腳本要難多了。PyTorch 的內(nèi)置變通方案長(zhǎng)期以來(lái)在 PyTorch 內(nèi)部實(shí)現(xiàn)了越來(lái)越多的操作符，這其中的許多操作符都只是將多個(gè)常用操作融合到一個(gè)更為復(fù)雜的函數(shù)中。

操作符的增加讓 PyTorch 中的模型創(chuàng)建更加輕松，隨著內(nèi)存讀寫(xiě)次數(shù)的減少，動(dòng)態(tài)圖模式的性能也更快。但隨之而來(lái)的代價(jià)是 PyTorch 中運(yùn)算符在短短幾年內(nèi)就膨脹至兩千有余。

人們常說(shuō)軟件開(kāi)發(fā)者是懶人，但老實(shí)說(shuō)又有誰(shuí)不是呢。在熟悉了 PyTorch 中某個(gè)新增的操作符后，開(kāi)發(fā)者們通常只會(huì)覺(jué)得這個(gè)新操作符能讓自己少些點(diǎn)代碼，而完全沒(méi)有意識(shí)到其中的性能提高。

此外，并不是所有的操作都能融合。大部分時(shí)間我們都在決定要融合哪些操作，又要將哪些操作分配給芯片和集群層面特定的計(jì)算資源。雖然一般來(lái)說(shuō)算子融合的策略都多少有些類(lèi)似，但根據(jù)架構(gòu)的不同也會(huì)有所區(qū)別。

英偉達(dá)稱(chēng)王

操作符數(shù)量的發(fā)展和其默認(rèn)的首選地位讓英偉達(dá)受益很多，每個(gè)被迅速優(yōu)化的操作符都是針對(duì)英偉達(dá)架構(gòu)的，且對(duì)其他硬件并不適用。任何想要完整實(shí)現(xiàn) PyTorch 的人工智能硬件初創(chuàng)公司，都得靠高性能才能原生支持兩千多個(gè)且還在不斷增長(zhǎng)的操作符。

在 GPU 上訓(xùn)練高 FLOPS 利用率的大型模型所需的技術(shù)水平越來(lái)越高，為達(dá)到最佳性能所需花樣也越發(fā)繁多。動(dòng)態(tài)圖模式的執(zhí)行和算子融合，意味著軟件、技術(shù)、模型的開(kāi)發(fā)都要被迫適應(yīng)最新一代 GPU 的計(jì)算和內(nèi)存比例范圍。

內(nèi)存墻是任何機(jī)器學(xué)習(xí)芯片的開(kāi)發(fā)者都逃不開(kāi)的命運(yùn)。ASIC 必須要能支持常用框架的同時(shí)，也要能支持混合使用英偉達(dá)及外部庫(kù)的、基于 GPU 優(yōu)化的 PyTorch 代碼這一默認(rèn)開(kāi)發(fā)策略。因此，為圖更高 FLOPS 及更嚴(yán)格的編程模型，而主動(dòng)放棄 GPU 的各種非計(jì)算包袱這種行為是非常沒(méi)有意義的。

易用性是王道。

要打破這種惡性循環(huán)的唯一方式是將英偉達(dá) GPU 上運(yùn)行模型的軟件，盡可能地?zé)o縫轉(zhuǎn)移至其他硬件上。隨著 PyTorch 2.0、OpenAI Triton，以及諸如MosaicML的MLOps公司所提供的模型架構(gòu)穩(wěn)定性和抽象逐漸得到主流承認(rèn)，芯片解決方案的架構(gòu)和性?xún)r(jià)比逐漸取代了英偉達(dá)卓越的軟件所帶來(lái)的易用性，成為驅(qū)動(dòng)購(gòu)買(mǎi)力的主要因素。

PyTorch 2.0

數(shù)月之前剛剛成立的PyTorch基金會(huì)正式脫離了Meta的掌握。在向開(kāi)放式開(kāi)發(fā)和管理模式轉(zhuǎn)變的同時(shí)，2.0 的早期測(cè)試版本已經(jīng)發(fā)布，并預(yù)計(jì)于 2023 年三月全年上市。PyTorch 2.0 與前代最主要的區(qū)別在于，新增的一個(gè)支持圖執(zhí)行模型的編譯解決方案，讓各種硬件資源的利用更加輕松。

PyTorch 2.0 在英偉達(dá) A100 上的訓(xùn)練性能有了86%的提升，CPU上的推理則有26%的提升，極大地縮減了模型訓(xùn)練所需的計(jì)算時(shí)間和成本。而這種性能提升也可以類(lèi)推至包括AMD、英特爾、Tenstorrent、Luminous Computing、特斯拉、谷歌、亞馬遜、微軟、Marvell、Meta、Graphcore、Cerebras、SambaNova 在內(nèi)的多個(gè) GPU 和加速器上。

PyTorch 2.0 的性能提升在未經(jīng)優(yōu)化的硬件上更為明顯。Meta 及其他公司對(duì) PyTorch 的大量貢獻(xiàn)背后，是希望能在他們數(shù)十億美元的訓(xùn)練集群上，以最小的努力實(shí)現(xiàn)更高的 FLOPS 利用率，讓他們的軟件棧更易于移植到其他硬件上，從而為機(jī)器學(xué)習(xí)領(lǐng)域引入新的競(jìng)爭(zhēng)力。

分布式訓(xùn)練也受益于 PyTorch 2.0，數(shù)據(jù)并行、分片、管道并行及張量并行均得到了更優(yōu)秀的 API 支持。除此之外，PyTorch 2.0 也通過(guò)全棧提供了對(duì)動(dòng)態(tài)圖形的原生支持，讓LLM不同序列長(zhǎng)度等更易于支持。這也是主流編譯器第一次支持從訓(xùn)練到推理的動(dòng)態(tài)形狀。

PrimTorch

對(duì)任何非英偉達(dá) GPU 之外的任何機(jī)器學(xué)習(xí) ASIC 來(lái)說(shuō)，想要編寫(xiě)一個(gè)完整支持全部?jī)汕в鄠€(gè)操作符的高性能后端是非常具有挑戰(zhàn)性的。而 PrimTorch 卻可以在保障 PyTorch 終端用戶(hù)可用性不變的前提下，將操作符數(shù)量減少至約 250 個(gè)原始操作符，讓非英偉達(dá)的 PyTorch 后端實(shí)現(xiàn)更簡(jiǎn)單容易，定制硬件和操作系統(tǒng)的供應(yīng)商也更容易提出自己的軟件棧。

TorchDynamo

穩(wěn)健的圖定義是向圖模式轉(zhuǎn)變的必需品，而過(guò)去五年間 Meta 和 PyTorch 在這方面解決方案的嘗試都有著明顯的缺陷，直到 TorchDynamo 的出現(xiàn)。TorchDynamo 可接收任何 PyTorch 用戶(hù)腳本并生成FX圖，甚至是調(diào)用三方外部庫(kù)的腳本也可以。

Dynamo 將所有復(fù)雜操作都?jí)嚎s為 PrimTorch 中約 250 個(gè)原始操作符。圖成型后所有未使用的操作會(huì)被棄置，成型的圖決定了有哪些中間操作需要被存儲(chǔ)或?qū)懭雰?nèi)存，有哪些可以被融合。這種方式極大地削減了模型內(nèi)開(kāi)銷(xiāo)，對(duì)用戶(hù)而言也是無(wú)感的。

目前在不修改任何源碼的前提下，TorchDynamo已在超過(guò)七千個(gè)PyTorch模型上通過(guò)了可行性測(cè)試，其中不乏來(lái)自 OpenAI、HuggingFace、Meta、英偉達(dá)、Stability.AI 的模型。這七千多個(gè)模型是從 GitHub 上熱度最高的 PyTorch 項(xiàng)目中直接選取的。

谷歌的 TensorFlow、Jax 及其他圖模式的執(zhí)行管道，通常需要用戶(hù)自行保障模型對(duì)編譯器架構(gòu)的兼容性，才能確保圖可以被捕獲。而 Dynamo 通過(guò)啟用部分圖捕獲、受保護(hù)的圖捕獲及即時(shí)重新捕獲進(jìn)行改善。

• 部分圖捕獲允許模型包括不支持或非 Python 的結(jié)構(gòu)。在無(wú)法生成圖的模型構(gòu)造部分插入圖斷點(diǎn)，并在部分圖之間以動(dòng)態(tài)圖模式執(zhí)行。
• 受保護(hù)圖捕獲校驗(yàn)被捕獲的圖是否可有效執(zhí)行。保護(hù)是指需要重新編譯的代碼變更，畢竟多次重復(fù)執(zhí)行的同一段代碼并不會(huì)重新編譯。
• 即時(shí)重新捕獲允許無(wú)效執(zhí)行的圖重新被捕獲。

PyTorch 意圖創(chuàng)建一個(gè)依賴(lài) Dynamo 圖生成的、統(tǒng)一且流暢的 UX。這項(xiàng)解決方案在不改變用戶(hù)體驗(yàn)的同時(shí)顯著提高性能，而圖捕獲意味著在大型計(jì)算資源的基礎(chǔ)上執(zhí)行可以更高效地并行進(jìn)行。

Dynamo 及AOT自動(dòng)求導(dǎo)會(huì)在之后將優(yōu)化后的 FX 圖傳入 PyTorch 本地編譯器層，即 TorchInductor。其他硬件企業(yè)也可直接取用此時(shí)的圖并輸入至他們自己的后端編譯器中。

TorchInductor

作為原生的 Python 深度學(xué)習(xí)編譯器，TorchInductor 可為多個(gè)加速器和后端生成快速代碼。Inductor（電感器）可接收包含約 250 個(gè)操作符的 FX 圖，并進(jìn)一步將其操作符數(shù)量削減至 50 左右。在這之后，Inductor 會(huì)進(jìn)入調(diào)度階段，融合算子并確定內(nèi)存規(guī)劃。

之后，Inductor 會(huì)進(jìn)入“代碼封裝”階段，生成可在 CPU、GPU 及其他人工智能加速器上運(yùn)行的代碼。封裝后的代碼可調(diào)用內(nèi)核并分配內(nèi)存，取代了編譯器堆棧中解釋器的部分。其中，后端代碼的生成部分借助 OpenAI 的 GPU 語(yǔ)言 Triton，輸出 PTX 代碼。對(duì) CPU 而言，英特爾編譯器所生成的 C 代碼也可以在非英特爾的 CPU 上運(yùn)行。

后續(xù)還會(huì)新增更多對(duì)硬件的支持，但 Inductor 確實(shí)顯著降低了在編寫(xiě) AI 硬件加速器的編譯器時(shí)所需的工作量。此外，代碼性能也得到了優(yōu)化，對(duì)內(nèi)存帶寬和容量的要求也大大地降低了。

我們寫(xiě)的編譯器不能只支持 GPU，而要能擴(kuò)展到對(duì)各類(lèi)硬件后端的支持。C 及（OpenAI）Triton 迫使著我們一定要具備這種通用性?！狫ason Ansel – Meta AI

OpenAI Triton

OpenAI 的 Triton 語(yǔ)言對(duì)英偉達(dá)機(jī)器學(xué)習(xí)閉源軟件的護(hù)城河有著毀滅性的打擊。Triton 可直接接收 Python 腳本，或者更常見(jiàn)地接收通過(guò)PyTorch的Inductor堆棧的信息流。隨后 Triton 會(huì)將輸入轉(zhuǎn)換為 LLVM 中間表示并生成代碼，使用 cutlass 等開(kāi)源庫(kù)取代英偉達(dá)的閉源 CUDA 庫(kù)（如 cuBLAS）。

CUDA 在專(zhuān)職于加速計(jì)算的開(kāi)發(fā)者中更為常用，在機(jī)器學(xué)習(xí)研究者或數(shù)據(jù)科學(xué)家之間則沒(méi)什么知名度。高效地使用 CUDA 并不容易，需要使用者對(duì)硬件架構(gòu)有深入理解，并可能會(huì)拖慢開(kāi)發(fā)過(guò)程。因此，機(jī)器學(xué)習(xí)專(zhuān)家們常常會(huì)依賴(lài) CUDA 專(zhuān)家對(duì)他們的代碼進(jìn)行修改、優(yōu)化，以及并行化。

而 Triton 則彌補(bǔ)了這一差距，讓高層語(yǔ)言達(dá)到與底層語(yǔ)言相媲美的性能水平。Triton 的內(nèi)核本身對(duì)一般的機(jī)器學(xué)習(xí)研究者而言具備可讀性，這一點(diǎn)對(duì)語(yǔ)言可用性非常重用。Triton 將內(nèi)存凝聚、共享內(nèi)存管理，以及 SM 內(nèi)部的調(diào)度全部自動(dòng)化，但對(duì)元素層面的矩陣乘法沒(méi)什么太大幫助，后者本身已經(jīng)足夠高效了。此外，Triton 在昂貴的逐點(diǎn)操作方面很有效，對(duì)涉及矩陣乘法的大型算子融合而言，也可明顯削減復(fù)雜如Flash注意力等操作的開(kāi)銷(xiāo)。

時(shí)至今日，OpenAI 的 Triton 才正式支持英偉達(dá)的 GPU，不過(guò)很快就要不同了。數(shù)個(gè)其他硬件供應(yīng)商都會(huì)在后續(xù)對(duì)其提供支持，這項(xiàng)開(kāi)源項(xiàng)目的前途一片光明。其他硬件加速器能夠直接集成至 Triton 中 LLVM IR，意味著在新硬件上建立人工智能編譯器堆棧的時(shí)間將大幅縮短。

英偉達(dá)龐大的軟件組織缺乏遠(yuǎn)見(jiàn)，沒(méi)能利用自己在機(jī)器學(xué)習(xí)軟硬件方面的巨大優(yōu)勢(shì)，一舉成為機(jī)器學(xué)習(xí)的默認(rèn)編譯器。英偉達(dá)對(duì)可用性關(guān)注的缺失讓外界中 OpenAI 及 Meta 得以開(kāi)發(fā)出向其他硬件方向移植的軟件棧。他們?yōu)槭裁礇](méi)能為機(jī)器學(xué)習(xí)研究者們開(kāi)發(fā)出一個(gè)像是 Triton 之類(lèi)的*簡(jiǎn)化版*CUDA？為什么像是Flash注意力一類(lèi)的技術(shù)是出自一個(gè)博士生而不是英偉達(dá)本身？

本篇報(bào)告中的剩余部分將會(huì)列出能讓微軟拿下一城的具體硬件加速器，以及目前正被快速集成至 PyTorch 2.0 或 OpenAI Trion 軟件棧中多家公司的硬件產(chǎn)品。此外，報(bào)告中也將列出相反觀點(diǎn)，為英偉達(dá)在人工智能培訓(xùn)領(lǐng)域的護(hù)城河或?qū)嵙μ峁┺q護(hù)。

查看英文原文：How Nvidia’s CUDA Monopoly In Machine Learning Is Breaking – OpenAI Triton And PyTorch 2.0

本文轉(zhuǎn)載來(lái)源：

https://www.infoq.cn/article/9TFWk1rM2h8hLJUvbUG0