Cómo acelerar el entrenamiento RL multimodal con SkyRL en SageMaker HyperPod

Resumen y por qué importa

El post-entrenamiento por refuerzo (RL) se está convirtiendo en un paso habitual para crear agentes basados en modelos de lenguaje que razonan y actúan en secuencias de pasos. Ejecutar este tipo de entrenamientos a escala, con cientos de GPU-horas y múltiples nodos, exige infraestructura persistente que soporte trabajos largos, recupere fallas de hardware sin perder progreso y permita visibilidad en tiempo real de la dinámica de entrenamiento.

Amazon SageMaker HyperPod provee precisamente esa capa de infraestructura sobre Amazon EKS. Con capacidades de resiliencia de clúster, reemplazo automático de nodos defectuosos y soporte para checkpointing, HyperPod permite que trabajos multi-nodo de larga duración retomen desde su último punto guardado en vez de reiniciar desde cero. En esta nota explicamos cómo combinar HyperPod con Ray y SkyRL para entrenar un modelo vision-language Qwen3-VL-8B en una tarea de navegación de laberintos usando Group Relative Policy Optimization (GRPO).

Conceptos clave: RL multietapa y GRPO

El RL convencional suele asignar recompensas a salidas individuales. El RL multi-turn, en cambio, entrena al agente sobre episodios completos: el agente observa un estado, actúa, recibe retroalimentación y continúa hasta terminar el episodio. La señal de aprendizaje proviene de la recompensa acumulada durante todo el episodio, no de decisiones aisladas.

En problemas con recompensas escasas, como navegar un laberinto 2D, los episodios solo generan recompensa al alcanzar la meta; no existe una clave por paso que indique si una acción puntual fue correcta. SkyRL aborda esto con GRPO: para cada posición inicial, el agente ejecuta varias corridas con la política actual, compara esas ejecuciones entre sí y refuerza las que superan el promedio del grupo, penalizando las que están por debajo. Esa comparación intra-grupo es la señal de entrenamiento, lo que permite prescindir de un critic o modelo de valor adicional.

Topología de entrenamiento usada en el ejemplo

La solución descrita ejecuta SkyRL sobre un clúster Ray provisionado por SageMaker HyperPod con una arquitectura de un nodo head CPU y tres nodos worker GPU. SkyRL coloca tanto la generación de rollouts como el entrenamiento en las mismas GPUs: motores vLLM crean los rollouts completos mientras el modelo de política, shardeado con Fully Sharded Data Parallel (FSDP), aplica las actualizaciones de gradiente.

Después de cada paso del optimizador, los pesos adaptadores LoRA se sincronizan desde las réplicas de entrenamiento hacia los motores de inferencia mediante un sistema de archivos compartido Amazon FSx para Lustre montado en /shared. Esta configuración permite que la inferencia y el entrenamiento compartan rápidamente los últimos adaptadores sin necesidad de transferencias complicadas.

Instancias y componentes clave en el despliegue descrito:

• Workers: 3 x ml.g7e.12xlarge (cada uno con 2 GPUs NVIDIA RTX PRO 6000 Blackwell, total 6 GPUs)
• Head: ml.r5d.16xlarge (512 GB RAM, gestiona Ray GCS, dashboard y consolidación de adaptadores LoRA)
• Modelo de política: Qwen3-VL-8B con LoRA rango 32, shardeado entre las 6 GPUs mediante PyTorch FSDP
• Motores de rollout: 6 instancias vLLM colocadas, una por GPU
• Almacenamiento compartido: Amazon FSx para Lustre montado en /shared para sincronización de LoRA y output de evaluación

Resiliencia, checkpointing y monitoreo

HyperPod ofrece monitoreo continuo del estado de los nodos y reemplaza automáticamente nodos defectuosos. Combinado con checkpoints periódicos, esto evita perder horas de rollouts en caso de fallas de hardware. La integración con Ray facilita crear el clúster desde SageMaker Studio y enviar trabajos de forma remota a través del endpoint Ray de HyperPod.

Además, el add-on HyperPod Observability provisiona dashboards pre-construidos en Amazon Managed Grafana, lo que permite seguir métricas de entrenamiento y comportamiento del clúster sin montar paneles manualmente. Para equipos y tomadores de decisión en America Latina, estas capacidades son críticas al ejecutar entrenamientos costosos: reducen riesgo operativo y facilitan cumplimiento de ventanas de ejecución en nubes públicas.

Resultados prácticos del experimento

Partiendo de un checkpoint de fine-tuning supervisado (VisGym SFT), el post-entrenamiento con GRPO en HyperPod mejoró la tasa de resolución del laberinto en un conjunto de evaluación fijo de 64 laberintos, pasando de 43.75% a más de 95%. Ese salto demuestra cómo el aprendizaje por refuerzo post-SFT puede transformar la capacidad de un modelo multimodal para tareas que requieren secuencias de acciones.

Requisitos y preparación del entorno

Para reproducir este flujo se requieren varios componentes en el clúster y el entorno de SageMaker:

• Un clúster SageMaker HyperPod orquestado por Amazon EKS con al menos 3 instancias ml.g7e.12xlarge y una ml.r5d.16xlarge
• Operadores de Kubernetes: KubeRay operator, HyperPod Observability EKS add-on y HyperPod Ray Endpoint Operator para envío remoto de trabajos
• Amazon FSx para Lustre CSI driver instalado, un filesystem FSx para Lustre, un PersistentVolume respaldado por ese filesystem y un PersistentVolumeClaim ReadWriteMany que los pods monten en /shared
• Un dominio de SageMaker Studio con permisos para conectarse al HyperPod
• El paquete Python toolkit-for-ray-on-sagemaker-ai instalado

Estos componentes permiten que los motores vLLM y las réplicas de entrenamiento sincronicen adaptadores LoRA y que SageMaker Studio actúe como punto de control y envío de trabajos.

Flujo de trabajo recomendado

El flujo general para entrenar con SkyRL en HyperPod sigue estos pasos:

1. Preparar e implementar el clúster HyperPod con los operadores y drivers necesarios.
2. Provisionar y montar el sistema de archivos FSx para Lustre para checkpoints y sincronización LoRA.
3. Crear el Ray cluster desde SageMaker Studio y configurar el endpoint para envío remoto.
4. Lanzar el trabajo SkyRL que coloca los rollouts vLLM y la política shardeada en los nodos GPU.
5. Monitorear métricas de entrenamiento y salud del clúster en Grafana via HyperPod Observability.
6. Usar checkpoints continuos para garantizar que, ante fallas, el trabajo pueda reanudarse desde el último estado guardado.
7. Consolidar y hospedar el modelo entrenado para inferencia o despliegue.

Relevancia para equipos en America Latina

Para organizaciones y grupos de investigación en la región, esta combinación de herramientas ofrece un camino para ejecutar entrenamientos RL multimodales sin invertir en infraestructura física a gran escala. El uso de recursos gestionados, sincronización eficiente con FSx y observabilidad integrada reduce la barrera operacional. Además, la capacidad de retomar trabajos largos es especialmente útil donde las ventanas de ejecución y los costos deben gestionarse cuidadosamente.

Conclusión

SkyRL sobre SageMaker HyperPod demuestra cómo orquestar entrenamiento RL multimodal a gran escala con resiliencia, sincronización eficiente y monitoreo listo para usar. El enfoque permite mejorar significativamente el desempeño en tareas secuenciales complejas sin depender de un critic adicional, usando GRPO y sincronización de adaptadores LoRA en un entorno Ray administrado. Para equipos que buscan escalar proyectos de RL en la nube, esta arquitectura ofrece una ruta práctica y robusta para experimentación y producción.