maderix/ane

ANE Training is a research project that demonstrates training neural networks directly on Apple's Neural Engine (ANE) using reverse-engineered private APIs (_ANEClient and _ANECompiler). It bypasses CoreML's inference-only restriction to enable pure ANE compute for backpropagation.

Key Capabilities:

• Transformer training (forward + backward pass) on Apple Silicon.
• Support for MHA (Multi-Head Attention) and GQA (Grouped-Query Attention) models.
• INT8 W8A8 quantization for increased throughput.
• GPU↔ANE zero-copy pipeline via shared IOSurface.

Important Limitations:

• It is a research proof-of-concept, not a production framework.
• Utilization is currently low (~5-9% of peak).
• Many element-wise operations fall back to the CPU.
• SDPA causal masking is handled via CPU decomposition because ANE hardware ignores attn_mask in SDPA ops.

Use the TUI dashboard to monitor the loss curve and power/CPU/memory graphs during training. You must run with sudo to access system metrics.

pip install blessed psutil numpy

# For static pipelines
sudo python3 dashboard.py

# For the dynamic pipeline
sudo python3 dashboard.py --dynamic

The dynamic pipeline is the recommended way to train. It uses shared ANE kernels with weights packed into spatial dimensions to avoid recompilation when weights change. You must select a model at build time.

Steps:

1. Navigate to training/training_dynamic.
2. Build the specific model using make MODEL=<model_name>.
3. Run the training executable with either --scratch (random initialization) or --resume (from a checkpoint).

# Dynamic pipeline (recommended)
cd training/training_dynamic
make MODEL=stories110m    # Stories110M (12L, MHA, 109M params)
# OR
make MODEL=qwen3_06b      # Qwen3-0.6B (28L, GQA, 596M params)

./train --scratch          # train from random init
./train --resume           # resume from checkpoint

The Apple Neural Engine (ANE) handles Model Intermediate Language (MIL) weight formats differently depending on the chip generation:

• M1 Generation: Does not support single-blob weight formats. You must use stories_mil.h (separate per-matrix weight blobs) instead of ane_mil_gen.h.
• M3, M4, and M5 Generations: Support both single-blob weights and per-matrix weight blobs. They also support BLOBFILE offset references.
• M3 Pro Specifics: Note that the M3 Pro has a fixed 512-wide lane structure in SRAM tiling; only ch=512 is known to compile successfully.

Training requires pretokenized TinyStories data. You can download it by running the following command from the training directory:

cd training && bash download_data.sh

To prepare for training, download the pretokenized TinyStories dataset (Llama 2 BPE, 32K vocab) from HuggingFace. This produces tinystories_data00.bin (~41 MB, ~20M tokens).

bash download_data.sh

The bridge directory contains a C-callable ANE API for compilation, evaluation, and I/O.

cd bridge && make

The Static Baseline (train_large) uses weights baked as constants in MIL kernels and requires a recompile via exec() every 10 steps. Use this for the original baseline implementation where the classifier and softmax run on the CPU.

# Build
make train_large

# Run with positional arguments (model, batch, steps, lr)
./train_large stories110M.bin 256 100 1e-4

# Run with CLI flags
./train_large --model stories110M.bin --steps 100 --lr 1e-4
./train_large --data ./tinystories_data00.bin --steps 100 --lr 1e-4

To measure INT8 W8A8 vs FP16 throughput on the ANE, use the following command:

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

The train_large_ane pipeline (PR#19) offloads the classifier forward, softmax, final RMSNorm, and RMSNorm backward to the ANE for improved throughput. You can disable these extras using the --no-ane-extras flag for debugging.

# Build
make train_large_ane

# Run with positional arguments
./train_large_ane stories110M.bin 256 100 1e-4

# Run with CLI flags
./train_large_ane --no-ane-extras --steps 100
./train_large_ane --data ./tinystories_data00.bin --steps 100 --lr 1e-4

The Dynamic Weight Pipeline (training_dynamic/) is the most efficient for long runs. It passes weights via IOSurface, requiring only a single compile at startup. The model is selected at build time using the MODEL variable.

cd training_dynamic

# Build for Qwen3-0.6B (default)
make MODEL=qwen3_06b

# Build for Stories110M
make MODEL=stories110m

# Training commands
./train --scratch              # train from random init
./train --resume               # resume from checkpoint
./train --steps 200 --lr 1e-4  # custom steps/lr

To contribute benchmark data, navigate to the training directory and use the make train_large command. You must provide the checkpoint file, sequence length, number of steps, and learning rate as arguments. When reporting results to the project's issue tracker, include your chip model, macOS version, and the full output containing JSON lines.

cd training && make train_large
./train_large ane_stories110M_ckpt.bin 256 20 1e-4