Training trên Apple Neural Engine qua private API

Dự án ANE ↗ của maderix ↗ là một research proof-of-concept đột phá: training neural network trực tiếp trên Apple Neural Engine (ANE) qua private API đã được reverse-engineer, hoàn toàn không dùng CoreML, Metal hay GPU.

ANE trên chip M4 đạt 15.8 TFLOPS FP16 (18.6 TOPS thực tế đo được), nhưng Apple chỉ cho phép dùng nó qua CoreML với giới hạn inference. Dự án này phá vỡ rào cản đó, cho thấy vấn đề nằm ở software support, không phải hardware capability.

Kiến trúc phần mềm#

CoreML không phải đường vào duy nhất. Class _ANEClient trong AppleNeuralEngine.framework cho phép truy cập trực tiếp pipeline compile → load → evaluate:

id client = [_ANEClient sharedConnection];
id model = [_ANEModel modelAtURL:compiledURL key:@"mykey"];

[client compileModel:model options:@{
    @"kANEFModelType": @"kANEFModelMIL",
    @"kANEFNetPlistFilenameKey": @"model.mil"
} qos:21 error:&err];

[client loadModel:model options:@{} qos:21 error:&err];

Thay vì compile ra file .mlmodelc rồi nạp lên ANE, _ANEInMemoryModelDescriptor cho phép compile MIL text trực tiếp trong memory — không cần disk round-trip. Đây là chìa khóa để training.

ANE dùng IOSurface để I/O — cơ chế shared memory tương tự GPU texture, mở đường cho zero-copy pipeline GPU↔ANE.

MIL — Ngôn ngữ intermediate của ANE#

ANE không nhận ONNX hay protobuf. Nó dùng MIL (Model Intermediate Language) — SSA representation với type và shape tường minh:

program(1.3) {
    func main<ios18>(
        tensor<fp16, [1, 1024, 1, 1024]> x,
        tensor<fp16, [1, 1024, 1, 1024]> w
    ) {
        tensor<fp16, [1, 1024, 1, 1024]> out =
            matmul(transpose_x = false, transpose_y = false,
                   x = x, y = w);
    } -> (out);
}

Tensor layout theo chuẩn NCDHW + Interleave: [Batch, Channels, Depth, Height, Width]. Ma trận 1024×1024 thành [1, 1024, 1, 1024] trong 4D.

Kiến trúc training#

Dynamic pipeline dùng shared ANE kernels với weights packed vào spatial dimension — không cần recompile khi weights thay đổi.

MHA models (Stories110M) — 6 kernels mỗi layer:

GQA models (Qwen3-0.6B) — 10 kernels mỗi layer với separate woFwd, qBwd, kvBwd cho grouped-query attention.

CPU đảm nhận: RMSNorm forward/backward, residual connections (DeepNet α scaling), loss computation, dW gradient accumulation (cblas_sgemm), Adam optimizer.

Kết quả hiệu năng#

Training throughput (M4):

INT8 W8A8 quantization — 1.88x speedup:

INT8 activations giảm một nửa L2 SRAM bandwidth giữa các tile nhờ quantize/dequantize.

Kỹ thuật tối ưu#

• Channel-first CPU layout — khớp format IOSurface [1,C,1,S], loại bỏ hoàn toàn transpose overhead
• vDSP vectorized RMSNorm — nhanh gấp 10 lần (6.7ms → 0.7ms)
• GCD async cblas overlap — dW gradient sgemms chạy song song với ANE evals
• Deferred cblas wait — đẩy wait vào forward pass của step sau
• ANE RMSNorm fusion — RMSNorm gộp vào forward kernels qua MIL ops
• Forward taps — Q, K, V, attention scores exposed qua concat outputs, tránh CPU recompute
• exec() restart — vượt giới hạn ~119 ANE compile mỗi process

Hạn chế#

• SDPA causal masking — ANE hardware bỏ qua attn_mask; causal attention phải decompose thành Q@K^T (ANE) → mask+softmax (CPU) → scores@V (ANE)
• ~119 compile limit — compiler leak resources; workaround bằng exec() restart với checkpoint
• FP16 gradient underflow — backward matmuls underflow fp16; fix bằng global loss scaling
• Utilization thấp — chỉ ~5-9% peak, nhiều element-wise ops vẫn fallback về CPU

Tài liệu tham khảo#

• GitHub — maderix/ANE ↗
• Inside the M4 Apple Neural Engine, Part 1: Reverse Engineering ↗
• Inside the M4 Apple Neural Engine, Part 2: Benchmarks ↗
• Inside the M4 Apple Neural Engine, Part 3: Training ↗
• hollance/neural-engine ↗ — Tài liệu cộng đồng về ANE
• apple/ml-ane-transformers ↗ — Reference transformers của Apple cho ANE