编辑：笑寒、陈陈
这或许是庞若鸣（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 训练的部分模型，现已被广泛应用于十亿级用户规模的产品功能中，包括：智能助手、多模态理解与生成、代码智能等关键能力。
了解更多内容，请参考原论文。
©
（文：机器之心）