Transferencia de ponderaciones para post-entrenamiento de RL en menos de 2 segundos

Recientemente logramos actualizaciones de parámetros entre máquinas de 1.3 segundos para Kimi-K2 (1T de parámetros), transfiriendo ponderaciones de 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 nuevas ponderaciones a los nodos de inferencia. Muchos marcos existentes tardan varios segundos, o incluso minutos, en modelos de un billón de parámetros.

Aprovechando la comunicación punto a punto RDMA, podemos hacer que la transferencia de ponderaciones sea sumamente 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, una primitiva unilateral donde el origen 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 de destino ni siquiera recibirá una notificación sobre la transferencia. Esto nos proporciona transferencias de baja latencia, alto rendimiento y sin copia controladas por los nodos de entrenamiento sin ninguna lógica de control en los nodos de inferencia.

Recopilación de metadatos: el controlador recopila los metadatos de los parámetros de todas las GPU de entrenamiento e inferencia.

Cómputo de programación: el controlador calcula una programación de transferencia de ponderaciones estática, mapeando 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 la programación 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.

Una vez definido el flujo de trabajo de alto nivel, el desafío clave es cómo ejecutar las transferencias de ponderaciones 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 ubicaciones de FSDP. Usando full_tensor(), todas las GPU en un DeviceMesh pueden reconstruir el parámetro completo, por lo que todas pueden servir como origen para la transferencia de ponderaciones.

Varios DeviceMesh disjuntos forman un grupo de mallas. Debido a que los DeviceMesh 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 imponer el orden.

El canal de transferencia de pesos superpone las copias de host a dispositivo, el ensamblaje de tensores, la fusión y cuantización, y las transferencias RDMA entre MeshGroups.

Tratamos la transferencia de cada tensor de parámetros como una tarea. El proceso de transferencia de ponderaciones utiliza múltiples tipos de recursos de hardware, por lo que dividimos una tarea de transferencia de ponderaciones en diferentes etapas de canalización que se superponen en el tiempo:

Copia de memoria (memcpy) del host al dispositivo — Si FSDP descarga la ponderación a la CPU

Preparación de parámetros — Reconstruir ponderación completa con full_tensor(), aplicar fusión de proyección, cuantizar si es necesario.

Transferencia RDMA — Escritura sin copia a la memoria de la GPU de inferencia remota

Barrera global — Una vez que se completan todas las llamadas a full_tensor(), sincronizar a través de los 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 canalización. Siempre que la tarea al frente de la cola completa la etapa, se mueve al final de 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 (out of memory), iniciamos la ejecución de una tarea solo si las tareas actuales en curso ocupan menos memoria temporal de GPU que un umbral 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 toda la infraestructura de red en lugar de un solo enlace.

La transferencia basada en RPC usa un búfer del lado de inferencia y llamadas del worker, mientras que la transferencia de un solo lado escribe los pesos directamente sin la participación del worker 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 las RPC, la serialización y la coordinación del plano de control.

Con la primitiva RDMA WRITE, actualizamos las ponderaciones de forma silenciosa en la memoria de la GPU de inferencia, sin copias adicionales. No interviene ningún mensaje del plano de control ni lógica de control de la CPU. No se requiere ninguna modificación en el motor de inferencia.

El proceso de transferencia de ponderaciones puede aprovechar cuatro tipos de recursos de hardware: (1) Movimiento de datos entre el host y el dispositivo (2) Cómputo de GPU para fusión de proyección y cuantizació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 ponderaciones en etapas de canalización, lo que permite una fácil superposición entre diferentes recursos de hardware.

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

Nuestra programación se calcula una vez en la inicialización. Cada iteración de entrenamiento simplemente repite el plan: el controlador emite una señal de "inicio" y las GPU siguen sus rutas preasignadas. La ejecución es predecible y ligera.

Es tentador entrelazar todo el proceso de actualización de ponderaciones en una función monolítica: recopilar metadatos, coincidencia de nombres, recopilación intranodo, fusión de proyección, cuantización, subsegmentación del mundo de comunicación, transferencia de red internodos. Es difícil de programar correctamente y aún más difícil de optimizar.

En nuestra ingeniería, separamos estos pasos como componentes individuales. Cada componente se puede probar de forma unitaria, analizar y optimizar de forma aislada.

La transferencia de ponderaciones rápida y confiable es un bloque de construcción crítico para el ajuste fino de RL a gran escala. Al combinar la primitiva RDMA WRITE, una programación de transferencia estática y una ejecución canalizada, redujimos las actualizaciones de un billón de parámetros a solo 1.3 segundos en Kimi-K2. El enfoque es fácil 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 unilaterales

Flujo de trabajo de alto nivel

Ejecución de transferencia de ponderaciones

DeviceMesh y grupos de mallas

Canalización de tareas

Control de uso de memoria de la GPU

Por qué es rápido y simple

Comunicación punto a punto

Transferencia de datos unilateral

Canalización

Programación estática

Separación limpia

Conclusión

Citación