獨家：誰在“掏空”深度學習框架PyTorch？

好幾年沒寫AI框架大稿了，

我先把結論放在標題里，

若想細看其中原委，

只能說，大幾千字閱讀量是跑不了的。

01.靈魂三問

PyTorch最近有點煩，

歷經多年碼海沉浮，

拼快慢，爭靜動，

較難易，掀生態。

看似地位穩如泰山：

“主流框架”

“事實工業標準”

“生態百花齊放”，

開發者“簞食壺漿以迎王師”，

“框架”作為AI這個熱門賽道里的基礎軟件，

如果愿意多花點時間觀察，

你就會發現，時間花完，

雖然沒看懂，但起碼累著了。

在此，要向這幾年參與AI框架技術的大牛們致敬，

還要向全球開源框架開發者致敬，

感謝你們貢獻了開源。

AI芯片需要框架翻譯才能釋放性能，

AI生產力需要框架解放。

人間為飲千歲，吞吐江海流霞。

近五年，我很多好朋友都出自這個領域，

感謝他們的認可，這是我寫作的不竭動力。

我長話短說，

不過，實在是說來話長。

變量太猛，戰局有變。

先來個靈魂三問：

1. 北美大模型玩家也被PyTorch統一了嗎？

2. 誰在掏空PyTorch？

3. 開發者視線轉移意味著什么？

02游戲結束了？

"The game is over.”

直譯：游戲結束了。

意譯：除了PyTorch，別的AI框架玩家都沒機會了。

這是某網友對AI框架戰局的評語。

我只能說，情緒到位了，邏輯還差點意思。

甚至說，是僅僅觀察了國內市場后得出的結論。

先看數據，

PyTorch一統江湖。

不得不說，真是帥啊。

再帥也有煩的時候。

因為整體市場占有率是一回事，

巨頭份額又是另外一回事。

當然，無論什么底層軟件，都想上大業務。

那么問題來了，

大模型巨頭玩家到底在用哪種AI框架？

尤其是，巨頭玩家的偏好和中小玩家不同。

巨頭玩家要的不是一個“現成的萬能工具箱”，

而是一個“按需定制的軍火庫”。

OpenAI用的是PyTorch，

但幾乎重寫了底層邏輯，

以適應自己的大模型架構。

谷歌用JAX，它為TPU量身定制。

華為用MindSpore框架，

背后是昇騰AI處理器。

為什么不繼續用PyTorch呢？

開源框架是基礎，不是答案。

開源框架適合普通開發者，

而巨頭們往往需要在此基礎上深度定制。

親愛的數據 原創漫畫

PyTorch作為框架“一哥”，

并不能保證它的原生版本，

就能直接滿足GPT-4或Llama3的全部需求。

所以，當你問“大模型巨頭用什么框架”時，

真正的問題應該是：

這些巨頭在訓練和推理大模型的時候，

到底用的什么秘密的框架技術？

答案可能不會直接寫在開源框架的市場份額里，

而是藏在巨頭們自家技術棧的深處。

這個時候，譚老師就出場了。

直接人到美國，現場做題，

當然，答案僅供參考。

美國巨頭的幾家，我們掰著手指頭數數看。

谷歌，Anthropic，xAI，蘋果，這四家訓練大模型用JAX，

另外兩家，Meta和OpenAI，用PyTorch。

這個大致的調查結果，非常神奇：

北美大模型巨頭的框架份額排第一的是JAX。

親愛的數據，原創漫畫

此時此刻，回望大洋彼岸，

國內模型框架仍然是PyTorch占據主流。

我只能感慨一句，

在框架這件事上，有些人保守了。

再看國內國產框架，幾年來，

MindSpore和“洋框架”打了不少張場仗，

雖然沒有占過上風，

但忍耐力還不算壞。

這時候，我又想給這個小章節重新擬個標題了，

《北美市場已生變數，國內持舊觀望》

——在不疑處生疑。

越是有人毫不懷疑地用PyTorch，

越要重新審視。

03.不吃生態這碗飯？

在我個人的寫作史上，

四年前那篇《搞深度學習框架的那幫人》，

其中大量篇幅在聊框架生態。

然而，時過境已遷，

身處大模型時代，

如果拿生態來衡量JAX，

實話實說，JAX也沒啥生態，

它的論文數量還沒有MindSpore多。

（此處不禮貌地笑了）

堂堂谷歌，堂堂華為，

框架生態都矮“一哥”一頭。

接著靈魂發問：

為啥JAX在大模型巨頭玩家中的份額這么高呢？

因為大模型框架對原來的那種生態，

可能要求沒那么高，

反倒對性能調試調優的效率要求高。

也就是，不吃生態這碗飯，而吃效率這碗飯。

往細里講，訓練時你寫個模型腳本不重要。

關鍵是你能把這個系統的性能，

調優的效率，精度的對齊，打出水平。

模型腳本是什么呢？

它是告訴框架“這個模型應該長什么樣”的說明書。

而在大模型時代，系統性能和訓練效率成了重中之重。框架需要提供強大的分布式訓練支持，穩定和性能壓倒一切。巨頭玩家搞大模型的時候，

對框架的要求和以前相比出現了巨大變化。

平靜的大海培養不出優秀的水手。

回到此前結論，重要的事情說三遍：

變量太猛，戰局有變。

04.為啥谷歌JAX框架這么行？

眾所周知，TensorFlow在大規模生產部署中，

仍具有技術優勢，市場份額也有。

尤其是其與TPU芯片深度整合，

繼續占據一席之地。

不過同為谷歌系的框架，

TensorFlow框架有很多歷史遺留問題，

而JAX框架低調地風生水起。

第一點，谷歌JAX框架設計簡潔，

不像TensorFlow框架設計復雜。

第二點，谷歌JAX一開始就是靜態圖，

且擅長靜態圖。

靜態圖天然適合科學計算和大模型的優化需求，

并延伸出自動并行。

同時，靜態圖也是自動并行的基礎，

為其提供了全局視圖和優化空間。

于是，水到渠成第三點，

JAX的自動并行非常能打，

是在靜態圖打下的好基礎上做到的，

有一套強大的自動并行工具。

對比PyTorch框架，早期沒重視并行，

以至于給了第三方背后開槍的機會。

（姑且按下，容后細述）

細看谷歌JAX，再看一眼昇思MindSpore，

氣質上很像，想法也總一樣，

它們都從一開始就考慮在框架上支持并行。

我甚至愿意下這樣一個結論，

當年，谷歌JAX和昇思MindSpore，

對并行考慮得多，

就是預測模型會變大。

而今看來，一切成真。

當然，還有一點非常重要，

是一窺當下框架競爭態勢的要點：Numpy。

Numpy是個好東西，甚至是越看越好。

你可以記不住它，

也可以永遠用不到它，

但請別忽略它的價值，

它的價值不是能用鈔票衡量的，

不僅在人類科學工具發展歷史上占有一席之地，

而且什么都不爭，克制低調，

實際卻是整個江湖的發源地，

甚至間接推動了AI框架的崛起。

在AI框架出現之前，

它是唯一的選擇，

Numpy為矩陣計算提供了最初的接口和工具。

然而，它不是PyTorch的選擇，

你看，JAX選了Numpy，

蘋果也選了Numpy，這可不是巧合。

為什么？

Numpy中立，且兼容性強，

甚至可以說它是一種科學計算領域的公共語言。

JAX選Numpy，

蘋果選Numpy，

再反觀PyTorch，雖已一統江湖，

但在設計上與Numpy漸行漸遠。

當谷歌JAX和蘋果這兩家代表未來趨勢的選手拉起Numpy站隊時，

PyTorch反而被推向了一個“孤立”的處境。

更關鍵的是，Numpy的中立讓它成為“公共資源”，誰都能拿來用。

但當你的對手都用這個公共資源來增強自己，

而你卻在構建別的，

這江湖里的對抗就不只是暗流涌動。

我們能看到，谷歌JAX和蘋果的選擇，

其實是對PyTorch的一個有力提醒。

還沒聊完。

還要深入靜態圖和動態圖的江湖。

不過且慢，我在硅谷遇見了一件小事，

于是，在話題出發前，

我們繞道去下這里“API”。

05誰敢小瞧框架API？

框架的API是啥？

API可翻譯成“接口”。

不過此接口，非彼接口，

再講那件小事。

某天，約見了一個程序員，

雖然他不搞AI，

但正巧我在寫這篇稿子，

聊一句框架的API。

他信心滿滿地跟我說：

“API不就是調用一個服務嘛，多老的東西了。”

而他說這話的表情吧，

就像揭露了你炒作技術概念。

當時，我真是……

所以，大家不要對美國硅谷全是濾鏡。

話說回來，

框架API它不再適用于之前那么單薄的理解了，

那是整層技術棧的總指揮。

縱觀歷史，唏噓不已，

框架每一輪競爭都離不開“接口的博弈”。

可不是讓“程序和程序對話”，這么簡單，

而是，從“開發工具”變成了，

“開發體驗”和“硬件性能”的代名詞。

誰的接口設計得好，誰首先就能抓住開發者。

框架接口的好壞，直接決定了一個框架的受歡迎程度。比如，屎堆型界面，誰也受不了。

早期TensorFlow，也因此跌倒。

因為太過復雜，被PyTorch搶了很多市場份額，

所以，你說，框架接口重不重要呢？

以前，調用一個服務。

現在，表面只是調用一個服務，

背后調度幾十臺GPU拼命算矩陣乘法。

甚至，框架的接口規模很大，

一個完整的生態系統可能需要上千個接口。

那么問題來了，

PyTorch有多少個接口？

答案是：一千多個。

說到這里，如果MindSpore和其全部對齊呢？

很困難，工作量很大。

所以，我觀察，華為先做了這樣一件事情：

大模型相關的大概就三四百個接口。

這也是一個不小的工作量。

先將大模型相關的框架接口對齊。

想全面對齊，就這個框架接口的數量級，

若不是體量大的公司，大舉揮師，

很難拿下這樣工程量的山頭。

另外，在大模型時代，

這些“流水線”已經復雜得很不像話了：

比如，大模型訓練需要把任務，

分布到幾十上百塊GPU上，

每塊GPU干啥，框架接口得安排得明明白白。

沒有這些優化，硬件就像一群分工混亂的工人，

干活慢，還亂用資源。

再比如，

框架接口如果調用得不好，

性能立馬掉一截。

不僅是個接口，

而是通過這些接口將硬件性能最大化釋放。

寫到這，大家也明白了。

有些技術追趕，

需要用大批量的、高質量的工程追趕。

沒有那么多彎道給你超車，

他們建設得早，路是他們的，車也是他們的。

建我們的道，

造我們的車。

超車需要的是長期耐心和投入。

框架接口的戰場打的是什么？

要我看，誰敢小瞧框架接口？

它打的是基本面。

06靜態圖與動態圖

說框架，繞不過接口，

更繞不過靜態圖和動態圖。

這里一直存在不少歷史遺留問題，

而且是個話題終結者。

一般都不愿意看這么細，這么深，

其實里面大有洞天。

這就像你是一個電影導演：

動態圖是邊拍，邊調劇本。

靜態圖是先定稿劇本，再去成片。

為什么說做框架難？

從表面上看，只需要告訴框架“我想訓練啥模型”，

它就會幫你自動（把代碼、數據、硬件資源組織好），讓模型跑出結果，而不用操心底層的數學運算和硬件適配。

這種“自動化”的背后，

繞不開靜態圖和動態圖。

但從更高的產業視角來看，它們不僅是技術選擇，

PyTorch的江山就是靠動態圖打下來，

門檻低，好用。

不是非得六邊形戰士才能用。

動態圖，好比，做題，做一步，

翻看一次答案，

錯了就是隨時改，

越做越有信心。

或者另一種方式是，

一道題算了三天三夜，才能對答案，改錯題。

后者更適合學霸。

這種玩法，

游戲難度太大，玩家就少了。

于是，誰靜態圖起家，

誰就吃了生態上的大虧。

很難講誰吃的虧更大。

因為谷歌JAX和昇思MindSpore在這件事情上，

它倆想法又一樣。

（此處又不禮貌地笑了）

它們都是在執行的時候盡量變成靜態圖去執行。

靜態圖主打“規劃好一切才準許行動”。

因為框架不僅要理解用戶的代碼，

還要判斷如何用最優的方式在硬件上執行。

而且靜態圖是自動并行的基礎，

自動并行是大模型對框架的核心需求之一。

再細看MindSpore，

到昇騰AI處理器第二代，

全面支持動態圖的時候，相當于有個負擔。

也就是說，昇思MindSpore針對昇騰第二代AI處理器去做設計的時候，

對靜態圖投入了大量精力。

動態圖雖然有，性能不太好，別人也不會用。

談易用性，是在性能還能用的情況下，

并不是在性能不能用的情況下奢談易用性。

也就是說，現在會繼續發展原先靜態圖的優勢，而短板也要補。

不過，一切都需要時間。

聊到這里，理解MindSpore在忙啥了吧？

也理解PyTorch在忙啥了吧？

再總結一下，PyTorch忙活兩件事，

第一要事，拼命搞它的靜態圖。

07誰在掏空PyTorch？

接上文，PyTorch還忙著搞另一件事，

拼命想把并行庫移到“框架里”。

猛一看，什么叫做“框架里”？

說來話長。

作為AI框架的“一哥”，

在設計初期，并沒有直接支持并行計算的能力。

為什么？

因為早期沒必要，現在很有必要。

早期還沒大模型呢，

軟件這件事就是，你不解決，有的是人幫你解決。

沒有原生并行能力的PyTorch，

吸引到Megatron前來。

這個有靠山的兄弟來了，看似局面穩了。

也就是說，當你需要讓模型跑在多張GPU上并行訓練時，PyTorch本身是幫不上太多忙的——這個活兒，靠的是像Megatron這樣的“專業外包選手”。

Megatron是什么？

是英偉達開發的一個庫，用于搞定大模型的并行計算。也可以叫它第三方庫。

當你要訓練大家伙（Transformer）的時候，

Megatron會幫你把模型的參數，

數據和計算任務切分到不同的GPU上，

同時保證結果一致。

不過該來的總該要來，

誰是底層，誰就容易無感。

底層軟件的宿命，就是被用得天經地義，卻無人留意。除非上面一層，仍然是自己人。

常言道，我和你一條船，不代表我和你一條心。

Megatron是英偉達的。

假如，我是說假如，Megatron這種并行庫，

連MindSpore都能對接了，

假如人家MindSpore還做得好，

那么MindSpore+Megatron行不行？

因為開發者只感受Megatron，底下用啥框架都可以。

終于，PyTorch意識到了，

戰略要地，怎么能留給別人插旗？

Megatron看似駐扎在PyTorch的地盤上，

但歸根結底，它還是英偉達的“親兵”。

別看現在在開源生態里打得風生水起，

關鍵時刻，這兵馬終究要聽號令回歸主陣營。

說白了，Megatron扎根在PyTorch，

平時幫著干活，一旦局勢有變，

還是要為自家陣地開路。

于是，PyTorch拼命補課，

把并行能力“嵌入”到框架里，成為原生功能。

這樣，不需要外接Megatron這樣的庫，

開發者也能直接用PyTorch寫出支持多GPU并行的大模型。

為什么并行能力對PyTorch很重要？

對誰都很重要，這是時代的要求，

世人皆知，模型太大，

但是模型大不是結果，而是需求的特點。

另一個變量，無聲登場，
而這個變量，才是真正的大殺器。

假如大模型的結構在逐漸固定，“逐漸”二字重讀，

打破“逐漸”則需要創新力。

若模型結構趨于不變，

開發者不再需要頻繁改動底層框架邏輯，

而是更多關注性能優化和并行計算。

這點對“AI芯片影響更大”，

但因為框架和芯片的“垂直整合”是另一層面的分析，

此處按下不表。

如果PyTorch的并行能力總是靠Megatron，

那么開發者的關注點會轉移到Megatron上。

久而久之，PyTorch的存在感會被削弱。

乍一看，AI框架的江湖，
PyTorch似乎已是“一代霸主”，
仿佛塵埃落定。
但事實遠非如此。
欲戴王冠，必承其重，
暗潮涌動，未必穩固。

你再看JAX框架，

技術路線是并行和編譯這些能力都放在框架本層里，且是靜態圖加自動并行。

巧了，谷歌JAX的思路又跟昇思MindSpore一樣。

剛才說了，PyTorch框架當下最重要的工作，

其一，做靜態圖（編譯），

由此可見，其二是把并行做回框架本層來。

誠如所見，

PyTorch新版本特性集中反映在兩個聚焦點上。

08開發者隨上層“遷移”

細想想，HuggingFace的崛起，

代表了AI開發從“框架中心”向“模型中心”遷移。

要我說，這種變化間接“架空”了PyTorch在開發者中的直接影響力。

HuggingFace的核心價值，

是把復雜的AI模型封裝成易用的工具庫和API，

開發者無需深入底層，

只需調用HuggingFace模型接口，就能干活。

PyTorch雖然強大，

但直接在上面開發頗有門檻，

不是說上手就會。

當HuggingFace的模型庫和工具，

成為開發者的入口，

PyTorch則淪為底層執行引擎，

成了看不見的“發動機”，

HuggingFace成了摸得到的“駕駛艙”。

這意味著，PyTorch將生態主導權部分讓出。

說實話，我怎么就感覺PyTorch越來越“薄”了呢？

你們說呢？

06其他沒機會了？

都說千秋霸業，

霸業再盛，若不隨勢而動，

也難保江湖恒久敬畏。

生態位的競爭變了，

開發者的關注點就轉移了。

PyTorch確實不綁定特定硬件，

但是，AI框架跟硬件強相關。

但芯片廠商為了吸引開發者，開發了特定的異構計算架構。比如，英偉達的CUDA，AMD的ROCm，甚至華為的CANN。

而且，事物有兩面性。

如果框架缺乏硬件支持，

其使用體驗和性能都會受到影響。

框架的廣泛應用離不開對不同硬件的適配，

而適配程度直接影響框架的競爭力。

PyTorch沒有硬件，而硬件架構演進很快，

比如，超節點，

UMA（統一內存架構），

Dataflow（美國AI芯片公司Sambanova架構）等等，框架如何快速適配？

多說一句，尤其在中國，

有很多國產芯片，正在備嘗艱苦。

假如模型結構固定的話，

玩幾個“經典款”就夠了。

Llama或千問很有潛力。

我目之所及，無論中國的，還是美國的，

很多大甲方的基礎模型就選這兩個。

只要你能給出最優解，我就選你。

只要AI框架自帶一個優秀的Llama組件，

只要這個標準組件的性能比PyTorch原生的Llama提升15%。大家都很“現實”，當“標準模型”成熟，開發者根本不關心它是用什么框架寫的。他們只想知道：哪種框架能讓我用的模型跑得更快？

若模型結構固定，

這個AI框架（軟件）和AI芯片（硬件）做垂直整合，就有機會做出新優勢。

因為硬件廠商了解自家芯片的每一個細節，

框架可以針對芯片特性做極致優化。

而PyTorch這種中立框架需要支持各種硬件，

優化時得考慮兼容性。

它也不可能為某一個硬件優化。

而JAX背后有谷歌TPU芯片，

MindSpore背后有昇騰AI處理器。

AI框架的未來，并不是單純拼生態廣度。

是誰能在“模型固定后”的賽道上，

跑得更快、更遠、更穩。

當下，模型是處于收斂和不收斂的中間狀態，

不是說模型結構不變，而是它還在演進，

中間狀態，多方戰況焦灼。

不得不說，蘋果的AI框架MLX，

Github上看，熱度漲得很快，不要小瞧。

蘋果MLX框架可面向個人開發者，或是小企業，

群體很大，抓住了，又有蘋果的硬件和全系統的支持，相對小（1TB）的模型就可在上面訓練，做微調。蘋果這個例子說明了什么呢？

“垂直整合”。

（AI大模型垂直整合：我有一條龍服務，我就是那條龍）

不過，蘋果MLX框架+服務器芯片，

才開始搞，還要再觀察，而且，蘋果在端上的優勢，是另一個故事了《AI推理篇》。將蘋果的MLX框架集成于M系列芯片中，猛搞終端推理優化。

這種策略雖然不搶占大模型訓練市場，但在iPhone、iPad上構筑強大封閉生態。放眼望去，這條路也能走出大生態。

同理也適用于谷歌JAX和昇思MindSpore。

至于PyTorch，最近有點煩。

（完）

One More Thing

2024年2月，我已有一個以PyTorch被“掏空”為選題的寫作計劃。

不過，因為寫作難度大，訪談難度高，久未成稿。

“新聞不是一天發生”，沒想到僅過一年，

局面基本可以確定如此，喟然一嘆，

這又是一個量變引起質變的故事，

世間萬事，莫過如此。