直接在 Apple Neural Engine 上训练神经网络，无需 CoreML 和 GPU

每一颗 Apple Silicon 处理器都包含一个专门用于神经网络运算的模块——Apple Neural Engine（ANE）。苹果官方声称其性能可达数十 TFLOPS，但官方工具将其使用限制为仅通过 CoreML 进行推理。模型训练只能在 CPU 或图形芯片上进行。

一位名为 maderix 的开发者决定测试 ANE 硬件在物理上是否真的无法执行反向传播，还是问题仅存在于苹果的软件限制。他用一个周末的时间深入分析了 macOS 的私有框架，成功让 ANE 不仅运行前向传播，还运行了 transformer 的反向传播。

该 GitHub 项目已获得超过 7000 颗星。让我们深入解析这个 hack 的工作原理、需要应用哪些技巧，以及实际的性能表现。

为什么要深入研究封闭框架

CoreML 等官方技术栈给人一种印象：ANE 是一个封闭的黑箱。你给它一个现成的模型，它输出结果。如果你想直接在客户端训练哪怕是一个很小的模型，也必须转向 Metal 或 MLX 框架，从而占用 GPU。

ANE 仓库的作者证明了这款芯片完全能够执行任意计算图。为此，他逆向工程了私有库 _ANEClient 和 _ANECompiler，以及内部模型描述语言 MIL（Model Intermediate Language）。

结果非常有趣：在没有一行 CoreML 或 Metal 代码的情况下，实现了完整的 transformer 训练。模型文本直接在 RAM 中组装，动态编译后发送到神经处理器。

在急于重写 Mac 上的训练脚本之前，值得看看真实的数据。作者本人也提前警告：这是一个学术实验，而非生产就绪的库。

目前还无法实现 100% 的 ANE 资源利用率。实际芯片利用率约为峰值的 5–9%。软件和硬件限制开始显现：

尽管如此，即使在这种模式下，项目在基本架构上也能提供不错的速度。

项目架构依赖于 ANE 和 CPU 之间的任务分配。神经网络加速器处理最繁重的矩阵乘法，而处理器负责周围逻辑和梯度累积。

前向传播和输入梯度计算（dx）完全在 ANE 上执行。权重梯度（dW）通过优化的 Accelerate 和 cblas_sgemm 库由 CPU 计算。Adam 优化器和 RMSNorm 层也在处理器上运行。

为避免在权重变化时每次都重新编译模型图，作者采用了一个技巧：将权重和激活通过空间维度打包成单个张量，在 MIL 内核内部再将其简单拆分回来。

CPU 和 ANE 之间的数据交换使用 IOSurface 内存。这使得张量可以在地址空间之间无额外复制地进行传输。数据被打包成 ANE 特定格式 [1, C, 1, S]，其中通道优先排列。这种方法消除了转置矩阵的开销。

在逆向工程过程中暴露了许多 Apple 硬件的陷阱。

首先，ANE 中的 SDPA（Scaled Dot-Product Attention）操作在硬件层面忽略因果掩码 attn_mask。注意力机制不得不分成三个阶段：在 ANE 上将 Q 和 K 相乘，在 CPU 上用 softmax 进行掩码，最后在 ANE 上与 V 完成最终乘法。

其次，内置编译器 _ANECompiler 包含内存泄漏。大约 119 次编译后，进程会因资源耗尽而崩溃。作者用激进的方式解决了这个问题：当计数器接近限制时，程序保存检查点并进行 exec()——在保持状态的情况下重启自身。

第三，反向传播期间的 FP16 计算很快导致下溢，使梯度变成零。问题通过用系数 256 * NLAYERS 对损失进行缩放来解决。

在 Apple M4 芯片上，结果相当有说服力。测试在两种架构上进行：

对于拥有 1.09 亿参数（12 层，经典 Multi-Head Attention）的 Stories110M 模型，单个训练步骤的时间为 91 毫秒。

更大的 Qwen3-0.6B 拥有 5.96 亿参数和 Grouped-Query Attention，每个步骤处理时间为 412 毫秒。

作者还测试了 INT8 W8A8 量化。使用 8 位权重和激活降低了芯片 L2 SRAM 内存的负载，并将吞吐量从 18.6 TOPS 提升到 M4 上的 35.1 TOPS——与 FP16 相比加速近 1.88 倍。

该项目不需要 PyTorch 或 Conda 等外部依赖。你只需要最新 macOS 15 和配备 Apple Silicon 的机器以及 Clang 编译器。私有 API 通过 objc_msgSend 在运行时引入。

要构建动态流水线，只需导航到项目文件夹并运行 make 命令：

cd training/training_dynamic
make MODEL=stories110m
./train --scratch

如果你想测试 INT8 量化或基准测试芯片的峰值 TOPS，仓库根目录中有单独的基准测试。

maderix 项目是 Apple 硬件"引擎盖下"高质量研究的典范。它表明 ANE 的限制仅存在于软件层面和生态系统的封闭性。

目前使用该仓库在生产环境中训练大型语言模型毫无意义——GPU 和 MLX 框架为此而存在。但如果你正在研究神经加速器的工作原理、为 Edge AI 编写自己的编译器，或想了解如何直接从 C 和 Objective-C 使用 macOS 私有 API，这段代码将是一个极好的学习资源。