Come fare fine-tuning di DiffusionGemma su PubMedQA con Unsloth

1. Apri un Jupyter Notebook su RunPod

Crea un nuovo pod RunPod con una GPU NVIDIA H100 e seleziona un template PyTorch/Jupyter. Configura almeno 100 GB di storage persistente così che i file del modello e gli output dell’addestramento non vadano persi quando il pod si arresta.

Aggiungi il tuo access token di Hugging Face come variabile d’ambiente:

HF_TOKEN=your_hugging_face_token

Questo permette di scaricare più velocemente modelli e dataset e ti consente di caricare l’adattatore LoRA salvato su Hugging Face senza autenticarti manualmente dal notebook.

Il pod configurato dovrebbe costare circa 3 $ all’ora, anche se il prezzo finale può variare in base alla disponibilità delle GPU e al tipo di pod selezionato.

Una volta che il pod è in esecuzione, apri JupyterLab o Jupyter Notebook dall’interfaccia di RunPod e crea un nuovo notebook chiamato diffusiongemma_pubmedqa.ipynb.

2. Installa i pacchetti necessari

Esegui i seguenti comandi nella prima cella del notebook per installare Unsloth e le librerie richieste per il caricamento, il fine-tuning e il salvataggio di 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"

Il comando %%capture nasconde il lungo output di installazione. Le versioni dei pacchetti sono bloccate per evitare problemi di compatibilità tra DiffusionGemma, Transformers, Unsloth e le librerie di training.

Quando l’installazione termina, riavvia il kernel del notebook prima di continuare.

3. Importa le librerie

Importa le librerie necessarie per la preparazione del dataset, il caricamento del modello, l’addestramento e la valutazione.

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 abilita download più veloci dall’Hugging Face Hub, mentre aumentare il limite di ricompilazione di Dynamo aiuta a prevenire interruzioni quando si lavora con il modello.

Quando la cella viene eseguita con successo, Unsloth applicherà le patch all’ambiente di training e l’output dovrebbe confermare che CUDA è disponibile e che la GPU H100 è stata rilevata.

🦥 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. Imposta la configurazione

Definisci in un unico punto modello, dataset, parametri di training, impostazioni di valutazione e directory di 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"

Useremo 3.000 esempi artificiali per l’addestramento e 200 esempi annotati manualmente per la valutazione. Per mantenere l’esperimento rapido, il modello verrà addestrato per 60 step e valuterà 50 esempi usando 16 passaggi di denoising.

Come funziona DiffusionGemma

Prima di caricare il modello, vale la pena visualizzare come DiffusionGemma produce effettivamente una risposta. Invece di scrivere i token uno dopo l’altro, parte da una canvas a lunghezza fissa di token e la affina in più passaggi di denoising, aggiornando molte posizioni in una volta finché il testo non si risolve in una risposta coerente. 

Il codice sotto riporta una lunghezza della canvas di 256, la dimensione di un singolo blocco. Per le nostre brevi risposte yes/no/maybe, una canvas è più che sufficiente, mentre output più lunghi vengono generati concatenando le canvas blocco dopo blocco. Il diagramma seguente mostra a grandi linee questo processo di affinamento:

5. Carica DiffusionGemma

Carica il modello DiffusionGemma ottimizzato per istruzioni in precisione bfloat16. Non useremo la quantizzazione a 4 bit perché la GPU H100 ha memoria sufficiente per caricare il modello a precisione più alta.

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)

La dimensione del vocabolario è usata quando si aggiunge rumore casuale durante l’addestramento a diffusione. La lunghezza della canvas determina il numero massimo di token che il modello può affinare in un singolo blocco di generazione.

Dovresti vedere un output simile a:

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

6. Aggiungi un adattatore LoRA

Aggiungi un adattatore LoRA così da addestrare solo un piccolo insieme di parametri aggiuntivi invece di aggiornare l’intero modello da 26 miliardi di parametri.

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

Questo riduce significativamente la memoria e il calcolo richiesti per il fine-tuning. Il gradient checkpointing è disabilitato perché l’H100 ha memoria GPU sufficiente per questo esperimento.

7. Carica PubMedQA

Carica il sottoinsieme artificiale di PubMedQA per l’addestramento e il sottoinsieme annotato manualmente per la valutazione.

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

Il sottoinsieme di training contiene esempi generati automaticamente, mentre il sottoinsieme di valutazione include domande biomediche annotate da esperti. 

Stampare la prima riga ci permette di ispezionare la domanda, il contesto dell’abstract e la decisione finale prima di formattare i dati.

Dovresti vedere:

Train size: 211269
Eval size: 1000

Ogni esempio contiene una domanda biomedica, uno o più passaggi di abstract di supporto e una risposta finale yes, no o maybe.

8. Converte il dataset

Converti ogni esempio di PubMedQA in un formato stile chat contenente un prompt dell’utente e una risposta dell’assistente.

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

I passaggi di contesto sono combinati in un’unica stringa e limitati a 2.500 caratteri per mantenere l’input gestibile. Ogni risposta è convertita in minuscolo e gli esempi con etichette diverse da yes, no o maybe vengono rimossi.

Stampare il primo esempio convertito aiuta a confermare che contesto, domanda e risposta siano stati formattati correttamente prima dell’addestramento.

9. Crea gli esempi per l’addestramento a diffusione

DiffusionGemma richiede che la risposta target sia collocata dentro una canvas a lunghezza fissa. Questa funzione tokenizza il prompt, converte la risposta in ID di token, la riempie (padding) fino alla lunghezza della canvas del modello e crea una maschera che indica quali token devono contribuire alla 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)

Il token di fine sequenza viene aggiunto dopo ogni risposta, mentre le posizioni rimanenti della canvas sono riempite con token di padding. La loss mask garantisce che l’addestramento si concentri solo sui token della risposta e di fine sequenza, e non sulle posizioni di padding.

10. Crea le funzioni di inferenza e valutazione

Definisci ora le funzioni usate per generare le risposte, pulire l’output del modello e calcolare l’accuratezza di valutazione.

Genera una risposta

La funzione answer_question() formatta il prompt, genera una risposta attraverso vari passaggi di denoising e decodifica i token generati in testo.

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

Estrai la previsione

Sebbene il prompt richieda una risposta di una sola parola, a volte il modello può generare testo aggiuntivo. Questa funzione estrae la prima previsione valida tra yes, no o 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"

Valuta l’accuratezza

La funzione di valutazione confronta ogni previsione pulita con la risposta corretta, stampa il risultato per ogni esempio e restituisce l’accuratezza complessiva insieme alle singole previsioni.

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. Valuta il modello prima del fine-tuning

Esegui la valutazione prima dell’addestramento per stabilire una baseline.

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

Questo valuta 50 esempi usando 16 passaggi di denoising per risposta. In questo esperimento, il modello base ha risposto correttamente a 30 domande su 50.

Questa baseline verrà poi confrontata con l’accuratezza del modello dopo il fine-tuning.

12. Configura l’addestramento

Passa il modello in modalità training, crea l’ottimizzatore e lo scheduler del learning rate e definisci come i token della risposta pulita verranno corrotti durante l’addestramento a diffusione.

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

Solo i parametri con requires_grad=True vengono passati all’ottimizzatore, il che significa che il processo di training aggiorna l’adattatore LoRA e non l’intero modello.

Successivamente, crea una funzione di corruzione che sostituisce una proporzione casuale della canvas della risposta con token casuali.

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)

La quantità di rumore cambia per ogni esempio. Durante l’addestramento, il modello impara a ricostruire la risposta originale a partire da questi token corrotti della canvas.

13. Addestra il modello

Il seguente loop addestra l’adattatore LoRA per 60 step usando l’accumulo del gradiente.

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

Per ogni esempio di training, il modello riceve il prompt biomedico e una canvas della risposta corrotta. La loss di entropia incrociata è calcolata solo per i token della risposta reale selezionati dalla loss mask.

L’accumulo del gradiente combina quattro esempi prima di aggiornare il modello. Viene applicato anche il gradient clipping per mantenere stabile l’addestramento.

In questo esperimento, l’addestramento si è concluso in circa due minuti:

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

La perdita in diminuzione indica che l’adattatore sta imparando a ricostruire le risposte attese dalla canvas corrotta. Durante il training puoi anche eseguire nvidia-smi nel terminale di RunPod per monitorare l’uso e l’utilizzo della memoria della GPU.

14. Valuta il modello fine-tunato

Esegui nuovamente la stessa valutazione dopo l’addestramento per misurare se il fine-tuning ha migliorato le prestazioni del modello.

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

Il modello fine-tunato viene valutato sugli stessi 50 esempi e usando gli stessi 16 passaggi di denoising della valutazione baseline.

Poi confronta l’accuratezza prima e dopo il 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 questo esperimento, l’accuratezza del modello è passata da 0,60 a 0,80.

Si tratta di un miglioramento di 20 punti percentuali: dopo il fine-tuning il modello ha risposto correttamente a 40 domande su 50, contro le 30 su 50 prima dell’addestramento.

15. Salva e carica l’adattatore fine-tunato

Salva l’adattatore LoRA addestrato e i file del processor nella directory di output definita in precedenza.

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

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

Dovresti vedere:

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

Questo salva solo il leggero adattatore LoRA, invece di un’altra copia completa del modello base da 26 miliardi di parametri.

Successivamente, carica l’adattatore e i file del processor sull’Hugging Face Hub:

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

Poiché la variabile d’ambiente HF_TOKEN è stata aggiunta durante la configurazione del pod su RunPod, Hugging Face dovrebbe autenticarsi automaticamente. Devi eseguire notebook_login() solo quando il token non è già configurato:

from huggingface_hub import notebook_login

notebook_login()

Al termine del caricamento, il repository conterrà l’adattatore LoRA e la configurazione del processor necessari per caricare in seguito il modello fine-tunato.

Fonte: kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

Considerazioni finali

Fare fine-tuning di DiffusionGemma con Unsloth è stato sorprendentemente semplice. La parte più lunga è stata installare le dipendenze corrette e capire come funziona il processo di training specifico per la diffusione. Una volta impostato l’ambiente, caricare il modello, addestrare l’adattatore LoRA, valutarlo e caricare i risultati su Hugging Face è stato tutto molto lineare.

Ho trovato DiffusionGemma particolarmente interessante perché non genera testo un token alla volta come un modello linguistico tradizionale. Invece lavora con una canvas fissa e affina gradualmente i token rumorosi attraverso step di denoising. Comprendere questo diverso processo di generazione e fare fine-tuning su un task di question answering biomedico ha reso l’esperimento particolarmente utile per me.

Anche con un setup ridotto, l’accuratezza sul campione da 50 esempi è salita da 0,60 a 0,80—anche se, con un campione di queste dimensioni, il margine d’errore è ampio e una baseline "sempre yes" ottiene già circa il 55% su questo set di valutazione. 

Vale anche la pena notare che il sottoinsieme di training artificiale contiene quasi nessuna etichetta "maybe", quindi il modello ha poche possibilità di apprendere quella classe, anche se compare nei dati di valutazione. Consideralo come un rapido esperimento sul funzionamento del processo di fine-tuning, non come prova che il modello sia pronto per un uso medico reale.

Pronto ad andare oltre un singolo fine-tuning run? Il nostro percorso Developing Large Language Models ti porta dalle basi di PyTorch e transformers alla creazione e al deployment dei tuoi LLM.