Вторая половина «стены памяти»: инференс моделей MoE на 35B с SSD с помощью обученного предсказания маршрутизации

Разбор статьи

Проблема и мотивация

Модели со структурой Mixture-of-Experts (MoE) масштабируют общее количество параметров, сохраняя число активных параметров на токен на низком уровне (например, активируя всего ~3B параметров на токен в модели на 35B). Однако инференс на потребительском оборудовании упирается в статическую память под веса: MoE на 35B требует около 19,5 ГБ даже при квантовании в int4. На устройстве с 24 ГБ это полностью вытесняет операционную систему и KV-кэш, приводя к жесткому своппингу страниц памяти и резкому падению скорости генерации (например, до 3,9 ток/с для резидентной модели на Mac mini с M4 Pro).

Выгрузка (оффлоадинг) весов экспертов на SSD снижает объем занимаемой оперативной памяти, но наивная загрузка по требованию создает непреодолимое бутылочное горлышко сериализации: выбор экспертов для слоя $N+1$ зависит от выходных данных слоя $N$. Чтение с SSD невозможно инициировать заранее, чтобы перекрыть его вычислениями, поэтому исполнение зависает на время латентности диска на каждом слое для каждого токена.

Ключевые идеи

• SSD как уровень хранения весов: Веса экспертов MoE хранятся на SSD в формате int4 и проецируются в память (mmap) в компактный потоковый пул. Пиковый объем памяти определяется активным рабочим набором, а не общим размером параметров модели.
• Предсказание как маршрутизация (Prediction as Routing): Вместо ожидания завершения вычислений слоя $N$ послойный прероутер прогнозирует маршрутизацию слоя $N+1$ на целый токен вперед. Важно, что на этапе генерации модель использует это предсказание непосредственно как саму маршрутизацию. Это полностью исключает отбрасывание токенов или аварийные дозагрузки по определению, так как подготовленный набор экспертов всегда строго равен маршрутизируемому.
• Необъединенная LoRA восстановления (Unmerged Recovery LoRA): Квантование 4-бит и приближенная маршрутизация снижают качество модели. Edge0 обучает низкоранговый адаптер на пути маршрутизации ученика с помощью дистилляции. Поскольку слияние LoRA с базой int4 и последующее переквантование стирают большую часть дельты (на уровне логитов сохраняется лишь 18% эффекта), Edge0 запускает адаптер параллельно в виде необъединенной дельты с пренебрежимо малыми накладными расходами.

Метод

Архитектура прероутера

Каждая голова прероутера на слое $N$ запускается на токене $t$ после слоя нормализации (post-attention layer norm) и предсказывает маршрутизацию для слоя $N+1$ на токен $t+1$. Входной вектор объединяет скрытое состояние с двумя one-hot векторами top-$k$ маршрутизации (текущего и предыдущего токенов):

$$\text{dim}=d_{\text{model}}+2E$$

Для модели 35B это $2048+2\times 256=2560$. Голова вычисляет:

$$\text{head}(x)={\mathrm{fc}}_2(\mathrm{erf}\text{-}\mathrm{gelu}({\mathrm{fc}}_1(x)))+\ell (x)$$

где $\ell (x)$ — линейный остаточный путь, инициализированный весами роутера следующего слоя.

Выполнение необъединенного адаптера

Слой MoE сохраняет квантованные базовые байты и параллельно применяет LoRA-дельту:

$$y=W_{int4}(x)+\frac{\alpha }{r}BAx$$

Трехэтапное обучение

1. Дистилляция голов прероутера: Головы прероутера обучаются имитировать истинные решения роутера следующего слоя модели-учителя при замороженной остальной части сети.
2. SFT на траектории ученика: Адаптеры LoRA восстановления (прикрепленные к проекциям внимания, линейного внимания и общих экспертов, но не потоковых экспертов) обучаются на сгенерированном учителем тексте, при этом маршрутизация принудительно направляется через прероутер.
3. Онлайн-дистилляция (On-Policy Distillation): Ученик генерирует текст в режиме премаршрутизации, базовый учитель в fp16 оценивает генерации, а параметры LoRA оптимизируются с помощью обратной дивергенции Кульбака — Лейблера с целевой функцией вероятностей top-$k$ с учетом «хвоста» распределения.

Результаты

• Пропускная способность и память (Mac mini M4 Pro, 24 GB):
  • edge0-35b ($K=4$): Достигает скорости генерации 20,4 ток/с, используя всего 2,9 ГиБ пиковой активной памяти из чекпоинта int4 размером 19,5 ГБ (по сравнению с 3,9 ток/с и 18,2 ГиБ при полном размещении в RAM).
  • edge0-8b ($K=8$): Достигает 28,0 ток/с генерации при пиковой памяти всего 1,5 ГиБ.
• Ускорение предвыборки (MacBook M2, 16 GB):
  • По сравнению с чисто реактивной потоковой передачей по запросу, предварительная выборка прероутером на 1 токен вперед увеличивает скорость декодирования на +80% ... +84% при различных параметрах маршрутизации ($K=2$: с 4,8 до 8,6 ток/с; $K=4$: с 3,5 до 6,4 ток/с; $K=8$: с 1,8 до 3,3 ток/с).
  • Время блокировки основного потока при загрузке резко сокращается (например, с 244,0 мс до 101,9 мс на шаг при $K=4$).
• Качество модели (OpenCompass):
  • По пяти бенчмаркам (AIME 2026, HumanEval, GPQA-Diamond, MMLU-Pro, IFBench) edge0-35b набирает в среднем 79,2 балла против 83,2 у базовой модели Qwen3.6 в fp16 (разрыв 3,9 балла).
  • edge0-8b достигает 69,9 против 72,7 у базовой Ling 3.0 tiny в fp16 (разрыв 2,8 балла).

Ограничения и открытые вопросы

• Конкурентность обслуживания: Edge0 обрабатывает одиночные запросы в порядке очереди (FIFO); пакетирование (батчинг) меняет рабочий набор экспертов и динамику кэширования, что не учитывалось в однопоточных профилях.
• Накладные расходы CPU на построение графа: Скорость декодирования упирается в накладные расходы фреймворка на стороне CPU (~44 мс/шаг при построении вычислительного графа MLX для 40 слоев), которые невозможно устранить оптимизацией SSD.
• Снижение качества в задачах на рассуждение: Потеря производительности в основном сосредоточена на сложных бенчмарках рассуждения (например, падение на 6,1 балла на AIME для 35B и на 10,0 баллов для 8B).
• Зависимость от оборудования: Относительный прирост скорости от прероутера зависит от задержек диска и снижается при прогретом кэше весов или при использовании сверхбыстрых хранилищ.