Cara Fine-tune DiffusionGemma pada PubMedQA Menggunakan Unsloth

1. Buka RunPod Jupyter Notebook

Buat pod RunPod baru dengan GPU NVIDIA H100 dan pilih template PyTorch/Jupyter. Konfigurasikan setidaknya 100 GB penyimpanan persisten agar file model dan keluaran pelatihan Anda tidak hilang saat pod berhenti.

Tambahkan token akses Hugging Face Anda sebagai variabel lingkungan:

HF_TOKEN=your_hugging_face_token

Ini memungkinkan model dan dataset diunduh lebih cepat dan memudahkan Anda mengunggah adapter LoRA yang disimpan ke Hugging Face tanpa masuk secara manual dari notebook.

Pod yang dikonfigurasi diperkirakan berbiaya sekitar $3 per jam, meskipun harga akhir dapat bervariasi tergantung pada ketersediaan GPU dan jenis pod yang dipilih.

Setelah pod berjalan, buka JupyterLab atau Jupyter Notebook dari antarmuka RunPod dan buat notebook baru bernama diffusiongemma_pubmedqa.ipynb.

2. Instal Paket yang Diperlukan

Jalankan perintah berikut di sel notebook pertama untuk memasang Unsloth dan pustaka yang diperlukan untuk memuat, fine-tune, dan menyimpan 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"

Perintah %%capture menyembunyikan keluaran instalasi yang panjang. Versi paket dipatok untuk menghindari masalah kompatibilitas antara DiffusionGemma, Transformers, Unsloth, dan pustaka pelatihan.

Setelah instalasi selesai, mulai ulang kernel notebook sebelum melanjutkan.

3. Impor Pustaka

Impor pustaka yang diperlukan untuk menyiapkan dataset, memuat model, melatih, dan mengevaluasi.

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 mengaktifkan unduhan yang lebih cepat dari Hugging Face Hub, sementara peningkatan batas re-kompilasi Dynamo membantu mencegah interupsi saat bekerja dengan model.

Saat sel berjalan dengan sukses, Unsloth akan menambal lingkungan pelatihan, dan keluaran harus mengonfirmasi bahwa CUDA tersedia dan GPU H100 terdeteksi.

🦥 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. Atur Konfigurasi

Tentukan model, dataset, parameter pelatihan, pengaturan evaluasi, dan direktori keluaran di satu tempat.

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"

Kita akan menggunakan 3.000 contoh artifisial untuk pelatihan dan 200 contoh berlabel manual untuk evaluasi. Agar eksperimen tetap cepat, model akan dilatih selama 60 langkah dan mengevaluasi 50 contoh menggunakan 16 langkah denoising.

Cara Kerja DiffusionGemma

Sebelum memuat model, ada baiknya membayangkan bagaimana DiffusionGemma benar-benar menghasilkan jawaban. Alih-alih menulis token satu per satu, model ini mulai dengan kanvas berukuran tetap berisi token dan menyempurnakannya selama beberapa langkah denoising, memperbarui banyak posisi sekaligus hingga teks menjadi respons yang koheren. 

Kode di bawah melaporkan panjang kanvas 256, ukuran satu blok. Untuk jawaban singkat ya/tidak/mungkin, satu kanvas sudah lebih dari cukup, sedangkan keluaran yang lebih panjang dihasilkan dengan merangkaikan kanvas blok demi blok. Diagram di bawah menunjukkan proses penyempurnaan ini pada tingkat tinggi:

5. Muat DiffusionGemma

Muat model DiffusionGemma yang telah di-tuning untuk instruksi dalam presisi bfloat16. Kita tidak akan menggunakan kuantisasi 4-bit karena GPU H100 memiliki memori yang cukup untuk memuat model pada presisi lebih tinggi.

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)

Ukuran kosakata digunakan saat menambahkan noise acak selama pelatihan difusi. Panjang kanvas menentukan jumlah maksimum token yang dapat disempurnakan model dalam satu blok generasi.

Anda akan melihat keluaran serupa dengan:

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

6. Tambahkan Adapter LoRA

Tambahkan adapter LoRA sehingga hanya sejumlah kecil parameter tambahan yang dilatih alih-alih memperbarui seluruh model 26 miliar parameter.

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

Ini secara signifikan mengurangi memori dan komputasi yang diperlukan untuk fine-tuning. Gradient checkpointing dinonaktifkan karena H100 memiliki memori GPU yang cukup untuk eksperimen ini.

7. Muat PubMedQA

Muat subset PubMedQA artifisial untuk pelatihan dan subset berlabel manual untuk evaluasi.

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

Subset pelatihan berisi contoh yang dihasilkan secara otomatis, sedangkan subset evaluasi berisi pertanyaan biomedis berlabel pakar. 

Mencetak baris pertama memungkinkan kita memeriksa pertanyaan, konteks abstrak, dan keputusan akhir sebelum memformat data.

Anda akan melihat:

Train size: 211269
Eval size: 1000

Setiap contoh berisi pertanyaan biomedis, satu atau lebih potongan abstrak pendukung, dan jawaban akhir yes, no, atau maybe.

8. Konversi Dataset

Konversikan setiap contoh PubMedQA ke format gaya percakapan yang berisi prompt pengguna dan jawaban asisten.

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

Potongan konteks digabung menjadi satu string dan dibatasi hingga 2.500 karakter agar input tetap terkelola. Setiap jawaban diubah menjadi huruf kecil, dan contoh dengan label di luar yes, no, atau maybe dihapus.

Mencetak contoh pertama yang telah dikonversi membantu memastikan bahwa konteks, pertanyaan, dan jawaban telah diformat dengan benar sebelum pelatihan.

9. Bangun Contoh Pelatihan Difusi

DiffusionGemma memerlukan jawaban target ditempatkan dalam kanvas berukuran tetap. Fungsi ini melakukan tokenisasi prompt, mengonversi jawaban menjadi ID token, melakukan padding hingga panjang kanvas model, dan membuat mask yang menunjukkan token mana yang berkontribusi terhadap 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 akhir-sekuens ditambahkan setelah setiap jawaban, sementara posisi kanvas yang tersisa diisi dengan token padding. Loss mask memastikan pelatihan hanya berfokus pada token jawaban dan token akhir-sekuens, bukan posisi yang dipadding.

10. Buat Fungsi Inferensi dan Evaluasi

Selanjutnya, definisikan fungsi untuk menghasilkan jawaban, membersihkan keluaran model, dan menghitung akurasi evaluasi.

Hasilkan jawaban

Fungsi answer_question() memformat prompt, menghasilkan respons melalui beberapa langkah denoising, dan mendekode token yang dihasilkan menjadi teks.

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

Ekstrak prediksi

Meskipun prompt meminta jawaban satu kata, model sesekali dapat menghasilkan teks tambahan. Fungsi ini mengekstrak yes, no, atau maybe pertama yang valid.

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"

Evaluasi akurasi

Fungsi evaluasi membandingkan setiap prediksi yang telah dibersihkan dengan jawaban yang benar, mencetak hasil untuk setiap contoh, dan mengembalikan akurasi keseluruhan beserta prediksi individual.

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. Evaluasi Model Sebelum Fine-Tuning

Jalankan evaluasi sebelum pelatihan untuk menetapkan baseline.

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

Ini mengevaluasi 50 contoh menggunakan 16 langkah denoising per jawaban. Dalam eksperimen ini, model dasar menjawab benar 30 dari 50 pertanyaan.

Baseline ini nantinya akan dibandingkan dengan akurasi model setelah fine-tuning.

12. Siapkan Pelatihan

Alihkan model ke mode pelatihan, buat optimizer dan scheduler laju pembelajaran, dan definisikan bagaimana token jawaban bersih akan dikorupsi selama pelatihan difusi.

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

Hanya parameter dengan requires_grad=True yang diteruskan ke optimizer, yang berarti proses pelatihan memperbarui adapter LoRA, bukan seluruh model.

Selanjutnya, buat fungsi korupsi yang mengganti proporsi acak dari kanvas jawaban dengan token acak.

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)

Jumlah noise berubah untuk setiap contoh. Selama pelatihan, model belajar merekonstruksi jawaban asli dari token kanvas yang dikorupsi ini.

13. Latih Model

Loop berikut melatih adapter LoRA selama 60 langkah menggunakan akumulasi gradien.

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

Untuk setiap contoh pelatihan, model menerima prompt biomedis dan kanvas jawaban yang dikorupsi. Loss cross-entropy dihitung hanya untuk token jawaban asli yang dipilih oleh loss mask.

Akumulasi gradien menggabungkan empat contoh sebelum memperbarui model. Gradient clipping juga diterapkan untuk menjaga pelatihan tetap stabil.

Dalam eksperimen ini, pelatihan selesai dalam waktu sekitar dua menit:

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

Penurunan loss yang konsisten menunjukkan bahwa adapter belajar merekonstruksi jawaban yang diharapkan dari kanvas yang dikorupsi. Selama pelatihan, Anda juga dapat menjalankan nvidia-smi di terminal RunPod untuk memantau penggunaan dan pemanfaatan memori GPU.

14. Evaluasi Model yang Telah Di-fine-tune

Jalankan evaluasi yang sama lagi setelah pelatihan untuk mengukur apakah fine-tuning meningkatkan kinerja model.

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

Model yang telah di-fine-tune dievaluasi pada 50 contoh yang sama dan menggunakan 16 langkah denoising yang sama seperti pada evaluasi baseline.

Selanjutnya, bandingkan akurasi sebelum dan sesudah 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

Dalam eksperimen ini, akurasi model meningkat dari 0,60 menjadi 0,80.

Ini merepresentasikan peningkatan 20 poin persentase, dengan model menjawab benar 40 dari 50 pertanyaan setelah fine-tuning, dibandingkan 30 dari 50 sebelum pelatihan.

15. Simpan dan Unggah Adapter yang Telah Di-fine-tune

Simpan adapter LoRA yang telah dilatih dan file prosesor ke direktori keluaran yang telah ditentukan sebelumnya.

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

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

Anda akan melihat:

Saved LoRA adapter to: diffusiongemma_pubmedqa_lora

Ini hanya menyimpan adapter LoRA yang ringan, bukan salinan lengkap lain dari model dasar 26 miliar parameter.

Selanjutnya, unggah file adapter dan prosesor ke Hugging Face Hub:

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

Karena variabel lingkungan HF_TOKEN telah ditambahkan selama konfigurasi pod RunPod, Hugging Face seharusnya melakukan autentikasi secara otomatis. Anda hanya perlu menjalankan notebook_login() saat token belum dikonfigurasi:

from huggingface_hub import notebook_login

notebook_login()

Setelah unggahan selesai, repositori akan berisi adapter LoRA dan konfigurasi prosesor yang diperlukan untuk memuat model yang telah di-fine-tune nanti.

Sumber: kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

Pemikiran Akhir

Melakukan fine-tuning DiffusionGemma dengan Unsloth ternyata cukup mudah. Bagian paling memakan waktu adalah memasang dependensi yang benar dan memahami cara kerja proses pelatihan khusus difusi. Setelah lingkungan siap, memuat model, melatih adapter LoRA, mengevaluasinya, dan mengunggah hasil ke Hugging Face berjalan sangat mulus.

Saya merasa DiffusionGemma sangat menarik karena tidak menghasilkan teks satu per satu token seperti model bahasa tradisional. Sebaliknya, ia bekerja dengan kanvas tetap dan secara bertahap menyempurnakan token bising melalui langkah-langkah denoising. Memahami proses generasi yang berbeda ini dan melakukan fine-tuning pada tugas tanya jawab biomedis membuat eksperimen ini sangat berharga bagi saya.

Bahkan dengan setup kecil, akurasi pada sampel 50 contoh naik dari 0,60 menjadi 0,80—meskipun pada ukuran sampel seperti itu, margin of error cukup lebar, dan baseline "selalu ya" sudah mencetak sekitar 55% pada set evaluasi ini. 

Perlu dicatat juga bahwa subset pelatihan artifisial hampir tidak mengandung label "maybe", sehingga model memiliki sedikit kesempatan untuk mempelajari kelas tersebut meskipun muncul dalam data evaluasi. Perlakukan ini sebagai eksperimen cepat tentang cara kerja proses fine-tuning, bukan sebagai bukti bahwa model siap untuk penggunaan medis nyata.

Siap melampaui satu sesi fine-tuning tunggal? Jalur Developing Large Language Models kami membawa Anda dari dasar-dasar PyTorch dan transformers hingga membangun dan menerapkan LLM Anda sendiri.