實戰練習：使用 Unsloth 進行 GRPO

在此練習中，您將使用 Unsloth 對模型進行 GRPO（Group Relative Policy Optimization）微調，以提高模型的推理能力。我們在第 3 章中介紹過 GRPO。

Unsloth 是一個加速大型語言模型微調的庫，可以更快地訓練模型並減少計算資源。Unsloth 集成了 TRL，因此我們將基於前面幾節所學的內容進行構建，並根據 Unsloth 的具體情況進行調整。

此練習可以在免費的 Google Colab T4 GPU 上執行。為了獲得最佳體驗，請按照上面連結的筆記本進行操作並親自嘗試。

安裝依賴項

首先，讓我們安裝必要的庫。我們將需要 Unsloth 用於加速微調，以及 vLLM 用於快速推理。

pip install unsloth vllm
pip install --upgrade pillow

設定 Unsloth

Unsloth 提供了一個類（FastLanguageModel），它將 transformers 與 Unsloth 最佳化整合在一起。讓我們匯入它。

from unsloth import FastLanguageModel

現在，讓我們載入 Google 的 Gemma 3 1B Instruct 模型並對其進行微調配置。

from unsloth import FastLanguageModel
import torch

max_seq_length = 1024  # Can increase for longer reasoning traces
lora_rank = 32  # Larger rank = smarter, but slower

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="google/gemma-3-1b-it",
    max_seq_length=max_seq_length,
    load_in_4bit=True,  # False for LoRA 16bit
    fast_inference=True,  # Enable vLLM fast inference
    max_lora_rank=lora_rank,
    gpu_memory_utilization=0.6,  # Reduce if out of memory
)

model = FastLanguageModel.get_peft_model(
    model,
    r=lora_rank,  # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],  # Remove QKVO if out of memory
    lora_alpha=lora_rank,
    use_gradient_checkpointing="unsloth",  # Enable long context finetuning
    random_state=3407,
)

此程式碼以 4 位量化載入模型以節省記憶體，並應用 LoRA（低秩適應）以實現高效微調。target_modules 引數指定要微調的模型層，use_gradient_checkpointing 啟用更長上下文的訓練。

本章我們不會詳細介紹 LoRA，您可以在第 11 章中瞭解更多資訊。

資料準備

在此練習中，我們將使用 GSM8K 資料集，其中包含小學數學問題。我們將格式化資料，以鼓勵模型在提供答案之前展示其推理過程。

首先，我們將定義提示和答案的格式。

# Define the system prompt that instructs the model to use a specific format
SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
...
</reasoning>
<answer>
...
</answer>
"""

XML_COT_FORMAT = """\
<reasoning>
{reasoning}
</reasoning>
<answer>
{answer}
</answer>
"""

現在，讓我們準備資料集。

import re
from datasets import load_dataset, Dataset


# Helper functions to extract answers from different formats
def extract_xml_answer(text: str) -> str:
    answer = text.split("<answer>")[-1]
    answer = answer.split("</answer>")[0]
    return answer.strip()


def extract_hash_answer(text: str) -> str | None:
    if "####" not in text:
        return None
    return text.split("####")[1].strip()


# Function to prepare the GSM8K dataset
def get_gsm8k_questions(split="train") -> Dataset:
    data = load_dataset("openai/gsm8k", "main")[split]
    data = data.map(
        lambda x: {
            "prompt": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": x["question"]},
            ],
            "answer": extract_hash_answer(x["answer"]),
        }
    )
    return data


dataset = get_gsm8k_questions()

資料集透過從資料集中提取答案並將其格式化為字串來準備。

定義獎勵函式

正如我們在前一頁討論的，GRPO 可以使用獎勵函式來根據可驗證的標準（如長度和格式）指導模型的學習。

在此練習中，我們將定義幾個獎勵函式，鼓勵良好推理的不同方面。例如，我們將獎勵模型提供整數答案，並遵循嚴格的格式。

# Reward function that checks if the answer is correct
def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
    responses = [completion[0]["content"] for completion in completions]
    q = prompts[0][-1]["content"]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    print(
        "-" * 20,
        f"Question:\n{q}",
        f"\nAnswer:\n{answer[0]}",
        f"\nResponse:\n{responses[0]}",
        f"\nExtracted:\n{extracted_responses[0]}",
    )
    return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]


# Reward function that checks if the answer is an integer
def int_reward_func(completions, **kwargs) -> list[float]:
    responses = [completion[0]["content"] for completion in completions]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]


# Reward function that checks if the completion follows the strict format
def strict_format_reward_func(completions, **kwargs) -> list[float]:
    pattern = r"^<reasoning>\n.*?\n</reasoning>\n<answer>\n.*?\n</answer>\n$"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, r) for r in responses]
    return [0.5 if match else 0.0 for match in matches]


# Reward function that checks if the completion follows a more relaxed format
def soft_format_reward_func(completions, **kwargs) -> list[float]:
    pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, r) for r in responses]
    return [0.5 if match else 0.0 for match in matches]


# Reward function that counts XML tags and penalizes extra content
def count_xml(text) -> float:
    count = 0.0
    if text.count("<reasoning>\n") == 1:
        count += 0.125
    if text.count("\n</reasoning>\n") == 1:
        count += 0.125
    if text.count("\n<answer>\n") == 1:
        count += 0.125
        count -= len(text.split("\n</answer>\n")[-1]) * 0.001
    if text.count("\n</answer>") == 1:
        count += 0.125
        count -= (len(text.split("\n</answer>")[-1]) - 1) * 0.001
    return count


def xmlcount_reward_func(completions, **kwargs) -> list[float]:
    contents = [completion[0]["content"] for completion in completions]
    return [count_xml(c) for c in contents]

這些獎勵函式有不同的用途：

獎勵函式	目的
`correctness_reward_func`	當模型的答案與正確答案匹配時獎勵模型
`int_reward_func`	獎勵模型提供數字答案
`strict_format_reward_func` 和 `soft_format_reward_func`	獎勵模型遵循指定的格式
`xmlcount_reward_func`	獎勵正確的 XML 標籤使用並懲罰關閉標籤後多餘的內容

使用 GRPO 進行訓練

現在我們將使用我們的模型、分詞器和獎勵函式設定 GRPO 訓練器。這部分與之前的練習遵循相同的方法。

from trl import GRPOConfig, GRPOTrainer

max_prompt_length = 256

training_args = GRPOConfig(
    learning_rate=5e-6,
    adam_beta1=0.9,
    adam_beta2=0.99,
    weight_decay=0.1,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine",
    optim="paged_adamw_8bit",
    logging_steps=1,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,  # Increase to 4 for smoother training
    num_generations=6,  # Decrease if out of memory
    max_prompt_length=max_prompt_length,
    max_completion_length=max_seq_length - max_prompt_length,
    # num_train_epochs = 1, # Set to 1 for a full training run
    max_steps=250,
    save_steps=250,
    max_grad_norm=0.1,
    report_to="none",  # Can use Weights & Biases
    output_dir="outputs",
)

trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=[
        xmlcount_reward_func,
        soft_format_reward_func,
        strict_format_reward_func,
        int_reward_func,
        correctness_reward_func,
    ],
    args=training_args,
    train_dataset=dataset,
)

GRPOConfig 設定了各種訓練超引數：

use_vllm: 使用 vLLM 啟用快速推理
learning_rate: 控制模型學習的速度
num_generations: 為每個提示生成的完成數量
max_steps: 執行的總訓練步數

現在讓我們開始訓練

trainer.train()

訓練可能需要一些時間。您可能不會立即看到獎勵增加——可能需要 150-200 步才能開始看到改進。請耐心！

測試模型

訓練完成後，讓我們測試模型，看看它的表現如何。首先，我們將儲存 LoRA 權重。

model.save_lora("grpo_saved_lora")

現在，讓我們用一個新問題測試模型。

from vllm import SamplingParams

text = tokenizer.apply_chat_template(
    [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": "Calculate pi."},
    ],
    tokenize=False,
    add_generation_prompt=True,
)

sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=1024,
)
output = (
    model.fast_generate(
        text,
        sampling_params=sampling_params,
        lora_request=model.load_lora("grpo_saved_lora"),
    )[0]
    .outputs[0]
    .text
)

print(output)

您應該會看到模型現在遵循指定的格式，在提供答案之前顯示其推理過程。

儲存模型

Unsloth 提供了多種儲存微調模型的方法，但我們將重點介紹最常用的方法。

# Save to 16-bit precision
model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")

推送到 Hugging Face Hub

我們將使用 push_to_hub_merged 方法將模型推送到 Hugging Face Hub。此方法允許我們以多種量化格式推送模型。

# Push to Hugging Face Hub (requires a token)
model.push_to_hub_merged(
    "your-username/model-name", tokenizer, save_method="merged_16bit", token="your-token"
)

Unsloth 還支援儲存為 GGUF 格式以與 llama.cpp 一起使用。

model.push_to_hub_gguf(
    "your-username/model-name",
    tokenizer,
    quantization_method=["q4_k_m", "q8_0", "q5_k_m"],
    token="your-token",
)

GGUF 檔案可與 llama.cpp 或基於 UI 的系統（如 Jan 或 Open WebUI）一起使用。

結論

在此練習中，您已經學會了如何：

設定 Unsloth 以進行加速微調
準備用於 GRPO 訓練的資料
定義自定義獎勵函式以指導模型學習
使用 GRPO 訓練模型
測試微調模型
以各種格式儲存模型

GRPO 是一種將語言模型與特定行為對齊的強大技術，Unsloth 使其即使在有限的硬體上也能使用。透過結合多個獎勵函式，您可以引導模型遵循特定格式，同時提高其推理能力。

有關更多資訊和資源，請檢視：

< > 在 GitHub 上更新

LLM 課程