DiffusionGemma mit Unsloth auf PubMedQA feinabstimmen

1. Öffne ein RunPod Jupyter Notebook

Erstelle ein neues RunPod-Pod mit einer NVIDIA H100-GPU und wähle ein PyTorch/Jupyter-Template. Konfiguriere mindestens 100 GB persistenten Speicher, damit Modell-Dateien und Trainingsergebnisse beim Stoppen des Pods nicht verloren gehen.

Füge dein Hugging Face Access Token als Umgebungsvariable hinzu:

HF_TOKEN=your_hugging_face_token

So lassen sich Modelle und Datensätze schneller herunterladen und du kannst den gespeicherten LoRA-Adapter direkt aus dem Notebook zu Hugging Face hochladen, ohne dich manuell anzumelden.

Das konfigurierte Pod kostet etwa 3 $ pro Stunde, der Endpreis kann je nach GPU-Verfügbarkeit und gewähltem Pod-Typ variieren.

Sobald das Pod läuft, öffne JupyterLab oder Jupyter Notebook über die RunPod-Oberfläche und erstelle ein neues Notebook namens diffusiongemma_pubmedqa.ipynb.

2. Installiere die benötigten Pakete

Führe die folgenden Befehle in der ersten Notebook-Zelle aus, um Unsloth sowie die Bibliotheken zu installieren, die zum Laden, Fine-Tunen und Speichern von DiffusionGemma benötigt werden.

%%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"

Der Befehl %%capture unterdrückt die langen Installationsausgaben. Die Paketversionen sind fixiert, um Kompatibilitätsprobleme zwischen DiffusionGemma, Transformers, Unsloth und den Trainingsbibliotheken zu vermeiden.

Starte nach Abschluss der Installation den Notebook-Kernel neu, bevor du weitermachst.

3. Bibliotheken importieren

Importiere die Bibliotheken für Datenaufbereitung, Modell-Laden, Training und Evaluation.

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 aktiviert schnellere Downloads aus dem Hugging Face Hub. Die höhere Dynamo-Recompile-Grenze hilft, Unterbrechungen bei der Modellausführung zu vermeiden.

Wenn die Zelle erfolgreich läuft, patcht Unsloth die Trainingsumgebung. Die Ausgabe sollte bestätigen, dass CUDA verfügbar ist und die H100-GPU erkannt wurde.

🦥 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. Konfiguration festlegen

Lege Modell, Datensatz, Trainingsparameter, Evaluations-Settings und Ausgabeverzeichnis zentral fest.

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"

Wir verwenden 3.000 künstlich erzeugte Beispiele fürs Training und 200 manuell gelabelte Beispiele für die Evaluation. Um das Experiment zügig zu halten, trainiert das Modell 60 Schritte und evaluiert 50 Beispiele mit 16 Denoising-Schritten.

So funktioniert DiffusionGemma

Bevor wir das Modell laden, lohnt sich ein Bild davon, wie DiffusionGemma tatsächlich antwortet. Statt Token nacheinander zu schreiben, startet es mit einer Canvas fester Länge und verfeinert sie über mehrere Denoising-Schritte. Viele Positionen werden gleichzeitig aktualisiert, bis ein stimmiger Text entsteht. 

Der folgende Code meldet eine Canvas-Länge von 256, die Größe eines Blocks. Für unsere kurzen yes/no/maybe-Antworten reicht eine Canvas locker aus. Längere Ausgaben entstehen, indem mehrere Canvas-Blöcke hintereinandergeschaltet werden. Das Diagramm zeigt den Verfeinerungsprozess auf hoher Ebene:

5. DiffusionGemma laden

Lade das instruktionsoptimierte DiffusionGemma-Modell in bfloat16-Präzision. Auf 4-Bit-Quantisierung verzichten wir, da die H100-GPU genug Speicher für höhere Präzision bietet.

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)

Die Vokabulargröße wird beim Hinzufügen von Zufallsrauschen während des Diffusionstrainings verwendet. Die Canvas-Länge bestimmt, wie viele Token das Modell in einem Generierungsblock maximal verfeinern kann.

Die Ausgabe sollte in etwa so aussehen:

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

6. Einen LoRA-Adapter hinzufügen

Füge einen LoRA-Adapter hinzu, sodass nur ein kleiner zusätzlicher Parametersatz trainiert wird, statt das gesamte 26-Milliarden-Parameter-Modell zu aktualisieren.

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

Das reduziert Speicher- und Rechenbedarf für das Fine-Tuning erheblich. Gradient Checkpointing ist deaktiviert, da die H100 für dieses Experiment genügend GPU-Speicher hat.

7. PubMedQA laden

Lade das künstliche PubMedQA-Subset fürs Training und das manuell gelabelte Subset für die Evaluation.

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])

Das Trainings-Subset enthält automatisch generierte Beispiele, während das Evaluations-Subset von Expertinnen und Experten gelabelte biomedizinische Fragen umfasst. 

Durch das Ausgeben der ersten Zeile prüfen wir Frage, Abstract-Kontext und finale Entscheidung, bevor wir die Daten formatieren.

Du solltest sehen:

Train size: 211269
Eval size: 1000

Jedes Beispiel enthält eine biomedizinische Frage, einen oder mehrere unterstützende Abstract-Passagen und eine finale Antwort yes, no oder maybe.

8. Datensatz konvertieren

Wandle jedes PubMedQA-Beispiel in ein Chat-Format mit User-Prompt und Assistant-Antwort um.

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

Die Kontextpassagen werden zu einem String zusammengeführt und auf 2.500 Zeichen begrenzt, um den Input handhabbar zu halten. Jede Antwort wird kleingeschrieben. Beispiele mit Labels außerhalb von yes, no oder maybe werden entfernt.

Das Ausgeben des ersten konvertierten Beispiels bestätigt, dass Kontext, Frage und Antwort korrekt formatiert sind, bevor das Training startet.

9. Diffusions-Trainingsexemplare bauen

DiffusionGemma erwartet die Zielantwort innerhalb einer Canvas fester Länge. Diese Funktion tokenisiert den Prompt, wandelt die Antwort in Token-IDs um, polstert auf die Canvas-Länge auf und erzeugt eine Maske, die anzeigt, welche Token zum Loss beitragen sollen.

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)

Nach jeder Antwort wird das End-of-Sequence-Token angehängt, freie Canvas-Positionen werden mit Padding-Token gefüllt. Die Loss-Maske stellt sicher, dass das Training nur auf Antwort- und EOS-Token fokussiert, nicht auf Padding.

10. Inferenz- und Evaluationsfunktionen erstellen

Definiere als Nächstes die Funktionen zum Generieren von Antworten, Bereinigen der Modellausgaben und Berechnen der Genauigkeit.

Eine Antwort generieren

Die Funktion answer_question() formatiert den Prompt, erzeugt eine Antwort über mehrere Denoising-Schritte und dekodiert die generierten Token zu Text.

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

Die Vorhersage extrahieren

Auch wenn der Prompt eine Ein-Wort-Antwort fordert, kann das Modell gelegentlich Zusatztext erzeugen. Diese Funktion extrahiert die erste gültige Vorhersage yes, no oder maybe.

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"

Genauigkeit evaluieren

Die Evaluationsfunktion vergleicht jede bereinigte Vorhersage mit der korrekten Antwort, druckt das Ergebnis je Beispiel und liefert die Gesamtgenauigkeit sowie die Einzelergebnisse zurück.

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. Modell vor dem Fine-Tuning evaluieren

Führe die Evaluation vor dem Training aus, um eine Basislinie zu erhalten.

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

Es werden 50 Beispiele mit 16 Denoising-Schritten pro Antwort evaluiert. In diesem Experiment beantwortete das Basismodell 30 von 50 Fragen korrekt.

Diese Basislinie wird später mit der Genauigkeit nach dem Fine-Tuning verglichen.

12. Training einrichten

Schalte das Modell in den Trainingsmodus, erstelle Optimierer und Learning-Rate-Scheduler und definiere, wie saubere Antwort-Token während des Diffusionstrainings korruptiert werden.

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

Nur Parameter mit requires_grad=True werden an den Optimierer übergeben. Damit wird der LoRA-Adapter trainiert, nicht das komplette Modell.

Als Nächstes erzeugen wir eine Korruptionsfunktion, die einen zufälligen Anteil der Antwort-Canvas durch zufällige Token ersetzt.

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)

Die Rauschmenge variiert je Beispiel. Beim Training lernt das Modell, die Originalantwort aus der korrupten Canvas zu rekonstruieren.

13. Modell trainieren

Die folgende Schleife trainiert den LoRA-Adapter 60 Schritte lang mit 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"
        )

Für jedes Trainingsexemplar erhält das Modell den biomedizinischen Prompt und eine korrupte Antwort-Canvas. Der Cross-Entropy-Loss wird nur für die echten Antwort-Token berechnet, die die Loss-Maske auswählt.

Gradient Accumulation fasst vier Beispiele zusammen, bevor das Modell aktualisiert wird. Gradient Clipping stabilisiert zusätzlich das Training.

In diesem Experiment dauerte das Training rund zwei Minuten:

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

Der kontinuierlich sinkende Loss zeigt, dass der Adapter lernt, die erwarteten Antworten aus der korrupten Canvas zu rekonstruieren. Während des Trainings kannst du im RunPod-Terminal auch nvidia-smi ausführen, um GPU-Speicher und Auslastung zu überwachen.

14. Das feinabgestimmte Modell evaluieren

Führe nach dem Training dieselbe Evaluation erneut aus, um zu messen, ob das Fine-Tuning die Leistung verbessert hat.

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

Das feinabgestimmte Modell wird auf denselben 50 Beispielen und mit denselben 16 Denoising-Schritten evaluiert wie in der Basislinie.

Vergleiche anschließend die Genauigkeit vor und nach dem Fine-Tuning.

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 diesem Experiment stieg die Genauigkeit des Modells von 0,60 auf 0,80.

Das entspricht einem Plus von 20 Prozentpunkten: Nach dem Fine-Tuning beantwortete das Modell 40 von 50 Fragen korrekt, zuvor 30 von 50.

15. Den feinabgestimmten Adapter speichern und hochladen

Speichere den trainierten LoRA-Adapter und die Processor-Dateien im zuvor definierten Ausgabeverzeichnis.

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

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

Du solltest sehen:

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

Es wird nur der schlanke LoRA-Adapter gespeichert, nicht eine weitere Kopie des 26-Milliarden-Parameter-Basemodells.

Lade anschließend Adapter und Processor-Dateien in den Hugging Face Hub hoch:

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

Da die Umgebungsvariable HF_TOKEN bei der RunPod-Konfiguration gesetzt wurde, sollte Hugging Face automatisch authentifizieren. Du musst nur notebook_login() ausführen, wenn noch kein Token konfiguriert ist:

from huggingface_hub import notebook_login

notebook_login()

Nach dem Upload enthält das Repository den LoRA-Adapter und die Processor-Konfiguration, um das feinabgestimmte Modell später zu laden.

Quelle: kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

Fazit

Das Fine-Tuning von DiffusionGemma mit Unsloth war überraschend unkompliziert. Am meisten Zeit kostete die richtige Abhängigkeitsinstallation und das Verständnis des diffusionsspezifischen Trainingsprozesses. War die Umgebung erst eingerichtet, liefen Modell laden, LoRA-Training, Evaluation und Upload zu Hugging Face sehr reibungslos.

Spannend an DiffusionGemma finde ich vor allem, dass es nicht wie ein klassisches Sprachmodell Token für Token generiert, sondern mit einer festen Canvas arbeitet und verrauschte Token über Denoising-Schritte verfeinert. Dieses andere Generationsprinzip zu verstehen und auf eine biomedizinische QA-Aufgabe zu fine-tunen, war für mich besonders lehrreich.

Selbst mit kleinem Setup stieg die Genauigkeit im 50er-Sample von 0,60 auf 0,80—wobei bei so kleiner Stichprobe der Fehlerbereich groß ist und eine „immer yes“-Baseline in diesem Evaluationsset bereits rund 55 % erreicht. 

Beachte auch: Das künstliche Trainings-Subset enthält fast keine „maybe“-Labels, das Modell kann diese Klasse also kaum lernen, obwohl sie in den Evaluationsdaten vorkommt. Sieh das Ganze als schnelles Experiment zum Fine-Tuning-Prozess, nicht als Beleg dafür, dass das Modell für echte medizinische Nutzung bereit ist.

Bereit, über ein einzelnes Fine-Tuning hinauszugehen? Unser Developing Large Language Models-Lernpfad führt dich von PyTorch- und Transformers-Grundlagen bis zum Bauen und Deployen eigener LLMs.