Orion: Обход CoreML для запуска и обучения LLM напрямую на Apple Neural Engine

Прямой доступ к ANE для работы с LLM

Orion предоставляет сквозную систему на Objective-C, которая полностью обходит CoreML для запуска и обучения LLM непосредственно на Apple Neural Engine (ANE). Этот подход даёт разработчикам прямой контроль над ANE, который ранее CoreML обрабатывал как «чёрный ящик»-планировщик, лишая любого прямого управления или возможности обучения.

Техническая реализация и ограничения

Проект основан на работе по реверс-инжинирингу, в ходе которой были сопоставлены приватные API ANEClient и ANECompiler. ANE представляет собой то, что разработчик называет «аппаратным несоответствием импеданса» с 17 общими программными ограничениями, 11 из которых были полностью недокументированы. Ключевые ограничения включают:

• Операция concat вызывает немедленный, тихий сбой компилятора
• Веса BLOBFILE требуют смещения в 64 байта от заголовка чанка, иначе происходит тихое числовое искажение
• ANE поддерживает внутреннее состояние, которое жёстко ограничивается примерно 119 компиляциями на процесс перед тихим сбоем

Решения проблем обучения

Предыдущие попытки обучения на ANE сталкивались с расхождением NaN после одного шага. Orion решает это с помощью:

• Настройки конвейера отложенной компиляции
• Реализации строгого ограничения активаций для предотвращения каскадного переполнения fp16 (ограничение активаций от -65504 до +65504)
• Использования цикла перезапуска процесса exec() после каждого шага обучения для обхода ограничения в 119 компиляций

Результаты производительности

Компилятор преобразует граф IR из 27 операций через пять этапов оптимизации в нативный для ANE MIL. Текущая производительность включает:

• Более 170 токенов/с для декодирования GPT-2 124M
• Механически стабильное многошаговое обучение трансформера с 110 млн параметров («потолок связности» аппаратного обеспечения)
• За более чем 1000 шагов потери снизились с 12,3 до 6,2 без единого NaN

Текущие ограничения

ANE запекает веса во время компиляции, что означает, что каждое обновление обучения требует штрафа за перекомпиляцию около 4,2 с. ANE выдаёт примерно 19 TFLOPS в fp16, но фундаментальным ограничением для его использования была не вычислительная мощность, а полное отсутствие нативного уровня оркестрации.