Las organizaciones que buscan expandir su infraestructura de inteligencia artificial para manejar modelos de billones de parámetros enfrentan un dilema crítico: reducir el tiempo de entrenamiento a un costo más bajo o acelerarlo a un costo superior. Una estrategia comúnmente utilizada para minimizar los tiempos de recuperación es el «checkpointing»; sin embargo, este método puede elevar significativamente los costos de almacenamiento. Por el contrario, una frecuencia menor de «checkpointing» podría reducir estos costos, pero incrementa el riesgo de perder progresos valiosos durante fallos, frecuentes en entornos distribuidos con miles de aceleradores.
Un caso reciente durante el entrenamiento del modelo Meta Llama 3 evidenció problemas, con un fallo cada tres horas, donde las GPUs fueron responsables del 60% de los incidentes. Otros problemas surgieron en redes, CPUs y discos. Estas interrupciones pueden acarrear la pérdida de días de trabajo, incrementando los costos y retrasando el tiempo de salida al mercado. A pesar de que los checkpoints frecuentes pueden saturar redes y sobrecargar el almacenamiento, comprometiendo el rendimiento, encontrar un balance es vital.
Para enfrentar estos desafíos, AWS ha introducido el «checkpointing» en capas gestionadas en Amazon SageMaker HyperPod. Esta infraestructura está diseñada para escalar y acelerar el desarrollo de modelos de inteligencia artificial generativa, utilizando la memoria del CPU para almacenar checkpoints de alto rendimiento. Además, replicas automáticas en nodos adyacentes mejoran la fiabilidad. SageMaker HyperPod detecta problemas en los nodos y los reemplaza para reanudar el entrenamiento de manera eficiente, ayudando a maximizar el rendimiento del modelo.
Esta innovación ha sido probada en grandes clústeres distribuidos, desde cientos hasta más de 15,000 GPUs, logrando guardar checkpoints en segundos. Su implementación no requiere experiencia técnica avanzada, facilitando su integración en scripts de entrenamiento de PyTorch.
Asimismo, el «checkpointing» en capas gestionadas permite a las organizaciones establecer políticas de retención personalizadas tanto para almacenamiento en memoria como persistente, utilizando Amazon S3 como respaldo. Esta tecnología optimiza significativamente la gestión de checkpoints en comparación con métodos tradicionales que dependen del almacenamiento remoto persistente.
Los mejores resultados se logran configurando la escritura de checkpoints en la memoria con frecuencia, mientras que las copias en Amazon S3 se realizan menos frecuentemente. Combinando el «managed tiered checkpointing» con SageMaker HyperPod, se mantiene un alto rendimiento en el entrenamiento, incluso en entornos a gran escala propensos a fallos.