Розробники зворотнього інжинірингу приватного API Apple Neural Engine вперше реалізували навчання нейронних мереж на ANE

Згідно з повідомленням CoinWorld, 3 березня розробник Manjeet Singh (GitHub: maderix) у співпраці з Claude Opus вперше реалізували тренування нейронної мережі з зворотнім поширенням на Apple Neural Engine (ANE) у чипі M4, використовуючи зворотне інженерне дослідження закритих приватних API Apple. ANE — це прискорювач, створений спеціально для інференції, офіційно не підтримує тренування, і розробники можуть викликати його функції лише опосередковано через CoreML. Проєкт обійшов CoreML, безпосередньо відобразивши понад 40 приватних класів, таких як _ANEClient, _ANECompiler, у повний програмний стек, що включає драйвер ядра IOKit, і виявив інтерфейс _ANEInMemoryModelDescriptor, який дозволяє компілювати моделі безпосередньо у пам’яті — ключовий момент для тренування, оскільки кожне оновлення ваг вимагає повторної компіляції. На даний момент реалізовано тренування одного шару трансформера (dim=768, seq=512), кожен крок займає 9.3 мс на M4, використання ANE становить 11.2% (1.78 TFLOPS, теоретична пікова продуктивність 15.8 TFLOPS). Вхідні градієнти для прямого і зворотного проходу обчислюються на ANE, градієнти ваг і оптимізатор Adam — на CPU. Проєкт також виявив, що основна обчислювальна примітивна операція ANE — це згортка, а не матричне множення; використання 1x1 згортки для вираження матричного множення дає приблизно 3-кратне підвищення пропускної здатності, а обхід CoreML для безпосереднього виклику дає додатковий приріст у 2-4 рази. Офіційна реклама Apple про «38 TOPS» є вводить в оману. Проєкт наразі знаходиться на ранній стадії: підтримує лише однорівневе тренування, використовує синтетичні дані, має близько 119 витоків ресурсів після компіляції, що вимагає перезапуску процесу для уникнення витоків, підтримка багаторівневого тренування та реальних даних ще розробляється. Проєкт відкритий за ліцензією MIT і за 5 днів зібрав близько 2800 зірок.