Cómo ajustar finamente DiffusionGemma en PubMedQA con Unsloth

1. Abre un Jupyter Notebook en RunPod

Crea un nuevo pod en RunPod con una GPU NVIDIA H100 y selecciona una plantilla de PyTorch/Jupyter. Configura al menos 100 GB de almacenamiento persistente para que los archivos del modelo y las salidas de entrenamiento no se pierdan al detener el pod.

Añade tu token de acceso de Hugging Face como variable de entorno:

HF_TOKEN=your_hugging_face_token

Así, los modelos y datasets se descargarán más rápido y podrás subir el adaptador LoRA guardado a Hugging Face sin iniciar sesión manualmente desde el notebook.

El coste estimado del pod configurado es de unos 3 $ por hora, aunque el precio final puede variar según la disponibilidad de GPU y el tipo de pod seleccionado.

Cuando el pod esté en ejecución, abre JupyterLab o Jupyter Notebook desde la interfaz de RunPod y crea un notebook nuevo llamado diffusiongemma_pubmedqa.ipynb.

2. Instala los paquetes necesarios

Ejecuta estos comandos en la primera celda del notebook para instalar Unsloth y las bibliotecas necesarias para cargar, ajustar finamente y guardar 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"

El comando %%capture oculta la salida larga de la instalación. Las versiones de los paquetes están fijadas para evitar problemas de compatibilidad entre DiffusionGemma, Transformers, Unsloth y las librerías de entrenamiento.

Cuando termine la instalación, reinicia el kernel del notebook antes de continuar.

3. Importa las librerías

Importa las librerías necesarias para preparar el dataset, cargar el modelo, entrenar y evaluar.

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 activa descargas más rápidas desde Hugging Face Hub, y aumentar el límite de recompilación de Dynamo ayuda a evitar interrupciones al trabajar con el modelo.

Si la celda se ejecuta correctamente, Unsloth parcheará el entorno de entrenamiento y la salida confirmará que CUDA está disponible y que se ha detectado la 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. Define la configuración

Define en un solo lugar el modelo, el dataset, los parámetros de entrenamiento, los ajustes de evaluación y el directorio de salida.

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"

Usaremos 3.000 ejemplos artificiales para entrenar y 200 ejemplos etiquetados manualmente para evaluar. Para mantener el experimento ágil, el modelo entrenará durante 60 pasos y evaluará 50 ejemplos usando 16 pasos de eliminación de ruido.

Cómo funciona DiffusionGemma

Antes de cargar el modelo, conviene visualizar cómo produce una respuesta DiffusionGemma. En lugar de escribir tokens uno tras otro, parte de un lienzo de longitud fija y lo va refinando en varios pasos de denoising, actualizando muchas posiciones a la vez hasta obtener una respuesta coherente. 

El siguiente código informa de una longitud de lienzo de 256, el tamaño de un bloque. Para nuestras respuestas cortas yes/no/maybe, un lienzo es más que suficiente, mientras que las salidas largas se generan encadenando lienzos bloque a bloque. El diagrama muestra el proceso de refinamiento a alto nivel:

5. Carga DiffusionGemma

Carga el modelo DiffusionGemma ajustado con instrucciones en precisión bfloat16. No usaremos cuantización a 4 bits porque la H100 tiene memoria suficiente para cargar el modelo con mayor precisión.

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)

El tamaño del vocabulario se usa al añadir ruido aleatorio durante el entrenamiento por difusión. La longitud del lienzo determina el número máximo de tokens que el modelo puede refinar en un único bloque de generación.

Deberías ver una salida similar a:

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

6. Añade un adaptador LoRA

Añade un adaptador LoRA para que solo se entrene un pequeño conjunto de parámetros adicionales en lugar de actualizar todo el modelo de 26 mil millones de parámetros.

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

Esto reduce de forma significativa la memoria y el cómputo necesarios para el ajuste fino. Desactivamos el gradient checkpointing porque la H100 tiene memoria suficiente para este experimento.

7. Carga PubMedQA

Carga el subconjunto artificial de PubMedQA para entrenamiento y el subconjunto etiquetado manualmente para evaluación.

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

El subconjunto de entrenamiento contiene ejemplos generados automáticamente, mientras que el de evaluación incluye preguntas biomédicas etiquetadas por expertos. 

Imprimir la primera fila nos permite inspeccionar la pregunta, el resumen (abstract) como contexto y la decisión final antes de dar formato a los datos.

Deberías ver:

Train size: 211269
Eval size: 1000

Cada ejemplo incluye una pregunta biomédica, uno o varios pasajes de resúmenes como soporte y una respuesta final yes, no o maybe.

8. Convierte el dataset

Convierte cada ejemplo de PubMedQA a un formato tipo chat con un prompt del usuario y una respuesta del asistente.

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

Los pasajes de contexto se combinan en una sola cadena y se limitan a 2.500 caracteres para mantener la entrada manejable. Cada respuesta se pasa a minúsculas y se eliminan los ejemplos con etiquetas fuera de yes, no o maybe.

Imprimir el primer ejemplo convertido ayuda a confirmar que el contexto, la pregunta y la respuesta se han formateado correctamente antes de entrenar.

9. Crea los ejemplos para entrenamiento por difusión

DiffusionGemma requiere colocar la respuesta objetivo dentro de un lienzo de longitud fija. Esta función tokeniza el prompt, convierte la respuesta en IDs de tokens, la rellena hasta la longitud del lienzo del modelo y crea una máscara que indica qué tokens contribuyen a la pérdida.

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)

Se añade el token de fin de secuencia tras cada respuesta, mientras que las posiciones restantes del lienzo se rellenan con tokens de padding. La máscara de pérdida garantiza que el entrenamiento se centre solo en los tokens de la respuesta y el fin de secuencia, y no en el padding.

10. Crea las funciones de inferencia y evaluación

A continuación, define las funciones para generar respuestas, limpiar la salida del modelo y calcular la precisión de la evaluación.

Genera una respuesta

La función answer_question() formatea el prompt, genera una respuesta mediante varios pasos de denoising y decodifica los tokens generados a texto.

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

Extrae la predicción

Aunque el prompt pide una respuesta de una sola palabra, a veces el modelo puede generar texto adicional. Esta función extrae la primera predicción válida entre 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"

Evalúa la precisión

La función de evaluación compara cada predicción limpiada con la respuesta correcta, imprime el resultado de cada ejemplo y devuelve la precisión global junto con las predicciones individuales.

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. Evalúa el modelo antes del ajuste fino

Ejecuta la evaluación antes de entrenar para establecer una línea base.

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

Esto evalúa 50 ejemplos usando 16 pasos de denoising por respuesta. En este experimento, el modelo base respondió correctamente 30 de 50 preguntas.

Más adelante compararemos esta línea base con la precisión del modelo tras el ajuste fino.

12. Prepara el entrenamiento

Pon el modelo en modo entrenamiento, crea el optimizador y el programador de tasa de aprendizaje, y define cómo se corromperán los tokens de la respuesta durante el entrenamiento por difusión.

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 se pasan al optimizador los parámetros con requires_grad=True, lo que significa que el entrenamiento actualiza el adaptador LoRA y no el modelo completo.

A continuación, crea una función de corrupción que reemplace una proporción aleatoria del lienzo de la respuesta con tokens aleatorios.

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 cantidad de ruido cambia en cada ejemplo. Durante el entrenamiento, el modelo aprende a reconstruir la respuesta original a partir de estos tokens corrompidos del lienzo.

13. Entrena el modelo

El siguiente bucle entrena el adaptador LoRA durante 60 pasos usando acumulación de gradientes.

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

En cada ejemplo de entrenamiento, el modelo recibe el prompt biomédico y un lienzo de respuesta corrompido. La pérdida de entropía cruzada se calcula solo para los tokens reales de la respuesta seleccionados por la máscara de pérdida.

La acumulación de gradientes combina cuatro ejemplos antes de actualizar el modelo. También se aplica recorte de gradientes para mantener el entrenamiento estable.

En este experimento, el entrenamiento terminó en aproximadamente dos minutos:

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

La pérdida decreciente indica que el adaptador aprende a reconstruir las respuestas esperadas a partir del lienzo corrompido. Durante el entrenamiento, también puedes ejecutar nvidia-smi en la terminal de RunPod para monitorizar el uso y la ocupación de memoria de la GPU.

14. Evalúa el modelo ajustado finamente

Ejecuta de nuevo la misma evaluación tras entrenar para medir si el ajuste fino ha mejorado el rendimiento.

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

El modelo ajustado se evalúa en los mismos 50 ejemplos y con los mismos 16 pasos de denoising que en la evaluación base.

Ahora, compara la precisión antes y después del ajuste fino.

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

En este experimento, la precisión del modelo subió de 0,60 a 0,80.

Esto supone una mejora de 20 puntos porcentuales: el modelo respondió correctamente 40 de 50 preguntas tras el ajuste fino, frente a 30 de 50 antes de entrenar.

15. Guarda y sube el adaptador ajustado

Guarda el adaptador LoRA entrenado y los archivos del procesador en el directorio de salida definido anteriormente.

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

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

Deberías ver:

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

Esto guarda solo el adaptador LoRA, ligero, en lugar de otra copia completa del modelo base de 26 mil millones de parámetros.

Después, sube el adaptador y los archivos del procesador a Hugging Face Hub:

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

Como añadiste la variable de entorno HF_TOKEN al configurar el pod de RunPod, Hugging Face debería autenticarse automáticamente. Solo necesitas ejecutar notebook_login() cuando el token no esté ya configurado:

from huggingface_hub import notebook_login

notebook_login()

Tras la subida, el repositorio contendrá el adaptador LoRA y la configuración del procesador necesarios para volver a cargar el modelo ajustado más adelante.

Fuente: kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

Reflexiones finales

Ajustar DiffusionGemma con Unsloth ha sido sorprendentemente sencillo. Lo más laborioso fue instalar las dependencias correctas y entender cómo funciona el proceso de entrenamiento específico de difusión. Una vez listo el entorno, cargar el modelo, entrenar el adaptador LoRA, evaluarlo y subir los resultados a Hugging Face fue muy fluido.

DiffusionGemma me ha parecido especialmente interesante porque no genera texto token a token como un modelo tradicional. Trabaja con un lienzo fijo y va refinando tokens ruidosos con pasos de denoising. Entender este proceso de generación distinto y ajustarlo finamente a una tarea de preguntas y respuestas biomédicas hizo que el experimento fuera especialmente valioso para mí.

Incluso con una configuración pequeña, la precisión en la muestra de 50 ejemplos pasó de 0,60 a 0,80; aunque con un tamaño así, el margen de error es amplio y una línea base “siempre yes” ya ronda el 55% en este conjunto de evaluación. 

También conviene señalar que el subconjunto de entrenamiento artificial contiene casi ninguna etiqueta "maybe", por lo que el modelo apenas tiene oportunidad de aprender esa clase aunque aparezca en los datos de evaluación. Tómalo como un experimento rápido para entender el proceso de ajuste fino, no como prueba de que el modelo esté listo para uso médico real.

¿Listo para ir más allá de un único ajuste fino? Nuestro itinerario Developing Large Language Models te lleva desde los fundamentos de PyTorch y transformers hasta crear y desplegar tus propios LLMs.