DiffusionGemma fijn-afstemmen op PubMedQA met Unsloth

1. Open een RunPod Jupyter Notebook

Maak een nieuwe RunPod-pod met een NVIDIA H100-GPU en selecteer een PyTorch/Jupyter-template. Configureer minstens 100 GB persistente opslag zodat je modelbestanden en trainingsoutputs niet verdwijnen als de pod stopt.

Voeg je Hugging Face access token toe als omgevingsvariabele:

HF_TOKEN=your_hugging_face_token

Hiermee kunnen modellen en datasets sneller downloaden en kun je de opgeslagen LoRA-adapter naar Hugging Face uploaden zonder handmatig in te loggen vanuit het notebook.

De geconfigureerde pod kost ongeveer $3 per uur, al kan de uiteindelijke prijs variëren afhankelijk van de GPU-beschikbaarheid en het geselecteerde podtype.

Zodra de pod draait, open je JupyterLab of Jupyter Notebook via de RunPod-interface en maak je een nieuw notebook met de naam diffusiongemma_pubmedqa.ipynb.

2. Installeer de vereiste pakketten

Voer de volgende commando’s uit in de eerste notebookcel om Unsloth en de libraries voor het laden, fijnstemmen en opslaan van DiffusionGemma te installeren.

%%capture
%pip install --upgrade pip wheel setuptools packaging ninja
%pip install unsloth
%pip install --no-deps --upgrade --force-reinstall git+https://github.com/unslothai/unsloth-zoo.git git+https://github.com/unslothai/unsloth.git
%pip install sentencepiece protobuf "datasets==4.3.0" "huggingface_hub>=0.34.0" hf_transfer
%pip install --no-deps bitsandbytes accelerate peft trl triton
%pip install --no-deps --upgrade "torchao>=0.16.0"
%pip install --no-deps transformers==5.11.0 "tokenizers>=0.22.0,<=0.23.0"

Het commando %%capture verbergt de lange installatie-uitvoer. De pakketversies zijn vastgezet om compatibiliteitsproblemen tussen DiffusionGemma, Transformers, Unsloth en de traininglibraries te voorkomen.

Start na de installatie de notebook-kernel opnieuw voordat je verdergaat.

3. Importeer de libraries

Importeer de libraries die nodig zijn voor datapreparatie, modelladen, training en evaluatie.

import copy
import os
import random
import time

import torch
from datasets import load_dataset
from unsloth import FastModel

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
torch._dynamo.config.recompile_limit = 64

print("Torch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print(
    "GPU:",
    torch.cuda.get_device_name(0)
    if torch.cuda.is_available()
    else "None",
)

HF_HUB_ENABLE_HF_TRANSFER schakelt snellere downloads van de Hugging Face Hub in, terwijl het verhogen van de Dynamo-recompilatiegrens helpt om onderbrekingen te voorkomen bij het werken met het model.

Als de cel succesvol draait, patcht Unsloth de trainingsomgeving en zou de uitvoer moeten bevestigen dat CUDA beschikbaar is en dat de H100-GPU is gedetecteerd.

🦥 Unsloth: Will patch your computer to enable 2x faster free fine-tuning.
🦥 Unsloth Zoo will now patch everything to make training faster!
Torch: 2.10.0+cu128
CUDA available: True
GPU: NVIDIA H100 80GB HBM3

4. Stel de configuratie in

Definieer het model, de dataset, de trainingsparameters, de evaluatie-instellingen en de outputmap op één plek.

MODEL_NAME = "unsloth/diffusiongemma-26B-A4B-it"
DATASET_NAME = "qiaojin/PubMedQA"

TRAIN_SUBSET = "pqa_artificial"
EVAL_SUBSET = "pqa_labeled"

N_TRAIN = 3000
N_EVAL = 200

MAX_CONTEXT_CHARS = 2500

STEPS = 60
GRAD_ACCUM = 4
LR = 1e-4
T_LO = 0.1

EVAL_TOTAL = 50
EVAL_DENOISING_STEPS = 16

OUTPUT_DIR = "diffusiongemma_pubmedqa_lora"

We gebruiken 3.000 kunstmatige voorbeelden voor training en 200 handmatig gelabelde voorbeelden voor evaluatie. Om het experiment snel te houden, traint het model 60 stappen en evalueert het 50 voorbeelden met 16 denoising-stappen.

Hoe DiffusionGemma werkt

Voordat we het model laden, is het de moeite waard om je voor te stellen hoe DiffusionGemma daadwerkelijk een antwoord produceert. In plaats van tokens één voor één te schrijven, begint het met een canvas met vaste lengte van tokens en verfijnt die over verschillende denoising-stappen, waarbij veel posities tegelijk worden geüpdatet totdat de tekst uitkristalliseert tot een coherent antwoord. 

De onderstaande code rapporteert een canvaslengte van 256, de grootte van één blok. Voor onze korte yes/no/maybe-antwoorden is één canvas meer dan voldoende, terwijl langere outputs blok voor blok worden gegenereerd door canvassen aaneen te schakelen. Het onderstaande diagram toont dit verfijningsproces op hoofdlijnen:

5. Laad DiffusionGemma

Laad het instruction-tuned DiffusionGemma-model in bfloat16-precisie. We gebruiken geen 4-bit quantization omdat de H100-GPU genoeg geheugen heeft om het model in hogere precisie te laden.

model, tokenizer = FastModel.from_pretrained(
    model_name=MODEL_NAME,
    dtype=torch.bfloat16,
    load_in_4bit=False,
)

processor = tokenizer
tok = processor.tokenizer if hasattr(processor, "tokenizer") else processor

vocab = model.config.text_config.vocab_size
canvas_len = model.config.canvas_length

dev = next(
    (p.device for p in model.parameters() if p.device.type != "meta"),
    torch.device("cuda"),
)

print("Vocab size:", vocab)
print("Canvas length:", canvas_len)
print("Model device:", dev)

De vocabulairegrootte wordt gebruikt bij het toevoegen van willekeurige ruis tijdens diffusion-training. De canvaslengte bepaalt het maximum aantal tokens dat het model in één generatieblok kan verfijnen.

Je zou een uitvoer moeten zien zoals:

Vocab size: 262144
Canvas length: 256
Model device: cuda:0

6. Voeg een LoRA-adapter toe

Voeg een LoRA-adapter toe zodat slechts een kleine set extra parameters wordt getraind in plaats van het volledige model met 26 miljard parameters te updaten.

model = FastModel.get_peft_model(
    model,
    r=64,
    lora_alpha=128,
    use_gradient_checkpointing=False,
)

Dit vermindert het geheugen- en rekkenverbruik voor fijnstemmen aanzienlijk. Gradient checkpointing is uitgeschakeld omdat de H100 voor dit experiment voldoende GPU-geheugen heeft.

7. Laad PubMedQA

Laad de kunstmatige PubMedQA-subset voor training en de handmatig gelabelde subset voor evaluatie.

train_data = load_dataset(
    DATASET_NAME,
    TRAIN_SUBSET,
    split="train",
)

eval_data = load_dataset(
    DATASET_NAME,
    EVAL_SUBSET,
    split="train",
)

print("Train size:", len(train_data))
print("Eval size:", len(eval_data))
print(train_data[0])

De trainingssubset bevat automatisch gegenereerde voorbeelden, terwijl de evaluatiesubset expert-gelabelde biomedische vragen bevat. 

Door de eerste rij te printen kunnen we de vraag, abstracte context en uiteindelijke beslissing inspecteren voordat we de data formatteren.

Je zou dit moeten zien:

Train size: 211269
Eval size: 1000

Elk voorbeeld bevat een biomedische vraag, één of meer ondersteunende abstractpassages en een eindantwoord van yes, no of maybe.

8. Converteer de dataset

Zet elk PubMedQA-voorbeeld om naar een chatstijl-formaat met een gebruikersprompt en een assistentantwoord.

def make_prompt(row):
    context = " ".join(row["context"]["contexts"])
    context = context[:MAX_CONTEXT_CHARS]
    question = row["question"]

    return f"""Answer the biomedical research question using only the context.

Context:
{context}

Question:
{question}

Answer with only one word: yes, no, or maybe."""


def make_answer(row):
    return row["final_decision"].strip().lower()


def convert_row(row):
    answer = make_answer(row)

    if answer not in ["yes", "no", "maybe"]:
        return None

    return {
        "messages": [
            {"role": "user", "content": make_prompt(row)},
            {"role": "assistant", "content": answer},
        ]
    }


train_rows = []

for row in train_data.select(range(N_TRAIN)):
    item = convert_row(row)

    if item is not None:
        train_rows.append(item)


eval_rows = []

for row in eval_data.select(range(N_EVAL)):
    item = convert_row(row)

    if item is not None:
        eval_rows.append(item)


print("Prepared train examples:", len(train_rows))
print("Prepared eval examples:", len(eval_rows))
print(train_rows[0]["messages"][0]["content"])
print("Answer:", train_rows[0]["messages"][1]["content"])

De contextpassages worden samengevoegd tot één string en beperkt tot 2.500 tekens om de input beheersbaar te houden. Elk antwoord wordt omgezet naar kleine letters en voorbeelden met labels buiten yes, no of maybe worden verwijderd.

Door het eerste geconverteerde voorbeeld te printen kun je bevestigen dat de context, vraag en het antwoord correct zijn geformatteerd voordat je gaat trainen.

9. Bouw de diffusion-trainingsvoorbeelden

DiffusionGemma vereist dat het doelantwoord in een canvas met vaste lengte wordt geplaatst. Deze functie tokenizeert de prompt, zet het antwoord om naar token-ID’s, padt het tot de canvaslengte van het model en maakt een masker dat aangeeft welke tokens aan de loss moeten bijdragen.

eos = model.generation_config.eos_token_id or [1]
eos = eos[0] if isinstance(eos, (list, tuple)) else eos

pad = tok.pad_token_id if tok.pad_token_id is not None else eos


def build_examples(rows):
    examples = []

    for row in rows:
        user_message = row["messages"][0]
        assistant_message = row["messages"][1]

        prompt_ids = processor.apply_chat_template(
            [user_message],
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
        )[0]

        answer_ids = tok.encode(
            assistant_message["content"],
            add_special_tokens=False,
        )

        content = answer_ids + [eos]
        n = len(content)

        if n > canvas_len:
            continue

        x0 = torch.tensor(
            content + [pad] * (canvas_len - n),
            dtype=torch.long,
        )

        loss_mask = torch.zeros(canvas_len, dtype=torch.bool)
        loss_mask[:n] = True

        examples.append((prompt_ids, x0, loss_mask))

    return examples


examples = build_examples(train_rows)

De end-of-sequence-token wordt na elk antwoord toegevoegd, terwijl de resterende canvasposities worden gevuld met padding-tokens. Het loss-masker zorgt ervoor dat training zich richt op de antwoord- en end-of-sequence-tokens en niet op de gepadde posities.

10. Maak de inferentie- en evaluatiefuncties

Definieer vervolgens de functies om antwoorden te genereren, de uitvoer van het model op te schonen en de evaluatienauwkeurigheid te berekenen.

Genereer een antwoord

De functie answer_question() formatteert de prompt, genereert een respons via meerdere denoising-stappen en decodeert de gegenereerde tokens naar tekst.

def answer_question(prompt, steps=64):
    input_ids = processor.apply_chat_template(
        [{"role": "user", "content": prompt}],
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to(dev)

    gen_config = copy.deepcopy(model.generation_config)
    gen_config.max_denoising_steps = steps
    gen_config.max_new_tokens = canvas_len

    model.eval()

    with torch.no_grad():
        output = model.generate(
            input_ids=input_ids,
            generation_config=gen_config,
        )

    generated = output.sequences[0, input_ids.shape[1]:]
    text = tok.decode(
        generated.tolist(),
        skip_special_tokens=True,
    )

    return text.strip().lower()

Haal de voorspelling eruit

Hoewel de prompt om een eenwoordig antwoord vraagt, kan het model soms extra tekst genereren. Deze functie haalt de eerste geldige yes, no of maybe-voorspelling eruit.

def clean_prediction(text):
    text = text.lower().strip()

    if text.startswith("yes"):
        return "yes"

    if text.startswith("no"):
        return "no"

    if text.startswith("maybe"):
        return "maybe"

    words = text.replace(".", " ").replace(",", " ").split()

    for word in words:
        if word in ["yes", "no", "maybe"]:
            return word

    return "unknown"

Evalueer de nauwkeurigheid

De evaluatiefunctie vergelijkt elke opgeschoonde voorspelling met het juiste antwoord, print het resultaat voor elk voorbeeld en retourneert de algehele nauwkeurigheid samen met de individuele voorspellingen.

def evaluate_model(
    rows,
    total=50,
    steps=64,
    title="Evaluation",
):
    correct = 0
    results = []
    total = min(total, len(rows))

    print(title)
    print("-" * len(title))

    for i, row in enumerate(rows[:total], start=1):
        prompt = row["messages"][0]["content"]
        gold = row["messages"][1]["content"]

        raw_pred = answer_question(prompt, steps=steps)
        pred = clean_prediction(raw_pred)

        is_correct = pred == gold
        correct += int(is_correct)

        results.append({
            "index": i,
            "gold": gold,
            "prediction": pred,
            "raw_prediction": raw_pred,
            "correct": is_correct,
        })

        print(
            f"{i:02d}. Gold: {gold} | "
            f"Pred: {pred} | Correct: {is_correct}"
        )

    accuracy = correct / total if total else 0

    print()
    print("Accuracy:", accuracy)
    print()

    return {
        "accuracy": accuracy,
        "correct": correct,
        "total": total,
        "results": results,
    }

11. Evalueer het model vóór het fijnstemmen

Voer de evaluatie uit vóór de training om een basislijn vast te stellen.

before_eval = evaluate_model(
    eval_rows,
    total=EVAL_TOTAL,
    steps=EVAL_DENOISING_STEPS,
    title="Before Fine-Tuning Evaluation",
)

Dit evalueert 50 voorbeelden met 16 denoising-stappen per antwoord. In dit experiment beantwoordde het basismodel 30 van de 50 vragen correct.

Deze basislijn vergelijken we later met de nauwkeurigheid van het model na het fijnstemmen.

12. Stel de training in

Zet het model in trainingsmodus, maak de optimizer en learning-rate-scheduler aan, en definieer hoe schone antwoordtokens worden gecorrumpeerd tijdens diffusion-training.

model.config.use_cache = True
model.train()

opt = torch.optim.AdamW(
    [p for p in model.parameters() if p.requires_grad],
    lr=LR,
    betas=(0.9, 0.95),
    weight_decay=0.0,
)

sched = torch.optim.lr_scheduler.OneCycleLR(
    opt,
    max_lr=LR,
    total_steps=STEPS,
    pct_start=0.03,
    anneal_strategy="cos",
)

Alleen parameters met requires_grad=True worden aan de optimizer doorgegeven, wat betekent dat het trainingsproces de LoRA-adapter bijwerkt en niet het volledige model.

Maak vervolgens een corruptiefunctie die een willekeurig deel van het antwoordcanvas vervangt door willekeurige tokens.

def corrupt(x0):
    noise_level = random.uniform(T_LO, 1.0)

    xt = x0.to(dev).clone()

    noise_mask = (
        torch.rand(canvas_len, device=dev) < noise_level
    )

    xt[noise_mask] = torch.randint(
        0,
        vocab,
        (canvas_len,),
        device=dev,
    )[noise_mask]

    return xt.unsqueeze(0)

De hoeveelheid ruis verandert per voorbeeld. Tijdens de training leert het model het originele antwoord te reconstrueren uit deze gecorrumpeerde canvastokens.

13. Train het model

De volgende lus traint de LoRA-adapter 60 stappen met gebruik van gradient accumulation.

order = list(range(len(examples)))
ptr = 0
start_time = time.time()

opt.zero_grad(set_to_none=True)

for step in range(1, STEPS + 1):
    step_loss = 0.0

    for _ in range(GRAD_ACCUM):
        if ptr >= len(order):
            random.shuffle(order)
            ptr = 0

        prompt_ids, x0, loss_mask = examples[order[ptr]]
        ptr += 1

        output = model(
            input_ids=prompt_ids.unsqueeze(0).to(dev),
            canvas_ids=corrupt(x0),
            self_conditioning_logits=None,
        )

        logits = output.logits[0].float()
        mask = loss_mask.to(dev)

        loss = torch.nn.functional.cross_entropy(
            logits[mask],
            x0.to(dev)[mask],
        )

        (loss / GRAD_ACCUM).backward()
        step_loss += loss.item() / GRAD_ACCUM

    torch.nn.utils.clip_grad_norm_(
        [
            p
            for p in model.parameters()
            if p.requires_grad
        ],
        1.0,
    )

    opt.step()
    sched.step()
    opt.zero_grad(set_to_none=True)

    if step % 20 == 0:
        elapsed = time.time() - start_time

        print(
            f"step {step}/{STEPS} | "
            f"loss {step_loss:.4f} | "
            f"{elapsed:.0f}s"
        )

Voor elk trainingsvoorbeeld ontvangt het model de biomedische prompt en een gecorrumpeerd antwoordcanvas. Cross-entropy-loss wordt alleen berekend voor de echte antwoordtokens die door het loss-masker zijn geselecteerd.

Gradient accumulation combineert vier voorbeelden voordat het model wordt geüpdatet. Gradient clipping wordt ook toegepast om de training stabiel te houden.

In dit experiment was de training in ongeveer twee minuten voltooid:

step 20/60 | loss 0.0019 | 43s
step 40/60 | loss 0.0003 | 85s
step 60/60 | loss 0.0001 | 126s

De gestaag dalende loss geeft aan dat de adapter leert de verwachte antwoorden te reconstrueren uit het gecorrumpeerde canvas. Tijdens de training kun je ook nvidia-smi draaien in de RunPod-terminal om het GPU-geheugenverbruik en de benutting te monitoren.

14. Evalueer het fijngestemde model

Voer na de training dezelfde evaluatie opnieuw uit om te meten of het fijnstemmen de prestaties van het model heeft verbeterd.

after_eval = evaluate_model(
    eval_rows,
    total=EVAL_TOTAL,
    steps=EVAL_DENOISING_STEPS,
    title="After Fine-Tuning Evaluation",
)

Het fijngestemde model wordt op dezelfde 50 voorbeelden geëvalueerd en met dezelfde 16 denoising-stappen als in de basisevaluatie.

Vergelijk daarna de nauwkeurigheid vóór en na het fijnstemmen.

before_accuracy = before_eval["accuracy"]
after_accuracy = after_eval["accuracy"]
improvement = after_accuracy - before_accuracy

print("Before fine-tuning accuracy:", before_accuracy)
print("After fine-tuning accuracy:", after_accuracy)
print("Improvement:", improvement)

Before fine-tuning accuracy: 0.6
After fine-tuning accuracy: 0.8
Improvement: 0.2

In dit experiment steeg de nauwkeurigheid van het model van 0,60 naar 0,80.

Dit is een verbetering van 20 procentpunt: het model beantwoordde 40 van de 50 vragen correct na het fijnstemmen, vergeleken met 30 van de 50 vóór de training.

15. Sla de fijngestemde adapter op en upload die

Sla de getrainde LoRA-adapter en processorbestanden op in de eerder gedefinieerde outputmap.

model.save_pretrained(OUTPUT_DIR)
processor.save_pretrained(OUTPUT_DIR)

print(f"Saved LoRA adapter to: {OUTPUT_DIR}")

Je zou dit moeten zien:

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

Hiermee sla je alleen de lichtgewicht LoRA-adapter op en niet nog een volledige kopie van het basismodel met 26 miljard parameters.

Upload daarna de adapter- en processorbestanden naar de Hugging Face Hub:

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

model.push_to_hub(REPO_ID)
processor.push_to_hub(REPO_ID)

Omdat de omgevingsvariabele HF_TOKEN tijdens de RunPod-podconfiguratie is toegevoegd, zou Hugging Face automatisch moeten authenticeren. Je hoeft alleen notebook_login() uit te voeren als het token nog niet is geconfigureerd:

from huggingface_hub import notebook_login

notebook_login()

Na het uploaden bevat de repository de LoRA-adapter en de processorconfiguratie die nodig zijn om het fijngestemde model later te laden.

Bron: kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

Tot slot

DiffusionGemma fijnstemmen met Unsloth was verrassend eenvoudig. Het meeste werk zat in het installeren van de juiste dependencies en begrijpen hoe het specifieke diffusion-trainingsproces werkt. Zodra de omgeving was ingesteld, verliepen het model laden, de LoRA-adapter trainen, evalueren en de resultaten naar Hugging Face uploaden heel soepel.

Ik vond DiffusionGemma vooral interessant omdat het geen tekst token-voor-token genereert zoals een traditioneel taalmodel. In plaats daarvan werkt het met een vast canvas en verfijnt het noisy tokens geleidelijk via denoising-stappen. Dit andere generatieproces begrijpen en het fijnstemmen op een biomedische vraag-antwoordtaak maakte het experiment voor mij extra waardevol.

Zelfs met een kleine opzet steeg de nauwkeurigheid op de steekproef van 50 voorbeelden van 0,60 naar 0,80—al is bij zo’n steekproefgrootte de foutmarge groot, en scoort een "altijd yes"-baseline al rond de 55% op deze evaluatieset. 

Het is ook goed om te weten dat de kunstmatige trainingssubset bijna geen "maybe"-labels bevat, waardoor het model die klasse nauwelijks kan leren, ook al komt die in de evaluatiedata voor. Zie dit dus als een snel experiment in hoe het fijnstemmingsproces werkt, niet als bewijs dat het model klaar is voor echt medisch gebruik.

Klaar om verder te gaan dan één fijne afstemming? Onze track Developing Large Language Models brengt je van PyTorch- en transformers-basics naar het bouwen en deployen van je eigen LLM’s.