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

Наскоро постигнахме актуализации на параметри между машини за 1,3 секунди за Kimi-K2 (1T параметри), като прехвърлихме тегла от 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 могат да възстановят пълния параметър, следователно всички могат да служат като източник за прехвърляне на тегла.

Множество несъвпадащи DeviceMeshes формират мрежова група. Тъй като DeviceMeshes в една и съща група са несъвпадащи, техните трансфери не си пречат и могат да се изпълняват напълно паралелно. Между мрежовите групи вмъкваме глобална бариера за налагане на ред.

Пайплайнът за прехвърляне на тегла припокрива копиранията host-to-device, сглобяването на тензори, сливането и квантизирането, както и RDMA прехвърлянията през MeshGroups.

Третираме прехвърлянето на всеки тензор от параметри като задача. Процесът на прехвърляне на тегла използва множество видове хардуерни източници, затова разделяме задачата за прехвърляне на тегла на различни етапи на конвейера, които се припокриват във времето:

Memcpy от хост към устройство — Ако FSDP извежда теглото към CPU

Подготовка на параметрите — Възстановяване на пълното тегло с full_tensor(), прилагане на сливане на проекции, квантоване при необходимост.

RDMA трансфер — Запис без копиране (zero-copy) в отдалечената памет на инферентния GPU

Глобална бариера — След като всички извиквания на full_tensor() приключат, синхронизиране през мрежовите групи с помощта на GLOO през Ethernet.

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

full_tensor() и други GPU операции въвеждат допълнително използване на GPU памет. За да избегнем грешка за липса на памет (out of memory), стартираме изпълнението на задача само ако текущите задачи в ход заемат по-малко временна GPU памет от конфигурируем праг (watermark).

Няколко дизайнерски решения правят нашата система значително по-бърза за изпълнение и по-лесна за поддръжка от общите решения с отворен код.

Прехвърлянето на тегла, базирано на Rank0, събира теглата, преди да ги разпредели към inference възела.

Комуникацията point-to-point прехвърля теглата от двата тренировъчни възела към мрежовите интерфейси на inference възела.

Често срещан модел е да се прекарват всички параметри през GPU с ранг-0: събиране на тренировъчен ранг-0, изпращане към инферентен ранг-0, след което разпръскване отново. Това бързо се превръща в тесно място, ограничено от PCIe честотната лента и мрежовата карта на едно GPU (напр. 400 Gbps ≈ 50 GB/s).

За разлика от това, нашата точка-до-точка настройка позволява на всеки тренировъчен GPU да изпраща директно към всеки инферентен GPU, насищайки цялата мрежова структура, вместо само една връзка.

Прехвърлянето, базирано на RPC, използва буфер от страната на inference и извиквания от работниците, докато едностранното прехвърляне записва теглата директно, без участие на inference worker-а.

Някои системи разчитат на извикване на метода update_weight() на инферентния двигател за всеки тензор. Това означава инвазивни промени в инферентния код, плюс надхвърляне (overhead) от 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 и мрежови групи

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

Контрол на използването на GPU памет

Защо е бързо и просто

Точка-до-точка комуникация

Едностранен прехвърляне на данни

Конвейерна обработка (Pipelining)

Статичен график

Чисто разделяне

Заключение

Цитиране