Sliter du med AI eller fullstack-utvikling? Ekspertene våre er her for å veilede deg: skreddersydde råd, teknisk integrasjon og mer. Ta kontakt på [email protected].

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

LLaMA, OpenLLaMA og XGen er banebrytende generative AI-modeller. Disse modellene gir enda bedre resultater når de finjusteres på dine egne data. I denne artikkelen skal vi se hvordan du finjusterer disse modellene på både en GPU og en TPU ved hjelp av JAX og EasyLM-biblioteket.

LLaMA, OpenLLaM og XGen

LLaMA-modellen ble lansert av Meta i februar 2023. Denne generative AI-modellen er en åpen kildekode-modell som tilbys i flere størrelser: 7B parametere, 13B parametere, 33B parametere og 65B parametere.

I juni 2021, da GPT-J ble lansert, begynte verden å innse at generative AI-modeller med åpen kildekode for alvor kunne konkurrere med OpenAI GPT-3. Nå, med LLaMA, har listen tydeligvis blitt hevet igjen, og denne modellen ser ut til å være et svært godt open source-alternativ til OpenAI ChatGPT og GPT-4.

LLaMAs lisens er imidlertid ikke forretningsvennlig: denne modellen kan ikke brukes til kommersielle formål... Men den gode nyheten er at det nå finnes andre modeller.

OpenLLaMA, som ble lansert i juni 2023, er en alternativ versjon av LLaMA, utviklet av Berkeley AI Research-teamet, som gir svært gode resultater og kan brukes i næringslivet. 2 versjoner er tilgjengelige i skrivende stund: 7B-parametere og 13B-parametere.

XGen, som ble lansert av Salesforce i juni 2023, er en annen svært kraftig grunnmodell som kan brukes i kommersielle applikasjoner. I skrivende stund er bare en versjon med 7B-parametere tilgjengelig. Det er verdt å merke seg at denne modellen støtter en kontekst på 8k tokens, mens LLaMA og OpenLLaMA bare støtter et innhold på 2k tokens.

Hvorfor finjustere din egen modell?

De ovennevnte modellene er grunnleggende modeller, noe som betyr at de er trent uten tilsyn på et stort tekstkorpus.

Disse grunnleggende AI-modellene er vanligvis et godt utgangspunkt, men de må finjusteres for å forstå hva du ønsker og gi gode resultater. Den enkleste måten å oppnå dette på er å bruke "few-shot learning" (også kjent som "prompt engineering"). Les gjerne vår egen læringsguide med få bilder her.

Few-shot-læring er praktisk, ettersom det kan utføres i farten uten å måtte lage en ny versjon av den generative AI-modellen, men noen ganger er det ikke nok.

For å få toppmoderne resultater må du finjustere en AI-modell for ditt eget bruksområde. Finjustering betyr at du endrer noen parametere i modellen basert på dine egne data og deretter får din egen versjon av modellen.

Finjustering er mye billigere enn å trene opp en generativ AI-modell fra bunnen av, men det krever fortsatt regnekraft, så du trenger avansert maskinvare for å finjustere din egen modell. Noen nyere alternative finjusteringsteknikker krever mindre regnekraft (se p-tuning, prompt tuning, soft tuning, parametereffektiv finjustering, adaptere, LoRA, QLoRA ...), men så langt har vi ikke klart å oppnå samme kvalitetsnivå med disse teknikkene, så vi kommer ikke til å nevne dem i denne veiledningen.

Finjustering av LLaMA på en TPU med JAX og EasyLM

I denne veiledningen fokuserer vi på finjustering av LLaMA med EasyLM-biblioteket, utgitt av 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 hjelp av samme teknikk.

Vi finjusterer LLaMA 7B på en Google TPU V3-8 her, men du kan fint gjøre det samme på en A100 GPU (bare les nøye gjennom "Installasjon"-delen i EasyLM-dokumentasjonen, som er litt annerledes). Du kan selvfølgelig også finjustere større versjoner av LLaMA (13B, 33B, 65B...), men da trenger du mye mer enn en TPU V3-8 eller en enkelt A100 GPU.

Da setter vi i gang!

Først oppretter du et tekstgenereringsdatasett for bruksområdet ditt, i JSONL-format, med "text" som nøkkel for hvert eksempel. Her er et enkelt datasett for sentimentanalyse:

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

Vær oppmerksom på et par viktige ting. For det første inneholder dette datasettet bare 3 eksempler for enkelhets skyld, men i virkeligheten vil du trenge mange flere eksempler. 300 eksempler er vanligvis en god start. For det andre, når du skal bruke den finjusterte modellen din til inferens, må du følge den samme formateringen og bruke prefiksene "[Content]:" og "[Sentiment]:". Til slutt er symbolet "</s>" viktig fordi det betyr at modellen skal slutte å generere her. Du finner flere eksempler på datasett i NLP Cloud-dokumentasjonen: Les mer her.

Opprett en TPU V3-8 VM på Google Cloud med programvareversjonen V2 Alpha:

SSH inn i den virtuelle maskinen og installer EasyLM:

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

Du kan nå laste ned og konvertere LLaMA-vektene. Det første alternativet er å be Meta om de offisielle vektene: https://ai.facebook.com/blog/large-language-model-llama-meta-ai/. Deretter konverterer du vektene til EasyLM med dette skriptet: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_torch_to_easylm.py. Det andre alternativet er å bruke LLaMA-vektene på HuggingFace: https://huggingface.co/decapoda-research/llama-7b-hf. Deretter konverterer du vektene til EasyLM med dette skriptet: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_hf_to_easylm.py.

Last opp datasettet til den virtuelle maskinen, tell hvor mange tokens det inneholder ved hjelp av 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 trener modellen for en kontekst med 1024 tokens, må du dele det returnerte antallet tokens med 1024.

Hvis du trener modellen for en kontekst med 2048 tokens, må du dele det returnerte antallet tokens med 2048.

Dette tallet angir antall trinn per epoke. Hvis du for eksempel vil trene i 5 epoker (noe som vanligvis er en god innstilling), må du multiplisere dette tallet med 5 og angi den resulterende verdien i --total_steps nedenfor.

Her er et konkret eksempel: Hvis datasettet inneholder 100 000 tokens, og du ønsker en kontekst på 1024 tokens og 5 epoker, blir det totale antallet trinn (100 000/1024)*5 = 488.

Avhengig av kontekstens lengde angir du --train_dataset.json_dataset.seq_length som 1024 eller 2048 nedenfor. Merk at finjustering av en modell for en kontekst på 2048 tokens krever mer minne, så hvis dette ikke er strengt nødvendig, anbefaler vi at du holder deg til en kontekst på 1024 tokens.

Du kan nå starte finjusteringsprosessen:

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 &

Noen forklaringer:

--save_model_freq: hvor ofte du vil lagre modellen i løpet av prosessen. Hvis du bare finjusterer på et lite datasett, kan du bare lagre på slutten av prosessen, og i så fall vil denne verdien være lik --total_steps.

--optimizer.adamw_optimizer.lr_warmup_steps: 10 % av det totale antall trinn er vanligvis en god verdi.

--tokenizer.vocab_file: banen til tokenizer.model-filen. Hvis du for eksempel bruker decapoda-repositoriet på HuggingFace, finner du lenken til tokenizeren her: https://huggingface.co/decapoda-research/llama-7b-hf/resolve/main/tokenizer.model

--logger.output_dir: sti til den endelige modellen og loggene

De andre parametrene kan stå urørt.

Når finjusteringsprosessen er ferdig, kan du hente modellen på banen du har angitt i --logger.output_dir.

Bruk av den finjusterte modellen for inferens

Nå har du din egen finjusterte modell, og du vil selvfølgelig bruke den!

En første strategi er å bruke EasyLM-biblioteket til inferens. I så fall kan du starte inferensserveren på denne måten:

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'

Da er det bare å sende forespørsler med cURL på denne måten:

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 annen strategi er å eksportere modellen til HuggingFace-formatet for å utføre inferens med et annet rammeverk. Slik eksporterer du modellen:

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'

Konklusjon

2023 har vært en stor milepæl for generative AI-modeller med åpen kildekode. I skrivende stund kan alle bruke gode modeller som LLaMA, OpenLLaMA, XGen, Orca, Falcon ...

Finjustering av disse modellene er den beste måten å oppnå banebrytende resultater, skreddersydd til ditt eget bruksområde, som kan overgå de beste proprietære AI-modellene som ChatGPT (GPT-3.5), GPT-4, Claude ...

I denne guiden viste jeg deg hvordan du finjusterer LLaMA, OpenLLaMA og XGen. Hvis du har spørsmål, er du velkommen til å ta kontakt med meg, og hvis du enkelt vil finjustere og distribuere avanserte generative AI-modeller uten teknisk kompleksitet, have a look at the NLP Cloud dedicated documentation!

Mark
Maskinlæringsingeniør hos NLP Cloud