Энтузиасты нашли способ запускать обучение ИИ-моделей на Neural Engine в Apple M4

Neural Engine оказался способен не только на инференс

Обычно Neural Engine в процессорах Apple M4 используется исключительно для инференса — выполнения уже обученных моделей ИИ. Именно для таких задач Apple и позиционирует этот блок, заявляя производительность до 38 TOPS.

Однако исследователь безопасности и разработчик под ником @0x0SojalSec сообщил в X, что ему удалось провести обратный инжиниринг существующих ограничений. В результате команда смогла задействовать до 15,8 TFLOPS вычислительной мощности Neural Engine для полноценных задач обучения, включая обратное распространение ошибки (backpropagation) в трансформерных моделях.

Как удалось обойти ограничения Apple

Вместо использования стандартных фреймворков Apple, таких как Core ML и Metal, разработчики создали собственный промежуточный язык моделей (Model Intermediate Language, MIL), который позволяет напрямую взаимодействовать с Neural Engine.

Для повышения производительности все данные хранятся в оперативной памяти, что избавляет систему от медленных операций записи на накопитель. Благодаря этому процесс обучения оказался не только быстрым, но и достаточно стабильным.

Кроме того, авторы проекта нашли оригинальный способ обхода зависаний во время обучения. При возникновении проблем используется команда exec(), которая перезапускает процесс, позволяя сохранить промежуточное состояние модели и продолжить обучение с контрольной точки.

Среди ключевых особенностей проекта:

• прямой доступ к Neural Engine без Core ML и Metal;
• использование собственного языка MIL для работы с ANE;
• хранение данных исключительно в оперативной памяти;
• поддержка обучения трансформерных моделей с backpropagation;
• механизм восстановления процесса через контрольные точки.

Первые результаты выглядят впечатляюще

Разработчик уже опубликовал проект на GitHub, продемонстрировав возможность обучения нейронных сетей непосредственно на аппаратуре Neural Engine с минимальным количеством дополнительных зависимостей.

Первые тесты показали многообещающие результаты. В частности, отдельные шаги обучения трансформерных моделей на чипах M4 выполняются всего за несколько миллисекунд.

Хотя Apple официально не поддерживает обучение моделей на Neural Engine и позиционирует его исключительно как ускоритель инференса, эксперимент показал, что потенциал этого блока значительно выше, чем считалось ранее.