機器之心報道

編輯：笑寒、陳陳

這或許是龐若鳴（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 訓練的部分模型，現已被廣泛應用于十億級用戶規模的產品功能中，包括：智能助手、多模態理解與生成、代碼智能等關鍵能力。

了解更多內容，請參考原論文。