Las organizaciones que están ampliando su infraestructura de inteligencia artificial para soportar modelos de billones de parámetros enfrentan un dilema complicado: reducir el tiempo de entrenamiento a un costo más bajo o acelerar el proceso a un costo superior. Frecuentemente, la técnica de «checkpointing» se utiliza para acelerar el tiempo de recuperación y minimizar las pérdidas de tiempo de entrenamiento. Sin embargo, esta estrategia se traduce en costos de almacenamiento significativamente más altos. Por otro lado, si el «checkpointing» se realiza con poca frecuencia, se pueden reducir los costos, pero esto aumenta el riesgo de perder un progreso valioso en el entrenamiento en caso de fallos, que son comunes en entornos de entrenamiento distribuidos con miles de aceleradores.

Durante el entrenamiento del modelo Meta Llama 3, se reportó un fallo cada tres horas, siendo los problemas relacionados con GPUs responsables del 60% de las fallas. Los otros problemas fueron atribuidos a redes, CPUs y discos. Esta inestabilidad puede resultar en la pérdida de días de progreso de entrenamiento, lo que incrementa los costos y el tiempo de salida al mercado. Aunque checkpoints frecuentes pueden saturar las redes y sobrecargar el almacenamiento, lo que impacta en el rendimiento, la necesidad de encontrar un balance se vuelve crucial.

Para abordar estos desafíos, AWS ha anunciado el «checkpointing» en capas gestionadas en Amazon SageMaker HyperPod, una infraestructura diseñada específicamente para escalar y acelerar el desarrollo de modelos de IA generativa. Esta solución utiliza la memoria del CPU para el almacenamiento de checkpoints de alto rendimiento, replicando automáticamente los datos en nodos de cálculo adyacentes para mejorar la fiabilidad. SageMaker HyperPod no solo identifica automáticamente los problemas de los nodos y reemplaza los afectados para reanudar el entrenamiento, sino que también ayuda a implementar la mejor estrategia de «checkpointing», maximizando el rendimiento del entrenamiento.

Esta nueva función ya ha sido probada en grandes clústeres de entrenamiento distribuidos, con cantidades que varían desde cientos hasta más de 15,000 GPUs, logrando guardar checkpoints en cuestión de segundos. Implementar esta característica no requiere una experiencia técnica profunda, y puede integrarse fácilmente en scripts de entrenamiento de PyTorch.

Adicionalmente, el «checkpointing» en capas gestionadas permite a las organizaciones establecer la frecuencia y las políticas de retención tanto para el almacenamiento en memoria como para el almacenamiento persistente, utilizando Amazon S3 como una opción de respaldo. Esta tecnología mejora notablemente el tiempo de recuperación y optimiza la gestión de checkpoints en comparación con los enfoques tradicionales que dependen del almacenamiento remoto persistente.

Los mejores resultados se pueden lograr configurando la escritura de checkpoints en la capa en memoria frecuentemente, mientras que las copias en Amazon S3 se pueden realizar de manera menos frecuente. Con estas capacidades, la combinación de «managed tiered checkpointing» y SageMaker HyperPod promete mantener un alto rendimiento en el entrenamiento, incluso en entornos a gran escala propensos a fallos.
vía: AWS machine learning blog