跳至内容

如何使用 Unsloth 在 PubMedQA 上微调 DiffusionGemma

学习如何使用 Unsloth、LoRA、Hugging Face 和 RunPod H100 GPU,在 PubMedQA 上微调 DiffusionGemma 26B-A4B,以完成生物医学问答。
更新 2026年8月3日  · 9分钟

用 AI 探索

在 ChatGPT 中打开在 Claude 中打开在 Perplexity 中打开

DiffusionGemma来自 Google DeepMind 的一款实验性语言模型,其生成文本的方式不同于传统的大型语言模型。它并非按从左到右逐个预测 token,而是从一个固定的画布(由噪声 token 组成)出发,并通过多次去噪逐步精炼。这使模型能够并行更新多个 token 位置,并在生成过程中修订部分回答。

在本指南中,我们将使用 微调 diffusiongemma-26B-A4B-it,在 NVIDIA H100 GPU 上基于 PubMedQA 数据集进行训练。模型将接收生物医学问题与支撑性上下文,并预测yesnomaybe。我们将准备数据、训练一个 LoRA 适配器,在微调前后对模型进行评估,并将最终适配器上传到 Hugging Face。

我也发布了完整的笔记本,方便您查看原始代码、跟进步骤并自行运行实验。

注意:本项目仅用于学习与实验,不应用于实际医疗决策。

1. 打开 RunPod Jupyter Notebook

创建一个新的 RunPod pod,使用 NVIDIA H100 GPU,并选择 PyTorch/Jupyter 模板。至少配置100 GB 的持久化存储,以避免 pod 停止时丢失模型文件和训练输出。

将您的 Hugging Face 访问令牌添加为环境变量:

HF_TOKEN=your_hugging_face_token

Editing the Runpod Pytorch template

这样可以更快下载模型与数据集,并允许您在不手动登录笔记本的情况下将已保存的 LoRA 适配器上传到 Hugging Face。

按此配置的 pod 费用约为每小时 3 美元,但最终价格可能因 GPU 供给和所选 pod 类型而异。

Runpod H100 GPU summary

当 pod 运行后,从 RunPod 界面打开 JupyterLab 或 Jupyter Notebook,并创建一个名为 diffusiongemma_pubmedqa.ipynb 的新笔记本。

2. 安装所需软件包

在第一个笔记本单元中运行以下命令,安装 Unsloth 以及加载、微调和保存 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"

%%capture 命令会隐藏冗长的安装输出。为避免 DiffusionGemma、Transformers、Unsloth 与训练库之间的兼容性问题,我们固定了软件包版本。

安装完成后,请重启笔记本内核再继续。

3. 导入库

导入准备数据集、加载模型、训练与评估所需的库。

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 可加速从 Hugging Face Hub 的下载;提高 Dynamo 的重新编译上限,有助于在与模型交互时避免中断。

当单元成功运行时,Unsloth 会对训练环境进行补丁处理,输出应确认 CUDA 可用且检测到 H100 GPU。

🦥 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. 设置配置

在一个位置定义模型、数据集、训练参数、评估设置和输出目录。

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"

我们将使用 3,000 个人工样本进行训练,使用 200 个人工标注样本进行评估。为加快实验,模型训练 60 步,并在评估中使用 16 个去噪步处理 50 个示例。

DiffusionGemma 的工作原理

在加载模型前,值得先了解 DiffusionGemma 实际如何生成答案。它不是逐个写出 token,而是以一个固定长度的画布开始,通过多次去噪逐步精炼,同时更新多个位置,直到文本演化为连贯的回答。

下面的代码会报告画布长度为 256,即单个区块的大小。对于我们短小的 yes/no/maybe 回答,一个画布已绰绰有余;更长的输出则通过区块接龙的方式逐块生成。下图高层示意了这一精炼过程:

Text Diffusion Explanation

5. 加载 DiffusionGemma

bfloat16 精度加载经过指令微调的 DiffusionGemma 模型。我们不使用 4-bit 量化,因为 H100 GPU 具备足够显存以较高精度加载模型。

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)

词表大小用于在扩散训练过程中添加随机噪声;画布长度决定模型在单个生成区块中可精炼的最大 token 数。

您应会看到类似如下输出:

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

6. 添加 LoRA 适配器

添加 LoRA 适配器,使训练仅更新一小部分新增参数,而非整个 260 亿参数的模型。

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

这能显著降低微调所需的内存与算力。已关闭梯度检查点,因为 H100 在本实验中拥有足够的显存。

7. 加载 PubMedQA

加载用于训练的人工 PubMedQA 子集,以及用于评估的人工标注子集。

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

训练子集包含自动生成的样本,而评估子集包含由专家标注的生物医学问题。

打印首行数据便于在格式化数据前检查问题、摘要上下文与最终决策。

您应会看到:

Train size: 211269
Eval size: 1000

每个样本包含一个生物医学问题、一个或多个支撑的摘要片段,以及一个最终答案(yesnomaybe)。

Loading the PubMedQA

8. 转换数据集

将每个 PubMedQA 样本转换为包含用户提示与助手回答的聊天式格式。

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

上下文片段会合并为一个字符串,并限制为 2,500 个字符以保持输入可控。每个答案都会转为小写,且会移除不在 yesnomaybe 之外的标签。

打印第一个转换后的样本,有助于在训练前确认上下文、问题与答案已正确格式化。

Formatting the PubMedQA dataset in chat format

9. 构建扩散训练样本

DiffusionGemma 要求将目标答案放入固定长度的画布中。该函数会对提示进行分词、将答案转换为 token ID、将其填充到模型的画布长度,并创建一个掩码以标明哪些 token 应参与损失计算。

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)

每个答案之后会添加序列结束 token,画布剩余位置由 padding token 填充。损失掩码确保训练仅关注答案与序列结束 token,而非填充位置。

10. 创建推理与评估函数

接下来,定义用于生成答案、清洗模型输出并计算评估准确率的函数。

生成答案

answer_question() 函数会格式化提示,通过若干去噪步生成响应,并将生成的 token 解码为文本。

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

提取预测

尽管提示要求单词作答,模型有时仍可能生成额外文本。此函数会提取第一个有效的 yesnomaybe 预测。

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"

评估准确率

评估函数会将每个清洗后的预测与正确答案进行比较,打印每个示例的结果,并返回整体准确率及各条预测结果。

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. 微调前评估模型

在训练前先运行评估以建立基线。

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

此处对 50 个示例进行评估,每个答案使用 16 个去噪步。在本实验中,基础模型在 50 个问题中答对了 30 个。

Baseline result of the DiffusionGemma

稍后我们将此基线与微调后的模型准确率进行比较。

12. 设置训练

将模型切换到训练模式,创建优化器与学习率调度器,并定义在扩散训练中如何损坏干净的答案 token。

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

仅将 requires_grad=True 的参数传入优化器,这意味着训练过程更新的是 LoRA 适配器而非完整模型。

接着,创建一个损坏函数,用随机 token 替换答案画布中随机比例的内容。

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)

每个样本的噪声量都会变化。训练中,模型学习从这些被损坏的画布 token 中重构原始答案。

13. 训练模型

以下循环使用梯度累积训练 LoRA 适配器 60 步。

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

对于每个训练样本,模型接收生物医学提示与一个被损坏的答案画布。交叉熵损失只对损失掩码选中的真实答案 token 进行计算。

梯度累积会在更新模型前合并四个样本的梯度。还应用了梯度裁剪以保持训练稳定。

在本实验中,训练约耗时两分钟:

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

损失值稳步下降表明适配器正在学习从损坏的画布中重构期望答案。训练期间,您也可以在 RunPod 终端运行 nvidia-smi 来监控 GPU 显存使用与利用率。

Nvidia model statistic while model training

14. 评估微调后的模型

在训练后再次运行相同评估,以衡量微调是否提升了模型表现。

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

微调后的模型在与基线相同的 50 个示例上评估,且同样使用 16 个去噪步。

Fine-tune result of the DiffusionGemma

接下来比较微调前后的准确率。

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

在本实验中,模型准确率从 0.60 提升至 0.80。

这意味着提升了 20 个百分点:微调后模型在 50 个问题中答对 40 个,而训练前为 30 个。

15. 保存并上传微调适配器

将已训练的 LoRA 适配器和处理器文件保存到先前定义的输出目录。

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

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

您应会看到:

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

这样只会保存轻量的 LoRA 适配器,而不是再保存一份 260 亿参数的基础模型副本。

接下来,将适配器与处理器文件上传到 Hugging Face Hub:

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

由于在 RunPod pod 配置时已添加 HF_TOKEN 环境变量,Hugging Face 应会自动完成身份验证。仅在尚未配置令牌时,才需要运行 notebook_login()

from huggingface_hub import notebook_login

notebook_login()

上传完成后,代码库将包含 LoRA 适配器与处理器配置,便于后续加载微调模型。

Finetuned lora on Hugging face: diffusiongemma_pubmedqa

来源kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

结语

使用 Unsloth 微调 DiffusionGemma 出乎意料地简单。最耗时的部分是安装正确的依赖,以及理解扩散特有的训练流程。一旦环境就绪,加载模型、训练 LoRA 适配器、评估并将结果上传到 Hugging Face 都非常顺畅。

我之所以对 DiffusionGemma 特别感兴趣,是因为它不像传统语言模型那样逐 token 生成文本,而是处理一个固定画布,并通过去噪步骤逐步精炼噪声 token。理解这种不同的生成过程,并将其在生物医学问答任务上进行微调,让这次实验对我而言格外有价值。

即便是在小规模设置下,50 个样本上的准确率也从 0.60 升至 0.80——不过在如此小的样本量下,误差幅度较大,且在该评估集上“始终回答 yes”的基线已能达到约 55%。

还需注意,人工训练子集中几乎没有 “maybe” 标签,因此即便评估数据中出现该类别,模型也很难学到。请将此视为快速了解微调流程的实验,而非说明模型已能用于真实医疗场景。

准备好超越单次微调运行了吗?我们的Developing Large Language Models 学习路径将从 PyTorch 与 transformers 基础带您走向自研与部署 LLM。

主题

Top LLM Courses

Tracks

开发大型语言模型

16小时
学习使用 PyTorch 和 Hugging Face 开发大型语言模型(LLMs),运用最新的深度学习和 NLP 技术。
查看详情Right Arrow
开始课程
查看更多Right Arrow