Trasferimento pesi per post-training RL in meno di 2 secondi

Abbiamo recentemente ottenuto aggiornamenti dei parametri tra macchine in 1,3 secondi per Kimi-K2 (1T di parametri), trasferendo i pesi da 256 GPU di addestramento (BF16) a 128 GPU di inferenza (FP8).

Nel fine-tuning dell'apprendimento per rinforzo asincrono, l'addestramento e l'inferenza vengono eseguiti su GPU separate. Dopo ogni passaggio di addestramento, i nuovi pesi devono essere inviati ai nodi di inferenza. Molti framework esistenti impiegano diversi secondi, o addirittura minuti, per modelli con trilioni di parametri.

Sfruttando la comunicazione punto-a-punto RDMA, siamo in grado di rendere il trasferimento dei pesi estremamente rapido, senza modificare il motore di inferenza, e di rendere il codice più facile da scrivere e gestire.

La nostra soluzione si basa su RDMA WRITE, una primitiva unidirezionale in cui la sorgente scrive direttamente nella memoria GPU di destinazione.

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

Il lato di destinazione non riceverà nemmeno una notifica per il trasferimento. Questo ci garantisce trasferimenti a bassa latenza, throughput elevato e zero-copy gestiti dai nodi di addestramento senza alcuna logica di controllo sui nodi di inferenza.

Raccolta dei metadati – Il controller raccoglie i metadati dei parametri da tutte le GPU di addestramento e di inferenza.

Calcolo della pianificazione – Il controller calcola una pianificazione statica per il trasferimento dei pesi, definendo quale GPU di addestramento invia quale parametro a quale GPU di inferenza e in quale ordine.

Distribuzione della pianificazione – Il controller invia la pianificazione a tutte le GPU di addestramento.

Esecuzione – Dopo ogni passaggio di addestramento, il controller segnala alle GPU di addestramento di avviare i trasferimenti.

Definito il flusso di lavoro di alto livello, la sfida principale è come eseguire i trasferimenti dei pesi in modo efficiente su scala di trilioni di parametri. Qui descriviamo i dettagli del percorso di esecuzione.

I parametri nell'addestramento sono distribuiti secondo i posizionamenti FSDP. Utilizzando full_tensor(), tutte le GPU in una DeviceMesh possono ricostruire l'intero parametro, potendo quindi fungere da sorgente per il trasferimento dei pesi.

Più DeviceMesh disgiunte formano un gruppo di mesh. Poiché le DeviceMesh nello stesso gruppo sono disgiunte, i loro trasferimenti non interferiscono e possono essere eseguiti completamente in parallelo. Tra i gruppi di mesh, inseriamo una barriera globale per imporre l'ordinamento.

La pipeline di trasferimento dei pesi sovrappone le copie host-to-device, l'assemblaggio dei tensor, la fusione e la quantizzazione, e i trasferimenti RDMA tra MeshGroups.

Trattiamo il trasferimento di ogni tensore di parametri come un'attività. Il processo di trasferimento dei pesi utilizza molteplici tipi di risorse hardware, pertanto suddividiamo un'attività di trasferimento pesi in diverse fasi di pipeline che si sovrappongono nel tempo:

Memcpy da host a dispositivo — Se FSDP scarica il peso sulla CPU

Preparazione dei parametri — Ricostruire l'intero peso con full_tensor(), applicare la fusione di proiezione, quantizzare se necessario.

Trasferimento RDMA — Scrittura zero-copy nella memoria remota della GPU di inferenza

Barriera globale — Una volta completate tutte le chiamate a full_tensor(), sincronizzare tra i gruppi di mesh utilizzando GLOO via Ethernet.

Nell'implementazione, manteniamo una coda FIFO di attività per ogni fase della pipeline. Ogni volta che l'attività in testa alla coda completa la fase, viene spostata in coda alla fase successiva.

full_tensor() e altre operazioni GPU introducono un utilizzo extra della memoria GPU. Per evitare errori di memoria esaurita (out of memory), avviamo l'esecuzione di un'attività solo se le attività attualmente in corso occupano meno memoria GPU temporanea rispetto a una soglia configurabile.

Diverse scelte di progettazione rendono il nostro sistema significativamente più veloce da eseguire e più facile da mantenere rispetto alle comuni soluzioni open source.

Il trasferimento dei pesi basato su Rank0 raccoglie i pesi prima di distribuirli al nodo di inferenza.

La comunicazione point-to-point trasferisce i pesi dalle due nodi di training alle interfacce di rete del nodo di inferenza.

Un pattern comune consiste nel convogliare tutti i parametri attraverso le GPU di rango 0: raccolta sul rango 0 di addestramento, invio al rango 0 di inferenza, quindi nuova distribuzione. Questo diventa rapidamente un collo di bottiglia, limitato dalla larghezza di banda PCIe e dalla NIC di una singola GPU (ad es. 400 Gbps ≈ 50 GB/s).

Al contrario, la nostra configurazione punto-a-punto consente a ogni GPU di addestramento di inviare direttamente a ogni GPU di inferenza, saturando l'intera rete anziché un singolo collegamento.

Il trasferimento basato su RPC usa un buffer sul lato di inferenza e chiamate worker, mentre il trasferimento one-sided scrive i pesi direttamente senza la partecipazione del worker di inferenza.

Alcuni sistemi si basano sulla chiamata del metodo update_weight() del motore di inferenza per ogni tensore. Ciò comporta modifiche invasive al codice di inferenza, oltre al sovraccarico dovuto a RPC, serializzazione e coordinamento del piano di controllo.

Con la primitiva RDMA WRITE, aggiorniamo i pesi silenziosamente nella memoria della GPU di inferenza, senza copie extra. Non sono coinvolti messaggi del piano di controllo né logica di controllo CPU. Non è richiesta alcuna modifica al motore di inferenza.

Il processo di trasferimento dei pesi può sfruttare quattro tipi di risorse hardware: (1) Spostamento dati host-dispositivo (2) Calcolo GPU per fusione di proiezione e quantizzazione (3) Rete RDMA per il piano dati (4) Ethernet per il piano di controllo.

Il nostro design suddivide le attività di trasferimento pesi in fasi di pipeline, consentendo una facile sovrapposizione tra diverse risorse hardware.

Alcune implementazioni ricalcolano una pianificazione di trasferimento a ogni passaggio di addestramento, raccogliendo ripetutamente metadati e distribuendo istruzioni. Ciò aggiunge una latenza inutile al piano di controllo.

La nostra pianificazione viene calcolata una sola volta all'inizializzazione. Ogni iterazione di addestramento riproduce semplicemente il piano: il controller emette un segnale di "via" e le GPU seguono i percorsi preassegnati. L'esecuzione è prevedibile e leggera.

È allettante intrecciare l'intero processo di aggiornamento dei pesi in un'unica funzione monolitica: raccolta metadati, corrispondenza nomi, raccolta intra-nodo, fusione di proiezione, quantizzazione, suddivisione del mondo di comunicazione, trasferimento di rete inter-nodo. È difficile da programmare correttamente e ancora più difficile da ottimizzare.

Nella nostra ingegneria, separiamo questi passaggi come componenti individuali. Ciascun componente può essere testato unitariamente, analizzato e ottimizzato in isolamento.

Un trasferimento dei pesi veloce e affidabile è un elemento costitutivo fondamentale per il fine-tuning RL su larga scala. Combinando la primitiva RDMA WRITE, una pianificazione di trasferimento statica e l' esecuzione in pipeline, abbiamo ridotto gli aggiornamenti di trilioni di parametri a soli 1,3 secondi su Kimi-K2. L'approccio è semplice da analizzare, facile da mantenere ed evita i colli di bottiglia dei design tradizionali.

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: trasferimenti unidirezionali

Flusso di lavoro di alto livello

Esecuzione del trasferimento dei pesi

DeviceMesh e gruppi di mesh

Pipeline delle attività

Controllo dell'utilizzo della memoria GPU

Perché è veloce e semplice

Comunicazione punto-a-punto

Trasferimento dati unidirezionale

Pipelining

Pianificazione statica

Separazione netta

Conclusione

Citazione