Gemma 2 fijn-afstemmen en lokaal gebruiken

In deze tutorial leren we hoe je het Gemma 2-model fijn-afstemt op een dataset met gesprekken tussen patiënt en arts. We converteren het model ook naar GGUF-indeling zodat je het offline lokaal op je laptop kunt gebruiken. 

We bekijken eerst Gemma 2 en de verbeteringen ten opzichte van eerdere generaties, voordat we het Gemma 2-model fijn-afstemmen op de dataset met patiënt-artsconversaties. Daarna voegen we de opgeslagen adapter samen met het basismodel, pushen we het volledige model naar de Hugging Face Hub en gebruiken we een Hugging Face Space om het model te converteren en te kwantiseren. Tot slot downloaden we het gekwantiseerde model en gebruiken we het lokaal met de Jan-applicatie. 

Afbeelding door auteur

Gemma 2 begrijpen

Gemma 2 is de nieuwste release in de Gemma-familie van open large language models (LLM’s) van Google. Het is beschikbaar voor onderzoekers en ontwikkelaars in twee groottes, 9 miljard (9B) en 27 miljard (27B) parameters, onder een commercieel-vriendelijke licentie. Dit betekent dat je het kunt fijn-afstemmen op je eigen private dataset en het fijnafstemde model zonder beperkingen in productie kunt inzetten. 

Bron: Google lanceert Gemma 2

Gemma 2 biedt aanzienlijke verbeteringen in prestaties en inference-efficiëntie. De herontworpen architectuur zorgt voor razendsnelle inference op diverse hardwareopstellingen en naadloze integratie met belangrijke AI-frameworks zoals Hugging Face Transformers, JAX, PyTorch en TensorFlow. 

Gemma 2 bevat ook robuuste veiligheidsmaatregelen en tools voor verantwoorde AI-implementatie. Het presteert beter dan Llama 3 en Groq-1 op verschillende benchmarks en wordt geleverd met verbeterde Keras 3-integratie voor naadloos fijn-afstemmen en modelinference. 

Als je meer wilt leren over de eerste generatie Gemma-modellen, is onze tutorial Fine Tuning Google Gemma: Enhancing LLMs with Customized Instructions een uitstekend startpunt. Deze handleiding geeft uitgebreide uitleg over het fijn-afstemmen van Gemma-modellen, zoals de Gemma 7b-it, op specifieke datasets en taken. 

Toegang krijgen tot Gemma 2

In deze sectie downloaden we het model, laden we het in 4-bit-kwantisatie en draaien we vervolgens inference op een GPU. 

1. Installeer alle benodigde Python-pakketten. 

%%capture
%pip install -U bitsandbytes
%pip install -U transformers
%pip install -U accelerate

2. We gebruiken Kaggle als onze ontwikkelomgeving. We moeten een omgevingsvariabele aanmaken via de functie "Secrets" en die vervolgens activeren. De omgevingsvariabele bevat een Hugging Face-token dat je kunt genereren door in te loggen op de Hugging Face-website. 

3. Haal het Hugging Face-token veilig op uit Kaggle Secrets en gebruik het om in te loggen op de Hugging Face CLI. Inloggen helpt ons om de beperkte modellen en datasets te downloaden. 

from huggingface_hub import login
from kaggle_secrets import UserSecretsClient
user_secrets = UserSecretsClient()

hf_token = user_secrets.get_secret("HUGGINGFACE_TOKEN")
login(token = hf_token)

4. Laad de tokenizer met de Gemma 2-modelrepositorylink op Hugging Face. 
5. Laad het model met de modelnaam, 4-bit-kwantisatie-instellingen en device map. Het downloadt eerst de volledige modelfiles en laad daarna het model in 4-bit om rekenkracht en geheugen te optimaliseren. 

Opmerking: Het Gemma 2 9B-It-model is groot, en zelfs met 16GB GPU-geheugen kunnen we het volledige model niet laden. Daarom laden we het model in 4-bit kwantisatie. 

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoConfig

modelName = "google/gemma-2-9b-it"

bnbConfig = BitsAndBytesConfig(
    load_in_4bit = True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained(modelName)

model = AutoModelForCausalLM.from_pretrained(
    modelName,
    device_map = "auto",
    quantization_config=bnbConfig
)

6. Maak de prompt met de systeeminstructies en de vraag van de gebruiker.
7. Zet de prompts om in tokens met de tokenizer.
8. Genereer het antwoord en decodeer het vervolgens naar leesbare Engelse tekst.
9. Geef het resultaat weer door de Markdown-stijl om te zetten naar HTML-formaat. 

from IPython.display import Markdown, display

system =  "You are a skilled software architect who consistently creates system designs for various applications."
user = "Design a system with the ASCII diagram for the customer support application."

prompt = f"System: {system} \n User: {user} \n AI: "
    
inputs = tokenizer(prompt, return_tensors='pt', padding=True, truncation=True).to("cuda")

outputs = model.generate(**inputs, max_length=500, num_return_sequences=1)

text = tokenizer.decode(outputs[0], skip_special_tokens=True)

Markdown(text.split("AI:")[1])

Zoals je ziet, levert het Gemma 2-model uitstekend werk. 

Heb je moeite met het draaien van de bovenstaande code? Raadpleeg dan de Kaggle-notebook: Gemma 2 Simple Inference on GPU.

Gemma 2 fijn-afstemmen met LoRA

In deze sectie stemmen we het Gemma 2 9B-It-model fijn op de zorgdataset, die bestaat uit gesprekken tussen patiënten en artsen. We laden het model en de tokenizer, laden de dataset, converteren de dataset, stellen het model in met trainingsargumenten en volgen de modelprestaties met de Weights and Biases-API.  

Wil je begrijpen hoe de theorie achter fijn-afstemmen werkt? Lees dan onze gids An Introductory Guide to Fine-Tuning LLMs.

1. De omgeving instellen

Installeer de benodigde Python-pakketten om het model op de medische dataset te laden, fijn af te stemmen en te evalueren. 

%%capture
%pip install -U transformers 
%pip install -U datasets 
%pip install -U accelerate 
%pip install -U peft 
%pip install -U trl 
%pip install -U bitsandbytes 
%pip install -U wandb

Laad de benodigde Python-pakketten en functies. 

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import (
    LoraConfig,
    PeftModel,
    prepare_model_for_kbit_training,
    get_peft_model,
)
import os, torch, wandb
from datasets import load_dataset
from trl import SFTTrainer, setup_chat_format

Log in op de Hugging Face CLI met de API-sleutel die we via Kaggle Secrets hebben opgeslagen. 

from huggingface_hub import login
from kaggle_secrets import UserSecretsClient
user_secrets = UserSecretsClient()

hf_token = user_secrets.get_secret("HUGGINGFACE_TOKEN")
login(token = hf_token)

Laad de Weights and Biases API-sleutel uit Kaggle Secrets om het project voor prestatiebewaking te starten. 

wb_token = user_secrets.get_secret("wandb")

wandb.login(key=wb_token)
run = wandb.init(
    project='Fine-tune Gemma-2-9b-it on HealthCare Dataset', 
    job_type="training", 
    anonymous="allow"
)

Stel het model- en dataset-ID in zodat we ze van de Hugging Face Hub kunnen laden. Daarnaast moeten we de naam van het fijnafstemde model instellen om een modelrepository op Hugging Face aan te maken en het fijnafstemde model te pushen.

base_model = "google/gemma-2-9b-it"
dataset_name = "lavita/ChatDoctor-HealthCareMagic-100k"
new_model = "Gemma-2-9b-it-chat-doctor"

2. Het model en de tokenizer laden

Stel het datatype en de attention-implementatie in op basis van de GPU. 

if torch.cuda.get_device_capability()[0] >= 8:
    !pip install -qqq flash-attn
    torch_dtype = torch.bfloat16
    attn_implementation = "flash_attention_2"
else:
    torch_dtype = torch.float16
    attn_implementation = "eager"

We moeten de QLoRA-configuratie maken zodat we het model met 4-bit-precisie kunnen laden, wat het geheugenverbruik vermindert en het fijnafstemproces versnelt.

# QLoRA config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch_dtype,
    bnb_4bit_use_double_quant=True,
)

Laad met de model-URL, LoRA-configuratie en attention-implementatie het Gemma 2 9B-It-model en de tokenizer.

# Load model
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=bnb_config,
    device_map="auto",
    attn_implementation=attn_implementation
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)

3. De adapter aan de laag toevoegen

Maak de Python-functie die het model gebruikt en de namen van alle lineaire modules extraheert. 

import bitsandbytes as bnb

def find_all_linear_names(model):
    cls = bnb.nn.Linear4bit
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split('.')
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
    if 'lm_head' in lora_module_names:  # needed for 16 bit
        lora_module_names.remove('lm_head')
    return list(lora_module_names)

modules = find_all_linear_names(model)

Het volledig fijn-afstemmen van het model kost veel tijd, dus om het trainingsproces te versnellen maken we een adapterlaag aan en koppelen we die, wat resulteert in een sneller en geheugenefficiënter proces. 

De adapterlaag wordt gemaakt met de doelmodules en het taaktype. Vervolgens richten we het chatformaat in voor het model en de tokenizer. Ten slotte koppelen we het basismodel aan de adapter om een Parameter Efficient Fine-Tuning (PEFT)-model te maken.

# LoRA config
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=modules
)
model, tokenizer = setup_chat_format(model, tokenizer)
model = get_peft_model(model, peft_config)

4. De dataset laden

We laden nu de lavita/ChatDoctor-HealthCareMagic-100k-dataset van de Hugging Face Hub. De dataset bestaat uit drie kolommen:

1. instruction: Bevat systeeminstructies. 
2. input: Uitgebreide vraag van de patiënt.
3. output: Het antwoord van de arts op de vraag van de patiënt. 

Na het laden van de dataset schudden we deze en selecteren we 1000 samples om de trainingstijd verder te verkorten. Tot slot maken we het chatformaat met de standaard chatt template en gebruiken dat om de “text”-kolom te maken. 

#Importing the dataset
dataset = load_dataset(dataset_name, split="all")
dataset = dataset.shuffle(seed=65).select(range(1000)) # Only use 1000 samples for quick demo

def format_chat_template(row):
    row_json = [{"role": "system", "content": row["instruction"]},
               {"role": "user", "content": row["input"]},
               {"role": "assistant", "content": row["output"]}]
    row["text"] = tokenizer.apply_chat_template(row_json, tokenize=False)
    return row

dataset = dataset.map(
    format_chat_template,
    num_proc= 4,
)

dataset

Output: 

Dataset({
    features: ['instruction', 'input', 'output', 'text'],
    num_rows: 1000
})

Laten we de “text”-kolom van rij 3 bekijken. 

dataset['text'][3]

De “text”-kolom bevat instructies, de vraag van de patiënt en het antwoord van de arts in OpenAI-stijl.

Voor model­evaluatie splitsen we de dataset in een train- en testsplit. 

dataset = dataset.train_test_split(test_size=0.1)

5. Het model configureren en trainen

We stellen nu de trainingsargumenten en SFT-parameters in en starten vervolgens het trainingstraject. 

Je kunt de verschillende hyperparameters aanpassen op basis van je omgeving, rekenkracht en geheugen. De onderstaande hyperparameters zijn geoptimaliseerd voor de Kaggle-notebook. Als je hetzelfde op Google Colab wilt draaien, experimenteer dan met trainingsinstellingen. 

# Setting Hyperparamter
training_arguments = TrainingArguments(
    output_dir=new_model,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=2,
    optim="paged_adamw_32bit",
    num_train_epochs=1,
    eval_strategy="steps",
    eval_steps=0.2,
    logging_steps=1,
    warmup_steps=10,
    logging_strategy="steps",
    learning_rate=2e-4,
    fp16=False,
    bf16=False,
    group_by_length=True,
    report_to="wandb"
)
# Setting sft parameters
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    peft_config=peft_config,
    max_seq_length= 512,
    dataset_text_field="text",
    tokenizer=tokenizer,
    args=training_arguments,
    packing= False,
)

model.config.use_cache = False
trainer.train()

Het fijn-afstemmen duurde bijna 25 minuten en zoals te zien is, daalde de training- en validatieverlies geleidelijk. Voor betere prestaties kun je het model op de volledige dataset minimaal drie epochs fijn-afstemmen.

Fijn-afstemmen kan best technisch zijn voor beginners en niet-technische lezers. Als je een eenvoudigere oplossing zoekt, bekijk dan de tutorial Fine-Tuning OpenAI's GPT-4: A Step-by-Step Guide. Die behandelt een eenvoudige manier om een state-of-the-art model te fijn-afstemmen via de OpenAI API.

6. Modellevaluatie

We ronden het Weights and Biases-experiment af, waarmee een evaluatierapport wordt gegenereerd. 

wandb.finish()
model.config.use_cache = True

Voor het gedetailleerde rapport ga je naar je Weights and Biases-account en klik je op de projectnaam "Fine-tune Gemma-2-9b-it on Healthcare Dataset" project name.

Screenshot is van wandb.ai

7. De wijzigingen opslaan

We slaan nu de fijnafstem-adapter lokaal op en pushen hem ook naar de Hugging Face Hub. 

trainer.model.save_pretrained(new_model)
trainer.model.push_to_hub(new_model, use_temp_dir=False)

Om de opgeslagen adapter in een andere Kaggle-notebook te gebruiken, moeten we de notebook opslaan. Dat kan door rechtsboven op “Save Version” te klikken. Selecteer daarna “Quick Save” en “Always save output when creating a Quick Save”, en klik vervolgens op de knop “Save”. 

Heb je problemen met het uitvoeren van de code, kloon dan de Kaggle-notebook en voer die uit. Je moet de Hugging Face- en Weights & Biases-API-sleutel instellen via Kaggle Secrets.

Om te bepalen of fijn-afstemmen of Retrieval-Augmented Generation (RAG) beter past bij jouw usecase, raad ik je ten zeerste aan om de blog RAG vs Fine-Tuning te lezen. 

Het basismodel met de adapter samenvoegen

We voegen nu de adapter samen met het basismodel en pushen het volledige model naar de Hugging Face Hub.

Maak een nieuwe Kaggle-notebook met CPU als accelerator en installeer de benodigde Python-pakketten.

Waarom gebruiken we de CPU als accelerator? Is dat niet traag? Ja, maar de GPU-machine op Kaggle biedt slechts 16 GB GPU-geheugen, wat genoeg is voor een model met 7 miljard parameters maar niet voor 9 miljard. De CPU-machine biedt intussen 30 GB, wat voldoende is om zowel het basismodel als de adapter te laden.

%%capture
%pip install -U bitsandbytes
%pip install -U transformers
%pip install -U accelerate
%pip install -U peft
%pip install -U trl

Haal de API-sleutel op uit Kaggle Secrets en log in op de Hugging Face CLI. 

from huggingface_hub import login
from kaggle_secrets import UserSecretsClient
user_secrets = UserSecretsClient()

hf_token = user_secrets.get_secret("HUGGINGFACE_TOKEN")
login(token = hf_token)

Om toegang te krijgen tot het opgeslagen model, moet je de opgeslagen Kaggle-notebook importeren. Klik daarvoor rechtsboven op de knop “Add Input” en selecteer het tabblad “Your Work”. Klik daarna op de plusknop van de opgeslagen notebook om toegang te krijgen tot alle bijbehorende bestanden. 

Stel de basis­model-URL in door de naam van de Hugging Face-repository op te geven. Configureer daarnaast de adapter-URL door de lokale map te specificeren waar de adapter is opgeslagen.

base_model_url = "google/gemma-2-9b-it"
new_model_url = "/kaggle/input/fine-tune-gemma-2-9b-it-on-healthcare-dataset/Gemma-2-9b-it-chat-doctor/"

Laad de tokenizer en het volledige model met de basis­model-URL. Zorg dat je het device op “cpu” en dtype op “float16” instelt.

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
from peft import PeftModel
import torch
from trl import setup_chat_format


# Reload tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(base_model_url)

base_model_reload= AutoModelForCausalLM.from_pretrained(
    base_model_url,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map="cpu",
)

Stel het chatformaat in op het zojuist geladen basismodel en combineer het met de adapter. Tot slot laden en mergen we de adapter met het basismodel. 

De functie merge_and_unload() helpt ons om de adaptergewichten samen te voegen met het basismodel en het als standalone model te gebruiken.

base_model_reload, tokenizer = setup_chat_format(base_model_reload, tokenizer)
model = PeftModel.from_pretrained(base_model_reload, new_model_url)

model = model.merge_and_unload()

Sla het volledige model en de tokenizer lokaal op.

model.save_pretrained("Gemma-2-9b-it-chat-doctor")
tokenizer.save_pretrained("Gemma-2-9b-it-chat-doctor")

Push ook alle modelfiles en de tokenizer naar de Hugging Face Hub. 

model.push_to_hub("Gemma-2-9b-it-chat-doctor", use_temp_dir=False)
tokenizer.push_to_hub("Gemma-2-9b-it-chat-doctor", use_temp_dir=False)

Je kunt naar je Hugging Face-modelrepository gaan en alle bestanden bekijken. 

Bron: kingabzpro/Gemma-2-9b-it-chat-doctor · Hugging Face

Loop je tegen problemen aan bij het mergen van modellen met de CPU, raadpleeg dan de Kaggle-notebook.

Kwantiseren met Hugging Face Space

Met de GGUF My Repo-Hugging Face Space zijn modelconversie en -kwantisatie eenvoudig en snel geworden. Je hoeft alleen maar in te loggen en de model-ID op te geven. 

Bron: GGUF My Repo - a Hugging Face Space by ggml-org

Er wordt een nieuwe modelrepository voor je aangemaakt met een gekwantiseerd modelbestand dat je later kunt downloaden om lokaal te gebruiken. Zo simpel is het. 

Wil je de theorie achter kwantisatie leren, lees dan ons artikel Quantization for Large Language Models (LLMs): Reduce AI Model Sizes Efficiently.

Bron: kingabzpro/Gemma-2-9b-it-chat-doctor-Q4_K_M-GGUF · Hugging Face

Maar als je liever zelf aan de slag gaat met llama.cpp-scripts, bekijk dan de tutorial Fine-Tuning Llama 3 and Using It Locally.

Het fijnafstemde model lokaal gebruiken

Om het gekwantiseerde model lokaal te gebruiken:

1. Download en installeer eerst de Jan-applicatie via de officiële website. 
2. Start de applicatie en ga naar de “Model Hub”. Je opent de menu’s door op de Windows- “🪟”-knop linksboven te klikken. 
3. Plak de link naar de fijnafstemde modelrepository in de zoekbalk en druk op Enter: "kingabzpro/Gemma-2-9b-it-chat-doctor-Q4_K_M-GGUF"

4. Je wordt doorgestuurd naar een nieuw menu, waar je het model direct kunt downloaden. 

5. Ga naar het menu “Thread” en selecteer het fijnafstemde model via het tabblad “Model” rechtsboven.  

6. Begin vragen te stellen aan de AI-assistent en krijg een aangepast antwoord. 

7. Zoals je ziet, is de eerste reactie onder de maat. We moeten een paar modelparameters aanpassen. Ga naar het tabblad “Model” en pas de volgende argumenten aan:

• Stop: <endofstring>, Met vriendelijke groet, Chat Doctor.
• Frequency Penalty: 0,5
• Max Tokens:  680

8. Ga naar het tabblad “Assistant” en verander de Instructions in: “Als je een arts bent, beantwoord dan de medische vragen op basis van de beschrijving van de patiënt.”
9. Test het model met meerdere vragen en ervaar resultaten die van hoge kwaliteit en to the point zijn.

Conclusie

Gemma 2 is een sterk model, en het kan nog beter worden als je het Keras 3-framework gebruikt voor gedistribueerd fijn-afstemmen en modelinference. Je krijgt snellere training en inference vergeleken met het gebruik van het Transformers-framework. 

Je kunt leren hoe je het Gemma-model fijn-afstemt en modelinference draait met Keras 3 door de tutorial Fine-Tune and Run Inference on Google's Gemma Model with TPU te volgen.

In deze tutorial hebben we kennisgemaakt met Gemma 2-modellen en hoe je er toegang toe krijgt met de Transformer-bibliotheek. We hebben ook geleerd hoe je het model fijn-afstemt op gesprekken tussen patiënt en arts, de adapter samenvoegt met het basismodel, het volledige model naar de Hugging Face Hub pusht, het model converteert en kwantiseert met Hugging Face Space en het model downloadt en lokaal op je laptop draait met de Jan-applicatie. 

Het is een leuk project dat AI-enthousiastelingen zouden moeten proberen. Het verbetert je vermogen om problemen op te lossen die ontstaan tijdens het fijn-afstemmen van large language models. De belangrijkste uitdagingen hebben vaak te maken met geheugen- en rekenbeperkingen, en het minimaliseren van verlies.

Je kunt ook leren je eigen LLM’s te ontwikkelen met PyTorch en Hugging Face in de cursus Developing Large Language Models op DataCamp.