M：eta为他豪掷2亿美元 上交校友、庞若鸣 “晒出在苹”果的最新论文

机器之心报道

这或许是庞若鸣（RuomingPang）在苹果参与的最后一篇论文。

庞若鸣，苹果基础模型团队负责人、杰出工程师，即将成为Meta新成立的超级智能团队的最新成员。他本科毕业于上海交通大学，在谷歌工作了15年，此后加入苹果。另据彭博社最新消息，Meta更是开出了2亿美金的天价来邀请庞若鸣加入。

虽然即将跨入另一段人生旅程，但庞若鸣还在为苹果站好最后一班岗。

7月9日，庞若鸣在X上宣传了自己参与的一项研究《AXLearn:ModularLargeModelTrainingonHeterogeneousInfrastructure》，据了解，这项研究是构建AppleFoundation模型的基础代码库。

具体而言，本文设计并实现了AXLearn，一个用于大规模深度学习模型训练的生产级系统，其具备良好的可扩展性和高性能。与其他先进的深度学习系统相比，AXLearn具有独特的优势：高度模块化和对异构硬件基础设施的全面支持。

AXLearn内部的软件组件接口遵循严格的封装原则，使得不同组件能够灵活组合，从而在异构计算环境中快速进行模型开发和实验。

此外，本文还提出了一种用于衡量模块化程度的新方法：基于代码行数的复杂度（LoC-complexity）指标。实验表明，AXLearn在系统扩展时可以保持恒定的复杂度，而其他系统则呈现出线性甚至二次增长的复杂度。

例如，将RotaryPositionEmbeddings（RoPE）这类功能集成到AXLearn的上百个模块中仅需约10行代码，而在其他系统中可能需要数百行代码才能实现相同效果。同时，AXLearn也保持了与主流高性能训练系统相当的训练性能。

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

论文标题：AXLearn:ModularLargeModelTrainingonHeterogeneousInfrastructure

AXLearn介绍

现阶段，像ChatGPT、Gemini这样的聊天机器人都是由大模型驱动的。这种深度学习系统会优先考虑性能和可扩展性。

作为全球最大的消费电子和在线服务公司之一，苹果已经将许多AI模型集成到自家产品中，服务于全球数十亿用户。

除了训练性能和可扩展性外，苹果对深度学习系统还有两个额外的要求。首先是赋能模型工程师，只需编写最少的代码，就能配置复杂的模型定义和训练方法。其次，作为一家大型科技公司，他们不能依赖单一的硬件供应商，因而他们的设计目标是兼容异构后端，如GPU、TPU和AWSTrainium。

为了达到上述目的，AXLearn被开发出来。

为了促进模块化，AXLearn的核心设计决策是强制执行严格的封装。此外，本文还通过将旋转位置嵌入（RoPE）和专家混合模型（MoE）集成到AXLearn中的案例研究，展示了该框架与传统代码行数计数方法的一致性。

图2显示了AXLearn的系统架构和工作流程。AXLearn有两个关键组件：

（1）AXLearn组合器（AXLearncomposer）和（2）AXLearn执行框架（AXLearnruntime）。

用户通常使用AXLearn内置的层库和第三方层组件来定义训练配置。基于该配置脚本，AXLearn组合器会首先生成完整的JAX程序。

这一过程包含以下关键步骤：包括为目标加速器实例选择合适的网格形状、为特定层应用分片注释、为目标硬件自动调优XLA编译选项、为后端选择合适的注意力内核，并根据模块层次中的标记点应用适当的重计算策略。这些注释对于训练的高效运行至关重要。

然后，JAX程序和编译选项被传递给XLA编译器，以生成加速器程序（例如，CUDA内核），该程序随后通过AXLearn运行时在分布式硬件（例如Kubernetes）上进行调度，并使用特定于加速器的运行时（例如CUDA运行时）。

AXLearn执行框架监控加速器程序的执行，并提供额外的功能，如高效的检查点、监控和容错能力。

实验评估

下表展示了不同系统的代码量复杂度（LoC-Complexities）汇总。

在AXLearn中，RoPE和MoE被严格封装。本文提供了一个10行的代码片段，可以将这两个功能集成到任何实验配置中。

在本文的内部实践中，正是通过类似的代码片段，成功配置了超过1000个实验，用于启用RoPE、MoE，或两者同时使用。随着模块数量或RoPE或MoE变体的增加，无需对任何现有接口进行更改，实现了恒定的代码复杂性。

在异构硬件上的性能

本文将AXLearn的训练性能与以下系统进行了对比：PyTorchFSDP、Megatron-LM以及MaxText，这些系统在GPU与TPU上均实现了先进的训练性能。

本文在三种硬件平台上评估了两个模型：Llama27B与Llama270B：

1.256/512H100GPU（分别对应32/64个AWSP5d实例，每个实例含8张H100）；

2.TPU-v5p-512/1024（分别对应64/128个GCPCloudTPU主机，每个主机含4颗芯片）；

3.1024颗Trainium2芯片（64个AWStrn2实例，每个实例含16颗Trainium2芯片）。

下表总结了性能结果。

为验证AXLearn的可扩展性，本文对两个实际部署的模型进行了弱扩展性（weak-scaling）实验。

这些结果表明，AXLearn接近线性扩展性，如图4所示。

AXLearn在TPU上同样展现出业界领先的推理性能。本文对比了AXLearn与vLLM在Llama27B与70B参数模型上的推理表现。

如表4和图5所示，AXLearn在延迟和吞吐量方面均显著优于vLLM：

在延迟方面，AXLearn在TTFT和TPOT上分别实现了500倍和6倍加速；

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

在生产中的使用体验

如今，AXLearn已从最初仅有数位开发者、训练百万级参数模型的工具，发展为支持数百位开发者训练十亿至万亿参数规模模型的大型平台。

它在任意时刻可并行支持超过10,000个实验的开发与运行，并部署在数十种异构硬件集群上。

借助AXLearn训练的部分模型，现已被广泛应用于十亿级用户规模的产品功能中，包括：智能助手、多模态理解与生成、代码智能等关键能力。