En la formación de modelos de lenguaje grande (LLM), la orquestación eficaz y la gestión de recursos informáticos plantean un desafío importante. La automatización del aprovisionamiento, el escalado y la gestión del flujo de trabajo de recursos es vital para optimizar el uso de recursos y agilizar flujos de trabajo complejos, logrando así procesos eficientes de capacitación en aprendizaje profundo. La orquestación simplificada permite a los investigadores y profesionales centrarse más en la experimentación de modelos, el ajuste de hiperparámetros y el análisis de datos, en lugar de lidiar con tareas engorrosas de gestión de infraestructura. La orquestación directa también acelera la innovación, acorta el tiempo de comercialización de nuevos modelos y aplicaciones y, en última instancia, mejora la eficiencia y eficacia general de los esfuerzos de investigación y desarrollo de LLM.

Esta publicación explora la perfecta integración de Capacitación en AWS con Lote de AWS, que muestra cómo las poderosas capacidades de aceleración del aprendizaje automático (ML) de Trainium se pueden aprovechar junto con las eficientes funcionalidades de orquestación que ofrece AWS Batch. Trainium proporciona una escalabilidad masiva, permite escalar sin esfuerzo los trabajos de capacitación desde modelos pequeños hasta LLM y ofrece acceso rentable a potencia computacional, lo que hace que la capacitación de LLM sea asequible y accesible. AWS Batch es un servicio administrado que facilita cargas de trabajo de computación por lotes en la nube de AWS, manejando tareas como la administración de infraestructura y la programación de trabajos, mientras le permite concentrarse en el desarrollo de aplicaciones y el análisis de resultados. AWS Batch proporciona funciones integrales, que incluyen computación por lotes administrada, cargas de trabajo en contenedores, entornos informáticos personalizados y colas de trabajos priorizados, junto con una integración perfecta con otros servicios de AWS.

Descripción general de la solución

El siguiente diagrama ilustra la arquitectura de la solución.

El proceso de formación se desarrolla de la siguiente manera:

1. El usuario crea una imagen de Docker configurada para adaptarse a las demandas de la tarea de capacitación subyacente.
2. La imagen es empujada a Registro de contenedores elásticos de Amazon (Amazon ECR) para prepararlo para su implementación.
3. El usuario envía el trabajo de capacitación a AWS Batch con la imagen de Docker.

Profundicemos en esta solución para ver cómo puede integrar Trainium con AWS Batch. El siguiente ejemplo demuestra cómo entrenar el modelo Llama 2-7B utilizando AWS Batch con Trainium.

Requisitos previos

Se recomienda no ejecutar los siguientes scripts en su máquina local. En lugar de eso, clona el repositorio de GitHub y ejecute los scripts proporcionados en una instancia basada en x86_64, preferiblemente usando un tipo de instancia C5.xlarge con el sistema operativo Linux/Ubuntu. Para esta publicación, ejecutamos el ejemplo en una instancia de Amazon Linux 2023.

Debe tener los siguientes recursos y herramientas antes de comenzar con la capacitación en AWS Batch:

sudo yum install -y docker 
sudo yum install -y jq

Clonar el repositorio

Clona el repositorio de GitHub y navega hasta el directorio requerido:

git clone https://github.com/aws-neuron/aws-neuron-samples.git 
cd aws-neuron-samples/torch-neuronx/training/aws-batch/llama2

Actualizar la configuración

Primero, actualice el config.txt archivo para especificar valores para las siguientes variables:

REGION                          # your aws region 
SUBNET                          # your subnet in which the Trainium instances would be launched 
SG                              # your security group you want to associate with your instances 
ECR_REPO                        # your ECR repo where the docker container image will be pushed to 
INSTANCE_ROLE                   # Instance profile ARN for your IAM Instance Role 
DO_PRE_COMPILATION              # boolean value (true\false) indicating if you want to do neuron pre-compilation for your training job 
TOKENIZED_DATASET_URI           # s3 uri to store the tokenized dataset 
NEURON_COMPILE_CACHE_URI        # s3 uri to store the neuron compile caches 
CHECKPOINT_SAVE_URI             # s3 uri to store the checkpoints

Después de proporcionar estos valores, su archivo config.txt debería verse similar al siguiente código

REGION=us-east-1
SUBNET=subnet-012345abcd5689
SG=sg-012345abcd5689
ECR_REPO=1010101010.dkr.ecr.us-east-1.amazonaws.com/your-docker-repo
INSTANCE_ROLE=arn:aws:iam::1010101010:instance-profile/your-instance-role
DO_PRE_COMPILATION=true
TOKENIZED_DATASET_URI=s3://your/s3/location/to/store/tokenized/dataset/
NEURON_COMPILE_CACHE_URI=s3://your/s3/location/to/store/neuron-compile-cache/
CHECKPOINT_SAVE_URI=s3://your/s3/location/to/store/checkpoints/

Consigue el tokenizador Llama

Para tokenizar el conjunto de datos, deberá obtener el tokenizador de Hugging Face. Siga el instrucciones para acceder al tokenizador Llama. (Debe reconocer y aceptar los términos de la licencia). Una vez que se le haya concedido acceso, podrá descargar el tokenizador desde Hugging Face. Después de una descarga exitosa, coloque el tokenizer.model archivo en el directorio raíz (\llama2).

Configurar el entrenamiento de llamas

Ejecute el setup.sh script, que simplifica los pasos previos para iniciar la capacitación de AWS Batch. Este script descarga los archivos Python necesarios para entrenar el Llama 2-7B modelo. Además, realiza la sustitución de variables de entorno dentro del formato proporcionado. plantillas y guiones diseñado para establecer recursos de AWS Batch. Cuando se ejecuta, se asegura de que la estructura de su directorio se ajuste a la siguiente configuración:

.
├── build
│ ├── compute_env.json
│ ├── job_def.json
│ ├── job_queue.json
│ └── launch_template.json
├── build_and_push_docker_image.sh
├── cleanup.sh
├── config.txt
├── create_resources.sh
├── data
│ ├── get_dataset.py
│ ├── config.json
│ └── tokenizer.model
├── docker
│ ├── Dockerfile
│ ├── llama2
│ │ ├── adamw_fp32_optim_params.py
│ │ ├── config.json
│ │ ├── llama_batch_training.sh
│ │ ├── modeling_llama_nxd.py
│ │ ├── requirements.txt
│ │ └── tp_zero1_llama2_7b_hf_pretrain.py
│ └── llama_batch_training.sh
├── download_and_tokenize_data.sh
├── images
│ └── aws-batch.png
├── README.md
├── scripts
│ ├── build_and_push_docker_image.sh
│ ├── cleanup.sh
│ ├── create_resources.sh
│ ├── download_and_tokenize_data.sh
│ └── submit_batch_job.sh
├── setup.sh
├── submit_batch_job.sh
└── templates
├── compute_env.json
├── job_def.json
├── job_queue.json
└── launch_template.json

Tokenizar el conjunto de datos

A continuación, ejecute el download_and_tokenize_data.sh script para completar los pasos de preprocesamiento de datos para el entrenamiento Llama 2-7B. En este caso utilizamos el wikicorpus conjunto de datos procedente de Hugging Face. Después de la recuperación del conjunto de datos, el script realiza la tokenización y carga el conjunto de datos tokenizado en la ubicación S3 predefinida especificada dentro del config.txt archivo de configuración. Las siguientes capturas de pantalla muestran los resultados del preprocesamiento.

Aprovisionar recursos

A continuación, ejecute el create_resources.sh script, que organiza el aprovisionamiento de los recursos necesarios para la tarea de formación. Esto incluye la creación de un grupo de ubicación, una plantilla de lanzamiento, un entorno informático, una cola de trabajos y una definición de trabajos. Las siguientes capturas de pantalla ilustran este proceso.

Construya y envíe la imagen de Docker

Ahora puedes ejecutar el script. build_and_push_docker_image.sh, que construye una imagen de contenedor Docker personalizada para su tarea de capacitación específica. Este script utiliza un Imagen del contenedor de aprendizaje profundo publicado por el equipo de Neuron, que contiene la pila de software requerida, y luego agregó instrucciones para ejecutar el entrenamiento Llama 2-7B encima. El guión de entrenamiento utiliza el neuronx_distributed biblioteca con paralelismo tensorial junto con el Cero-1 Optimizador. Posteriormente, la imagen del contenedor Docker recién generada se carga en su repositorio ECR designado según lo especificado por la variable ECR_REPO en el archivo de configuración config.txt.

Si desea modificar cualquiera de los hiperparámetros de entrenamiento de Llama, realice los cambios necesarios en ./docker/llama_batch_training.sh antes de correr build_and_push_docker_image.sh.

Las siguientes capturas de pantalla ilustran el proceso para crear y enviar la imagen de Docker.

Enviar el trabajo de formación

Ejecute el submit_batch_job.sh script para iniciar el trabajo de AWS Batch e iniciar el entrenamiento del modelo Llama2, como se muestra en las siguientes capturas de pantalla.

Tras el envío del trabajo por lotes, un Servicio de contenedor elástico de Amazon (Amazon ECS) el clúster se aprovisiona dinámicamente. Cuando esté operativo, puede navegar hasta el clúster para monitorear todas las tareas que se ejecutan activamente en el Trn1.32xl instancias, lanzadas a través de este trabajo. De forma predeterminada, este ejemplo está configurado para utilizar 4 instancias trn1.32xl. Para personalizar esta configuración, puede modificar el numNodos parámetro en el submit_batch_job.sh guion.

Registros y seguimiento

Después del envío del trabajo, puede utilizar Registros de Amazon CloudWatch para monitoreo, almacenamiento y visualización integrales de todos los registros generados por AWS Batch. Complete los siguientes pasos para acceder a los registros:

1. En la consola de CloudWatch, elija Grupos de registro bajo Registros en el panel de navegación.
2. Elegir /aws/batch/job para ver los registros de trabajos por lotes.
3. Busque grupos de registros que coincidan con los nombres o definiciones de trabajos de AWS Batch.
4. Elija el trabajo para ver sus detalles.

La siguiente captura de pantalla muestra un ejemplo.

Puntos de control

Los puntos de control generados durante el entrenamiento se almacenarán en la ubicación S3 predefinida especificada como CHECKPOINT_SAVE_URI en el config.txt archivo. De forma predeterminada, el punto de control se guarda cuando se completa el entrenamiento. Sin embargo, puede ajustar este comportamiento optando por guardar el punto de control después de cada N pasos dentro del ciclo de entrenamiento. Para obtener instrucciones detalladas sobre esta personalización, consulte Puntos de control.

Limpiar

Cuando hayas terminado, ejecuta el cleanup.sh script para gestionar la eliminación de recursos creados durante la publicación. Este script se encarga de eliminar varios componentes, como la plantilla de inicio, el grupo de ubicación, la definición de trabajo, la cola de trabajos y el entorno informático. AWS Batch maneja automáticamente la limpieza de la pila ECS y las instancias de Trainium, por lo que no es necesario eliminarlas ni detenerlas manualmente.

Conclusión

La integración perfecta de Trainium con AWS Batch representa un avance significativo en el ámbito de la capacitación en ML. Al combinar las capacidades incomparables de Trainium con las potentes funcionalidades de orquestación de AWS Batch, podrá beneficiarse de numerosas maneras. En primer lugar, obtiene acceso a una escalabilidad masiva, con la capacidad de escalar sin esfuerzo trabajos de capacitación desde modelos pequeños hasta LLM. Con hasta 16 chips Trainium por instancia y el potencial de capacitación distribuida en decenas de miles de aceleradores, puede abordar incluso las tareas de capacitación más exigentes con facilidad gracias a las instancias Trainium. Además, ofrece una solución rentable que le ayuda a aprovechar la energía que necesita a un precio atractivo. Con el servicio totalmente administrado que ofrece AWS Batch para cargas de trabajo informáticas, puede descargar complejidades operativas como el aprovisionamiento de infraestructura y la programación de trabajos, lo que le permite centrar sus esfuerzos en crear aplicaciones y analizar resultados. En última instancia, la integración de Trainium con AWS Batch le permite acelerar la innovación, acortar el tiempo de comercialización de nuevos modelos y aplicaciones y mejorar la eficiencia y eficacia generales de sus esfuerzos de aprendizaje automático.

Ahora que ha aprendido a orquestar Trainium mediante AWS Batch, le animamos a que lo pruebe en su próximo trabajo de formación en aprendizaje profundo. Puedes explorar más tutoriales que lo ayudará a obtener experiencia práctica con AWS Batch y Trainium, y le permitirá administrar sus cargas de trabajo y recursos de capacitación de aprendizaje profundo para un mejor rendimiento y rentabilidad. Entonces, ¿por qué esperar? ¡Comience a explorar estos tutoriales hoy y lleve su capacitación en aprendizaje profundo al siguiente nivel con Trainium y AWS Batch!

Sobre los autores

Scott Perry es arquitecto de soluciones en el equipo del acelerador Annapurna ML en AWS. Con sede en Canadá, ayuda a los clientes a implementar y optimizar cargas de trabajo de inferencia y capacitación en aprendizaje profundo utilizando AWS Inferentia y AWS Trainium. Sus intereses incluyen grandes modelos de lenguaje, aprendizaje por refuerzo profundo, IoT y genómica.

Sadaf Rasool es ingeniero de aprendizaje automático en el equipo Annapurna ML Accelerator en AWS. Como profesional entusiasta y optimista de IA/ML, se mantiene firme en la creencia de que la aplicación ética y responsable de la IA tiene el potencial de mejorar la sociedad en los años venideros, fomentando tanto el crecimiento económico como el bienestar social.