Orion: Characterizing and Programming Apple's Neural Engine for LLM Training and Inference

本文介绍了 Orion，这是首个能够绕过 CoreML 直接利用苹果神经引擎（ANE）进行大语言模型训练与推理的开源端到端系统，它通过揭示并规避 ANE 的 20 项限制、实现权重热更新以将每步训练时间缩短 8.5 倍，并在 M4 Max 设备上成功实现了 GPT-2 的高效推理及 1.1 亿参数模型的稳定训练。

要点

这篇论文介绍了一个名为 Orion 的开源项目，它的目标非常大胆：让苹果设备（iPhone、Mac）里那个被“雪藏”的超级大脑——神经引擎（ANE），真正跑起来，不仅能用来聊天（推理），还能用来学习（训练）人工智能模型。

为了让你更容易理解，我们可以把整个故事想象成在一个拥有顶级厨房（ANE）的餐厅里，如何打破规定，让厨师真正大展身手。

1. 背景：被锁住的“超级厨房”

想象一下，苹果公司的每一款设备（超过 20 亿台）里都藏着一个超级厨房（这就是神经引擎 ANE）。这个厨房配备了 16 个顶级厨师，速度极快，专门处理复杂的烹饪任务（比如生成文字、识别图像）。

• 现状：虽然厨房设备一流，但苹果给它的操作手册（CoreML 框架） 写得像天书一样，而且是个“黑盒”。
• 问题：开发者只能告诉厨房“做顿饭”，但不知道具体怎么切菜、怎么炒菜。更糟糕的是，这个厨房不允许“边做边改食谱”。如果你想让厨师学会新菜（训练模型），按照旧规矩，每学一步，就得把整个厨房拆了重装一遍，重新发号施令。这太慢了，慢到没人愿意用它来学习。

2. Orion 的突破：拿到“后门钥匙”

Orion 团队做了一件大胆的事：他们逆向工程了苹果的内部系统，找到了一把“后门钥匙”（私有 API），直接跳过了那个笨重的操作手册，直接指挥超级厨房。

他们不仅拿到了钥匙，还写了一本全新的“厨房生存指南”，列出了 20 条以前没人知道的“潜规则”（比如：切菜必须按特定顺序、盘子大小必须一致等）。

3. 核心魔法：如何“边学边改”？（Delta Compilation）

这是 Orion 最厉害的地方。

• 以前的做法（全量重编译）：
  想象厨师每学会一个新动作（比如把盐放多一点），你就得把整个厨房清空，重新装修，重新发一遍所有指令。这要花 4.2 秒。如果你要学 1000 步，大部分时间都花在“装修”上，而不是“做饭”上。

• Orion 的做法（Delta Compilation / 增量编译）：
  Orion 发现了一个秘密：虽然厨房规定“食谱”是印在墙上的（编译时固化），但食材（权重数据） 其实是放在地上的袋子里的。
  于是，Orion 发明了一种“手术式更新”：

  1. 让厨师暂停工作（卸载程序）。
  2. 悄悄把地上的食材袋换掉（更新磁盘上的权重文件）。
  3. 让厨师继续工作（重新加载程序）。

  整个过程不需要重新装修厨房，也不需要重新发指令。时间从 4.2 秒缩短到了 0.5 秒！这让训练速度直接提升了 3.8 倍。

4. 其他亮点：像换衣服一样换“技能包”（LoRA）

通常，如果你想让一个厨师学会做川菜，你得重新培训他。但 Orion 引入了 LoRA 技术。

• 比喻：厨师（基础模型）还是那个厨师，但他可以随时换上一件“川菜围裙”（适配器）。
• 效果：你不需要重新培训厨师，只需要在烹饪时递给他不同的围裙（通过输入数据），他就能立刻做川菜。换完围裙，再递给他一件“粤菜围裙”，他又能做粤菜。这实现了热插拔，无需重新编译。

5. 解决“崩溃”问题：让厨房不“发疯”

在之前的尝试中，这个厨房很容易在学到第 2 步时就“发疯”（产生 NaN，即数字错误导致崩溃）。Orion 团队像侦探一样，找到了三个导致发疯的 bug：

1. 旧食谱残留：重启时用了旧的食谱，导致新动作做错了。
2. 数字溢出：厨师太激动，把数字算得太大（超过了浮点数的极限），导致崩溃。Orion 给数字加了“安全带”（截断）。
3. 食材变质：食材袋里的数据格式不对，导致厨师吃到了垃圾。Orion 增加了“质检员”，确保食材干净。

结果：现在，这个厨房可以稳定地连续工作 1000 步，不再发疯，而且没有任何数字错误。

6. 实际表现：快还是慢？

• 推理（聊天）：在 M4 Max 芯片上，Orion 每秒能生成 170 个单词。虽然比直接用 CPU 或 GPU 稍慢一点点（因为换食材袋有点耗时），但它不占用 CPU 和 GPU，而且不耗电（闲置时完全关闭）。
• 训练（学习）：以前需要 85 分钟才能学完 1000 步，现在只需要 22 分钟。

总结

Orion 是什么？
它是一个开源工具包，把苹果设备里那个原本只能“听指挥”的神经引擎，变成了一个可以自主学习和进化的超级大脑。

为什么这很重要？
以前，我们只能在苹果设备上“用”AI，不能“练”AI。Orion 证明了，只要打破软件限制，你的手机和电脑本身就是一台强大的 AI 训练机。这不仅让隐私数据留在本地，还让未来的 AI 应用更加智能和个性化。

简单来说，Orion 就是给苹果设备的“超级厨房”解除了封印，并教它如何高效地“边做边学”。

技术摘要

Orion：面向 Apple 神经引擎（ANE）的 LLM 训练与推理系统技术总结

1. 研究背景与问题 (Problem)

Apple 设备中广泛部署了神经处理单元（NPU），即 Apple Neural Engine (ANE)。自 A11 芯片以来，超过 20 亿台活跃设备搭载了 ANE，且在 M4 芯片上其峰值算力（INT8）可达 38 TOPS。然而，尽管硬件基础庞大，ANE 在大语言模型（LLM）领域几乎未被利用，原因如下：

• 软件栈封闭：Apple 的公开框架 CoreML 是一个黑盒调度器，开发者无法强制指定 ANE 执行，也无法直接控制梯度计算或进行模型训练。
• 缺乏直接编程接口：ANE 的原生指令集（MIL IR）未公开，编译和评估 API 位于私有框架 AppleNeuralEngine.framework 中。
• 现有方案局限：
  • CoreML 不支持训练，且调度不可控。
  • 先前的逆向工程工作（如 maderix）仅完成了硬件基准测试和简单的 C 语言调用，未构建完整的 LLM 运行时。
  • ANEgpt 项目尝试了 Transformer 训练，但存在训练步数后 NaN 发散、缺乏编译器优化、依赖 Python 编排以及无法断点续训等问题。
• 训练瓶颈：ANE 采用“编译后分发”（compile-then-dispatch）模式，权重在编译时固化（baked）。传统的训练方法意味着每次权重更新都需要重新编译整个程序，导致极高的延迟（约 4.2 秒/步），且受限于每进程约 119 次的编译上限。

2. 方法论 (Methodology)

Orion 是首个开源的端到端系统，通过完全绕过 CoreML，直接利用 Apple 的私有 API（ANEClient, ANECompiler 等）实现 ANE 的直接编程。系统架构分为五层：

2.1 系统架构

1. 编译器 (Compiler)：将包含 27 种操作的高层图 IR 转换为 ANE 原生的 MIL 代码。
2. 运行时 (Runtime)：管理程序编译、缓存、评估及基于 IOSurface 的零拷贝 I/O。
3. 模型层：处理配置、分词器和权重加载。
4. CLI 应用：支持推理、训练和基准测试。

2.2 核心技术创新

A. ANE 特性刻画 (Characterization)

Orion 扩展了先前的硬件基准测试，整理出 20 项 ANE 编程约束，其中 14 项为首次发现。关键约束包括：

• MIL IR 限制：不支持 concat 操作（需拆分为独立程序）；gelu 激活函数必须替换为 tanh 近似；矩阵乘法转置标志需命名常量等。
• 内存与 I/O：多输出/输入缓冲区必须具有统一的字节大小；输入/输出需按 MIL 变量名字母顺序排列；最小 IOSurface 大小约为 49KB（需填充序列维度）。
• 编译限制：每进程编译次数上限约 119 次；权重在编译时固化。

B. 增量编译 (Delta Compilation)

针对“权重固化导致每次训练步需全量重编译”的瓶颈，Orion 提出了增量编译技术：

• 原理：利用 ANEModel 的 unload 和 load 接口。当权重更新时，无需重新解析 MIL 文本或调用 ANECompiler，而是直接卸载现有程序，在磁盘上修补权重文件（BLOBFILE），然后重新加载。
• 效果：由于 MIL 文本和权重字典键未变，ANE 会保留程序的内部标识（hexStringIdentifier），仅刷新权重数据。这完全绕过了编译器，将单步重编译时间从 4,200ms 降低至 494ms。

C. 数值稳定性 (Numerical Stability)

解决了导致 ANEgpt 在训练第一步后 NaN 发散的三个关键 Bug：

1. 延迟编译：确保在加载检查点权重后再编译程序，避免权重不匹配。
2. FP16 溢出钳制：在 Softmax 和 LayerNorm 前将激活值钳制在 [-65504, +65504] 范围内，防止溢出。
3. 梯度清洗：在写入 BLOBFILE 前，将 NaN 替换为 0，将无穷大替换为边界值，并增加验证机制。

D. LoRA 适配器作为输入 (LoRA Adapter-as-Input)

为了在不重新编译的情况下切换适配器，Orion 将 LoRA 的低秩矩阵（A 和 B）作为 IOSurface 输入 传递，而非固化在程序中的权重。这实现了适配器的热插拔（Hot-swap）。

3. 关键贡献 (Key Contributions)

1. 首个开源 ANE LLM 系统：实现了从直接 ANE 执行、编译器流水线到稳定多步训练（含断点续训）的完整原生运行时。
2. ANE 约束目录：发布了包含 20 项约束的详细目录，填补了公开知识的空白，特别是 14 项新发现的 MIL IR 和内存布局限制。
3. 增量编译技术：通过“卸载 - 修补 - 重载”机制，将训练中的重编译开销降低了 8.5 倍，并消除了每进程编译次数限制。
4. 稳定的训练实现：在 M4 Max 上成功训练了 1.1 亿参数的 Transformer 模型（Stories110M），在 1000 步训练中零 NaN 发生，并实现了断点续训。
5. LoRA 热插拔：通过 IOSurface 输入实现适配器切换，无需重新编译。
6. 完全开源：以 MIT 许可证发布，核心运行时为 Objective-C，仅在权重转换时使用 Python。

4. 实验结果 (Results)

实验在配备 M4 Max 芯片（16 核 ANE）的 Mac Studio 上进行。

4.1 推理性能 (Inference)

• GPT-2 124M：
  • 解码模式：ANE 达到 170 tokens/s（CPU 基线为 283 tokens/s，受限于 IOSurface 往返开销）。
  • 预填充模式：缓存后达到 165 tokens/s。
  • 精度：与 CPU FP32 基线相比，Token 级完全一致，最大 Logit 误差仅为 0.073。

4.2 训练性能 (Training)

• Stories110M (TinyStories 数据集)：
  • v1.0 (全量重编译)：每步耗时 5,108ms（其中重编译占 4,200ms），1000 步需 85 分钟。
  • v2.0 (增量编译)：每步耗时 1,345ms（重编译降至 494ms），1000 步仅需 22.4 分钟。
  • 加速比：总训练时间加速 3.8 倍。
  • 稳定性：在 5 个独立的断点续训链（共 25 步）中，未出现任何 NaN 或 Inf，损失单调下降。

4.3 其他发现

• 特定操作加速：ANE 在大规模词表 Softmax 操作上比 CPU 快 33.8 倍（2.40ms vs 81.11ms）。
• 功耗：ANE 在空闲时硬断电（Zero idle power），适合始终在线的推理场景。

5. 意义与影响 (Significance)

1. 解锁被忽视的算力：证明了数十亿台 Apple 设备中的 ANE 硬件可以被直接用于 LLM 的训练和推理，而不仅仅依赖 GPU 或 CPU。
2. 打破训练壁垒：通过增量编译技术，解决了 ANE“权重固化”这一看似不可逾越的训练瓶颈，使得在端侧设备进行多步迭代训练成为可能。
3. 方法论价值：展示了通过逆向工程私有 API 来探索 NPU 约束的系统性方法，为其他厂商的 NPU（如 Qualcomm Hexagon, Samsung NPU）研究提供了参考。
4. 开源生态：Orion 的开源为社区提供了构建 NPU 原生 AI 系统的基础设施，推动了端侧大模型（On-Device AI）的发展。

总结：Orion 不仅是一个技术实现，更是一次对 Apple 封闭生态的“破壁”尝试。它通过深入理解硬件约束、创新编译优化策略（增量编译）以及解决数值稳定性问题，成功将 ANE 从单纯的推理加速器转变为支持完整训练流程的通用 AI 加速器。