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

Согласно новостям BlockBeats, 3 марта разработчик Манджит Сингх (GitHub: maderix) совместно с Клодом Опусом проработал обратное проектирование нераскрытых приватных API Apple для внедрения обучения нейронных сетей с обратной распространенией на Apple Neural Engine (ANE) чипа M4 впервые. ANE — это ускоритель проектирования выводов от Apple, и официальная возможность обучения никогда не была открыта, а разработчики могут вызывать её функцию вывода только косвенно через фреймворк CoreML.

Проект обошёл CoreML и напрямую сопоставил с полным стеком программного обеспечения из более чем 40 приватных классов, таких как _ANEClient и _ANECompiler, в ядро IOKit, при этом обнаружил интерфейс _ANEInMemoryModelDescriptor, который компилирует модель непосредственно в память — ключ к обучению, поскольку каждое обновление веса требует перекомпилирования. В настоящее время обучен один слой трансформатора (dim=768, seq=512), с 9,3 мс на шаг на M4, 11,2% использованием ANE (1,78 TFLOPS, теоретический пик 15,8 TFLOPS), на ANE рассчитываются градиенты вперёд и обратного распространения, а на CPU выполняются градиенты веса и оптимизаторы Адама.

Проект также обнаружил, что основные вычислительные примитивы ANE — это свёртки, а не умножение матриц, и выражение матричного умножения с помощью 1x1 свёрток может обеспечить примерно в 3 раза больше пропускной способности, а обход прямых вызовов CoreML даёт дополнительное усиление в 2-4 раза, а официальная реклама Apple «38 TOPS» вводит в заблуждение. Проект всё ещё находится на ранней стадии: используется только одноуровневое обучение, синтетические данные, около 119 скомпилированных утечек ресурсов, которые необходимо перезапустить, чтобы избежать процесса, а многоуровневое обучение и поддержка реальных данных всё ещё находятся в стадии разработки. Проект является открытым исходным кодом по протоколу MIT и получил около 2800 звёзд за 5 дней после релиза.