Why this matters
Directly running training workloads on device NPUs has been largely blocked by vendor SDKs that treat NPUs as inference-only. This project shows that the hardware capability exists: by reverse-engineering private ANE APIs and the MIL format, you can compile and dispatch custom forward/backward kernels to the ANE and move significant pieces of the gradient pipeline off the CPU. That opens a new axis for edge/desktop research into low-power or private-model training, even if the current implementation remains a research prototype.
What Sets It Apart
-
Practical ANE backprop proof-of-concept: implements transformer forward and backward passes (dx on ANE, dW accumulation on CPU) and trains small Transformer configs end-to-end on Apple Silicon M4.
- So what: demonstrates feasibility of NPU training beyond vendor inference stacks instead of just a theoretical claim.
-
Engineering optimizations tuned to ANE: channel-first IOSurface layout, MIL kernel packing of weights into spatial dims (dynamic weights), INT8 W8A8 path that can nearly 1.9× throughput on microbenchmarks, and fused kernels (SDPA/wo/FFN) to reduce CPU work.
- So what: these are concrete performance levers for anyone trying to squeeze training through ANE constraints (SRAM, single-input format, compile limits).
-
Transparent limits and mitigations: documents FP16 underflow fixes (global loss scaling), a ~119-compile leak workaround (exec() restart), causal-mask handling (CPU mask+softmax), and measured utilization (~5–9% peak).
- So what: gives realistic expectations and actionable mitigations rather than optimistic marketing claims.
Who It's For and Tradeoffs
Great fit if you are an engineer or researcher exploring NPU-based training primitives, want detailed benchmarks of ANE behavior, or need a reference implementation for custom MIL kernels and IOSurface I/O on Apple Silicon. The repo contains dynamic and static training pipelines, example models (Stories110M, Qwen3-0.6B config), and reproducible benchmarks.
Look elsewhere if you need a production training stack, full GPU-scale training, or cross-platform support: utilization is low, many element-wise ops still fall back to CPU, it depends on undocumented private APIs (breaking changes possible), and it requires macOS on Apple Silicon (tested on M4).
Where It Fits
Positioned as a research toolkit and benchmark reference at the intersection of edge AI, NPU tooling, and transformer training. It informs efforts to unlock NPUs for training work—useful for people building compilers, runtime hacks, or specialized low-power training flows—rather than teams who need maintained production tooling.
