这很可能是庞若鸣在苹果公司所参与的论文中的最后一篇。

庞若鸣，这位苹果公司基础模型团队的负责人，同时也是一位杰出的工程师，即将加入Meta新组建的超级智能团队。他曾在上海交通大学完成本科学业，随后在谷歌度过了15年的职业生涯，并在那之后加入了苹果公司。根据彭博社的最新报道，Meta为了邀请庞若鸣加盟，甚至开出了高达2亿美金的丰厚薪酬。

尽管即将踏入人生的新篇章，庞若鸣依然坚守在苹果的岗位上，为最后一段旅程画上圆满的句号。

7月9日，庞若鸣在平台X上推广了一项他参与的研究成果——《AXLearn：基于异构基础设施的模块化大型模型训练》，据悉，该研究构成了Apple Foundation模型的基础代码库。

本文设计并成功构建了AXLearn，这是一套适用于大规模深度学习模型训练的生产级系统，该系统不仅拥有出色的可扩展性，而且性能卓越。相较于其他前沿的深度学习系统，AXLearn展现出其独特的优势，主要体现在其高度模块化的设计以及对异构硬件基础设施的全面兼容性。

AXLearn 的软件组件接口严格执行封装规范，确保了各组件间能够实现灵活的搭配，进而使得在异构计算场景中，模型开发与实验能够高效进行。

此外，本研究还引入了一种全新的评估模块化水平的技术：即依据代码行数来测定的复杂度（LoC-complexity）标准。实验结果显示，AXLearn 在系统功能扩展过程中，其复杂度能够保持稳定不变，相比之下，其他系统在扩展过程中复杂度则表现为线性甚至二次方的增长趋势。

将类似RoPE（旋转位置嵌入）这样的功能融入AXLearn的众多模块中，仅需大约10行代码即可完成，而在其他系统里，实现同样的功能可能需要数百行代码。与此同时，AXLearn的训练性能与主流的高性能训练系统相媲美。

AXLearn 介绍

目前，诸如GPT、Gemini等聊天机器人均依赖于大型模型进行驱动。这些深度学习系统在设计和开发过程中，首先注重的是其性能表现和扩展能力。

苹果公司，作为全球领先的消费电子产品及在线服务提供商之一，已将众多人工智能模型融入自家的产品线，为全球数十亿用户提供服务。

除了对训练效能和系统扩展性提出要求，苹果公司对深度学习系统还设定了两个特别条件。首先，他们期望模型工程师能够通过编写少量代码，轻松完成复杂模型定义和训练流程的配置。其次，鉴于其作为一家规模庞大的科技公司，苹果不希望过度依赖单一硬件供应商，因此他们的设计宗旨是确保系统能够与多种异构后端相兼容，包括GPU、TPU以及AWS Trainium。

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

为了推动模块化发展，AXLearn 的核心设计理念是实施严格的封装规定。此外，通过在 AXLearn 中融入旋转位置嵌入（RoPE）技术与专家混合模型（MoE），本文通过案例研究证明了该框架与传统的代码行数统计方法保持了高度的一致性。

图 2 展示了 AXLearn 的系统构造及其运作流程。AXLearn 包含两个核心部分：

AXLearn 的组合器（AXLearn composer）以及AXLearn 的执行框架（AXLearn runtime）。

用户一般会利用AXLearn自带的层级库以及外部的层级组件来设定训练的参数。依托于这个配置脚本，AXLearn的整合器将率先构建出完整的JAX程序。

此流程涉及以下核心环节：需挑选适合目标加速器实例的网格形态、对特定层级实施分片标记、自动调整 XLA 编译器的相关选项以适应目标硬件、为后端系统挑选恰当的注意力处理核心，以及依据模块层级中的关键点实施相应的重计算措施。这些标记对训练效率的提升具有至关重要的作用。

随后，JAX程序及编译参数被提交给XLA编译器，用以构建加速器程序，如CUDA内核。该程序接着在AXLearn运行时驱动下，于分布式硬件平台（如Kubernetes）中执行调度任务，并借助针对加速器定制的运行时环境（如CUDA运行时）进行操作。

AXLearn框架负责对加速器程序的运行进行监控，同时赋予其额外的性能，包括实施高效的检查点机制、执行监控以及具备容错功能。

实验评估

本表格详细呈现了各个系统代码复杂度的总体数据，具体表现为代码行数复杂度（LoC-Complexities）。

在 AXLearn 平台中，RoPE 与 MoE 的应用被严密隔离。本文提供了一段仅有十行的代码示例，该示例能够将这两种功能无缝融入任何实验的配置之中。

在本文的研究过程中，我们正是借助此类代码片段，成功搭建了超过一千项实验，旨在激活RoPE、MoE功能，亦或两者兼而用之。伴随着模块数量的增多，或是RoPE或MoE变种的增加，我们无需对现有的接口进行调整，从而保持了代码的复杂度恒定。

在异构硬件上的性能

本文对比了 AXLearn 的训练效果，与诸如 PyTorch FSDP、Megatron-LM 以及 MaxText 等系统，这些系统在 GPU 和 TPU 上均展现了卓越的训练效能。

本篇文章对两个模型进行了测试，分别在三种不同的硬件设备上运行：一个是Llama2 7B模型，另一个则是Llama2 70B模型。

256/512个H100 GPU分别对应32/64个AWS P5d实例，而每个AWS P5d实例均配备8张H100显卡。

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 所示，AXLearn 在延迟方面表现显著；图 5 显示，其在吞吐量上也明显超越 vLLM。

在生产中的使用体验

如今，AXLearn 已从最初仅有少数几位开发者参与，用于训练百万级参数模型的工具，演变成为一个庞大的平台，能够支持数百位开发者进行训练，并能够处理十亿至万亿参数规模的大型模型。

该系统在任何时间点都能同时处理并支持超过一万项实验的研发与执行，并且已经在多种不同构型的硬件集群中进行部署。

利用AXLearn训练而成的部分模型，目前已广泛融入覆盖十亿用户规模的产品功能之中，诸如：智能客服、多渠道信息解读与创作、代码自动生成等核心功能。