初始化项目，由ModelHub XC社区提供模型
Model: krogoldAI/QueryRefiner-0.5B-v0.1-GRPO Source: Original Platform
This commit is contained in:
ModelHub XC
247
Codes/01 - GRPO Training.py
Normal file
247
Codes/01 - GRPO Training.py
Normal file
@@ -0,0 +1,247 @@
|
||||
import torch
|
||||
import wandb
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from trl import GRPOConfig, GRPOTrainer
|
||||
from datasets import load_dataset, Dataset
|
||||
from huggingface_hub import login
|
||||
from common.constants import SYSTEM_PROMPT
|
||||
from common.rewards import SmoothXMLRewardEvaluator
|
||||
|
||||
##########
|
||||
# CONFIG #
|
||||
##########
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
DATASET = "krogoldAI/rag-ambiguous-queries"
|
||||
OUTPUT_MODEL = "model/repo"
|
||||
HF_TOKEN = "<MY_HF_TOKEN>"
|
||||
|
||||
# Training hyperparameters
|
||||
NUM_EPOCHS = 1
|
||||
PER_DEVICE_BATCH = 8
|
||||
GRADIENT_ACCUMULATION = 1 # 2
|
||||
LEARNING_RATE = 5e-6 # Low for GRPO stability
|
||||
NUM_GENERATIONS = 2
|
||||
MAX_NEW_TOKENS = 512
|
||||
WARMUP_STEPS = 50
|
||||
|
||||
#########
|
||||
# LOGIN #
|
||||
#########
|
||||
|
||||
print("Logging into Hugging Face...")
|
||||
login(HF_TOKEN)
|
||||
|
||||
##################
|
||||
# LOAD TOKENIZER #
|
||||
##################
|
||||
|
||||
print(f"Loading tokenizer from {MODEL_NAME}...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
||||
tokenizer.padding_side = "left"
|
||||
|
||||
# Ensure pad token is set
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
################
|
||||
# LOAD DATASET #
|
||||
################
|
||||
|
||||
print(f"Loading dataset {DATASET}...")
|
||||
raw_dataset = load_dataset(DATASET, "default", split="train")
|
||||
|
||||
# Extract the list of questions
|
||||
questions_list = raw_dataset["questions"]
|
||||
print(f"Dataset size: {len(questions_list)} examples")
|
||||
|
||||
# Create a proper dataset from the list
|
||||
dataset = Dataset.from_dict({"question": questions_list})
|
||||
dataset = dataset.shuffle(seed=42)
|
||||
|
||||
# Take on 3000 examples
|
||||
dataset = dataset.select(range(3000))
|
||||
|
||||
###########################
|
||||
# PREPROCESSING FUNCTION #
|
||||
###########################
|
||||
|
||||
def preprocess(example):
|
||||
"""
|
||||
Apply chat template to create proper prompt format.
|
||||
"""
|
||||
messages = [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": example["question"]},
|
||||
]
|
||||
|
||||
# Apply chat template without tokenization to get the prompt string
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)
|
||||
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"query": example["question"] # Keep for reference
|
||||
}
|
||||
|
||||
print("Preprocessing dataset with chat template...")
|
||||
dataset = dataset.map(preprocess, remove_columns=dataset.column_names)
|
||||
|
||||
##############
|
||||
# LOAD MODEL #
|
||||
##############
|
||||
|
||||
print(f"\nLoading model {MODEL_NAME}...")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL_NAME,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="auto",
|
||||
attn_implementation="flash_attention_2",
|
||||
use_cache=False,
|
||||
)
|
||||
|
||||
# If Liger Kernel available, apply now
|
||||
try:
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_qwen2
|
||||
try:
|
||||
apply_liger_kernel_to_qwen2(model)
|
||||
except TypeError:
|
||||
apply_liger_kernel_to_qwen2()
|
||||
print("Liger Kernel applied successfully for Qwen2 optimization")
|
||||
except Exception as e:
|
||||
print(f"Liger Kernel not available or failed to apply: {e}")
|
||||
|
||||
print(f"Model loaded. Parameters: {model.num_parameters() / 1e9:.3f}B")
|
||||
|
||||
model.generation_config.max_new_tokens = MAX_NEW_TOKENS
|
||||
model.generation_config.temperature = 0.7
|
||||
model.generation_config.do_sample = True
|
||||
|
||||
###################
|
||||
# REWARD FUNCTION #
|
||||
###################
|
||||
|
||||
# Initialize reward evaluator
|
||||
evaluator = SmoothXMLRewardEvaluator()
|
||||
|
||||
def reward_function(completions, prompts, **kwargs):
|
||||
rewards = []
|
||||
|
||||
for i, (completion, prompt) in enumerate(zip(completions, prompts)):
|
||||
# Extract generated portion
|
||||
if completion.startswith(prompt):
|
||||
response = completion[len(prompt):].strip()
|
||||
else:
|
||||
response = completion
|
||||
try:
|
||||
analysis = evaluator.get_detailed_analysis(response)
|
||||
score = analysis["final_score"]
|
||||
rewards.append(score)
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR in reward evaluation: {e}")
|
||||
print(f"Response that failed: {response[:200]}")
|
||||
rewards.append(0.0)
|
||||
|
||||
return rewards
|
||||
|
||||
|
||||
##############
|
||||
# MONITORING #
|
||||
##############
|
||||
|
||||
# Monitoring the LLM from WandB
|
||||
wandb.login(key = "<MY_WANDB_KEY>")
|
||||
run = wandb.init(project='QueryRefiner-GRPO-Phase1', job_type="training", anonymous="allow")
|
||||
|
||||
######################
|
||||
# TRAINING ARGUMENTS #
|
||||
######################
|
||||
|
||||
print("Configuring training arguments...")
|
||||
training_args = GRPOConfig(
|
||||
output_dir="./qwen_grpo_checkpoints",
|
||||
run_name="qwen-grpo-query-analysis",
|
||||
num_train_epochs=NUM_EPOCHS,
|
||||
per_device_train_batch_size=PER_DEVICE_BATCH,
|
||||
gradient_accumulation_steps=GRADIENT_ACCUMULATION,
|
||||
learning_rate=LEARNING_RATE,
|
||||
num_generations=NUM_GENERATIONS,
|
||||
optim="adamw_torch_fused",
|
||||
warmup_steps=WARMUP_STEPS,
|
||||
lr_scheduler_type="cosine",
|
||||
weight_decay=0.01,
|
||||
max_grad_norm=1.0,
|
||||
bf16=True,
|
||||
gradient_checkpointing=True,
|
||||
use_liger_kernel=True,
|
||||
logging_steps=10,
|
||||
save_steps=200,
|
||||
save_total_limit=3,
|
||||
push_to_hub=True,
|
||||
hub_model_id=OUTPUT_MODEL,
|
||||
hub_private_repo=True,
|
||||
hub_strategy="checkpoint",
|
||||
dataloader_num_workers=4,
|
||||
dataloader_pin_memory=True,
|
||||
remove_unused_columns=False,
|
||||
report_to="wandb",
|
||||
seed=42,
|
||||
)
|
||||
|
||||
###########
|
||||
# TRAINER #
|
||||
###########
|
||||
|
||||
print("Initializing GRPOTrainer...")
|
||||
trainer = GRPOTrainer(
|
||||
model=model,
|
||||
reward_funcs=reward_function,
|
||||
args=training_args,
|
||||
train_dataset=dataset, # Uses by default the "prompt" column, all other columns are ignored
|
||||
)
|
||||
|
||||
#########
|
||||
# TRAIN #
|
||||
#########
|
||||
|
||||
print("=" * 50)
|
||||
print("Starting GRPO training...")
|
||||
print(f"Model: {MODEL_NAME}")
|
||||
print(f"Dataset size: {len(dataset)}")
|
||||
print(f"Epochs: {NUM_EPOCHS}")
|
||||
print(f"Effective batch size: {PER_DEVICE_BATCH * GRADIENT_ACCUMULATION}")
|
||||
print(f"Generations per prompt: {NUM_GENERATIONS}")
|
||||
print("=" * 50)
|
||||
|
||||
trainer.train()
|
||||
|
||||
####################
|
||||
# SAVE FINAL MODEL #
|
||||
####################
|
||||
|
||||
print("\nSaving final model...")
|
||||
final_model = trainer.model
|
||||
final_model.config.use_cache = True # Re-enable cache for inference
|
||||
final_model.save_pretrained("./final_grpo_model")
|
||||
tokenizer.save_pretrained("./final_grpo_model")
|
||||
|
||||
###############
|
||||
# PUSH TO HUB #
|
||||
###############
|
||||
|
||||
print(f"\nPushing final model to {OUTPUT_MODEL}...")
|
||||
try:
|
||||
final_model.push_to_hub(OUTPUT_MODEL, private=True)
|
||||
tokenizer.push_to_hub(OUTPUT_MODEL, private=True)
|
||||
print("Successfully pushed model and tokenizer to Hub!")
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to push to hub: {e}")
|
||||
print("Model saved locally at ./final_grpo_model")
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("GRPO Training Complete!")
|
||||
print("=" * 50)
|
||||
Reference in New Issue
Block a user