Unsloth を使って PubMedQA で DiffusionGemma をファインチューニングする方法

1. RunPod の Jupyter Notebook を開く

PyTorch/Jupyter テンプレートを選択し、NVIDIA H100 GPU を搭載した新しい RunPod ポッドを作成してください。ポッド停止時にモデルファイルや学習結果が失われないよう、永続ストレージを 100 GB 以上に設定します。

Hugging Face のアクセストークンを環境変数として追加します:

HF_TOKEN=your_hugging_face_token

これによりモデルやデータセットのダウンロードが高速化され、ノートブックから手動ログインすることなく、保存した LoRA アダプターを Hugging Face にアップロードできるようになります。

構成したポッドの費用は概ね 1 時間あたり $3 程度ですが、最終的な価格は GPU の空き状況や選択したポッドタイプにより変動する可能性があります。

ポッドが起動したら、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 がどのように回答を生成するかをイメージしておくとよいでしょう。トークンを順次生成するのではなく、固定長のキャンバスを用意し、複数回のノイズ除去ステップで多くの位置を一度に更新しながら、整合した応答へと収束させます。

以下のコードでは、1 ブロックあたりのサイズである 256 のキャンバス長が報告されます。今回のような yes/no/maybe の短い回答なら 1 枚のキャンバスで十分です。より長い出力は、ブロック単位でキャンバスを連結して生成されます。下図は高レベルでの洗練プロセスの概念図です。

5. DiffusionGemma を読み込む

bfloat16 精度で命令追従済みの DiffusionGemma モデルを読み込みます。H100 GPU は十分なメモリを備えているため、4 ビット量子化は使用しません。

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)

語彙サイズは拡散学習時にランダムノイズを加える際に使用します。キャンバス長は、単一生成ブロック内でモデルが洗練できる最大トークン数を決定します。

出力は次のようになるはずです。

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

6. LoRA アダプターを追加する

260 億パラメータの全体を更新する代わりに、少数の追加パラメータのみを学習するよう LoRA アダプターを追加します。

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

これにより、ファインチューニングに必要なメモリと計算量が大幅に削減されます。H100 は十分な GPU メモリを持つため、本実験では勾配チェックポイントは無効化しています。

7. 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

各例には、生物医学の質問、1 つ以上の補助的な要約テキスト、そして最終回答として yes、no、maybe のいずれかが含まれます。

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

コンテキストの複数段落は 1 つの文字列に結合し、入力を扱いやすくするため 2,500 文字に制限します。回答は小文字化し、yes、no、maybe 以外のラベルを持つ例は除外します。

学習前に、変換後の最初の例を出力して、コンテキスト・質問・回答が正しく整形されていることを確認します。

9. 拡散学習用の例を構築する

DiffusionGemma では、ターゲット回答を固定長のキャンバス内に配置する必要があります。以下の関数は、プロンプトをトークナイズし、回答をトークン ID に変換し、モデルのキャンバス長までパディングし、損失に寄与すべきトークンを示すマスクを作成します。

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)

各回答の後ろに終端トークンを追加し、残りのキャンバス位置はパディングトークンで埋めます。損失マスクにより、学習はパディングではなく、回答と終端トークンにのみ焦点を当てます。

10. 推論と評価の関数を作成する

次に、回答の生成、モデル出力のクリーニング、評価精度の算出に用いる関数を定義します。

回答を生成する

answer_question() 関数は、プロンプトを整形し、複数のノイズ除去ステップで応答を生成し、生成トークンをテキストにデコードします。

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

予測を抽出する

プロンプトでは 1 語のみの回答を求めていますが、モデルが追加テキストを生成する場合があります。この関数は、最初に現れる有効な yes、no、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"

精度を評価する

評価関数は、クリーニングした予測と正解を比較し、各例の結果を表示するとともに、全体の精度と個々の予測結果を返します。

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

ここでは 1 回の回答につき 16 のノイズ除去ステップで 50 例を評価します。本実験では、ベースモデルは 50 問中 30 問に正答しました。

このベースラインは、後でファインチューニング後の精度と比較します。

12. 学習をセットアップする

モデルを学習モードに切り替え、オプティマイザーと学習率スケジューラを作成し、拡散学習中にクリーンな回答トークンをどのように汚染するかを定義します。

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 アダプターです。

次に、回答キャンバスのランダムな割合をランダムトークンに置き換える汚染関数を作成します。

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)

ノイズ量は例ごとに変化します。学習中、モデルはこのように汚染されたキャンバストークンから元の回答を再構成することを学びます。

13. モデルを学習する

以下のループでは、勾配の蓄積を用いて 60 ステップにわたり LoRA アダプターを学習します。

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

各学習例で、モデルは生物医学プロンプトと汚染された回答キャンバスを受け取ります。クロスエントロピー損失は、損失マスクで選択された実際の回答トークンに対してのみ計算されます。

勾配の蓄積により、4 例分をまとめてからモデルを更新します。学習の安定化のため、勾配クリッピングも適用します。

本実験では、学習は約 2 分で完了しました。

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 メモリ使用量や稼働状況を監視することもできます。

14. 学習済みモデルを評価する

ファインチューニング後に同一の評価を再実行し、性能が向上したかを測定します。

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

ベースラインと同様に、同じ 50 例・同じ 16 のノイズ除去ステップで評価します。

続いて、ファインチューニング前後の精度を比較します。

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

これにより、260 億パラメータのベースモデルをもう 1 つ保存するのではなく、軽量な LoRA アダプターのみを保存します。

続いて、アダプターとプロセッサのファイルを Hugging Face Hub にアップロードします。

REPO_ID = "kingabzpro/diffusiongemma_pubmedqa"

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

HF_TOKEN 環境変数を RunPod のポッド構成時に追加している場合、Hugging Face への認証は自動で行われます。トークンが未設定のときのみ、notebook_login() を実行してください。

from huggingface_hub import notebook_login

notebook_login()

アップロードが完了すると、リポジトリには、後で学習済みモデルを読み込むために必要な LoRA アダプターとプロセッサ設定が含まれます。

出典: kingabzpro/diffusiongemma_pubmedqa · Hugging Face 

まとめ

Unsloth を使った DiffusionGemma のファインチューニングは意外なほど簡単でした。最も時間がかかったのは、適切な依存関係のインストールと、拡散特有の学習プロセスの理解です。環境が整ってしまえば、モデルの読み込み、LoRA アダプターの学習、評価、そして Hugging Face へのアップロードまで、とてもスムーズに進みました。

DiffusionGemma は、従来の言語モデルのように 1 トークンずつ生成するのではなく、固定キャンバスをノイズ除去ステップで徐々に洗練していく点が特に興味深いと感じました。この異なる生成プロセスを理解し、生物医学 QA タスクでファインチューニングすることは、私にとって非常に有意義な実験でした。

今回の小規模セットアップでも、50 例のサンプルにおける精度は 0.60 から 0.80 に上昇しました。ただし、このサンプルサイズでは誤差範囲が大きく、「常に yes」と答えるだけでもこの評価セットでは約 55% 得点できる点には留意が必要です。

また、人工の学習サブセットには「maybe」ラベルがほとんど含まれていないため、評価データに登場するにもかかわらず、そのクラスを学習する機会が限られている点にも注意してください。本稿はファインチューニングの手順を理解するための簡易実験であり、実医療用途に直ちに適用できるという証拠ではありません。

単発のファインチューニングを超えて取り組む準備はできていますか？Developing Large Language Models トラックでは、PyTorch と transformers の基礎から始めて、独自の LLM を構築・デプロイするところまで学べます。