Entrenamiento asíncrono con LoRA y vLLM en Hugging Face Jobs: un flujo sin NCCL

Resumen

TRL (v1.14) incorporó soporte para entrenar adaptadores LoRA desde AsyncGRPOTrainer (PR #7017). Esto permite sincronizar únicamente el adaptador —no el modelo completo— con vLLM, lo cual cambia las exigencias de infraestructura: un adaptador de rango 1 pesa solo unos megabytes frente a varios gigabytes del modelo completo. Aprovechando Hugging Face Jobs y Storage Buckets montados como sistema de archivos, es posible ejecutar el entrenador y las réplicas de inferencia en máquinas separadas sin usar NCCL ni compartir disco local.

Por qué LoRA facilita separar entrenamiento e inferencia

LoRA (Low-Rank Adaptation) es especialmente apropiado en escenarios de RL: investigaciones y experiencias prácticas muestran que un adaptador de rango bajo (incluso 1) puede capturar la información necesaria para políticas basadas en gradiente de política. En la práctica, un adaptador de rango 1 para un modelo de ~1.5B queda en el orden de megabytes, mientras que el modelo completo puede ocupar ~3 GB. En vez de redistribuir gigabytes tras cada actualización, basta enviar unos megabytes.

Esto no solo reduce el tráfico entre nodos, sino que permite a vLLM mantener múltiples adaptadores cargados: rollouts antiguos continúan con la política con la que iniciaron, y los nuevos usan la versión más reciente.

Arquitectura propuesta sobre Hugging Face Jobs

La solución real probada consta de componentes ligeros y replicables:

• Un Job que corre AsyncGRPOTrainer con LoRA (y opcionalmente FSDP en el entrenamiento).
• Dos Jobs que sirven vLLM, cada uno en su propia GPU, con la imagen estándar vllm/vllm-openai.
• Un Storage Bucket montado en todos los Jobs en la misma ruta, actuando como sistema de archivos compartido.
• Un proxy que recibe peticiones del trainer y enruta rollouts a la réplica más apropiada; además, se encarga de difundir las cargas de adaptadores a las réplicas.

Cada Hugging Face Job es un contenedor en una VM independiente (no puede crear múltiples nodos desde un único Job). Por eso no podemos depender de NCCL ni de un filesystem local compartido entre procesos: la sincronización de modelos completos dejaría la solución inviable por el volumen de datos.

Cómo funciona el flujo de sincronización de adaptadores

El entrenador guarda periódicamente el adaptador en disco bajo <output_dir>/.vllm_lora/trl-policy-v{N}, y publica el directorio con un rename atómico. A continuación el trainer llama al endpoint /v1/load_lora_adapter de vLLM enviando la ruta del adaptador (no tensores). vLLM lee directamente desde el disco y puede servir requests con model=“trl-policy-v{N}”.

En entornos tradicionales esto exige un filesystem compartido (NFS, por ejemplo). En Hugging Face Jobs se consigue montando un Storage Bucket como volumen POSIX dentro de cada Job mediante hf-mount. Un ejemplo de montaje (simplificado) en la ejecución de Job es:

hf jobs run … -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora …

De ese modo la ruta enviada por el entrenador es válida dentro de cada contenedor y los servidores pueden cargar el adaptador desde la misma ubicación.

El rol del proxy: KV cache y broadcasting

Separar trainer y réplicas implica dos retos operativos:

1. Mantener eficiencia en la generación: vLLM usa cachés KV para acelerar generación. Si un rollout se dirige a una réplica que no tiene el prefijo de KV correspondiente, se pierde eficiencia.
2. Asegurar que todas las réplicas carguen los nuevos adaptadores cuando el trainer publica una versión.

La solución incluye un pequeño proxy que actúa entre el trainer y las réplicas de vLLM. Sus responsabilidades son:

• Añadir el header de autenticación a las peticiones hacia las réplicas (porque el trainer habla con el proxy por localhost y éste con réplicas vía HTTPS).
• Enrutar cada rollout a la réplica que probablemente ya tenga el prefijo de KV correspondiente (mejorando latencia y uso de GPU).
• Difundir (broadcast) las cargas de adaptador para que cada réplica obtenga la versión recién publicada.

Este proxy mantiene la lógica de afinitidad necesaria para que la separación física no degrade fuertemente el rendimiento de generación.

Configuración de réplicas y staleness

En AsyncGRPOTrainer cada sincronización aumenta la versión de la política. max_staleness define cuántas versiones puede atrasarse un rollout antes de descartarse. Por ejemplo, con max_staleness=4, una muestra generada con trl-policy-v3 sigue siendo válida mientras el trainer esté en v7; por tanto las réplicas deben ser capaces de servir varias versiones simultáneamente. vLLM permite reservar ranuras de adaptador para ese propósito.

Esto explica por qué cada réplica debe tener espacio para múltiples adaptadores en memoria y por qué el proxy debe saber qué versión tiene cada réplica.

Persistencia y tolerancia a preemptions

Los Jobs en HF son efímeros; sin embargo, al almacenar checkpoints y el adaptador final en el bucket, un entrenador preemptado puede retomar sin perder el trabajo. El bucket actúa como almacenamiento persistente entre ejecuciones.

Resultados prácticos y eficiencia

En pruebas reales usando este patrón, con el mismo recipe se observaron tiempos de entrenamiento para 500 pasos que variaron entre 3 h 27 min y 53 min en cinco ejecuciones. Esto ejemplifica cómo la arquitectura y la separación de roles —junto con la eficiencia de LoRA— pueden acelerar flujos de trabajo en RL cuando el bottleneck se resuelve adecuadamente.

Consideraciones para equipos en América Latina

Para equipos en la región, esta aproximación tiene ventajas relevantes:

• Reduce la necesidad de clústeres con comunicación de alta velocidad (NCCL), lo cual puede simplificar despliegues en clouds o proveedores locales.
• Minimiza el tráfico de datos entre máquinas al sincronizar solo unos megabytes por actualización.
• Permite reutilizar infra con Jobs independientes (entrenamiento e inferencia) y persistencia en buckets, facilitando la recuperación ante interrupciones.

Al planear adopción, consideren la latencia entre Jobs, el costo de almacenamiento y las políticas de I/O del proveedor. El patrón también facilita replicar la solución en entornos híbridos o multi-cloud, siempre que el bucket sea accesible desde los contenedores.

Conclusión

El soporte a LoRA en AsyncGRPOTrainer (TRL v1.14) abre la puerta a un modo de trabajo donde entrenamiento e inferencia conviven en máquinas separadas sin necesidad de NCCL ni filesystem compartido local, gracias a adaptadores ligeros y Storage Buckets montados con hf-mount. Añadiendo un proxy para afinar el enrutamiento y la difusión de adaptadores, se consigue una arquitectura sencilla, tolerante a preemptions y eficiente para entrenamiento RL con políticas adaptativas. Para equipos que buscan escalar sin invertir en clústeres complejos, esta estrategia ofrece una vía pragmática y reproducible.