Transferència de pesos per a postentrenament d'AR en menys de 2 segons

Recentment hem aconseguit actualitzacions de paràmetres entre màquines en 1,3 segons per a Kimi-K2 (1 bilió de paràmetres), transferint pesos de 256 GPU d'entrenament (BF16) a 128 GPU d'inferència (FP8).

En l'ajust fi d'aprenentatge per reforç asíncron, l'entrenament i la inferència s'executen en GPU separades. Després de cada pas d'entrenament, s'han de enviar nous pesos als nodes d'inferència. Molts marcs existents triguen diversos segons, o fins i tot minuts, per a models de bilions de paràmetres.

Aprofitant la comunicació punt a punt RDMA, podem fer que la transferència de pesos sigui extremadament ràpida, sense modificar el motor d'inferència, i fer que el codi sigui més fàcil d'escriure i mantenir.

La nostra solució es basa en RDMA WRITE, una primitiva unilateral on l'origen escriu directament a la memòria de la GPU de destinació.

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.
    ...

El costat de destinació ni tan sols rebrà una notificació de la transferència. Això ens proporciona transferències de zero còpies, baixes latències i alt rendiment impulsades pels nodes d'entrenament sense cap lògica de control als nodes d'inferència.

Recopilació de metadades: el controlador recull metadades de paràmetres de totes les GPU d'entrenament i d'inferència.

Càlcul de planificació: el controlador calcula una planificació estàtica de transferència de pesos, assignant quina GPU d'entrenament envia quin paràmetre a quina GPU d'inferència i en quin ordre.

Distribució de planificació: el controlador envia la planificació a totes les GPU d'entrenament.

Execució: després de cada pas d'entrenament, el controlador indica a les GPU d'entrenament que iniciïn les transferències.

Un cop definit el flux de treball d'alt nivell, el repte principal és com executar les transferències de pesos de manera eficient a escala de bilions de paràmetres. Aquí descrivim els detalls del camí d'execució.

Els paràmetres en l'entrenament es distribueixen segons les ubicacions de FSDP. Utilitzant full_tensor(), totes les GPU d'un DeviceMesh poden reconstruir el paràmetre complet, per la qual cosa totes poden servir com a origen per a la transferència de pesos.

Múltiples DeviceMesh disjunts formen un grup de malla. Com que els DeviceMesh del mateix grup són disjunts, les seves transferències no interfereixen i es poden executar completament en paral·lel. Entre grups de malla, inserim una barrera global per aplicar l'ordenació.

La canalització de transferència de pesos solapa les còpies de l’amfitrió al dispositiu, el muntatge de tensors, la fusió i la quantificació, i les transferències RDMA a través de MeshGroups.

Tractem la transferència de cada tensor de paràmetres com una tasca. El procés de transferència de pesos utilitza diversos tipus de recursos de maquinari, per la qual cosa dividim una tasca de transferència de pesos en diferents etapes de pipeline que se superposen en el temps:

Còpia de memòria de l'amfitrió al dispositiu (memcpy): si FSDP descarrega el pes a la CPU

Preparació de paràmetres: reconstrueix el pes complet amb full_tensor(), aplica la fusió de projecció i quantifica si cal.

Transferència RDMA: escriptura de zero còpies a la memòria remota de la GPU d'inferència

Barrera global: un cop realitzades totes les crides a full_tensor(), se sincronitza entre els grups de malla mitjançant GLOO a través d'Ethernet.

En la implementació, mantenim una cua FIFO de tasques per a cada etapa del pipeline. Sempre que la tasca del cap de la cua completi l'etapa, es mou a la cua de la cua de la següent etapa.

full_tensor() i altres operacions de GPU introdueixen un ús addicional de memòria de la GPU. Per evitar un error de manca de memòria (OOM), iniciem l'execució d'una tasca només si les tasques actuals en marxa ocupen menys memòria temporal de GPU que un llindar configurable.

Diverses decisions de disseny fan que el nostre sistema sigui significativament més ràpid d'executar i més fàcil de mantenir que les solucions de codi obert habituals.

La transferència de pesos basada en Rank0 recopila els pesos abans de distribuir-los al node d’inferència.

La comunicació punt a punt transfereix els pesos de tots dos nodes d’entrenament a les interfícies de xarxa del node d’inferència.

Un patró habitual és canalitzar tots els paràmetres a través de les GPU de rang 0: recopilar al rang 0 d'entrenament, enviar al rang 0 d'inferència i tornar a distribuir. Això es converteix ràpidament en un punt d'estrangulament, limitat per l'amplada de banda PCIe d'una sola GPU i la targeta de xarxa (per exemple, 400 Gbps ≈ 50 GB/s).

En contrast, la nostra configuració punt a punt permet que cada GPU d'entrenament enviï directament a cada GPU d'inferència, saturant tota la infraestructura de xarxa en lloc d'un sol enllaç.

La transferència basada en RPC fa servir una memòria intermèdia del costat de la inferència i crides de treballador, mentre que la transferència d’una sola banda escriu els pesos directament sense participació del treballador d’inferència.

Alguns sistemes depenen de cridar al mètode update_weight() del motor d'inferència per a cada tensor. Això significa canvis intrusius al codi d'inferència, a més de la sobrecàrrega de les RPC, la serialització i la coordinació del pla de control.

Amb la primitiva RDMA WRITE, actualitzem els pesos silenciosament a la memòria de la GPU d'inferència, sense còpies addicionals. No hi ha cap missatge de pla de control ni lògica de control de CPU implicada. No cal cap modificació al motor d'inferència.

El procés de transferència de pesos pot aprofitar quatre tipus de recursos de maquinari: (1) moviment de dades amfitrió-dispositiu (2) càlcul de GPU per a la fusió de projecció i quantificació (3) xarxa RDMA per al pla de dades (4) Ethernet per al pla de control.

El nostre disseny divideix les tasques de transferència de pesos en etapes de pipeline, permetent una fàcil superposició entre diferents recursos de maquinari.

Algunes implementacions recalculen una planificació de transferència a cada pas d'entrenament, recopilant repetidament metadades i distribuint instruccions. Això afegeix una latència innecessària al pla de control.

La nostra planificació es calcula un cop a la inicialització. Cada iteració d'entrenament simplement reprodueix el pla: el controlador emet un senyal de «començar» i les GPU segueixen les seves rutes preassignades. L'execució és predictible i lleugera.

És temptador enredar tot el procés d'actualització de pesos en una sola funció monolítica: recopilar metadades, conciliar noms, recopilació intranode, fusió de projecció, quantificació, subsegmentació del món de comunicació, transferència de xarxa internode. És difícil de programar correctament i encara més difícil d'optimitzar.

En la nostra enginyeria, separem aquests passos com a components individuals. Cada component es pot provar de manera unitària, analitzar i optimitzar de manera aïllada.

Una transferència de pesos ràpida i fiable és un bloc de construcció crucial per a l'ajust fi d'AR a gran escala. Combinant la primitiva RDMA WRITE, una planificació de transferència estàtica i una execució en pipeline, hem reduït les actualitzacions de bilions de paràmetres a només 1,3 segons a Kimi-K2. L'enfocament és fàcil d'entendre, senzill de mantenir i evita els colls d'ampolla dels dissenys tradicionals.

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: transferències unilaterals

Flux de treball d'alt nivell

Execució de la transferència de pesos

DeviceMesh i grups de malles

Pipelines de tasques

Control d'ús de memòria de la GPU

Per què és ràpid i senzill

Comunicació punt a punt

Transferència de dades unilateral

Pipelines

Planificació estàtica

Separació neta

Conclusió

Citació