Como ajustar o DiffusionGemma no PubMedQA usando Unsloth

1. Abra um Jupyter Notebook na RunPod

Crie um novo pod na RunPod com uma GPU NVIDIA H100 e selecione um template de PyTorch/Jupyter. Configure pelo menos 100 GB de armazenamento persistente para que os arquivos do modelo e os resultados de treino não se percam quando o pod for encerrado.

Adicione seu token de acesso do Hugging Face como variável de ambiente:

HF_TOKEN=your_hugging_face_token

Isso acelera o download de modelos e datasets e permite enviar o adaptador LoRA salvo para o Hugging Face sem precisar fazer login manualmente pelo notebook.

O pod configurado deve custar aproximadamente US$ 3 por hora, embora o preço final possa variar conforme a disponibilidade da GPU e o tipo de pod selecionado.

Quando o pod estiver rodando, abra o JupyterLab ou Jupyter Notebook pela interface da RunPod e crie um novo notebook chamado diffusiongemma_pubmedqa.ipynb.

2. Instale os pacotes necessários

Execute os comandos abaixo na primeira célula do notebook para instalar o Unsloth e as bibliotecas necessárias para carregar, ajustar e salvar o 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"

O comando %%capture oculta a saída extensa da instalação. As versões estão fixadas para evitar problemas de compatibilidade entre DiffusionGemma, Transformers, Unsloth e as bibliotecas de treino.

Quando a instalação terminar, reinicie o kernel do notebook antes de continuar.

3. Importe as bibliotecas

Importe as bibliotecas necessárias para preparar o dataset, carregar o modelo, treinar e avaliar.

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 habilita downloads mais rápidos do Hugging Face Hub, enquanto aumentar o limite de recompilação do Dynamo ajuda a evitar interrupções ao trabalhar com o modelo.

Quando a célula rodar com sucesso, o Unsloth fará patches no ambiente de treino, e a saída deve confirmar que o CUDA está disponível e que a GPU H100 foi detectada.

🦥 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. Defina a configuração

Defina, em um só lugar, o modelo, o dataset, os parâmetros de treino, as configurações de avaliação e o diretório de saída.

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 exemplos artificiais para treino e 200 exemplos rotulados manualmente para avaliação. Para manter o experimento ágil, o modelo vai treinar por 60 etapas e avaliar 50 exemplos usando 16 etapas de denoising.

Como o DiffusionGemma funciona

Antes de carregar o modelo, vale visualizar como o DiffusionGemma realmente produz uma resposta. Em vez de escrever tokens em sequência, ele começa com uma tela de comprimento fixo e a refina ao longo de várias etapas de denoising, atualizando muitas posições de uma vez até chegar a um texto coerente. 

O código abaixo informa uma tela de comprimento 256, o tamanho de um único bloco. Para nossas respostas curtas de sim/não/talvez, uma tela é mais do que suficiente, enquanto saídas mais longas são geradas encadeando telas bloco a bloco. O diagrama abaixo mostra esse processo de refinamento em alto nível:

5. Carregue o DiffusionGemma

Carregue o modelo DiffusionGemma ajustado para instruções em precisão bfloat16. Não usaremos quantização em 4 bits porque a H100 tem memória suficiente para carregar o modelo em maior precisão.

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)

O tamanho do vocabulário é usado ao adicionar ruído aleatório durante o treino por difusão. O comprimento da tela determina o número máximo de tokens que o modelo consegue refinar em um único bloco de geração.

Você deve ver uma saída semelhante a:

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

6. Adicione um adaptador LoRA

Adicione um adaptador LoRA para treinar apenas um conjunto pequeno de parâmetros adicionais em vez de atualizar todo o modelo de 26 bilhões de parâmetros.

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

Isso reduz significativamente a memória e o computacional necessários para o fine-tuning. O gradient checkpointing está desativado porque a H100 tem memória de GPU suficiente para este experimento.

7. Carregue o PubMedQA

Carregue o subconjunto artificial do PubMedQA para treino e o subconjunto rotulado manualmente para avaliação.

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

O subconjunto de treino contém exemplos gerados automaticamente, enquanto o de avaliação contém perguntas biomédicas rotuladas por especialistas. 

Imprimir a primeira linha permite inspecionar a pergunta, o contexto do resumo e a decisão final antes de formatar os dados.

Você deve ver:

Train size: 211269
Eval size: 1000

Cada exemplo contém uma pergunta biomédica, um ou mais trechos de resumos de apoio e uma resposta final de yes, no ou maybe.

8. Converta o dataset

Converta cada exemplo do PubMedQA em um formato de chat contendo um prompt do usuário e a resposta do 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"])

Os trechos de contexto são combinados em uma única string e limitados a 2.500 caracteres para manter a entrada gerenciável. Cada resposta é convertida para minúsculas, e exemplos com rótulos fora de yes, no ou maybe são removidos.

Imprimir o primeiro exemplo convertido ajuda a confirmar que o contexto, a pergunta e a resposta foram formatados corretamente antes do treino.

9. Construa os exemplos para treino por difusão

O DiffusionGemma exige que a resposta alvo fique dentro de uma tela de comprimento fixo. Esta função tokeniza o prompt, converte a resposta em IDs de tokens, preenche até o comprimento da tela do modelo e cria uma máscara indicando quais tokens entram no cálculo da perda.

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)

O token de fim de sequência é adicionado após cada resposta, e as posições restantes da tela são preenchidas com tokens de padding. A máscara de perda garante que o treino foque apenas nos tokens da resposta e no token de fim de sequência, e não nas posições preenchidas.

10. Crie as funções de inferência e avaliação

Em seguida, defina as funções usadas para gerar respostas, limpar a saída do modelo e calcular a acurácia de avaliação.

Gerar uma resposta

A função answer_question() formata o prompt, gera uma resposta por várias etapas de denoising e decodifica os tokens gerados em 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()

Extrair a previsão

Embora o prompt solicite uma resposta de uma palavra, às vezes o modelo pode gerar texto adicional. Esta função extrai a primeira previsão válida de 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"

Avaliar a acurácia

A função de avaliação compara cada previsão limpa com a resposta correta, imprime o resultado de cada exemplo e retorna a acurácia geral junto com as previsões individuais.

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. Avalie o modelo antes do fine-tuning

Execute a avaliação antes do treino para estabelecer uma linha de base.

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

Isso avalia 50 exemplos usando 16 etapas de denoising por resposta. Neste experimento, o modelo base respondeu corretamente 30 de 50 perguntas.

Essa linha de base será comparada depois com a acurácia do modelo após o fine-tuning.

12. Configure o treino

Coloque o modelo em modo de treino, crie o otimizador e o scheduler de taxa de aprendizado e defina como os tokens de resposta limpos serão corrompidos durante o treino por difusão.

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

Apenas parâmetros com requires_grad=True são passados ao otimizador, o que significa que o processo de treino atualiza o adaptador LoRA e não o modelo completo.

Em seguida, crie uma função de corrupção que substitui uma proporção aleatória da tela de resposta por tokens aleatórios.

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)

A quantidade de ruído muda a cada exemplo. Durante o treino, o modelo aprende a reconstruir a resposta original a partir desses tokens corrompidos na tela.

13. Treine o modelo

O loop a seguir treina o adaptador LoRA por 60 etapas usando acumulação de 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"
        )

Para cada exemplo de treino, o modelo recebe o prompt biomédico e uma tela de resposta corrompida. A loss de entropia cruzada é calculada apenas para os tokens reais da resposta selecionados pela máscara de perda.

A acumulação de gradiente combina quatro exemplos antes de atualizar o modelo. O clipping de gradiente também é aplicado para manter o treino estável.

Neste experimento, o treino foi concluído em cerca de dois minutos:

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

A redução constante da perda indica que o adaptador está aprendendo a reconstruir as respostas esperadas a partir da tela corrompida. Durante o treino, você também pode executar nvidia-smi no terminal da RunPod para monitorar o uso e a utilização da memória da GPU.

14. Avalie o modelo ajustado

Execute a mesma avaliação após o treino para medir se o fine-tuning melhorou o desempenho do modelo.

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

O modelo ajustado é avaliado nos mesmos 50 exemplos e usando as mesmas 16 etapas de denoising da avaliação de base.

Em seguida, compare a acurácia antes e depois do ajuste.

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

Neste experimento, a acurácia do modelo subiu de 0,60 para 0,80.

Isso representa um ganho de 20 pontos percentuais, com o modelo respondendo corretamente 40 de 50 perguntas após o fine-tuning, contra 30 de 50 antes do treino.

15. Salve e envie o adaptador ajustado

Salve o adaptador LoRA treinado e os arquivos do processor no diretório de saída definido anteriormente.

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

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

Você deve ver:

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

Isso salva apenas o adaptador LoRA leve, e não outra cópia completa do modelo base de 26 bilhões de parâmetros.

Em seguida, envie o adaptador e os arquivos do processor para o Hugging Face Hub:

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

Como a variável de ambiente HF_TOKEN foi adicionada durante a configuração do pod na RunPod, o Hugging Face deve autenticar automaticamente. Você só precisa executar notebook_login() quando o token ainda não estiver configurado:

from huggingface_hub import notebook_login

notebook_login()

Depois que o upload terminar, o repositório terá o adaptador LoRA e a configuração do processor necessários para carregar o modelo ajustado posteriormente.

Fonte: kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

Considerações finais

Fazer o fine-tuning do DiffusionGemma com o Unsloth foi surpreendentemente simples. A parte mais demorada foi instalar as dependências corretas e entender como funciona o processo de treino específico de difusão. Depois que o ambiente estava pronto, carregar o modelo, treinar o adaptador LoRA, avaliá-lo e enviar os resultados para o Hugging Face foi bem tranquilo.

Achei o DiffusionGemma especialmente interessante porque ele não gera texto token a token como um modelo tradicional. Em vez disso, trabalha com uma tela fixa e vai refinando tokens ruidosos por etapas de denoising. Entender esse processo de geração diferente e ajustá-lo para a tarefa de perguntas e respostas biomédicas tornou o experimento especialmente valioso para mim.

Mesmo com uma configuração pequena, a acurácia na amostra de 50 exemplos subiu de 0,60 para 0,80 — embora, com um tamanho de amostra assim, a margem de erro seja grande, e um baseline de "sempre sim" já marque cerca de 55% nesse conjunto de avaliação. 

Também vale notar que o subconjunto artificial de treino tem quase nenhum rótulo "maybe", então o modelo quase não tem chance de aprender essa classe, mesmo ela aparecendo nos dados de avaliação. Trate isto como um experimento rápido para entender o processo de fine-tuning, não como evidência de que o modelo está pronto para uso médico real.

Pronto para ir além de um único rodada de ajuste? Nossa trilha Developing Large Language Models leva você dos fundamentos de PyTorch e transformers à construção e implantação dos seus próprios LLMs.