Передача весов для RL пост-обучения менее чем за 2 секунды

Недавно мы добились обновления параметров между машинами за 1,3 секунды для Kimi-K2 (1 трлн параметров), передав веса с 256 обучающих GPU (BF16) на 128 GPU для инференса (FP8).

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

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

Наше решение основано на RDMA WRITE — одностороннем примитиве, при котором источник выполняет прямую запись в память GPU назначения.

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

Узел назначения даже не получает уведомления о передаче. Это обеспечивает низкую задержку, высокую пропускную способность и передачу с нулевым копированием (zero-copy), управляемую обучающими узлами без какой-либо логики управления на узлах инференса.

Сбор метаданных — контроллер собирает метаданные параметров со всех обучающих и инференс-GPU.

Вычисление расписания — контроллер вычисляет статическое расписание передачи весов, определяя, какой обучающий GPU передает определенный параметр на конкретный инференс-GPU и в какой последовательности.

Распространение расписания — контроллер отправляет расписание на все обучающие GPU.

Выполнение — после каждого шага обучения контроллер подает сигнал обучающим GPU начать передачу.

После определения высокоуровневого рабочего процесса ключевой задачей становится эффективное выполнение передачи весов в масштабе триллионов параметров. Здесь мы описываем детали процесса выполнения.

Параметры при обучении распределяются согласно размещению FSDP. Используя full_tensor(), все GPU в DeviceMesh могут восстановить полный параметр, поэтому все они могут служить источниками для передачи весов.

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

Конвейер передачи весов перекрывает копирование с хоста на устройство, сборку тензоров, слияние и квантизацию, а также передачи RDMA между MeshGroup.

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

Копирование из хоста на устройство (memcpy) — если FSDP выгружает веса в CPU

Подготовка параметров — восстановление полных весов с помощью full_tensor(), применение проекционного слияния, квантование при необходимости.

Передача по RDMA — запись в память удаленного инференс-GPU с нулевым копированием

Глобальный барьер — после завершения всех вызовов full_tensor() выполняется синхронизация между группами сетей с использованием GLOO через Ethernet.

При реализации мы поддерживаем очередь FIFO для задач на каждой стадии конвейера. Как только задача в начале очереди завершает стадию, она перемещается в конец очереди следующей стадии.

Функция full_tensor() и другие операции GPU вызывают дополнительное использование памяти GPU. Чтобы избежать ошибки нехватки памяти (OOM), мы запускаем задачу, только если текущие выполняемые задачи занимают меньше временной памяти GPU, чем настроенный порог (watermark).

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

Передача весов на основе Rank0 сначала собирает веса, а затем распределяет их на инференс-узел.

Обмен по схеме point-to-point передаёт веса из обоих тренировочных узлов в сетевые интерфейсы инференс-узла.

Обычно параметры проходят через GPU с рангом 0: сбор на обучающем узле 0, передача на инференс-узел 0, затем повторное распределение. Это быстро становится узким местом, ограниченным пропускной способностью PCIe и сетевой карты одного GPU (например, 400 Гбит/с ≈ 50 ГБ/с).

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

Передача на основе RPC использует буфер на стороне инференса и вызовы worker, тогда как односторонняя передача записывает веса напрямую без участия inference-worker.

Некоторые системы полагаются на вызов метода update_weight() движка инференса для каждого тензора. Это означает интрузивные изменения кода инференса, а также накладные расходы на RPC, сериализацию и координацию плоскости управления.

Благодаря примитиву RDMA WRITE мы обновляем веса без лишнего шума непосредственно в памяти инференс-GPU, без лишнего копирования. Никакие сообщения плоскости управления и логика управления на стороне CPU не задействуются. Модификация движка инференса не требуется.

Процесс передачи весов может использовать четыре типа аппаратных ресурсов: (1) Перемещение данных между хостом и устройством (2) Вычисления GPU для проекционного слияния и квантования (3) Сеть RDMA для плоскости данных (4) Ethernet для плоскости управления.

Наш дизайн разделяет задачи передачи весов на стадии конвейера, обеспечивая простое перекрытие использования различных аппаратных ресурсов.

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

Наше расписание вычисляется один раз при инициализации. Каждая итерация обучения просто воспроизводит план: контроллер подает сигнал «пуск», и GPU следуют заранее заданным маршрутам. Выполнение предсказуемо и требует малых ресурсов.

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

В нашей разработке мы разделили эти этапы на отдельные компоненты. Каждый компонент можно изолированно протестировать, проанализировать и оптимизировать.

Быстрая и надежная передача весов — критически важный блок для масштабного дообучения с подкреплением. Комбинируя примитив RDMA WRITE, статическое расписание передачи и конвейерное выполнение, мы сократили обновление триллионных параметров до 1,3 секунды на Kimi-K2. Этот подход прост для понимания, удобен в обслуживании и позволяет избежать «узких мест», характерных для традиционных архитектур.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: односторонние передачи

Общий рабочий процесс

Выполнение передачи весов

DeviceMesh и группы сетей (Mesh Groups)

Конвейер задач

Контроль использования памяти GPU

Почему это быстро и просто

Одноранговая коммуникация

Односторонняя передача данных

Конвейеризация

Статическое расписание

Четкое разделение

Заключение

Цитирование