ANE

The Problem

There is no Apple‑provided training API for the Neural Engine (ANE). The ANE is a 15.8 TFLOPS inference accelerator on Apple Silicon, but developers who want to run custom back‑propagation graphs must reverse‑engineer private _ANEClient/_ANECompiler APIs and the Model Intermediate Language (MIL) format. This repo fills that gap, offering a from‑scratch forward+backward pipeline that runs on the ANE while the CPU handles norms, residuals, loss and optimizer updates.

What This Does

The core training loop lives in training/train_large.m and pulls in the MIL generators from training/forward.h, training/backward.h and the ANE wrapper training/ane_runtime.h. Six ANE kernels are dispatched per step:

Kernel	Purpose	Key files
`kFwdAttn`	RMSNorm + QKV + SDPA + output projection	`forward.h`
`kFwdFFN`	RMSNorm + SwiGLU FFN	`forward.h`
`kFFNBwd`	FFN backward (W₂ᵀ + SiLU_bwd + W₁ᵀ + W₃ᵀ)	`backward.h`
`kSdpaBwd1`	Woᵀ + SDPA backward part 1	`backward.h`
`kSdpaBwd2`	SDPA backward part 2 (softmax grad, dQ, dK)	`backward.h`
`kQKVb`	QKV backward → dx	`backward.h`

The CPU (Accelerate cblas) accumulates dW gradients, performs Adam updates, and runs RMSNorm backward. IOSurface memory is laid out channel‑first [1,C,1,S] to eliminate transposes, and vDSP‑vectorised RMSNorm cuts the naive 6.7 ms to 0.7 ms. An exec() restart bypasses the ~119‑compile limit per process.

How It Is Wired

Entry point: training/train_large.m – main() parses hyper‑parameters, builds the MIL program via ANECompiler (resolved at runtime with objc_msgSend), and enters the step loop.
Build: Makefile in training/ invokes clang -O2 -framework Foundation -framework IOSurface -framework CoreML -framework Accelerate -ldl -lobjc -o train_large training/train_large.m.
Import graph: static analysis of 11 files (9 C/C++ header, 2 Python) shows 2 internal modules (training/dashboard, training/tokenize) with zero import edges and no circular dependencies.
External touches: the process links against Foundation, IOSurface, CoreML and Accelerate; it reads/write check‑point blobs and writes training/tiny_train.m logs. No network or database calls are made.
Call‑flow cost: the deepest nesting (5 levels) appears in training/backward.h, training/stories_io.h and training/dashboard.py; control‑flow is hard to follow without guard‑clause refactoring.

How To Use It

Setup (macOS 15+, Apple Silicon, tested on M4):

# Build the training binary
xcrun clang -O2 -framework Foundation -framework IOSurface \
  -framework CoreML -framework Accelerate -ldl -lobjc \
  -o train_large training/train_large.m

Run

./train_large

No environment variables or extra config files are required beyond the system frameworks. The README also lists optional per‑kernel test binaries (test_*.m) that can be compiled individually with the same clang flags.

Real‑World Use

A research team prototyping a transformer‑training pipeline on‑device can invoke ./train_large to obtain 9.3 ms/step throughput (M4, dim 768, seq 512) while keeping all heavy compute on the ANE. The resulting check‑point can be shipped to a GPU‑based production system for final fine‑tuning, or used to evaluate quantisation‑aware training tricks that need low‑precision ANE kernels.

Code Health & Issues

Static analysis (11 files) surface the following measured findings:

HIGH/cognitive_load – Deep nesting ×5 in training/backward.h, training/stories_io.h, training/dashboard.py; control flow is hard to follow.
MEDIUM/clarity – Duplicated 6‑line blocks (8 occurrences) across training/ane_runtime.h, training/stories_config.h, training/backward.h, training/forward.h; extract shared helpers.
MEDIUM/resilience – Broad except in training/dashboard.py swallows errors indiscriminately; catch specific exceptions.
MEDIUM/cognitive_load – training/dashboard.py is 760 lines; split into cohesive units by responsibility.
MEDIUM/cognitive_load – training/forward.h has 41 branch points over 144 lines; decompose decision‑heavy logic (strategy dispatch table).

SDLC observations (outside the measured block): no CI/CD pipeline (.github/ or CI config absent), no Dockerfile, licence file present, no lockfile, no committed secrets.

The Bottom Line

This repo delivers a working, low‑level ANE training pipeline where Apple does not, achieving ~1.8 TFLOPS sustained on a single transformer layer. The codebase is functional but suffers from deep nesting, duplicated logic and an oversized dashboard file that together raise maintenance risk. It is best suited for experimenters and hardware‑research teams comfortable with private APIs and willing to refactor the identified hot‑spots for long‑term stability.