Kæmper du med AI eller full-stack-udvikling? Vores eksperter er her for at vejlede dig: skræddersyet rådgivning, teknisk integration og meget mere. Kontakt os på [email protected].

Sådan finjusterer du LLaMA, OpenLLaMA og XGen med JAX på en GPU eller en TPU

LLaMA, OpenLLaMA og XGen er banebrydende generative AI-modeller. Disse modeller giver endnu bedre resultater, når de finjusteres på dine egne data. I denne artikel skal vi se, hvordan man finjusterer disse modeller på både en GPU og en TPU ved hjælp af JAX og EasyLM-biblioteket.

LLaMA, OpenLLaM og XGen

LLaMA-modellen blev udgivet af Meta i februar 2023. Denne generative AI-model er en open source-model, der foreslås i flere størrelser: 7B parametre, 13B parametre, 33B parametre og 65B parametre.

I juni 2021, da GPT-J blev frigivet, begyndte verden at indse, at open source generative AI-modeller for alvor kunne konkurrere med OpenAI GPT-3. Nu med LLaMA er barren klart blevet hævet igen, og denne model ser ud til at være et meget godt open source-alternativ til OpenAI ChatGPT og GPT-4.

LLaMAs licens er dog ikke forretningsvenlig: denne model kan ikke bruges til kommercielle formål... Men den gode nyhed er, at der nu findes andre modeller.

OpenLLaMA, udgivet i juni 2023, er en alternativ version af LLaMA, udviklet af Berkeley AI Research team, der giver meget gode resultater og kan bruges til erhvervslivet. 2 versioner er tilgængelige i skrivende stund: 7B parametre og 13B parametre.

XGen, der blev udgivet af Salesforce i juni 2023, er en anden meget stærk grundmodel, der kan bruges i kommercielle applikationer. Kun en version med 7B-parametre er tilgængelig i skrivende stund. Det er værd at bemærke, at denne model understøtter en 8k tokens-kontekst, mens LLaMA og OpenLLaMA kun understøtter et 2k tokens-indhold.

Hvorfor finjustere din egen model?

Ovenstående modeller er grundlæggende modeller, hvilket betyder, at de er blevet trænet på en uovervåget måde på et stort korpus af tekster.

Disse grundlæggende AI-modeller er normalt et godt grundlag, men de skal indstilles, så de forstår, hvad du ønsker, og giver gode resultater. Den nemmeste måde at opnå dette på er ved at bruge few-shot learning (også kendt som "prompt engineering"). Du er velkommen til at læse vores dedikerede læringsguide til få skud her.

Few-shot learning er praktisk, da det kan udføres på farten uden at skulle oprette en ny version af den generative AI-model, men det er nogle gange ikke nok.

For at få state-of-the-art-resultater vil du gerne finjustere en AI-model til din egen brugssag. Finjustering betyder, at du ændrer nogle parametre i modellen baseret på dine egne data og derefter får din egen version af modellen.

Finjustering er meget billigere end at træne en generativ AI-model fra bunden, men det kræver stadig regnekraft, så du har brug for avanceret hardware for at kunne finjustere din egen model. Nogle nyere alternative finjusteringsteknikker kræver mindre regnekraft (se p-tuning, prompt tuning, soft tuning, parameter efficient fine-tuning, adapters, LoRA, QLoRA...), men indtil videre er det ikke lykkedes os at opnå det samme kvalitetsniveau med disse teknikker, så vi vil ikke nævne dem i denne vejledning.

Finjustering af LLaMA på en TPU med JAX og EasyLM

I denne tutorial fokuserer vi på at finjustere LLaMA med EasyLM-biblioteket, som er udgivet af Berkeley AI Research-teamet: 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.

Du kan også finjustere OpenLLaMA eller XGen ved hjælp af den samme teknik.

Vi finjusterer LLaMA 7B på en Google TPU V3-8 her, men du kan sagtens gøre det samme på en A100 GPU (læs blot omhyggeligt delen "Installation" i EasyLM-dokumentationen, der er lidt anderledes). Selvfølgelig kan du også finjustere større versioner af LLaMA (13B, 33B, 65B...), men du skal bruge meget mere end en TPU V3-8 eller en enkelt A100 GPU.

Så er vi i gang!

Først skal du oprette et tekstgenereringsdatasæt til din use case i JSONL-format med "text" som nøgle for hvert eksempel. Her er et simpelt sentimentanalyse-datasæt:

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

Bemærk venligst et par vigtige ting. For det første indeholder dette datasæt kun 3 eksempler for enkelthedens skyld, men i det virkelige liv har du brug for mange flere eksempler. 300 eksempler er normalt en god start. For det andet, når du bruger din finjusterede model til inferens, skal du nøje følge den samme formatering ved at bruge præfikserne "[Content]:" og "[Sentiment]:". Sidst, men ikke mindst, er "</s>"-tokenet vigtigt, fordi det betyder, at modellen skal stoppe med at generere her. Du kan finde flere eksempler på datasæt i NLP Cloud-dokumentationen: Læs mere her.

Opret en TPU V3-8 VM på Google Cloud med softwareversionen V2 Alpha:

SSH ind i VM'en og installer EasyLM:

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

Du kan nu downloade og konvertere LLaMA-vægtene. Den første mulighed er at bede Meta om de officielle vægte: https://ai.facebook.com/blog/large-language-model-llama-meta-ai/. Konverter derefter vægtene til EasyLM med dette script: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_torch_to_easylm.py. Den anden mulighed er at bruge LLaMA-vægtene på HuggingFace: https://huggingface.co/decapoda-research/llama-7b-hf. Konverter derefter vægtene til EasyLM med dette script: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_hf_to_easylm.py.

Upload dit datasæt til VM'en, tæl, hvor mange tokens det indeholder, ved hjælp af HF LLaMA tokenizer:

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())))"

Hvis du træner din model til en kontekst med 1024 tokens, skal du dividere det returnerede antal tokens med 1024.

Hvis du træner din model til en kontekst med 2048 tokens, skal du dividere det returnerede antal tokens med 2048.

Dette tal vil være antallet af trin pr. epoke. Så hvis du f.eks. vil træne i 5 epoker (hvilket normalt er en god indstilling), skal du gange dette tal med 5 og sætte den resulterende værdi ind i --total_steps nedenfor.

Her er et konkret eksempel: Hvis dit datasæt indeholder 100.000 tokens, og du ønsker en kontekst på 1024 tokens og 5 epoker, vil dit samlede antal trin være (100.000/1024)*5 = 488.

Afhængigt af din kontekstlængde skal du indstille --train_dataset.json_dataset.seq_length som 1024 eller 2048 nedenfor. Bemærk, at finjustering af en model til en kontekst på 2048 tokens kræver mere hukommelse, så hvis det ikke er strengt nødvendigt, anbefaler vi, at du holder dig til en kontekst på 1024 tokens.

Du kan nu starte finjusteringsprocessen:

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 &

Nogle forklaringer:

--save_model_freq: hvor ofte du ønsker at gemme din model under processen. Hvis du kun finjusterer på et lille datasæt, kan du kun gemme i slutningen af processen, og i det tilfælde vil denne værdi være lig med --total_steps.

--optimizer.adamw_optimizer.lr_warmup_steps: 10% af det samlede antal trin er normalt en god værdi.

--tokenizer.vocab_file: stien til tokenizer.model-filen. Hvis du for eksempel bruger decapoda-repositoriet på HuggingFace, er her linket til tokenizeren: https://huggingface.co/decapoda-research/llama-7b-hf/resolve/main/tokenizer.model

--logger.output_dir: sti til den endelige model og logfiler

De andre parametre kan forblive uberørte.

Når finjusteringsprocessen er færdig, kan du hente din model på den sti, du har angivet i --logger.output_dir.

Brug af den finjusterede model til inferens

Nu har du din egen finjusterede model, og du vil selvfølgelig gerne bruge den!

En første strategi er at bruge EasyLM-biblioteket til inferens. I så fald kan du starte inferensserveren på denne måde:

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'

Så skal du blot sende dine anmodninger med cURL på denne måde:

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

En anden strategi er at eksportere din model til HuggingFace-formatet for at kunne udføre inferens med et andet framework. Her kan du se, hvordan du eksporterer den:

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'

Konklusion

2023 har været en stor milepæl for open source generative AI-modeller. I skrivende stund kan alle bruge fantastiske modeller som LLaMA, OpenLLaMA, XGen, Orca, Falcon...

Finjustering af disse modeller er den bedste måde at opnå banebrydende resultater, skræddersyet til din egen brugssag, der kan overgå de bedste proprietære AI-modeller som ChatGPT (GPT-3.5), GPT-4, Claude...

I denne guide viste jeg dig, hvordan du finjusterer LLaMA, OpenLLaMA og XGen. Hvis du har spørgsmål, så tøv ikke med at kontakte mig, og hvis du nemt vil finjustere og implementere avancerede generative AI-modeller uden nogen teknisk kompleksitet, have a look at the NLP Cloud dedicated documentation!

Mark
Maskinlæringsingeniør hos NLP Cloud