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

全球已有超过二十亿台苹果设备出厂预装神经网络处理单元（NPU）——即苹果神经引擎（Apple Neural Engine, ANE），但该加速器目前在大语言模型（LLM）相关任务中仍基本处于闲置状态。苹果官方机器学习框架 Core ML 采用了一层不透明的抽象封装，既不允许开发者直接编程控制 ANE，也不支持设备端训练（on-device training）。为此，我们提出 Orion 系统：据我们所知，这是首个开源、端到端的完整系统，它将 ANE 的直接执行能力、专用编译器流水线，以及支持断点续训（checkpoint resume）的稳定多步训练全部集成于单一原生运行时之中，并完全绕过 Core ML，转而通过苹果私有的 `_ANEClient` 和 `_ANECompiler` API 实现底层调用。 在前人（maderix 团队）对 ANE 进行特征刻画（characterization）工作的基础上，我们进一步扩展了公众对 ANE 硬件约束的认知，系统梳理出涵盖 MIL 中间表示（MIL IR）程序、内存布局、编译限制及数值行为等维度的共 20 条约束规则；其中 14 条为 Orion 开发过程中首次发现、此前从未公开披露的约束。 Orion 包含一个专用编译器，可将图结构中间表示（graph IR）经由五轮优化遍历，逐步降级（lower）为 ANE 原生支持的 MIL 格式；同时配备一个轻量高效运行时，负责管理基于 IOSurface 的零拷贝张量输入/输出（zero-copy tensor I/O）、已编译程序缓存（program caching），以及面向权重更新的差分编译（delta compilation）。 由于 ANE 要求权重在编译阶段即完成固化（bake weights at compile time），若采用常规方式训练，每一步迭代均需重新完整编译整个程序，耗时约 4.2 秒。我们发现，实际可采用“卸载—修补权重文件—重载”的方式动态更新已编译程序，从而彻底绕过耗时的 `ANECCompile()` 调用，将单步重编译时间从 4200 毫秒大幅缩短至 494 毫秒（提速 8.5 倍），最终实现整体训练速度提升 3.8 倍。 在搭载 M4 Max 芯片的设备上，Orion 实现了 GPT-2（1.24 亿参数）模型推理吞吐量超 170 tokens/秒；同时成功在 TinyStories 数据集上，以零 NaN（非数值）异常的稳定表现，完成一个 1.1 亿参数 Transformer 模型共计 1000 步的端到端训练，全程仅耗时 22 分钟。此外，我们还提出“LoRA 适配器作为输入”（LoRA adapter-as-input）的新范式：用户可通过 IOSurface 输入动态热切换（hot-swap）不同 LoRA 适配器，全程无需任何重新编译。

Apple设备广泛搭载的神经引擎（ANE）因CoreML框架的抽象层和缺乏底层控制，长期无法用于大语言模型（LLM）的端侧训练与高效推理；现有方案无法支持直接ANE编程、编译时权重更新、多步稳定训练或热插拔适配器，导致ANE在生成式AI场景中严重未被利用。

绕过CoreML，首次通过逆向工程调用Apple私有_ANEClient和_ANECompiler API构建端到端原生运行时Orion；提出delta编译技术——通过IOSurface零拷贝卸载/补丁/重载权重文件替代全量重编译，将单步训练延迟从4200ms降至494ms；并实现MIL IR层级的五阶段编译优化与20条ANE硬约束建模（含14条新发现），使ANE首次支持稳定多步训练与LoRA热交换。

实验在M4 Max上完成：GPT-2 124M达170+ tokens/s推理吞吐；TinyStories上110M参数Transformer训练1000步仅22分钟且全程无NaN；首创LoRA adapter-as-input机制，通过IOSurface输入动态切换适配器无需重编译；完整开源（代码已发布）；基于maderix前期逆向工作深化，系统性揭示ANE内存布局、数值精度（如FP16梯度截断）、编译内存上限（<1.8GB）等20条约束；delta编译为硬件级增量学习提供新范式。

maderix, 'Reverse Engineering Apple's Neural Engine' (2023); Apple, 'Core ML Framework Documentation' (2022–2024); Wu et al., 'LLM on Edge: A Survey of On-Device Large Language Models' (arXiv:2402.13135); Chen et al., 'DeltaTuning: A General Framework for Parameter-Efficient Fine-Tuning' (ACL 2023); Intel, 'OpenVINO LLM Inference on CPU/GPU' (2024); Google, 'Gemini Nano on Pixel 8' (2023)