Google добавила в llm-d чередование RL-задач на общих ускорителях

Google Cloud выпустила для llm-d стек кооперативного тайм-слайсинга, который позволяет чередовать фазы независимых задач постобучения больших языковых моделей с подкреплением на одних ускорителях. По результатам первых тестов компании, совокупный коэффициент занятости ускорителей вырос примерно с 40% до 70% без ухудшения сходимости и точности моделей.

Google Cloud представила механизм кооперативного тайм-слайсинга в инфраструктурном проекте llm-d. Он предназначен для распределённого постобучения больших языковых моделей с подкреплением, где генерация траекторий и обновление весов чередуются, оставляя часть ускорителей без работы.

Система рассматривает отдельные фазы RL-конвейера как планируемые операции. Клиентская библиотека предоставляет вызовы gRPC acquire() и yield(), с помощью которых задача запрашивает и освобождает исключительный доступ к группе ускорителей. Оркестратор после освобождения ресурсов передаёт их следующей задаче в очереди.

Переключение выполняет привилегированный Snapshot Agent: он приостанавливает процессы, переносит состояние устройства из памяти ускорителя в оперативную память узла, а затем восстанавливает состояние другой задачи. В каждый момент на ускорителях находится состояние только одной нагрузки. При обратном переключении выполнение продолжается без перезапуска контейнера, повторной инициализации фреймворка или загрузки модели с накопителя.

В выпущенный стек вошли Snapshot Agent, Accelerator Orchestrator и клиентские библиотеки Python; руководства по интеграции опубликованы на GitHub. Первым механизмом создания снимков служит cuda-checkpoint. Агент рассчитан как на Kubernetes, так и на автономную работу в средах bare metal и Slurm.

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