Está a ter dificuldades com a IA ou com o desenvolvimento full-stack? Os nossos especialistas estão aqui para o orientar: aconselhamento personalizado, integração técnica e muito mais. Entre em contacto com [email protected].

Como ajustar LLaMA, OpenLLaMA e XGen, com JAX em uma GPU ou TPU

LLaMA, OpenLLaMA e XGen são modelos de IA generativos de ponta. Esses modelos fornecem resultados ainda melhores quando ajustados em seus próprios dados. Neste artigo, vamos ver como ajustar esses modelos em uma GPU e uma TPU, usando JAX e a biblioteca EasyLM.

LLaMA, OpenLLaM e XGen

O modelo LLaMA foi lançado pela Meta em fevereiro de 2023. Este modelo de IA generativa é um modelo de código aberto proposto em vários tamanhos: 7B parâmetros, 13B parâmetros, 33B parâmetros e 65B parâmetros.

Em junho de 2021, quando o GPT-J foi lançado, o mundo começou a perceber que os modelos de IA generativa de código aberto poderiam competir seriamente com o OpenAI GPT-3. Agora, com o LLaMA, a fasquia foi claramente elevada novamente e este modelo parece ser uma alternativa de código aberto muito boa ao OpenAI ChatGPT e GPT-4.

No entanto, a licença do LLaMA não é favorável às empresas: este modelo não pode ser utilizado para fins comerciais... Mas a boa notícia é que agora existem outros modelos.

O OpenLLaMA, lançado em junho de 2023, é uma versão alternativa do LLaMA, desenvolvida pela equipa de investigação de IA de Berkeley, que apresenta resultados muito bons e pode ser utilizada para fins comerciais. Atualmente, estão disponíveis 2 versões: 7B parâmetros e 13B parâmetros.

O XGen, lançado pela Salesforce em junho de 2023, é outro modelo de base muito poderoso que pode ser utilizado em aplicações comerciais. Até ao momento, apenas está disponível uma versão com parâmetros 7B. Vale a pena notar que este modelo suporta um contexto de 8k tokens, enquanto o LLaMA e o OpenLLaMA apenas suportam um conteúdo de 2k tokens.

Porquê afinar o seu próprio modelo?

Os modelos acima referidos são modelos de base, o que significa que foram treinados de forma não supervisionada num grande corpus de textos.

Estes modelos fundamentais de IA são normalmente uma boa base, mas precisam de ser afinados para compreender corretamente o que se pretende e produzir bons resultados. A forma mais fácil de o conseguir é utilizando a aprendizagem de poucos disparos (também conhecida como "engenharia de prontidão"). Não hesite em ler o nosso guia de aprendizagem dedicado de poucos disparos aqui.

A aprendizagem de poucos disparos é conveniente, uma vez que pode ser efectuada em tempo real sem ter de criar uma nova versão do modelo de IA generativo, mas por vezes não é suficiente.

Para obter resultados de ponta, é necessário afinar um modelo de IA para o seu próprio caso de utilização. Afinar significa que modificará alguns parâmetros do modelo com base nos seus próprios dados e, em seguida, obterá a sua própria versão do modelo.

A afinação é muito mais barata do que treinar um modelo generativo de IA a partir do zero, mas continua a exigir poder de computação, pelo que é necessário hardware avançado para afinar o seu próprio modelo. Algumas técnicas de afinação alternativas recentes requerem menos poder de computação (ver p-tuning, prompt tuning, soft tuning, parameter efficient fine-tuning, adapters, LoRA, QLoRA...), mas até agora não conseguimos obter o mesmo nível de qualidade com estas técnicas, pelo que não as vamos mencionar neste tutorial.

Ajuste fino do LLaMA em uma TPU com JAX e EasyLM

Neste tutorial, concentramo-nos no ajuste fino do LLaMA com a biblioteca EasyLM, lançada pela equipa de investigação de IA de Berkeley: 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.

Também pode afinar o OpenLLaMA ou o XGen utilizando a mesma técnica.

Aqui afinamos o LLaMA 7B numa Google TPU V3-8, mas pode perfeitamente fazer o mesmo numa GPU A100 (basta ler atentamente a parte da "Instalação" na documentação do EasyLM que é ligeiramente diferente). Claro que também pode afinar versões maiores do LLaMA (13B, 33B, 65B...), mas vai precisar de muito mais do que um TPU V3-8 ou um único GPU A100.

Aqui vamos nós!

Primeiro, crie um conjunto de dados de geração de texto para o seu caso de utilização, no formato JSONL, utilizando "text" como chave para cada exemplo. Aqui está um conjunto de dados simples de análise de sentimentos:

{"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"}

É de salientar alguns aspectos importantes. Em primeiro lugar, este conjunto de dados contém apenas 3 exemplos por uma questão de simplicidade, mas na vida real precisará de muitos mais exemplos. 300 exemplos é normalmente um bom começo. Em segundo lugar, quando utilizar o seu modelo aperfeiçoado para inferência, terá de seguir rigorosamente a mesma formatação, utilizando os prefixos "[Conteúdo]:" e "[Sentimento]:". Por último, o símbolo "</s>" é importante porque significa que o modelo deve parar de gerar aqui. Pode encontrar mais exemplos de conjuntos de dados na documentação do NLP Cloud: Saiba mais aqui.

Crie uma VM TPU V3-8 no Google Cloud com a versão de software V2 Alpha:

SSH na VM e instale o EasyLM:

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

Pode agora descarregar e converter os pesos LLaMA. A primeira opção é pedir os pesos oficiais à Meta: https://ai.facebook.com/blog/large-language-model-llama-meta-ai/. Em seguida, converta os pesos para EasyLM com este script: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_torch_to_easylm.py. A segunda opção é utilizar os pesos LLaMA no HuggingFace: https://huggingface.co/decapoda-research/llama-7b-hf. Em seguida, converta os pesos para EasyLM com este script: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_hf_to_easylm.py.

Carregue o seu conjunto de dados para a VM, conte quantos tokens contém, utilizando o 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())))"

Se treinar o seu modelo para um contexto de 1024 tokens, terá de dividir o número de tokens devolvido por 1024.

Se treinar o seu modelo para um contexto de 2048 tokens, terá de dividir o número de tokens devolvido por 2048.

Este número será o número de passos por época. Assim, por exemplo, se quiser treinar durante 5 épocas (o que é normalmente uma boa definição), terá de multiplicar este número por 5 e colocar o valor resultante em --total_steps abaixo.

Eis um exemplo concreto: se o seu conjunto de dados contém 100 000 tokens e pretende um contexto de 1024 tokens e 5 épocas, o número total de passos será (100 000/1024)*5 = 488.

Dependendo do comprimento do contexto, defina --train_dataset.json_dataset.seq_length como 1024 ou 2048 abaixo. Observe que o ajuste fino de um modelo para um contexto de 2048 tokens requer mais memória, portanto, se isso não for estritamente necessário, recomendamos que você se atenha a um contexto de 1024 tokens.

Pode agora iniciar o processo de afinação:

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 &

Algumas explicações:

--save_model_freq: a frequência com que pretende guardar o seu modelo durante o processo. Se só fizer um ajuste fino num pequeno conjunto de dados, pode guardar apenas no fim do processo e, nesse caso, este valor será igual a --total_steps.

--optimizer.adamw_optimizer.lr_warmup_steps: 10% do total de passos é normalmente um bom valor.

--tokenizer.vocab_file: o caminho para o ficheiro tokenizer.model. Por exemplo, se usar o repositório decapoda no HuggingFace, aqui está o link para o tokenizador: https://huggingface.co/decapoda-research/llama-7b-hf/resolve/main/tokenizer.model

--logger.output_dir: caminho para o modelo final e registos

Os outros parâmetros podem ser mantidos inalterados.

Quando o processo de afinação estiver concluído, pode recuperar o seu modelo no caminho especificado em --logger.output_dir.

Utilização do modelo aperfeiçoado para inferência

Agora tem o seu próprio modelo aperfeiçoado e quer utilizá-lo, claro!

Uma primeira estratégia consiste em utilizar a biblioteca EasyLM para a inferência. Nesse caso, pode lançar o servidor de inferência da seguinte forma:

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'

Em seguida, basta enviar os seus pedidos com cURL desta forma:

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]:"]}'

Uma segunda estratégia é exportar o seu modelo para o formato HuggingFace, a fim de efetuar a inferência com outra estrutura. Aqui está como exportá-lo:

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'

Conclusão

2023 foi um grande marco para os modelos de IA generativa de código aberto. A partir deste momento, toda a gente pode utilizar modelos fantásticos como LLaMA, OpenLLaMA, XGen, Orca, Falcon...

O aperfeiçoamento destes modelos é a melhor forma de obter resultados de vanguarda, adaptados ao seu próprio caso de utilização, que podem superar significativamente os melhores modelos de IA proprietários, como o ChatGPT (GPT-3.5), GPT-4, Claude...

Neste guia, mostrei como fazer o ajuste fino do LLaMA, do OpenLLaMA e do XGen. Se tiver dúvidas, não hesite em entrar em contato comigo e se quiser ajustar e implantar facilmente modelos avançados de IA generativa sem nenhuma complexidade técnica, have a look at the NLP Cloud dedicated documentation!

Mark
Engenheiro de aprendizagem automática na NLP Cloud