PyTorch Monarch llega a GPUs AMD con tolerancia a fallos en ROCm

Entrenar modelos de lenguaje (LLMs) con miles de millones de parámetros requiere entrenamiento distribuido en cientos o miles de GPUs. A esa escala, los fallos de hardware no son excepciones sino eventos esperados. Un error de memoria en una sola GPU, una partición de red o el crash de un nodo pueden tirar abajo un job de entrenamiento que llevaba días o semanas. Aunque el trabajo previo del equipo mostró escalado casi lineal de FP8 (96,16% de eficiencia sobre 1.024 GPUs MI325 con DeepSeekV3-671B), el desafío clave sigue siendo la confiabilidad a escala.

Para atacar este problema, el equipo llevó PyTorch Monarch a las GPUs AMD Instinct con ROCm, extendiendo el modelo single-controller más allá del universo CUDA.

Este artículo revisa la arquitectura de PyTorch Monarch, el esfuerzo de ingeniería para portar su runtime GPU y el stack de comunicación distribuida a ROCm, y cómo el sistema se recupera dinámicamente de fallos de nodos sin detener el entrenamiento completo.

El desafío: confiabilidad a escala

Las estrategias tradicionales de tolerancia a fallos dependen fuertemente del checkpointing periódico: guardar el estado completo del modelo a almacenamiento persistente a intervalos regulares. Cuando ocurre un fallo, el job entero reinicia desde el último checkpoint. Simple en el papel, este enfoque tiene desventajas significativas en compute perdido y GPU idle.

Para entrenamiento verdaderamente a gran escala, escalar no basta: el entrenamiento también debe recuperarse de fallos. Se necesita un enfoque dinámico, uno que permita que los nodos sanos sigan entrenando mientras los fallados se recuperan y se reincorporan, minimizando el cómputo desperdiciado y maximizando la utilización de GPU.

¿Qué es PyTorch Monarch?

PyTorch Monarch introduce un nuevo paradigma de programación distribuida que permite orquestar clusters completos de GPU desde un solo programa Python. Con su runtime basado en actores, su abstracción de process mesh y su modelo de ejecución asíncrono, Monarch simplifica el entrenamiento distribuido a gran escala.

La arquitectura opera en varios niveles:

• API Python: interfaz single-program donde el desarrollador escribe código Python simple para obtener ejecución GPU distribuida
• Runtime Monarch: administra actores y meshes, árboles de supervisión y sharding tensorial
• Runtime Rust (Tokio): garantiza alto rendimiento y seguridad de memoria
• Infraestructura: integra con RDMA, RCCL/NCCL, SLURM, Kubernetes y SkyPilot

Al desacoplar la estrategia de paralelismo dentro de cada réplica del mecanismo de tolerancia a fallos entre réplicas, Monarch entrega un modelo más limpio. Los fallos quedan aislados (los actores tienen estado privado y los crashes no propagan), son jerárquicos (se manejan al nivel más bajo posible) y la recuperación es rápida (segundos para reinicio local, minutos solo si escalan).

¿Cómo se portó Monarch a ROCm?

Llevar Monarch a GPUs AMD requirió un esfuerzo de ingeniería significativo para portar el runtime GPU y el stack de comunicación distribuida a ROCm. El equipo implementó con éxito tres caminos principales:

• Comunicaciones colectivas: usaron hipify_torch para convertir el código puente C++ de CUDA a HIP y linkearon contra RCCL, que espeja la API de NCCL.
• Gestión de memoria GPU: extendieron el build system para autodetectar la plataforma y rutear las llamadas del CUDA driver API por sus equivalentes HIP.
• Integración RDMA: configurar GPU_PLATFORM=rocm deja intacto el path RDMA basado en libibverbs y solo intercambia los bindings del lado GPU de CUDA a HIP para transferencias GPU-direct.

Dos temas transversales moldearon el puerto:

• Sin link estático para el runtime HIP: NVIDIA distribuye libcudart_static.a, así que el path CUDA linkea cudart_static directo. ROCm no tiene equivalente estático para libamdhip64, entonces el build ROCm linkea amdhip64 dinámicamente. Ambas plataformas además hacen dlopen de las funciones del driver API GPU.
• Shim de compatibilidad en Rust en lugar de fork de bindings: cuando hipify_torch reescribe los headers C/C++, bindgen emite tipos con nombres HIP como hipError_t, hipDeviceptr_t y hipStream_t. En vez de agregar ramas #ifdef en cada punto de llamada Rust, agregaron un módulo rocm_compat en nccl-sys y rdmaxcel-sys que re-exporta símbolos HIP bajo nombres CUDA. Ejemplo: pub type cudaError_t = hipError_t.

El puerto culminó con 1.171 tests pasando y soporte completo para ROCm 7.0+. Las contribuciones están upstream en las PRs #2393 y #2891.

Caso de estudio: entrenamiento sin checkpoint global

Para demostrar el poder de Monarch en GPUs AMD, el equipo lo integró con TorchTitan y TorchFT. La arquitectura tiene tres capas:

• Monarch: orquestador de procesos y cluster. Spawnea ReplicaActors y un servicio Lighthouse, organizando GPUs en Process Meshes.
• TorchFT: maneja tolerancia a fallos a nivel de step. Contacta al Lighthouse para coordinar quorum, ejecuta Quorum AllReduce y saltea nodos fallados.
• TorchTitan: motor de entrenamiento con Forward (FSDP), Backward y Optimizer, y manejo de checkpoints y métricas.

Cuando se inyecta un fallo en los actores de entrenamiento, el Lighthouse lo detecta y TorchFT lo maneja. Las réplicas sanas continúan entrenando de forma independiente pese a los fallos de sus pares, sin requerir una interrupción global.

Flujo de recuperación dinámica

Con cuatro grupos de réplicas, la secuencia es:

• Entrenamiento normal: OrchestrationManager spawnea 4 ReplicaActors y un Lighthouse. Cada ReplicaActor spawnea una réplica con 8 procesos GPU corriendo trainers TorchTitan. Sincronización DiLoCo cada 20 steps.
• Detección de fallo: un proceso GPU en la Réplica 0 crashea. El supervisor Monarch captura report_training_error con traceback completo antes de que el proceso muera. Las Réplicas 1, 2 y 3 quedan marcadas como no afectadas.
• Reinicio local: ReplicaActor 0 inicia un reinicio en el lugar (_stop_and_restart()). Las otras 3 réplicas siguen sincronizando (quorum_id=2).
• Transferencia de checkpoint peer-to-peer: el Lighthouse elige a la Réplica 1 como donante. Se transfiere modelo, optimizer, scheduler y estado del trainer de la Réplica 1 a la Réplica 0 en recuperación.
• Entrenamiento reanudado: cuando la Réplica 0 está sincronizada, se establece un nuevo quorum (quorum_id=3) con las 4 réplicas y la sincronización DiLoCo se reanuda.

Todo el proceso de recuperación se completa sin intervención manual, sin reinicio completo por checkpoint y con mínima disrupción al throughput global.

¿Qué tan bien escala? Resultados en MI300 y MI355

El enfoque fue validado en SLURM y Kubernetes con clusters MI300 y MI355.

Cluster SLURM de 16 nodos MI300 (128 GPUs)

Entrenaron un modelo Llama 3 8B en un cluster SLURM de 16 nodos (128 GPUs MI300), inyectando fallos RCCL cada 180 segundos con sincronización de quorum cada 20 steps:

• El número de workers activos fluctuó dinámicamente entre 8 y 16
• El entrenamiento continuó sin reinicios completos
• La curva de loss mostró convergencia estable, cercana al baseline sin inyección de fallos

Cluster Kubernetes de 32 nodos MI355 (256 GPUs)

Escalaron el experimento a un cluster Kubernetes de 32 nodos (256 GPUs MI355). El número de participantes se mantuvo altamente estable (fluctuando levemente entre 30 y 32 durante eventos de recuperación), y el loss promedio global bajó suavemente de 12 a aproximadamente 4.

¿Qué representa para el ecosistema AMD?

Este trabajo es la primera validación a gran escala de tolerancia a fallos avanzada sobre AMD Instinct. Confirma que el stack ROCm alcanza paridad funcional con CUDA en workloads de entrenamiento resiliente. Para operadores de infraestructura de IA en LatAm que evalúan alternativas a NVIDIA H100/H200 (con precios spot en AWS todavía por sobre USD 4 por hora), la validación de MI300/MI355 con TorchFT abre una ventana concreta para probar clusters más baratos sin renunciar a las garantías de resiliencia que hacen viable entrenar modelos de miles de millones de parámetros.