ANE — 在 Apple 神經網路引擎上訓練神經網路

GitHub: https://github.com/maderix/ANE

🔍 總覽

ANE 是一個研究專案，首次公開展示了可以在 Apple 的 神經網路引擎 (ANE) — 內建於 Apple Silicon 晶片（M1/M2/M4 系列）的專用 AI 加速器 — 上直接執行完整的 神經網路訓練（前向 + 後向傳播）。

Apple 透過 CoreML 框架，官方將 ANE 限制為 僅限推論 使用。此專案 逆向工程 Apple 的私有 API（_ANEClient、_ANECompiler、_ANEInMemoryModelDescriptor），以繞過此限制，並在 ANE 硬體上原生運行自訂計算圖，包括反向傳播。

這不是一個生產級框架。它是一個 概念驗證 和基準測試參考，證明此限制是軟體強加的，而非硬體限制。

🎯 重要性

Apple Silicon 晶片（尤其是 M4）包含一個額定功率為 15.8 TFLOPS 的 ANE — 大量的專用 AI 計算能力，Apple 將其鎖定為僅限推論。每次運行 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 grad, 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%
通道優先佈局	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 毫秒/步，維持 1.78 TFLOPS（M4 ANE 15.8 TFLOPS 峰值的 11.2%）。

🔑 主要優化

通道優先 CPU 佈局 — 匹配 ANE 的原生 [1,C,1,S] IOSurface 格式，消除了所有轉置開銷。
vDSP 向量化 RMSNorm — 比樸素實現快 10 倍（6.7ms → 0.7ms）。
GCD 非同步 cblas 重疊 — 權重梯度 sgemm 操作與 ANE 評估在序列分派隊列上並行運行。
延遲 cblas 等待 — dW sgemm 的等待被推遲到下一步的前向傳播，以實現最大的計算重疊。
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+ 上的 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 層	多層流水線調度尚未實現
僅限合成數據	真實標記數據支援正在進行中
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. 無關聯或未經其認可。