機器之心報道

這或許是龐若鳴（RuomingPang）在蘋果參與的最后一篇論文。

龐若鳴，蘋果基礎(chǔ)模型團隊負責人、杰出工程師，即將成為Meta新成立的超級智能團隊的最新成員。他本科畢業(yè)于上海交通大學，在谷歌工作了15年，此后加入蘋果。另據(jù)彭博社最新消息，Meta更是開出了2億美金的天價來邀請龐若鳴加入。

雖然即將跨入另一段人生旅程，但龐若鳴還在為蘋果站好最后一班崗。

7月9日，龐若鳴在X上宣傳了自己參與的一項研究《AXLearn:ModularLargeModelTrainingonHeterogeneousInfrastructure》，據(jù)了解，這項研究是構(gòu)建AppleFoundation模型的基礎(chǔ)代碼庫。

具體而言，本文設(shè)計并實現(xiàn)了AXLearn，一個用于大規(guī)模深度學習模型訓練的生產(chǎn)級系統(tǒng)，其具備良好的可擴展性和高性能。與其他先進的深度學習系統(tǒng)相比，AXLearn具有獨特的優(yōu)勢：高度模塊化和對異構(gòu)硬件基礎(chǔ)設(shè)施的全面支持。

AXLearn內(nèi)部的軟件組件接口遵循嚴格的封裝原則，使得不同組件能夠靈活組合，從而在異構(gòu)計算環(huán)境中快速進行模型開發(fā)和實驗。

此外，本文還提出了一種用于衡量模塊化程度的新方法：基于代碼行數(shù)的復雜度（LoC-complexity）指標。實驗表明，AXLearn在系統(tǒng)擴展時可以保持恒定的復雜度，而其他系統(tǒng)則呈現(xiàn)出線性甚至二次增長的復雜度。

例如，將RotaryPositionEmbeddings（RoPE）這類功能集成到AXLearn的上百個模塊中僅需約10行代碼，而在其他系統(tǒng)中可能需要數(shù)百行代碼才能實現(xiàn)相同效果。同時，AXLearn也保持了與主流高性能訓練系統(tǒng)相當?shù)挠柧毿阅堋?/p>

開源地址：https://github.com/apple/axlearn

論文標題：AXLearn:ModularLargeModelTrainingonHeterogeneousInfrastructure

AXLearn介紹

現(xiàn)階段，像ChatGPT、Gemini這樣的聊天機器人都是由大模型驅(qū)動的。這種深度學習系統(tǒng)會優(yōu)先考慮性能和可擴展性。

作為全球最大的消費電子和在線服務(wù)公司之一，蘋果已經(jīng)將許多AI模型集成到自家產(chǎn)品中，服務(wù)于全球數(shù)十億用戶。

除了訓練性能和可擴展性外，蘋果對深度學習系統(tǒng)還有兩個額外的要求。首先是賦能模型工程師，只需編寫最少的代碼，就能配置復雜的模型定義和訓練方法。其次，作為一家大型科技公司，他們不能依賴單一的硬件供應(yīng)商，因而他們的設(shè)計目標是兼容異構(gòu)后端，如GPU、TPU和AWSTrainium。

為了達到上述目的，AXLearn被開發(fā)出來。

為了促進模塊化，AXLearn的核心設(shè)計決策是強制執(zhí)行嚴格的封裝。此外，本文還通過將旋轉(zhuǎn)位置嵌入（RoPE）和專家混合模型（MoE）集成到AXLearn中的案例研究，展示了該框架與傳統(tǒng)代碼行數(shù)計數(shù)方法的一致性。

圖2顯示了AXLearn的系統(tǒng)架構(gòu)和工作流程。AXLearn有兩個關(guān)鍵組件：

（1）AXLearn組合器（AXLearncomposer）和（2）AXLearn執(zhí)行框架（AXLearnruntime）。

用戶通常使用AXLearn內(nèi)置的層庫和第三方層組件來定義訓練配置?；谠撆渲媚_本，AXLearn組合器會首先生成完整的JAX程序。

這一過程包含以下關(guān)鍵步驟：包括為目標加速器實例選擇合適的網(wǎng)格形狀、為特定層應(yīng)用分片注釋、為目標硬件自動調(diào)優(yōu)XLA編譯選項、為后端選擇合適的注意力內(nèi)核，并根據(jù)模塊層次中的標記點應(yīng)用適當?shù)闹赜嬎悴呗浴＿@些注釋對于訓練的高效運行至關(guān)重要。

然后，JAX程序和編譯選項被傳遞給XLA編譯器，以生成加速器程序（例如，CUDA內(nèi)核），該程序隨后通過AXLearn運行時在分布式硬件（例如Kubernetes）上進行調(diào)度，并使用特定于加速器的運行時（例如CUDA運行時）。

AXLearn執(zhí)行框架監(jiān)控加速器程序的執(zhí)行，并提供額外的功能，如高效的檢查點、監(jiān)控和容錯能力。

實驗評估

下表展示了不同系統(tǒng)的代碼量復雜度（LoC-Complexities）匯總。

在AXLearn中，RoPE和MoE被嚴格封裝。本文提供了一個10行的代碼片段，可以將這兩個功能集成到任何實驗配置中。

在本文的內(nèi)部實踐中，正是通過類似的代碼片段，成功配置了超過1000個實驗，用于啟用RoPE、MoE，或兩者同時使用。隨著模塊數(shù)量或RoPE或MoE變體的增加，無需對任何現(xiàn)有接口進行更改，實現(xiàn)了恒定的代碼復雜性。

在異構(gòu)硬件上的性能

本文將AXLearn的訓練性能與以下系統(tǒng)進行了對比：PyTorchFSDP、Megatron-LM以及MaxText，這些系統(tǒng)在GPU與TPU上均實現(xiàn)了先進的訓練性能。

本文在三種硬件平臺上評估了兩個模型：Llama27B與Llama270B：

1.256/512H100GPU（分別對應(yīng)32/64個AWSP5d實例，每個實例含8張H100）；

2.TPU-v5p-512/1024（分別對應(yīng)64/128個GCPCloudTPU主機，每個主機含4顆芯片）；

3.1024顆Trainium2芯片（64個AWStrn2實例，每個實例含16顆Trainium2芯片）。

下表總結(jié)了性能結(jié)果。

為驗證AXLearn的可擴展性，本文對兩個實際部署的模型進行了弱擴展性（weak-scaling）實驗。

這些結(jié)果表明，AXLearn接近線性擴展性，如圖4所示。

AXLearn在TPU上同樣展現(xiàn)出業(yè)界領(lǐng)先的推理性能。本文對比了AXLearn與vLLM在Llama27B與70B參數(shù)模型上的推理表現(xiàn)。

如表4和圖5所示，AXLearn在延遲和吞吐量方面均顯著優(yōu)于vLLM：

在延遲方面，AXLearn在TTFT和TPOT上分別實現(xiàn)了500倍和6倍加速；

在吞吐量方面，AXLearn在7B模型推理上快2.8倍，在70B模型上快1.6倍。

在生產(chǎn)中的使用體驗

如今，AXLearn已從最初僅有數(shù)位開發(fā)者、訓練百萬級參數(shù)模型的工具，發(fā)展為支持數(shù)百位開發(fā)者訓練十億至萬億參數(shù)規(guī)模模型的大型平臺。

它在任意時刻可并行支持超過10,000個實驗的開發(fā)與運行，并部署在數(shù)十種異構(gòu)硬件集群上。

借助AXLearn訓練的部分模型，現(xiàn)已被廣泛應(yīng)用于十億級用戶規(guī)模的產(chǎn)品功能中，包括：智能助手、多模態(tài)理解與生成、代碼智能等關(guān)鍵能力。