LLM 课程文档
实践练习:使用 Unsloth 的 GRPO
并获得增强的文档体验
开始使用
实践练习:使用 Unsloth 的 GRPO
在本练习中,你将使用 Unsloth 通过 GRPO(Group Relative Policy Optimization,组相对策略优化)微调模型,以提高模型的推理能力。我们在第 3 章中介绍了 GRPO。
Unsloth 是一个加速 LLM 微调的库,使模型训练更快,计算资源更少成为可能。Unsloth 插件式地接入 TRL,因此我们将基于之前章节所学的知识,并针对 Unsloth 的特性进行调整。
本练习可以在免费的 Google Colab T4 GPU 上运行。为了获得最佳体验,请跟随上面链接的 notebook 并亲自尝试。
安装依赖
首先,让我们安装必要的库。我们将需要 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(Low-Rank Adaptation,低秩适配)以实现高效微调。“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 上更新