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

Entrenar modelos de lenguaje (LLMs) con miles de millones de parametros 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 particion de red o el crash de un nodo pueden tirar abajo un job de entrenamiento que llevaba dias o semanas. Aunque el trabajo previo del equipo mostro escalado casi lineal de FP8 (96,16% de eficiencia sobre 1.024 GPUs MI325 con DeepSeekV3-671B), el desafio clave sigue siendo la confiabilidad a escala.

Para atacar este problema, el equipo llevo PyTorch Monarch a las GPUs AMD Instinct con ROCm, extendiendo el modelo single-controller mas alla del universo CUDA.

Este articulo revisa la arquitectura de PyTorch Monarch, el esfuerzo de ingenieria para portar su runtime GPU y el stack de comunicacion distribuida a ROCm, y como el sistema se recupera dinamicamente de fallos de nodos sin detener el entrenamiento completo.

El desafio: confiabilidad a escala

Las estrategias tradicionales de tolerancia a fallos dependen fuertemente del checkpointing periodico: guardar el estado completo del modelo a almacenamiento persistente a intervalos regulares. Cuando ocurre un fallo, el job entero reinicia desde el ultimo 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 tambien debe recuperarse de fallos. Se necesita un enfoque dinamico, uno que permita que los nodos sanos sigan entrenando mientras los fallados se recuperan y se reincorporan, minimizando el computo desperdiciado y maximizando la utilizacion de GPU.

¿Que es PyTorch Monarch?

PyTorch Monarch introduce un nuevo paradigma de programacion distribuida que permite orquestar clusters completos de GPU desde un solo programa Python. Con su runtime basado en actores, su abstraccion de process mesh y su modelo de ejecucion asincrono, Monarch simplifica el entrenamiento distribuido a gran escala.

La arquitectura opera en varios niveles:

• API Python: interfaz single-program donde el desarrollador escribe codigo Python simple para obtener ejecucion GPU distribuida
• Runtime Monarch: administra actores y meshes, arboles de supervision 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 replica del mecanismo de tolerancia a fallos entre replicas, Monarch entrega un modelo mas limpio. Los fallos quedan aislados (los actores tienen estado privado y los crashes no propagan), son jerarquicos (se manejan al nivel mas bajo posible) y la recuperacion es rapida (segundos para reinicio local, minutos solo si escalan).

¿Como se porto Monarch a ROCm?

Llevar Monarch a GPUs AMD requirio un esfuerzo de ingenieria significativo para portar el runtime GPU y el stack de comunicacion distribuida a ROCm. El equipo implemento con exito tres caminos principales:

• Comunicaciones colectivas: usaron hipify_torch para convertir el codigo puente C++ de CUDA a HIP y linkearon contra RCCL, que espeja la API de NCCL.
• Gestion de memoria GPU: extendieron el build system para autodetectar la plataforma y rutear las llamadas del CUDA driver API por sus equivalentes HIP.
• Integracion 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 estatico para el runtime HIP: NVIDIA distribuye libcudart_static.a, asi que el path CUDA linkea cudart_static directo. ROCm no tiene equivalente estatico para libamdhip64, entonces el build ROCm linkea amdhip64 dinamicamente. Ambas plataformas ademas 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 modulo rocm_compat en nccl-sys y rdmaxcel-sys que re-exporta simbolos HIP bajo nombres CUDA. Ejemplo: pub type cudaError_t = hipError_t.

El puerto culmino con 1.171 tests pasando y soporte completo para ROCm 7.0+. Las contribuciones estan 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 integro 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 metricas.

Cuando se inyecta un fallo en los actores de entrenamiento, el Lighthouse lo detecta y TorchFT lo maneja. Las replicas sanas continuan entrenando de forma independiente pese a los fallos de sus pares, sin requerir una interrupcion global.

Flujo de recuperacion dinamica

Con cuatro grupos de replicas, la secuencia es:

• Entrenamiento normal: OrchestrationManager spawnea 4 ReplicaActors y un Lighthouse. Cada ReplicaActor spawnea una replica con 8 procesos GPU corriendo trainers TorchTitan. Sincronizacion DiLoCo cada 20 steps.
• Deteccion de fallo: un proceso GPU en la Replica 0 crashea. El supervisor Monarch captura report_training_error con traceback completo antes de que el proceso muera. Las Replicas 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 replicas siguen sincronizando (quorum_id=2).
• Transferencia de checkpoint peer-to-peer: el Lighthouse elige a la Replica 1 como donante. Se transfiere modelo, optimizer, scheduler y estado del trainer de la Replica 1 a la Replica 0 en recuperacion.
• Entrenamiento reanudado: cuando la Replica 0 esta sincronizada, se establece un nuevo quorum (quorum_id=3) con las 4 replicas y la sincronizacion DiLoCo se reanuda.

Todo el proceso de recuperacion se completa sin intervencion manual, sin reinicio completo por checkpoint y con minima disrupcion al throughput global.

¿Que 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 sincronizacion de quorum cada 20 steps:

• El numero de workers activos fluctuo dinamicamente entre 8 y 16
• El entrenamiento continuo sin reinicios completos
• La curva de loss mostro convergencia estable, cercana al baseline sin inyeccion de fallos

Cluster Kubernetes de 32 nodos MI355 (256 GPUs)

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

¿Que representa para el ecosistema AMD?

Este trabajo es la primera validacion 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 evaluan alternativas a NVIDIA H100/H200 (con precios spot en AWS todavia por sobre USD 4 por hora), la validacion de MI300/MI355 con TorchFT abre una ventana concreta para probar clusters mas baratos sin renunciar a las garantias de resiliencia que hacen viable entrenar modelos de miles de millones de parametros.