Cómo Amazon Search M5 ahorró un 30 % en costos de capacitación de LLM mediante el uso de AWS Trainium | Servicios web de Amazon

Durante décadas, Amazon ha sido pionero e innovado en el aprendizaje automático (ML), brindando experiencias agradables a sus clientes. Desde los primeros días, Amazon ha utilizado el aprendizaje automático para diversos casos de uso, como recomendaciones de libros, búsquedas y detección de fraude. Al igual que el resto de la industria, los avances del hardware acelerado han permitido a los equipos de Amazon buscar arquitecturas modelo utilizando redes neuronales y aprendizaje profundo (DL).

El programa M5 dentro de Amazon Search posee la estrategia de aprendizaje por descubrimiento para Amazon y crea modelos a gran escala en multilingües, multilocales, multientidades, multitarea y multimodales, como texto, imágenes y videos. El programa M5 ha estado brindando incorporaciones universales y modelos básicos a gran escala a cientos de equipos de aprendizaje automático en todo Amazon, al tiempo que mantiene controles estrictos sobre la optimización de costos. Para lograrlo, el equipo de M5 evalúa periódicamente nuevas técnicas para reducir costes.

Como muchas organizaciones de ML, los aceleradores se utilizan en gran medida para acelerar el entrenamiento y la inferencia de DL. Cuando AWS lanzó aceleradores especialmente diseñados con el primer lanzamiento de Inferencia de AWS En 2020, el equipo M5 rápidamente comenzó a Utilícelos para implementar cargas de trabajo de producción de manera más eficiente., ahorrando costos y reduciendo la latencia. El año pasado, AWS lanzó su tren de AWS aceleradores, que optimizan el rendimiento por costo para desarrollar y construir modelos DL de próxima generación. En esta publicación, analizamos cómo M5 pudo reducir el costo de entrenar sus modelos en un 30 % y compartimos algunas de las mejores prácticas que aprendimos a lo largo del camino.

Instancias de tren

Con los avances en aceleradores especialmente diseñados, Amazon también ofrece aceleradores atractivos en forma de AWS Inferentia y Trainium. Como sus nombres lo indican, estos chips están optimizados para superar las necesidades de cargas de trabajo de inferencia y entrenamiento, respectivamente. Para el entrenamiento a gran escala de modelos básicos que alcanzan miles de millones de parámetros de tamaño, Trainium Instancias Trn1 y Trn1n Son opciones ideales por sus características. Las instancias Trn1 cuentan con la tecnología más avanzada NeuronCore-v2y tiene una gran cantidad de memoria y cálculo de acelerador. Las instancias Trn1n también se pueden elegir para una mayor cantidad de ancho de banda de red (1,600 Gbs), por lo que son ideales para una capacitación de alto rendimiento teniendo en cuenta la optimización de costos.

Para utilizar aceleradores, necesita una capa de software que los admita. Con chips Trn e Inf, el SDK de AWS Neuron desbloquea aceleradores especialmente diseñados por Amazon con la ayuda de PyTorch XLA. PyTorch XLA convierte el modo ansioso de PyTorch en una implementación basada en gráficos en modo diferido. Estos gráficos luego se utilizan y se compilan para usarlos con el acelerador. PyTorch Neuron (parte del SDK de Neuron) permite a los usuarios de PyTorch entrenar sus modelos en Trainium NeuronCores con unas pocas líneas de código.

Modelo y carga de trabajo

El equipo de M5 entrena e implementa modelos fundamentales y representaciones universales para ayudar a varios equipos de Amazon a brindar placer a Amazon.com clientes. Uno de esos modelos es un modelo de codificador de texto seguido de un perceptrón multicapa (MLP) con interacciones de características explícitas o implícitas definidas por la arquitectura de la red neuronal con cientos de millones de parámetros entrenables. Este modelo se entrena con miles de millones de tokens y se utiliza para generar millones de incrustaciones en una configuración de inferencia por lotes fuera de línea. Estas incorporaciones son entradas para un servicio de Amazon de nivel 1 orientado al cliente.

Usos de la infraestructura para el oleoducto de producción. Lote de AWS estrategias de cola de participación justa, utilizando un clúster trn1.32xlarge de múltiples nodos habilitado para EFA como cálculo para el entrenamiento del modelo. Funcionalmente, el proceso de producción realiza entrenamiento incremental del modelo, evaluación del modelo entrenado e inferencia por lotes fuera de línea en el modelo entrenado, todo usando PyTorch como biblioteca DL subyacente.

Goals

Deleitar a nuestros clientes es un principio fundamental. Dada la naturaleza del proceso de cara al cliente, es fundamental que todos los acuerdos de nivel de servicio (SLA) se cumplan sin regresiones. Identificamos dos criterios de aceptación críticos para adaptar nuestro proceso de producción de GPU existente y realizar la transición a Trainium:

• Calidad del modelo – La calidad de nuestros modelos impacta directamente en la experiencia del cliente. Requerimos que haya menos del 0.1% de diferencia en la calidad del modelo entre GPU y Trainium.
• Rendimiento del entrenamiento – Entrenamos nuestros modelos de forma iterativa periódicamente para brindar la experiencia más reciente a nuestros clientes. Requerimos que la convergencia del modelo se logre dentro de un período de tiempo predefinido (como 1 semana) para cumplir con nuestros SLA de producción.

En las siguientes secciones, compartimos nuestro viaje de trabajo hacia atrás a partir de este criterio y nuestros aprendizajes para respaldar cargas de trabajo de producción a escala de Amazon.

Guion de entrenamiento

Antes de comenzar con el entrenamiento del modelo, debemos realizar cambios en el script de entrenamiento para que cumpla con XLA. Dado el tamaño del modelo, utilizamos datos distribuidos paralelos (DDP) para entrenar el modelo. DDP nos permite aumentar el rendimiento del entrenamiento de modelos al aumentar la cantidad de máquinas utilizadas para ejecutar el entrenamiento de modelos, sin ningún cambio de código. Seguimos las instrucciones proporcionadas en el Tutorial de entrenamiento de Neuron PyTorch MLP para agregar construcciones específicas de XLA en nuestros scripts de capacitación. Estos cambios de código son sencillos de implementar. Los siguientes son algunos aprendizajes técnicos importantes del ejercicio que mejoraron enormemente el rendimiento de nuestro modelo:

• Colocación de xm.mark_step() – xm.mark_step() compila y ejecuta los gráficos de cálculo recopilados de forma perezosa. Invocando mark_step demasiadas veces conducirá a una mayor cantidad de gráficos pequeños, mientras que invocarlo muy pocas veces conducirá a pocos gráficos, pero grandes. Dependiendo de su aplicación, el rendimiento y la implementación de su modelo de capacitación variarán según la ubicación de su xm.mark_step(). Nuestra implementación coloca uno xm.mark_step() después de un pase hacia adelante y hacia atrás, y uno después del paso del optimizador.
• Embalaje del cargador de datos con el cargador de dispositivos multiprocesamiento XLA – Este es un paso crítico que puede pasarse por alto fácilmente. El cargador de dispositivos multiprocesamiento torch_xla.distributed.parallel_loader.MpDeviceLoader carga datos de entrenamiento en cada dispositivo XLA con opciones para precargar y superponer la carga de datos con las ejecuciones del dispositivo para mejorar el rendimiento. El cargador de dispositivos también invoca xm.mark_step() y, por lo tanto, es capaz de crear gráficos para la carga de datos al dispositivo desde el host.

Compilación para Trainium

Tradicionalmente, el ciclo de desarrollo de modelos con GPU implica realizar cambios en el modelo o script de entrenamiento y ejecutarlo directamente en el dispositivo GPU. Los aceleradores como Trainium que utilizan XLA requieren un paso adicional antes de poder ejecutar el entrenamiento del modelo en el acelerador. Los gráficos de cálculo XLA sólo se pueden ejecutar después de haber sido compilados. Generalmente, hay dos formas de realizar esta compilación: Antes de tiempo (AOT), donde primero rastrea y compila todos los gráficos y luego los ejecuta, o Justo a tiempo (JIT), donde los gráficos se rastrean, compilan y ejecutan a medida que aparecen. se encuentran. El SDK de Neuron proporciona ambos listos para usar. Normalmente, la compilación AOT se realiza primero. Luego se ejecutan los gráficos después de esta compilación. Si se encuentran nuevos gráficos, el tiempo de ejecución de Neuron invoca una compilación JIT antes de ejecutarlos. Para realizar la compilación AOT, Neuron SDK proporciona neuron_parallel_compile, una utilidad de compilación que extrae gráficos de una ejecución de prueba del script de entrenamiento y realiza una compilación AOT paralela.

Un aspecto importante de la compilación de AOT es garantizar que no se creen nuevos gráficos de cálculo durante el transcurso de la capacitación. Una fuente de nuevos gráficos de cálculo (y, por lo tanto, de recompilaciones) son las formas dinámicas de los lotes de entrenamiento durante el entrenamiento del modelo. Descubrimos que el uso de formas estáticas y lotes de tamaño fijo elimina las compilaciones de tiempo de entrenamiento y mejora en gran medida el rendimiento del entrenamiento sin ningún efecto en la precisión del modelo. Al imponer tales restricciones en el entrenamiento, observamos que solo se requieren de 4 a 5 pasos de entrenamiento del modelo, un paso de validación del modelo y puntos de control del modelo una vez para rastrear todos los gráficos durante la compilación de AOT. Es importante tener en cuenta que Neuron SDK está en constante evolución y, en el futuro, también admitirá formas dinámicas.

Además, los gráficos compilados se almacenan en el Caché persistente de neuronas en disco o en un Servicio de almacenamiento simple de Amazon (Amazon S3) cubo. Esto es especialmente útil para cargas de trabajo de producción donde la arquitectura del modelo y la configuración de entrenamiento no cambian. Por lo tanto, la sobrecarga de compilación se produce sólo una vez. Usar el caché es tan simple como configurar un indicador de entorno:

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

El compilador Neuron también proporciona tres opciones de optimización a nivel de compilador (O1, O2, O3) para equilibrar el tiempo de compilación y el rendimiento de ejecución del modelo. O1 permite optimizaciones centrales en el gráfico de cálculo y minimiza el tiempo de compilación, O3 proporciona un rendimiento de ejecución del modelo mejorado a costa de un mayor tiempo de compilación y O2 (opción predeterminada) es un equilibrio entre los dos. Para nuestro caso de uso, utilizamos la optimización O1 y observamos una reducción del 86 % en el tiempo de compilación sin cambios en las métricas de precisión del modelo, mientras observamos aproximadamente una reducción del 5 al 7 % en el rendimiento en comparación con la optimización predeterminada (O2). Dependiendo del caso de uso, puede elegir diferentes niveles de optimización.

Para resumir, utilizamos las siguientes banderas para la compilación:

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

Compatibilidad de puntos de control

Cuando la compilación se complete con éxito, podemos proceder a entrenar nuestros modelos en Trainium. Como se mencionó anteriormente, entrenamos nuestros modelos de manera incremental, lo que significa que cargamos un punto de control del modelo previamente entrenado y continuamos entrenando con nuevos datos. PyTorch y PyTorch XLA permiten una transición perfecta entre aceleradores a través de la interoperabilidad de puntos de control. Tener la flexibilidad de movernos entre GPU y Trainium nos permitió cargar sin problemas el modelo de GPU anterior y entrenar en máquinas Trainium. Esto fue fundamental para garantizar que pudiéramos inicializar nuestro modelo con el mejor modelo previamente entrenado sin ningún tiempo de inactividad en la producción ni pérdida de precisión del modelo.

Debido a que el modelo de GPU se guardó usando las utilidades estándar de guardado de modelos de PyTorch, pudimos usar la utilidad de carga del punto de control de PyTorch para cargar el modelo de GPU en dispositivos Trainium.

Por ejemplo, en GPU/CPU, puede guardar el modelo con el siguiente código:

torch.save(model.state_dict(), PATH)

Luego vuelves a cargar el modelo en Trainium:

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

De manera similar, puedes guardar el modelo en Trainium con el siguiente código:

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

Y cargue el modelo nuevamente en GPU/CPU:

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

De hecho, debido a que usamos DDP para el entrenamiento del modelo, la carga del modelo es independiente de la cantidad de máquinas utilizadas para entrenar el punto de control anterior. Esto nos permite escalar horizontalmente la flota Trn1 sin cambios de código ni efectos adversos en el entrenamiento del modelo. Estos puntos de control basados ​​en PyTorch se pueden usar directamente o incluso programar mediante scripts para casos de uso de inferencia en AWS Inferentia2 u otros aceleradores.

Estabilidad operativa

No se puede enfatizar lo suficiente que ejecutar cargas de trabajo en producción requiere que se cumplan múltiples SLA. Para nuestro caso de uso, además de los SLA de calidad del modelo y rendimiento de la capacitación, es imperativo que el proceso de producción sea operativamente estable, lo que significa un tiempo de inactividad e interrupciones mínimos durante la capacitación, evaluación e inferencia del modelo.

Al igual que con el canal existente basado en GPU, agregamos numerosos mecanismos para hacer que el canal sea operativamente estable. Antes de comenzar el entrenamiento del modelo, realizamos múltiples pruebas de cordura para evaluar el estado de las máquinas. Estas pruebas generalmente incluyen operaciones tensoriales simples para verificar el estado de los dispositivos aceleradores. Hemos observado que para la capacitación distribuida, también es importante ejecutar pruebas para verificar la comunicación colectiva entre instancias. Usamos el Conjunto de pruebas NCCOM del SDK de Neuron para lograr esto, ejecutando una variedad de operaciones como recopilación total, reducción total y reducción-dispersión.

Incluso después de seguir las sugerencias que hemos mencionado, hemos observado que los problemas transitorios son inevitables en cualquier tubería, independientemente del acelerador subyacente. Para generar resiliencia en cualquier proceso de capacitación, recomendamos incorporar mecanismos de reintento para resolver estos posibles problemas. Usamos Reintentos automatizados de AWS Batch para reintentar trabajos que encuentran un error transitorio durante el entrenamiento del modelo. Estos reinicios pueden resultar costosos si se produce una falla hacia el final del entrenamiento. Para contrarrestar este problema, hemos adaptado nuestros scripts de entrenamiento para cargar un punto de control del modelo previamente entrenado y continuar entrenando desde ese punto. Con esta funcionalidad, podemos reiniciar agresivamente trabajos de capacitación fallidos con una sobrecarga mínima.

Con estos mecanismos de resiliencia implementados, pudimos lograr tasas de éxito del 98.5 % para nuestras cargas de trabajo en Trn1, comparables a las tasas de éxito de nuestra canalización de GPU existente.

Resultados

Para validar la precisión de nuestros modelos, inicializamos dos modelos desde el mismo punto de control de GPU y entrenamos uno en Trainium y el otro en una GPU comparable. Ambos modelos fueron entrenados con los mismos hiperparámetros de entrenamiento. El conjunto de datos utilizado para el cálculo de métricas es un conjunto de datos reservado y evaluamos la precisión del modelo en este conjunto de datos cada N pasos globales. El eje X es el paso global y el eje Y es la precisión del modelo. Observamos menos del 0.1 % de diferencia en la precisión del modelo en cada punto del siguiente gráfico.

Además, para evaluar la rentabilidad del entrenamiento del modelo, preferimos comparar el tiempo necesario para alcanzar la convergencia del modelo. Creemos que esto proporciona una visión más práctica del ahorro de costos en comparación con medidas como el costo por token, los FLOPS/dólar logrados y otros factores. Considerando el tiempo de entrenamiento de trn1.32xl y comparable Nube informática elástica de Amazon (Amazon EC2), hemos observado que Trainium ofrece un costo hasta un 30 % más económico para modelar la convergencia.

Conclusión

Hay muchos factores a considerar al evaluar diferentes aceleradores para sus cargas de trabajo de DL. Algunos de los más importantes son la calidad, el rendimiento, el costo y la disponibilidad del modelo. Es fundamental asegurarse de que la calidad y el rendimiento de su modelo no se sacrifiquen en función del acelerador que elija.

Gracias a nuestra asociación y colaboración con el equipo de Annapurna Neuron, el equipo de Amazon Search M5 ha podido ahorrar hasta un 30 % en costos al migrar a Trainium. El equipo puede utilizar Trainium y lograr la calidad del modelo y la paridad de rendimiento con aceleradores comparables en el mercado. La interoperabilidad de los puntos de control y los cambios mínimos de código con soporte para XLA han permitido a M5 elegir entre múltiples aceleradores para sus cargas de trabajo. Esto ha permitido al equipo de M5 aprovechar la gran potencia informática de Trainium y crear soluciones independientes del acelerador para deleitar a los clientes de Amazon.com. Desde un punto de vista operativo, se ha demostrado que Trainium es capaz de respaldar servicios de nivel 1 a escala de Amazon. El equipo de M5 continúa trasladando más cargas de trabajo a Trainium para proporcionar los mejores modelos para Amazon al menor costo.

En resumen, el equipo de M5 ha podido realizar una formación de aprendizaje automático rentable y de nivel de producción añadiendo Trainium a la flota de aceleradores. Le recomendamos que eche un vistazo a Trainium y otros dispositivos Neuron como AWS Inferentia para aprovechar los beneficios del silicio de Amazon diseñado específicamente para cargas de trabajo de aprendizaje automático. Comience fácilmente con uno de los muchos tutoriales que presentan diferentes modelos, como Llama 2, disponible en Trainium.

Acerca de los autores

Abhinandan Patni es ingeniero de software sénior en Amazon Search. Se centra en la creación de sistemas y herramientas para la formación de aprendizaje profundo distribuido escalable y la inferencia en tiempo real.

James Park es arquitecto de soluciones en Amazon Web Services. Trabaja con Amazon.com para diseñar, construir e implementar soluciones tecnológicas en AWS y tiene un interés particular en la inteligencia artificial y el aprendizaje automático. En su tiempo libre le gusta buscar nuevas culturas, nuevas experiencias y mantenerse al día con las últimas tendencias tecnológicas. Puedes encontrarlo en Etiqueta LinkedIn.

Jerry Manil es ingeniero de software en Amazon Search. Trabaja para mejorar la eficiencia, solidez y escalabilidad de la infraestructura de formación distribuida.

Ken Su es ingeniero de software en Amazon Search. Trabaja para mejorar la eficiencia de la capacitación y el flujo de trabajo de capacitación distribuido escalable. Fuera del trabajo, le gusta el senderismo y el tenis.

RJ Es Ingeniero dentro de Amazon. Construye y optimiza sistemas para sistemas distribuidos para capacitación y trabaja en la optimización de la adopción de sistemas para reducir la latencia de la inferencia de ML. Fuera del trabajo, está explorando el uso de IA generativa para crear recetas de comida.

• Distribución de relaciones públicas y contenido potenciado por SEO. Consiga amplificado hoy.
• PlatoData.Network Vertical Generativo Ai. Empodérate. Accede Aquí.
• PlatoAiStream. Inteligencia Web3. Conocimiento amplificado. Accede Aquí.
• PlatoESG. Carbón, tecnología limpia, Energía, Ambiente, Solar, Gestión de residuos. Accede Aquí.
• PlatoSalud. Inteligencia en Biotecnología y Ensayos Clínicos. Accede Aquí.
• Fuente: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/