Обучение с подкреплением (Reinforcement Learning, RL) становится неотъемлемой частью современного развития больших языковых моделей (LLM). В этой области эффективность обновления параметров напрямую влияет на общую производительность и скорость тренировки моделей с триллионными параметрами. В июле этого года команда Kimi представила модель Kimi K2, которая кардинально изменила подход к обновлению параметров, сократив время этого процесса с 10 минут до всего 20 секунд для модели с масштабом порядка триллиона параметров. Эта революционная оптимизация способствовала устранению значимого узкого места в процессе RL-тренировки и стала базой для новых стандартов в индустрии искусственного интеллекта. Рассмотрим подробно, как был достигнут такой уровень эффективности, какие технологические проблемы пришлось преодолеть и какие решения применить.
Основные архитектурные подходы в RL-тренировке крупных языковых моделей делятся на два типа: колокейты и дисагрегацию. Колокейты означают совместное использование ресурсов GPU для обучения и выводов с чередованием их использования, в то время как дисагрегация подразумевает полное разделение устройств GPU между тренировкой и инференсом. В обоих случаях критическим этапом является обновление параметров после каждой сессии обучения перед началом rollout - так называемая фаза синхронизации параметров. Затягивание этого процесса ведёт к простоям GPU, снижает эффективность и увеличивает длительность полного цикла обучения. Поэтому оптимизация обновления параметров является одной из самых актуальных задач в контексте RL.
Команда Kimi ориентировалась на ко-локейт архитектуру, поскольку она чаще используется в их внутренних экспериментах. При этом их подход отличается от большинства открытых решений: обучение и инференс развернуты в разных контейнерах одного физического сервера, что позволяет абсолютно изолировать эти процессы, снижая взаимное влияние и упрощая разработку. Такой дизайн повторяет подходы, применяемые в онлайн-деплоя Kimi, позволяя использовать сервисную инфраструктуру повторно. Однако этот метод влечёт сложность - отсутствие прямого взаимодействия процесс групп обучения и инференса усложняет передачу параметров. Для решения этой проблемы была разработана промежуточная компонента - checkpoint-engine, которая выполняет роль моста, обеспечивающего эффективный обмен параметрами между тренировкой и инференсом при минимальных изменениях в логике самих процессов.
Изначально при разработке версии k1.5 идея заключалась в использовании CUDA IPC (межпроцессного обмена памятью GPU) для передачи тензоров. Каждому ранку обучения соответствовал checkpoint-engine, который после тренировки распространял обновлённые параметры через сети или NVLink, упаковывая их в IPC хендлы и передавая инференсу. Передача по слоям и по экспертам (EP) позволяла избежать проблем с ограничениями GPU памяти. Такая архитектура в целом удовлетворяла требованиям, но с ростом масштабов моделей и числа устройств возникали серьезные проблемы с производительностью - процесс обновления параметров доходил до 10 минут, что было неприемлемо для постоянной работы.
Подробный анализ выявил несколько проблем: нестабильное использование GPU памяти из-за передачи по слоям и EP, возникающие из-за этого перебои; многочисленные мелкие асинхронные передачи данных, вызывавшие накладные расходы; длительная сериализация и десериализация IPC хендлов, а также последовательное выполнение операций передачи данных и обновления весов в vLLM, что препятствовало параллелизации. Для борьбы с этими недостатками была применена стратегия формирования больших буферов фиксированного размера - так называемых бакетов. Тензоры сгруппировались в буферы, что позволило сократить количество мелких коммуникаций и стабилизировать распределение памяти GPU. Использование одного совместно используемого буфера для передачи между checkpoint-engine и vLLM избавило от постоянных накладных расходов на IPC хендлы - первая операция передачи содержала хендлы, последующие - только метаинформацию. Кроме того, алгоритм сбора IPC хендлов оптимизировали, ограничив операции их объединения только первоначальным этапом и исключив глобальный all_gather для всех ранков, поскольку в vLLM к запросам обычно обращается единственный ранк.
Еще был внедрён метод двойного буферизации, позволяющий перекрывать передачу данных и обновление параметров в vLLM, создавая двухэтапный конвейер, что существенно повысило эффективную пропускную способность. Подобные улучшения сократили время обновления параметров с 10 минут до примерно 2 минут на крупных кластерах из GPU H800, что уже практически удовлетворяло внутренние требования к RL. Однако дальнейший анализ показал, что теоретически данный показатель может быть улучшен, учитывая, что пропускная способность межсвязей типа RDMA и NVLink достигает 100 ГиБ/с, что при 1 терабайте данных должно обеспечивать обновление менее чем за 10 секунд. Выяснилось, что основным узким местом является последовательное выполнение передачи данных из CPU в GPU (Host to Device, H2D) до этапа широковещательной передачи (Broadcast). Скорость H2D ограничена около 40-50 ГиБ/с, что не позволяет полностью раскрыть потенциал сети и PCIe.
К тому же, операции обновления весов в vLLM страдали из-за повторных вызовов Python-функций, синхронизации между CPU и GPU, особенно когда происходила активация expert weight update с использованием .item(), что делало процесс непредсказуемо медленным. Для решения этих проблем были внедрены три дополнительных оптимизации. Во-первых, делалась попытка пересечения процессов H2D и Broadcast, чтобы максимально эффективно использовать пропускную способность PCIe. Правда, практическая реализация показала, что при этом они конкурируют за PCIe, снижают взаимную производительность.
Второе и третье улучшения - кэширование часто запрашиваемых данных, таких как dict(self.named_parameters()) и отображения expert_map, что позволило сократить CPU-GPU синхронизацию и улучшить стабильность обновления. В итоге разработчики изменили логику распределения задач: сначала все узлы параллельно загружают данные H2D, используя множество PCIe-линий для максимального пропускания данных, а затем быстро выполняется Device-to-Device (D2D) копирование уже в GPU памяти для организации широковещательной передачи. Этот подход позволил эффективно раскрыть пропускную способность узлов, повысив общую скорость обновления параметров. Дополнительной проблемой стало ненадёжное поведение инференс-движка во время RL.
Падения работы инференса приводили к сбоям всей тренировки. Классический подход заключается в простом рестарте упавших процессов, но в RL при прямой передаче параметров из тренировки в инференс без хранения на диске такой метод исполняется с задержкой, связанной с повторной загрузкой чекпоинтов. Идеальным стало решение организации онлайн-обновления весов именно аварийно перезапускаемых экземпляров инференса. Для этого команда Kimi интегрировала mooncake-transfer-engine, которая позволила прямое чтение весов с CPU RDMA работающих экземпляров и передачу данных только упавшим узлам без дополнительного расхода GPU-памяти в активных экземплярах. Такой метод усовершенствовал отказоустойчивость системы - восстановление весов бинарно-феноменально снизилось до 40 секунд, что является очень быстрым для RL систем с такими масштабами моделей.
Отдельный аспект - ускоренный старт инференс-сервисов. В Kimi разработали подход предварительной загрузки весов в разделяемую память (/dev/shm), что было значительно быстрее чтения с диска, однако при больших моделях расходует немало памяти и требует времени на прогрев. Благодаря оптимизациям checkpoint-engine и возможностям параллельной регистрации чекпоинтов во время загрузки vLLM с его torch.compile и CUDA Capture Graph, время запуска удалось сократить до уровня совместимого с самим процессом старта vLLM, не увеличивая время ожидания. В результате в крупных сервисах внутренняя латентность запуска снизилась и пользовательский опыт улучшился.
Важной вехой для сообщества стала публикация open source версии checkpoint-engine. Она предоставляет удобный и гибкий интерфейс ParameterServer, который позволяет разработчикам легко реализовать обновление весов в различных инференс-движках. Кроме того, поддерживаются две основные стратегии обновления - полностью широковещательное (broadcast) и peer-to-peer (P2P), что позволяет эффективно использовать систему как в колокейтах, так и в разделённых по функциям кластерах. Поддерживается регистрация нескольких чекпоинтов с возможностью динамической смены моделей в онлайн и оффлайн условиях. Несмотря на то что P2P реализация сейчас достаточно базовая (с последовательным чтением и распространением от ранка 0), команда Kimi планирует совершенствовать этот функционал, включая параллельное чтение с разных ранков и выполнение аналогичных двухэтапных конвейерных процессов, повышая производительность и гибкость.
Таким образом, революционные решения в сфере эффективной передачи и обновления параметров, внедренные в Kimi K2, прокладывают путь для более масштабируемого и производительного RL тренинга в больших LLM. Высокооптимизированный checkpoint-engine и интеграция с vLLM открывают новые горизонты для исследователей и инженеров, стремящихся быстро и надёжно обучать и обновлять модели на самых последних аппаратных платформах. Возможность гибкого взаимодействия с инференс-движками и поддержка отказоустойчивости делают платформу Kimi удобным инструментом для гибких и масштабных проектов, а предоставление open source решений гарантирует развитие и совершенствование технологий сообща с мировой AI-сообществом. В эпоху, когда скорости обработки данных и обновления моделей критически важны для конкурентоспособности, работа команды Kimi демонстрирует, что внимание к инженерной архитектуре и глубокий анализ узких мест позволяет совершать реальные прорывы, значительно меняющие правила игры в обучении искусственного интеллекта. .
