Jak dostroić DiffusionGemma na PubMedQA za pomocą Unsloth

1. Otwórz notatnik Jupyter na RunPod

Utwórz nowy pod RunPod z GPU NVIDIA H100 i wybierz szablon PyTorch/Jupyter. Skonfiguruj co najmniej 100 GB trwałej pamięci, aby pliki modelu i wyniki trenowania nie zginęły po zatrzymaniu poda.

Dodaj swój token dostępu Hugging Face jako zmienną środowiskową:

HF_TOKEN=your_hugging_face_token

Dzięki temu modele i zbiory danych będą pobierać się szybciej, a zapisany adapter LoRA wyślesz na Hugging Face bez ręcznego logowania z notatnika.

Skonfigurowany pod powinien kosztować około 3 USD na godzinę, choć ostateczna cena może się różnić w zależności od dostępności GPU i wybranego typu poda.

Gdy pod będzie działać, otwórz JupyterLab lub Jupyter Notebook z interfejsu RunPod i utwórz nowy notatnik o nazwie diffusiongemma_pubmedqa.ipynb.

2. Zainstaluj wymagane pakiety

Uruchom poniższe polecenia w pierwszej komórce notatnika, aby zainstalować Unsloth oraz biblioteki potrzebne do wczytywania, dostrajania i zapisywania DiffusionGemma.

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

Polecenie %%capture ukrywa długi wynik instalacji. Wersje pakietów są przypięte, aby uniknąć problemów ze zgodnością między DiffusionGemma, Transformers, Unsloth i bibliotekami treningowymi.

Po zakończeniu instalacji zrestartuj jądro notatnika, zanim przejdziesz dalej.

3. Zaimportuj biblioteki

Zaimportuj biblioteki potrzebne do przygotowania zbioru danych, wczytania modelu, trenowania i ewaluacji.

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 włącza szybsze pobieranie z Hugging Face Hub, a zwiększenie limitu rekompilacji Dynamo pomaga uniknąć przerw podczas pracy z modelem.

Jeśli komórka wykona się poprawnie, Unsloth załata środowisko treningowe, a w wyjściu zobaczysz potwierdzenie dostępności CUDA i wykrycia GPU H100.

🦥 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. Ustaw konfigurację

Zdefiniuj w jednym miejscu model, zbiór danych, parametry treningu, ustawienia ewaluacji i katalog wyjściowy.

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"

Użyjemy 3000 sztucznych przykładów do treningu i 200 ręcznie oznaczonych przykładów do ewaluacji. Aby eksperyment był szybki, model będzie trenowany przez 60 kroków i oceni 50 przykładów, używając 16 kroków odszumiania.

Jak działa DiffusionGemma

Zanim wczytamy model, warto wyobrazić sobie, jak DiffusionGemma faktycznie tworzy odpowiedź. Zamiast pisać tokeny jeden po drugim, zaczyna od płótna o stałej długości i udoskonala je w kilku krokach odszumiania, aktualizując wiele pozycji naraz, aż tekst przekształci się w spójną odpowiedź. 

Poniższy kod raportuje długość płótna równą 256, co odpowiada rozmiarowi jednego bloku. Dla naszych krótkich odpowiedzi yes/no/maybe jedno płótno w zupełności wystarcza, natomiast dłuższe wyjścia generowane są przez łączenie płócien blok po bloku. Poniższy diagram pokazuje ten proces udoskonalania na wysokim poziomie:

5. Wczytaj DiffusionGemma

Wczytaj dostrojony instrukcjami model DiffusionGemma w precyzji bfloat16. Nie będziemy używać kwantyzacji 4-bitowej, ponieważ GPU H100 ma wystarczająco pamięci, by załadować model w wyższej precyzji.

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)

Rozmiar słownika jest używany podczas dodawania losowego szumu w trakcie treningu dyfuzyjnego. Długość płótna określa maksymalną liczbę tokenów, które model może udoskonalać w pojedynczym bloku generacji.

Powinieneś zobaczyć wyjście podobne do:

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

6. Dodaj adapter LoRA

Dodaj adapter LoRA, aby trenować jedynie niewielki zestaw dodatkowych parametrów zamiast aktualizować cały 26‑miliardowy model.

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

To znacząco zmniejsza zapotrzebowanie na pamięć i obliczenia podczas dostrajania. Checkpointing gradientów jest wyłączony, bo H100 ma wystarczającą pamięć GPU do tego eksperymentu.

7. Wczytaj PubMedQA

Wczytaj sztuczny podzbiór PubMedQA do treningu i ręcznie oznaczony podzbiór do ewaluacji.

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

Podzbiór treningowy zawiera automatycznie wygenerowane przykłady, a podzbiór ewaluacyjny — pytania biomedyczne oznaczone przez ekspertów. 

Wydruk pierwszego wiersza pozwala przejrzeć pytanie, kontekst z abstraktu i końcową decyzję przed formatowaniem danych.

Powinieneś zobaczyć:

Train size: 211269
Eval size: 1000

Każdy przykład zawiera pytanie biomedyczne, jeden lub więcej fragmentów abstraktu jako wsparcie oraz końcową odpowiedź: yes, no lub maybe.

8. Przekonwertuj zbiór danych

Przekonwertuj każdy przykład PubMedQA do formatu czatu zawierającego podpowiedź użytkownika i odpowiedź asystenta.

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

Fragmenty kontekstu są łączone w jeden ciąg i ograniczane do 2500 znaków, aby utrzymać rozsądny rozmiar wejścia. Każda odpowiedź jest zamieniana na małe litery, a przykłady z etykietami innymi niż yes, no lub maybe są usuwane.

Wydruk pierwszego przekonwertowanego przykładu pomaga potwierdzić, że kontekst, pytanie i odpowiedź zostały poprawnie sformatowane przed treningiem.

9. Zbuduj przykłady do treningu dyfuzyjnego

DiffusionGemma wymaga umieszczenia docelowej odpowiedzi wewnątrz płótna o stałej długości. Ta funkcja tokenizuje podpowiedź, konwertuje odpowiedź na identyfikatory tokenów, dopełnia ją do długości płótna modelu i tworzy maskę wskazującą, które tokeny mają wpływać na stratę.

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)

Po każdej odpowiedzi dodawany jest token końca sekwencji, a pozostałe pozycje płótna wypełniane są tokenami dopełniającymi. Maska straty sprawia, że trening koncentruje się wyłącznie na tokenach odpowiedzi i końca sekwencji, a nie na pozycjach dopełnienia.

10. Utwórz funkcje wnioskowania i ewaluacji

Zdefiniuj teraz funkcje służące do generowania odpowiedzi, czyszczenia wyjścia modelu i obliczania dokładności ewaluacji.

Wygeneruj odpowiedź

Funkcja answer_question() formatuje podpowiedź, generuje odpowiedź w kilku krokach odszumiania i dekoduje wygenerowane tokeny do tekstu.

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

Wyodrębnij predykcję

Choć podpowiedź prosi o jednowyrazową odpowiedź, model może czasem wygenerować dodatkowy tekst. Ta funkcja wyodrębnia pierwszą poprawną predykcję yes, no lub 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"

Oceń dokładność

Funkcja ewaluacji porównuje każdą wyczyszczoną predykcję z poprawną odpowiedzią, drukuje wynik dla każdego przykładu i zwraca ogólną dokładność wraz z poszczególnymi predykcjami.

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. Oceń model przed dostrajaniem

Uruchom ewaluację przed treningiem, aby ustalić punkt odniesienia.

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

To ocenia 50 przykładów, używając 16 kroków odszumiania na odpowiedź. W tym eksperymencie model bazowy poprawnie odpowiedział na 30 z 50 pytań.

Ten wynik bazowy porównamy później z dokładnością modelu po dostrojeniu.

12. Skonfiguruj trening

Przełącz model w tryb treningu, utwórz optymalizator i scheduler tempa uczenia oraz zdefiniuj, jak czyste tokeny odpowiedzi będą korumpowane podczas treningu dyfuzyjnego.

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

Tylko parametry z requires_grad=True są przekazywane do optymalizatora, co oznacza, że trenowany jest adapter LoRA, a nie cały model.

Następnie utwórz funkcję korupcji, która zastępuje losową część płótna odpowiedzi losowymi tokenami.

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)

Ilość szumu zmienia się dla każdego przykładu. Podczas treningu model uczy się odtwarzać oryginalną odpowiedź z tych skorumpowanych tokenów płótna.

13. Wytrenuj model

Pętla poniżej trenuje adapter LoRA przez 60 kroków, wykorzystując akumulację gradientów.

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

Dla każdego przykładu treningowego model otrzymuje podpowiedź biomedyczną i skorumpowane płótno odpowiedzi. Strata cross-entropy jest liczona tylko dla prawdziwych tokenów odpowiedzi wybranych przez maskę straty.

Akumulacja gradientów łączy cztery przykłady przed aktualizacją modelu. Przycinanie gradientu również jest stosowane, aby utrzymać stabilność treningu.

W tym eksperymencie trening zakończył się w około dwie minuty:

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

Systematycznie malejąca strata wskazuje, że adapter uczy się odtwarzać oczekiwane odpowiedzi ze skorumpowanego płótna. Podczas treningu możesz też uruchomić nvidia-smi w terminalu RunPod, aby monitorować użycie i obciążenie pamięci GPU.

14. Oceń dostrojony model

Uruchom tę samą ewaluację ponownie po treningu, aby sprawdzić, czy dostrajanie poprawiło wydajność modelu.

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

Dostrojony model jest oceniany na tych samych 50 przykładach i z użyciem tych samych 16 kroków odszumiania co w ewaluacji bazowej.

Następnie porównaj dokładność przed i po dostrojeniu.

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

W tym eksperymencie dokładność modelu wzrosła z 0,60 do 0,80.

To oznacza wzrost o 20 punktów procentowych: po dostrojeniu model poprawnie odpowiedział na 40 z 50 pytań, w porównaniu z 30 z 50 przed treningiem.

15. Zapisz i wyślij dostrojony adapter

Zapisz wytrenowany adapter LoRA i pliki procesora do wcześniej zdefiniowanego katalogu wyjściowego.

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

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

Powinieneś zobaczyć:

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

Zapisywany jest tylko lekki adapter LoRA, a nie kolejna pełna kopia bazowego modelu z 26 miliardami parametrów.

Następnie wyślij adapter i pliki procesora na Hugging Face Hub:

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

Ponieważ zmienna środowiskowa HF_TOKEN została dodana podczas konfiguracji poda RunPod, Hugging Face powinien uwierzytelnić się automatycznie. Musisz uruchomić notebook_login() tylko wtedy, gdy token nie został wcześniej skonfigurowany:

from huggingface_hub import notebook_login

notebook_login()

Po zakończeniu wysyłania repozytorium będzie zawierać adapter LoRA i konfigurację procesora wymaganą do późniejszego wczytania dostrojonego modelu.

Źródło: kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

Wnioski końcowe

Dostrajanie DiffusionGemma za pomocą Unsloth okazało się zaskakująco proste. Najwięcej czasu zajęła instalacja właściwych zależności i zrozumienie, jak działa specyficzny dla dyfuzji proces treningowy. Gdy środowisko było gotowe, wczytanie modelu, trenowanie adaptera LoRA, ewaluacja i wysłanie wyników na Hugging Face przebiegły bardzo sprawnie.

DiffusionGemma wydała mi się szczególnie interesująca, ponieważ nie generuje tekstu token po tokenie jak tradycyjny model językowy. Zamiast tego pracuje na stałym płótnie i stopniowo udoskonala zaszumione tokeny w krokach odszumiania. Zrozumienie tego odmiennego procesu generacji i dostrojenie go do zadania odpowiadania na pytania biomedyczne uczyniło eksperyment szczególnie wartościowym.

Nawet przy małej konfiguracji dokładność na 50‑elementowej próbce wzrosła z 0,60 do 0,80 — choć przy takiej wielkości próbki margines błędu jest duży, a bazowa strategia „zawsze yes” osiąga około 55% na tym zbiorze ewaluacyjnym. 

Warto też zauważyć, że sztuczny podzbiór treningowy zawiera prawie żadnych etykiet „maybe”, więc model ma niewielką szansę nauczyć się tej klasy, choć pojawia się ona w danych ewaluacyjnych. Potraktuj to jako szybki eksperyment pokazujący, jak działa proces dostrajania, a nie dowód, że model nadaje się do realnych zastosowań medycznych.

Chcesz pójść dalej niż pojedyncze dostrajanie? Nasza ścieżka Developing Large Language Models prowadzi cię od podstaw PyTorcha i transformers do budowania i wdrażania własnych LLM‑ów.