ANE Documentation

File: README.md Copy Chapter

ANE Training — Backpropagation on Apple Neural Engine

Training neural networks directly on Apple's Neural Engine (ANE) via reverse-engineered private APIs. No CoreML training APIs, no Metal, no GPU — pure ANE compute.

Project Scope & Intent Copy Chapter

I'm genuinely grateful for all the attention this project has received — I never expected a weekend research hack to blow up like this. Thank you to everyone who starred, forked, ran benchmarks on their own hardware, and shared the work. It means a lot.

That said, I want to set clear expectations about what this project is and isn't.

This is a research project, not a production framework.

The goal was to demonstrate that training on the Apple Neural Engine — and potentially other NPUs — is possible, and that the barrier has always been software support, not hardware capability. The ANE is a remarkably capable piece of silicon that Apple restricts to inference-only use through CoreML. This project bypasses that restriction using reverse-engineered private APIs to show what's possible when you give the hardware a chance.

What This Project Is Copy Section

• A proof of concept for ANE training via _ANEClient and _ANECompiler private APIs
• A set of benchmarks documenting real ANE performance characteristics (throughput, power, SRAM behavior)
• A reference for anyone exploring direct ANE access outside CoreML
• Research code that I update when I find something interesting

What This Project Is Not Copy Section

• A maintained framework or library
• A replacement for CoreML, MLX, llama.cpp, or any production inference stack
• A path to training large models on consumer hardware (yet)

On The Hype Copy Section

Some coverage of this project has overstated its implications. To be clear:

• Training works, but utilization is low (~5-9% of peak) with significant engineering challenges remaining
• Many element-wise operations still fall back to CPU
• This does not replace GPU training for anything beyond small research models today

The honest results — including all limitations — are documented in the accompanying articles:

• Part 1: Reverse Engineering
• Part 2: Benchmarks
• Part 3: Training

On Maintenance Copy Section

I don't intend to grow this into a large community project. My focus is on original research (compiler infrastructure for edge AI optimization), and maintaining an open-source framework takes time away from that.

That said:

• I'll keep pushing updates when I discover something interesting
• Bug fixes and benchmark contributions (especially on hardware I don't own) are welcome
• Feature requests will likely go unaddressed — but feel free to fork
• PRs will be merged at a relatively slow pace, otherwise I become the bottleneck for community growth around this tech

Fork it, build on it Copy Section

This is MIT licensed for a reason. Everyone now has access to AI-assisted development tools that can adapt and extend code in hours. If this project is useful to you — take it, modify it, build something better. If you do something cool with it, I'd love to hear about it.If in future, community decides to maintain one source of truth repo, I'm in full support of that.

What This Is Copy Chapter

A from-scratch implementation of transformer training (forward + backward pass) running on the ANE in Apple Silicon. The ANE is a 15.8 TFLOPS FP16 (M4) inference accelerator that Apple does not expose for training. This project reverse-engineers the _ANEClient / _ANECompiler private APIs and the MIL (Model Intermediate Language) format to run custom compute graphs — including backpropagation — directly on ANE hardware.

Current results:

• All forward and backward dx passes on ANE, dW gradients on CPU (Accelerate cblas)
• Adam optimizer, gradient accumulation, checkpoint/resume via exec() restart
• GQA (Grouped-Query Attention) support with per-head tiling/reduction
• GPU↔ANE zero-copy pipeline via shared IOSurface (GPU prefill → ANE decode)

INT8 W8A8 quantization — 1.88x throughput (M4, H16G):

INT8 activations halve L2 SRAM bandwidth between tiles via MIL quantize/dequantize ops. Weights use constexpr_affine_dequantize (int8 stored, fp16 at compile time).

Architecture Copy Chapter

The dynamic pipeline uses shared ANE kernels with weights packed into spatial dimensions (no recompilation when weights change):

MHA models (Stories110M) — 6 kernels per layer:

GQA models (Qwen3-0.6B) — 10 kernels per layer:
Adds separate woFwd, qBwd, kvBwd kernels for grouped-query attention (Q_DIM ≠ DIM).

CPU handles: RMSNorm forward/backward, residual connections (DeepNet α scaling), loss computation, dW gradient accumulation (cblas_sgemm), Adam optimizer updates.

Key optimizations:

• Channel-first CPU layout — matches ANE IOSurface [1,C,1,S] format, eliminates all transpose overhead
• vDSP vectorized RMSNorm — 10x faster than naive (6.7ms → 0.7ms)
• GCD async cblas overlap — dW gradient sgemms run in parallel with ANE evals on a serial dispatch queue
• Deferred cblas wait — wait pushed into next step's forward pass for maximum overlap
• ANE RMSNorm fusion — RMSNorm folded into forward kernels as MIL ops (reduce_sum + pow + mul)
• Wo^T fusion — output projection backward merged into SDPA backward kernel
• Forward taps — Q, K, V, attention scores, hidden states exposed via concat outputs, avoiding CPU recompute
• exec() restart — bypasses ~119 ANE compile limit per process

File Structure Copy Chapter

├── 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)
├── ane_int8_bench.m            # INT8 W8A8 vs FP16 throughput benchmark
├── sram_bench.m                # ANE SRAM bandwidth probing
├── sram_probe.m                # SRAM size/layout exploration
├── gpu_ane_share.m             # GPU↔ANE zero-copy IOSurface demo
├── gpu_prefill_ane_decode.m    # GPU prefill → ANE decode pipeline
├── bridge/
│   ├── ane_bridge.h            # C-callable ANE API (compile, eval, I/O)
│   ├── ane_bridge.m            # Bridge implementation (int8 + fp16 weight blobs)
│   └── Makefile
└── training/
    ├── ane_runtime.h           # ANE private API wrapper (compile, eval, IOSurface)
    ├── ane_classifier.h        # Classifier fwd (32K conv), softmax, rmsnorm on ANE
    ├── train_large.m           # Static pipeline (weights as constants, recompiles)
    ├── training_dynamic/
    │   ├── train.m             # Dynamic training loop (model-agnostic)
    │   ├── config.h            # Derived sizes, structs, alloc helpers
    │   ├── mil_dynamic.h       # MIL generators for dynamic weight kernels (GQA-aware)
    │   ├── io.h                # IOSurface I/O, weight staging, GQA tile/reduce
    │   ├── models/
    │   │   ├── stories110m.h   # Stories110M config (12L, MHA)
    │   │   └── qwen3_06b.h    # Qwen3-0.6B config (28L, GQA)
    │   └── Makefile
    ├── dashboard.py            # Live training dashboard (blessed TUI)
    └── Makefile

Training Data Copy Chapter

Training requires pretokenized TinyStories data. To download:

cd training && bash download_data.sh

See training/README.md for detailed training instructions.

Building Copy Chapter

Requires macOS 15+ on Apple Silicon (tested on M4).

# Dynamic pipeline (recommended) — model selected at build time
cd training/training_dynamic
make MODEL=stories110m    # Stories110M (12L, MHA, 109M params)
make MODEL=qwen3_06b      # Qwen3-0.6B (28L, GQA, 596M params)
./train --scratch          # train from random init
./train --resume           # resume from checkpoint

# Static pipeline (legacy — recompiles weights each step)
cd training && make train_large
./train_large ane_stories110M_ckpt.bin 256 100 1e-4

# INT8 benchmark
xcrun clang -O2 -fobjc-arc -framework Foundation -framework IOSurface -ldl \
  -o ane_int8_bench ane_int8_bench.m
./ane_int8_bench

# Bridge library (C-callable ANE API)
cd bridge && make

No external dependencies. Uses only system frameworks + private ANE APIs resolved at runtime via objc_msgSend.

How It Works Copy Chapter

1. MIL generation — Objective-C code constructs MIL program text at runtime, specifying convolutions (for linear layers), matmul (for attention), softmax, element-wise ops
2. In-memory compilation — _ANEInMemoryModelDescriptor compiles MIL text + weight blobs directly to ANE programs, no disk mlmodelc needed
3. IOSurface I/O — Input/output tensors passed via IOSurface shared memory in [1, channels, 1, spatial] format (fp16 or fp32; fp16 direct I/O is ~37% faster)
4. Dynamic weights — Activations and weights packed into a single spatial input dimension, sliced apart inside the MIL kernel. Weights change without recompilation.
5. Gradient flow — Forward taps expose intermediates needed for backward; backward kernels compute dx (input gradients) on ANE; dW (weight gradients) computed on CPU via cblas
6. INT8 quantization — constexpr_affine_dequantize for int8 weights, quantize/dequantize between layers for int8 activation caching in L2 SRAM (1.88x throughput)

Limitations Copy Chapter

• SDPA causal masking — ANE hardware ignores attn_mask in SDPA ops; causal attention is decomposed into separate Q@K^T (ANE) → mask+softmax (CPU) → scores@V (ANE)
• ~119 compile limit — ANE compiler leaks resources; worked around via exec() restart with checkpoint
• FP16 gradient underflow — backward matmuls underflow in fp16; fixed with global loss scaling (256 * NLAYERS)
• Single-input constraint — multi-input ANE requests cause 0x1d error; inputs packed into spatial dimension instead

Performance Copy Chapter

Training throughput (M4):

ANE peak throughput (M4, H16G):

GPU↔ANE inference pipeline (M4, seq=256):

Disclaimer Copy Chapter

This project uses Apple's private, undocumented APIs (_ANEClient, _ANECompiler, _ANEInMemoryModelDescriptor). These APIs are not covered by any public stability guarantee and may change or break with any macOS update. This is independent research into Apple Neural Engine architecture, using APIs discovered through runtime introspection for research and educational purposes under fair use and interoperability provisions (see Sega v. Accolade, 1992; DMCA §1201(f)). No Apple proprietary code or binaries are included in this repository. This project is not affiliated with or endorsed by Apple Inc. Use at your own risk.

License Copy Chapter

MIT — see LICENSE

Built by a human + Claude, one weekend at a time.

File: training/README.md Copy Chapter

ANE Training — On-Device Training on Apple Neural Engine

Training transformer models directly on Apple's Neural Engine using private ANE APIs. Supports multiple architectures including GQA (Grouped-Query Attention).

Supported Models Copy Chapter

Model configs live in training_dynamic/models/*.h. To add a new model, create a header with the architecture defines (see below).

Architecture Copy Chapter

• SDPA causal mask workaround: ANE hardware ignores attn_mask — decompose into Q@K^T (ANE conv) + mask+softmax (CPU) + scores@V (ANE conv)
• GQA support: K/V heads tiled to match Q heads for SDPA, reduced back after backward pass

Three Training Pipelines Copy Chapter

1. Static Baseline (`train_large`) Copy Section

Original pipeline. Weights baked as constants in MIL kernels — recompile every 10 steps via exec() restart.

• 60 weight-bearing + 12 weight-free kernels = 72 per compile batch
• Classifier + softmax + RMSNorm backward on CPU
• 106.7 ms/step, 7.6s compile per restart

2. Static + ANE Extras (`train_large_ane`) — PR#19 Copy Section

Offloads classifier forward (32K conv), softmax, final RMSNorm, and RMSNorm backward to ANE. Bridge API for C-callable ANE access.

• 86 kernels per compile batch (+24 rmsnorm_bwd, +1 classifier, +1 finalRms)
• 91.8 ms/step (14% faster), 9.6s compile per restart
• Use --no-ane-extras to disable and fall back to CPU (for debugging)

3. Dynamic Weight Pipeline (`training_dynamic/`) Copy Section

Weights passed via IOSurface spatial dimension — compile 10 kernels once at startup, no recompilation needed. Supports multiple models via make MODEL=xxx.

• 10 shared kernels across all layers (GQA-aware: split sdpaFwd/woFwd, split qBwd/kvBwd)
• ~115 ms/step (Stories110M) / ~412 ms/step (Qwen3-0.6B), 0.4s one-time compile
• No exec() restart, no compile limit issues

Performance Comparison (20 Steps) Copy Chapter

Key insights:

• Dynamic wins on wall time for any practical run length (3.9x faster at 20 steps)
• PR#19 has the best per-step throughput (92ms) but compile overhead dominates short runs
• Static restarts every 10 steps, so dynamic's zero-recompile advantage compounds

Files Copy Chapter

Usage Copy Chapter

1. Download Training Data Copy Section

bash download_data.sh

Downloads pretokenized TinyStories (Llama 2 BPE, 32K vocab) from HuggingFace. Produces tinystories_data00.bin (~41 MB, ~20M tokens).

2. Build & Train Copy Section

# Static baseline (classifier + softmax on CPU)
make train_large
./train_large stories110M.bin 256 100 1e-4
./train_large --model stories110M.bin --steps 100 --lr 1e-4
./train_large --data ./tinystories_data00.bin --steps 100 --lr 1e-4

# PR#19: ANE-offloaded classifier + softmax + rmsnorm_bwd
make train_large_ane
./train_large_ane stories110M.bin 256 100 1e-4
./train_large_ane --no-ane-extras --steps 100    # disable ANE extras
./train_large_ane --data ./tinystories_data00.bin --steps 100 --lr 1e-4

# Dynamic pipeline (model selected at build time)
cd training_dynamic
make MODEL=qwen3_06b           # default — Qwen3-0.6B (28L, GQA, 596M)
make MODEL=stories110m         # Stories110M (12L, MHA, 109M)
./train --scratch              # train from random init
./train --resume               # resume from checkpoint
./train --steps 200 --lr 1e-4  # custom steps/lr

CLI flags (train_large / train_large_ane):

• --steps N (default 10000)
• --lr F (default 3e-4)
• --model PATH — pretrained weights file
• --data PATH — tokenized TinyStories .bin file (default: tinystories_data00.bin)
• --ckpt PATH — checkpoint file (preserved across exec() restarts)
• --resume — resume from checkpoint
• --no-ane-extras — (train_large_ane only) disable ANE classifier/softmax/rmsnorm_bwd

3. Monitor with Dashboard Copy Section

pip install blessed psutil numpy
sudo python3 dashboard.py          # static pipeline
sudo python3 dashboard.py --dynamic # dynamic pipeline

4. Benchmarking Copy Section

All programs print an Efficiency Report at completion:

=== Efficiency Report ===
Total steps:     20
Wall time:       11738 ms (11.7 s)
Compile time:    9583 ms (81.6%)
Train time:      1835 ms (15.6%)
Avg train:       91.8 ms/step
ANE TFLOPS:      1.15 sustained

Adding a New Model Copy Chapter

Create training_dynamic/models/mymodel.h:

#pragma once
#define MODEL_NAME "MyModel-1B"

#define DIM 2048        // model hidden dim
#define HIDDEN 5504     // FFN intermediate dim
#define HEADS 32        // number of query heads
#define KV_HEADS 8      // number of KV heads (= HEADS for MHA)
#define HD 64           // head dim (can differ from DIM/HEADS)
#define SEQ 256         // sequence length
#define NLAYERS 22      // number of transformer layers
#define VOCAB 32000     // vocabulary size

#define CKPT_PATH "ane_mymodel_dyn_ckpt.bin"
#define DEFAULT_DATA_PATH "../tinystories_data00.bin"

Everything else is derived automatically: GQA_RATIO, Q_DIM, KV_DIM, weight sizes, IOSurface layouts, MIL kernels.

Build with: make MODEL=mymodel

Constraints:

• HEADS must be divisible by KV_HEADS
• HD is explicit (not necessarily DIM/HEADS — Qwen3 uses HD=128 with DIM/HEADS=64)
• For MHA (no GQA), set KV_HEADS = HEADS

Key Techniques Copy Chapter

• NEON vectorized fp16↔fp32: ARM NEON intrinsics for fast IOSurface data transfer
• vDSP cross-entropy: vDSP_mtrans + vvexpf + vDSP_sve — 8x faster than scalar
• Async weight gradients: cblas_sgemm dispatched to background queue, overlapped with ANE
• Vocab compaction (dynamic): 32K–152K → 9.2K active tokens, up to 16.5x reduction in classifier work
• Dynamic weight packing: Activations + weights concatenated in IOSurface spatial dimension — one kernel serves all layers
• GQA tile/reduce: K/V tiled from KV_HEADS→HEADS on CPU before SDPA backward, gradients reduced HEADS→KV_HEADS after
• exec() restart: Workaround for ANE ~119 compile limit per process