Come mettere a punto LLaMA, OpenLLaMA e XGen con JAX su una GPU o una TPU

LLaMA, OpenLLaMA e XGen sono modelli di intelligenza artificiale generativa all'avanguardia. Questi modelli danno risultati ancora migliori se messi a punto sui propri dati. In questo articolo vedremo come mettere a punto questi modelli sia su una GPU che su una TPU, utilizzando JAX e la libreria EasyLM.

LLaMA, OpenLLaM e XGen

Il modello LaMa è stato rilasciato da Meta nel febbraio 2023. Questo modello di intelligenza artificiale generativa è un modello open-source proposto in diverse dimensioni: 7B parametri, 13B parametri, 33B parametri e 65B parametri.

Nel giugno 2021, quando è stato rilasciato GPT-J, il mondo ha iniziato a rendersi conto che i modelli di IA generativa open-source potevano seriamente competere con OpenAI GPT-3. Ora, con LLaMA, l'asticella è stata chiaramente alzata di nuovo e questo modello sembra essere un'ottima alternativa open-source a OpenAI ChatGPT e GPT-4.

La licenza di LLaMA non è però favorevole alle imprese: questo modello non può essere utilizzato per scopi commerciali... Ma la buona notizia è che ora esistono altri modelli.

OpenLLaMA, rilasciato nel giugno 2023, è una versione alternativa di LLaMA, sviluppata dal team di Berkeley AI Research, che fornisce ottimi risultati e può essere utilizzata per le aziende. Al momento sono disponibili 2 versioni: 7B parametri e 13B parametri.

XGen, rilasciato da Salesforce nel giugno 2023, è un altro modello fondamentale molto potente che può essere utilizzato in applicazioni commerciali. Al momento è disponibile solo una versione con parametri 7B. Vale la pena notare che questo modello supporta un contesto di 8k tokens, mentre LLaMA e OpenLLaMA supportano solo un contenuto di 2k tokens.

Perché mettere a punto il proprio modello?

I modelli sopra descritti sono modelli fondazionali, il che significa che sono stati addestrati in modo non supervisionato su un ampio corpus di testi.

Questi modelli fondamentali di intelligenza artificiale sono di solito una buona base, ma devono essere messi a punto per capire correttamente ciò che si vuole e per fornire buoni risultati. Il modo più semplice per raggiungere questo obiettivo è utilizzare l'apprendimento a pochi colpi (noto anche come "prompt engineering"). Non esitate a leggere la nostra guida all'apprendimento di pochi scatti qui.

L'apprendimento a pochi colpi è comodo perché può essere eseguito al volo senza dover creare una nuova versione del modello generativo dell'IA, ma a volte non è sufficiente.

Per ottenere risultati all'avanguardia, è necessario mettere a punto un modello di intelligenza artificiale per il proprio caso d'uso. Per messa a punto si intende la modifica di alcuni parametri del modello in base ai propri dati, per poi ottenere la propria versione del modello.

La sintonizzazione fine è molto più economica dell'addestramento di un modello generativo di intelligenza artificiale da zero, ma richiede comunque potenza di calcolo, quindi è necessario un hardware avanzato per sintonizzare il proprio modello. Alcune recenti tecniche alternative di fine-tuning richiedono meno potenza di calcolo (vedi p-tuning, prompt tuning, soft tuning, parameter efficient fine-tuning, adapters, LoRA, QLoRA...) ma finora non siamo riusciti a ottenere lo stesso livello di qualità con queste tecniche, quindi non le menzioneremo in questo tutorial.

Messa a punto di LLaMA su una TPU con JAX e EasyLM

In questo tutorial ci concentriamo sulla messa a punto di LLaMA con la libreria EasyLM, rilasciata dal team 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.

Con la stessa tecnica è possibile mettere a punto OpenLLaMA o XGen.

Qui abbiamo messo a punto LLaMA 7B su una TPU V3-8 di Google, ma è possibile fare lo stesso su una GPU A100 (basta leggere attentamente la parte "Installazione" nella documentazione di EasyLM, che è leggermente diversa). Naturalmente è possibile mettere a punto anche versioni più grandi di LLaMA (13B, 33B, 65B...), ma è necessario molto più di una TPU V3-8 o di una singola GPU A100.

Eccoci qua!

Per prima cosa, creare un set di dati per la generazione di testo per il proprio caso d'uso, in formato JSONL, utilizzando "text" come chiave per ogni esempio. Ecco un semplice set di dati per l'analisi del sentiment:

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

Si prega di notare un paio di cose importanti. Innanzitutto, questo set di dati contiene solo 3 esempi per semplicità, ma nella vita reale ne serviranno molti di più. 300 esempi sono di solito un buon inizio. In secondo luogo, quando utilizzerete il vostro modello perfezionato per l'inferenza, dovrete seguire rigorosamente la stessa formattazione, utilizzando i prefissi "[Content]:" e "[Sentiment]:". Infine, il token "</s>" è importante perché significa che il modello deve smettere di generare qui. Ulteriori esempi di set di dati sono disponibili nella documentazione di NLP Cloud: per saperne di più qui.

Creare una VM TPU V3-8 su Google Cloud con la versione software V2 Alpha:

SSH nella macchina virtuale e installare EasyLM:

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

È ora possibile scaricare e convertire i pesi LLaMA. La prima opzione è chiedere a Meta i pesi ufficiali: https://ai.facebook.com/blog/large-language-model-llama-meta-ai/. Quindi convertire i pesi in EasyLM con questo script: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_torch_to_easylm.py. La seconda opzione è quella di utilizzare i pesi LLaMA su HuggingFace: https://huggingface.co/decapoda-research/llama-7b-hf. Quindi convertire i pesi in EasyLM con questo script: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_hf_to_easylm.py.

Caricare il set di dati sulla macchina virtuale, contare quanti token contiene, utilizzando il tokenizer 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 si addestra il modello per un contesto di 1024 token, è necessario dividere il numero di token restituito per 1024.

Se si addestra il modello per un contesto di 2048 tokens, è necessario dividere il numero di tokens restituito per 2048.

Questo numero sarà il numero di passi per epoche. Quindi, ad esempio, se si vuole allenare per 5 epoche (che di solito è una buona impostazione), occorre moltiplicare questo numero per 5 e inserire il valore risultante in --total_steps.

Ecco un esempio concreto: se il vostro set di dati contiene 100.000 token e volete un contesto di 1024 token e 5 epoche, il numero totale di passi sarà (100.000/1024)*5 = 488.

A seconda della lunghezza del contesto, impostare --train_dataset.json_dataset.seq_length come 1024 o 2048. Si noti che la messa a punto di un modello per un contesto di 2048 tokens richiede più memoria, quindi se non è strettamente necessario si consiglia di attenersi a un contesto di 1024 tokens.

A questo punto è possibile avviare il processo di messa a punto:

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 &

Alcune spiegazioni:

--save_model_freq: la frequenza con cui si desidera salvare il modello durante il processo. Se si esegue la messa a punto solo su un piccolo insieme di dati, si può salvare solo alla fine del processo; in tal caso questo valore sarà uguale a --total_steps.

--optimizer.adamw_optimizer.lr_warmup_steps: Il 10% dei passi totali è di solito un buon valore.

--tokenizer.vocab_file: il percorso del file tokenizer.model. Per esempio, se si usa il repository decapoda su HuggingFace, ecco il link al tokenizer: https://huggingface.co/decapoda-research/llama-7b-hf/resolve/main/tokenizer.model

--logger.output_dir: percorso del modello finale e dei registri

Gli altri parametri possono essere lasciati inalterati.

Una volta terminato il processo di messa a punto, è possibile recuperare il modello nel percorso specificato in --logger.output_dir.

Utilizzo del modello ottimizzato per l'inferenza

Ora avete il vostro modello perfezionato e volete usarlo, naturalmente!

Una prima strategia consiste nell'utilizzare la libreria EasyLM per l'inferenza. In questo caso si può lanciare il server di inferenza in questo modo:

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'

Quindi è sufficiente inviare le richieste con cURL in questo modo:

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 seconda strategia consiste nell'esportare il modello nel formato HuggingFace per eseguire l'inferenza con un altro framework. Ecco come esportarlo:

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'

Conclusione

Il 2023 è stata una grande pietra miliare per i modelli di intelligenza artificiale generativa open-source. A partire da questo momento, tutti possono utilizzare grandi modelli come LLaMA, OpenLLaMA, XGen, Orca, Falcon...

La messa a punto di questi modelli è il modo migliore per ottenere risultati all'avanguardia, adattati al proprio caso d'uso, che possono superare in modo significativo i migliori modelli AI proprietari come ChatGPT (GPT-3.5), GPT-4, Claude...

In questa guida vi ho mostrato come mettere a punto LLaMA, OpenLLaMA e XGen. Se avete domande non esitate a contattarmi, e se volete mettere a punto e distribuire facilmente modelli avanzati di IA generativa senza alcuna complessità tecnica, have a look at the NLP Cloud dedicated documentation!

Mark
Ingegnere dell'apprendimento automatico presso NLP Cloud