Transfer de ponderi pentru post-antrenare RL în sub 2 secunde

Am obținut recent actualizări de parametri între mașini în 1,3 secunde pentru Kimi-K2 (1T parametri), transferând ponderi de la 256 de GPU-uri de antrenare (BF16) la 128 de GPU-uri de inferență (FP8).

În reglarea fină prin învățare prin consolidare asincronă, antrenarea și inferența rulează pe GPU-uri separate. După fiecare pas de antrenare, noile ponderi trebuie trimise către nodurile de inferență. Multe cadre existente necesită câteva secunde sau chiar minute pentru modelele cu un trilion de parametri.

Prin utilizarea comunicării punct la punct RDMA, reușim să facem transferul de ponderi extrem de rapid, fără a schimba motorul de inferență și facem codul mai ușor de scris și de întreținut.

Soluția noastră se bazează pe RDMA WRITE, o primitivă unidirecțională în care sursa scrie direct în memoria GPU a destinației.

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

Partea de destinație nici măcar nu va fi notificată cu privire la transfer. Acest lucru ne oferă transferuri cu latență redusă, debit ridicat, fără copiere, conduse de nodurile de antrenare fără nicio logică de control pe nodurile de inferență.

Colectarea de metadate – Controlerul adună metadatele parametrilor de la toate GPU-urile de antrenare și inferență.

Calculul programului – Controlerul calculează un program static de transfer de ponderi, mapând ce GPU de antrenare trimite ce parametru către ce GPU de inferență și în ce ordine.

Distribuirea programului – Controlerul trimite programul către toate GPU-urile de antrenare.

Execuție – După fiecare pas de antrenare, controlerul semnalează GPU-urilor de antrenare să înceapă transferurile.

Odată definit fluxul de lucru la nivel înalt, provocarea principală este cum să se execute eficient transferurile de ponderi la scara unui trilion de parametri. Aici descriem detaliile căii de execuție.

Parametrii din antrenare sunt distribuiți conform plasărilor FSDP. Folosind full_tensor(), toate GPU-urile dintr-un DeviceMesh pot reconstrui parametrul complet, putând servi astfel ca sursă pentru transferul de ponderi.

Mai multe DeviceMesh-uri disjuncte formează un grup de rețele (mesh group). Deoarece DeviceMesh-urile din același grup sunt disjuncte, transferurile lor nu interferează și pot rula complet în paralel. Între grupurile de rețele, inserăm o barieră globală pentru a impune ordonarea.

Pipelina de transfer a ponderilor suprapune copiile de la gazdă la dispozitiv, asamblarea tensorilor, fuziunea și cuantizarea și transferurile RDMA între MeshGroups.

Tratăm transferul fiecărui tensor de parametri ca pe o sarcină. Procesul de transfer de ponderi utilizează mai multe tipuri de surse hardware, motiv pentru care împărțim o sarcină de transfer de ponderi în diferite etape de conductă care se suprapun în timp:

Memcpy de la gazdă la dispozitiv — Dacă FSDP descarcă ponderea pe CPU

Pregătirea parametrilor — Reconstruirea ponderii complete cu full_tensor(), aplicarea fuziunii proiecțiilor, cuantificarea dacă este necesar.

Transfer RDMA — Scriere fără copiere în memoria GPU de inferență la distanță

Barieră globală — După ce toate apelurile full_tensor() s-au terminat, sincronizarea între grupurile de rețele utilizând GLOO prin Ethernet.

În implementare, menținem o coadă FIFO de sarcini pentru fiecare etapă a conductei. Ori de câte ori sarcina din capul cozii finalizează etapa, aceasta este mutată la coada etapei următoare.

full_tensor() și alte operațiuni GPU introduc o utilizare suplimentară a memoriei GPU. Pentru a evita eroarea de memorie insuficientă (out of memory), începem execuția unei sarcini numai dacă sarcinile curente în desfășurare ocupă mai puțină memorie GPU temporară decât un prag (watermark) configurabil.

Mai multe alegeri de design fac ca sistemul nostru să ruleze semnificativ mai rapid și să fie mai ușor de întreținut decât soluțiile obișnuite open-source.

Transferul de ponderi bazat pe Rank0 colectează ponderile înainte de a le distribui către nodul de inferență.

Comunicarea punct-la-punct transferă ponderile de la ambele noduri de antrenare către interfețele de rețea ale nodului de inferență.

Un model comun este direcționarea tuturor parametrilor prin GPU-urile de rang 0: colectarea pe rangul 0 de antrenare, trimiterea către rangul 0 de inferență, apoi redistribuirea. Acest lucru devine rapid un punct de blocaj, limitat de lățimea de bandă PCIe a unui singur GPU și de placa NIC (de exemplu, 400 Gbps ≈ 50 GB/s).

În schimb, configurația noastră punct la punct permite fiecărui GPU de antrenare să trimită direct către fiecare GPU de inferență, saturând întreaga țesătură de rețea în loc de o singură legătură.

Transferul bazat pe RPC folosește un buffer pe partea de inferență și apeluri de worker, în timp ce transferul unidirecțional scrie ponderile direct, fără participarea worker-ului de inferență.

Unele sisteme se bazează pe apelarea metodei update_weight() a motorului de inferență pentru fiecare tensor. Aceasta înseamnă modificări invazive ale codului de inferență, plus overhead de la apelurile RPC, serializare și coordonarea planului de control.

Cu primitiva RDMA WRITE, actualizăm ponderile în mod silențios în memoria GPU de inferență, fără copii suplimentare. Nu este implicat niciun mesaj în planul de control și nicio logică de control CPU. Nu este necesară nicio modificare a motorului de inferență.

Procesul de transfer de ponderi poate utiliza patru tipuri de resurse hardware: (1) Deplasarea datelor gazdă-dispozitiv (2) Calculul GPU pentru fuziunea proiecțiilor și cuantificare (3) Rețeaua RDMA pentru planul de date (4) Ethernet pentru planul de control.

Designul nostru a împărțit sarcinile de transfer de ponderi în etape de conductă, permițând o suprapunere ușoară între diferite resurse hardware.

Unele implementări reec calculează un program de transfer la fiecare pas de antrenare, colectând în mod repetat metadate și distribuind instrucțiuni. Acest lucru adaugă o latență nenecesară în planul de control.

Programul nostru este calculat o singură dată la inițializare. Fiecare iterație de antrenare pur și simplu reia planul: controlerul emite un semnal de „pornire”, iar GPU-urile își urmează rutele preatribuite. Execuția este predictibilă și ușoară.

Este tentant să încâlcești întregul proces de actualizare a ponderilor într-o singură funcție monolitică: colectarea de metadate, potrivirea denumirilor, colectarea intra-nod, fuziunea proiecțiilor, cuantificarea, subslicing-ul universului de comunicare, transferul de rețea inter-nod. Este greu de programat corect și și mai greu de optimizat.

În ingineria noastră, separăm aceste etape ca componente individuale. Fiecare componentă poate fi testată Unitar, analizată și optimizată izolat.

Transferul rapid și fiabil de ponderi este o componentă de bază esențială pentru reglarea fină RL la scară largă. Combinând primitiva RDMA WRITE, un program de transfer static și execuția în conductă (pipelined), am redus actualizările pentru un trilion de parametri la doar 1,3 secunde pe Kimi-K2. Abordarea este ușor de înțeles, ușor de întreținut și evită blocajele modelelor tradiționale.

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: transferuri unidirecționale

Flux de lucru la nivel înalt

Executarea transferului de ponderi

DeviceMesh și grupuri de rețele (Mesh Groups)

Conductă de sarcini (Task pipeline)

Controlul utilizării memoriei GPU

De ce este rapid și simplu

Comunicare punct la punct

Transfer de date unidirecțional

Conductă (Pipelining)

Program static

Separare curată

Concluzie

Citare