Bezpośrednie trenowanie sieci neuronowych na Apple Neural Engine bez CoreML i GPU

Każdy procesor Apple Silicon zawiera dedykowany blok operacji sieci neuronowych — Apple Neural Engine (ANE). Producent deklaruje dziesiątki teraflopsów wydajności, ale oficjalne narzędzia ograniczają jego użycie wyłącznie do wnioskowania poprzez CoreML. Trenowanie modeli oferowane jest albo na CPU, albo na karcie graficznej.

Deweloper o nicku maderix postanowił sprawdzić, czy sprzęt ANE jest fizycznie niezdolny do wstecznej propagacji, czy też problem wynika wyłącznie z ograniczeń programowych Apple. W ciągu jednego weekendu rozłożył na części prywatne frameworki macOS i zmusił ANE do wykonywania nie tylko przebiegu forward, ale także backward dla transformerów.

Projekt zdobył ponad 7 tysięcy gwiazdek na GitHubie. Przyjrzyjmy się, jak działa ten hack, jakie triki trzeba było zastosować i jak wygląda rzeczywista wydajność.

Dlaczego warto zajrzeć do zamkniętych frameworków

Oficjalny stos, jak CoreML, tworzy wrażenie, że ANE to zamknięta czarna skrzynka. Podajesz mu gotowy model, a on zwraca wynik. Jeśli chcesz trenować nawet niewielki model bezpośrednio na kliencie, musiałeś sięgnąć po Metal lub framework MLX, obciążając GPU.

Autor repozytorium ANE wykazał, że chip jest w pełni zdolny do wykonywania dowolnych grafów obliczeniowych. W tym celu zreverse-engineerował prywatne biblioteki _ANEClient i _ANECompiler, a także wewnętrzny język opisu modeli MIL (Model Intermediate Language).

Rezultat okazał się interesujący: pełnoprawne trenowanie transformerów bez ani jednej linijki CoreML czy Metal. Model jest składany bezpośrednio w RAM, kompilowany w locie i wysyłany do procesora neuronowego.

Ograniczenia i surowa rzeczywistość

Zanim rzucisz się do przepisywania skryptów trenujących na Maca, warto przyjrzeć się uczciwym liczbom. Sam autor ostrzega z góry: to eksperyment akademicki, nie biblioteka gotowa do produkcji.

Nie udało się jeszcze osiągnąć 100% wykorzystania zasobów ANE. Rzeczywiste wykorzystanie chipa to około 5–9% szczytowej wydajności. Ujawniają się ograniczenia programowe i sprzętowe:

• Część operacji matematycznych nie jest wspierana przez chip w wymaganej formie i wraca do CPU.
• Przebieg wsteczny dla wag (dW) nadal musi być obliczany przez procesor.
• Kompilator ANE ma wycieki pamięci, więc po około stu iteracjach potrzebne są obejścia.

Mimo wszystko nawet w tym trybie projekt oferuje przyzwoitą szybkość na podstawowych architekturach.

Jak działa pipeline trenowania

Architektura projektu opiera się na rozkładaniu zadań między ANE a CPU. Akcelerator sieci neuronowych obsługuje najcięższe mnożenia macierzy, podczas gdy procesor zajmuje się otaczającą logiką i akumulacją gradientów.

Przebieg forward i obliczanie gradientu wejściowego (dx) są w pełni wykonywane na ANE. Gradienty wag (dW) są obliczane przez CPU za pomocą zoptymalizowanych bibliotek Accelerate i cblas_sgemm. Optymalizator Adam i warstwa RMSNorm również działają na procesorze.

Aby uniknąć rekompilacji grafu modelu przy każdym kroku, gdy wagi się zmieniają, autor zastosował trik: wagi i aktywacje są pakowane w jeden tensor przez wymiary przestrzenne, a wewnątrz jądra MIL są one po prostu rozdzielane z powrotem.

Do wymiany danych między CPU a ANE wykorzystywana jest pamięć IOSurface. Umożliwia to transfer tensorów bez niepotrzebnego kopiowania między przestrzeniami adresowymi. Dane są pakowane w specyficzny dla ANE format [1, C, 1, S], gdzie kanały są na pierwszym miejscu. To podejście wyeliminowało narzut transpozycji macierzy.

Nieoczywiste problemy i obejścia

Podczas reverse engineeringu ujawniło się wiele pułapek sprzętu Apple.

Po pierwsze, operacja SDPA (Scaled Dot-Product Attention) w ANE ignoruje maskę kauzalną attn_mask na poziomie sprzętowym. Mechanizm attention musiał zostać podzielony na trzy etapy: mnożenie Q i K na ANE, maskowanie z softmax na CPU i końcowe mnożenie przez V z powrotem na ANE.

Po drugie, wbudowany kompilator _ANECompiler zawiera wyciek pamięci. Po około 119 kompilacjach proces ulega awarii z powodu wyczerpania zasobów. Autor rozwiązał problem radykalnie: gdy licznik zbliża się do limitu, program zapisuje checkpoint i wykonuje exec() — restartuje się z zachowaniem stanu.

Po trzecie, obliczenia FP16 podczas przebiegu wstecznego szybko prowadzą do underflow, powodując zamianę gradientów w zera. Problem został rozwiązany przez skalowanie lossa ze współczynnikiem 256 * NLAYERS.

Wydajność na M4

Na chipie Apple M4 wyniki okazały się dość wymowne. Testy przeprowadzono na dwóch architekturach:

Dla modelu Stories110M o 109 milionach parametrów (12 warstw, klasyczny Multi-Head Attention), czas jednego kroku trenowania wynosił 91 milisekund.

Większy Qwen3-0.6B z 596 milionami parametrów i Grouped-Query Attention przetwarza jeden krok w 412 milisekund.

Autor przetestował również kwantyzację INT8 W8A8. Użycie 8-bitowych wag i aktywacji zmniejsza obciążenie pamięci L2 SRAM chipa i zwiększa przepustowość z 18,6 TOPS do 35,1 TOPS na M4 — przyspieszenie prawie 1,88x w porównaniu z FP16.

Jak uruchomić projekt

Projekt nie wymaga zewnętrznych zależności jak PyTorch czy Conda. Wszystko, czego potrzebujesz, to najnowszy macOS 15 na maszynie z Apple Silicon i kompilator Clang. Prywatne API są pobierane w czasie działania poprzez objc_msgSend.

Aby zbudować dynamiczny pipeline, wystarczy przejść do folderu projektu i uruchomić polecenie make:

cd training/training_dynamic
make MODEL=stories110m
./train --scratch

Jeśli chcesz sprawdzić kwantyzację INT8 lub zmierzyć szczytowe TOPS swojego chipa, w folderze głównym repozytorium znajdują się osobne benchmarki.

Końcowe przemyślenia

Projekt maderix to przykład jakościowego badania „pod maską