ANE

长期以来，Apple Silicon 芯片中的 神经引擎 (ANE, Apple Neural Engine) 被视为专为推理加速设计的“黑盒”。苹果通过 CoreML 等官方框架，严格限制其仅用于模型推断，而将训练任务完全交由 GPU (Metal) 处理。

ANE Training 的开发者通过逆向工程苹果的私有 API，成功实现了直接在 ANE 硬件上进行神经网络的反向传播训练。这意味着，无需 CoreML 训练接口，无需 Metal，甚至无需 GPU 参与计算，纯靠 ANE 即可完成训练流程。

项目核心：证明“不可能”成为可能

该项目并非旨在取代现有的生产级框架（如 MLX 或 llama.cpp），而是一个概念验证 (PoC)。其核心目标非常明确：

• 打破软件壁垒：证明 ANE 硬件本身具备训练能力，唯一的障碍是苹果的软件限制。
• 逆向工程突破：利用 _ANEClient 和 _ANECompiler 等未文档化的私有 API，绕过官方限制，直接下发自定义计算图。
• 探索硬件极限：通过基准测试记录 ANE 的真实性能特征，为未来研究提供参考。

开发者声明：“这是一个研究项目，而非生产级框架。目标是证明在 NPU 上训练是可能的。许多报道夸大了其意义，目前它并不能取代 GPU 训练大型模型，但它展示了当我们给硬件一个机会时可能实现的功能。”

技术实现：纯 ANE 计算的奇迹

该项目从头开始实现了一个在 Apple Silicon ANE 上运行 Transformer 单层训练 的系统。

关键架构细节

• 私有 API 调用：通过 objc_msgSend 在运行时动态解析私有 API，构建 MIL (Model Intermediate Language) 程序文本。
• 内存编译：使用 _ANEInMemoryModelDescriptor 将 MIL 文本和权重直接编译为 ANE 可执行程序，无需写入磁盘。
• 零拷贝 I/O：输入/输出张量通过 IOSurface 共享内存传递，极大降低了数据传输开销。
• 混合计算流：
  • ANE 负责：所有前向传播、反向传播中的 dx (输入梯度) 计算。
  • CPU 负责：RMSNorm 反向、残差连接、损失计算、dW (权重梯度) 累积及 Adam 优化器更新。

极致优化策略

为了在有限的资源下榨干性能，项目实施了多项硬核优化：

1. 通道优先布局：匹配 ANE IOSurface 格式，消除所有张量转置开销。
2. vDSP 向量化：CPU 端的 RMSNorm 反向计算速度提升 10 倍。
3. 异步重叠执行：利用 GCD 将 CPU 端的 cblas (dW 梯度计算) 与 ANE 的评估过程并行运行。
4. 操作融合：将 RMSNorm 融合进前向内核，将输出投影反向合并到 SDPA 反向内核中。
5. 绕过编译限制：ANE 编译器存在每进程约 119 次编译的资源泄漏限制，项目通过带检查点的 exec() 重启机制巧妙绕过。

性能表现：潜力巨大，挑战尚存

目前的基准测试结果显示了显著的性能提升轨迹，但也揭示了现实差距：

• 当前性能：9.3 毫秒/步，ANE 利用率达到 11.2%。
• 优化历程：从最初的 33.5 毫秒/步 (3.1% 利用率) 优化至今，性能提升了近 4 倍。
• 调度效率：每个训练步仅需 6 次 ANE 内核调度。

已知局限

• 利用率偏低：11.2% 的利用率意味着仍有大量算力未被释放，主要受限于逐元素操作回退到 CPU 执行。
• 单层限制：目前仅支持单个 Transformer 层的训练，多层模型需要复杂的流水线调度。
• 编译瓶颈：权重变化需重新编译程序，且受限于编译器的资源泄漏问题。
• 因果掩码缺失：ANE 硬件在 SDPA 操作中忽略注意力掩码，导致因果注意力需分解为多步执行。

项目定位与未来

开发者明确表示，ANE Training 不是一个维护型的开源库，而是一个持续探索的研究代码库：

• 不承诺长期维护：不会处理功能请求，但欢迎 Bug 修复和基准测试贡献。
• MIT 许可证：鼓励社区 Fork、修改和构建更好的工具。"如果这个项目对你有用——拿走它，修改它，构建更好的东西。"
• 风险提示：项目依赖私有 API，随时可能因 macOS 更新而失效，仅供研究和教育目的。