和硕网站建设baidu百度网盘

说明

unsloth框架可以进行各种sft训练，包括lora和grpo训练。我参考官方方法，使用模型Qwen2.5-3B-Instruct和数据集gsm8k，写了一个grpo训练的例子。

代码

这个代码加载模型Qwen2.5-3B-Instruct和数据集gsm8k。训练完成后先保存lora模型然后保存合并后的模型。

import os
from unsloth import FastLanguageModel, PatchFastRL
from unsloth import is_bfloat16_supported
import torch
import re
from datasets import load_dataset, DatasetPatchFastRL("GRPO", FastLanguageModel)
from trl import GRPOConfig, GRPOTrainermax_seq_length = 1024 # Can increase for longer reasoning traces
lora_rank = 64 # Larger rank = smarter, but slowerprint(">>>>>>>>>>>>>>>FastLanguageModel.from_pretrained:")
model, tokenizer = FastLanguageModel.from_pretrained(model_name = "./Qwen2.5-3B-Instruct",max_seq_length = max_seq_length,load_in_4bit = True, # False for LoRA 16bitfast_inference = True, # Enable vLLM fast inferencemax_lora_rank = lora_rank,gpu_memory_utilization = 0.8, # Reduce if out of memory
)print(">>>>>>>>>>>>>>>FastLanguageModel.get_peft_model:")
model = FastLanguageModel.get_peft_model(model,r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128target_modules = ["q_proj", "k_proj", "v_proj", "o_proj","gate_proj", "up_proj", "down_proj",], # Remove QKVO if out of memorylora_alpha = lora_rank,use_gradient_checkpointing = "unsloth", # Enable long context finetuningrandom_state = 3407,
)# Load and prep dataset
SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
...
</reasoning>
<answer>
...
</answer>
"""XML_COT_FORMAT = """\
<reasoning>
{reasoning}
</reasoning>
<answer>
{answer}
</answer>
"""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 Nonereturn text.split("####")[1].strip()# uncomment middle messages for 1-shot prompting
def get_gsm8k_questions(split = "train") -> Dataset:print(f">>>>>>>>>>>>>>>_get_gsm8k_questions, split:{split}")data = load_dataset('./gsm8k', 'main')[split] # type: ignoredata = data.map(lambda x: { # type: ignore'prompt': [{'role': 'system', 'content': SYSTEM_PROMPT},{'role': 'user', 'content': x['question']}],'answer': extract_hash_answer(x['answer'])}) # type: ignorereturn data # type: ignoreprint(">>>>>>>>>>>>>>>get_gsm8k_questions:")
dataset = get_gsm8k_questions()# Reward functions
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)]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]def strict_format_reward_func(completions, **kwargs) -> list[float]:"""Reward function that checks if the completion has a specific format."""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]def soft_format_reward_func(completions, **kwargs) -> list[float]:"""Reward function that checks if the completion has a specific format."""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]def count_xml(text) -> float:count = 0.0if text.count("<reasoning>\n") == 1:count += 0.125if text.count("\n</reasoning>\n") == 1:count += 0.125if text.count("\n<answer>\n") == 1:count += 0.125count -= len(text.split("\n</answer>\n")[-1])*0.001if text.count("\n</answer>") == 1:count += 0.125count -= (len(text.split("\n</answer>")[-1]) - 1)*0.001return countdef xmlcount_reward_func(completions, **kwargs) -> list[float]:contents = [completion[0]["content"] for completion in completions]return [count_xml(c) for c in contents]print(">>>>>>>>>>>>>>>training_args:")
training_args = GRPOConfig(use_vllm = True, # use vLLM for fast inference!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 = "adamw_8bit",logging_steps = 10,bf16 = is_bfloat16_supported(),fp16 = not is_bfloat16_supported(),per_device_train_batch_size = 1,gradient_accumulation_steps = 1, # Increase to 4 for smoother trainingnum_generations = 8, # Decrease if out of memorymax_prompt_length = 300,max_completion_length = 300,# num_train_epochs = 1, # Set to 1 for a full training runmax_steps = 100,save_steps = 50,max_grad_norm = 0.1,report_to = "none", # Can use Weights & Biasesoutput_dir = "outputs_2",
)
print(f"training_args:{training_args}")print(">>>>>>>>>>>>>>>GRPOTrainer:")
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,
)
print(">>>>>>>>>>>>>>>trainer.train:")
trainer.train()print(">>>>>>>>>>>>>>>model.save_lora:")
model.save_lora("grpo_saved_lora_2")print(">>>>>>>>>>>>>>>model.save_lora:")
model.save_pretrained_merged("./model_merged", tokenizer, save_method = "merged_16bit",)print(">>>>>>>>>>>>>>>DONE<<<<<<<<<<<<<<")

参考消息

使用unsloth进行grpo训练报错及解决方法