La infraestructura moderna de IA atiende múltiples cargas de trabajo simultáneas en el mismo clúster, desde el preentrenamiento y el ajuste del modelo básico (FM) hasta la inferencia y evaluación de la producción. En este entorno compartido, la demanda de aceleradores de IA fluctúa continuamente a medida que las cargas de trabajo de inferencia aumentan con los patrones de tráfico y los experimentos se completan y liberan recursos. A pesar de esta disponibilidad dinámica de los aceleradores de IA, las cargas de trabajo de capacitación tradicionales permanecen atrapadas en su asignación informática inicial, incapaces de aprovechar la capacidad informática inactiva sin intervención manual.
Amazon SageMaker HyperPod ahora admite el entrenamiento elástico, lo que permite que sus cargas de trabajo de aprendizaje automático (ML) escale automáticamente según la disponibilidad de recursos. En esta publicación, demostramos cómo el entrenamiento elástico lo ayuda a maximizar la utilización de la GPU, reducir costos y acelerar el desarrollo de modelos mediante la adaptación dinámica de recursos, mientras mantiene la calidad del entrenamiento y minimiza la intervención manual.
Cómo la asignación estática afecta la utilización de la infraestructura
Considere un clúster de 256 GPU que ejecuta cargas de trabajo de entrenamiento e inferencia. Durante las horas de menor actividad nocturna, la inferencia puede liberar 96 GPU. Eso deja 96 GPU inactivas y disponibles para acelerar el entrenamiento. Los trabajos de formación tradicionales se ejecutan a una escala fija; Estos trabajos no pueden absorber la capacidad informática inactiva. Como resultado, un único trabajo de entrenamiento que comienza con 32 GPU se bloquea en esta configuración inicial, mientras que 96 GPU adicionales permanecen inactivas; Esto se traduce en 2304 horas de GPU desperdiciadas por día, lo que representa miles de dólares gastados diariamente en inversiones en infraestructura infrautilizadas. El problema se agrava a medida que aumenta el tamaño del grupo.
Escalar dinámicamente la capacitación distribuida es técnicamente complejo. Incluso con una infraestructura que admita la elasticidad, es necesario detener trabajos, reconfigurar recursos, ajustar la paralelización y volver a bloquear los puntos de control. Esta complejidad se ve agravada por la necesidad de mantener el progreso del entrenamiento y la precisión del modelo a lo largo de estas transiciones. A pesar del soporte subyacente de SageMaker HyperPod con Amazon EKS y marcos como PyTorch y NeMo, la intervención manual aún puede consumir horas de ingeniería de aprendizaje automático. La necesidad de ajustar repetidamente las ejecuciones de entrenamiento en función de la disponibilidad del acelerador distrae a los equipos de su trabajo real en el desarrollo de modelos.
El intercambio de recursos y la preferencia por la carga de trabajo añaden otra capa de complejidad. Los sistemas actuales carecen de la capacidad de manejar con elegancia solicitudes de recursos parciales de cargas de trabajo de mayor prioridad. Considere un escenario en el que un trabajo de ajuste crítico requiere 8 GPU de un clúster donde una carga de trabajo previa al entrenamiento ocupa las 32 GPU. Los sistemas actuales obligan a tomar una decisión binaria: detener todo el trabajo de preentrenamiento o negar recursos a la carga de trabajo de mayor prioridad, aunque 24 GPU serían suficientes para continuar el preentrenamiento a escala reducida. Esta limitación lleva a las organizaciones a aprovisionar en exceso la infraestructura para evitar la contención de recursos, lo que genera colas más grandes de trabajos pendientes, mayores costos y una reducción de la eficiencia del clúster.
Descripción general de la solución
SageMaker HyperPod ahora ofrece entrenamiento elástico. Las cargas de trabajo de capacitación pueden ampliarse automáticamente para utilizar los aceleradores disponibles y contraerse elegantemente cuando se necesitan recursos en otros lugares, todo ello manteniendo la calidad de la capacitación. SageMaker HyperPod gestiona la compleja orquestación de la gestión de puntos de control, la reasignación de rangos y la coordinación de procesos, minimizando la intervención manual y ayudando a los equipos a centrarse en el desarrollo de modelos en lugar de en la gestión de infraestructura.
El operador de capacitación SageMaker HyperPod se integra con el plano de control y el programador de recursos de Kubernetes para tomar decisiones de escala. Supervisa los eventos del ciclo de vida del pod, la disponibilidad de los nodos y las señales de prioridad del programador. Esto le permite detectar oportunidades de escalamiento casi instantáneamente, ya sea de recursos recientemente disponibles o de nuevas solicitudes de cargas de trabajo de mayor prioridad. Antes de iniciar cualquier transición, el operador evalúa las posibles acciones de escalamiento con respecto a las políticas configuradas (límites mínimos y máximos de nodos, límites de frecuencia de escalamiento) antes de iniciar las transiciones.
Flujo de trabajo de eventos de escalamiento de entrenamiento elástico
El entrenamiento elástico agrega o elimina réplicas paralelas de datos mientras mantiene constante el tamaño del lote global. Cuando los recursos están disponibles, se unen nuevas réplicas y aceleran el rendimiento sin afectar la convergencia. Cuando una carga de trabajo de mayor prioridad necesita recursos, el sistema elimina las réplicas en lugar de finalizar todo el trabajo. El entrenamiento continúa a capacidad reducida.
Cuando ocurre un evento de escala, el operador transmite una señal de sincronización a todos los rangos. Cada proceso completa su paso actual y guarda el estado utilizando PyTorch Distributed Checkpoint (DCP). A medida que se unen nuevas réplicas o salen réplicas existentes, el operador vuelve a calcular las asignaciones de rango e inicia el reinicio del proceso durante todo el trabajo de capacitación. Luego, DCP carga y redistribuye los datos del punto de control para que coincidan con el nuevo recuento de réplicas, asegurándose de que cada trabajador tenga el modelo y el estado del optimizador correctos. La capacitación se reanuda con réplicas ajustadas y el tamaño de lote global constante garantiza que la convergencia no se vea afectada.
Para los clústeres que utilizan Kueue (incluido el gobierno de tareas de SageMaker HyperPod), el entrenamiento elástico implementa una gestión inteligente de la carga de trabajo a través de múltiples solicitudes de admisión. El operador primero solicita los recursos mínimos requeridos con alta prioridad y luego solicita incrementalmente capacidad adicional con menor prioridad. Este enfoque permite una preferencia parcial: cuando las cargas de trabajo de mayor prioridad necesitan recursos, solo se revocan las réplicas de menor prioridad, lo que permite que el entrenamiento continúe en la base garantizada en lugar de terminar por completo.
Empezando con el entrenamiento elástico
En las siguientes secciones, lo guiaremos a través de la instalación y configuración del entrenamiento elástico en SageMaker HyperPod.
Requisitos previos
Antes de integrar el entrenamiento elástico en su carga de trabajo de entrenamiento, asegúrese de que su entorno cumpla con los siguientes requisitos:
Configurar el aislamiento del espacio de nombres y los controles de recursos
Si utiliza el escalado automático del clúster (como Karpenter), establezca ResourceQuotas a nivel de espacio de nombres. Sin ellos, las solicitudes de recursos de la capacitación elástica pueden desencadenar un aprovisionamiento ilimitado de nodos. ResourceQuotas limita el máximo de recursos que los trabajos pueden solicitar y al mismo tiempo permite un comportamiento elástico dentro de límites definidos.
El siguiente código es un ejemplo de ResourceQuota para un espacio de nombres limitado a 8 instancias ml.p5.48xlarge (cada instancia tiene 8 GPU NVIDIA H100, 192 vCPU y 640 GiB de memoria, por lo que 8 instancias = 64 GPU, 1536 vCPU y 5120 GiB de memoria):
Recomendamos organizar las cargas de trabajo en espacios de nombres separados por equipo o proyecto, con asignaciones de control de acceso basado en roles (RBAC) de AWS Identity and Access Management (IAM) para admitir un control de acceso y un aislamiento de recursos adecuados.
Construir un contenedor de entrenamiento HyperPod
El operador de capacitación de HyperPod utiliza un iniciador PyTorch personalizado del paquete HyperPod Elastic Agent Python para detectar eventos de escala, coordinar operaciones de puntos de control y administrar el proceso de encuentro cuando cambia el tamaño del mundo. Instale el agente elástico, luego reemplace torchrun con hyperpodrun en su comando de inicio. Para obtener más detalles, consulte Agente elástico HyperPod.
El siguiente código es un ejemplo de configuración de contenedor de entrenamiento:
Habilite el escalado elástico en el código de entrenamiento:
Complete los siguientes pasos para habilitar el escalado elástico en su código de entrenamiento:
Agregue la importación del agente elástico HyperPod a su secuencia de comandos de entrenamiento para detectar cuándo ocurren eventos de escalado:
Modifique su ciclo de entrenamiento para verificar eventos elásticos después de cada lote de entrenamiento. Cuando se detecta un evento de escala, su proceso de capacitación debe guardar un punto de control y salir correctamente, lo que permite al operador reiniciar el trabajo con un nuevo tamaño mundial:
El patrón clave aquí es verificar elastic_event_detected() durante el ciclo de entrenamiento y regresar de la función de entrenamiento después de guardar un punto de control. Esto permite al operador de capacitación coordinar la transición de escala entre todos los trabajadores.
Finalmente, implemente funciones de carga y guardado de puntos de control utilizando PyTorch DCP. DCP es esencial para el entrenamiento elástico porque vuelve a fragmentar automáticamente los estados del modelo y del optimizador cuando el trabajo se reanuda con un número diferente de réplicas:
Para escenarios de entrenamiento de una sola época en los que cada muestra de datos debe verse exactamente una vez, debe conservar el estado del cargador de datos en todos los eventos de escalado. Sin esto, cuando su trabajo se reanude con un tamaño mundial diferente, las muestras procesadas previamente pueden repetirse u omitirse, lo que afecta la calidad de la capacitación. Un cargador de datos con estado guarda y restaura la posición del cargador de datos durante los puntos de control, asegurándose de que el entrenamiento continúe desde el punto exacto donde se detuvo. Para obtener detalles de implementación, consulte la guía del cargador de datos con estado en la documentación.
Enviar trabajo de entrenamiento elástico
Con su contenedor de capacitación creado y el código instrumentado, está listo para enviar un trabajo de capacitación elástico. La especificación del trabajo define cómo escala su carga de trabajo de capacitación en respuesta a la disponibilidad de recursos del clúster a través de la configuración de elasticPolicy.
Cree una especificación HyperPodPyTorchJob que defina su comportamiento de escalado elástico usando el siguiente código:
La configuración de elasticPolicy controla cómo responde su trabajo de capacitación a los cambios de recursos:
minReplicas y maxReplicas: definen los límites de escala. Su trabajo siempre mantendrá al menos minReplicas y nunca excederá maxReplicas, manteniendo un uso de recursos predecible. replicaIncrementStep frente a replicaDiscreteValues: elija un enfoque para escalar la granularidad. Utilice replicaIncrementStep para un escalado uniforme (por ejemplo, un paso de 2 significa escalar a 2, 4, 6, 8 nodos). Utilice replicaDiscreteValues: [2, 4, 8] para especificar las configuraciones permitidas exactas. Esto es útil cuando ciertos tamaños de mundo funcionan mejor para la estrategia de paralelización de su modelo. GracefulShutdownTimeoutInSeconds: esto le da a su proceso de capacitación tiempo para completar los puntos de control antes de que el operador fuerce un apagado. Configúrelo según el tamaño de su punto de control y el rendimiento del almacenamiento. scalingTimeoutInSeconds: esto introduce un retraso de estabilización antes de la ampliación para evitar la destrucción cuando los recursos fluctúan rápidamente. El operador espera este período después de detectar recursos disponibles antes de desencadenar un evento de ampliación. falloyScaleDownTimeoutInSeconds: cuando los pods fallan o fallan, el operador espera este período de recuperación antes de reducirlos. Esto evita reducciones innecesarias debido a fallas transitorias.
El entrenamiento elástico incorpora mecanismos anti-golpe para mantener la estabilidad en entornos con disponibilidad de recursos que fluctúa rápidamente. Estas protecciones incluyen períodos de estabilidad mínimos obligatorios entre eventos de escala y una estrategia de retroceso exponencial para transiciones frecuentes. Al evitar fluctuaciones excesivas, el sistema garantiza que los trabajos de capacitación puedan lograr un progreso significativo en cada punto de escala en lugar de verse abrumados por operaciones frecuentes en los puntos de control. Puede ajustar estas políticas anti-thrashing en la configuración de políticas elásticas, lo que permite un enfoque equilibrado entre el escalado responsivo y la estabilidad del entrenamiento que se alinea con la dinámica específica del clúster y los requisitos de carga de trabajo.
Luego puede enviar el trabajo usando kubectl o la CLI de SageMaker HyperPod, como se describe en la documentación:
Uso de recetas de SageMaker HyperPod
Hemos creado recetas SageMaker HyperPod para entrenamiento elástico para FM disponibles públicamente, incluidos Llama y GPT-OSS. Estas recetas proporcionan configuraciones validadas previamente que manejan la estrategia de paralelización, los ajustes de hiperparámetros y la administración de puntos de control automáticamente, lo que requiere solo cambios de configuración de YAML para especificar la política elástica sin modificaciones de código. Los equipos simplemente especifican los límites mínimos y máximos de los nodos en sus especificaciones de trabajo, y el sistema gestiona toda la coordinación de escalamiento a medida que fluctúan los recursos del clúster.
Las recetas también admiten configuraciones específicas de escala a través del campo scale_config, por lo que puede definir diferentes hiperparámetros (tamaño de lote, tasa de aprendizaje) para cada tamaño mundial. Esto es particularmente útil cuando el escalado requiere ajustar la distribución de lotes o permitir tamaños de lotes desiguales. Para ver ejemplos detallados, consulte el repositorio de recetas de SageMaker HyperPod.
Resultados de rendimiento
Para demostrar el impacto del entrenamiento elástico, ajustamos un modelo Llama-3 70B en el conjunto de datos TAT-QA utilizando un clúster SageMaker HyperPod con hasta 8 instancias ml.p5.48xlarge. Este punto de referencia ilustra cómo se desempeña el entrenamiento elástico en la práctica cuando se escala dinámicamente en respuesta a la disponibilidad de recursos, simulando un entorno realista donde las cargas de trabajo de entrenamiento e inferencia comparten la capacidad del clúster.
Evaluamos el entrenamiento elástico en dos dimensiones clave: rendimiento del entrenamiento y convergencia del modelo durante las transiciones de escala. Observamos una mejora constante en el rendimiento en diferentes configuraciones de escala de 1 nodo a 8 nodos, como se muestra en las siguientes figuras. El rendimiento del entrenamiento mejoró de 2000 tokens/segundo en 1 nodo y hasta 14000 tokens/segundo en 8 nodos. A lo largo de la ejecución del entrenamiento, la pérdida continuó disminuyendo a medida que el entrenamiento del modelo continuó convergiendo.
Rendimiento del entrenamiento con Elastic Training
Convergencia de modelos con Elastic Training
Integración con las capacidades de SageMaker HyperPod
Más allá de sus capacidades básicas de escalamiento, la capacitación elástica aprovecha la integración con las capacidades de infraestructura de SageMaker HyperPod. Las políticas de gobernanza de tareas activan automáticamente eventos de escalamiento cuando cambian las prioridades de la carga de trabajo, lo que permite que la capacitación genere recursos para cargas de trabajo de inferencia o evaluación de mayor prioridad. El soporte para los planes de capacitación de SageMaker permite que la capacitación se escale de manera oportunista utilizando tipos de capacidad con costos optimizados mientras se mantiene la resiliencia a través de la reducción automática cuando se recuperan instancias puntuales. El complemento de observabilidad SageMaker HyperPod complementa estas capacidades al proporcionar información detallada sobre la ampliación de eventos, el rendimiento de los puntos de control y la progresión del entrenamiento, lo que ayuda a los equipos a monitorear y optimizar sus implementaciones de entrenamiento elástico.
Conclusión
La capacitación elástica en SageMaker HyperPod aborda el problema de los recursos desperdiciados en los clústeres de IA. Los trabajos de capacitación ahora pueden escalar automáticamente a medida que los recursos estén disponibles sin requerir ajustes manuales de la infraestructura. La arquitectura técnica del entrenamiento elástico mantiene la calidad del entrenamiento durante las transiciones de escala. Al preservar el tamaño del lote global y la tasa de aprendizaje en diferentes configuraciones de datos paralelos, el sistema mantiene propiedades de convergencia consistentes independientemente de la escala actual.
Puede esperar tres beneficios principales. Primero, desde una perspectiva operativa, la reducción de los ciclos de reconfiguración manual cambia fundamentalmente la forma en que trabajan los equipos de ML. Los ingenieros pueden centrarse en la innovación y el desarrollo de modelos en lugar de en la gestión de la infraestructura, lo que mejora significativamente la productividad del equipo y reduce los gastos operativos. En segundo lugar, la eficiencia de la infraestructura experimenta mejoras espectaculares a medida que las cargas de trabajo de capacitación consumen dinámicamente la capacidad disponible, lo que lleva a reducciones sustanciales en las horas de inactividad de la GPU y los correspondientes ahorros de costos. En tercer lugar, el tiempo de comercialización se acelera considerablemente a medida que los trabajos de capacitación se escalan automáticamente para utilizar los recursos disponibles, lo que permite ciclos de desarrollo e implementación de modelos más rápidos.
Para comenzar, consulte la guía de documentación. Implementaciones de muestra y recetas están disponibles en el repositorio de GitHub.
Acerca de los autores
Roy Allela es arquitecto senior de soluciones especializado en IA/ML en AWS. Ayuda a los clientes de AWS, desde pequeñas empresas emergentes hasta grandes empresas, a capacitar e implementar modelos básicos de manera eficiente en AWS. Tiene experiencia en ingeniería de microprocesadores y le apasionan los problemas de optimización computacional y la mejora del rendimiento de las cargas de trabajo de IA. Puedes conectarte con Roy en LinkedIn.
Anirudh Viswanathan es gerente técnico senior de productos en AWS con el equipo de SageMaker, donde se enfoca en aprendizaje automático. Tiene una Maestría en Robótica de la Universidad Carnegie Mellon y un MBA de la Wharton School of Business. Anirudh es un inventor reconocido en más de 50 patentes de IA/ML. Le gusta correr largas distancias, explorar galerías de arte y asistir a espectáculos de Broadway. Puedes conectarte con Anirudh en LinkedIn.
Arun Kumar Lokanatha es arquitecto senior de soluciones de aprendizaje automático en Amazon SageMaker AI. Tiene una maestría de la UIUC con especialización en ciencia de datos. Se especializa en cargas de trabajo de IA generativa, ayudando a los clientes a crear e implementar LLM utilizando SageMaker HyperPod, trabajos de capacitación de SageMaker y capacitación distribuida de SageMaker. Fuera del trabajo, le gusta correr, hacer senderismo y cocinar.
Oleg Talalov es ingeniero sénior de desarrollo de software en AWS y trabaja en el equipo SageMaker HyperPod, donde se centra en el aprendizaje automático y la infraestructura informática de alto rendimiento para la formación en aprendizaje automático. Tiene una maestría de la Universidad Politécnica de San Petersburgo Pedro el Grande. Oleg es un inventor de múltiples tecnologías de IA/ML y le gusta andar en bicicleta, nadar y correr. Puedes conectarte con Oleg en LinkedIn
Qianlin Liang es ingeniero de desarrollo de software en AWS en el equipo de SageMaker, donde se enfoca en sistemas de inteligencia artificial. Tiene un doctorado. en Ciencias de la Computación de la Universidad de Massachusetts Amherst. Su investigación desarrolla técnicas de sistemas para un aprendizaje automático eficiente y resiliente. Fuera del trabajo, le gusta correr y fotografiar. Puede conectarse con Qianlin en LinkedIn.
Trevor Harvey es especialista principal en IA generativa en Amazon Web Services (AWS) y arquitecto de soluciones profesional certificado por AWS. En AWS, Trevor trabaja con clientes para diseñar e implementar soluciones de aprendizaje automático y lidera estrategias de comercialización de servicios de IA generativa.
Anirban Roy es ingeniero principal en AWS del equipo de SageMaker y se centra principalmente en la infraestructura de entrenamiento de IA, la resiliencia y la observabilidad. Tiene una maestría en Ciencias de la Computación del Instituto de Estadística de la India en Calcuta. Anirban es un experimentado creador de sistemas de software distribuido con más de 20 años de experiencia y múltiples patentes y publicaciones. Le gusta andar en bicicleta, leer no ficción, hacer jardinería y viajar por la naturaleza. Puedes conectarte con Anirban en LinkedIn
Arun Nagarajan es ingeniero principal del equipo de inteligencia artificial de Amazon SageMaker, donde actualmente se enfoca en la capacitación distribuida en toda la pila. Desde que se unió al equipo de SageMaker durante su año de lanzamiento, Arun ha contribuido a múltiples productos dentro de SageMaker AI, incluidas soluciones de inferencia en tiempo real y MLOps. Cuando no está trabajando en infraestructura de aprendizaje automático, le gusta explorar el aire libre en el noroeste del Pacífico y esquiar en las pistas.