[ad_1]
Los modelos de lenguajes grandes (LLM) se han convertido en un tema de conversación diaria. Su rápida adopción queda demostrada por el tiempo que lleva llegar a 100 millones de usuarios, que ha caído de «4,5 años en Facebook» a un mínimo histórico de sólo «2 meses en ChatGPT». Un transformador generativo preentrenado (GPT) utiliza actualizaciones autorregresivas causales para hacer predicciones. Se ha demostrado que estas arquitecturas modelo ofrecen un rendimiento sorprendente para una variedad de tareas como el reconocimiento de voz, la generación de texto y la respuesta a preguntas. Varios modelos más nuevos, como NeoX, Falcon y Llama, utilizan la arquitectura GPT como columna vertebral. La formación de LLM requiere enormes cantidades de tiempo de computación y cuesta millones de dólares. En esta publicación, resumimos el proceso de capacitación de GPT NeoX en AWS Trainium, un acelerador de aprendizaje automático (ML) especialmente diseñado y optimizado para la capacitación en aprendizaje profundo. Explicamos cómo entrenamos dichos modelos de manera rentable (3,2 millones de tokens por dólar) utilizando AWS Trainium sin sacrificar la calidad del modelo.
Descripción general de la solución
Modelos GPT NeoX y Pythia
GPT NeoX y Pythia son modelos de lenguaje causal de código abierto de Eleuther-AI con aproximadamente 20 mil millones de parámetros en NeoX y 6,9 mil millones en Pythia. Ambos son modelos de decodificadores con un diseño arquitectónico similar al Chat GPT3. Sin embargo, también cuentan con varios añadidos que también están muy extendidos en los modelos más nuevos como el Llama. En particular, presentan una incrustación de posición rotacional (ROPE) con rotación parcial a lo largo de las dimensiones de la cabeza. Los modelos originales (NeoX y Pythia 6.9B) se entrenan en un conjunto de datos de pila disponible abiertamente con deduplicación y utilizando el backend de Megatron y Deepspeed.
Demostramos el entrenamiento previo y el ajuste de estos modelos en instancias Trn1 basadas en AWS Trainium utilizando la biblioteca Neuron NeMo. Para garantizar una prueba de concepto y una reproducción rápida, utilizamos un subconjunto más pequeño del conjunto de datos de Wikipedia tokenizado mediante el tokenizador de codificación de par de bytes (BPE) GPT2.
solución completa
Descargue el conjunto de datos de Wikipedia previamente tokenizado como se muestra:
Tanto NeoX 20B como Pythia 6.9B usan CUERDA con rotación parcial, por ejemplo, se gira el 25% de las dimensiones del cabezal y el resto se deja sin girar. Para implementar de manera eficiente la rotación parcial en el acelerador AWS Trainium, en lugar de concatenar las dimensiones giratorias y no giratorias, agregamos frecuencias cero para las dimensiones no giratorias y luego rotamos todo el conjunto de dimensiones del cabezal. Este sencillo truco nos ayudó a mejorar el rendimiento (secuencias procesadas por segundo) en AWS Trainium.
Pasos de entrenamiento
Para realizar la capacitación, utilizamos el clúster Trn1 de múltiples nodos de Amazon Elastic Compute Cloud (Amazon EC2) administrado por SLURM, donde cada nodo contiene una instancia trn1.32xl. Cada trn1.32xl tiene 16 aceleradores con dos trabajadores por acelerador. Después de descargar el último paquete de Neuron NeMo, utilice los scripts de ajuste y preentrenamiento de Neox y Pythia proporcionados con hiperparámetros optimizados y ejecute lo siguiente para el entrenamiento de cuatro nodos.
- Compilar: precompilar el modelo con tres iteraciones de tren para generar y guardar los gráficos:
- Ejecutar: ejecuta el entrenamiento cargando los diagramas en caché de los primeros pasos.
- Monitorear los resultados
Para operar el modelo Pythia 6.9B con repuesto se deben seguir los mismos pasos neox_20B_slurm.sh
de pythia_6.9B_slurm.sh
.
Experimentos de preentrenamiento y ajuste fino.
Demostramos el entrenamiento previo de los modelos GPT-NeoX y Pythia en AWS Trainium utilizando la biblioteca Neuron NeMo durante 10 000 iteraciones y también demostramos el ajuste fino de estos modelos durante 1000 pasos. Para el entrenamiento previo, utilizamos el tokenizador GPT2-BPE en NeMo y seguimos la misma configuración que en el modelo original. El ajuste fino en AWS Trainium requiere cambiar algunos parámetros (p. ej. Factor de división del tamaño del vocabulario), que se proporcionan en los scripts de ajuste para tener en cuenta las diferencias entre Megatron y NeMo, así como los cambios entre GPU y AWS Trainium. El rendimiento del entrenamiento distribuido con múltiples nodos y diferente número de nodos se muestra en la Tabla 1.
Modelo | Tensor paralelo | Tubería paralela | Número de instancias | Costo ($/hora) | Longitud de la secuencia | Tamaño de lote global | Rendimiento (seq/s) | Relación costo-rendimiento (tokens/$) |
Pitia 6.9B | octavo | 1 | 1 | 7,59 | 2048 | 256 | 10.4 | 10.102.387 |
octavo | 1 | 4 | 30.36 | 2048 | 256 | 35,8 | 8.693.881 | |
NeoX 20B | octavo | 4 | 4 | 30.36 | 2048 | 16384 | 13.60 | 3.302.704 |
octavo | 4 | octavo | 60,72 | 2048 | 16384 | 26,80 | 3.254.134 | |
octavo | 4 | dieciséis | 121,44 | 2048 | 16384 | 54,30 | 3.296.632 | |
octavo | 4 | 32 | 242,88 | 2048 | 16384 | 107,50 | 3.263.241 | |
octavo | 4 | 64 | 485.76 | 2048 | 16384 | 212.00 | 3.217.708 |
Tabla 1. Comparación del rendimiento medio de los modelos GPT NeoX y Pythia para entrenar hasta 500 pasos con un número variable de nodos. El precio de trn1.32xl se basa en la tarifa por hora de reserva de 3 años.
A continuación, también evaluamos el historial de pérdidas del entrenamiento del modelo en AWS Trainium y lo comparamos con la ejecución correspondiente en un clúster P4d (núcleos de GPU Nvidia A100). Además de la pérdida de entrenamiento, también comparamos indicadores útiles como la norma de gradiente, que es la norma 2 de los gradientes del modelo calculada en cada iteración de entrenamiento para monitorear el progreso del entrenamiento. Los resultados del entrenamiento se muestran en las Figuras 1, 2 y el ajuste fino de NeoX 20B se muestra en la Figura 3.
Ilustración 1. Pérdida de capacitación promedio entre todos los trabajadores (izquierda) y norma de gradiente (derecha) en cada paso de capacitación. NeoX 20B se entrena en 4 nodos con un pequeño conjunto de datos Wiki en GPU y Trainium con los mismos hiperparámetros de entrenamiento (tamaño de lote global = 256). La GPU utiliza BF16 y precisión mixta estándar, mientras que AWS Trainium utiliza BF16 completo con redondeo estocástico. Las curvas de norma de pérdida y gradiente coinciden para GPU y AWS Trainium.
Figura 2. Pérdida de capacitación promedio entre todos los trabajadores (izquierda) y norma de gradiente (derecha) en cada paso de capacitación. De manera similar a GPT NeoX en la Figura 1, Pythia 6.9B se entrena en 4 nodos con un pequeño conjunto de datos Wiki en GPU y Trainium con los mismos hiperparámetros de entrenamiento (tamaño de lote global = 256). Las curvas de pérdida y de gradiente coinciden para GPU y Trainium.
Figura 3. Ajuste del modelo GPT NeoX 20B en GPU y AWS Trainium con pérdida de entrenamiento promediada entre todos los trabajadores (izquierda) y norma de gradiente (derecha). Se utiliza un pequeño conjunto de datos Wiki para la demostración de ajuste. Las curvas de norma de pérdida y gradiente coinciden para GPU y AWS Trainium.
En esta publicación, demostramos la capacitación rentable de LLM en hardware de aprendizaje profundo de AWS. Entrenamos los modelos GPT NeoX 20B y Pythia 6.9B en AWS Trn1 utilizando la biblioteca Neuron NeMo. El rendimiento normalizado de costos para 20 mil millones de modelos que utilizan AWS Trainium es de aproximadamente 3,2 millones de tokens por cada dólar gastado. Junto con la capacitación rentable en AWS Trainium, obtenemos una precisión del modelo similar como resultado de la pérdida de pasos de capacitación y la trayectoria de la norma de gradiente. También hemos optimizado los puntos de control disponibles para el modelo NeoX 20B en AWS Trainium. Para obtener más información sobre la capacitación distribuida con NeMo Megatron en AWS Trainium, consulte la referencia de AWS Neuron para NeMo Megatron. Un buen recurso para comenzar a ajustar el modelo Llama se puede encontrar aquí: Llama2 Fine-Tuning. Para comenzar con AWS Trainium administrado en Amazon SageMaker, consulte Entrene sus modelos de aprendizaje automático con AWS Trainium y Amazon SageMaker.
Sobre los autores
Gaurav Gupta Actualmente es científico aplicado en los laboratorios de inteligencia artificial de Amazon Web Services (AWS). Dr. Gupta completó su doctorado en la USC Viterbi. Sus intereses de investigación incluyen modelado de datos secuenciales, aprendizaje de ecuaciones diferenciales parciales, teoría de la información para el aprendizaje automático, modelos dinámicos fraccionarios y redes complejas. Actualmente trabaja en problemas matemáticos y aplicados relacionados con el comportamiento de entrenamiento de LLM, modelos de visión con PDE y modelos multimodales de teoría de la información. Dr. Gupta ha publicado en revistas y conferencias líderes como Neurips, ICLR, ICML, Nature, IEEE Control Society y ACM Cyber-Physical Society.
Ben Snyder es un científico aplicado en AWS Deep Learning. Sus intereses de investigación incluyen modelos fundamentales, aprendizaje por refuerzo y optimización asincrónica. Fuera del trabajo, le gusta andar en bicicleta y acampar en el campo.
Amith (R) Mamidala es ingeniero senior de aplicaciones para aprendizaje automático en AWS Annapurna Labs. Dr. Mamidala obtuvo su doctorado en informática y comunicaciones de alto rendimiento en la Universidad Estatal de Ohio. Mientras trabajaba en IBM Research, el Dr. Mamidala contribuyó al desarrollo de la clase de computadoras BlueGene, que frecuentemente encabezaba el ranking de las 500 supercomputadoras más potentes y energéticamente eficientes. El proyecto recibió la Medalla Nacional de Tecnología e Innovación en 2009. Después de un breve período como ingeniero de inteligencia artificial en un fondo de cobertura financiero, el Dr. Mamidala se mudó a los laboratorios de Annapurna y se centró en entrenar modelos de lenguaje de gran tamaño.
Jun (Lucas) Huan es un científico senior en AWS AI Labs. Dr. Huan está involucrado en inteligencia artificial y ciencia de datos. Ha publicado más de 180 artículos revisados por pares en conferencias y revistas líderes. Recibió el Premio al Desarrollo Profesional Temprano de la Facultad de la NSF en 2009. Antes de unirse a AWS, trabajó en Baidu Research como científico distinguido y director del Laboratorio de Big Data de Baidu. Fundó StylingAI Inc., una nueva empresa de inteligencia artificial, y se desempeñó como director ejecutivo y científico jefe de 2019 a 2021. Antes de ingresar a la industria, fue profesor Charles E. y Mary Jane Spahr en el Departamento EECS de la Universidad de Kansas.
Sruti Koparkar es gerente senior de marketing de productos en AWS. Ayuda a los clientes a explorar, evaluar y adoptar la infraestructura de computación acelerada de Amazon EC2 para sus necesidades de aprendizaje automático.
[ad_2]