Comment affiner DiffusionGemma sur PubMedQA avec Unsloth

1. Ouvrir un notebook Jupyter sur RunPod

Créez un nouveau pod RunPod avec un GPU NVIDIA H100 et sélectionnez un modèle PyTorch/Jupyter. Configurez au moins 100 Go de stockage persistant pour que vos fichiers de modèle et sorties d'entraînement ne soient pas perdus lorsque le pod s'arrête.

Ajoutez votre jeton d'accès Hugging Face comme variable d'environnement :

HF_TOKEN=your_hugging_face_token

Cela accélère le téléchargement des modèles et jeux de données et vous permet de téléverser l'adaptateur LoRA enregistré sur Hugging Face sans vous connecter manuellement depuis le notebook.

Le coût du pod configuré sera d'environ 3 $ de l'heure, le prix final pouvant varier selon la disponibilité du GPU et le type de pod sélectionné.

Une fois le pod démarré, ouvrez JupyterLab ou Jupyter Notebook depuis l'interface RunPod et créez un nouveau notebook nommé diffusiongemma_pubmedqa.ipynb.

2. Installer les paquets requis

Exécutez les commandes suivantes dans la première cellule du notebook pour installer Unsloth et les bibliothèques nécessaires pour charger, affiner et enregistrer 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"

La commande %%capture masque la sortie verbeuse de l'installation. Les versions des paquets sont figées pour éviter les problèmes de compatibilité entre DiffusionGemma, Transformers, Unsloth et les bibliothèques d'entraînement.

Une fois l'installation terminée, redémarrez le noyau du notebook avant de continuer.

3. Importer les bibliothèques

Importez les bibliothèques nécessaires à la préparation du jeu de données, au chargement du modèle, à l'entraînement et à l'évaluation.

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 active des téléchargements plus rapides depuis le Hub Hugging Face, tandis que l'augmentation de la limite de recompilation de Dynamo aide à prévenir les interruptions lors de l'utilisation du modèle.

Lorsque la cellule s'exécute correctement, Unsloth patche l'environnement d'entraînement, et la sortie doit confirmer que CUDA est disponible et que le GPU H100 a été détecté.

🦥 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. Définir la configuration

Définissez au même endroit le modèle, le jeu de données, les paramètres d'entraînement, les réglages d'évaluation et le répertoire de sortie.

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"

Nous utiliserons 3 000 exemples artificiels pour l'entraînement et 200 exemples annotés manuellement pour l'évaluation. Pour garder l'expérience rapide, le modèle s'entraînera pendant 60 étapes et évaluera 50 exemples en utilisant 16 étapes de débruitage.

Comment fonctionne DiffusionGemma

Avant de charger le modèle, il est utile d'imaginer comment DiffusionGemma produit réellement une réponse. Au lieu d'écrire des jetons les uns après les autres, il part d'une toile de longueur fixe et l'affine sur plusieurs étapes de débruitage, en mettant à jour de nombreuses positions à la fois jusqu'à obtenir une réponse cohérente. 

Le code ci-dessous indique une longueur de toile de 256, la taille d'un bloc. Pour nos réponses courtes yes/no/maybe, une seule toile suffit largement, tandis que des sorties plus longues sont générées en enchaînant les toiles bloc par bloc. Le schéma ci-dessous illustre ce processus d'affinage à haut niveau :

5. Charger DiffusionGemma

Chargez le modèle DiffusionGemma ajusté aux instructions en précision bfloat16. Nous n'utiliserons pas la quantification 4 bits, car le GPU H100 dispose de suffisamment de mémoire pour charger le modèle en plus haute précision.

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 taille du vocabulaire est utilisée lors de l'ajout de bruit aléatoire pendant l'entraînement par diffusion. La longueur de la toile détermine le nombre maximal de jetons que le modèle peut affiner dans un seul bloc de génération.

Vous devriez obtenir une sortie similaire :

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

6. Ajouter un adaptateur LoRA

Ajoutez un adaptateur LoRA afin de n'entraîner qu'un petit ensemble de paramètres supplémentaires au lieu de mettre à jour l'intégralité du modèle de 26 milliards de paramètres.

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

Cela réduit significativement la mémoire et le calcul nécessaires à l'affinage. Le gradient checkpointing est désactivé car le H100 dispose de suffisamment de mémoire GPU pour cette expérience.

7. Charger PubMedQA

Chargez le sous-ensemble artificiel de PubMedQA pour l'entraînement et le sous-ensemble annoté manuellement pour l'évaluation.

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

Le sous-ensemble d'entraînement contient des exemples générés automatiquement, tandis que le sous-ensemble d'évaluation contient des questions biomédicales annotées par des experts. 

L'impression de la première ligne nous permet d'inspecter la question, le contexte de l'abstract et la décision finale avant de mettre en forme les données.

Vous devriez voir :

Train size: 211269
Eval size: 1000

Chaque exemple contient une question biomédicale, un ou plusieurs passages d'abstract en appui, et une réponse finale yes, no ou maybe.

8. Convertir le jeu de données

Convertissez chaque exemple PubMedQA en un format de type chat contenant une invite utilisateur et une réponse d'assistant.

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

Les passages de contexte sont combinés en une seule chaîne et limités à 2 500 caractères pour garder l'entrée gérable. Chaque réponse est convertie en minuscules, et les exemples avec des étiquettes autres que yes, no ou maybe sont supprimés.

Imprimer le premier exemple converti permet de vérifier que le contexte, la question et la réponse ont été correctement formatés avant l'entraînement.

9. Construire les exemples d'entraînement par diffusion

DiffusionGemma nécessite que la réponse cible soit placée dans une toile de longueur fixe. Cette fonction tokenise l'invite, convertit la réponse en identifiants de jetons, la padde à la longueur de toile du modèle et crée un masque indiquant quels jetons doivent contribuer à la perte.

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)

Le jeton de fin de séquence est ajouté après chaque réponse, tandis que les positions restantes de la toile sont remplies de jetons de bourrage. Le masque de perte garantit que l'entraînement se concentre uniquement sur les jetons de réponse et de fin de séquence, et non sur les positions remplies.

10. Créer les fonctions d'inférence et d'évaluation

Définissez ensuite les fonctions utilisées pour générer les réponses, nettoyer la sortie du modèle et calculer la précision d'évaluation.

Générer une réponse

La fonction answer_question() met en forme l'invite, génère une réponse via plusieurs étapes de débruitage, puis décode les jetons générés en texte.

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

Extraire la prédiction

Même si l'invite demande une réponse d'un seul mot, le modèle peut parfois générer du texte supplémentaire. Cette fonction extrait la première prédiction valide yes, no ou 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"

Évaluer la précision

La fonction d'évaluation compare chaque prédiction nettoyée avec la bonne réponse, affiche le résultat pour chaque exemple et renvoie la précision globale ainsi que les prédictions individuelles.

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. Évaluer le modèle avant l'affinage

Exécutez l'évaluation avant l'entraînement pour établir une référence.

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

Cela évalue 50 exemples en utilisant 16 étapes de débruitage par réponse. Dans cette expérience, le modèle de base a répondu correctement à 30 questions sur 50.

Cette base de référence sera ensuite comparée à la précision du modèle après affinage.

12. Configurer l'entraînement

Passez le modèle en mode entraînement, créez l'optimiseur et le planificateur de taux d'apprentissage, puis définissez comment les jetons de réponse propres seront corrompus pendant l'entraînement par diffusion.

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

Seuls les paramètres avec requires_grad=True sont passés à l'optimiseur, ce qui signifie que le processus d'entraînement met à jour l'adaptateur LoRA plutôt que le modèle complet.

Créez ensuite une fonction de corruption qui remplace une proportion aléatoire de la toile de réponse par des jetons aléatoires.

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é de bruit change pour chaque exemple. Pendant l'entraînement, le modèle apprend à reconstruire la réponse originale à partir de ces jetons de toile corrompus.

13. Entraîner le modèle

La boucle suivante entraîne l'adaptateur LoRA pendant 60 étapes avec accumulation de gradients.

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

Pour chaque exemple d'entraînement, le modèle reçoit l'invite biomédicale et une toile de réponse corrompue. La perte d'entropie croisée est calculée uniquement pour les vrais jetons de réponse sélectionnés par le masque de perte.

L'accumulation de gradients combine quatre exemples avant de mettre à jour le modèle. Un clipping des gradients est également appliqué pour stabiliser l'entraînement.

Dans cette expérience, l'entraînement s'est terminé en environ deux minutes :

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

La baisse régulière de la perte indique que l'adaptateur apprend à reconstruire les réponses attendues à partir de la toile corrompue. Pendant l'entraînement, vous pouvez également exécuter nvidia-smi dans le terminal RunPod pour surveiller l'utilisation et la mémoire du GPU.

14. Évaluer le modèle affiné

Exécutez la même évaluation après l'entraînement pour mesurer si l'affinage a amélioré les performances du modèle.

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

Le modèle affiné est évalué sur les mêmes 50 exemples et avec les mêmes 16 étapes de débruitage que lors de l'évaluation de référence.

Comparez ensuite la précision avant et après l'affinage.

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

Dans cette expérience, la précision du modèle est passée de 0,60 à 0,80.

Cela représente un gain de 20 points de pourcentage, avec 40 bonnes réponses sur 50 après affinage, contre 30 sur 50 avant entraînement.

15. Enregistrer et téléverser l'adaptateur affiné

Enregistrez l'adaptateur LoRA entraîné et les fichiers du processeur dans le répertoire de sortie défini plus haut.

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

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

Vous devriez voir :

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

Cela enregistre uniquement le léger adaptateur LoRA, plutôt qu'une autre copie complète du modèle de base de 26 milliards de paramètres.

Téléversez ensuite l'adaptateur et les fichiers du processeur sur le Hub Hugging Face :

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

Comme la variable d'environnement HF_TOKEN a été ajoutée lors de la configuration du pod RunPod, Hugging Face devrait s'authentifier automatiquement. Vous n'avez besoin d'exécuter notebook_login() que si le jeton n'a pas déjà été configuré :

from huggingface_hub import notebook_login

notebook_login()

Après la fin du téléversement, le dépôt contiendra l'adaptateur LoRA et la configuration du processeur nécessaires pour charger ultérieurement le modèle affiné.

Source : kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

Dernières réflexions

L'affinage de DiffusionGemma avec Unsloth s'est révélé étonnamment simple. La partie la plus chronophage a été l'installation des bonnes dépendances et la compréhension du processus d'entraînement spécifique à la diffusion. Une fois l'environnement prêt, le chargement du modèle, l'entraînement de l'adaptateur LoRA, son évaluation et le téléversement des résultats sur Hugging Face se sont déroulés sans accroc.

J'ai trouvé DiffusionGemma particulièrement intéressant car il ne génère pas le texte un jeton à la fois comme un modèle de langage traditionnel. Il fonctionne plutôt avec une toile fixe et affine progressivement des jetons bruités via des étapes de débruitage. Comprendre ce processus de génération différent et l'affiner sur une tâche de questions-réponses biomédicales a rendu l'expérience particulièrement enrichissante pour moi.

Même avec un petit dispositif, la précision sur l'échantillon de 50 exemples est passée de 0,60 à 0,80 — bien que sur un échantillon de cette taille, la marge d'erreur soit large, et un baseline « toujours yes » atteigne déjà environ 55 % sur cet ensemble d'évaluation. 

Il convient aussi de noter que le sous-ensemble d'entraînement artificiel contient quasiment aucune étiquette « maybe », de sorte que le modèle a peu d'occasions d'apprendre cette classe, alors qu'elle apparaît dans les données d'évaluation. Considérez ceci comme une expérience rapide pour comprendre le processus d'affinage, et non comme la preuve que le modèle est prêt pour un usage médical réel.

Prêt à aller au-delà d'un seul affinage ? Notre parcours Developing Large Language Models vous amène des fondamentaux de PyTorch et transformers jusqu'à la création et au déploiement de vos propres LLM.