機器之心報道
編輯:笑寒、陳陳
這或許是龐若鳴(Ruoming Pang)在蘋果參與的最后一篇論文。
龐若鳴,蘋果基礎模型團隊負責人、杰出工程師,即將成為 Meta 新成立的超級智能團隊的最新成員。他本科畢業于上海交通大學,在谷歌工作了 15 年,此后加入蘋果。另據彭博社最新消息,Meta 更是開出了 2 億美金的天價來邀請龐若鳴加入。
雖然即將跨入另一段人生旅程,但龐若鳴還在為蘋果站好最后一班崗。
7 月 9 日,龐若鳴在 X 上宣傳了自己參與的一項研究《 AXLearn: Modular Large Model Training on Heterogeneous Infrastructure 》,據了解,這項研究是構建 Apple Foundation 模型的基礎代碼庫。
具體而言,本文設計并實現了AXLearn,一個用于大規模深度學習模型訓練的生產級系統,其具備良好的可擴展性和高性能。與其他先進的深度學習系統相比,AXLearn 具有獨特的優勢:高度模塊化和對異構硬件基礎設施的全面支持。
AXLearn 內部的軟件組件接口遵循嚴格的封裝原則,使得不同組件能夠靈活組合,從而在異構計算環境中快速進行模型開發和實驗。
此外,本文還提出了一種用于衡量模塊化程度的新方法:基于代碼行數的復雜度(LoC-complexity)指標。實驗表明,AXLearn 在系統擴展時可以保持恒定的復雜度,而其他系統則呈現出線性甚至二次增長的復雜度。
例如,將 Rotary Position Embeddings(RoPE)這類功能集成到 AXLearn 的上百個模塊中僅需約 10 行代碼,而在其他系統中可能需要數百行代碼才能實現相同效果。同時,AXLearn 也保持了與主流高性能訓練系統相當的訓練性能。
- 論文地址:https://arxiv.org/pdf/2507.05411
- 開源地址:https://github.com/apple/axlearn
- 論文標題: AXLearn: Modular Large Model Training on Heterogeneous Infrastructure
AXLearn 介紹
現階段,像 ChatGPT、Gemini 這樣的聊天機器人都是由大模型驅動的。這種深度學習系統會優先考慮性能和可擴展性。
作為全球最大的消費電子和在線服務公司之一,蘋果已經將許多 AI 模型集成到自家產品中,服務于全球數十億用戶。
除了訓練性能和可擴展性外,蘋果對深度學習系統還有兩個額外的要求。首先是賦能模型工程師,只需編寫最少的代碼,就能配置復雜的模型定義和訓練方法。其次,作為一家大型科技公司,他們不能依賴單一的硬件供應商,因而他們的設計目標是兼容異構后端,如 GPU、TPU 和 AWS Trainium。
為了達到上述目的,AXLearn 被開發出來。
為了促進模塊化,AXLearn 的核心設計決策是強制執行嚴格的封裝。 此外,本文還通過將旋轉位置嵌入(RoPE)和專家混合模型(MoE)集成到 AXLearn 中的案例研究,展示了該框架與傳統代碼行數計數方法的一致性。
圖 2 顯示了 AXLearn 的系統架構和工作流程。AXLearn 有兩個關鍵組件:
(1)AXLearn 組合器(AXLearn composer)和(2)AXLearn 執行框架(AXLearn runtime)。
用戶通常使用 AXLearn 內置的層庫和第三方層組件來定義訓練配置。基于該配置腳本,AXLearn 組合器會首先生成完整的 JAX 程序。
這一過程包含以下關鍵步驟: 包括為目標加速器實例選擇合適的網格形狀、為特定層應用分片注釋、為目標硬件自動調優 XLA 編譯選項、為后端選擇合適的注意力內核,并根據模塊層次中的標記點應用適當的重計算策略。這些注釋對于訓練的高效運行至關重要。
然后,JAX 程序和編譯選項被傳遞給 XLA 編譯器,以生成加速器程序(例如,CUDA 內核),該程序隨后通過 AXLearn 運行時在分布式硬件(例如 Kubernetes)上進行調度,并使用特定于加速器的運行時(例如 CUDA 運行時)。
AXLearn 執行框架監控加速器程序的執行,并提供額外的功能,如高效的檢查點、監控和容錯能力。
實驗評估
下表展示了不同系統的代碼量復雜度(LoC-Complexities)匯總。
在 AXLearn 中,RoPE 和 MoE 被嚴格封裝。本文提供了一個 10 行的代碼片段,可以將這兩個功能集成到任何實驗配置中。
在本文的內部實踐中,正是通過類似的代碼片段,成功配置了超過 1000 個實驗,用于啟用 RoPE、MoE,或兩者同時使用。隨著模塊數量或 RoPE 或 MoE 變體的增加,無需對任何現有接口進行更改,實現了恒定的代碼復雜性。
在異構硬件上的性能
本文將 AXLearn 的訓練性能與以下系統進行了對比:PyTorch FSDP、Megatron-LM 以及 MaxText,這些系統在 GPU 與 TPU 上均實現了先進的訓練性能。
本文在三種硬件平臺上評估了兩個模型:Llama2 7B 與 Llama2 70B:
1. 256/512 H100 GPU(分別對應 32/64 個 AWS P5d 實例,每個實例含 8 張 H100);
2. TPU-v5p-512/1024(分別對應 64/128 個 GCP Cloud TPU 主機,每個主機含 4 顆芯片);
3. 1024 顆 Trainium2 芯片(64 個 AWS trn2 實例,每個實例含 16 顆 Trainium2 芯片)。
下表總結了性能結果。
為驗證 AXLearn 的可擴展性,本文對兩個實際部署的模型進行了弱擴展性(weak-scaling)實驗。
這些結果表明,AXLearn 接近線性擴展性,如圖 4 所示。
AXLearn 在 TPU 上同樣展現出業界領先的推理性能。本文對比了 AXLearn 與 vLLM 在 Llama2 7B 與 70B 參數模型上的推理表現。
如表 4 和圖 5 所示,AXLearn 在延遲和吞吐量方面均顯著優于 vLLM:
- 在延遲方面,AXLearn 在 TTFT 和 TPOT 上分別實現了 500 倍和 6 倍加速;
- 在吞吐量方面,AXLearn 在 7B 模型推理上快 2.8 倍,在 70B 模型上快 1.6 倍。
在生產中的使用體驗
如今,AXLearn 已從最初僅有數位開發者、訓練百萬級參數模型的工具,發展為支持數百位開發者訓練十億至萬億參數規模模型的大型平臺。
它在任意時刻可并行支持超過 10,000 個實驗的開發與運行,并部署在數十種異構硬件集群上。
借助 AXLearn 訓練的部分模型,現已被廣泛應用于十億級用戶規模的產品功能中,包括:智能助手、多模態理解與生成、代碼智能等關鍵能力。
了解更多內容,請參考原論文。
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.