Инженерное исследование: перенос трансформера Laya на ANE

Среда исследования: Apple M3 Max (40 ядер GPU, 128 ГиБ объединённой памяти), macOS 27.2, Core ML Tools 9.0, PyTorch 2.7.0 и NumPy 2.1.3. Это независимый прототип в experiments/ane_engineering/; выпущенный рантайм не изменён.

Текущий результат

Многоязычный прототип с фиксированными B=1, L=96 успешно назначает полный энкодер, голову решений и скоринг на Neural Engine в ожидаемом плане вычислений Core ML: 6 390 неконстантных операций предпочитают ANE, с оценочными весами стоимости, суммирующимися примерно в 1. Остальные 3 809 записей — константы. Исходный экспорт с перечислимыми формами/SDPA предпочитал CPU для всех 1 318 назначенных операций при CPU_AND_NE, несмотря на то что 988 отдельных операций перечисляют ANE как поддерживаемое устройство.

Полный путь предсказания прототипа измерил 5.167 мс p50 на 50 отборочных вызовах, включая токенизацию, поиск эмбеддингов, маски внимания, инференс ANE, вычисление головы действий на CPU, калибровку и форматирование. Изолированное тело Core ML измерило 4.403 мс p50 на 30 вызовах с синтетическими эмбеддингами. Последнее — это измерение компонента, а не сквозное утверждение о скорости. Загрузка модели и компиляция исключены из обоих тёплых таймингов.

На подмножестве фиксированной длины исходного эталона FP32 59/59 сравнений ответов совпадают, включая восемь языков и вопросы choice/score/noul. Наибольшее изменение калиброванной вероятности — 0.002925, и 100 повторных публичных вызовов дают конечные и идентичные округлённые результаты. Исходная фикстура из 63 вопросов содержит три длинных входа по 1 024 токена и один вход на 147 токенов с 20 вариантами; эти четыре оценки явно пропускаются экспортом L96. В этой фикстуре встречаются повторяющиеся рубрики. Это регрессионные сравнения, а не 59 независимых размеченных примеров или доказательство неизменной общей точности на задачах.

Отдельные экспорты L192 и L1024 также помещают все 6 390 назначенных операций тела на ANE. L192 проходит 60/60 сравнений; L1024 проходит полную 63/63 golden-фикстуру. Оба проходят 100 повторных публичных вызовов, и максимальная ошибка калиброванной вероятности остаётся 0.002925. Само подмножество длинных входов имеет максимальную ошибку 0.001128. Все оценённые счётчики использования токенов совпадают с эталоном.

Это последовательные отборочные прогоны, а не парные сравнения между бэкендами. Тайминги тела используют 5 прогревных и 30 измеряемых вызовов; тайминги полного короткого вопроса используют 10 прогревных и 50 измеряемых вызовов. Полный путь включает работу хоста и проверки входов. Экран L96 предшествует финальным дополнительным проверкам валидации входа; финальное контролируемое сравнение использует текущий адаптер и записывает свой отпечаток исходников. Времена компиляции/загрузки измерены в каждом процессе после конвертации, а не обещание о первом в истории запуске системы с пустыми кешами фреймворка. Большие фиксированные графы выполняют дополненную работу даже для коротких запросов. Практичный адаптер выбирал бы отдельные корзины длин; маршрутизация каждого запроса через L1024 отбросила бы преимущество коротких входов.

Отдельный прогон настоящей фикстуры workload(1, long=True) подтверждает один запрос на 1 024 токена, а не короткий запрос, дополненный до этого размера. Он измеряет 91.703 мс p50 / 94.776 мс p95 на 50 полных предсказаниях после десяти прогревных вызовов; все округлённые выходы остаются стабильными. Исторический бенчмарк MLX на длинном входе составляет 51.98 мс p50. Это не парные измерения одного раунда, но этот экран не даёт свидетельств, что нынешний граф ANE ускоряет длинные входы. Хеш входа, фактическая длина токенов, отпечатки текущего эксперимента и сырые тайминги сохранены в long1024-performance.json.

Сырые свидетельства:

• План вычислений полного тела и тайминги компонентов
• Валидация против эталона FP32 на реальных входах и тайминги полного предсказания
• Валидация L192
• Валидация полной фикстуры L1024
• План одного слоя и числовые проверки
• План одного MLP и числовые проверки

MLComputePlan описывает ожидаемое размещение, а не трассу исполнения оборудования. CPU_AND_NE допускает CPU и ANE; это не переключатель «только ANE». В этом эксперименте все назначенные операции тяжёлого тела предпочитают ANE, но аппаратная телеметрия рантайма оценивается отдельно. Последующая диагностика Instruments записала аппаратную активность Neural Engine; её таблица глобальна и не может отнести каждое событие к этой модели. Парное сравнение с MLX, результаты интегрирования мощности и ограничения трассы описаны в ANE_BENCHMARKS.md. Нулевой счётчик ANE от инструмента мониторинга не может установить отсутствие активности ANE без валидации этого счётчика на данной ОС/устройстве.

Почему исходный граф был плохой целью для ANE

Базовая линия сохраняет обычную раскладку трансформера B×L×C, операции с динамической формой, батчевое внимание по всем головам и оператор SDPA из Core ML. При выборе CPU/ANE её план содержит 24 операции SDPA без сообщённого назначения устройства, плюс множество приведений, срезов, запросов формы, сборов и транспонирований. Некоторые отдельные операции поддерживают ANE, но граф в целом не разбит на неё. Поддержка устройства для отдельных операторов поэтому недостаточна как свидетельство полезного пути исполнения на ANE.

Успешный прототип меняет несколько вещей сразу. Это свидетельство, что комбинация позволяет размещение на ANE, а не завершённая абляция, выявляющая один единственный проблемный оператор. Управляющие эксперименты с SDPA в исходной раскладке фиксированной формы и с явным вниманием — полезные следующие различающие эксперименты.

Опубликованные рекомендации Apple по трансформерам советуют channel-first 4D-тензоры, свёртки 1×1 для проекций, внимание по головам и меньше копий раскладки. Эти принципы мотивировали реализацию; исторические ускорения DistilBERT у Apple не устанавливают улучшение в 10× против уже быстрой базовой линии MLX FP16 этого проекта. Статья Apple об ANE Transformer, эталонная реализация Apple.

Архитектура прототипа и численный контракт

model.py — отдельная экспортная модель, построенная из исходных параметров контрольной точки:

• Скрытые активации используют B,C,1,L. Каждый плотный вес W[out,in] становится ядром свёртки 1×1 K[out,in,0,0] без дообучения или аппроксимации весов.
• Внимание разделено на отдельные 64-канальные головы. Тензор ключей транспонируется один раз, и два явных einsum вычисляют QK и AV, сохраняя раскладку 4D. Softmax выполняется по оси ключей, измерение 1.
• RoPE делит каждую голову на её две 32-канальные половины. Её косинусы, синусы и основания берутся из исходной модели, включая многоязычную локальную theta 160000.
• Нормализация каналов сохраняет исходный порядок normalized * weight + bias и epsilon. Она не копирует по-другому упорядоченное affine-выражение или опциональный зажим из эталонного LayerNorm Apple.
• Первая норма внимания энкодера остаётся единичной; энкодер использует точный erf GELU, тогда как два FFN головы решений сохраняют ReLU.
• Маски полного внимания и скользящего окна сохраняют маскирование валидных ключей и правило дополненного запроса. Локальный радиус читается из local_attention // 2.
• Финальный сбор маркеров выражается через внешне подготовленный one-hot селектор и 4D einsum. Граф сохраняет 32 слота маркеров, включая неактивные слоты; хост заменяет неактивные логиты исходным значением -1e4.

Целый исходный слой энкодера FP32 и его аналог BC1S отличались не более чем на 2.29e-5 в проверке раскладки PyTorch. Ошибка выхода слоя Core ML FP16 была больше, как и ожидается для этого изменения точности/бэкенда. Зонд тела изначально содержал самосравнение; это недействительное свидетельство было удалено, а его проверка раскладки PyTorch на всём теле явно указана как not_measured. Реальная валидация полной модели вместо этого идёт против сохранённых исходных логитов и решений FP32.

Независимые регрессии на CPU в test_ane_layout.py дополнительно сравнивают полный крошечный ConvBody с исходным DecisionModel, используя явные и SDPA оракулы внимания, три типа вопросов, изменённые значения паддинга, ненулевые bias норм, несколько оснований RoPE и нестандартный epsilon нормы. Эти тесты покрывают семантику раскладки и маскирования, не смешивая их с точностью оборудования FP16 полной контрольной точки. Все пять тестов раскладки прошли локально.

Прототип внимания добавляет конечный bias маски -1e4. Это даёт намеченное поведение маски на конечных проверенных активациях, но это не побитовая идентичность замене маскированных оценок на -1e4 или -бесконечность для произвольных экстремальных входов. Точно так же исполнение Core ML FP16 не заявлено как побитово идентичное исходной модели FP32.

runtime.py обеспечивает одну границу CPU→ANE→CPU для всего трансформера, а не переход устройства на каждый слой:

1. CPU токенизирует каждый вопрос, собирает только запрошенные строки эмбеддингов и строит аддитивные маски фиксированной формы, векторы типов и селекторы маркеров. Он не переиспользует контекстные скрытые состояния или K/V между вопросами.
2. Один вызов Core ML выполняет нормализацию эмбеддингов, все 22 слоя энкодера, оба слоя головы решений и скоринг-свёртки.
3. CPU выводит признаки действий из некалиброванного softmax сырых логитов и запускает небольшую исходную голову действий в FP32 с erf GELU. Публичная калибровка и форматирование вывода затем используют существующую реализацию.

Экспортированное тело имеет вычисления FP16, тогда как небольшая голова действий на хосте — FP32. Поиск эмбеддингов использует исходные сохранённые веса. Исходный файл safetensors содержит 169 тензоров FP16 и один тензор FP32: «эталон FP32» описывает исходное исполнение PyTorch, а не утверждение, что исходная контрольная точка хранится целиком в FP32. Эта граница смешанной точности — часть численного контракта прототипа. Различия вероятностей действий, равные нулю на насыщенных фикстурах, не доказывают идентичность логитов действий.

Адаптер отклоняет несовместимые размеры пакета, ID/маркеры вне диапазона, недопустимые значения маски, пустые строки ключей внимания и длины входов, превышающие его фиксированную ёмкость. Его лимит входа по умолчанию — 96 токенов, а размер батча — один; несколько вопросов исполняются последовательно. Он никогда молча не усекает запрос, чтобы уместиться в более короткий экспорт.

Воспроизведение, происхождение и шлюзы качества

Запускайте из корня репозитория в закреплённом .venv. Сгенерированные пакеты игнорируются Git; дублирующий NPZ эмбеддингов или большой артефакт весов не требуется. probe.py отказывается от существующего выходного каталога. Новые экспорты записывают исходные значения SHA256 весов/конфигурации, хеши содержимого пакета, форму, версии инструментов и отпечатки исходников эксперимента в манифест.

Команды воспроизведения пишут под artifacts/ane-repro/, потому что закоммиченные каталоги экспериментов уже содержат отчёты и манифесты. Выберите другой новый каталог при повторении экспорта; ни один из экспортёров молча не переиспользует существующий пакет.

Установите зависимости конвертации, разработки и исследования сжатия через pip install -e '.[convert,dev,research]' (или соответствующие дополнения uv sync). Дополнение research закрепляет kmeans1d==0.4.0; эта опциональная зависимость используется для экспериментов группового K-means FP16 и записывается в новые манифесты.

По умолчанию конвертация разрешает laya-multilingual в закреплённую контрольную точку Hugging Face репозитория и скачивает её при необходимости. Передайте --source /path/to/checkpoint, чтобы использовать существующие локальные файлы. Валидация по умолчанию использует исходный каталог, записанный в манифесте пакета. Сам рантайм ANEAgent принимает только локальные файлы; он проверяет исходные хеши весов/конфигурации, фиксированную форму и хеш содержимого пакета перед загрузкой. Валидация также отклоняет эталон с другим хешем исходных весов. Измеренный SHA256 исходных весов многоязычной модели — 9d628fd971b700382ac6f65920a86f149777b2e748e0c955fb3b19695aa8f204.

# Small placement probes, then the complete model.
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind mlp --length 96 --output artifacts/ane-repro/mlp96
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind layer --length 96 --output artifacts/ane-repro/layer96
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 96 --output artifacts/ane-repro/body96
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body96/model.mlpackage \
  --length 96 --repeats 100 \
  --output artifacts/ane-repro/validation96.json

Валидатор требует, чтобы все оценённые решения argmax совпадали, ошибки калиброванных и action-вероятностей <=0.02, выходы были конечными, использование токенов неизменным, а повторные публичные выходы идентичными. Неудачный кандидат пишет passed: false и завершается неуспешно. Пропущенные случаи остаются непроверенными, даже если подмножество фиксированной формы проходит. Начальный отчёт L96 применил эти шлюзы явно после измерения; он помечен соответственно, без изменения записанных таймингов.

Более длинные фиксированные экспорты и их команды валидации полного эталона:

.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 192 --output artifacts/ane-repro/body192
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body192/model.mlpackage --length 192 \
  --output artifacts/ane-repro/validation192.json
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 1024 --output artifacts/ane-repro/body1024
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body1024/model.mlpackage --length 1024 \
  --output artifacts/ane-repro/validation1024.json
.venv/bin/python -m experiments.ane_engineering.benchmark \
  --package artifacts/ane-repro/body1024/model.mlpackage --length 1024 --long \
  --output artifacts/ane-repro/long1024-performance.json

Эти команды воспроизводят независимо проверенные более длинные экспорты, перечисленные выше. Их размещение и числовая точность проверялись отдельно от графа L96; дополнение коротких запросов до 1024 токенов не является предлагаемой продакшн-политикой.

Отбор сжатия и цель 10×

palettize.py готовит независимые варианты палитры только весов, с 8-битными равномерными таблицами поиска как недорогим первым отборочным кандидатом. Выбираются только веса свёрток размером более 2048 элементов; константы RoPE, нормализация, математика активаций и веса действий хоста остаются неизменными. Сгруппированные выходные каналы используют отдельные таблицы поиска. K-means — более дорогой режим. Он использует установленную реализацию kmeans1d для этих групп FP16. Core ML Tools 9.0 распараллеливает независимые группы через пул процессов starmap, когда num_kmeans_workers > 1; в экспериментах используются восемь воркеров для офлайн K-means и один поток математической библиотеки на воркер. Число воркеров меняет пропускную способность экспорта, но не намеченную цель кодбука. Младшие 6/4-битные кандидаты — отдельные приближённые артефакты, а не точные реализации.

Размеры сжатия относятся к экспортированному пакету тела трансформера. Исходная таблица эмбеддингов на 196.608 миллиона записей остаётся на хосте, при этом для каждого запроса ищутся только запрошенные строки, а небольшие веса действий хоста остаются неизменными. Пакет тела, сжимающийся примерно вдвое, — это не двукратное сокращение всей контрольной точки, памяти рантайма или энергии на запрос.

.venv/bin/python -m experiments.ane_engineering.palettize \
  --package artifacts/ane-repro/body96/model.mlpackage \
  --bits 8 --mode uniform --group-size 32 \
  --output artifacts/ane-repro/body96-w8
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body96-w8/model.mlpackage --length 96 \
  --output artifacts/ane-repro/validation96-w8.json

# Independent K-means candidates; inspect each validation exit status.
for bits in 8 6 4; do
  VECLIB_MAXIMUM_THREADS=1 OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 \
    .venv/bin/python -m experiments.ane_engineering.palettize \
    --package artifacts/ane-repro/body96/model.mlpackage \
    --bits "$bits" --mode kmeans --group-size 32 --workers 8 \
    --output "artifacts/ane-repro/body96-w${bits}km"
  .venv/bin/python -m experiments.ane_engineering.validate \
    --package "artifacts/ane-repro/body96-w${bits}km/model.mlpackage" --length 96 \
    --output "artifacts/ane-repro/validation96-w${bits}km.json"
done

Оба равномерных экрана W8 сохраняют все 59 решений argmax фикстуры и проходят 100 повторных вызовов, но не проходят шлюз вероятности: размер группы 32 достигает ошибки 0.023612, а размер группы 4 — 0.033858. Экран W8 K-means/group32 проходит с максимальной ошибкой 0.014393 и 59/59 решениями. Насыщенные вероятности действий скрывают различия логитов действий до 16.24 для этого кандидата K-means; прохождение этой небольшой регрессионной фикстуры не устанавливает сохранённую общую калибровку или точность на задачах. Сжатые кандидаты остаются отдельно идентифицированными приближёнными моделями.

Фиксированный кандидат W6 K-means/group32 также сохраняет 59/59 решений argmax и стабильные повторные выходы, но не проходит с максимальной ошибкой вероятности 0.052243. Его максимальная ошибка логитов действий — 76.62. Сокращение весов до шести бит поэтому не удовлетворяет неизменённому критерию приёмки, даже несмотря на то что его задержка отборочного экрана для коротких запросов остаётся близкой к FP16.

W4 K-means/group32 также сохраняет 59/59 решений, но максимальная ошибка вероятности возрастает до 0.200221, а максимальная ошибка логитов действий — до 605.30. Он не проходит тот же шлюз. Отсутствие изменений argmax во всех пяти сжатых экранах показывает, почему одних только насыщенных решений этой фикстуры недостаточно для адекватного теста приёмки.

Все сжатые варианты сохраняют 6 390 операций, назначенных устройству с предпочтением NE, 59/59 совпадений argmax на фикстуре и 100 стабильных повторных вызовов. Остальные записи плана включают константы и выражения восстановления весов из LUT; одних метаданных размещения недостаточно, чтобы доказать, сколько сжатых данных перемещается из DRAM во время запроса. Три экспорта K-means занимают 186.50, 56.13 и 23.90 секунды соответственно с восемью офлайн-воркерами. Настройка и рабочие процессы завершаются до каждого измерения инференса.

Эти последовательные экраны на 50 вызовов не устанавливают отношений ускорения. Экран FP16 предшествует последним проверкам валидации входа, и экраны не чередуются. W8 K-means — единственный сжатый финалист для более сильного внутрисессионного сравнения в ANE_BENCHMARKS.md. Отклонённые варианты сохранены как свидетельство границы точности, а не рекомендованные развёртывания. Сжатие валидировалось только для этого многоязычного подмножества L96; результат L1024 на 63 вопроса выше относится к отдельному экспорту FP16.

Палитра Core ML реконструирует плавающие веса из индексированных таблиц поиска; меньшие хранимые тензоры сами по себе не устанавливают более быстрый инференс или низкую энергию. Каждый вариант нуждается в одних и тех же шлюзах точности, свежем плане вычислений и парном сквозном сравнении скорости/энергии. Документация по палетизации Core ML.

Начальное успешное размещение на ANE устанавливает правдоподобный путь оптимизации, а не результат 10×. Сравнение должно использовать MLX FP16, включая его скомпилированный вариант там, где он быстрее, и сообщать равный объём задачи. Мощность должна интегрироваться по полным запросам. Как валовую энергию системы, так и любую оценку за вычетом простоя следует сообщать с их ограничениями измерения. Независимый математический обзор — ANE_MATH.md.

Что устанавливает ручная реализация

Полезная ручная работа здесь — полное, независимо проверенное переписывание графа вычислений в раскладку, которую Core ML может отобразить на ANE. Это меняет размещение исполнения, сохраняя обученные параметры. Это существенно эффективнее, чем изменение compute_units на исходном графе. Это не убирает 24 последовательных блока внимания/MLP или их плотную проектную работу.

Следующий точный кандидат для длинных запросов — внимание, которое действительно посещает только локальные окна, реализованное с фиксированными тайлами query/key и исходными правилами паддинга и RoPE. Текущий граф всё ещё вычисляет плотную матрицу оценок и применяет локальную маску. Тайловый граф мог бы сократить эту работу, но больше срезов, границ и небольших свёрток может подорвать планирование ANE; его размещение, точность и сквозная выгода остаются неизмеренными. Дальнейшее сжатие весов требует калибровки или восстановления качества после неудач выше. Дистилляция или меньше слоёв ввели бы новую модель и потребовали бы более широкой оценки качества задач. Ни одно из этих нереализованных направлений не даёт свидетельства выигрыша 10× сегодня.