Hur man finjusterar LLaMA, OpenLLaMA och XGen med JAX på en GPU eller en TPU

LLaMA, OpenLLaMA och XGen är banbrytande generativa AI-modeller. Dessa modeller ger till och med mycket bättre resultat när de finjusteras på dina egna data. I den här artikeln ska vi se hur man finjusterar dessa modeller på både en GPU och en TPU, med hjälp av JAX och EasyLM-biblioteket.

LLaMA, OpenLLaM och XGen

LLaMA-modellen släpptes av Meta i februari 2023. Denna generativa AI-modell är en öppen källkodsmodell som föreslås i flera storlekar: 7B parametrar, 13B parametrar, 33B parametrar och 65B parametrar.

I juni 2021, när GPT-J släpptes, började världen inse att generativa AI-modeller med öppen källkod på allvar kunde konkurrera med OpenAI GPT-3. Nu med LLaMA har ribban tydligt höjts igen och den här modellen verkar vara ett mycket bra open source-alternativ till OpenAI ChatGPT och GPT-4.

LLaMA:s licens är dock inte särskilt affärsvänlig: denna modell får inte användas för kommersiella ändamål... Men den goda nyheten är att det nu finns andra modeller.

OpenLLaMA, som släpptes i juni 2023, är en alternativ version av LLaMA som utvecklats av Berkeleys AI-forskningsteam och som ger mycket bra resultat och kan användas för affärsändamål. 2 versioner är tillgängliga i skrivande stund: 7B parametrar och 13B parametrar.

XGen, som Salesforce lanserade i juni 2023, är en annan mycket kraftfull grundmodell som kan användas i kommersiella applikationer. Endast en version med 7B-parametrar är tillgänglig när detta skrivs. Det är värt att notera att denna modell stöder en kontext på 8k tokens, medan LLaMA och OpenLLaMA endast stöder ett innehåll på 2k tokens.

Varför finjustera din egen modell?

Modellerna ovan är grundmodeller, vilket innebär att de har tränats på ett oövervakat sätt på en stor textkorpus.

Dessa grundläggande AI-modeller är vanligtvis en bra grund, men de måste ställas in för att korrekt förstå vad du vill ha och returnera bra resultat. Det enklaste sättet att uppnå detta är genom att använda "few-shot learning" (även känt som "prompt engineering"). Läs gärna vår dedikerade guide om några få skott här.

Few-shot learning är bekvämt eftersom det kan utföras i farten utan att man behöver skapa en ny version av den generativa AI-modellen, men ibland är det inte tillräckligt.

För att få toppmoderna resultat kommer du att vilja finjustera en AI-modell för ditt eget användningsfall. Finjustering innebär att du ändrar vissa parametrar i modellen baserat på dina egna data och sedan får din egen version av modellen.

Finjustering är mycket billigare än att lära upp en generativ AI-modell från grunden, men det kräver fortfarande beräkningskraft, så du behöver avancerad hårdvara för att finjustera din egen modell. Vissa nya alternativa finjusteringstekniker kräver mindre beräkningskraft (se p-tuning, prompt tuning, soft tuning, parameter efficient fine-tuning, adapters, LoRA, QLoRA...) men hittills har vi inte lyckats uppnå samma kvalitetsnivå med dessa tekniker så vi kommer inte att nämna dem i den här handledningen.

Finjustering av LLaMA på en TPU med JAX och EasyLM

I den här handledningen fokuserar vi på att finjustera LLaMA med EasyLM-biblioteket, som släppts av Berkeleys AI Research-team: 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 också finjustera OpenLLaMA eller XGen med samma teknik.

Vi finjusterar LLaMA 7B på en Google TPU V3-8 här, men du kan perfekt göra samma sak på en A100 GPU (läs bara noggrant delen "Installation" i EasyLM-dokumentationen som är något annorlunda). Naturligtvis kan du också finjustera större versioner av LLaMA (13B, 33B, 65B...) men du kommer att behöva mycket mer än en TPU V3-8 eller en enda A100 GPU.

Nu kör vi!

Skapa först en datauppsättning för textgenerering för ditt användningsfall, i JSONL-format, med "text" som nyckel för varje exempel. Här är en enkel datauppsättning för sentimentanalys:

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

Observera ett par viktiga saker. För det första innehåller detta dataset bara 3 exempel för enkelhetens skull, men i verkliga livet behöver du många fler exempel. 300 exempel är vanligtvis en bra start. För det andra, när du använder din finjusterade modell för inferens, måste du strikt följa samma formatering med prefixen "[Innehåll]:" och "[Sentiment]:". Sist men inte minst är tecknet "</s>" viktigt eftersom det innebär att modellen ska sluta generera här. Du hittar fler exempel på dataset i NLP Cloud-dokumentationen: Läs mer här.

Skapa en TPU V3-8 VM på Google Cloud med programvaruversionen V2 Alpha:

SSH till den virtuella datorn och installera EasyLM:

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

Du kan nu ladda ner och konvertera LLaMA-vikterna. Det första alternativet är att be Meta om de officiella vikterna: https://ai.facebook.com/blog/large-language-model-llama-meta-ai/. Konvertera sedan vikterna till EasyLM med detta skript: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_torch_to_easylm.py. Det andra alternativet är att använda LLaMA-vikterna på HuggingFace: https://huggingface.co/decapoda-research/llama-7b-hf. Konvertera sedan vikterna till EasyLM med detta skript: https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_hf_to_easylm.py.

Ladda upp ditt dataset till VM och räkna hur många tokens det innehåller med hjälp 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())))"

Om du tränar din modell för en kontext med 1024 tokens måste du dividera det returnerade antalet tokens med 1024.

Om du tränar din modell för en kontext med 2048 tokens måste du dividera det returnerade antalet tokens med 2048.

Detta tal kommer att vara antalet steg per epok. Så om du till exempel vill träna i 5 epoker (vilket vanligtvis är en bra inställning) måste du multiplicera detta nummer med 5 och ange det resulterande värdet i --total_steps nedan.

Här är ett konkret exempel: om ditt dataset innehåller 100 000 tokens, och du vill ha en kontext med 1024 tokens och 5 epoker, blir det totala antalet steg (100 000/1024)*5 = 488.

Beroende på kontextens längd anger du --train_dataset.json_dataset.seq_length som 1024 eller 2048 nedan. Observera att finjustering av en modell för en kontext med 2048 tokens kräver mer minne, så om detta inte är absolut nödvändigt rekommenderar vi att du håller dig till en kontext med 1024 tokens.

Du kan nu påbörja 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 &

Några förklaringar:

--save_model_freq: hur ofta du vill spara din modell under processen. Om du bara finjusterar i ett litet dataset kan du spara i slutet av processen, och i så fall kommer detta värde att vara lika med --total_steps.

--optimizer.adamw_optimizer.lr_warmup_steps: 10% av det totala antalet steg är vanligtvis ett bra värde.

--tokenizer.vocab_file: sökvägen till filen tokenizer.model. Om du till exempel använder decapoda-arkivet på HuggingFace är här länken till tokenizern: https://huggingface.co/decapoda-research/llama-7b-hf/resolve/main/tokenizer.model

--logger.output_dir: sökväg till slutlig modell och loggar

Övriga parametrar kan lämnas orörda.

När finjusteringsprocessen är klar kan du hämta din modell på den sökväg som du angav i --logger.output_dir.

Användning av den finjusterade modellen för inferens

Nu har du din egen finjusterade modell och du vill naturligtvis använda den!

En första strategi är att använda EasyLM-biblioteket för inferens. I så fall kan du starta inferensservern på följande sätt:

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'

Skicka sedan dina förfrågningar med cURL på följande sätt:

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 annan strategi är att exportera din modell till HuggingFace-formatet för att kunna utföra inferens med ett annat ramverk. Så här kan du exportera 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'

Slutsats

2023 har varit en stor milstolpe för generativa AI-modeller med öppen källkod. I skrivande stund kan alla använda fantastiska modeller som LLaMA, OpenLLaMA, XGen, Orca, Falcon...

Att finjustera dessa modeller är det bästa sättet att uppnå banbrytande resultat, skräddarsydda för ditt eget användningsområde, som avsevärt kan överträffa de bästa egenutvecklade AI-modellerna som ChatGPT (GPT-3.5), GPT-4, Claude...

I den här guiden visade jag dig hur du finjusterar LLaMA, OpenLLaMA och XGen. Om du har frågor tveka inte att nå ut till mig, och om du enkelt vill finjustera och distribuera avancerade generativa AI-modeller utan någon teknisk komplexitet, have a look at the NLP Cloud dedicated documentation!

Mark
Maskininlärningsingenjör på NLP Cloud