ANE 训练 — 在 Apple Neural Engine 上进行反向传播

通过逆向工程的私有 API，直接在 Apple 的 Neural Engine (ANE) 上训练神经网络。不使用 CoreML 训练 API，不使用 Metal，也不依赖 GPU——纯粹的 ANE 计算。

项目范围与意图

我由衷感谢大家对这个项目的关注——我从未想过一个周末的研究性尝试会如此火爆。感谢所有点赞、叉仓、在自己的硬件上运行基准测试并分享成果的人，这对我来说意义重大。

不过，我想明确一下这个项目的定位和边界。

这是一个研究项目，而非生产级框架。

我们的目标是证明在 Apple Neural Engine——以及潜在的其他 NPU 上进行训练是可行的，并且一直以来的瓶颈在于软件支持，而非硬件能力。ANE 是一块功能强大的芯片，但 Apple 通过 CoreML 将其限制为仅用于推理。本项目通过逆向工程的私有 API 绕过了这一限制，展示了只要给予硬件机会，就能实现怎样的可能性。

本项目是什么

• 基于 _ANEClient 和 _ANECompiler 私有 API 的 ANE 训练概念验证
• 一组记录真实 ANE 性能特征（吞吐量、功耗、SRAM 行为）的基准测试
• 为任何探索 CoreML 之外 ANE 直接访问方式的人提供的参考
• 我会在发现有趣内容时更新的研究代码

本项目不是什么

• 不是一个维护良好的框架或库
• 不是 CoreML、MLX、llama.cpp 或任何生产级推理栈的替代品
• 至少目前还不是在消费级硬件上训练大型模型的途径

关于炒作

一些关于该项目的报道夸大了其影响。澄清如下：

• 训练确实可以运行，但利用率较低（约峰值的 5–9%），且仍存在显著的工程挑战
• 许多逐元素操作仍然回退到 CPU
• 这并不意味着它可以取代 GPU 训练，至少在当前阶段还无法用于任何超出小型研究模型的任务

诚实的结果——包括所有局限性——都记录在配套的文章中：

• 第 1 部分：逆向工程
• 第 2 部分：基准测试
• 第 3 部分：训练

关于维护

我不打算将这个项目发展成一个大型社区项目。我的重心在于原创研究（边缘 AI 优化的编译器基础设施），而维护一个开源框架会占用我在这方面的精力。

尽管如此：

• 只要我发现有趣的内容，就会继续推送更新
• 欢迎提交 bug 修复和基准测试贡献（尤其是针对我没有的硬件）
• 功能请求可能不会被采纳——但欢迎大家 fork
• PR 的合并速度可能会比较慢，否则我会成为推动该技术社区发展的瓶颈。

Fork 它，基于它构建

之所以采用 MIT 许可证，是有原因的。如今，每个人都可以使用 AI 辅助开发工具，在几小时内就完成代码的适配和扩展。如果这个项目对你有用，那就拿走它，修改它，构建出更好的东西。如果你用它做出了什么酷炫的东西，我很乐意听听你的故事。如果未来社区决定维护一个事实来源仓库，我也完全支持。

---

这是什么

一个从头开始实现的 Transformer 训练框架（前向 + 反向传播），运行在 Apple Silicon 的 ANE 上。ANE 是一款 15.8 TFLOPS FP16（M4）的推理加速器，但 Apple 并未将其开放用于训练。本项目通过逆向工程 _ANEClient / _ANECompiler 私有 API 以及 MIL（Model Intermediate Language）格式，实现了在 ANE 硬件上直接运行自定义计算图——包括反向传播。

当前结果：

模型	参数量	ms/步	流水线
Stories110M（12 层，dim=768，MHA 12/12）	1.09 亿	91 ms	动态（无需重新编译）
Qwen3-0.6B（28 层，dim=1024，GQA 16/8）	5.96 亿	412 ms	动态（无需重新编译）

• 所有前向和反向 dx 传递都在 ANE 上执行，dW 梯度则在 CPU 上计算（Accelerate cblas）
• 使用 Adam 优化器，支持梯度累积，并通过 exec() 重启实现检查点和恢复
• 支持 GQA（分组查询注意力），并采用每头平铺与归约策略
• 通过共享 IOSurface 实现 GPU↔ANE 零拷贝流水线（GPU 预填充 → ANE 解码）

INT8 W8A8 量化——M4, H16G 下吞吐量提升 1.88 倍：

配置	FP16	INT8 W8A8	加速比
128x 卷积 512 通道 64x64	18.6 TOPS，14.8ms	35.1 TOPS，7.8ms	1.88x
64x 卷积 512 通道 64x64	18.4 TOPS，7.5ms	34.1 TOPS，4.0ms	1.85x

INT8 激活值通过 MIL quantize/dequantize 操作，将 tile 之间的 L2 SRAM 带宽减半。权重使用 constexpr_affine_dequantize（以 int8 存储，编译时转换为 fp16）。

架构

动态流水线使用共享 ANE 内核，权重打包进空间维度（权重变化时无需重新编译）：

MHA 模型（Stories110M）——每层 6 个内核：

内核	功能
sdpaFwd	QKV 投影 + SDPA + 输出投影
ffnFused	SwiGLU FFN（W1、W3、SiLU、W2）
ffnBwdW2t / ffnBwdW13t	FFN 反向传播（为节省内存拆分为两部分）
sdpaBwd1 / sdpaBwd2	SDPA 反向传播

GQA 模型（Qwen3-0.6B）——每层 10 个内核： 新增独立的 woFwd、qBwd 和 kvBwd 内核，用于处理分组查询注意力（Q_DIM ≠ DIM）。

CPU 负责：RMSNorm 的前向和反向计算、残差连接（DeepNet α 缩放）、损失计算、dW 梯度累积（cblas_sgemm）以及 Adam 优化器的更新。

关键优化：

• 通道优先的 CPU 布局——与 ANE IOSurface [1,C,1,S] 格式一致，消除所有转置开销
• vDSP 向量化 RMSNorm——比原生实现快 10 倍（6.7ms → 0.7ms）
• GCD 异步 cblas 重叠——dW 梯度 sgemms 在串行调度队列上与 ANE 评估并行执行
• 延迟 cblas 等待——将等待推迟到下一步的前向传播中，以实现最大重叠
• ANE RMSNorm 融合——将 RMSNorm 作为 MIL 操作融合进前向内核（reduce_sum + pow + mul）
• Wo^T 融合——输出投影的反向传播与 SDPA 反向传播内核合并
• 前向输出引出——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 矩阵乘法）
├── ane_int8_bench.m            # INT8 W8A8 与 FP16 吞吐量对比基准测试
├── sram_bench.m                # ANE SRAM 带宽探测
├── sram_probe.m                # SRAM 大小/布局探索
├── gpu_ane_share.m             # GPU↔ANE 零拷贝 IOSurface 示例
├── gpu_prefill_ane_decode.m    # GPU 预填充 → ANE 解码流水线
├── bridge/
│   ├── ane_bridge.h            # 可由 C 调用的 ANE API（编译、评估、I/O）
│   ├── ane_bridge.m            # 桥接实现（int8 + fp16 权重块）
│   └── Makefile
└── training/
    ├── ane_runtime.h           # ANE 私有 API 封装（编译、评估、IOSurface）
    ├── ane_classifier.h        # 分类器前向传播（32K 卷积）、softmax、rmsnorm 在 ANE 上执行
    ├── train_large.m           # 静态流水线（权重作为常量，需重新编译）
    ├── training_dynamic/
    │   ├── train.m             # 动态训练循环（模型无关）
    │   ├── config.h            # 衍生尺寸、结构体、内存分配辅助函数
    │   ├── mil_dynamic.h       # 用于动态权重核的 MIL 生成器（支持 GQA）
    │   ├── io.h                # IOSurface I/O、权重暂存、GQA 片段化/归约
    │   ├── models/
    │   │   ├── stories110m.h   # Stories110M 配置（12 层，MHA）
    │   │   └── qwen3_06b.h    # Qwen3-0.6B 配置（28 层，GQA）
    │   └── Makefile
    ├── dashboard.py            # 实时训练仪表板（受认可的 TUI）
    └── Makefile

训练数据

训练需要预分词的 TinyStories 数据。下载方法如下：

cd training && bash download_data.sh

详细训练说明请参阅 training/README.md。

构建

需要 macOS 15+ 且搭载 Apple Silicon 芯片（已在 M4 上测试）。

# 动态流水线（推荐）—— 模型在构建时选定
cd training/training_dynamic
make MODEL=stories110m    # Stories110M（12 层，MHA，1.09 亿参数）
make MODEL=qwen3_06b      # Qwen3-0.6B（28 层，GQA，5.96 亿参数）
./train --scratch          # 从随机初始化开始训练
./train --resume           # 从检查点恢复训练

# 静态流水线（遗留方案——每步都会重新编译权重）
cd training && make train_large
./train_large ane_stories110M_ckpt.bin 256 100 1e-4

# INT8 基准测试
xcrun clang -O2 -fobjc-arc -framework Foundation -framework IOSurface -ldl \
  -o ane_int8_bench ane_int8_bench.m
./ane_int8_bench

# 桥接库（可由 C 调用的 ANE API）
cd bridge && make

无外部依赖。仅使用系统框架及通过 objc_msgSend 在运行时解析的 ANE 私有 API。

工作原理

1. MIL 生成 — Objective-C 代码在运行时构造 MIL 程序文本，指定卷积（用于线性层）、矩阵乘法（用于注意力机制）、softmax 和逐元素操作。
2. 内存中编译 — _ANEInMemoryModelDescriptor 直接将 MIL 文本和权重块编译为 ANE 程序，无需磁盘上的 mlmodelc 文件。
3. IOSurface I/O — 输入输出张量通过 IOSurface 共享内存以 [1, 通道数, 1, 空间维度] 格式传递（fp16 或 fp32；fp16 直接 I/O 速度提升约 37%）。
4. 动态权重 — 激活值和权重打包到单个空间输入维度中，在 MIL 内核内部再进行切分。权重可在不重新编译的情况下变化。
5. 梯度流 — 前向传播会暴露反向传播所需的中间结果；反向传播内核在 ANE 上计算输入梯度 dx；权重梯度 dW 则由 CPU 通过 cblas 计算。
6. INT8 量化 — 使用 constexpr_affine_dequantize 处理 int8 权重，并在各层之间进行量化/反量化，以便在 L2 SRAM 中缓存 int8 激活值（吞吐量提升 1.88 倍）。

限制

• SDPA 因果掩码 — ANE 硬件会忽略 SDPA 操作中的 attn_mask；因果注意力被分解为单独的 Q@K^T（ANE）→ 掩码+softmax（CPU）→ 分数@V（ANE）。
• 约 119 次编译限制 — ANE 编译器存在资源泄漏问题，可通过重启进程并加载检查点来规避。
• FP16 梯度下溢 — 反向传播中的矩阵乘法在 fp16 下会发生下溢，已通过全局损失缩放（256 * NLAYERS）解决。
• 单输入约束 — ANE 请求多输入会导致 0x1d 错误；应将输入打包到空间维度中。

性能

训练吞吐量（M4）：

模型	参数	每步时间 (ms)	层数	每层内核数
Stories110M	1.09 亿	91	12	6（MHA）
Qwen3-0.6B	5.96 亿	412	28	10（GQA）

ANE 峰值吞吐量（M4，H16G）：

精度	峰值 TOPS	配置
FP16	18.6	128 个卷积，512 通道，64×64
INT8 W8A8	35.1	128 个卷积，512 通道，64×64

GPU↔ANE 推理流水线（M4，序列长度=256）：

模型	GPU 预填充	ANE 解码	总计
Stories110M	6.7 ms	1.9 ms	8.8 ms
Qwen3-0.6B	9.7 ms	2.3 ms	12.0 ms

免责声明

本项目使用了 Apple 的私有未文档化 API（_ANEClient、_ANECompiler、_ANEInMemoryModelDescriptor）。这些 API 不受任何公开稳定性保证，可能随 macOS 更新而更改或失效。本项目是对 Apple Neural Engine 架构的独立研究，所使用的 API 是通过运行时内省发现的，仅用于研究和教育目的，符合合理使用和互操作性条款（参见 Sega v. Accolade, 1992；DMCA §1201(f)）。本仓库不包含任何 Apple 专有代码或二进制文件。本项目与 Apple Inc. 无任何关联，亦未获其认可。请用户自行承担风险。

许可证

MIT — 详见 LICENSE

---

由人类与 Claude 共同打造，利用每个周末的时间完成。

ANE 快速上手指南

ANE 是一个研究性项目，旨在通过逆向工程 Apple 的私有 API，直接在 Apple Neural Engine (ANE) 上运行神经网络的反向传播（训练）。本项目不依赖 CoreML 训练接口或 Metal，纯粹利用 ANE 算力。

注意：这是一个概念验证（PoC）研究项目，并非生产级框架。目前主要用于探索硬件潜力和小规模模型的研究训练。

环境准备

• 操作系统：macOS 15+ (Sequoia)
• 硬件要求：Apple Silicon 芯片（已在 M4 上测试，其他 M 系列芯片可能兼容但性能不同）
• 开发工具：Xcode Command Line Tools (包含 clang, make 等)
• 依赖项：无外部依赖，仅使用系统框架和运行时解析的私有 ANE API

确保你的终端已安装必要的构建工具：

xcode-select --install

安装步骤

克隆仓库并进入训练目录：

git clone <repository-url>
cd ANE/training/training_dynamic

根据你想训练的模型选择编译目标（二选一）：

选项 1：编译 Stories110M 模型 (1.09 亿参数，MHA 架构)

make MODEL=stories110m

选项 2：编译 Qwen3-0.6B 模型 (5.96 亿参数，GQA 架构)

make MODEL=qwen3_06b

下载训练数据 项目需要预分词的 TinyStories 数据集。运行以下脚本自动下载：

bash download_data.sh

基本使用

编译完成后，可直接运行训练脚本。

从头开始训练（随机初始化权重）

./train --scratch

从检查点恢复训练 如果之前中断过，可以使用以下命令恢复：

./train --resume

查看实时训练仪表盘 项目包含一个基于终端的实时监控界面（需 Python 环境）：

python dashboard.py

进阶：INT8 量化基准测试

如果你想测试 ANE 的 INT8 W8A8 量化吞吐量，可在根目录编译并运行基准测试：

cd ../..
xcrun clang -O2 -fobjc-arc -framework Foundation -framework IOSurface -ldl \
  -o ane_int8_bench ane_int8_bench.m
./ane_int8_bench

核心机制说明

• 动态权重：权重被打包到空间维度中，无需在每次权重更新时重新编译模型。
• 混合计算：前向传播和输入梯度 (dx) 在 ANE 上计算；权重梯度 (dW) 和优化器步骤在 CPU (Accelerate/cblas) 上计算以规避当前限制。
• 显存优化：利用 IOSurface 实现 GPU 与 ANE 之间的零拷贝数据共享。