AI Feeds
开源项目

ANE

通过逆向工程Apple私有API，在Apple Neural Engine上直接进行神经网络训练的开源实现，支持Transformer模型的前向和反向传播。亮点在于突破了CoreML仅推理的限制，展示了ANE硬件在训练上的潜力，在M4上可达109M参数模型91ms/step的性能。实测INT8量化可带来1.88倍吞吐提升。研究向项目，使用私有API有风险，非生产级框架。

README

ANE Training — Apple Neural Engine 上的反向传播训练

通过逆向工程得到的私有 API，直接在 Apple Neural Engine（ANE）上训练神经网络。无需 CoreML 训练 API，无需 Metal，无需 GPU——纯 ANE 计算。

项目范围与意图

我由衷感谢这个项目获得的所有关注——我从未想到一个周末的研究黑客会如此火爆。感谢每一位 star、fork、在自己硬件上运行基准测试并分享成果的人。这意义重大。

话虽如此，我想明确说明这个项目是什么、不是什么。

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

目标是证明在 Apple Neural Engine（以及其他 NPU）上进行训练是可能的，并且障碍一直是软件支持，而非硬件能力。ANE 是一块性能卓越的硅片，但 Apple 通过 CoreML 将其限制为仅推理使用。本项目使用逆向工程得到的私有 API 绕过这一限制，展示当你给硬件一个机会时能实现什么。

这个项目是什么

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

这个项目不是什么

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

关于炒作

一些关于此项目的报道夸大了其意义。明确说明：

  • 训练确实可行，但利用率较低（约为峰值的 5-9%），仍有大量工程挑战
  • 许多逐元素运算仍回退到 CPU
  • 目前这不能取代 GPU 对除小型研究模型之外的任何训练

诚实的成果——包括所有限制——在配套文章中记录：

关于维护

我不打算将这个项目发展成一个大型社区项目。我的重点是最初的研究（边缘 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 硬件上运行自定义计算图——包括反向传播。

当前结果：

模型 参数量 毫秒/步 流水线
Stories110M（12层，dim=768，MHA 12/12） 109M 91 ms 动态（无需重新编译）
Qwen3-0.6B（28层，dim=1024，GQA 16/8） 596M 412 ms 动态（无需重新编译）
  • 所有前向和反向 dx 传递在 ANE 上，dW 梯度在 CPU 上（Accelerate cblas）
  • Adam 优化器，梯度累积，通过 exec() 重启进行 checkpoint/恢复
  • 支持 GQA（分组查询注意力）的每头 tile/归约
  • 通过共享 IOSurface 实现 GPU↔ANE 零拷贝流水线（GPU 预填充 → ANE 解码）

INT8 W8A8 量化——1.88 倍吞吐量（M4，H16G）：

配置 FP16 INT8 W8A8 加速比
128x conv 512ch 64x64 18.6 TOPS, 14.8ms 35.1 TOPS, 7.8ms 1.88x
64x conv 512ch 64x64 18.4 TOPS, 7.5ms 34.1 TOPS, 4.0ms 1.85x

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

架构

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

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

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

GQA 模型（Qwen3-0.6B）——每层 10 个 kernel： 为分组查询注意力（Q_DIM ≠ DIM）添加独立的 woFwd、qBwd、kvBwd kernel。

CPU 处理：RMSNorm 前向/反向、残差连接（DeepNet α 缩放）、损失计算、dW 梯度累积（cblas_sgemm）、Adam 优化器更新。

关键优化：

  • Channel-first CPU 布局——匹配 ANE IOSurface [1,C,1,S] 格式，消除所有转置开销
  • vDSP 向量化 RMSNorm——比朴素实现快 10 倍（6.7ms → 0.7ms）
  • GCD 异步 cblas 重叠——dW 梯度 sgemm 与 ANE 求值在串行调度队列上并行运行
  • 延迟 cblas 等待——等待被推到下一步的前向传递中，实现最大重叠
  • ANE RMSNorm 融合——通过 MIL 操作（reduce_sum + pow + mul）将 RMSNorm 融合到前向 kernel 中
  • Wo^T 融合——输出投影反向合并到 SDPA 反向 kernel 中
  • 前向 taps——Q、K、V、注意力分数、隐藏状态通过 concat 输出暴露，避免 CPU 重新计算
  • exec() 重启——绕过每个进程约 119 次的 ANE 编译限制

文件结构

├── api_exploration.m           # 初始 ANE API 发现
├── inmem_basic.m               # 内存中 MIL 编译概念验证
├── inmem_bench.m               # ANE 分发延迟基准测试
├── inmem_peak.m                # 峰值 TFLOPS 测量（2048x2048 矩阵乘法）
├── ane_int8_bench.m            # INT8 W8A8 vs 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            # Bridge 实现（int8 + fp16 权重块）
│   └── Makefile
└── training/
    ├── ane_runtime.h           # ANE 私有 API 包装器（编译、求值、IOSurface）
    ├── ane_classifier.h        # 分类器前向（32K 卷积）、softmax、rmsnorm on ANE
    ├── train_large.m           # 静态流水线（权重作为常量，每次重新编译）
    ├── training_dynamic/
    │   ├── train.m             # 动态训练循环（模型无关）
    │   ├── config.h            # 派生尺寸、结构体、分配辅助函数
    │   ├── mil_dynamic.h       # 动态权重 kernel 的 MIL 生成器（GQA 感知）
    │   ├── io.h                # IOSurface I/O、权重 staging、GQA tile/归约
    │   ├── models/
    │   │   ├── stories110m.h   # Stories110M 配置（12层，MHA）
    │   │   └── qwen3_06b.h    # Qwen3-0.6B 配置（28层，GQA）
    │   └── Makefile
    ├── dashboard.py            # 实时训练仪表盘（blessed 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，109M 参数）
make MODEL=qwen3_06b      # Qwen3-0.6B（28层，GQA，596M 参数）
./train --scratch          # 从随机初始化训练
./train --resume           # 从 checkpoint 恢复

# 静态流水线（旧版——每次步骤重新编译权重）
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

# Bridge 库（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, channels, 1, spatial]（fp16 或 fp32；fp16 直接 I/O 快约 37%）
  4. 动态权重——激活和权重被打包到单个空间输入维度中，在 MIL kernel 内部切片。权重变化无需重新编译。
  5. 梯度流——前向 taps 暴露后向需要的中继结果；反向 kernel 在 ANE 上计算 dx（输入梯度）；dW（权重梯度）在 CPU 上通过 cblas 计算
  6. INT8 量化——constexpr_affine_dequantize 用于 int8 权重，层间使用 quantize/dequantize 进行 int8 激活缓存在 L2 SRAM 中（1.88 倍吞吐量）

限制

  • SDPA 因果掩码——ANE 硬件在 SDPA 操作中忽略 attn_mask；因果注意力被分解为独立的 Q@K^T（ANE）→ mask+softmax（CPU）→ scores@V（ANE）
  • 约 119 次编译限制——ANE 编译器泄漏资源；通过使用 checkpoint 的 exec() 重启解决
  • FP16 梯度下溢——反向矩阵乘法在 fp16 下下溢；使用全局损失缩放（256 * NLAYERS）修复
  • 单输入约束——多输入 ANE 请求导致 0x1d 错误；输入改为打包到空间维度

性能

训练吞吐量（M4）：

模型 参数量 毫秒/步 层数 Kernel 数/层
Stories110M 109M 91 ms 12 6（MHA）
Qwen3-0.6B 596M 412 ms 28 10（GQA）

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

精度 峰值 TOPS 配置
FP16 18.6 128x conv 512ch 64x64
INT8 W8A8 35.1 128x conv 512ch 64x64

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

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

免责声明

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

许可

MIT —— 见 LICENSE

由一个人类 + Claude 共同构建，每个周末一点一滴。

开源项目maderix2026-07-29原文

相关内容