Cách fine-tune DiffusionGemma trên PubMedQA bằng Unsloth

1. Mở RunPod Jupyter Notebook

Tạo một pod RunPod mới với GPU NVIDIA H100 và chọn mẫu PyTorch/Jupyter. Cấu hình tối thiểu 100 GB bộ nhớ lưu trữ cố định để tệp mô hình và đầu ra huấn luyện không bị mất khi pod dừng.

Thêm token truy cập Hugging Face của bạn làm biến môi trường:

HF_TOKEN=your_hugging_face_token

Điều này cho phép tải nhanh hơn các mô hình và bộ dữ liệu, đồng thời giúp bạn tải lên bộ chuyển tiếp LoRA đã lưu lên Hugging Face mà không cần đăng nhập thủ công từ notebook.

Pod đã cấu hình sẽ có chi phí khoảng 3 đô-la mỗi giờ, tuy nhiên giá cuối cùng có thể thay đổi tùy theo tình trạng sẵn có của GPU và loại pod được chọn.

Khi pod đã chạy, mở JupyterLab hoặc Jupyter Notebook từ giao diện RunPod và tạo một notebook mới tên diffusiongemma_pubmedqa.ipynb.

2. Cài đặt các gói cần thiết

Chạy các lệnh sau trong ô đầu tiên của notebook để cài Unsloth và các thư viện cần cho việc tải, fine-tune và lưu 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"

Lệnh %%capture ẩn phần xuất cài đặt dài. Phiên bản các gói được cố định để tránh xung đột tương thích giữa DiffusionGemma, Transformers, Unsloth và các thư viện huấn luyện.

Sau khi cài đặt xong, hãy khởi động lại kernel của notebook trước khi tiếp tục.

3. Import các thư viện

Import các thư viện cần cho việc chuẩn bị dữ liệu, tải mô hình, huấn luyện và đánh giá.

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 cho phép tải xuống từ Hugging Face Hub nhanh hơn, trong khi tăng giới hạn biên dịch lại của Dynamo giúp tránh gián đoạn khi làm việc với mô hình.

Khi ô chạy thành công, Unsloth sẽ vá môi trường huấn luyện, và đầu ra sẽ xác nhận CUDA khả dụng cũng như GPU H100 đã được nhận diện.

🦥 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. Thiết lập cấu hình

Định nghĩa mô hình, bộ dữ liệu, tham số huấn luyện, thiết lập đánh giá và thư mục đầu ra ở một nơi.

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"

Chúng ta sẽ dùng 3.000 ví dụ nhân tạo cho huấn luyện và 200 ví dụ gán nhãn thủ công cho đánh giá. Để thí nghiệm nhanh, mô hình sẽ huấn luyện 60 bước và đánh giá 50 ví dụ với 16 bước khử nhiễu.

Cách DiffusionGemma hoạt động

Trước khi tải mô hình, đáng để hình dung cách DiffusionGemma thực sự tạo ra câu trả lời. Thay vì viết token tuần tự, mô hình bắt đầu với một canvas độ dài cố định các token và tinh chỉnh chúng qua nhiều bước khử nhiễu, cập nhật nhiều vị trí cùng lúc cho đến khi văn bản trở nên mạch lạc. 

Đoạn mã dưới báo cáo độ dài canvas là 256, kích thước của một khối. Với các câu trả lời ngắn yes/no/maybe, một canvas là quá đủ, trong khi đầu ra dài hơn được sinh bằng cách xâu chuỗi các canvas nối tiếp từng khối. Sơ đồ dưới đây minh họa quá trình tinh chỉnh ở mức khái quát:

5. Tải DiffusionGemma

Tải mô hình DiffusionGemma đã tinh chỉnh theo hướng dẫn ở độ chính xác bfloat16. Chúng ta sẽ không dùng lượng tử hóa 4-bit vì GPU H100 có đủ bộ nhớ để tải mô hình ở độ chính xác cao hơ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)

Kích thước từ vựng được dùng khi thêm nhiễu ngẫu nhiên trong huấn luyện khuếch tán. Độ dài canvas quyết định số token tối đa mà mô hình có thể tinh chỉnh trong một khối sinh duy nhất.

Bạn sẽ thấy đầu ra tương tự:

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

6. Thêm bộ chuyển tiếp LoRA

Thêm bộ chuyển tiếp LoRA để chỉ huấn luyện một tập tham số nhỏ thay vì cập nhật toàn bộ mô hình 26 tỷ tham số.

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

Điều này giúp giảm đáng kể bộ nhớ và tính toán cần cho fine-tuning. Tắt gradient checkpointing vì H100 có đủ bộ nhớ GPU cho thí nghiệm này.

7. Tải PubMedQA

Tải tập con PubMedQA nhân tạo để huấn luyện và tập gán nhãn thủ công để đánh giá.

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

Tập con huấn luyện chứa các ví dụ được tạo tự động, trong khi tập đánh giá chứa các câu hỏi y sinh học được chuyên gia gán nhãn. 

In hàng đầu tiên giúp chúng ta kiểm tra câu hỏi, ngữ cảnh tóm tắt và đáp án cuối cùng trước khi định dạng dữ liệu.

Bạn sẽ thấy:

Train size: 211269
Eval size: 1000

Mỗi ví dụ gồm một câu hỏi y sinh, một hoặc nhiều đoạn tóm tắt hỗ trợ, và đáp án cuối cùng là yes, no hoặc maybe.

8. Chuyển đổi bộ dữ liệu

Chuyển mỗi ví dụ PubMedQA sang định dạng hội thoại gồm một lời nhắc từ người dùng và câu trả lời của trợ lý.

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

Các đoạn ngữ cảnh được gộp thành một chuỗi và giới hạn ở 2.500 ký tự để đầu vào gọn hơn. Mỗi đáp án được chuyển thành chữ thường, và các ví dụ có nhãn ngoài yes, no hoặc maybe sẽ bị loại.

In ví dụ đã chuyển đổi đầu tiên giúp xác nhận ngữ cảnh, câu hỏi và câu trả lời đã được định dạng đúng trước khi huấn luyện.

9. Xây dựng ví dụ huấn luyện khuếch tán

DiffusionGemma yêu cầu đặt đáp án mục tiêu vào bên trong một canvas có độ dài cố định. Hàm này tokenize lời nhắc, chuyển đáp án thành ID token, đệm đến độ dài canvas của mô hình, và tạo một mặt nạ cho biết token nào góp phần vào 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)

Token kết thúc chuỗi được thêm sau mỗi đáp án, trong khi các vị trí còn lại của canvas được điền bằng token đệm. Mặt nạ loss đảm bảo huấn luyện chỉ tập trung vào các token của đáp án và token kết thúc chuỗi thay vì vị trí đệm.

10. Tạo hàm suy luận và đánh giá

Tiếp theo, định nghĩa các hàm dùng để sinh đáp án, làm sạch đầu ra của mô hình và tính độ chính xác đánh giá.

Sinh một câu trả lời

Hàm answer_question() định dạng lời nhắc, sinh phản hồi qua một số bước khử nhiễu, và giải mã token sinh ra thành văn bản.

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

Trích xuất dự đoán

Dù lời nhắc yêu cầu một từ đơn, đôi khi mô hình có thể sinh thêm văn bản. Hàm này trích xuất dự đoán hợp lệ đầu tiên là yes, no hoặc 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"

Đánh giá độ chính xác

Hàm đánh giá so sánh mỗi dự đoán đã làm sạch với đáp án đúng, in kết quả cho từng ví dụ, và trả về độ chính xác tổng thể kèm các dự đoán riêng lẻ.

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. Đánh giá mô hình trước khi fine-tuning

Chạy đánh giá trước khi huấn luyện để thiết lập đường cơ sở.

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

Đánh giá này chấm 50 ví dụ với 16 bước khử nhiễu mỗi câu trả lời. Trong thí nghiệm này, mô hình gốc trả lời đúng 30/50 câu hỏi.

Đường cơ sở này sau đó sẽ được so sánh với độ chính xác của mô hình sau khi fine-tuning.

12. Thiết lập huấn luyện

Chuyển mô hình sang chế độ huấn luyện, tạo optimizer và lịch học (scheduler), và định nghĩa cách token đáp án sạch sẽ bị làm nhiễu trong quá trình huấn luyện khuếch tá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",
)

Chỉ các tham số có requires_grad=True được truyền vào optimizer, nghĩa là quá trình huấn luyện cập nhật bộ chuyển tiếp LoRA thay vì toàn bộ mô hình.

Tiếp theo, tạo một hàm làm nhiễu thay thế một tỷ lệ ngẫu nhiên của canvas đáp án bằng các token ngẫu nhiên.

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)

Mức nhiễu thay đổi ở mỗi ví dụ. Trong huấn luyện, mô hình học cách khôi phục đáp án gốc từ các token canvas đã bị nhiễu này.

13. Huấn luyện mô hình

Vòng lặp sau huấn luyện bộ chuyển tiếp LoRA trong 60 bước với tích lũy gradient.

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

Với mỗi ví dụ huấn luyện, mô hình nhận lời nhắc y sinh và một canvas đáp án đã bị nhiễu. Loss cross-entropy chỉ được tính cho các token đáp án thật do mặt nạ loss lựa chọn.

Tích lũy gradient gộp bốn ví dụ trước khi cập nhật mô hình. Cắt gradient cũng được áp dụng để giữ huấn luyện ổn định.

Trong thí nghiệm này, huấn luyện hoàn tất khoảng hai phút:

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

Loss giảm đều cho thấy bộ chuyển tiếp đang học cách tái tạo các đáp án kỳ vọng từ canvas bị nhiễu. Trong khi huấn luyện, bạn cũng có thể chạy nvidia-smi trong terminal RunPod để theo dõi bộ nhớ và mức sử dụng GPU.

14. Đánh giá mô hình đã fine-tune

Chạy lại đánh giá sau huấn luyện để đo lường liệu fine-tuning có cải thiện hiệu năng mô hình không.

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

Mô hình đã fine-tune được đánh giá trên cùng 50 ví dụ và dùng 16 bước khử nhiễu như ở đánh giá đường cơ sở.

Tiếp theo, so sánh độ chính xác trước và sau 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

Trong thí nghiệm này, độ chính xác của mô hình tăng từ 0,60 lên 0,80.

Điều này tương đương mức cải thiện 20 điểm phần trăm, với 40/50 câu trả lời đúng sau fine-tuning, so với 30/50 trước huấn luyện.

15. Lưu và tải lên bộ chuyển tiếp đã fine-tune

Lưu bộ chuyển tiếp LoRA đã huấn luyện và các tệp processor vào thư mục đầu ra đã định trước.

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

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

Bạn sẽ thấy:

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

Thao tác này chỉ lưu bộ chuyển tiếp LoRA gọn nhẹ thay vì một bản sao đầy đủ khác của mô hình nền 26 tỷ tham số.

Tiếp theo, tải bộ chuyển tiếp và tệp processor lên Hugging Face Hub:

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

Vì biến môi trường HF_TOKEN đã được thêm khi cấu hình pod RunPod, Hugging Face sẽ tự động xác thực. Bạn chỉ cần chạy notebook_login() khi token chưa được cấu hình sẵn:

from huggingface_hub import notebook_login

notebook_login()

Sau khi tải lên hoàn tất, kho lưu trữ sẽ chứa bộ chuyển tiếp LoRA và cấu hình processor cần để tải mô hình đã fine-tune sau này.

Nguồn: kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

Suy nghĩ cuối

Fine-tuning DiffusionGemma bằng Unsloth diễn ra bất ngờ là dễ dàng. Phần tốn thời gian nhất là cài đúng phụ thuộc và hiểu cách hoạt động của quy trình huấn luyện đặc thù cho khuếch tán. Khi môi trường đã thiết lập xong, việc tải mô hình, huấn luyện bộ chuyển tiếp LoRA, đánh giá và tải kết quả lên Hugging Face đều rất trơn tru.

Tôi thấy DiffusionGemma đặc biệt thú vị vì nó không sinh văn bản từng token một như mô hình ngôn ngữ truyền thống. Thay vào đó, nó làm việc với một canvas cố định và dần tinh chỉnh các token nhiễu qua các bước khử nhiễu. Hiểu quy trình sinh khác biệt này và fine-tune nó cho nhiệm vụ hỏi đáp y sinh khiến thí nghiệm đặc biệt giá trị với tôi.

Ngay cả với thiết lập nhỏ, độ chính xác trên mẫu 50 ví dụ tăng từ 0,60 lên 0,80 — dù với cỡ mẫu như vậy, sai số biên khá lớn, và một đường cơ sở "luôn trả lời yes" đã đạt khoảng 55% trên tập đánh giá này. 

Cũng cần lưu ý rằng tập con huấn luyện nhân tạo hầu như không có nhãn "maybe", nên mô hình ít có cơ hội học lớp này dù nó xuất hiện trong dữ liệu đánh giá. Hãy coi đây là một thí nghiệm nhanh về cách hoạt động của quá trình fine-tuning, không phải bằng chứng cho thấy mô hình sẵn sàng dùng trong y khoa thực tế.

Sẵn sàng vượt ra ngoài một lần fine-tune? Lộ trình Phát triển Mô hình Ngôn ngữ Lớn của chúng tôi sẽ đưa bạn từ nền tảng PyTorch và transformers đến xây dựng và triển khai LLM của riêng bạn.