Перенос модели Google Gemma-4 31B на AWS Inferentia2 inf2.24xlarge обеспечил полное совпадение токенов (token-for-token match) с эталоном на CPU — однако каждое сгенерированное предложение было бессмыслицей. Разрыв между «совпадением» и «работоспособностью» теперь служит предупреждением для всех, кто пытается уместить массивные LLM на специализированных чипах для инференса от Amazon.

Почему посимвольного совпадения токенов недостаточно

Разработчик сравнил каждый выходной токен устройства Inferentia с токеном, полученным при запуске модели на CPU. Потоки были идентичны, поэтому казалось, что оборудование в точности воспроизвело эталонную реализацию. На самом деле оба потока подавали некорректный промпт в модель, лишенную чат-шаблона и снабженную неверными маркерами реплик. Отсутствие шаблона загнало модель в бесконечный цикл, выдающий бессмыслицу. Оборудование выполнило свою работу — оно воспроизвело баг, который уже присутствовал в эталонном коде.

Урок прост: SEQ_MATCH (последовательное равенство токенов) не означает корректность. Если эталонная реализация сломана, верная аппаратная копия унаследует ту же ошибку. Валидация должна выходить за рамки посимвольного паритета; необходимы сквозные функциональные проверки с правильно отформатированными входными данными.

Буферы, маскирующиеся под параметры

На этапе загрузки загрузчик модели пропустил компонент под названием layer_scalar. В определении модели PyTorch этот объект был зарегистрирован как буфер (buffer), а не как параметр (parameter). Буферы — это статические тензоры, которые не обновляются в процессе обучения, и многие загрузчики игнорируют их при конвертации в форматы, совместимые с Neuron. Пропуск этого компонента оставил коэффициенты масштабирования для нескольких слоев значениями по умолчанию, что исказило математические вычисления во всей сети. Ошибок не возникло; модель скомпилировалась, и конвейер инференса запустился, но численные результаты оказались неверными.

Всем, кто переносит большие модели на Inferentia, следует проводить аудит каждого непараметрического тензора. Даже если тензор не предназначен для обучения, он может быть критически важен для корректных вычислений при прямом проходе. Ручная проверка включения буферов может предотвратить скрытые ошибки масштабирования, которые трудно диагностировать.

Волатильность спот-инстансов и 39-минутная компиляция

Запуск модели с 31 миллиардом параметров на спот-инстансе кажется дешевым решением, но экономия сопряжена с непредсказуемыми случаями изъятия ресурсов (reclaim events). Время компиляции разработчика — около 39 минут на перевод модели в код, совместимый с Neuron — исчезло, когда AWS изъяла инстанс. Чтобы справиться с прерываниями, была создана трехступенчатая система защиты:

  • ModelBuilder удерживал использование памяти в пределах 384 ГБ хоста, избегая сбоев, требующих перезапуска.
  • Мгновенное зеркалирование в S3 как файлов необработанных весов, так и скомпилированных «neffs» (исполняемых файлов Neuron) позволяло новому инстансу продолжить работу ровно с того места, где остановился предыдущий.
  • Мультирегиональный поллер сканировал регионы AWS на наличие доступных мощностей спот-инстансов и запускал новый инстанс, как только тот появлялся.

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

Ловушки шардинга при смешанных конфигурациях внимания

Gemma-4 31B использует две конфигурации внимания. В некоторых слоях задействовано четыре головы key-value (KV), в других — другое количество. Равномерное разделение модели на восемь параллельных рангов (ranks) не срабатывает, если количество KV-голов в слое не делится нацело. Попытка распределить слой с 4 головами на восемь рангов заставила бы каждый ранг обрабатывать «полголовы» — математическая невозможность, вызывающая несоответствие размерностей и ошибки выполнения.

Решением стало дублирование глобально шардированных слоев (тех, у которых количество голов совместимо) на всех рангах и шардинг только «скользящих» слоев (sliding layers), чье количество голов позволяло равномерное разделение. Эта гибридная стратегия позволила сохранить эффективность тензорного параллелизма, избежав при этом недопустимого деления KV-голов и устранив ошибки тензорной параллелизации, которые мешали предыдущим попыткам.

Вывод

Перенос гигантской LLM на Inferentia — это не просто упражнение в стиле «скомпилировал и запустил». Он требует тщательного функционального тестирования, выходящего за рамки равенства токенов, скрупулезной проверки того, что каждый тензор — будь то параметр или буфер — обрабатывается правильно, и стратегии развертывания, учитывающей возможность изъятия спот-инстансов. Наконец, шардинг должен учитывать внутреннюю геометрию внимания модели; в противном случае параллелизм, обещающий скорость, станет источником скрытых сбоев.