Transferencia de pesos para post-adestramento de RL en

Recentemente conseguimos actualizacións de parámetros entre máquinas en 1,3 segundos para Kimi-K2 (1T de parámetros), transferindo os pesos desde 256 GPUs de adestramento (BF16) ata 128 GPUs de inferencia (FP8).

No axuste fino de aprendizaxe por reforzo asincrónica, o adestramento e a inferencia executaranse en GPUs separadas. Despois de cada paso de adestramento, deben enviarse novos pesos aos nós de inferencia. Moitos marcos existentes tardan varios segundos —ou mesmo minutos— en modelos de billóns de parámetros.

Ao aproveitar a comunicación punto a punto por RDMA, conseguimos que a transferencia de pesos sexa extremadamente rápida, sen cambiar o motor de inferencia, e facemos que o código sexa máis doado de escribir e manter.

A nosa solución está construída sobre RDMA WRITE, unha primitiva unilateral onde a orixe escribe directamente na memoria da GPU de destino.

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

O lado de destino nin sequera recibirá unha notificación da transferencia. Isto ofrécenos transferencias de baixa latencia, alto rendemento e sen copia, impulsadas polos nós de adestramento sen ningunha lóxica de control nos nós de inferencia.

Recollida de metadatos – O controlador recolle os metadatos dos parámetros de todas as GPUs de adestramento e inferencia.

Cálculo da planificación – O controlador calcula un calendario estático de transferencia de pesos, mapeando qué GPU de adestramento envía qué parámetro a qué GPU de inferencia e en qué orde.

Distribución da planificación – O controlador envía o calendario a todas as GPUs de adestramento.

Execución – Despois de cada paso de adestramento, o controlador sinaliza ás GPUs de adestramento que inicien as transferencias.

Unha vez definido o fluxo de traballo de alto nivel, o reto principal é como executar as transferencias de pesos de forma eficiente a escala de billóns de parámetros. Aquí describimos os detalles da ruta de execución.

Os parámetros no adestramento distribúense segundo as colocacións de FSDP. Usando full_tensor(), todas as GPUs dun DeviceMesh poden reconstruír o parámetro completo, polo que todas poden servir como orixe para a transferencia de pesos.

Varios DeviceMeshes disxuntos forman un grupo de malla. Debido a que os DeviceMeshes do mesmo grupo son disxuntos, as súas transferencias non interferirán e poden executarse totalmente en paralelo. Entre os grupos de malla, inserimos unha barreira global para impoñer a orde.

O canal de transferencia de pesos solápase coas copias de host a dispositivo, a ensamblaxe de tensores, a fusión e a cuantización, e as transferencias RDMA a través de MeshGroups.

Tratamos a transferencia de cada tensor de parámetros como unha tarefa. O proceso de transferencia de pesos utiliza varios tipos de recursos de hardware, polo que dividimos unha tarefa de transferencia de pesos en diferentes etapas de tubaxe que se superpoñen no tempo:

Copia de memoria (memcpy) de host a dispositivo — Se FSDP descarga o peso á CPU

Preparación de parámetros — Reconstruír o peso completo con full_tensor(), aplicar fusión de proxeccións, cuantificar se é necesario.

Transferencia RDMA — Escritura sen copia na memoria remota da GPU de inferencia

Barreira global — Despois de que todas as chamadas a full_tensor() rematen, sincronizar a través dos grupos de malla usando GLOO a través de Ethernet.

Na implementación, mantemos unha cola FIFO de tarefas para cada etapa da tubaxe. Sempre que a tarefa que está na cabeza da cola completa a etapa, móvese á cola da seguinte etapa.

full_tensor() e outras operacións de GPU introducen un uso adicional de memoria na GPU. Para evitar erros de falta de memoria (out of memory), iniciamos a execución dunha tarefa só se as tarefas actuais en curso ocupan menos memoria temporal de GPU que un limiar configurable.

Varias decisións de deseño fan que o noso sistema sexa significativamente máis rápido de executar e máis doado de manter que as solucións comúns de código aberto.

A transferencia de pesos baseada en Rank0 recolle os pesos antes de distribuílos ao nodo de inferencia.

A comunicación punto a punto transfire os pesos desde ambos os nodos de adestramento ás interfaces de rede do nodo de inferencia.

Un patrón común é canalizar todos os parámetros a través das GPUs de rango 0: recoller na GPU de rango 0 de adestramento, enviar á GPU de rango 0 de inferencia e despois distribuír de novo. Isto convértese rapidamente nun cuello de botella, limitado polo ancho de banda PCIe dunha única GPU e a tarxeta de rede (por exemplo, 400 Gbps ≈ 50 GB/s).

Pola contra, a nosa configuración punto a punto permite que cada GPU de adestramento envíe directamente a cada GPU de inferencia, saturando toda a estrutura de rede en lugar dunha única ligazón.

A transferencia baseada en RPC usa un búfer do lado da inferencia e chamadas de traballadores, mentres que a transferencia unidireccional escribe os pesos directamente sen participación do traballador de inferencia.

Algúns sistemas dependen de chamar ao método update_weight() do motor de inferencia para cada tensor. Iso significa cambios intrusivos no código de inferencia, ademais da sobrecarga de RPC, serialización e coordinación do plano de control.

coa primitiva RDMA WRITE, actualizamos os pesos en silencio na memoria da GPU de inferencia, sen copias adicionais. Non hai mensaxes no plano de control nin lóxica de control da CPU implicada. Non se require ningunha modificación no motor de inferencia.

O proceso de transferencia de pesos pode aproveitar catro tipos de recursos de hardware: (1) Movemento de datos entre host e dispositivo (2) Cómputo de GPU para a fusión de proxeccións e cuantificación (3) Rede RDMA para o plano de datos (4) Ethernet para o plano de control.

O noso deseño divide as tarefas de transferencia de pesos en etapas de tubaxe, permitindo unha fácil superposición entre diferentes recursos de hardware.

Algunhas implementacións recalculan un calendario de transferencia en cada paso de adestramento, recollendo repetidamente metadatos e distribuíndo instrucións. Isto engade unha latencia innecesaria ao plano de control.

O noso calendario calcúlase unha vez na inicialización. Cada iteración de adestramento simplemente reproduce o plan: o controlador emite un sinal de "inicio" e as GPUs seguen as súas rutas preasignadas. A execución é previsible e lixeira.

É tentador entrelazar todo o proceso de anovación de pesos nunha única función monolítica: recoller metadatos, coincidencia de nomes, recollida intra-nó, fusión de proxeccións, cuantificación, subdivisión do mundo de comunicación, transferencia de rede inter-nó. É difícil de programar correctamente e aínda máis difícil de optimizar.

Na nosa enxeñaría, separamos estes pasos como compoñentes individuais. Cada compoñente pode probarse unitariamente, analizarse e optimizarse de xeito illado.

A transferencia de pesos rápida e fiable é un bloque de construción crítico para o axuste fino de RL a grande escala. Combinando a primitiva RDMA WRITE, un calendario de transferencia estático e a execución en tubaxe, reducimos as actualizacións de billóns de parámetros a só 1,3 segundos en Kimi-K2. O enfoque é sinxelo de entender, doado de manter e evita os cuellos de botella dos deseños tradicionais.

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: transferencias unilaterais

Fluxo de traballo de alto nivel

Execución da transferencia de pesos

DeviceMesh e grupos de malla

Tubaxe de tarefas

Control de uso de memoria da GPU

Por que é rápido e sinxelo

Comunicación punto a punto

Transferencia de datos unilateral

Tubaxe (Pipelining)

Planificación estática

Separación limpa

Conclusión

Citación