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

BlockBeats повідомляє, 3 березня розробник Manjeet Singh (GitHub: maderix) у співпраці з Claude Opus здійснили зворотне інженерне дослідження приватних API Apple, які раніше не були опубліковані, і вперше на апаратній платформі Apple Neural Engine (ANE) на чипі M4 реалізували тренування нейронних мереж з зворотнім поширенням помилки. 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 згортки для представлення матричного множення дає приблизно трьохкратне підвищення пропускної здатності. Обхід CoreML і прямий виклик дають додатковий приріст у 2-4 рази. Офіційна реклама Apple про «38 TOPS» є вводом у оману. На даний час проект знаходиться на ранній стадії: підтримується лише однорівневе тренування, використовується синтетичні дані, є близько 119 витоків пам’яті після компіляції, що вимагає перезапуску процесу, підтримка багаторівневого тренування та реальних даних ще розробляється. Проект опубліковано під ліцензією MIT і за 5 днів зібрав близько 2800 зірок.