Cum să ajustezi fin DiffusionGemma pe PubMedQA folosind Unsloth

1. Deschide un Jupyter Notebook pe RunPod

Creează un pod RunPod nou cu un GPU NVIDIA H100 și selectează un template PyTorch/Jupyter. Configurează cel puțin 100 GB de stocare persistentă ca să nu pierzi fișierele modelului și rezultatele antrenării când podul se oprește.

Adaugă tokenul tău de acces Hugging Face ca variabilă de mediu:

HF_TOKEN=your_hugging_face_token

Asta permite descărcarea mai rapidă a modelelor și seturilor de date și îți permite să încarci adaptorul LoRA salvat pe Hugging Face fără să te autentifici manual din notebook.

Podul configurat ar trebui să coste aproximativ 3 $ pe oră, deși prețul final poate varia în funcție de disponibilitatea GPU-ului și de tipul de pod selectat.

După ce podul rulează, deschide JupyterLab sau Jupyter Notebook din interfața RunPod și creează un notebook nou numit diffusiongemma_pubmedqa.ipynb.

2. Instalează pachetele necesare

Rulează următoarele comenzi în prima celulă a notebookului pentru a instala Unsloth și bibliotecile necesare pentru încărcarea, reglajul fin și salvarea 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"

Comanda %%capture ascunde outputul lung al instalării. Versiunile pachetelor sunt blocate pentru a evita problemele de compatibilitate între DiffusionGemma, Transformers, Unsloth și bibliotecile de antrenare.

După ce instalarea se încheie, repornește kernelul notebookului înainte de a continua.

3. Importează bibliotecile

Importă bibliotecile necesare pentru pregătirea setului de date, încărcarea modelului, antrenare și evaluare.

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 activează descărcări mai rapide din Hugging Face Hub, iar creșterea limitei de recompilare pentru Dynamo ajută la prevenirea întreruperilor când lucrezi cu modelul.

Când celula rulează cu succes, Unsloth va aplica patch-uri mediului de antrenare, iar outputul ar trebui să confirme că CUDA este disponibil și că GPU-ul H100 a fost detectat.

🦥 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. Setează configurația

Definește într-un singur loc modelul, setul de date, parametrii de antrenare, setările de evaluare și directorul de output.

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"

Vom folosi 3.000 de exemple artificiale pentru antrenare și 200 de exemple etichetate manual pentru evaluare. Pentru a păstra experimentul rapid, modelul va fi antrenat timp de 60 de pași și va evalua 50 de exemple folosind 16 pași de de-zgomotare.

Cum funcționează DiffusionGemma

Înainte să încărcăm modelul, merită să ne imaginăm cum produce efectiv DiffusionGemma un răspuns. În loc să scrie tokenii unul după altul, pornește cu un canvas de lungime fixă și îl rafinează pe parcursul mai multor pași de de-zgomotare, actualizând multe poziții simultan până când textul se conturează într-un răspuns coerent. 

Codul de mai jos raportează o lungime a canvasului de 256, dimensiunea unui singur bloc. Pentru răspunsurile scurte yes/no/maybe, un canvas este mai mult decât suficient, în timp ce outputurile mai lungi sunt generate prin înlănțuirea canvasurilor bloc cu bloc. Diagrama de mai jos arată, la nivel înalt, acest proces de rafinare:

5. Încarcă DiffusionGemma

Încarcă modelul DiffusionGemma ajustat pe instrucțiuni în precizie bfloat16. Nu vom folosi cuanta pe 4 biți deoarece GPU-ul H100 are suficientă memorie pentru a încărca modelul la o precizie mai mare.

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)

Dimensiunea vocabularului este folosită când adăugăm zgomot aleator în timpul antrenării prin difuzie. Lungimea canvasului determină numărul maxim de tokeni pe care modelul îi poate rafina într-un singur bloc de generare.

Ar trebui să vezi un output similar cu:

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

6. Adaugă un adaptor LoRA

Adaugă un adaptor LoRA pentru ca doar un set mic de parametri suplimentari să fie antrenați, în loc să actualizăm întregul model cu 26 de miliarde de parametri.

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

Acest lucru reduce semnificativ memoria și calculul necesare pentru reglajul fin. Checkpointing-ul de gradient este dezactivat deoarece H100 are memorie GPU suficientă pentru acest experiment.

7. Încarcă PubMedQA

Încarcă subsetul artificial PubMedQA pentru antrenare și subsetul etichetat manual pentru evaluare.

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

Subsetul de antrenare conține exemple generate automat, în timp ce subsetul de evaluare conține întrebări biomedicale etichetate de experți. 

Afișarea primul rând ne permite să inspectăm întrebarea, contextul din abstract și decizia finală înainte de a formata datele.

Ar trebui să vezi:

Train size: 211269
Eval size: 1000

Fiecare exemplu conține o întrebare biomedicală, unul sau mai multe pasaje-suport din abstract și un răspuns final yes, no sau maybe.

8. Convertește setul de date

Convertește fiecare exemplu PubMedQA într-un format tip chat care conține un prompt al utilizatorului și un răspuns al asistentului.

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

Pasajele de context sunt combinate într-un singur șir și limitate la 2.500 de caractere pentru a păstra inputul gestionabil. Fiecare răspuns este convertit la litere mici, iar exemplele cu etichete în afara yes, no sau maybe sunt eliminate.

Afișarea primului exemplu convertit ajută la confirmarea că contextul, întrebarea și răspunsul au fost formate corect înainte de antrenare.

9. Construiește exemplele pentru antrenarea prin difuzie

DiffusionGemma necesită plasarea răspunsului-țintă într-un canvas de lungime fixă. Această funcție tokenizează promptul, convertește răspunsul în ID-uri de tokeni, îl potrivește la lungimea canvasului modelului și creează o mască ce arată care tokeni ar trebui să contribuie la loss.

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)

Tokenul de sfârșit de secvență este adăugat după fiecare răspuns, iar pozițiile rămase din canvas sunt umplute cu tokeni de padding. Masca de loss asigură că antrenarea se concentrează doar pe tokenii din răspuns și pe tokenul de sfârșit de secvență, nu pe pozițiile completate.

10. Creează funcțiile de inferență și evaluare

În continuare, definește funcțiile folosite pentru a genera răspunsuri, a curăța outputul modelului și a calcula acuratețea evaluării.

Generează un răspuns

Funcția answer_question() formatează promptul, generează un răspuns prin mai mulți pași de de-zgomotare și decodează tokenii generați în 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()

Extrage predicția

Deși promptul cere un răspuns dintr-un singur cuvânt, modelul poate genera ocazional text suplimentar. Această funcție extrage prima predicție validă yes, no sau 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"

Evaluează acuratețea

Funcția de evaluare compară fiecare predicție curățată cu răspunsul corect, afișează rezultatul pentru fiecare exemplu și returnează acuratețea generală împreună cu predicțiile individuale.

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. Evaluează modelul înainte de reglajul fin

Rulează evaluarea înainte de antrenare pentru a stabili un reper.

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

Aceasta evaluează 50 de exemple folosind 16 pași de de-zgomotare per răspuns. În acest experiment, modelul de bază a răspuns corect la 30 din 50 de întrebări.

Acest reper va fi comparat mai târziu cu acuratețea modelului după reglajul fin.

12. Configurează antrenarea

Comută modelul în modul de antrenare, creează optimizerul și schedulerul ratei de învățare și definește cum vor fi corupți tokenii de răspuns curați în timpul antrenării prin difuzie.

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

Doar parametrii cu requires_grad=True sunt trecuți către optimizer, ceea ce înseamnă că procesul de antrenare actualizează adaptorul LoRA, nu întregul model.

În continuare, creează o funcție de corupere care înlocuiește o proporție aleatoare din canvasul răspunsului cu tokeni aleatori.

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)

Cantitatea de zgomot se schimbă pentru fiecare exemplu. În timpul antrenării, modelul învață să reconstruiască răspunsul original din acești tokeni de canvas corupți.

13. Antrenează modelul

Următoarea buclă antrenează adaptorul LoRA timp de 60 de pași folosind acumularea gradientului.

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

Pentru fiecare exemplu de antrenare, modelul primește promptul biomedical și un canvas de răspuns corupt. Loss-ul cross-entropy este calculat doar pentru tokenii reali ai răspunsului selectați de masca de loss.

Acumularea gradientului combină patru exemple înainte de a actualiza modelul. Se aplică și tăierea gradientului pentru a menține antrenarea stabilă.

În acest experiment, antrenarea s-a încheiat în aproximativ două minute:

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

Scăderea constantă a loss-ului indică faptul că adaptorul învață să reconstruiască răspunsurile așteptate din canvasul corupt. În timpul antrenării, poți rula și nvidia-smi în terminalul RunPod pentru a monitoriza utilizarea memoriei și a GPU-ului.

14. Evaluează modelul reglat fin

Rulează aceeași evaluare din nou după antrenare pentru a măsura dacă reglajul fin a îmbunătățit performanța modelului.

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

Modelul reglat fin este evaluat pe aceleași 50 de exemple și folosind aceiași 16 pași de de-zgomotare ca în evaluarea de bază.

În continuare, compară acuratețea înainte și după reglajul fin.

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

În acest experiment, acuratețea modelului a crescut de la 0,60 la 0,80.

Aceasta reprezintă o îmbunătățire de 20 de puncte procentuale, modelul răspunzând corect la 40 din 50 de întrebări după reglajul fin, comparativ cu 30 din 50 înainte de antrenare.

15. Salvează și încarcă adaptorul reglat fin

Salvează adaptorul LoRA antrenat și fișierele procesorului în directorul de output definit anterior.

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

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

Ar trebui să vezi:

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

Aceasta salvează doar adaptorul LoRA ușor, nu o altă copie completă a modelului de bază cu 26 de miliarde de parametri.

În continuare, încarcă adaptorul și fișierele procesorului pe Hugging Face Hub:

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

Deoarece variabila de mediu HF_TOKEN a fost adăugată în timpul configurării podului RunPod, Hugging Face ar trebui să se autentifice automat. Trebuie să rulezi notebook_login() doar când tokenul nu a fost deja configurat:

from huggingface_hub import notebook_login

notebook_login()

După ce încărcarea se încheie, depozitul va conține adaptorul LoRA și configurația procesorului necesare pentru a încărca ulterior modelul reglat fin.

Sursă: kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

Gânduri finale

Reglarea fină a DiffusionGemma cu Unsloth a fost surprinzător de ușoară. Cea mai consumatoare parte de timp a fost instalarea dependențelor corecte și înțelegerea modului în care funcționează procesul de antrenare specific difuziei. Odată ce mediul a fost configurat, încărcarea modelului, antrenarea adaptorului LoRA, evaluarea și încărcarea rezultatelor pe Hugging Face au mers foarte lin.

Am găsit DiffusionGemma deosebit de interesant pentru că nu generează text un token pe rând ca un model lingvistic tradițional. În schimb, lucrează cu un canvas fix și rafinează treptat tokeni zgomotoși prin pași de de-zgomotare. Înțelegerea acestui proces diferit de generare și reglarea lui fină pe o sarcină de întrebări-răspunsuri biomedicale au făcut experimentul deosebit de valoros pentru mine.

Chiar și cu un setup mic, acuratețea pe eșantionul de 50 de exemple a crescut de la 0,60 la 0,80—deși pe un eșantion de dimensiunea asta, marja de eroare e mare, iar un baseline „întotdeauna yes” obține deja în jur de 55% pe acest set de evaluare. 

Merită menționat și că subsetul artificial de antrenare conține aproape deloc etichete „maybe”, așa că modelul are puține șanse să învețe acea clasă, deși apare în datele de evaluare. Tratează asta ca pe un experiment rapid despre cum funcționează procesul de reglaj fin, nu ca pe o dovadă că modelul e gata pentru uz medical real.

Gata să mergi mai departe de un singur reglaj fin? Traseul nostru Developing Large Language Models te poartă de la fundamentele PyTorch și transformers până la construirea și implementarea propriilor tale LLM-uri.