Las organizaciones que buscan expandir su infraestructura de inteligencia artificial enfrentan un dilema cada vez más común: optimizar el tiempo de entrenamiento de modelos de gran envergadura, compuestos por billones de parámetros, a un costo razonable. Mientras que acelerar este proceso puede aumentar los costos, reducir el tiempo de entrenamiento sin generar gastos excesivos es el objetivo deseado. Una técnica habitual es el «checkpointing», que permite acelerar la recuperación y minimizar pérdidas de tiempo. No obstante, esta práctica incrementa significativamente los costos de almacenamiento. Por el contrario, realizar checkpoints con menor frecuencia puede abaratar los costos, pero aumenta el riesgo de pérdida de datos en caso de fallos, una realidad común en entornos de entrenamiento distribuidos con miles de aceleradores.
Durante el entrenamiento del modelo Meta Llama 3, se registró un fallo cada tres horas, adjudicando el 60% de estos incidentes a problemas con las GPUs, mientras que el resto se debió a inconvenientes con las redes, CPUs y discos. Esta inestabilidad no solo encarece los costos de entrenamiento, sino que también retrasa el tiempo de lanzamiento al mercado. Aunque los checkpoints frecuentes pueden saturar redes y sobrecargar el almacenamiento, encontrar un equilibrio es crucial.
Para resolver estos problemas, AWS ha lanzado el «checkpointing» en capas gestionadas en Amazon SageMaker HyperPod, una infraestructura que promete escalar y acelerar el desarrollo de modelos de IA generativa. Esta innovación utiliza la memoria del CPU para almacenar checkpoints de alto rendimiento, replicando los datos automáticamente en nodos adyacentes para mejorar la fiabilidad. SageMaker HyperPod detecta automáticamente los problemas en los nodos y reemplaza los defectuosos, reanudando el entrenamiento y maximizando el rendimiento.
La nueva función ha sido probada en grandes clústeres distribuidos, con capacidades que van desde cientos hasta más de 15,000 GPUs, logrando guardar checkpoints en cuestión de segundos. No se requiere experiencia técnica avanzada para implementar esta solución, que puede integrarse fácilmente en scripts de PyTorch.
Además, el «checkpointing» en capas gestionadas permite a las organizaciones determinar la frecuencia y las políticas de retención, usando Amazon S3 como 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 métodos tradicionales que dependen de almacenamiento remoto persistente.
Para obtener los mejores resultados, se recomienda configurar la escritura de checkpoints en la capa de memoria frecuentemente, mientras que las copias en Amazon S3 pueden realizarse de forma menos regular. Con estas capacidades, la combinación de «managed tiered checkpointing» y SageMaker HyperPod prevé mantener un rendimiento alto en el entrenamiento, incluso en entornos distribuidos de gran escala propensos a fallos.
