Transferencia de pesos para post-entrenamiento RL en

Recientemente logramos actualizaciones de parámetros entre máquinas en 1.3 segundos para Kimi-K2 (1T de parámetros), transfiriendo pesos desde 256 GPU de entrenamiento (BF16) a 128 GPU de inferencia (FP8).

En el ajuste fino de aprendizaje por refuerzo asíncrono, el entrenamiento y la inferencia se ejecutan en GPU separadas. Después de cada paso de entrenamiento, se deben enviar los nuevos pesos a los nodos de inferencia. Muchos marcos existentes tardan varios segundos, o incluso minutos, en modelos de un billón de parámetros.

Al aprovechar la comunicación punto a punto RDMA, logramos que la transferencia de pesos sea extremadamente rápida, sin cambiar el motor de inferencia y facilitando la escritura y el mantenimiento del código.

Nuestra solución se basa en RDMA WRITE, un primitivo unidireccional donde la fuente escribe directamente en la memoria de la 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.
    ...

El lado del destino ni siquiera recibe notificación de la transferencia. Esto nos proporciona transferencias de baja latencia, alto rendimiento y cero copias impulsadas por los nodos de entrenamiento sin ninguna lógica de control en los nodos de inferencia.

Recopilación de metadatos – El controlador reúne los metadatos de los parámetros de todas las GPU de entrenamiento e inferencia.

Cómputo de programación – El controlador calcula un programa de transferencia de pesos estático, asignando qué GPU de entrenamiento envía qué parámetro a qué GPU de inferencia, y en qué orden.

Distribución de programación – El controlador envía el programa a todas las GPU de entrenamiento.

Ejecución – Después de cada paso de entrenamiento, el controlador indica a las GPU de entrenamiento que inicien las transferencias.

Con el flujo de trabajo de alto nivel definido, el desafío clave es cómo ejecutar transferencias de pesos de manera eficiente a una escala de un billón de parámetros. Aquí describimos los detalles de la ruta de ejecución.

Los parámetros en el entrenamiento se distribuyen de acuerdo con las colocaciones de FSDP. Usando full_tensor(), todas las GPU en un DeviceMesh pueden reconstruir el parámetro completo, por lo tanto, todas pueden servir como fuente para la transferencia de pesos.

Múltiples DeviceMeshes disjuntos forman un grupo de mallas. Debido a que los DeviceMeshes en el mismo grupo son disjuntos, sus transferencias no interfieren y pueden ejecutarse completamente en paralelo. Entre los grupos de mallas, insertamos una barrera global para hacer cumplir el orden.

La canalización de transferencia de pesos solapa las copias de host a dispositivo, el ensamblaje de tensores, la fusión y la cuantización, y las transferencias RDMA a través de MeshGroups.

Tratamos la transferencia de cada tensor de parámetros como una tarea. El proceso de transferencia de pesos utiliza múltiples tipos de fuentes de hardware, por lo tanto, dividimos una tarea de transferencia de pesos en diferentes etapas de tubería que se superponen en el tiempo:

Memcpy de host a dispositivo — Si FSDP descarga el peso a la CPU

Preparación de parámetros — Reconstruir el peso completo con full_tensor(), aplicar fusión de proyección, cuantificar si es necesario.

Transferencia RDMA — Escritura de copia cero en la memoria remota de la GPU de inferencia

Barrera global — Después de que todas las llamadas a full_tensor() se completen, sincronizar entre grupos de mallas usando GLOO a través de Ethernet.

En la implementación, mantenemos una cola FIFO de tareas para cada etapa de la tubería. Cada vez que la tarea en la cabecera de la cola completa la etapa, se mueve a la cola de la siguiente etapa.

full_tensor() y otras operaciones de GPU introducen un uso adicional de memoria de GPU. Para evitar errores de falta de memoria, iniciamos la ejecución de una tarea solo si las tareas en curso actuales ocupan menos memoria de GPU temporal que una marca de agua configurable.

Varias decisiones de diseño hacen que nuestro sistema sea significativamente más rápido de ejecutar y más fácil de mantener que las soluciones comunes de código abierto.

La transferencia de pesos basada en Rank0 recopila los pesos antes de distribuirlos al nodo de inferencia.

La comunicación punto a punto transfiere los pesos desde ambos nodos de entrenamiento a las interfaces de red del nodo de inferencia.

Un patrón común es canalizar todos los parámetros a través de GPU de rango 0: recopilar en el rango 0 de entrenamiento, enviar al rango 0 de inferencia, y luego volver a dispersar. Esto se convierte rápidamente en un cuello de botella, limitado por el ancho de banda PCIe de una sola GPU y la NIC (por ejemplo, 400 Gbps ≈ 50 GB/s).

Por el contrario, nuestra configuración punto a punto permite que cada GPU de entrenamiento envíe directamente a cada GPU de inferencia, saturando todo el tejido de red en lugar de un solo enlace.

La transferencia basada en RPC usa un búfer del lado de la inferencia y llamadas de trabajadores, mientras que la transferencia de un solo lado escribe los pesos directamente sin participación del trabajador de inferencia.

Algunos sistemas dependen de llamar al método update_weight() del motor de inferencia para cada tensor. Eso significa cambios intrusivos en el código de inferencia, además de la sobrecarga de RPC, serialización y coordinación del plano de control.

Con el primitivo RDMA WRITE, actualizamos los pesos silenciosamente en la memoria de la GPU de inferencia, sin copias adicionales. No se involucra ningún mensaje de plano de control ni lógica de control de CPU. No se requiere modificación del motor de inferencia.

El proceso de transferencia de pesos puede aprovechar cuatro tipos de recursos de hardware: (1) Movimiento de datos de host a dispositivo (2) Cómputo de GPU para fusión de proyección y cuantificación (3) Red RDMA para el plano de datos (4) Ethernet para el plano de control.

Nuestro diseño divide las tareas de transferencia de pesos en etapas de tubería, lo que permite una superposición sencilla entre diferentes recursos de hardware.

Algunas implementaciones vuelven a calcular un programa de transferencia en cada paso de entrenamiento, recopilando metadatos y distribuyendo instrucciones repetidamente. Esto añade una latencia innecesaria al plano de control.

Nuestro programa se calcula una vez en la inicialización. Cada iteración de entrenamiento simplemente repite el plan: el controlador emite una señal de “go”, y las GPU siguen sus rutas preasignadas. La ejecución es predecible y ligera.

Es tentador enredar todo el proceso de actualización de pesos en una función monolítica: recopilar metadatos, coincidencia de nombres, recopilación dentro del nodo, fusión de proyección, cuantificación, subsegmentación del mundo de comunicación, transferencia de red entre nodos. Es difícil programar correctamente, y aún más difícil de optimizar.

En nuestra ingeniería, separamos estos pasos como componentes individuales. Cada componente puede ser probado unitariamente, analizado y optimizado de forma aislada.

La transferencia de pesos rápida y fiable es un bloque de construcción crítico para el ajuste fino de RL a gran escala. Al combinar el primitivo RDMA WRITE, un programa de transferencia estático y ejecución en tubería, redujimos las actualizaciones de un billón de parámetros a solo 1.3 segundos en Kimi-K2. El enfoque es simple de analizar, fácil de mantener y evita los cuellos de botella de los diseños tradicionales.

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 unidireccionales

Flujo de trabajo de alto nivel

Ejecución de transferencia de pesos

DeviceMesh y Grupos de Mallas

Tubería de tareas

Control de uso de memoria de GPU

Por qué es rápido y sencillo

Comunicación punto a punto

Transferencia de datos unidireccional

Tubería (Pipelining)

Programa estático

Separación limpia

Conclusión

Citación