The Problem
Training neural networks on the Apple Neural Engine (ANE) is a pain in the neck. Apple decided you can only use the ANE for inference, leaving developers to jump through hoops for training. It’s like having a powerful engine but only being allowed to drive in reverse. This repo tackles that by reverse-engineering private APIs to enable training directly on ANE hardware.
What This Does
The ANE repo gives you a from-scratch implementation for training transformer models using the ANE. It’s not just a half-baked idea; it actually utilizes the private ANEClient and ANECompiler APIs to run custom compute graphs. Check out the training/trainlarge.m file for the main training loop, where the magic happens. Each training step involves six ANE kernels, as defined in the architecture section of the README, allowing for forward and backward passes right on the chip.
The training/Makefile sets everything up for you, so you can compile the training program without dealing with unnecessary dependencies. You can run ./trainlarge and watch your model train on the ANE without the need for CoreML or Metal. Pretty neat, right?
Real-World Use
Imagine you’re building a transformer model for a natural language processing task and you want to leverage the ANE's compute power. You can follow the setup in training/README.md to get started. Just tweak train.m or trainlarge.m to fit your model architecture. You can also peek at the training/test*.m files for unit tests on individual kernels to ensure everything is working as expected.
Here’s a quick snippet to kick things off:
The Bottom Line
This repo is a solid solution for those who want to push the limits of Apple’s hardware but don’t expect corporate-level support. It’s a bit niche and probably overkill for small projects, but if you're deep into machine learning and own Apple Silicon, you’ll find this useful. Just be prepared to dig into some reverse-engineering if you want to modify anything.