• 5 min de lectura
PyTorch Monarch lleva entrenamiento tolerante a fallos a GPUs AMD
PyTorch ha portado Monarch a AMD ROCm, permitiendo que el entrenamiento distribuido se recupere de fallos de GPU y de nodo sin reiniciar todo el trabajo.
Imagen: Hacker News
Entrenar grandes modelos de lenguaje en cientos o miles de GPUs hace que las fallas de hardware sean inevitables. Un solo error de memoria en una GPU, una partición de red o un fallo de nodo puede borrar días o semanas de trabajo.
PyTorch dice que ha llevado PyTorch Monarch a las GPUs AMD Instinct con ROCm, extendiendo el modelo de controlador único del runtime más allá de CUDA. El proyecto está diseñado para permitir que los trabajadores sanos sigan entrenando mientras los nodos fallidos se recuperan y se reincorporan, en lugar de forzar el reinicio completo del trabajo desde un checkpoint.
Cómo Monarch maneja el entrenamiento distribuido
Monarch permite a los desarrolladores orquestar un clúster de GPUs completo desde un solo programa Python. Su runtime basado en actores, la abstracción de process-mesh y el modelo de ejecución asíncrono soportan flujos de trabajo que combinan entrenamiento, evaluación y aprendizaje por refuerzo en un único script.
Recomendado
El apocalipsis laboral por la IA podría estar más lejos
La arquitectura separa la estrategia de paralelismo dentro de cada réplica de entrenamiento de la tolerancia a fallos entre réplicas. Los actores mantienen estado privado, las fallas están aisladas y los árboles de supervisión gestionan los errores en el nivel más bajo posible. Los reinicios locales pueden completarse en segundos, mientras que la recuperación escalada tarda minutos.
La pila incluye una API en Python, un runtime Monarch para actores, meshes, árboles de supervisión y particionado de tensores, y un runtime en Rust seguro en memoria construido sobre Tokio. Se integra con RDMA, RCCL/NCCL, SLURM, Kubernetes y SkyPilot.
Portar Monarch de CUDA a ROCm
El port a ROCm requirió cambios en la ejecución en GPU y en las comunicaciones distribuidas:
- Comunicaciones colectivas: PyTorch usó hipify_torch para convertir el puente en C++ de CUDA a HIP, y luego lo enlazó contra RCCL, que replica la API de NCCL.
- Gestión de memoria GPU: El sistema de compilación ahora detecta la plataforma y enruta las llamadas al driver de CUDA a través de sus equivalentes en HIP.
- RDMA: Configurar GPU_PLATFORM=rocm preserva la ruta RDMA basada en libibverbs mientras cambia los enlaces de GPU de CUDA a HIP para transferencias GPU-direct.
ROCm no proporciona un equivalente estático a libcudart_static.a de NVIDIA, por lo que la compilación para ROCm enlaza amdhip64 de forma dinámica. Ambas plataformas siguen cargando las funciones del driver de GPU con dlopen, manteniendo el contrato del runtime consistente.
El proyecto también añadió una capa de compatibilidad en Rust en lugar de bifurcar los enlaces. Tipos de HIP como hipError_t y hipStream_t se reexportan bajo nombres de CUDA, permitiendo que el resto del código en Rust se mantenga agnóstico a la plataforma. PyTorch informa que las 1.171 pruebas pasaron, con soporte para ROCm 7.0+. Los cambios se integraron en upstream a través del PR #2393 y el PR #2891.
Monarch en ROCm ahora soporta el runtime de Actores, RDMA, Supervisión y particionado de tensores en SLURM, Kubernetes y SkyPilot, con soporte posterior para TorchTitan y TorchFT.
Recuperación ante fallos con TorchTitan y TorchFT
Para su demostración, PyTorch combinó Monarch, TorchFT y TorchTitan en una configuración de entrenamiento distribuido sin checkpoints:
- Monarch orquesta el clúster, lanzando ReplicaActors y un servicio Lighthouse.
- TorchFT gestiona la tolerancia a fallos a nivel de paso, la coordinación de quórum, Quorum AllReduce y el salto de nodos fallidos.
- TorchTitan ejecuta los pasos forward, backward y del optimizador mientras gestiona checkpoints y métricas.
En una prueba con cuatro grupos de réplicas, cada ReplicaActor lanzó una réplica con ocho procesos de GPU. Cuando un proceso de GPU en la Réplica 0 falló, Monarch aisló el error mientras las Réplicas 1, 2 y 3 continuaron entrenando. La Réplica 0 se reinició en el mismo lugar, y el Lighthouse seleccionó la Réplica 1 para transferir el estado del modelo, el optimizador, el scheduler y el trainer.
Tras una breve pausa en el límite de quórum, las cuatro réplicas se reincorporaron y la sincronización DiLoCo se reanudó. No fue necesaria intervención manual ni recarga de un checkpoint global.
Resultados en clústeres MI300 y MI355
En un clúster SLURM de 16 nodos con 128 GPUs MI300, PyTorch entrenó un modelo Llama 3 8B mientras inyectaba fallos de RCCL cada 180 segundos y sincronizaba el quórum cada 20 pasos. Los trabajadores activos variaron entre 8 y 16, pero el entrenamiento continuó sin un reinicio completo y la curva de pérdida coincidió estrechamente con una referencia sin fallos.
Una segunda prueba usó un clúster Kubernetes de 32 nodos con 256 GPUs MI355. La participación se mantuvo entre 30 y 32 nodos, mientras que la pérdida media global descendió de forma suave de 12 a aproximadamente 4.
Los siguientes pasos de PyTorch incluyen ampliar el soporte de NIC, mejorar el rendimiento del runtime, añadir más frameworks de preentrenamiento y de aprendizaje por refuerzo en ROCm, reducir la latencia de recarga al reincorporarse, solapar la recuperación con el cálculo y continuar el trabajo de código abierto con la comunidad de PyTorch.
AI Editor
Ava covers the rapidly evolving world of artificial intelligence, from foundational models and research labs to the real-world economics of intelligence. With a background in computational linguistics, she cuts through the hype to find out what actually works. She firmly believes that benchmarks are just marketing until reproduced in the wild.
vía Hacker News
