Optimice el entrenamiento de modelos en Amazon SageMaker AI con NVIDIA Blackwell

La optimización del entrenamiento de modelos en Amazon SageMaker AI con GPU NVIDIA Blackwell cambia lo que es práctico para modelos de IA grandes. Si entrena modelos grandes hoy en día, probablemente esté trabajando con un conjunto familiar de restricciones: tamaños de lote limitados por la memoria de la GPU, longitudes de secuencia acortadas para evitar errores de falta de memoria y fragmentación del modelo que agrega una sobrecarga de comunicación a medida que escala. La memoria ampliada de Blackwell y los nuevos formatos de precisión reducen esas limitaciones directamente. Las instancias P6-B200 con 8 GPU Blackwell están disponibles en los trabajos de capacitación de IA de Amazon SageMaker y puede reservar la capacidad mediante un plan de capacitación flexible con acceso predecible, administración de costos y administración automatizada de recursos. Los trabajos de capacitación de IA de Amazon SageMaker le permiten entrenar modelos de aprendizaje automático a gran escala al aprovisionar y administrar automáticamente la infraestructura y los recursos informáticos subyacentes, para que pueda concentrarse en sus datos y algoritmos en lugar de en las operaciones de infraestructura.

Esta publicación le muestra cómo configurar trabajos de capacitación en Amazon SageMaker AI para aprovechar al máximo la arquitectura de Blackwell en AWS. Aprenderá a seleccionar tamaños de lote y longitudes de secuencia que aprovechen la memoria ampliada de Blackwell, elegirá el formato de precisión adecuado para el tamaño de su modelo (parámetros de 1B a 64B) y aplicará puntos de control de activación estratégicamente. Al final, tendrá un marco práctico para ajustar su configuración de capacitación y lanzar trabajos de capacitación distribuidos en instancias P6-B200.

Los trabajos de capacitación de Blackwell configurados correctamente pueden procesar lotes de mayor tamaño sin una fragmentación agresiva, lo que reduce la sobrecarga de comunicación y mejora el rendimiento. Las secuencias más largas se vuelven viables para tareas de dependencia de largo alcance. Con el formato de precisión adecuado, los modelos que anteriormente requerían configuraciones de múltiples nodos pueden ejecutarse en un solo nodo de 8 GPU, lo que significa ciclos de iteración más rápidos, menos gastos generales de red y menores costos de infraestructura.

Entendiendo NVIDIA Blackwell

Antes de configurar su trabajo de capacitación, es útil comprender qué diferencia a Blackwell de las generaciones anteriores de GPU. La arquitectura de doble chip de Blackwell y los Tensor Cores de quinta generación ofrecen ganancias mensurables para el entrenamiento con múltiples GPU desde el primer momento. La interconexión NVLink 5 proporciona hasta 1,8 TB/s de ancho de banda bidireccional de GPU a GPU, mientras que la mayor capacidad HBM del B200 y el mayor ancho de banda de memoria ayudan a reducir la presión de la memoria para lotes grandes, secuencias largas y cargas de trabajo de entrenamiento distribuidas.

Los ejemplos de esta publicación utilizan entrenamiento de 8 GPU de un solo nodo con modelos de transformadores que van desde parámetros de 1B a 64B. La configuración de entrenamiento utiliza PyTorch Fully Sharded Data Parallel (FSDP), una técnica de entrenamiento distribuido que fragmenta los parámetros del modelo, los gradientes y los estados del optimizador en las GPU para entrenar modelos más grandes que la memoria de una sola GPU. Los resultados cubren múltiples configuraciones con diferentes tamaños de lote, longitudes de secuencia y formatos de precisión para mostrar cuándo los diferentes enfoques ofrecen resultados óptimos.

Gestión de la memoria

La memoria ampliada de Blackwell (180 GB en B200, 268 GB en B300) le brinda espacio para optimizar en tres áreas: tamaños de lotes más grandes, fragmentación de modelos simplificada y longitudes de secuencia más largas.

Los tamaños de lote más grandes reducen la cantidad de pasos de sincronización de gradiente entre las GPU, lo que mejora el rendimiento general. La fragmentación de modelos simplificada es posible porque más memoria por GPU significa que es posible reducir el grado de paralelismo del modelo o eliminarlo por completo para algunos modelos. Menos fragmentos significan menos gastos generales de comunicación entre GPU. Las secuencias más largas permiten que los modelos procesen más contexto en una sola pasada, lo cual es fundamental para tareas de dependencia de largo alcance.

Si el rendimiento es su objetivo principal, comience con el ajuste del tamaño del lote. Si la sobrecarga de comunicación es el cuello de botella, primero simplifique la fragmentación. Si su tarea requiere un contexto de largo alcance, priorice la longitud de la secuencia. El tamaño del lote y la longitud de la secuencia aumentan el consumo de memoria y es importante encontrar un equilibrio efectivo.

Los puntos de control de activación le ayudan a equilibrar el uso de la memoria y la computación. Intercambia un mayor tiempo de cálculo (normalmente entre un 10% y un 30% de sobrecarga dependiendo de la arquitectura del modelo) por una reducción en el uso de la memoria de la GPU al recalcular las activaciones intermedias durante el paso hacia atrás en lugar de almacenarlas. La memoria liberada se puede reinvertir en lotes más grandes o secuencias más largas. Dado que la sobrecarga informática varía según la carga de trabajo, compare su configuración específica para comprender las ventajas y desventajas antes de comprometerse con los puntos de control.

Por ejemplo, en la Figura 1, comparamos tres configuraciones de entrenamiento para un LLM de parámetro 1B usando precisión MXFP8 con una longitud de secuencia de 8K. Sin puntos de control de activación (BS=1), el rendimiento es de ~6.000 tokens/seg, pero la memoria máxima es alta, 15,5 GB. Habilitar los puntos de control de activación en el mismo tamaño de lote reduce drásticamente la memoria a 2,3 GB (ya que las activaciones intermedias se recalculan en lugar de almacenarse), pero el rendimiento también disminuye ligeramente debido a esa sobrecarga de recálculo. La recompensa clave viene en la tercera barra: con los puntos de control de activación habilitados y el tamaño del lote aumentado a 16, la memoria liberada permite un lote mucho más grande, elevando el rendimiento a ~51 000 tokens/seg (aproximadamente 8 veces la línea base), mientras que la memoria máxima sube a 22,8 GB, aún dentro de los límites de la GPU.

Figura 1. Diferencia de rendimiento con y sin puntos de control de activación

Para decidir si los puntos de control de activación tienen sentido para su carga de trabajo, considere el tamaño de su modelo y el uso de memoria:

Modelos pequeños (hasta ~14B de parámetros): generalmente no se necesitan puntos de control de activación. Con la memoria ampliada de Blackwell, la mayoría de los modelos pequeños caben cómodamente sin ella. Si está ejecutando en el extremo superior de este rango y alcanza la presión de la memoria, el punto de control de activación agrega una sobrecarga de cómputo a cambio de ahorros significativos de memoria, que puede reinvertir en lotes más grandes. Modelos grandes (~14B+ parámetros): en este tamaño de modelo, el consumo de memoria oscila entre 87 y 171 GB, según el tamaño del lote y la longitud de la secuencia. Sin puntos de control de activación, la mayoría de las configuraciones fallan con errores de falta de memoria (OOM) de CUDA. Cuando agrega puntos de control, la memoria liberada le permite aumentar el tamaño del lote lo suficiente como para que el rendimiento mejore a pesar de la sobrecarga informática adicional. Para los modelos grandes, los puntos de control no son opcionales. Es un requisito previo para un entrenamiento estable.

Formatos de precisión

Los Tensor Cores de quinta generación de Blackwell proporcionan aceleración de hardware para formatos de precisión reducida (FP8, MXFP8 y NVFP4), lo que los convierte principalmente en optimizaciones de rendimiento en lugar de técnicas de ahorro de memoria. El uso de una precisión más baja reduce los requisitos de ancho de banda de la memoria y, al mismo tiempo, aumenta la cantidad de operaciones que la GPU puede ejecutar por ciclo. Sin embargo, el entrenamiento de precisión reducida es más o menos neutral en cuanto a la memoria de forma predeterminada, donde Transformer Engine mantiene pesos primarios de alta precisión (para actualizaciones del optimizador) y copias cuantificadas, por lo que los formatos de menor precisión no se traducen directamente en un menor uso de memoria. La cuantificación en sí misma introduce una sobrecarga (convertir entre formatos de precisión y mantener múltiples copias de pesos en la memoria), lo que significa que el beneficio neto depende del tamaño del modelo y de si el entrenamiento está vinculado a la computación o a la memoria. Si bien NVFP4 ofrece el mayor rendimiento, sus beneficios de rendimiento escalan principalmente con modelos grandes y cargas de trabajo de inferencia, donde no se necesitan pesos primarios.

Para cargas de trabajo vinculadas a la computación (normalmente modelos más pequeños), la velocidad de cálculo es el factor limitante y la sobrecarga de cuantificación compensa parcialmente las ganancias de rendimiento derivadas de una menor precisión. Para cargas de trabajo vinculadas a la memoria (normalmente modelos más grandes), el movimiento de datos es el cuello de botella, y la huella de memoria reducida de los formatos de menor precisión aborda directamente la restricción, generando ganancias más significativas:

Modelos pequeños (hasta ~14B de parámetros): en este tamaño de modelo, los formatos de precisión reducida (FP8, MXFP8, NVFP4) ofrecen mejoras de rendimiento modestas y similares con respecto a FP16, ya que la sobrecarga de cuantificación reduce la ventaja de velocidad. El ajuste del tamaño de lote tiende a ofrecer ganancias más significativas que la selección precisa del formato. Comience con FP8 para obtener un mayor rendimiento. Tiene menos gastos generales que MXFP8 o NVFP4 y, a menudo, es un buen valor predeterminado para la mayoría de las cargas de trabajo de modelos pequeños. Tenga en cuenta que con la configuración predeterminada de TransformerEngine, los formatos de precisión reducida utilizan más memoria que FP16, ya que TransformerEngine mantiene los pesos con mayor precisión y los lanza sobre la marcha. Si la memoria es una restricción y su optimizador la admite, use quantized_model_init para almacenar pesos directamente en FP8, reduciendo la memoria por debajo de los niveles de FP16. Modelos grandes (~14B+ parámetros): Aquí es donde la precisión reducida produce su mayor impacto. FP8 normalmente proporciona un sólido equilibrio entre rendimiento y eficiencia de la memoria. Si bien MXFP8 es teóricamente más eficiente en cuanto a memoria, su sobrecarga de transposición compensa parcialmente esa ventaja en la práctica. Sin embargo, si la estabilidad de la convergencia o la precisión numérica es una prioridad para su carga de trabajo, MXFP8 puede ser la mejor opción, ya que su esquema de cuantificación más fino tiende a preservar la precisión del modelo de manera más confiable que FP8. Para modelos grandes donde la memoria es el principal cuello de botella, NVFP4 puede ofrecer ganancias de rendimiento adicionales, ya que su ventaja en la velocidad de multiplicación de matrices aumenta con el tamaño del modelo. Para lograr esos beneficios se requiere una inversión significativa en ingeniería. Utilice recetas a nivel de marco de Megatron Core, que proporcionan configuraciones NVFP4 validadas, en lugar de implementarlas desde cero.

TransformerEngine de NVIDIA maneja la complejidad de la implementación: conmutación automática de precisión mixta, núcleos fusionados y escalado dinámico de pérdidas. Antes de pasar a producción, valide la convergencia mediante el seguimiento de las curvas de pérdida en todos los formatos para confirmar que la precisión elegida cumpla con los requisitos de exactitud.

No todas las cargas de trabajo se benefician de una optimización agresiva. Si su modelo se entrena cómodamente dentro de los límites de la memoria y cumple con sus requisitos de rendimiento con FP16, es posible que la complejidad adicional de los formatos de precisión reducida no valga la pena el esfuerzo de ingeniería. Comience con mediciones de referencia y luego optimice solo los cuellos de botella que pueda medir.

Introducción a la formación de Blackwell en Amazon SageMaker AI

Las secciones anteriores cubren las decisiones clave: cuánta memoria tiene para trabajar, si los puntos de control de activación tienen sentido para el tamaño de su modelo y qué formato de precisión se adapta a su carga de trabajo. Las siguientes secciones ponen en práctica esas decisiones mediante trabajos de capacitación en inteligencia artificial de Amazon SageMaker.

Amazon SageMaker AI proporciona un entorno totalmente administrado para la capacitación distribuida en instancias de Blackwell, el manejo del aprovisionamiento de instancias, la orquestación de contenedores y la integración con servicios de AWS como Amazon Simple Storage Service (Amazon S3), Amazon CloudWatch, Amazon Elastic Container Registry (Amazon ECR) y AWS Identity and Access Management (AWS IAM).

Requisitos previos

Antes de comenzar, confirme que tiene:

Lanzar un trabajo de formación

Para iniciar su trabajo de capacitación, complete los siguientes pasos.

Paso 1: crea tu guión

Descargue el archivo fsdp.py del ejemplo de FSDP del repositorio de NVIDIA TransformerEngine. Este script implementa el entrenamiento FSDP y acepta hiperparámetros como argumentos de la línea de comandos.

Paso 2: crear el script del punto de entrada

Prepare un archivo train.sh para configurar torchrun e iniciar el script de entrenamiento:

Paso 3: construye y empuja tu contenedor

Cree un contenedor Docker personalizado que amplíe los contenedores de aprendizaje profundo (DLC) de AWS, incluya fsdp.py y train.sh y tenga instalado TransformerEngine 2.11. El DLC proporciona una imagen base validada con PyTorch y las bibliotecas CUDA necesarias para la compatibilidad con Blackwell. Aquí está el Dockerfile que puedes usar:

Una vez compilado, cree un repositorio privado de Amazon ECR si aún no tiene uno y envíe la imagen a Amazon ECR (tenga en cuenta el repositoryUri del resultado para usarlo en la etiqueta de la ventana acoplable y en los comandos de inserción). Para obtener instrucciones de compilación detalladas, consulte Adaptación de su propio contenedor Docker para que funcione con Amazon SageMaker AI.

Paso 4: capacidad segura

Reserve capacidad a través de un plan de capacitación flexible para un acceso predecible a la tarifa estándar o utilice Managed Spot Training para cargas de trabajo con costos optimizados. Utilice planes de capacitación flexibles para ejecuciones de capacitación de producción que requieran reserva de capacidad y que estén diseñadas para brindar disponibilidad continua; Utilice Spot para cargas de trabajo de experimentación y tolerantes a fallas donde la reducción de costos supera el riesgo de interrupción. Las instancias puntuales están sujetas a interrupciones, así que asegúrese de que su script de entrenamiento guarde los puntos de control en Amazon S3 a intervalos regulares. Amazon SageMaker AI reanuda automáticamente un trabajo puntual interrumpido si proporciona un checkpoint_s3_uri en la configuración de su estimador. Cree un plan de capacitación flexible para obtener capacidad predecible en la consola SageMaker AI y seleccione "trabajo de capacitación" como recurso de destino. Si su trabajo puede tolerar reinicios y la prioridad es la reducción de costos, puede elegir Spot, donde primero debe aumentar su cuota para el uso del trabajo de capacitación puntual ml.p6-b200.48xlarge. Nota: Los planes de capacitación flexibles reservan capacidad e incurren en cargos durante la duración del plan, independientemente de si los trabajos de capacitación se están ejecutando activamente. Revise los detalles de precios antes de crear un plan.

Paso 5: envíe el trabajo de capacitación

Reemplace los valores del marcador de posición con el ARN de su plan de capacitación real y el URI de la imagen ECR, luego ejecute el siguiente código desde su entorno de desarrollo local o una instancia de cuaderno de SageMaker AI:

Para el entrenamiento puntual administrado, reemplace Training_plan con use_spot_instances=True, configure max_run y max_wait y agregue un checkpoint_s3_uri para la reanudación automática.

Paso 6: Supervise su trabajo de formación

Amazon SageMaker AI transmite registros a Amazon CloudWatch automáticamente. En la consola de SageMaker AI, navegue hasta Trabajos de capacitación y seleccione su trabajo para encontrar el grupo de registros de CloudWatch. Abra /aws/sagemaker/TrainingJobs y busque líneas [rango 0] para valores de pérdida y rendimiento. Para confirmar que ha cargado el formato de precisión, busque mensajes como "Usando la receta FP8" o "MXFP8 habilitado".

Si su trabajo se detiene debido a un error de falta de memoria (OOM) de CUDA, el registro muestra el tamaño de la asignación. Reduzca el tamaño del lote o la longitud de la secuencia o agregue puntos de control de activación si aún no lo ha hecho.

Limpieza

Para evitar cargos continuos después de sus pruebas, detenga cualquier trabajo de capacitación en ejecución en la consola de Amazon SageMaker AI (tenga en cuenta que esto no cancela un plan de capacitación flexible). Advertencia: Los siguientes pasos de limpieza eliminan permanentemente recursos y no se pueden deshacer. Verifique que haya realizado una copia de seguridad de todos los datos que necesita antes de continuar. Elimine su repositorio de Amazon ECR (esto elimina permanentemente las imágenes de contenedores), elimine cualquier artefacto de entrenamiento almacenado en Amazon S3 (esto elimina permanentemente los datos de entrenamiento y los puntos de control) y elimine los grupos de registros de CloudWatch (esto elimina permanentemente los registros de entrenamiento). Elimine el rol de ejecución de IAM creado para Amazon SageMaker AI.

Conclusión

En esta publicación, aprendió cómo optimizar el entrenamiento del modelo de IA en las GPU NVIDIA Blackwell mediante trabajos de entrenamiento de IA de Amazon SageMaker. Configuró tamaños de lotes y longitudes de secuencias para aprovechar la memoria ampliada de Blackwell, aplicó puntos de control de activación según el tamaño de su modelo y seleccionó formatos de precisión adecuados a su carga de trabajo. También configuró un contenedor personalizado con TransformerEngine, aseguró capacidad a través de un plan de capacitación flexible y lanzó un trabajo de capacitación distribuida en instancias ml.p6-b200.48xlarge.

Los modelos de transformadores con parámetros de 1B a 64B muestran ganancias consistentes cuando se combinan estas optimizaciones. La clave es comprender si su carga de trabajo está vinculada a la computación o a la memoria y luego aplicar los cambios de forma incremental para poder medir el impacto de cada uno.

Si está listo para comenzar, explore la documentación de Amazon SageMaker AI para revisar las opciones de instancia y los detalles de configuración, o compre un plan de capacitación flexible para reservar capacidad de Blackwell para su próxima ejecución de capacitación. Si tiene preguntas sobre su carga de trabajo específica, comuníquese con AWS Support o con el equipo de su cuenta de AWS.

Sobre los autores