Пренос тежина за RL пост-тренинг за мање од 2 секунде

Недавно смо постигли ажурирање параметара између машина за 1,3 секунде за Kimi-K2 (1T параметара), преносећи тежине са 256 тренинг GPU-ова (BF16) на 128 инференцијских GPU-ова (FP8).

У асинхроном подешавању кроз појачано учење (reinforcement learning), тренинг и инференција раде на одвојеним 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.
    ...

Одредишна страна неће чак ни добити обавештење о преносу. Ово нам омогућава преносе са ниском латенцијом, високим пропусным опсегом и без копирања, које покрећу чворови за тренинг без икакве контролне логике на чворовима за инференцију.

Прикупљање метаподатака – Контролер прикупља метаподатке о параметрима са свих тренинг и инференцијских GPU-ова.

Израчунавање распореда – Контролер израчунава статички распоред преноса тежина, мапирајући који тренинг GPU шаље који параметар ком инференцијском GPU-у, и којим редоследом.

Дистрибуција распореда – Контролер шаље распоред свим тренинг GPU-овима.

Извршавање – Након сваког корака тренинга, контролер сигнализира тренинг GPU-овима да започну преносе.

Са дефинисаним високим током рада, кључни изазов је како ефикасно извршити преносе тежина на скали од билион параметара. Овде описујемо детаље путање извршавања.

Параметри у тренингу су дистрибуирани према FSDP поставкама. Коришћењем full_tensor(), сви GPU-ови у DeviceMesh-у могу да реконструишу пуни параметар, стога сви могу служити као извор за пренос тежина.

Више дисјунктних DeviceMesh-ева чини mesh групу. Пошто су DeviceMesh-еви у истој групи дисјунктни, њихови преноси се не међусобно не ометају и могу се извршавати потпуно паралелно. Између mesh група уводимо глобалну баријеру да бисмо наметнули редослед.

Пиплајн за пренос тежина преклапа копирања са домаћина на уређај, склапање тензора, фузију и квантизацију, као и RDMA преносе преко MeshGroups.

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

Host-to-device memcpy — Ако FSDP пребацује тежину на CPU

Припрема параметара — Реконструкција пуне тежине са full_tensor(), примена фузије пројекције, квантизација ако је потребно.

RDMA пренос — Упис без копирања у удаљену инференцијску GPU меморију

Глобална баријера — Након што су сви full_tensor() позиви завршени, синхронизација кроз mesh групе користећи GLOO преко Етернета.

У имплементацији одржавамо FIFO ред задатака за сваку фазу цјевовода. Када год почетак реда задатака заврши фазу, он се помера на крај реда следеће фазе.

full_tensor() и друге GPU операције уводе додатну употребу GPU меморије. Да бисмо избегли грешку недостатка меморије (out of memory), започињемо извршавање задатка само ако тренутни задаци у току заузимају мање привремене GPU меморије од подесивог прага.

Неколико дизајнерских избора чини наш систем знатно бржим за рад и лакшим за одржавање од уобичајених отворених решења.

Пренос тежина заснован на Rank0 прикупља тежине пре него што их дистрибуира до инференцијског чвора.

Комуникација од тачке до тачке преноси тежине са оба тренинг чвора на мрежне интерфејсе инференцијског чвора.

Уобичајени образац је да се сви параметри усмере кроз rank-0 GPU-ове: прикупљање на тренинг rank-0, слање на инференцијски rank-0, па поновно дистрибуирање. Ово брзо постаје уски грло, ограничено PCIe пропусним опсегом једног GPU-а и NIC-ом (нпр., 400 Gbps ≈ 50 GB/s).

Насупрот томе, наше поставка од тачке до тачке омогућава сваком тренинг GPU-у да шаље директно сваком инференцијском GPU-у, засићујући комплетну мрежну структуру уместо једне везе.

Пренос заснован на RPC-у користи бафер на страни инференције и позиве радника, док једнострани пренос уписује тежине директно без учешћа инференцијског радника.

Неки системи се ослањају на позивање методе update_weight() инференцијског механизма за сваки тензор. То значи интрузивне промене у инференцијском коду, плус рескио од RPC-ова, сериализације и координације контролне раван.

Са RDMA WRITE примитивом, ажурирамо тежине тихо у инференцијској GPU меморији, без додатних копија. Никаква порука контролне равни и никаква CPU контролна логика није укључена. Није потребна никаква модификација инференцијског механизма.

Процес преноса тежина може користити четири врсте хардверских ресурса: (1) Кретање података хост-уређај (2) GPU рачунање за фузију пројекције и квантизацију (3) RDMA мрежу за раван података (4) Етернет за контролну раван.

Наш дизајн дели задатке преноса тежина у фазе цјевовода, омогућавајући лако преклапање преко различитих хардверских ресурса.

Неке имплементације поново израчунавају распоред преноса у сваком кораку тренинга, више пута прикупљајући метаподатке и дистрибуирајући инструкције. Ово додаје непотребну латенцију контролне равни.

Наш распоред се израчунава једном при иницијализацији. Свака итерација тренинга једноставно поновно покреће план: контролер издаје сигнал „идри“, а GPU-ови прате своје унапред додељене руте. Извршавање је предвидљиво и лагано.

Мамљиво је заплести цео процес ажурирања тежина у једну монолитну функцију: прикупљање метаподатака, подударање имена, унутар-чворно прикупљање, фузија пројекције, квантизација, подрезивање комуникацијског свијета, међучворни мрежни пренос. Тешко је програмирати исправно, а још теже оптимизовати.

У нашем инжењерингу, одвајамо ове кораке као појединачне компоненте. Свака компонента се може јединично тестирати, разматрати и оптимизовати изоловано.

Брз и поуздан пренос тежина је кључни градивни блок за велико скалирање RL финог подешавања. Комбиновањем 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 групе

Цјевовод задатака

Контрола коришћења GPU меморије

Зашто је брзо и једноставно

Комуникација од тачке до тачке

Једнострани пренос података

Цјевовод (Pipelining)

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

Чисто раздвајање

Закључак

Цитат