Cómo entrenar un Modelo de Lenguaje con Megatron-LM

Entrenamiento de Modelo de Lenguaje con Megatron-LM

Entrenar modelos de lenguaje grandes en Pytorch requiere más que un simple bucle de entrenamiento. Por lo general, se distribuye en múltiples dispositivos, con muchas técnicas de optimización para un entrenamiento estable y eficiente. La biblioteca Hugging Face 🤗 Accelerate se creó para admitir entrenamiento distribuido en GPUs y TPUs con una integración muy fácil en los bucles de entrenamiento. 🤗 Transformers también admite entrenamiento distribuido a través de la API Trainer, que proporciona un entrenamiento completo en PyTorch, sin necesidad de implementar un bucle de entrenamiento.

Otra herramienta popular entre los investigadores para preentrenar grandes modelos de transformadores es Megatron-LM, un poderoso marco desarrollado por el equipo de Investigación en Aprendizaje Profundo Aplicado de NVIDIA. A diferencia de accelerate y Trainer, el uso de Megatron-LM no es sencillo y puede resultar abrumador para principiantes. Pero está altamente optimizado para el entrenamiento en GPUs y puede ofrecer algunas aceleraciones. En esta publicación de blog, aprenderás cómo entrenar un modelo de lenguaje en GPUs de NVIDIA en Megatron-LM y cómo usarlo con transformers.

Intentaremos desglosar los diferentes pasos para entrenar un modelo GPT2 en este marco, que incluyen:

• Configuración del entorno
• Preprocesamiento de datos
• Entrenamiento
• Conversión del modelo a 🤗 Transformers

¿Por qué Megatron-LM?

Antes de adentrarnos en los detalles del entrenamiento, primero entendamos qué hace que este marco sea más eficiente que otros. Esta sección está inspirada en este excelente blog sobre el entrenamiento BLOOM con Megatron-DeepSpeed, por favor, consulta el blog para obtener más detalles, ya que esta publicación tiene la intención de brindar una introducción básica a Megatron-LM.

• SetFit Aprendizaje Eficiente de Pocos Ejemplos Sin Indicaciones
• Cómo 🤗 Accelerate ejecuta modelos muy grandes gracias a PyTorch
• Clasificación de imágenes con AutoTrain

DataLoader

Megatron-LM cuenta con un DataLoader eficiente donde los datos se tokenizan y se mezclan antes del entrenamiento. También divide los datos en secuencias numeradas con índices que se almacenan de manera que solo necesitan calcularse una vez. Para construir el índice, se calcula el número de épocas basado en los parámetros de entrenamiento y se crea un ordenamiento que luego se mezcla. Esto es diferente a la mayoría de los casos en los que se itera a través de todo el conjunto de datos hasta que se agota y luego se repite para la segunda época. Esto suaviza la curva de aprendizaje y ahorra tiempo durante el entrenamiento.

Núcleos de CUDA fusionados

Cuando se ejecuta una computación en la GPU, se obtienen los datos necesarios de la memoria, luego se realiza la computación y el resultado se guarda nuevamente en la memoria. En términos sencillos, la idea de los núcleos fusionados es que las operaciones similares, que generalmente se realizan por separado en Pytorch, se combinan en una sola operación de hardware. Por lo tanto, se reduce la cantidad de movimientos de memoria realizados en múltiples computaciones discretas al fusionarlas en una sola. La figura a continuación ilustra la idea de la fusión de núcleos. Está inspirada en este artículo, que discute el concepto en detalle.

Cuando f, g y h se fusionan en un solo núcleo, los resultados intermedios x’ e y’ de f y g se almacenan en los registros de la GPU y se utilizan inmediatamente por h. Pero sin la fusión, x’ e y’ tendrían que copiarse en la memoria y luego cargarse por h. Por lo tanto, la fusión de núcleos brinda una aceleración significativa a las computaciones. Megatron-LM también utiliza una implementación fusionada de AdamW de Apex que es más rápida que la implementación de Pytorch.

Aunque uno puede personalizar el DataLoader como Megatron-LM y usar el optimizador Fused de Apex con transformers, no es una tarea sencilla para principiantes construir Núcleos de CUDA fusionados personalizados.

Ahora que estás familiarizado con el marco y lo que lo hace ventajoso, ¡entremos en los detalles del entrenamiento!

¿Cómo entrenar con Megatron-LM?

Configuración

La forma más sencilla de configurar el entorno es obtener un Contenedor de PyTorch de NVIDIA que incluye todas las instalaciones requeridas desde NGC. Consulta la documentación para obtener más detalles. Si no deseas utilizar este contenedor, deberás instalar las últimas versiones de PyTorch, CUDA, NCCL, NVIDIA APEX y la biblioteca nltk.

Entonces, después de haber instalado Docker, puedes ejecutar el contenedor con el siguiente comando ( xx.xx representa tu versión de Docker) y luego clonar el repositorio de Megatron-LM dentro de él:

docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:xx.xx-py3
git clone https://github.com/NVIDIA/Megatron-LM

También debes agregar el archivo de vocabulario vocab.json y la tabla de fusiones merges.txt de tu tokenizador dentro de la carpeta Megatron-LM de tu contenedor. Estos archivos se pueden encontrar en el repositorio del modelo junto con los pesos, consulta este repositorio para GPT2. También puedes entrenar tu propio tokenizador utilizando transformers. Puedes revisar el proyecto CodeParrot para un ejemplo práctico. Ahora, si deseas copiar estos datos desde fuera del contenedor, puedes utilizar los siguientes comandos:

sudo docker cp vocab.json ID_DEL_CONTENEDOR:/workspace/Megatron-LM
sudo docker cp merges.txt ID_DEL_CONTENEDOR:/workspace/Megatron-LM

Preprocesamiento de datos

En el resto de este tutorial, utilizaremos el modelo y los datos de CodeParrot como ejemplo.

Los datos de entrenamiento requieren algún preprocesamiento. Primero, debes convertirlos a un formato json suelto, con un json que contenga una muestra de texto por línea. Si estás utilizando 🤗 Datasets, aquí tienes un ejemplo de cómo hacerlo (siempre dentro de la carpeta Megatron-LM):

from datasets import load_dataset

train_data = load_dataset('codeparrot/codeparrot-clean-train', split='train')
train_data.to_json("codeparrot_data.json", lines=True)  

Luego, los datos se tokenizan, se mezclan y se procesan en un formato binario para el entrenamiento utilizando el siguiente comando:

#si nltk no está instalado
pip install nltk
python tools/preprocess_data.py \
       --input codeparrot_data.json \
       --output-prefix codeparrot \
       --vocab vocab.json \
       --dataset-impl mmap \
       --tokenizer-type GPT2BPETokenizer \
       --merge-file merges.txt \
       --json-keys content \
       --workers 32 \
       --chunk-size 25 \
       --append-eod

Las opciones workers y chunk_size se refieren al número de trabajadores utilizados en el preprocesamiento y al tamaño del fragmento de datos asignado a cada uno. dataset-impl se refiere al modo de implementación de los conjuntos de datos indexados [‘lazy’, ‘cached’, ‘mmap’]. Esto genera dos archivos codeparrot_content_document.idx y codeparrot_content_document.bin que se utilizan en el entrenamiento.

Entrenamiento

Puedes configurar la arquitectura del modelo y los parámetros de entrenamiento como se muestra a continuación, o ponerlo en un script de bash que ejecutarás. Este comando ejecuta el preentrenamiento en 8 GPUs para un modelo CodeParrot de 110M de parámetros. Ten en cuenta que los datos se dividen por defecto en una proporción de 969:30:1 para los conjuntos de entrenamiento/validación/prueba.

GPUS_PER_NODE=8
MASTER_ADDR=localhost
MASTER_PORT=6001
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))
DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"
CHECKPOINT_PATH=/workspace/Megatron-LM/experiments/codeparrot-small
VOCAB_FILE=vocab.json
MERGE_FILE=merges.txt
DATA_PATH=codeparrot_content_document
GPT_ARGS="--num-layers 12
--hidden-size 768
--num-attention-heads 12
--seq-length 1024
--max-position-embeddings 1024
--micro-batch-size 12
--global-batch-size 192
--lr 0.0005
--train-iters 150000
--lr-decay-iters 150000
--lr-decay-style cosine
--lr-warmup-iters 2000
--weight-decay .1
--adam-beta2 .999
--fp16
--log-interval 10
--save-interval 2000
--eval-interval 200
--eval-iters 10
"
TENSORBOARD_ARGS="--tensorboard-dir experiments/tensorboard"
python3 -m torch.distributed.launch $DISTRIBUTED_ARGS \
        pretrain_gpt.py \
        --tensor-model-parallel-size 1 \
        --pipeline-model-parallel-size 1 \
        $GPT_ARGS \
        --vocab-file $VOCAB_FILE \
        --merge-file $MERGE_FILE \
        --save $CHECKPOINT_PATH \
        --load $CHECKPOINT_PATH \
        --data-path $DATA_PATH \
        $TENSORBOARD_ARGS

Con esta configuración, el entrenamiento tarda aproximadamente 12 horas.

Esta configuración utiliza Paralelismo de Datos, pero también es posible utilizar Paralelismo de Modelo para modelos muy grandes que no caben en una GPU. La primera opción consiste en Paralelismo de Tensor que divide la ejecución de un solo módulo transformador en varias GPUs, deberá cambiar el parámetro tensor-model-parallel-size al número deseado de GPUs. La segunda opción es el Paralelismo de Tuberías donde los módulos transformadores se dividen en etapas de tamaño igual. El parámetro pipeline-model-parallel-size determina el número de etapas en las que se divide el modelo. Para más detalles, por favor consulte este blog

Convirtiendo el modelo a 🤗 Transformers

Después del entrenamiento, queremos usar el modelo en transformers, por ejemplo, para evaluación o para implementarlo en producción. Puede convertirlo a un modelo de transformers siguiendo este tutorial. Por ejemplo, después de que el entrenamiento haya finalizado, puede copiar los pesos de la última iteración 150k y convertir el archivo model_optim_rng.pt a un archivo pytorch_model.bin que es compatible con transformers con los siguientes comandos:

# para ejecutar fuera del contenedor:
mkdir -p nvidia/megatron-codeparrot-small
# copie los pesos del contenedor
sudo docker cp CONTAINER_ID:/workspace/Megatron-LM/experiments/codeparrot-small/iter_0150000/mp_rank_00/model_optim_rng.pt nvidia/megatron-codeparrot-small
git clone https://github.com/huggingface/transformers.git
git clone https://github.com/NVIDIA/Megatron-LM.git
export PYTHONPATH=Megatron-LM
python transformers/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py nvidia/megatron-codeparrot-small/model_optim_rng.pt

Tenga cuidado, deberá reemplazar el archivo de vocabulario generado y la tabla de fusiones después de la conversión, con los originales que presentamos anteriormente si planea cargar el tokenizador desde allí.

No olvide subir su modelo al hub y compartirlo con la comunidad, solo se necesitan tres líneas de código 🤗:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("nvidia/megatron-codeparrot-small")
# esto crea un repositorio bajo su nombre de usuario con el nombre del modelo codeparrot-small
model.push_to_hub("codeparrot-small")

También puede usarlo fácilmente para generar texto:

from transformers import pipeline

pipe = pipeline("text-generation", model="your_username/codeparrot-small")
outputs = pipe("def hello_world():")
print(outputs[0]["generated_text"])

 def hello_world():   
    print("¡Hola Mundo!")

Tranfsormers también maneja eficientemente la inferencia de modelos grandes. En caso de que haya entrenado un modelo muy grande (por ejemplo, utilizando Paralelismo de Modelo), puede usarlo fácilmente para la inferencia con el siguiente comando:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("your_username/codeparrot-large", device_map="auto")

Esto utilizará la biblioteca de aceleración en segundo plano para distribuir automáticamente los pesos del modelo en los dispositivos disponibles (GPUs, RAM de la CPU) que tenga.

Aviso legal: Hemos demostrado que cualquiera puede usar Megatron-LM para entrenar modelos de lenguaje. La pregunta es cuándo usarlo. Este marco obviamente agrega cierta sobrecarga de tiempo debido a los pasos adicionales de preprocesamiento y conversión. Por lo tanto, es importante que decida qué marco es más adecuado para su caso y tamaño de modelo. Recomendamos probarlo para modelos de pre-entrenamiento o ajuste fino extendido, pero probablemente no para ajuste fino más corto de modelos del tamaño de VoAGI. La API del Trainer y la biblioteca de accelerate también son muy útiles para el entrenamiento del modelo, son independientes del dispositivo y brindan una flexibilidad significativa a los usuarios.

¡Felicidades 🎉 ahora sabe cómo entrenar un modelo GPT2 en Megatron-LM y hacerlo compatible con transformers!

We will continue to update Zepes; if you have any questions or suggestions, please contact us!