Запуск 26B MoE на дешёвом inf2.xlarge: упаковка без fp4

Зачем это читать?

Вы думаете, что для запуска 26B-модели нужен дорогой сервер? Я тоже так думал, пока не упаковал Gemma-4 с 128 экспертами на самый дешёвый инстанс AWS Inferentia2 — inf2.xlarge с 16 ГБ оперативной памяти и 32 ГБ HBM за $0.76/час. Это в 8.6 раз дешевле исходного решения. Как? Без fp4, только int8 и немного изобретательности.

Проблема: модель не влезает

Исходная версия на inf2.24xlarge (12 NeuronCores, 192 ГБ HBM) работала, но была избыточной для одного пользователя. Цель — inf2.xlarge (2 ядра, 32 ГБ HBM, 16 ГБ RAM). Две стены: 128 экспертов при bf16 весят ~45.6 ГБ, что не влезает в 32 ГБ HBM, а стандартная загрузка требует гигабайты host RAM.

Стена 1: математика памяти

Эксперты составляют ~93% весов. При bf16 128 экспертов — 45.6 ГБ. С репликацией на 2 ядра (TP=2) получаем ~22.8 ГБ на ядро при бюджете 16 ГБ. Top-8 роутинг уменьшает вычисления, но не память: все 128 экспертов должны быть загружены. Первая мысль — int8 для экспертов. NxD предоставляет QuantizedColumnParallel и QuantizedRowParallel. При int8 эксперты занимают ~11.4 ГБ на ранг. Должно поместиться. Но нет — ошибка Allocation Failure: не хватает 3–4 ГБ.

Ложный след: fp4

Я решил, что нужны fp4-эксперты. Потратил дни на изучение: F4E2M1FN_X4 существует как константа, но не работает — ядра только CPU или предпродакшн. В AWS документации чётко: NxD поддерживает только INT8 и FP8. FP4 — тупик.

Прорыв: квантование не того

Оказалось, 3–4 ГБ, которых не хватало, — не в экспертах (они уже int8), а в весах, оставленных в bf16. Главный виновник — lm_head (1.48 ГБ на ранг). Решение: int8 + шардинг lm_head через QuantizedColumnParallel(gather_output=True). Это уменьшило его до ~0.37 ГБ/ранг. Также int8 для shared dense MLP. Итог: ~12–13 ГБ на ранг — comfortably under 16 ГБ. Модель заработала на inf2.8xlarge (2 ядра, 32 ГБ HBM).

Стена 2: host RAM

inf2.8xlarge и inf2.xlarge имеют одинаковые 2 ядра и 32 ГБ HBM, но host RAM: 128 ГБ vs 16 ГБ. Компиляция требует ~180 ГБ (пик). Но для развёртывания нужен только neff-файл (веса на устройстве). Решение: минималистичная загрузка — структура через accelerate.init_empty_weights(), эмбеддинги из отдельного файла, neff загружается через torch.jit.load. На inf2.xlarge с файлом подкачки 40 ГБ neff загрузился за 112 секунд, пик RSS ~11 ГБ.

Баги: эмбеддинги с масштабированием

Первая попытка выдала пустой ответ. Оказалось, Gemma-4 использует Gemma4TextScaledWordEmbedding, который умножает эмбеддинги на sqrt(hidden_size) (~53). На хосте я использовал обычный nn.Embedding без масштаба. Одна строка кода исправила:

ie = emb(ids) * (hidden_size ** 0.5)  # как в Gemma4TextScaledWordEmbedding

После этого модель выдала правильный ответ.

Практический вывод

Что делать прямо сейчас:

• Профилируйте занятость памяти, а не гадайте — я потратил дни на fp4, а решение было в другом месте.
• Не верьте, что «почти работает» — проверяйте документацию по поддерживаемым фичам.
• Разделяйте компиляцию и развёртывание: компиляция требует много RAM, развёртывание — мало.
• При выносе эмбеддингов на хост наследуйте класс эмбеддинга, а не только таблицу.
• Проверяйте на чистом инстансе: docker run на свежей spot-машине — единственное доказательство.

Готовый образ на Docker Hub: xbill9/gemma4-optb-26b:xlarge. Модель на Hugging Face: xbill9/gemma-4-26B-A4B-it-inferentia2-xlarge (публичный).

Да, decode всего ~6 tok/s — цена за полный просчёт всех 128 экспертов вместо top-8. Но для одного пользователя, когда важна стоимость, это правильный трейд-офф. 26B-модель на самом дешёвом Inferentia2 — это реально.