Entrenamiento de PyTorch rápido y tolerante a fallos en AI Runtime

A gran escala, la eficiencia de tu entrenamiento está determinada por una única métrica: "goodput", la proporción de tiempo que tus GPU dedican a la computación productiva en lugar de esperar o recuperarse de fallos. Dado que los fallos de las GPU son el caso esperado a gran escala, la capacidad de recuperarse de forma rápida y automática de un fallo es la única manera de mantener un alto goodput y gestionar tu gasto total en GPU.

Dos subsistemas definen el éxito o el fracaso de esa recuperación, aunque ambos suelen considerarse como un aspecto secundario: el pipeline de datos que alimenta a tus aceleradores y el mecanismo de checkpointing que toma instantáneas del estado para que un trabajo pueda reanudarse. Si cometes un error en cualquiera de ellos, cada fallo te costará mucho más tiempo de inactividad de la GPU del que debería. Incluso fuera de los escenarios de fallo, un pipeline de datos que no pueda mantener el ritmo de tus aceleradores dejará a tus GPU sin recursos de forma silenciosa y reducirá el goodput con la misma seguridad que lo haría una caída del sistema. Analizaremos los mecanismos y las ventajas y desventajas de ambos, y cómo cada uno influye en tu goodput y en el gasto total de GPU. Consulta la Guía de rendimiento y resiliencia de entrenamiento complementaria para obtener indicaciones de código y ejemplos.

Para la parte de infraestructura del mismo problema, sobre cómo una flota detecta y aísla las GPU defectuosas antes de que interrumpan un trabajo, consulta la publicación complementaria, Cómo mantenemos la fiabilidad de las GPU en Databricks AI.

Por qué los fallos son el caso esperado a gran escala

A medida que aumenta el número de GPU en un trabajo, la probabilidad de que sobreviva a toda su duración sin interrupciones disminuye rápidamente. Un modelo de cálculo rápido y sencillo de la publicación complementaria de Databricks asume que cada GPU tiene una tasa de fallos anualizada de aproximadamente el 1%. Bajo esa suposición, la publicación señala que "un trabajo de 256 GPU que se ejecuta durante 30 días tiene aproximadamente un 19% de probabilidad de sufrir un fallo. Con 1,024 GPU, esa cifra aumenta al 57%", y estos son solo problemas a nivel de infraestructura.

Para fundamentar esa estimación en la realidad, la supercomputadora delta de 608 GPU H100 registró fallos cada 1.9 horas, lo que significa que para un trabajo de 32 GPU, el tiempo promedio hasta el fallo sería de 36 horas. La conclusión principal es que es probable que tu trabajo de entrenamiento falle en algún momento, y tomar las decisiones correctas puede hacer que tu modelo sea resiliente y reducir el tiempo total perdido cuando esto ocurra.

Impacto 1: El formato de checkpoint determina con qué frecuencia puedes permitirte guardar

El checkpointing es donde se gana o se pierde la resiliencia, y el mecanismo que elijas tiene un efecto de primer orden en la frecuencia con la que puedes guardar. Esta es la palanca más importante para tu goodput: si creas un checkpoint una vez al día, un fallo requerirá volver a ejecutar un promedio de 12 horas de trabajo duplicado para volver al estado en el que se encontraba cuando ocurrió el fallo.

El cuello de botella monolítico de torch.save

El primer checkpoint que la mayoría de los equipos escriben es un simple torch.save en el rango 0. Dependiendo de cómo se entrene tu modelo, existen potencialmente dos problemas:

1. Para el entrenamiento distribuido, recopila todos los estados en el rango 0 y escribe un único archivo.
2. Un único proceso escribe todo el checkpoint para sincronizarse de forma síncrona. Esto puede bloquearse por aspectos como las transferencias de red al guardar en almacenes de objetos remotos como Unity Catalog (UC).

Este comportamiento de bloqueo deja a tus GPU inactivas, lo que reduce tu goodput. Pero hay una manera de reducir la cantidad de tiempo que tu GPU dedica al checkpointing: la API de checkpoint distribuido de Torch.

Checkpoint distribuido (DCP): cada rango escribe su propio shard

El checkpoint distribuido de PyTorch invierte el diseño. Cada rango escribe su propio shard distinto en paralelo, junto con un pequeño archivo .metadata que describe cómo se componen los shards en los tensores completos.

El tiempo de guardado disminuye aproximadamente como 1/N con el número de rangos y, debido a que el archivo .metadata registra la distribución global, el mismo checkpoint se puede volver a cargar en un número diferente de GPU. DCP vuelve a planificar qué bytes necesita cada nuevo rango, por lo que la recuperación en un clúster de capacidad reducida después de perder nodos funciona sin problemas.

DCP vale la pena incluso para trabajos simples de paralelismo de datos

Una suposición común es que DCP es solo para modelos fragmentados (sharded), y que un trabajo de paralelismo de datos (DDP), donde cada rango contiene una réplica idéntica de los pesos, no tiene nada que ganar. No es así, DCP fragmenta el estado del modelo y lo escribe en paralelo en cada worker, incluso para tareas de entrenamiento DDP.

También es la misma API que necesitarás el día que pases a FSDP o al paralelismo de tensores, por lo que adoptarla temprano significa que nunca tendrás que reescribir el código de resiliencia en el peor momento posible.

Los guardados asíncronos hacen que la frecuencia sea casi gratuita

Incluso con escrituras en paralelo, un guardado síncrono bloquea el entrenamiento hasta que los bytes estén seguros en el almacenamiento, lo que para un checkpoint grande en un volumen remoto representa decenas de segundos de tiempo de inactividad del acelerador. async_save divide la operación: una copia rápida a un búfer de almacenamiento temporal (staging buffer) y luego una carga en segundo plano que se superpone con la continuación del entrenamiento.

El bucle de entrenamiento (training loop) solo paga por la copia temporal, no por la carga. Un checkpoint que solía costar decenas de segundos de tiempo de inactividad ahora no cuesta casi nada, que es exactamente lo que hace que el checkpointing frecuente de la siguiente sección sea viable.

En AI Runtime, UCVolumeWriter y UCVolumeReader implementan DCP en volúmenes de UC, realizando el staging de la I/O a través de NVMe local y marcando un checkpoint como completado solo una vez que sus datos se han guardado por completo. Consulta la guía de rendimiento y resiliencia para obtener todos los detalles y ejemplos de código.

Lo anterior excluye el tiempo de almacenamiento en red para torch.save.

Impacto 2: La frecuencia de los checkpoints determina tu coste de recuperación

Aquí es donde las piezas se acumulan. Cuando un trabajo falla, pierde todo desde el último checkpoint válido y debe volver a computarlo. Por lo tanto, el trabajo perdido esperado por fallo es aproximadamente la mitad del intervalo de checkpoint, y los guardados asíncronos económicos te permiten reducir ese intervalo.

Reducir el intervalo en un factor de 10 reduce el tiempo esperado de recuperación en un factor de 10. Recuerda la cifra de Llama 3 de ~8.6 interrupciones por día: con esa tasa de fallos, crear checkpoints cada 2 horas significa que esperas perder 8.6 horas al día en reentrenamiento, un goodput del 64%. Si creas checkpoints cada 30 minutos, solo perderás 2.15 horas, un goodput del 91%.

La recuperación también debe ser automática. Al reiniciar, el trabajo debe encontrar el checkpoint más reciente que haya terminado de escribirse, omitiendo cualquiera que haya quedado a medias debido a la caída, y reanudarse desde allí sin intervención humana (no human in the loop). DCP hace que esto sea fiable: el archivo .metadata se escribe solo después de que se hayan guardado todos los shards, por lo que su presencia es un marcador confiable de "este guardado está completo" para seleccionar.

Impacto 3: La carga de datos (dataloading) determina si tus GPU llegan a estar inactivas

Un trabajo de entrenamiento avanza a la velocidad de su entrada más lenta. Cuando los aceleradores esperan el siguiente lote (batch), tu goodput disminuye ya que tus GPU simplemente están inactivas. La única manera de solucionar este problema es asegurarse de que tu pipeline de entrada superponga la preparación de datos para el siguiente paso con la computación del paso actual, como se muestra en la siguiente figura:

A menudo vemos que los clientes que pasan a superponer la carga de datos con la computación experimentan una reducción del 20 al 50% en el tiempo de ejecución real (wall-clock time).

El coste de leer directamente desde el almacenamiento remoto

En una plataforma gobernada, los datos de entrenamiento residen en un almacenamiento de objetos remoto. En AI Runtime, los volúmenes de Unity Catalog (UC) se presentan como montajes de red.

Leer archivos directamente desde ese montaje en cada acceso vincula el tiempo de tu paso a la latencia de la red y vuelve a descargar los mismos archivos en cada época (epoch). La solución es un dataloader que copia cada archivo a un almacenamiento local rápido en el primer acceso, atiende las lecturas posteriores desde esa caché local y obtiene los siguientes archivos en paralelo mientras la GPU realiza la computación.

Con AI Runtime, UCVolumeDataset y DataLoader hacen exactamente esto (consulta la guía para ver ejemplos de código). UCVolumeDataset transmite archivos desde un volumen de UC, almacenando cada uno en caché en el NVMe local en el primer acceso, y particiona los archivos entre rangos y workers para que cada acelerador obtenga una porción disjunta y sin superposiciones. Nuestro DataLoader es una subclase directa de DataLoader de PyTorch cuyos valores predeterminados están optimizados para esta ruta, por lo que los archivos se recuperan y se almacenan en caché de forma simultánea mientras la GPU realiza los cálculos, en lugar de hacerlo uno a uno en el hilo de entrenamiento.

Ejemplo: entrenar un modelo de imágenes a partir de archivos de UC

Considera una carga de trabajo sencilla de clasificación de imágenes: decodificar archivos JPEG desde un volumen de UC, aplicar aumento de datos y entrenar un modelo de visión. Veamos dos formas de hacer esto en la misma GPU, con el mismo modelo y tamaño de lote: el Dataset estándar de PyTorch leyendo desde un volumen de UC frente a UCVolumeDataset más los valores predeterminados de DataLoader de Databricks.

No tienes que adivinar a dónde se va el tiempo

Como parte del desarrollo de DataLoader, nos hemos asegurado de que registre sus métricas en MLflow, lo que facilita ver de un vistazo si tu canalización de datos está bloqueando el entrenamiento.

La métrica fetch_seconds mide explícitamente cuánto tarda el dataloader en producir un lote y, durante este tiempo, tu GPU permanece inactiva.

Impacto 4: Olvidar la canalización de datos corrompe tu modelo de forma silenciosa

Existe un último error de resiliencia que no produce ningún mensaje de error, ninguna caída ni ningún trabajo fallido, solo un modelo que es sutilmente peor de lo que debería ser. Ocurre cuando guardas un punto de control (checkpoint) del modelo, el optimizador y el paso, pero no de la posición de tu canalización de datos dentro del conjunto de datos.

Considera un trabajo que se interrumpe a mitad de una época. Restaura el modelo correctamente y reanuda el bucle de entrenamiento, pero el dataloader vuelve a empezar desde el principio del conjunto de datos.

El trabajo reanudado vuelve a entrenarse con ejemplos que ya vio en esta época y, potencialmente, omite los que aún no había alcanzado. A lo largo de los muchos reinicios que la escala convierte en algo habitual, esto sesga silenciosamente la distribución de tus datos. El modelo sigue entrenándose; solo que lo hace con una muestra incorrecta de tus datos, precisamente el tipo de fallo silencioso que resulta más costoso, porque el trabajo se completa y nadie ve ningún problema hasta que las métricas resultan decepcionantes.

La solución es tratar la posición de los datos como parte del punto de control. Dependiendo de tu canalización, eso significa realizar el seguimiento de un desplazamiento de muestra o fragmento (shard) y avanzar al reanudar, hacer que un conjunto de datos personalizado serialice su propia posición o crear puntos de control en los límites de la época. Todo esto se basa en un requisito previo: el determinismo. El barajado y el aumento de datos se basan en generadores de números aleatorios, por lo que esas semillas y estados de RNG también deben formar parte del punto de control; de lo contrario, el orden de los datos después de un reinicio no coincidirá con el orden anterior, y una posición guardada apuntará a las muestras incorrectas.

La semilla, el orden reproducible y la canalización de datos reanudable son tres expresiones de una misma idea. La guía cubre cada estrategia con código.

Resumen

El entrenamiento rápido y tolerante a fallos proviene de un puñado de decisiones que se potencian mutuamente:

1. Utiliza Distributed Checkpoint en lugar de torch.save, incluso para DDP, de modo que los guardados sean paralelos y de bajo coste en lugar de un cuello de botella en serie.
2. Guarda de forma asíncrona para que los puntos de control sean casi gratuitos, lo que te permite guardar con frecuencia.
3. Recupérate automáticamente al punto de control válido más reciente, de modo que un fallo cueste minutos de recálculo, no horas.
4. Superpón la carga de datos con el cálculo mediante el almacenamiento en caché y la precarga (prefetching) desde el almacenamiento remoto para que los aceleradores nunca se queden inactivos esperando la entrada. Esto supone un ahorro recurrente de horas de GPU en cada paso.
5. Crea puntos de control de la canalización de datos y del estado de RNG, para que un trabajo reanudado continúe con los datos correctos en lugar de corromper silenciosamente tu modelo.

El principio unificador: los puntos de control frecuentes, económicos y completos convierten un fallo de hardware de un evento que pone fin al trabajo en un error de redondeo, y una canalización de entrada superpuesta mantiene ocupados a los aceleradores en el intervalo. Los guardados económicos (asíncronos) hacen que la frecuencia sea asequible; los guardados completos (modelo, datos y RNG) hacen que la recuperación sea correcta. Con ambos elementos implementados, y una flota que detecta e aísla el hardware que falla, tu tiempo de entrenamiento efectivo se acerca al límite máximo que permite el hardware, independientemente de lo inestable que sea el clúster subyacente.

Referencias

• Meta, The Llama 3 Herd of Models (2024): §3.3.1, fiabilidad del entrenamiento y desglose de interrupciones en hasta 16 384 GPU H100.
• Characterizing the Resilience of Hopper H100 and Ampere A100 GPUs (2025): estudio de campo de 2,5 años del sistema Delta; MTBE y probabilidad de fallo por tamaño de trabajo.
• Databricks, How we keep GPUs reliable across Databricks AI: comprobación del estado de las GPU a nivel de flota y el modelo probabilístico de fallos.

¿Listo para probarlo? Consulta la Guía de rendimiento y resiliencia del entrenamiento en la documentación de Databricks AI Runtime para ver el código completo, y lee How we keep GPUs reliable across Databricks AI para conocer la perspectiva de la infraestructura.

(Esta entrada del blog ha sido traducida utilizando herramientas basadas en inteligencia artificial) Publicación original