ANE — 在 Apple Neural Engine 上训练神经网络
GitHub: https://github.com/maderix/ANE
🔍 概述
ANE 是一个研究项目，首次公开展示了完整的**神经网络训练（前向 + 后向传播）**可以直接在 Apple 的 Neural Engine (ANE) 上执行——这是 Apple Silicon 芯片（M1/M2/M4 系列）内置的专用 AI 加速器。
Apple 通过 CoreML 框架官方将 ANE 限制为仅用于推理。本项目逆向工程 Apple 的私有 API（_ANEClient、_ANECompiler、_ANEInMemoryModelDescriptor）来绕过此限制，并在 ANE 硬件上原生运行自定义计算图——包括反向传播。
它不是一个生产级框架。它是一个概念验证和一个基准参考，证明了该限制是软件强制的，而非硬件限制。
🎯 重要意义
Apple Silicon 芯片（尤其是 M4）包含一个额定功率为15.8 TFLOPS 的 ANE——这是 Apple 锁定用于推理的大量专用 AI 计算能力。每次运行 CoreML 时，您只能获得推理。训练总是回退到 GPU 或 CPU。
本项目回答了这个问题：“能否在 ANE 上进行训练？”
答案是肯定的——而这个仓库就是证明。
⚙️ 工作原理
该项目通过以下方式实现完整的Transformer 层训练循环：
-
MIL 程序生成 — 在运行时使用 Objective-C 构建 Apple 的模型中间语言 (MIL) 程序，定义卷积（用于线性层）、矩阵乘法（用于注意力）、softmax 和逐元素运算。
-
内存中编译 — 使用
_ANEInMemoryModelDescriptor将 MIL 文本 + 权重 blob 直接编译成 ANE 程序，无需将.mlmodelc文件写入磁盘。 -
IOSurface I/O — 通过
IOSurface共享内存以[1, channels, 1, spatial]fp16 格式传递输入/输出张量——这是 ANE 硬件原生期望的格式。 -
权重嵌入 — 权重作为
BLOBFILE常量嵌入到 ANE 程序中，并在权重更新时为每个批次重新编译。 -
梯度流 — 前向“钩子”暴露了后向传播所需的中间激活。后向核在 ANE 上计算输入梯度（
dx）；权重梯度（dW）通过cblas_sgemm在 CPU 上计算。
🧱 架构：每个训练步骤 6 个 ANE 核
| 核函数 | 功能 |
|---|---|
kFwdAttn |
RMSNorm + QKV 投影 + SDPA + 输出投影 |
kFwdFFN |
RMSNorm + SwiGLU FFN (W1, W3, SiLU, W2) |
kFFNBwd |
FFN 后向（W2ᵀ + SiLU_bwd + W1ᵀ + W3ᵀ） |
kSdpaBwd1 |
Woᵀ + SDPA 后向部分 1 (dV, probs, dp) |
kSdpaBwd2 |
SDPA 后向部分 2 (softmax 梯度, dQ, dK) |
kQKVb |
QKV 后向（Wqᵀ + Wkᵀ + Wvᵀ → dx） |
CPU 处理： RMSNorm 后向、残差连接、损失计算、dW 梯度累加 (cblas_sgemm) 和 Adam 优化器更新。
📊 基准测试结果 (M4 Mac, 单个 Transformer 层, dim=768, seq=512)
| 优化 | ms/步 | ANE 利用率 |
|---|---|---|
| 基线 (vDSP 转置) | 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% |
最佳结果：9.3 ms/步，持续 1.78 TFLOPS（M4 ANE 15.8 TFLOPS 峰值的 11.2%）。
🔑 关键优化
- Channel-first CPU 布局 — 匹配 ANE 原生
[1,C,1,S]IOSurface 格式，消除了所有转置开销。 - vDSP 向量化 RMSNorm — 比朴素实现快 10 倍（6.7ms → 0.7ms）。
- GCD 异步 cblas 重叠 — 权重梯度
sgemm操作与 ANE 评估在串行调度队列上并行运行。 - 延迟 cblas 等待 —
dWsgemm 的等待被推迟到下一步的前向传播中，以实现最大的计算重叠。 - ANE RMSNorm 融合 — RMSNorm 直接作为 MIL 操作折叠到前向 ANE 核中。
- 前向钩子 — Q、K、V、注意力分数和隐藏状态通过连接输出来暴露，避免在后向传播期间进行 CPU 重计算。
- exec() 重启 — 通过检查点/恢复绕过了每个进程约 119 次 ANE 编译限制。
📁 仓库结构
├── api_exploration.m # 初始 ANE API 发现实验
├── inmem_basic.m # 内存中 MIL 编译概念验证
├── inmem_bench.m # ANE 分派延迟基准测试
├── inmem_peak.m # 峰值 TFLOPS 测量 (2048×2048 矩阵乘法)
├── sram_bench.m # ANE SRAM 带宽探测
├── sram_probe.m # SRAM 大小/布局探索
└── training/
├── ane_runtime.h # ANE 私有 API 包装器 (编译, 评估, IOSurface)
├── ane_mil_gen.h # MIL 程序生成辅助函数
├── model.h # 模型权重初始化和 blob 构建器
├── forward.h # 前向传播 MIL 生成器
├── backward.h # 后向传播 MIL 生成器
├── train.m # 最小训练循环 (早期原型)
├── tiny_train.m # 2 层小型模型训练
├── train_large.m # 主程序: 单层 dim=768 训练 (优化版)
├── test_*.m # 单个核的单元测试
└── Makefile
🛠️ 构建和运行
要求： macOS 15+ on Apple Silicon (已在 M4 上测试)。
# 构建
xcrun clang -O2 -framework Foundation -framework IOSurface \
-framework CoreML -framework Accelerate -ldl -lobjc \
-o train_large training/train_large.m
# 运行
./train_large
无外部依赖 — 仅使用系统框架以及在运行时通过 objc_msgSend 解析的私有 ANE API。
⚠️ 已知限制
| 限制 | 详情 |
|---|---|
| 低 ANE 利用率 | 峰值的约 11.2%；许多逐元素运算仍回退到 CPU |
| 约 119 次编译限制 | ANE 编译器资源泄露；通过 exec() 重启解决 |
| 单个 Transformer 层 | 多层流水线调度尚未实现 |
| 仅合成数据 | 支持真实 token 化数据是进行中的工作 |
| SDPA 因果掩码 | ANE 忽略 SDPA 中的 attn_mask；通过手动分解解决 |
| 私有 API | 使用未公开的 API，任何 macOS 更新都可能导致其失效 |
📚 相关研究文章
🏷️ 技术栈
- 语言： Objective-C, C
- 平台： macOS 15+ / Apple Silicon (M1/M2/M4)
- 框架： Foundation, IOSurface, CoreML, Accelerate
- 私有 API：
_ANEClient,_ANECompiler,_ANEInMemoryModelDescriptor - 格式： MIL (模型中间语言), fp16
⚖️ 法律声明
本项目使用 Apple 的私有、未公开 API。这些 API 不保证稳定性，并且任何 macOS 更新都可能导致其失效。本项目是基于合理使用和互操作性条款（Sega v. Accolade, 1992; DMCA §1201(f)）的独立研究。不包含任何 Apple 的专有代码或二进制文件。与 Apple Inc. 无关，也未得到其认可。