Cómo ajustar LLaMA, OpenLLaMA y XGen con JAX en una GPU o TPU

LLaMA, OpenLLaMA y XGen son modelos de IA generativa de última generación. Estos modelos ofrecen resultados aún mejores cuando se ajustan con tus propios datos. En este artículo veremos cómo afinar estos modelos tanto en una GPU como en una TPU, utilizando JAX y la librería EasyLM.

LLaMA, OpenLLaM y XGen

El modelo LLaMA fue lanzado por Meta en febrero de 2023. Este modelo generativo de IA es un modelo de código abierto propuesto en varios tamaños: 7B parámetros, 13B parámetros, 33B parámetros y 65B parámetros.

En junio de 2021, cuando se publicó GPT-J, el mundo empezó a darse cuenta de que los modelos de IA generativa de código abierto podían competir seriamente con OpenAI GPT-3. Ahora, con LLaMA, el listón ha vuelto a subir claramente y este modelo parece ser una muy buena alternativa de código abierto a OpenAI ChatGPT y GPT-4.

Sin embargo, la licencia de LLaMA no es favorable a las empresas: este modelo no puede utilizarse con fines comerciales... Pero la buena noticia es que ahora existen otros modelos.

OpenLLaMA, publicado en junio de 2023, es una versión alternativa de LLaMA, desarrollada por el equipo Berkeley AI Research, que da muy buenos resultados y puede utilizarse para los negocios. Hay 2 versiones disponibles en el momento de escribir este artículo: 7B parámetros y 13B parámetros.

XGen, lanzado por Salesforce en junio de 2023, es otro modelo fundacional muy potente que puede utilizarse en aplicaciones comerciales. En el momento de redactar este documento sólo está disponible una versión con parámetros 7B. Cabe destacar que este modelo admite un contexto de 8k tokens, mientras que LLaMA y OpenLLaMA solo admiten un contenido de 2k tokens.

¿Por qué afinar su propio modelo?

Los modelos anteriores son modelos fundacionales, lo que significa que se han entrenado de forma no supervisada en un gran corpus de textos.

Estos modelos fundacionales de IA suelen ser una buena base, pero hay que afinarlos para que entiendan bien lo que quieres y devuelvan buenos resultados. La forma más sencilla de conseguirlo es utilizar el aprendizaje de pocos disparos (también conocido como "ingeniería de avisos"). No dude en leer aquí nuestra guía de aprendizaje dedicada a los pocos disparos.

El aprendizaje en pocos pasos es cómodo, ya que puede realizarse sobre la marcha sin tener que crear una nueva versión del modelo generativo de IA, pero a veces no es suficiente.

Para obtener los resultados más avanzados, querrás ajustar un modelo de IA a tu propio caso de uso. Ajustar significa modificar algunos parámetros del modelo a partir de tus propios datos y obtener así tu propia versión del modelo.

El ajuste fino es mucho más barato que entrenar un modelo generativo de IA desde cero, pero sigue requiriendo potencia de cálculo, por lo que se necesita un hardware avanzado para poder ajustar el modelo. Algunas técnicas alternativas recientes de ajuste fino requieren menos potencia de cálculo (véase p-tuning, prompt tuning, soft tuning, ajuste fino eficiente de parámetros, adaptadores, LoRA, QLoRA...), pero hasta ahora no hemos conseguido obtener el mismo nivel de calidad con estas técnicas, por lo que no vamos a mencionarlas en este tutorial.

Ajuste de LLaMA en una TPU con JAX y EasyLM

En este tutorial nos centraremos en la puesta a punto de LLaMA con la biblioteca EasyLM, publicada por el equipo de Berkeley AI Research: https://github.com/young-geng/EasyLM. This library is based on JAX which makes the fine-tuning process fast and compatible with both GPUs and Google TPUs.

También puede ajustar OpenLLaMA o XGen utilizando la misma técnica.

Aquí afinamos LLaMA 7B en una Google TPU V3-8, pero puedes hacer perfectamente lo mismo en una GPU A100 (simplemente lee cuidadosamente la parte de "Instalación" en la documentación de EasyLM que es ligeramente diferente). Por supuesto, también puedes afinar versiones mayores de LLaMA (13B, 33B, 65B...) pero necesitarás mucho más que una TPU V3-8 o una única GPU A100.

¡Allá vamos!

En primer lugar, cree un conjunto de datos de generación de texto para su caso de uso, en formato JSONL, utilizando "texto" como clave para cada ejemplo. Aquí tienes un sencillo conjunto de datos de análisis de sentimiento:

{"text":"[Content]: I love NLP Cloud, this company is awesome!\n[Sentiment]: Positive"}
{"text":"[Content]: Training LLMs is a complex but rewarding process.\n[Sentiment]: Neutral"}
{"text":"[Content]: My fine-tuning keeps crashing because of an OOM error! It just does not work at all!\n[Sentiment]: Negative"}

Tenga en cuenta un par de cosas importantes. En primer lugar, este conjunto de datos sólo contiene 3 ejemplos para simplificar, pero en la vida real necesitarás muchos más ejemplos. 300 ejemplos suele ser un buen comienzo. En segundo lugar, cuando utilice su modelo ajustado para la inferencia, deberá seguir estrictamente el mismo formato, utilizando los prefijos "[Contenido]:" y "[Sentimiento]:". Por último, el símbolo "</s>" es importante porque significa que el modelo debe dejar de generarse aquí. Puede encontrar más ejemplos de conjuntos de datos en la documentación de NLP Cloud: más información aquí.

Cree una TPU V3-8 VM en Google Cloud con la versión de software V2 Alpha:

SSH en la máquina virtual e instalar EasyLM:

git clone https://github.com/young-geng/EasyLM
cd EasyLM
bash ./scripts/tpu_vm_setup.sh

Ahora puedes descargar y convertir los pesos LLaMA. La primera opción es pedir a Meta los pesos oficiales: https://ai.facebook.com/blog/large-language-model-llama-meta-ai/. A continuación, convierta los pesos a EasyLM con este script: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_torch_to_easylm.py. La segunda opción es utilizar los pesos LLaMA en HuggingFace: https://huggingface.co/decapoda-research/llama-7b-hf. A continuación, convierta los pesos a EasyLM con este script: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_hf_to_easylm.py.

Cargue su conjunto de datos en la VM, cuente cuántos tokens contiene, utilizando el tokenizador HF LLaMA:

pip install -U transformers
python -c "from transformers import LlamaTokenizer; tokenizer = LlamaTokenizer.from_pretrained('decapoda-research/llama-7b-hf'); f = open('/path/to/your/dataset', 'r'); print(len(tokenizer.encode(f.read())))"

Si entrena su modelo para un contexto de 1024 tokens, tendrá que dividir el número de tokens devuelto por 1024.

Si entrena su modelo para un contexto de 2048 tokens, tendrá que dividir el número devuelto de tokens por 2048.

Este número será el número de pasos por epoch. Así, por ejemplo, si desea entrenar durante 5 épocas (que suele ser un buen ajuste) tendrá que multiplicar este número por 5 y poner el valor resultante en --total_steps a continuación.

He aquí un ejemplo concreto: si su conjunto de datos contiene 100.000 tokens, y desea un contexto de 1024 tokens y 5 épocas, su número total de pasos será (100.000/1024)*5 = 488.

Dependiendo de la longitud de su contexto, establezca --train_dataset.json_dataset.seq_length como 1024 o 2048 a continuación. Tenga en cuenta que el ajuste fino de un modelo para un contexto de 2048 tokens requiere más memoria, por lo que si esto no es estrictamente necesario le recomendamos que se limite a un contexto de 1024 tokens.

Ahora puede iniciar el proceso de ajuste:

nohup python -u EasyLM/EasyLM/models/llama/llama_train.py \
--total_steps=your number of steps \
--save_model_freq=usually same as your number of steps \
--optimizer.adamw_optimizer.lr_warmup_steps=usually 10% of total steps \
--train_dataset.json_dataset.path='/path/to/your/dataset' \
--train_dataset.json_dataset.seq_length=1024 or 2048 \
--load_checkpoint='params::/path/to/converted/model' \
--tokenizer.vocab_file='/path/to/tokenizer' \
--logger.output_dir=/path/to/output  \
--mesh_dim='1,4,2' \
--load_llama_config='7b' \
--train_dataset.type='json' \
--train_dataset.text_processor.fields='text' \
--optimizer.type='adamw' \
--optimizer.accumulate_gradient_steps=1 \
--optimizer.adamw_optimizer.lr=0.002 \
--optimizer.adamw_optimizer.end_lr=0.002 \
--optimizer.adamw_optimizer.lr_decay_steps=100000000 \
--optimizer.adamw_optimizer.weight_decay=0.001 \
--optimizer.adamw_optimizer.multiply_by_parameter_scale=True \
--optimizer.adamw_optimizer.bf16_momentum=True &

Algunas explicaciones:

--save_model_freq: la frecuencia con la que desea guardar el modelo durante el proceso. Si sólo realiza un ajuste fino en un conjunto de datos pequeño, puede guardar sólo al final del proceso, y en ese caso este valor será igual a --total_steps.

--optimizer.adamw_optimizer.lr_warmup_steps: El 10% del total de pasos suele ser un buen valor.

--tokenizer.vocab_file: la ruta al archivo tokenizer.model. Por ejemplo, si utiliza el repositorio decapoda en HuggingFace, este es el enlace al tokenizador: https://huggingface.co/decapoda-research/llama-7b-hf/resolve/main/tokenizer.model

--logger.output_dir: ruta al modelo final y a los registros

Los demás parámetros pueden dejarse intactos.

Una vez finalizado el proceso de ajuste, puede recuperar su modelo en la ruta especificada en --logger.output_dir.

Utilización del modelo ajustado para la inferencia

Ya tiene su propio modelo afinado y, por supuesto, quiere utilizarlo.

Una primera estrategia consiste en utilizar la biblioteca EasyLM para la inferencia. En ese caso, puede iniciar el servidor de inferencia de la siguiente manera:

python EasyLM/EasyLM/models/llama/llama_serve.py  \
--mesh_dim='1,1,-1' \
--load_llama_config='7b' \
--load_checkpoint='params::/path/to/your/model' \
--tokenizer.vocab_file='/path/to/tokenizer'

A continuación, simplemente envíe sus peticiones con cURL de la siguiente manera:

curl "http://0.0.0.0:5007/generate" \
-H "Content-Type: application/json" \
-X POST -d '{"prefix_text":["[Content]: EasyLM works really well!\n[Sentiment]:"]}'

Una segunda estrategia consiste en exportar su modelo al formato HuggingFace para realizar inferencias con otro marco. A continuación te explicamos cómo puedes exportarlo:

python EasyLM/EasyLM/models/llama/convert_easylm_to_hf.py \
--load_checkpoint='params::/path/to/output/model/streaming_params' \
--tokenizer_path='/path/to/tokenizer' \
--model_size='7b' \
--output_dir='/path/to/converted/model'

Conclusión

2023 ha sido un gran hito para los modelos de IA generativa de código abierto. A partir de este momento, todo el mundo puede utilizar grandes modelos como LLaMA, OpenLLaMA, XGen, Orca, Falcon...

Afinar estos modelos es la mejor manera de obtener resultados de vanguardia, adaptados a su propio caso de uso, que pueden superar significativamente a los mejores modelos de IA patentados como ChatGPT (GPT-3.5), GPT-4, Claude...

En esta guía te mostré cómo poner a punto LLaMA, OpenLLaMA y XGen. Si tienes preguntas no dudes en ponerte en contacto conmigo, y si quieres poner a punto y desplegar fácilmente modelos avanzados de IA generativa sin ninguna complejidad técnica, have a look at the NLP Cloud dedicated documentation!

Mark
Ingeniero de aprendizaje automático en NLP Cloud