Сегодня мы публикуем материал от ZenteiQ, стартапа в области искусственного интеллекта, выделившегося из Индийского института науки (IISc) и выбранного в рамках инициативы IndiaAI Mission для создания суверенных научных фундаментальных моделей с нуля. Чтобы предварительно обучать модели, понимающие физические законы, управляющие уравнения и численные методы, компания ZenteiQ спроектировала и создала крупномасштабный конвейер пакетной генерации синтетических данных, который работает с несколькими семействами моделей и передает проверенные наборы данных на этап обучения. Google Kubernetes Engine (GKE) и GKE Inference Gateway предоставили ZenteiQ уровень маршрутизации вывода в рамках ее конвейера, что позволило увеличить совокупную пропускную способность вывода в 4,07 раза и нарастить ежедневный объем генерации с 1,47 миллиарда до 5,97 миллиарда токенов на каждые 1000 ускорителей. — Фисайо Фейисетан (Fisayo Feyisetan), менеджер по продукту, Google Cloud

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

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

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

Когда генерация данных замедляет предварительное обучение

Чтобы максимизировать скорость генерации данных, наш конвейер разделяет подготовку промптов на рабочих подах и обслуживание моделей на vLLM в многоузловых пулах ускорителей. Каждый пакет проходит через четыре этапа: расширение научных исходных концепций в пакеты промптов, генерация кандидатных траекторий рассуждений с помощью моделей-генераторов, оценка ответов с помощью моделей-оценщиков и проверка механизмов безопасности (guardrails) с помощью отдельного семейства моделей перед упаковкой проверенных наборов данных для предварительного обучения.

Изначально мы запускали эти пулы моделей на Kubernetes со стандартной балансировкой нагрузки ClusterIP по круговому алгоритму (round-robin) между конечными точками вывода. По мере масштабирования нашего конвейера для обеспечения графиков обучения наших моделей с 9B и 35B параметрами два узких места в маршрутизации начали замедлять создание наборов данных.

Неравномерные очереди движка: наши восходящие рабочие поды разворачивают исходные концепции в промпты сильно различающейся сложности и длины ответов. Поскольку стандартная балансировка нагрузки round-robin не имела информации об очередях движка vLLM, она продолжала отправлять новые запросы на обслуживающие поды, которые уже были заняты обработкой длинных траекторий рассуждений, в то время как на соседних узлах имелись свободные мощности.
Избыточное вычисление префиксов: промпты в пакете генерации разделяют примерно от 70% до 80% своего инструктивного контекста и спецификаций задач. Маршрутизация round-robin распределяла запросы с идентичными инструкциями по разным экземплярам обслуживания, заставляя каждый под заново пересчитывать один и тот же общий префикс промпта с нуля вместо повторного использования горячих кэшей в памяти ускорителя.

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

Создание на базе GKE Inference Gateway с поддержкой llm-d

Мы уже выполняем подготовку данных и задачи предварительного обучения MaxText на GKE, что позволяет нашей команде управлять генерацией данных и обучением моделей на одной и той же платформе Kubernetes. Следуя рекомендации Google Cloud устранить узкие места маршрутизации без изменения нашей базовой научной методологии или архитектуры конвейера, мы развертываем GKE Inference Gateway на базе открытого llm-d фреймворка распределенного вывода перед нашими пулами ускорителей vLLM.

В нашем конвейере рабочие поды извлекают научные исходные документы и доменные таксономии из Cloud Storage и разворачивают их в пакеты многошаговых промптов. Эти пакеты проходят через единое развертывание GKE Inference Gateway, которое управляет трафиком во всех наших семействах моделей и обслуживающих узлах.

Вместо поддержания изолированной маршрутизации для каждого уровня моделей мы направляем запросы на генерацию, запросы на оценку и проверки guardrails по их соответствующим пулам моделей через один и тот же шлюз. Как только кандидатные ответы проходят оценку и проверки безопасности, рабочие поды дедуплицируют и упаковывают проверенные наборы данных для непосредственного потребления последующими задачами предварительного обучения MaxText на GKE.

GKE Inference Gateway направляет запросы пакетной генерации, оценки и проверки guardrails по многоузловым пулам ускорителей, работающим под управлением vLLM, перед тем как проверенные наборы данных поступают на предварительное обучение MaxText.

От круговой маршрутизации (round-robin) к маршрутизации с учетом нагрузки и кэша

С GKE Inference Gateway и llm-d решения о маршрутизации принимаются на основе сигналов обслуживания в реальном времени от наших экземпляров vLLM, а не статического чередования round-robin. GKE Inference Gateway непрерывно отслеживает глубину очереди в реальном времени и использование кэша ключ-значение (KV) в наших конечных точках vLLM. Когда обслуживающий под занят декодированием вывода длинного уравнения, шлюз направляет входящие запросы на узлы со свободными мощностями, поэтому тяжелые промпты больше не скапливаются за занятыми воркерами, пока соседние узлы простаивают.

В то же время шлюз направляет запросы с совпадающими научными инструкциями на тот узел обслуживания, где кэш KV этого префикса уже прогрет в памяти ускорителя. Поскольку 70–80% каждого промпта составляют общий контекст инструкций, в то время как генерируемые траектории рассуждений различаются от запроса к запросу, сочетание локальности KV-кэша с балансировкой очереди в реальном времени позволяет нашим подам пропускать повторяющуюся работу по префиллингу без перегрузки какого-либо отдельного реплика.

Более высокая пропускная способность, предсказуемые запуски обучения

Внедрение GKE Inference Gateway на базе llm-d изменило скорость и предсказуемость нашего конвейера генерации данных. В наших производственных кластерах интеллектуальная маршрутизация обеспечила:

В 4 раза более высокую пропускную способность парка: совокупная пропускная способность подскочила с 16 964 до 69 114 токенов в секунду на 1000 ускорителей, увеличив ежедневный объем генерации с 1,47 миллиарда до почти 6 миллиардов сгенерированных токенов.
Увеличение производительности в 3,5–5,5 раза для отдельных пулов моделей: балансировка очередей в реальном времени поддерживала равномерную нагрузку на узлы обслуживания на протяжении каждого прогона, предотвращая замедления в середине прогона, вызванные длинными выводами рассуждений.
На 14 процентных пунктов выше коэффициент попадания в префикс-кэш: маршрутизация промптов с совпадающими научными инструкциями на прогретые кэши позволила избежать избыточной работы по префиллу до того, как каждый запрос разветвлялся в собственное ветвление.

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

Что дальше для физического ИИ в ZenteiQ

Поскольку мы масштабируем наши научные базовые модели с 9B до 35B параметров и выше, мы планируем использовать ту же настройку GKE Inference Gateway и маршрутизацию llm-d для посткабельного обучения с подкреплением (RL) и циклов оценки автоматических численных решателей. Поскольку каждый шаг в этих многошаговых развертываниях опирается на контекст предыдущих шагов, они еще сильнее зависят от балансировки очередей в реальном времени и прогретых префикс-кэшей.

Чтобы опробовать эти возможности на собственных рабочих нагрузках, изучите документацию GKE Inference Gateway и ознакомьтесь с открытым проектом llm-d. Вы также можете прочитать о том, как команда платформы Gemini сократила задержку на 35% с помощью GKE Inference Gateway.