深入Apple Neural Engine: 逆向工程实现ANE上的神经网络训练
自 Apple Silicon M1 芯片问世以来，其内置的 Apple Neural Engine (ANE) 就以惊人的 AI 推理性能吸引了业界的广泛关注。在最新的 M4 芯片中，ANE 的算力已达到 38 TOPS (INT8) / 15.8 TFLOPS (FP16)，这是一个与数据中心级加速卡相媲美的数字。然而，Apple 始终将 ANE 定位为一个推理加速器，仅通过高层框架 CoreML 提供 API，开发者无法直接访问其底层接口，更不用说利用它进行模型训练。
这引出了一个核心问题：ANE 的硬件架构是否原生支持训练所需的反向传播等操作？还是说它仅仅是一个被高度特化的前向推理引擎？
一个名为 ANE Training 的开源项目给了我们答案。该项目作者通过对 Apple 私有框架的精湛逆向工程，成功地在 ANE 上实现了完整的 Transformer 模型训练（包括前向和反向传播）。这不仅展示了 ANE 硬件的真实潜力，也揭示了它远不止是一个推理引擎。
核心技术架构
项目的核心在于绕过 CoreML，直接与 ANE 的底层驱动和编译器进行交互。这主要依赖于对 AppleNeuralEngine.framework 和 CoreML.framework 中未公开的私有类的逆向分析。
私有 API 逆向
ANE Training 没有链接任何私有框架，而是选择在运行时动态加载它们，并通过 Objective-C 的动态消息派发机制 objc_msgSend 来调用私有 API。这种方法避免了编译时的直接依赖，保证了代码的灵活性。
关键的私有类和它们扮演的角色如下：
_ANECompilerService: 负责将模型编译为 ANE 可以执行的二进制程序_ANEInMemoryModel/_ANEInMemoryModelDescriptor: 允许直接从内存中的 MIL 文本和权重数据创建模型，无需经过文件系统_ANEClient: 与 ANE 硬件守护进程aned通信的客户端_ANERequest: 封装一次具体的计算请求，包含输入输出的IOSurface和要执行的模型
整个流程可以概括为三个关键函数：
ane_init() — 使用 dlopen 在运行时动态加载 AppleNeuralEngine.framework，然后通过 NSClassFromString 获取私有类的句柄：
// 动态API解析的概念示例
typedef id (*objc_msgSend_t)(id, SEL, ...);
objc_msgSend_t msgSend = (objc_msgSend_t)objc_msgSend;
// 动态获取 _ANEClient 类
Class ANEClient = NSClassFromString(@"_ANEClient");
id client = msgSend(ANEClient, sel_getUid("sharedClient"));
ane_compile() — 接收 MIL 文本字符串和权重二进制 blob，通过 _ANEInMemoryModelDescriptor 创建模型描述符，再经 _ANECompilerService 编译为 ANE 可执行程序。加载失败时会进行 100ms 重试（等待 ANE 资源回收）。
ane_eval() — 创建 _ANERequest，绑定输入输出的 IOSurface 缓冲区，通过 evaluateWithQoS:options:request:error: 提交给 ANE 硬件执行。
MIL (Model Intermediate Language)
MIL 是 ANE Training 与 ANE 编译器沟通的语言。它是一种类似于文本汇编的中间表示，用于描述神经网络的计算图。项目在运行时动态生成这些 MIL 程序文本。
与 MIL 文本相伴的是权重二进制 blob，它的格式非常特殊：
- 一个 64 字节的全局头部（
buf[0]=0x01, buf[4]=0x02） - 随后是多个权重块(chunk)，每个块以 64 字节头部开始
- 每个 chunk 头部包含魔数
0xDEADBEEF、数据大小和偏移量 - 权重数据以 FP16 格式存储（F32 通过
(_Float16)截断转换）
在训练过程中，权重需要不断更新。项目实现了 F32（CPU 端 Adam 优化器使用）到 FP16（ANE 硬件使用）的高效转换。MIL 支持的操作包括卷积（用于线性层）、matmul（用于注意力）、softmax 等。
IOSurface 张量 I/O
CPU 和 ANE 之间的数据交换是性能的关键瓶颈。为了避免昂贵的数据拷贝，ANE Training 全面采用 IOSurface——macOS 中用于在不同进程或设备之间共享内存的框架，实现了真正的零拷贝传输。
当 CPU 需要读写 IOSurface 中的张量数据时，必须先调用 IOSurfaceLock 获得访问权并确保数据同步，操作完成后再调用 IOSurfaceUnlock。
一个值得注意的细节是 ANE 对张量布局的偏好。ANE 内部偏爱 Channels-First 的数据格式，具体表现为 [1, Channels, 1, Spatial]。为了最大化性能，项目在 CPU 端的所有计算中都维持了这个布局，从而避免了在每次 ANE 调用前后进行耗时的 transpose 操作。
训练实现
一个完整的 Transformer Block 训练步骤被分解为多个 ANE 内核（Kernel）和 CPU 上的计算任务。
前向传播 (6 个 ANE 内核)
前向传播被分解为 6 个独立的 ANE 内核：
| 内核 | 功能 | 权重 |
|---|---|---|
kFwdAttn |
RMSNorm + QKV投影 + SDPA + 输出投影 | Wq, Wk, Wv, Wo, rms1, mask |
kFwdFFN |
RMSNorm + SwiGLU FFN (W1, W3, SiLU, W2) | W1, W2, W3, rms2 |
kFFNBwd |
FFN反向 (W2^T + SiLU_bwd + W1^T + W3^T) | W2^T, W1^T, W3^T |
kSdpaBwd1 |
Wo^T + SDPA反向第1部分 (dV, probs, dp) | Wo^T, mask |
kSdpaBwd2 |
SDPA反向第2部分 (softmax grad, dQ, dK) | — |
kQKVb |
QKV反向 (Wq^T + Wk^T + Wv^T → dx) | Wq^T, Wk^T, Wv^T |
这种分解并非任意，而是为了"截取"(tap)反向传播所需的中间激活值。前向内核不仅输出最终结果，还通过 concat 将 Q、K、V、注意力分数等中间结果写入 IOSurface，直接传递给反向传播内核，避免 CPU 介入。
操作如 RMSNorm 被融合进这些内核。归一化逻辑作为 MIL 操作（reduce_sum + pow + mul）预置在 kFwdAttn 和 kFwdFFN 程序中，减少了内核启动开销和内存带宽。
反向传播 (混合 CPU/ANE)
反向传播是一个混合计算过程：
ANE 负责 — 计算密集的转置矩阵乘法（dx = W^T @ dy），即输入梯度的计算。所有 kFFNBwd、kSdpaBwd1/2、kQKVb 内核都在 ANE 上运行。
CPU 负责 — 以下操作仍在 CPU 上执行：
- RMSNorm 反向: 通过 vDSP 向量化实现的链式法则
- 残差连接: 简单的梯度加法
- Loss 计算: 数值稳定的 softmax + log-prob（带 max 减法）
- 梯度累积 (dW): 通过 Accelerate 框架的
cblas_sgemm计算dW += dy @ x^T - Adam 优化器: 完整的优化器步骤（m/v 状态、偏差校正）以 FP32 精度执行
- 梯度裁剪: 全局 L2 范数计算，超过
max_norm时进行缩放
这是一个典型的异构计算模式：ANE 负责并行的、计算密集的张量操作，而 CPU 负责串行的、控制流复杂或 ANE 不支持的操作。
关键优化
为了榨干性能，项目采用了多项系统级优化：
1. Channel-First CPU 布局
通过在 CPU 端始终维持 [1, C, 1, S] 的数据布局，项目消除了所有 transpose 操作。对于 110M 模型，一次 transpose 可能耗时数毫秒，在每一步训练中累加起来是非常可观的开销。
2. vDSP 向量化 RMSNorm
最初，RMSNorm 的反向传播在 CPU 上是一个简单的 C 语言循环，耗时 6.7ms。通过使用 Accelerate 框架中的 vDSP 指令进行重写，利用了 CPU 的 SIMD 单元，将耗时骤降至 0.7ms，实现了近 10 倍的加速。
3. GCD 异步 cblas 重叠
梯度累积（cblas_sgemm）是一个 CPU 密集型任务。项目利用 Grand Central Dispatch (GCD) 将 sgemm 计算提交到一个后台串行队列，使其与 ANE 内核的执行并行。当 ANE 处理 kFwdAttn 时，CPU 可以同时计算上一步的 dW 梯度。
4. 延迟 cblas 等待
CPU 不需要在提交 sgemm 任务后立刻等待完成。等待操作被推迟到下一步的前向传播中真正需要该结果时才执行。这最大化了 ANE 和 CPU 的并行度。
5. ANE RMSNorm 融合
RMSNorm 实现为 MIL 操作，融合进 kFwdAttn 和 kFwdFFN 内核的第一个卷积层，完全消除了这一步的 CPU 计算和数据往返开销。
6. Wo^T 融合
在 Attention 反向传播中，dScores 需要乘以 Wo 的转置。这一步被融合进了 kSdpaBwd1 内核，将内核数从 7 个减少到 6 个。
7. exec() 重启
逆向工程发现一个奇特的限制：一个进程大约编译了 119 个 ANE 模型后，_ANECompilerService 就会开始返回错误。这似乎是 ANE 守护进程 aned 中的某种资源泄漏。解决方案简单而粗暴但有效：在达到限制前，保存二进制 checkpoint，使用 execl() 系统调用重启自身进程，并从 checkpoint 恢复训练状态。对用户来说，整个过程看起来像一次连续的训练运行。
优化效果
| 优化 | ms/step | ANE利用率 |
|---|---|---|
| 基线 (vDSP transpose) | 33.5 | 3.1% |
| Channel-first 布局 | 20.3 | 5.2% |
| vDSP 向量化 RMSNorm | 14.2 | 7.4% |
| GCD 异步 cblas 重叠 | 11.4 | 9.2% |
| ANE RMSNorm 融合 | 11.4 | 9.2% |
| Wo^T 融合 (7→6 内核) | 11.4 | 9.2% |
| 延迟 cblas 等待 | 9.3 | 11.2% |
动态权重管线
上述的训练流程被称为"静态管线"，因为模型权重被硬编码在编译好的 ANE 程序中。这意味着每一次权重更新后，都必须重新编译所有相关的 ANE 内核。分析显示，76% 的训练时间消耗在了模型编译上。
为了解决这个问题，项目设计了"动态管线"。其核心思想是：将权重视为模型的输入，而非静态数据。
实现方式
ANE 内核只编译一次，权重被定义为输入占位符。在每一步训练中，输入 x 和当前的模型权重被打包进一个巨大的 IOSurface 输入张量中：
输入张量布局 (以 sdpaFwd 为例):
[1, DIM, 1, SEQ + 4*DIM] fp32
[0:SEQ] = xnorm (激活值)
[SEQ:SEQ+DIM] = Wq (权重矩阵)
[SEQ+DIM:SEQ+2D] = Wk
[SEQ+2D:SEQ+3D] = Wv
[SEQ+3D:SEQ+4D] = Wo
ANE 内核从这个输入张量中切分（slice）出所需的权重来进行计算。权重更新只需通过 IOSurfaceLock → memcpy → IOSurfaceUnlock 完成，代价可以忽略不计。
对比
| 特性 | 静态管线 | 动态管线 |
|---|---|---|
| 编译 | 每次权重更新重编译 | 启动时编译一次 |
| 内核数 | 72 (每层独立) | 9 (所有层共享) |
| 编译开销 | 76% 训练时间 | 启动 2-3 秒 |
| 单步时间 | 更快 (编译器可优化常量权重) | 略慢 (额外打包/切分开销) |
SRAM 与性能分析
ANE 内部拥有高速的片上 SRAM，项目通过实验估算出其大小约为 16MB。当一个 ANE 内核的工作集能够完全放入 SRAM 时，性能极高。一旦超出这个阈值，数据就需要与系统 DRAM 交换，导致性能急剧下降（性能悬崖）。
峰值性能测量通过链式卷积基准测试完成，M4 ANE 可达到约 15.8 TFLOPS (FP16)。然而在实际的 Transformer 训练中，利用率仅为 5-9%（约 1-2 TFLOPS）。这说明整个训练流程的瓶颈不在 ANE 的计算单元，而在于频繁的 CPU 回退、CPU 与 ANE 之间的同步开销、以及多个内核之间的调度延迟。
Stories110M 训练结果
项目在一个 109M 参数的 Transformer 模型上进行了端到端训练测试（12 层，维度 768，12 头，序列长度 256）：
| 平台 | 管线类型 | 每步时间 (ms) | ANE内核数 |
|---|---|---|---|
| M3 Ultra | 静态 | 91 | 72 |
| M4 | 静态 | 106 | 72 |
| M3 Ultra | 动态 | 110 | 9 (共享) |
M3 Ultra 比 M4 更快这一点值得关注——原因是 M3 Ultra 拥有更高的 DRAM 带宽和更强的 CPU 性能，而当前训练流程受这些因素的制约远大于 ANE 原始算力。
工程实现
实现体现了以下几点特征：
- 纯 Objective-C 实现: 约 23,600 行代码，零外部依赖
- 仅使用系统框架: Foundation, IOSurface, CoreML, Accelerate——具有极佳的可移植性
- Bridge API: 提供 C-style 桥接接口，支持 Python 通过
ctypes调用，扩展了可用性 - 完备的测试: 13 个测试文件，覆盖从底层 API 到端到端训练的完整矩阵
- 二进制 Checkpoint: 自定义格式，包含权重、Adam 状态和累计训练时间，支持中断恢复
局限性与展望
项目揭示了当前利用 ANE 进行训练的几个边界条件：
- 因果掩码: ANE 硬件似乎忽略 SDPA 中的
attn_mask输入，导致因果注意力需要分解为 Q@K^T → mask+softmax → scores@V 三个独立操作 - ~119 次编译限制:
exec()重启是一个变通方案；真正的解决方案需要理解 Apple 编译器服务内部的资源泄漏 - 低硬件利用率: 5-9% 的利用率说明优化空间依然巨大，未来需要将更多计算（如梯度累积、优化器）移到 ANE 上
可行性和硬件约束
ANE Training 证明了在 Apple Neural Engine 上进行训练是可行的。当前的性能瓶颈主要源于软件层面——私有 API 的限制、不完整的硬件功能暴露、以及 CPU 与 ANE 之间的协作开销——而非硬件本身的根本性缺陷。
如果 Apple 愿意开放更底层的 ANE 编程接口，或者在 CoreML 中提供官方的训练支持，Apple Silicon 设备将有可能成为强大且能效比高的端侧 AI 训练平台。