maderix/ANE

# 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 或任何生产推理堆栈的替代品 - （目前）在消费级硬件上训练大模型的途径 ### 关于炒作 一些关于该项目的报道夸大了其影响。明确来说： - 训练是可行的，但利用率很低（峰值的 ~2-3%），仍存在重大工程挑战 - 许多逐元素运算仍然回退到 CPU - 这**并不**取代目前用于小型研究模型以外的任何 GPU 训练 诚实的结果 —— 包括所有限制 —— 记录在随附的文章中： - [第一部分：逆向工程](https://maderix.substack.com/p/inside-the-m4-apple-neural-engine) - [第二部分：基准测试](https://maderix.substack.com/p/inside-the-m4-apple-neural-engine-615) ### 关于维护 我不打算将其发展为一个大型社区项目。我的重点是原创研究（用于边缘 AI 优化的编译器基础设施），而维护一个开源框架会占用这些时间。 话虽如此： - 当我发现有趣的内容时，我会继续推送更新 - 欢迎错误修复和基准测试贡献（尤其是在我没有的硬件上） - 功能请求很可能无法得到处理 —— 但请随意 Fork - PR 将以相对较慢的速度合并，否则我会成为这项技术社区增长的瓶颈 ### Fork 它，基于它构建 它是 MIT 许可的，这是有原因的。每个人现在都可以使用 AI 辅助开发工具，在几小时内调整和扩展代码。如果这个项目对你有用 —— 拿去，修改它，构建更好的东西。如果你用它做了很酷的事情，我很乐意听听。如果未来社区决定维护一个权威源仓库，我完全支持。 ## 这是什么 在 Apple Silicon 的 ANE 上从头实现的 transformer 训练（前向 + 反向传播）实现。ANE 是一个 15.8 TFLOPS (M4) 的推理加速器，Apple 不将其暴露用于训练。该项目逆向工程了 `_ANEClient` / `_ANECompiler` 私有 API 和 MIL (Model Intermediate Language) 格式，以直接在 ANE 硬件上运行自定义计算图 —— 包括反向传播。 **当前结果 (M4，单 transformer 层，dim=768，seq=512)：** - 9.3 毫秒/步，11.2% ANE 利用率（持续 1.78 TFLOPS） - 每个训练步骤 6 次 ANE kernel 调度 - 所有前向和反向 dx 传递在 ANE 上，dW 梯度在 CPU 上（Accelerate cblas） - Adam 优化器，梯度累积，检查点/恢复 ## 架构 训练循环每步使用 6 个 ANE kernel： | Kernel | 功能 | 权重 | |--------|----------|---------| | `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 梯度, dQ, dK) | — | | `kQKVb` | QKV 反向 (Wq^T + Wk^T + Wv^T → dx) | Wq^T, Wk^T, Wv^T | CPU 负责：RMSNorm 反向、残差连接、loss 计算、dW 梯度累积 (cblas_sgemm)、Adam 优化器更新。 关键优化： - **通道优先 CPU 布局** —— 匹配 ANE IOSurface `[1,C,1,S]` 格式，消除所有转置开销 - **vDSP 向量化 RMSNorm** —— 比原生快 10 倍 (6.7ms → 0.7ms) - **GCD 异步 cblas 重叠** —— dW 梯度 sgemm 与 ANE eval 在串行调度队列上并行运行 - **延迟 cblas 等待** —— 等待推迟到下一步的前向传递以实现最大重叠 - **ANE RMSNorm 融合** —— RMSNorm 作为 MIL 算子折叠进前向 kernel (reduce_sum + pow + mul) - **Wo^T 融合** —— 输出投影反向合并入 SDPA 反向 kernel - **前向 taps** —— Q、K、V、注意力分数、隐藏状态通过 concat 输出暴露，避免 CPU 重计算 - **exec() 重启** —— 绕过每个进程约 119 次 ANE 编译限制 ## 文件结构 ``` ├── api_exploration.m # Initial ANE API discovery ├── inmem_basic.m # In-memory MIL compilation proof-of-concept ├── inmem_bench.m # ANE dispatch latency benchmarks ├── inmem_peak.m # Peak TFLOPS measurement (2048x2048 matmul) ├── sram_bench.m # ANE SRAM bandwidth probing ├── sram_probe.m # SRAM size/layout exploration └── training/ ├── ane_runtime.h # ANE private API wrapper (compile, eval, IOSurface) ├── ane_mil_gen.h # MIL program generation helpers ├── model.h # Model weight initialization and blob builders ├── forward.h # Forward pass MIL generators ├── backward.h # Backward pass MIL generators ├── train.m # Minimal training loop (early prototype) ├── tiny_train.m # 2-layer tiny model training ├── train_large.m # Main: single-layer dim=768 training (optimized) ├── test_*.m # Unit tests for individual kernels └── Makefile ``` ## 构建 需要 Apple Silicon 上的 macOS 15+（在 M4 上测试）。 ``` # 构建主 training 程序 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。 ## 工作原理 1. **MIL 生成** —— Objective-C 代码在运行时构建 MIL 程序文本，指定卷积（用于线性层）、matmul（用于注意力）、softmax、逐元素运算 2. **内存编译** —— `_ANEInMemoryModelDescriptor` 将 MIL 文本 + 权重 blob 直接编译为 ANE 程序，无需磁盘 mlmodelc 3. **IOSurface I/O** —— 输入/输出张量通过 IOSurface 共享内存以 `[1, channels, 1, spatial]` 格式 (fp16) 传递 4. **权重嵌入** —— 权重作为 BLOBFILE 常量嵌入 ANE 程序；当权重更改时每批次重新编译 5. **梯度流** —— 前向 taps 暴露反向传播所需的中间值；反向 kernel 在 ANE 上计算 dx（输入梯度）；dW（权重梯度）通过 cblas 在 CPU 上计算 ## 限制 - **SDPA 因果掩码** —— ANE 硬件在 SDPA 算子中忽略 `attn_mask`；因果注意力被分解为独立的 Q@K^T (ANE) → mask+softmax (ANE 通过 add+softmax) → scores@V (ANE) - **~119 次编译限制** —— ANE 编译器泄漏资源；通过带检查点的 `exec()` 重启解决 - **单层** —— 目前训练单个 transformer 层；多层需要流水线调度 - **合成数据** —— 目前使用随机数据进行基准测试；真实 tokenized 数据支持正在进行中 ## 性能历史 | 优化 | ms/step | ANE 利用率 | |---|---|---| | 基线 (vDSP 转置) | 33.5 | 3.1% | | 通道优先布局 | 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 kernels) | 11.4 | 9.2% | | 延迟 cblas 等待 | **9.3** | **11.2%** | ## 免责声明 该项目使用 Apple 私有的、未文档化的 API（`_ANEClient`、`_ANECompiler`、`_ANEInMemoryModelDescriptor`）。这些 API 不受任何公开稳定性保证的约束，并且可能随任何 macOS 更改或中断。这是对 Apple Neural Engine 架构的独立研究，使用通过运行时内省发现的 API，用于研究和教育目的，依据合理使用和互操作性条款（参见 *Sega v. Accolade*, 1992; DMCA §1201(f)）。本仓库不包含任何 Apple 专有代码或二进制文件。该项目未隶属于或受 Apple Inc. 认可。使用风险自负。 ## 许可证 MIT —— 见 [LICENSE](LICENSE) *由人类 + Claude 构建，一个周末接一个周末。*

标签：ANE, Apex, Apple Neural Engine, CoreML, CVE监控, iOS安全, NPU, Proof of Concept, 云资产清单, 反向传播, 底层开发, 性能基准测试, 机器学习, 深度学习, 硬件加速, 神经网络训练, 私有API, 系统架构, 计算机科学, 逆向工程, 配置审计