In the previous issue, we introduced the use of LoRA to train instruction models. In this issue, we will continue to train our own R1 model. LoRA has a wide range of applications in vertical fields such as law, healthcare, image processing, education, and programming. We will take education as an example to train a logical reasoning model.
Practical Implementation of LoRA Fine-Tuning
Jun Jun AI, WeChat Official Account: Jun Jun AI Practical Implementation of LoRA Fine-Tuning (1) – Training Instruction Models
Training ObjectiveTrain the Qwen/Qwen3-0.6B model on a dataset of elementary school math application problems, enabling it to possess capabilities in mathematical reasoning and multi-step problem solving.
Train a large model for structured reasoning, outputting the reasoning process and final answer through specific formats (reasoning and answer tags).
• Utilized the GRPO training method
• Employed LoRA for efficient fine-tuning
• Supports various quantization and saving formats
• Includes detailed reward function design
Dataset
The GSM8K (Grade School Math 8K) dataset is a high-quality dataset of elementary school math application problems, primarily used to evaluate and train AI models in mathematical reasoning and multi-step problem solving. Download link:https://hf-mirror.com/datasets/openai/gsm8k
GRPO (Group Relative Policy Optimization)
Group Relative Policy Optimization is a reinforcement learning algorithm used for training LLMs and is one of the core technologies of the DeepSeek-R1 model.
• The core of GRPO is to optimize the policy model through relative rewards of samples within a group, rather than relying on traditional value function models (like the critic model in PPO). It simplifies the training process by sampling a group of outputs and using their reward values to calculate relative advantages.
Working Principle:
• Sampling and Reward Calculation: For each input question, GRPO samples a group of outputs from the current policy and calculates the reward value for each output.
• Relative Advantage Estimation: By comparing the reward values of each output with the average reward value within the group, the relative advantage of each output is calculated.
• Policy Update: Based on the relative advantages, GRPO updates the policy model, prioritizing outputs with higher relative advantages. It also controls the magnitude of policy updates through KL divergence constraints to ensure the stability of the policy distribution.
Base Model Download
For practical implementation, we can choose a smaller model with lower resource requirements. We will take the Qwen3-0.6B model as an example and use modelscope to download the model.
#cache_dir is the model storage path
from modelscope import snapshot_download
snapshot_download('Qwen/Qwen3-0.6B', cache_dir="/content/models")R1 Model Training
Load the pre-trained model and LoRA. FastLanguageModel is an optimized LLM loader and processor provided by the unsloth library, with the following main functions:
• Provides faster model loading and inference speed
• Supports 4-bit quantization for model loading
• Integrates vLLM for fast inference
• Supports LoRA (Low-Rank Adaptation) fine-tuning
PatchFastRL is a patch function for reinforcement learning (RL):
• Adds reinforcement learning-related functionalities to language models
• Specifically supports the GRPO (Guided Reward Policy Optimization) algorithm
• Optimizes the model training process
"""Using the Unsloth framework for GRPO (Generative Reward-Paired Optimization) fine-tuning"""
from unsloth import FastLanguageModel
import torch
# Set model parameters
max_seq_length = 1024 # Can be increased for longer reasoning trajectories
lora_rank = 8 # A larger rank will make the model smarter but slower to train
# Load pre-trained model
model, tokenizer = FastLanguageModel.from_pretrained(
local_model_name = "/content/models/Qwen/Qwen3-0.6B", # Model path
max_seq_length = max_seq_length,
load_in_4bit = True, # Load using 4-bit quantization
fast_inference = True, # Enable vLLM fast inference, requires vllm to be installed. pip install vllm
max_lora_rank = lora_rank,
gpu_memory_utilization = 0.6, # Can be lowered if GPU memory is insufficient
)
# Configure LoRA parameters
model = FastLanguageModel.get_peft_model(
model,
r = lora_rank, # LoRA rank, recommended values are 8, 16, 32, 64, 128
target_modules = [
"q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj",
], # Can remove QKVO if GPU memory is insufficient
lora_alpha = lora_rank,
use_gradient_checkpointing = "unsloth", # Enable long context fine-tuning
random_state = 3407,
)Define system prompts to allow the large model to respond in the form of <reasoning> and <answer> tags. Reasoning is the reasoning process, answer is the final answer. Load the dataset and preprocess it.
# Data preparation section
import refrom datasets import load_dataset, Dataset
# Define system prompt
SYSTEM_PROMPT = """Respond in the following format:<reasoning>...</reasoning><answer>...</answer>"""
# Define XML format template
XML_COT_FORMAT = """
<reasoning>{reasoning}</reasoning><answer>{answer}</answer>
"""
def extract_xml_answer(text: str) -> str:
"""Extract the answer part from XML formatted text"""
answer = text.split("<answer>")[-1]
answer = answer.split("</answer>")[0]
return answer.strip()
def extract_hash_answer(text: str) -> str | None:
"""Extract the answer from text marked with ####"""
if "####" not in text:
return None
return text.split("####")[1].strip()
def get_gsm8k_questions(split = "train") -> Dataset:
"""Load the GSM8K dataset and preprocess it"""
#data = load_dataset('openai/gsm8k', 'main')[split]
data = load_dataset('/datasets/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
# Load dataset
dataset = get_gsm8k_questions()Define various reward functions, such as correctness, XML structure compliance, and completeness of tag structure as quantitative metrics. The group relative optimization algorithm utilizes these quantitative scores to obtain the best answer within the group.
# Define various reward functions
def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
"""Correctness reward function: checks if the answer is correct"""
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]:
"""Integer reward function: checks if the answer is an integer"""
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]:
"""Strict format reward function: checks if it fully complies with XML 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]:
"""Loose format reward function: checks if it generally complies with XML 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:
"""Calculate the completeness score of XML tags"""
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>[-1]")[-1])*0.001
if text.count("\n</answer>") == 1:
count += 0.125
count -= (len(text.split("\n</answer>[-1]")[-1]) - 1)*0.001
return count
def xmlcount_reward_func(completions, **kwargs) -> list[float]:
"""XML tag counting reward function"""
contents = [completion[0]["content"] for completion in completions]
return [count_xml(c) for c in contents]Configure GRPO parameters and initialize the trainer to start training.
# Training configuration
max_prompt_length = 256
from trl import GRPOConfig, GRPOTrainer
training_args = GRPOConfig(
learning_rate = 5e-6, # Learning rate
adam_beta1 = 0.9, # Adam optimizer parameter
adam_beta2 = 0.99,
weight_decay = 0.1, # Weight decay
warmup_ratio = 0.1, # Warmup ratio
lr_scheduler_type = "cosine", # Learning rate scheduler type
optim = "paged_adamw_8bit", # Optimizer type
logging_steps = 1, # Logging step count
per_device_train_batch_size = 1, # Training batch size per device
gradient_accumulation_steps = 1, # Gradient accumulation steps
num_generations = 6, # Number of generations
max_prompt_length = max_prompt_length, # Maximum prompt length
max_completion_length = max_seq_length - max_prompt_length, # Maximum completion length
max_steps = 250, # Maximum training steps
save_steps = 250, # Save steps
max_grad_norm = 0.1, # Maximum gradient norm
report_to = "none", # Reporting target
output_dir = "outputs", # Output directory
)
# Initialize trainer and configure reward functions
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,
)
# Start training
trainer.train()After training is complete, save the model and perform a test. Introduce sampling parameters from vllm to control the sampling temperature and top-p of the model output. These two parameters can be referenced:
The Principles and Functions of Temperature and Top P
Jun Jun AI, WeChat Official Account: Jun Jun AI The Principles and Functions of Temperature and Top P
# Save the trained LoRA model
model.save_lora("grpo_saved_lora")
# Test the model
text = tokenizer.apply_chat_template([
{"role" : "system", "content" : SYSTEM_PROMPT},
{"role" : "user", "content" : "The probability of event x occurring given that event a occurs is 60%, and the probability of event x occurring given that event b occurs is 80%. Events a and b are independent. What is the probability of event x occurring given that both events a and b occur?"},
], tokenize = False, add_generation_prompt = True)
from vllm import SamplingParams
sampling_params = SamplingParams(
temperature = 0.8, # Sampling temperature
top_p = 0.95, # Top-p sampling parameter
max_tokens = 1024, # Maximum number of generated tokens
)
# Use the saved LoRA for inference
output = model.fast_generate(
text,
sampling_params = sampling_params,
lora_request = model.load_lora("grpo_saved_lora"),
)[0].outputs[0].text
print(output)Finally, save the LoRA model in the appropriate format according to your needs and deployment requirements. push_to_hub_merged indicates pushing to Hugging Face Hub (if conditions allow).
# Model saving options
# Save as 16-bit float
if True: model.save_pretrained_merged("model", tokenizer, save_method = "merged_16bit") #model.push_to_hub_merged("hf/model", tokenizer, save_method = "merged_16bit", token = "")
# Save as 4-bit integer
if False: model.save_pretrained_merged("model", tokenizer, save_method = "merged_4bit") model.push_to_hub_merged("hf/model", tokenizer, save_method = "merged_4bit", token = "")
# Only save LoRA adapter
if False: model.save_pretrained_merged("model", tokenizer, save_method = "lora") model.push_to_hub_merged("hf/model", tokenizer, save_method = "lora", token = "")
# GGUF/llama.cpp conversion options
# Save as 8-bit Q8_0
if False: model.save_pretrained_gguf("model", tokenizer) model.push_to_hub_gguf("hf/model", tokenizer, token = "")
# Save as 16-bit GGUF
if False: model.save_pretrained_gguf("model", tokenizer, quantization_method = "f16") model.push_to_hub_gguf("hf/model", tokenizer, quantization_method = "f16", token = "")
# Save as q4_k_m GGUF
if False: model.save_pretrained_gguf("model", tokenizer, quantization_method = "q4_k_m") model.push_to_hub_gguf("hf/model", tokenizer, quantization_method = "q4_k_m", token = "")
# Save multiple GGUF options
if False: model.push_to_hub_gguf(
"hf/model",
tokenizer,
quantization_method = ["q4_k_m", "q8_0", "q5_k_m"],
token = "",
)That’s all for this issue on LoRA fine-tuning of large models. If you have any questions, feel free to discuss in the comments section. Follow us for more updates! #LoRA #Model Fine-Tuning #vllm #deepseek #R1 #GRPO #Group Relative Policy Optimization