Разработчику удалось запустить обучения нейросетей на NPU от Apple, который доступен только для инференса

Разработчик Манжит Сингх (Manjeet Singh) опубликовал ^[1] на GitHub проект ANE Training — это способ запускать обучение ^[2] нейросетей на NPU от Apple. Официально компания позволяет использовать NPU в своих устройствах только для запуска моделей.

Neural Engine — нейропроцессор, встроенный в чипы Apple серий A и M. Он обрабатывает ML-задачи на устройствах компании: распознаёт пользователя с помощью Face ID, работает с командами Siri, улучшает фотографии и запускает функции Apple Intelligence. С помощью фреймворка Core ML сторонние разработчики могут запускать на Neural Engine собственные предобученные модели.

Сингх отмечает, что он провёл реверс-инжиниринг приватных API Apple (_ANEClient, _ANECompiler), чтобы запускать на нейропроцессоре обучения собственных моделей. Итоговый проект оказался рабочим, но загрузка ускорителя остаётся низкой — до 9% от пика. Часть вычислений по-прежнему переходит на CPU.

Автор проекта провёл эксперимент, в рамках которого обучил на Neural Engine небольшую модель Stories110M на 110 млн параметров. Архитектурно она похожа на Llama 2. Один шаг обучения на M3 Ultra занял ~0,09 секунды, а на M4 — ~0,11 секунды.

Сингх отмечает, что у способа пока есть множество проблем и ограничений. Например, у Neural Engine есть лимит примерно в 119 компиляций на процесс. Из-за этого обучение приходится перезапускать с чекпойнта. Кроме того, больше всего времени занимает не само обучение, а подготовка к нему. У способа есть и риски: всё построено на базе недокументированных функций API. Apple не собирается их поддерживать и в любом обновлении может закрыть доступ.

Автор отмечает, что не планирует развивать проект и просто делится своим исследованием. В репозитории на GitHub ^[1] доступен код и инструкции по запуску, а в блоге разработчика есть подробные статьи по реверс-инжинирингу и тестированию Neural Engine (1 ^[3], 2 ^[4]).