革命性突破：直接在 Apple Neural Engine 上训练 Transformer

Apple 的 Neural Engine (ANE) 在 M4 芯片中提供 15.8 TFLOPS 的推理性能，但训练呢？官方声称不可能。直到现在。

突破：纯 ANE 训练

ANE Training 是一个从零实现的完整 Transformer 训练循环——前向和反向传播——直接在 ANE 硬件上运行，使用逆向工程的私有 API。没有 CoreML 训练 API。没有 Metal 着色器。没有 GPU 回退。纯 ANE 计算。

当前基准测试 (M4, dim=768, seq=512)： - 9.3 毫秒/步 - 11.2% ANE 利用率 (1.78 TFLOPS 持续) - 每个训练步 6 个 ANE 内核调度

架构深入解析

系统在每个训练步协调 6 个专用 ANE 内核：

CPU 仅处理： RMSNorm 反向、残差、损失、dW 累积 (Accelerate cblas)、Adam 更新。

关键优化措施

1. 通道优先布局 - 匹配 ANE IOSurface [1,C,1,S]，消除转置开销
2. vDSP RMSNorm - 10 倍加速 (6.7ms → 0.7ms)
3. ANE RMSNorm 融合 - 内置于前向内核
4. 延迟 cblas - 最大化 ANE/CPU 重叠
5. exec() 重启 - 绕过 119 编译限制

性能演进： 通过系统优化从 33.5ms → 9.3ms。

几分钟内上手

# 需要 macOS 15+ Apple Silicon
xcrun clang -O2 -framework Foundation -framework IOSurface \
-framework CoreML -framework Accelerate -ldl -lobjc \
-o train_large training/train_large.m

./train_large

零依赖，仅需系统框架。

文件结构亮点

• api_exploration.m - 私有 API 发现
• inmem_bench.m - ANE 调度延迟
• sram_probe.m - SRAM 带宽探索
• training/train_large.m - 生产级单层训练器

限制与路线图

✅ 通过分解实现因果掩码 ✅ 梯度累积/检查点 ✅ Adam 优化器

🔄 多层流水线 🔄 真实分词数据集 🔄 完整模型训练

法律声明

使用运行时自省未文档化 API 用于研究/教育目的 (DMCA §1201(f))。不包含 Apple 专有代码。

2.1k 星标，362 分叉 - 加入 Apple Silicon ML 革命：https://github.com/maderix/ANE