Запуск модели с триллионом параметров в Sakana AI

Источник: Modal•

Запуск модели с триллионом параметров в Sakana AI

Развертывание модели Mixture-of-Experts (MoE) с ~1 триллионом параметров для инференса через API — классическая системная задача. При полной точности FP16 или BF16 объем памяти требует очень больших кластеров. Стандартные методы посттренировочной квантизации (PTQ) уменьшают этот объем, но могут…

Развертывание модели Mixture-of-Experts (MoE) с ~1 триллионом параметров для инференса через API — классическая системная задача. При полной точности FP16 или BF16 объем памяти требует очень больших кластеров. Стандартные методы посттренировочной квантизации (PTQ) уменьшают этот объем, но могут ухудшить качество модели, особенно для неанглоязычных моделей, где обычные калибровочные датасеты для квантизации не учитывают языковой регистр или поведение вызова инструментов.

В преддверии запуска Sakana Namazu, модели с ~1T параметров, оптимизированной для японского бизнес-контекста, живого поиска в интернете и выполнения кода, инженерная команда Modal по работе с клиентами на передовой работала напрямую с Sakana AI, чтобы запустить их модель на Modal.

Мы настроили квантизацию NVFP4 для сохранения японского языка и обучили кастомную модель спекулятивного предсказания на сырых потоках токенов. Это работало в глобальном пуле GPU Modal, что позволило Sakana запустить публичный API без необходимости резервирования вычислительных мощностей.

Вот как мы оптимизировали стек обслуживания вплоть до уровня ядер.

Квантизация без компромиссов

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

Работая вместе, Modal смогла запустить Sakana Namazu на NVFP4 — формате пониженной точности, созданном для GPU NVIDIA Blackwell, что снизило требования к памяти и увеличило пропускную способность.

Вместо применения готового метода квантизации инженерная команда Modal по работе с клиентами на передовой начала с измерения того, как веса Sakana Namazu реагируют на сжатие в формате NVFP4, а затем построила подход на основе полученных результатов:

  • Выборочная квантизация экспертов: квантизировались только маршрутизируемые эксперты, на которых приходится основная часть параметров. Механизм внимания (attention), маршрутизатор и разделяемые эксперты остались с более высокой точностью (BF16).
  • Релевантные для предметной области калибровочные данные: калибровка проводилась на диалогах в стиле Sakana Namazu с большим акцентом на японский язык, чтобы каждый из 23 040 экспертов модели получил репрезентативный трафик до фиксации коэффициентов.

Результатом стала модель, которая соответствовала контрольной точке (чекпоинту) Sakana для продакшна при следовании инструкциям на японском языке в пределах погрешности собственных опубликованных бенчмарков Sakana, работая при этом со скоростью и потреблением памяти NVFP4.

Модель предсказания, обученная на репрезентативном трафике Sakana Namazu

Спекулятивное декодирование — это метод ускорения генерации, при котором небольшая модель предсказания угадывает несколько токенов вперед, а основная модель проверяет их за один шаг. Чем больше таких догадок принимает основная модель, тем быстрее выполняется каждый запрос. Как мы уже говорили ранее, спекуляции — это все, что вам нужно.

Обычные модели предсказания обучаются на общем тексте, поэтому они плохо угадывают японский бизнес-регистр Sakana Namazu (что приводит к уровню принятия ниже 40%). Вместо этого команда Modal обучила кастомную модель предсказания DFlash на репрезентативном трафике Sakana Namazu.

Ключевой деталью реализации было обучение на точных идентификаторах токенов (token IDs), созданных моделью, а не на тексте, который был декодирован и токенизирован заново. На странице они выглядят одинаково, но повторная токенизация разрушает точное выравнивание границ токенов внутри вызовов инструментов, схем JSON и маркеров рассуждений модели. Этот разрыв снижает процент принятых токенов. Modal создала конвейер сбора данных так, чтобы сырые ID токенов поступали напрямую из продакшна в обучающий корпус объемом более 30 миллиардов токенов. Вот результаты:

  • Увеличение коэффициента принятия спекулятивных токенов на +25% по сравнению с обычной моделью предсказания
  • Падение задержки декодирования на один токен примерно на 35%
  • Прирост трафика на один репликат примерно на 40% до начала роста задержек очереди, что позволило Sakana обслужить спрос в день запуска с меньшим количеством GPU

С момента запуска в работу модели предсказания были внесены дополнительные улучшения. Первая версия была обучена до того, как Sakana Namazu обработала хотя бы один публичный запрос, поэтому она училась на репрезентативном трафике, а не на реальном. Как только пошел продакшн-трафик, команда Modal обучила вторую модель предсказания с нуля примерно на 1 млн реальных запросов.

В ней используется DFlash2 — более новая архитектура, которая лучше справляется с ситуациями, когда клиенты используют ненулевую температуру сэмплирования, что охватывает большую часть трафика Sakana Namazu. Для оценки мы запустили две одинаковые копии среды обслуживания Sakana Namazu — каждую со своей моделью предсказания. Затем мы прогнали через обе 1 000 отложенных запросов, используя настройки сэмплирования, которые реально отправляют клиенты Sakana. Новая модель предсказания никогда не видела эти запросы во время обучения. Средняя длина принятия (сколько предсказанных токенов принимает главная модель за один шаг) выросла с 3,51 до 5,01 (прирост на 43%), а задержка p99 упала с 68 с до 52 с.

Работа над производительностью при выходе новых версий движка

Стеки обслуживания развиваются быстро, и каждый релиз меняет ядра, планировщик и арифметику таким образом, который проявляется только при реальном продакшн-трафике. Modal и Sakana работали вместе, чтобы тестировать и проводить бенчмаркинг этих изменений, а также безопасно развертывать их для рабочих нагрузок Sakana, выполняя A/B тестирование на том же оборудовании и том же пакете запросов.

Наибольший прирост дали следующие изменения:

  • Устранение задержки синхронизации CPU и GPU в цикле спекулятивного декодирования. Каждый шаг декодирования приостанавливался, пока CPU считывал результаты от GPU. Устранение этого ожидания сократило время простоя GPU на каждом шаге более чем на 90%.
  • Маршрутизация вычислений смеси экспертов (MoE) NVFP4 на более быстрые ядра. Новые пути ядер для матричных умножений экспертов, на которые приходится около 60% каждого шага декодирования, сократили время шага и, как побочный эффект более точной арифметики, повысили принятие предсказаний на 9% с той же моделью предсказания.
  • Кеширование диспетчеризации префилла (prefill) и настройки исполнителя экспертов. Работа, которая повторялась при каждом запросе, теперь выполняется один раз и используется повторно, что главным образом улучшило показатель времени до первого токена (time to first token).

Корректность работы на разных GPU

Sakana Namazu работает на четырех GPU с тензорным параллелизмом, и спекулятивное декодирование добавляет тонкий нюанс: каждый GPU должен прийти к согласию относительно того, какие предсказанные токены были приняты. Modal запрограммировала этап верификации для всех рангов, обнаружила редкий сценарий расхождения и сделала принятие единым общим решением. В результате целый класс редких задержек был полностью устранен без какого-либо влияния на качество вывода.

Команда также распространила спекулятивное декодирование на структурированный вывод. API-клиенты Sakana Namazu полагаются на ответы в формате JSON-схемы для использования инструментов, и Modal обеспечила генерацию с грамматическими ограничениями внутри спекулятивного конвейера, чтобы такие запросы получали то же ускорение без единого нарушения схемы.

Более умный маршрутизатор между репликатами и автомасштабирование

После запуска Modal включила собственный маршрутизатор с учетом KV-кеша (KV-aware router) в развертывание Sakana без необходимости внесения каких-либо изменений в код со стороны Sakana.

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

Влияние на задержки хвоста было мгновенным. Скачки времени до первого тонака (TTFT), достигавшие нескольких минут, исчезли, очереди на реплику упали почти до нуля, а сквозная задержка хвоста улучшилась, в то время как скорость декодирования на токен осталась прежней.

Управление планированием емкости с помощью GPU по требованию

Перенос критически важной рабочей нагрузки на новую инфраструктуру — трудная задача для любой передовой ИИ-лаборатории. Для обеспечения необходимого количества GPU часто требуется долгосрочный контракт или крупное бронирование.

Поскольку новый коммерческий API-сервис для Sakana Namazu еще не получал запросов, этот путь был бы рискованным. Modal позволил им эластично масштабироваться для удовлетворения спроса, возникшего во время запуска, не беспокоясь о переобеспечении или медленном автомасштабировании.

«Для Sakana Namazu нам потребовалась гибкость масштабирования емкости GPU вверх или вниз без ущерба для качества модели. Modal предоставил нам эту гибкость, а их инженеры работали непосредственно с нами над решением проблем с производительностью, вызванных квантованием».

Что это значит за пределами Sakana Namazu

Теперь Sakana Namazu работает в продакшене на инфраструктуре, которая была совместно оптимизирована под него: квантование NVFP4, настроенное на сохранение японского языка, спекулятивная черновая модель, обученная на потоках сырых токенов Sakana, и бессерверная емкость, которая динамически масштабируется в зависимости от спроса на запуск, а не по забронированному за месяцы вперед тарифу.

Будучи одной из ведущих ИИ-лабораторий Японии, исследовательская группа Sakana теперь изучает возможность запуска дополнительных рабочих нагрузок на Modal помимо исходного API.

Мы рады видеть, что Sakana выпустит дальше!

О чём эта статья