機(jī)器之心報(bào)道

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

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

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

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

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

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

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

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

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

論文標(biāo)題：AXLearn:ModularLargeModelTrainingonHeterogeneousInfrastructure

AXLearn介紹

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

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

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

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

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

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

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

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

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

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

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

實(shí)驗(yàn)評估

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

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

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

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

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

本文在三種硬件平臺(tái)上評估了兩個(gè)模型：Llama27B與Llama270B：

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

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

3.1024顆Trainium2芯片（64個(gè)AWStrn2實(shí)例，每個(gè)實(shí)例含16顆Trainium2芯片）。

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

為驗(yàn)證AXLearn的可擴(kuò)展性，本文對兩個(gè)實(shí)際部署的模型進(jìn)行了弱擴(kuò)展性（weak-scaling）實(shí)驗(yàn)。

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

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

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

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

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

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

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

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

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