2 секундтан аз уақытта RL кейінгі оқытуға арналған салмақты тасымалдау

Біз жақында Kimi-K2 (1T параметрлері) үшін 1,3 секундтық машинааралық параметр жаңартуларына қол жеткіздік, бұл ретте салмақтар 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.
    ...

Мақсат тіпті тасымалдау туралы хабарландыру алмайды. Бұл бізге инференция түйіндерінде басқару логикасынсыз оқыту түйіндері арқылы басқарылатын төмен кідірісті, жоғары өткізу қабілетін және нөлдік көшіру арқылы тасымалдауды қамтамасыз етеді.

Метадеректерді жинау – Контроллер барлық оқыту және инференция GPU құрылғыларынан параметр метадеректерін жинайды.

Жоспарды есептеу – Контроллер қай оқыту GPU қай параметрді қай инференция GPU құрылғысына және қандай ретпен жіберетінін көрсететін статикалық салмақты тасымалдау кестесін есептейді.

Жоспарды тарату – Контроллер кестені барлық оқыту GPU құрылғыларына жібереді.

Орындау – Әрбір оқыту қадамынан кейін контроллер оқыту GPU құрылғыларына тасымалдауды бастау туралы сигнал береді.

Жоғары деңгейдегі жұмыс процесі анықталғаннан кейін, негізгі міндет – триллион параметрлік ауқымда салмақты тасымалдауды тиімді орындау. Мұнда біз орындау жолының мәліметтерін сипаттаймыз.

Оқытудағы параметрлер FSDP орналастыруларына сәйкес бөлінген. full_tensor() көмегімен DeviceMesh ішіндегі барлық GPU толық параметрді қайта құра алады, демек, олардың барлығы салмақты тасымалдау үшін деректер көзі бола алады.

Бір-бірімен қиылыспайтын бірнеше DeviceMesh mesh тобын құрайды. Бір топтағы DeviceMesh құрылғылары қиылыспайтындықтан, олардың тасымалдау әрекеттері бір-біріне кедергі жасамайды және толығымен параллель жұмыс істей алады. Mesh топтары арасында біз реттілікті қамтамасыз ету үшін жаһандық тосқауыл енгіземіз.

Салмақты тасымалдау құбыры host-to-device көшірмелерін, тензорларды жинауды, біріктіру мен кванттауды және MeshGroups арқылы RDMA тасымалдарын қабаттастырады.

Біз әрбір параметр тензорын тасымалдауды тапсырма ретінде қарастырамыз. Салмақты тасымалдау процесі аппараттық көздердің бірнеше түрін пайдаланады, сондықтан біз салмақты тасымалдау тапсырмасын уақыт бойынша қиылысатын әртүрлі конвейер кезеңдеріне бөлеміз:

Хосттан құрылғыға memcpy — Егер FSDP салмақты CPU-ға түсірсе

Параметрді дайындау — full_tensor() көмегімен толық салмақты қалпына келтіру, проекциялық біріктіруді қолдану, қажет болса кванттау.

RDMA тасымалдауы — Қашықтағы инференция GPU жадына нөлдік көшіру арқылы жазу

Жаһандық тосқауыл — Барлық full_tensor() шақырулары аяқталғаннан кейін Ethernet арқылы GLOO көмегімен mesh топтары бойынша синхрондау.

Іске асыру барысында біз әр конвейер кезеңі үшін FIFO тапсырмалар кезегін сақтаймыз. Кезек басындағы тапсырма кезеңді аяқтаған сайын, ол келесі кезең кезегінің соңына жылжытылады.

full_tensor() және басқа GPU операциялары қосымша GPU жадын пайдалануды тудырады. Жадтың жетіспеушілігі қатесін болдырмау үшін біз тапсырманың орындалуын ағымдағы орындалып жатқан тапсырмалар конфигурацияланатын су белгісінен аз уақытша GPU жадын алған жағдайда ғана бастаймыз.

Бірнеше дизайн шешімдері біздің жүйемізді жалпыға қолжетімді ашық бастапқы кодты шешімдерге қарағанда әлдеқайда жылдам жұмыс істеуге және техникалық қызмет көрсетуді жеңілдетуге мүмкіндік береді.

Rank0 негізіндегі салмақ тасымалдау оларды inference түйініне тарату алдында салмақтарды жинайды.

Нүктеден нүктеге байланыс салмақтарды екі training түйінінен де inference түйінінің желілік интерфейстеріне тасымалдайды.

Жалпы үлгі – барлық параметрлерді 0-ші дәрежелі GPU арқылы өткізу: оқытудың 0-ші дәрежесінде жинау, инференцияның 0-ші дәрежесіне жіберу, содан кейін қайтадан шашырату. Бұл бір GPU-ның PCIe өткізу қабілеті мен NIC (мысалы, 400 Гбит/с ≈ 50 ГБ/с) шектелетін шектеулі нүктеге айналады.

Керісінше, біздің нүкте-нүкте конфигурациямыз әрбір оқу GPU-ының әрбір инференс GPU-ына тікелей дерек жіберуіне мүмкіндік береді, осылайша бір ғана қосылуды емес, бүкіл желі матасын толық жүктейді.

RPC-негізіндегі тасымалдау inference жағындағы буферді және worker шақыруларын пайдаланады, ал біржақты тасымалдау салмақтарды inference-worker қатысуынсыз тікелей жазады.

Кейбір жүйелер әрбір тензор үшін инференция қозғалтқышының update_weight() әдісін шақыруға сүйенеді. Бұл инференция кодына тікелей өзгерістер енгізуді, сондай-ақ RPC, сериализация және басқару жазықтығын үйлестіруден қосымша шығындарды білдіреді.

RDMA WRITE примитивінің көмегімен біз қосымша көшірмелерсіз инференция GPU жадындағы салмақтарды үнсіз жаңартамыз. Басқару жазықтығының хабарламасы және CPU басқару логикасы қатыспайды. Инференция қозғалтқышына ешқандай өзгерту қажет емес.

Салмақты тасымалдау процесі аппараттық ресурс түрлерінің төрт түрін пайдалана алады: (1) Хост пен құрылғы арасындағы деректерді жылжыту (2) Проекцияны біріктіру және кванттау үшін GPU есептеулері (3) Деректер жазықтығы үшін RDMA желісі (4) Басқару жазықтығы үшін Ethernet.

Біздің дизайнымыз салмақты тасымалдау тапсырмаларын конвейерлік кезеңдерге бөледі, бұл әртүрлі аппараттық ресурстар арасында оңай қиылысуға мүмкіндік береді.

Кейбір іске асырулар әрбір оқыту қадамында тасымалдау кестесін қайта есептейді, метадеректерді қайталап жинайды және нұсқауларды таратады. Бұл қажетсіз басқару жазықтығының кідірісін қосады.

Біздің кесте инициализация кезінде бір рет есептеледі. Әрбір оқыту итерациясы жоспарды жай ғана қайталайды: контроллер «жүр» сигналын береді, ал GPU өздеріне алдын ала тағайындалған маршруттарды орындайды. Орындалуы болжамды және жеңіл.

Бүкіл салмақты жаңарту процесін бір монолитті функцияға біріктіру қызықты көрінуі мүмкін: метадеректерді жинау, атауларды сәйкестендіру, түйінішілік жинау, проекцияны біріктіру, кванттау, коммуникация әлемін қосалқы кесу, түйінаралық желіні тасымалдау. Оны дұрыс бағдарламалау қиын, ал оңтайландыру одан да қиын.

Біздің инженерлік тәжірибемізде біз бұл қадамдарды жеке компоненттер ретінде бөлеміз. Әрбір компонентті бөлек сынақтан өткізуге, талдауға және оңтайландыруға болады.

Жылдам, сенімді салмақты тасымалдау – бұл ауқымды RL дәл баптау үшін маңызды құрылыс блогы. RDMA WRITE примитивін, статикалық тасымалдау кестесін және конвейерлік орындалуды біріктіре отырып, біз Kimi-K2 жүйесінде триллион параметрлік жаңартуларды небары 1,3 секундқа дейін қысқарттық. Бұл тәсілді түсіну оңай, техникалық қызмет көрсету ыңғайлы және дәстүрлі конструкциялардың кедергілерін болдырмайды.

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 жадын пайдалануды басқару

Неліктен ол жылдам әрі қарапайым

Нүктеден нүктеге коммуникация

Бір жақты деректерді тасымалдау

Конвейерлеу

Статикалық кесте

Таза бөлу

Қорытынды

Дәйексөз