Оглавление

По данным Hugging Face, разработчики из проекта ANEMLL нашли остроумное решение для запуска современных моделей с RMSNorm на устаревшем аппаратном обеспечении Apple Neural Engine.

Проблема аппаратной совместимости

Apple Neural Engine (ANE), представленный в 2017 году с чипом A11 Bionic, был разработан с фиксированным набором предопределенных операций для популярных машинно-обучаемых workloads. В то время доминирующим методом нормализации в трансформерных моделях была Layer Normalization (LayerNorm), на которой Apple и основала набор операций ANE.

Однако примерно с 2019 года более легковесная альтернатива — Root-Mean-Square Normalization (RMSNorm) — начала набирать популярность. Она предлагает схожие преимущества с LayerNorm, но с более простыми вычислениями и меньшим количеством параметров. К сожалению, ANE не поддерживает RMSNorm как нативную операцию, и его аппаратный набор операций оставался практически неизменным на протяжении поколений.

Поскольку ANE является «черным ящиком», разработчики не могут видеть реализацию или модифицировать определенные операции, такие как Softmax, ReLU или даже LayerNorm. Но реализации Apple обычно сильно оптимизированы, и критически важно использовать нативные операторы там, где они доступны, для достижения лучшей производительности и точности.

Математическое решение

Идея заключается в манипуляции входными данными так, чтобы LayerNorm, уже доступный на ANE, мог выполнять ту же функцию, что и RMSNorm. Поскольку LayerNorm нормализует на основе среднего и дисперсии, а RMSNorm нуждается только в среднеквадратичном значении (форма стандартного отклонения без вычитания среднего), мы стремимся привести входные данные в форму, где среднее равно нулю, а дисперсия соответствует среднеквадратичному значению.

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

Реализация на PyTorch

import torch
import torch.nn.functional as F

def ANEMLL_RMS(x, weight, eps=1e-5):
 # ❶ Сделать среднее последнего измерения равным нулю.
 doubled = torch.cat([x, -x], dim=-1)

 # ❷ Запустить оптимизированный для ANE LayerNorm на удвоенном тензоре.
 normed = F.layer_norm(
 doubled,
 normalized_shape=doubled.shape[-1:],
 weight=None,
 bias=None,
 eps=eps
 )

 # ❸ Отбросить зеркальную половину → корректно нормализованные RMS активации.
 normed = normed[..., :x.shape[-1]]

 # ❹ Применить обучаемый коэффициент усиления (γ)
 return normed * weight

Этот элегантный математический трюк демонстрирует, что иногда самые эффективные решения оказываются и самыми простыми. Вместо того чтобы ждать аппаратных обновлений от Apple, разработчики нашли способ «обмануть» существующее железо, заставив его выполнять современные операции. Ирония в том, что для реализации более простого RMSNorm приходится использовать более сложный LayerNorm — типичная ситуация в мире обратной совместимости.

Перспективы развития

Если Core ML добавит RMSNorm как встроенную операцию, компилятор мог бы внутренне свести её к этому трюку с двойным LayerNorm на основе конкатенации на данный момент, одновременно прозрачно отображая её на оптимизированные аппаратные инструкции, как только будущие поколения ANE добавят нативную поддержку. Эта конструкция гарантирует, что модели, построенные сегодня, будут продолжать эффективно работать как на текущем, так и на будущих аппаратных устройствах Apple без необходимости изменений со стороны разработчиков.

Реализацию можно найти в библиотеке ANEMLL: https://github.com/Anemll.