[2025.06.22] add reward model、PPO、GRPO train code
This commit is contained in:
hanluzhi0429@gmail.com
+25
-12
@@ -1,15 +1,17 @@
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
from modeling_dolly import DollyForCausalLM
|
||||
from datasets import load_dataset
|
||||
from trl import SFTConfig, SFTTrainer, setup_chat_format
|
||||
import torch
|
||||
import os
|
||||
import deepspeed
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
deepspeed.init_distributed()
|
||||
|
||||
model_path = "./model"
|
||||
tokenizer_path = "./model"
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(model_path).to(device)
|
||||
model = DollyForCausalLM.from_pretrained(model_path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
||||
|
||||
try:
|
||||
@@ -42,7 +44,8 @@ finetune_name = "Dolly-SFT"
|
||||
sft_config = SFTConfig(
|
||||
output_dir="./sft_output",
|
||||
max_steps=1000000,
|
||||
per_device_train_batch_size=4,
|
||||
per_device_train_batch_size=16,
|
||||
per_device_eval_batch_size=2,
|
||||
learning_rate=5e-5,
|
||||
logging_steps=10,
|
||||
save_steps=1000,
|
||||
@@ -50,8 +53,11 @@ sft_config = SFTConfig(
|
||||
evaluation_strategy="steps",
|
||||
eval_steps=500,
|
||||
load_best_model_at_end=True,
|
||||
resume_from_checkpoint=True,
|
||||
resume_from_checkpoint=True,
|
||||
hub_model_id=finetune_name,
|
||||
deepspeed="./config/ds_config.json",
|
||||
fp16=True,
|
||||
gradient_accumulation_steps=8,
|
||||
)
|
||||
|
||||
trainer = SFTTrainer(
|
||||
@@ -73,15 +79,22 @@ if os.path.exists(sft_config.output_dir):
|
||||
checkpoint_dir = latest_checkpoint
|
||||
print(f"Resuming training from checkpoint: {latest_checkpoint}")
|
||||
|
||||
|
||||
trainer.train(resume_from_checkpoint=checkpoint_dir)
|
||||
|
||||
trainer.save_model(f"./{finetune_name}")
|
||||
if trainer.is_world_process_zero():
|
||||
trainer.save_model(f"./{finetune_name}")
|
||||
|
||||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
model = AutoModelForCausalLM.from_pretrained(f"./{finetune_name}").to(device)
|
||||
tokenizer = AutoTokenizer.from_pretrained(f"./{finetune_name}")
|
||||
|
||||
prompt = "Write a haiku about programming"
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)
|
||||
prompt = "Write a haiku about programming"
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)
|
||||
|
||||
inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
|
||||
outputs = model.generate(**inputs, max_new_tokens=100)
|
||||
print("\nAfter training:")
|
||||
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
|
||||
inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
|
||||
outputs = model.generate(**inputs, max_new_tokens=100)
|
||||
print("\nAfter training:")
|
||||
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
|
||||
#deepspeed --num_gpus 8 step4_cold_start.py
|
||||
@@ -0,0 +1,152 @@
|
||||
import os
|
||||
import torch
|
||||
from accelerate import PartialState
|
||||
from datasets import load_from_disk
|
||||
from transformers.trainer_utils import get_last_checkpoint
|
||||
from modeling_dolly import DollyForSequenceClassification
|
||||
from modeling_dolly import DollyForCausalLM
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoModelForSequenceClassification,
|
||||
AutoTokenizer,
|
||||
HfArgumentParser,
|
||||
)
|
||||
|
||||
from trl import (
|
||||
ModelConfig,
|
||||
PPOConfig,
|
||||
PPOTrainer,
|
||||
ScriptArguments,
|
||||
get_kbit_device_map,
|
||||
get_peft_config,
|
||||
get_quantization_config,
|
||||
)
|
||||
from trl.trainer.utils import SIMPLE_CHAT_TEMPLATE
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = HfArgumentParser((ScriptArguments, PPOConfig, ModelConfig))
|
||||
script_args, training_args, model_args = parser.parse_args_into_dataclasses()
|
||||
if script_args.deepspeed_config:
|
||||
training_args.deepspeed = script_args.deepspeed_config
|
||||
checkpoint_dir = None
|
||||
if os.path.exists(training_args.output_dir):
|
||||
checkpoint_dir = get_last_checkpoint(training_args.output_dir)
|
||||
if checkpoint_dir is not None:
|
||||
training_args.resume_from_checkpoint = checkpoint_dir
|
||||
print(f"恢复训练自检查点: {checkpoint_dir}")
|
||||
|
||||
torch_dtype = (
|
||||
model_args.torch_dtype if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype)
|
||||
)
|
||||
quantization_config = get_quantization_config(model_args)
|
||||
model_kwargs = dict(
|
||||
revision=model_args.model_revision,
|
||||
attn_implementation=model_args.attn_implementation,
|
||||
torch_dtype=torch_dtype,
|
||||
device_map=get_kbit_device_map() if quantization_config is not None else None,
|
||||
quantization_config=quantization_config,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
os.path.join("model", model_args.model_name_or_path),
|
||||
padding_side="left",
|
||||
trust_remote_code=model_args.trust_remote_code
|
||||
)
|
||||
|
||||
|
||||
if tokenizer.chat_template is None:
|
||||
tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE
|
||||
|
||||
reward_model_path = os.path.join("model", "reward_model")
|
||||
value_model = DollyForSequenceClassification.from_pretrained(
|
||||
reward_model_path,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
num_labels=1,
|
||||
**model_kwargs
|
||||
)
|
||||
|
||||
reward_model = DollyForSequenceClassification.from_pretrained(
|
||||
reward_model_path,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
num_labels=1,
|
||||
**model_kwargs
|
||||
)
|
||||
|
||||
policy_path = os.path.join("model", "policy_model")
|
||||
policy = DollyForCausalLM.from_pretrained(
|
||||
policy_path,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
**model_kwargs
|
||||
)
|
||||
|
||||
peft_config = get_peft_config(model_args)
|
||||
ref_policy = None
|
||||
if peft_config is None:
|
||||
ref_policy = DollyForCausalLM.from_pretrained(
|
||||
policy_path,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
**model_kwargs
|
||||
)
|
||||
|
||||
dataset_path = os.path.join("datas", script_args.dataset_name)
|
||||
if not os.path.exists(dataset_path):
|
||||
raise ValueError(f"数据集路径不存在: {dataset_path}")
|
||||
|
||||
dataset = load_from_disk(dataset_path)
|
||||
|
||||
eval_samples = 100
|
||||
train_dataset = dataset.select(range(len(dataset) - eval_samples))
|
||||
eval_dataset = dataset.select(range(len(dataset) - eval_samples, len(dataset)))
|
||||
dataset_text_field = "prompt"
|
||||
|
||||
def prepare_dataset(dataset, tokenizer):
|
||||
def tokenize(element):
|
||||
outputs = tokenizer(
|
||||
element[dataset_text_field],
|
||||
padding=False,
|
||||
)
|
||||
return {"input_ids": outputs["input_ids"]}
|
||||
|
||||
return dataset.map(
|
||||
tokenize,
|
||||
batched=True,
|
||||
remove_columns=dataset.column_names,
|
||||
num_proc=training_args.dataset_num_proc,
|
||||
)
|
||||
|
||||
with PartialState().local_main_process_first():
|
||||
train_dataset = prepare_dataset(train_dataset, tokenizer)
|
||||
eval_dataset = prepare_dataset(eval_dataset, tokenizer)
|
||||
|
||||
trainer = PPOTrainer(
|
||||
args=training_args,
|
||||
tokenizer=tokenizer,
|
||||
model=policy,
|
||||
ref_model=ref_policy,
|
||||
reward_model=reward_model,
|
||||
value_model=value_model,
|
||||
train_dataset=train_dataset,
|
||||
eval_dataset=eval_dataset,
|
||||
peft_config=peft_config,
|
||||
)
|
||||
|
||||
trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
|
||||
|
||||
trainer.save_model(training_args.output_dir)
|
||||
|
||||
trainer.generate_completions()
|
||||
"""
|
||||
torchrun --nproc_per_node=8 \
|
||||
train_ppo.py \
|
||||
--model_name_or_path policy_model \
|
||||
--dataset_name my_ppo_dataset \
|
||||
--output_dir ./ppo_output \
|
||||
--per_device_train_batch_size 8 \
|
||||
--gradient_accumulation_steps 2 \
|
||||
--learning_rate 1e-5 \
|
||||
--num_train_epochs 3 \
|
||||
--deepspeed_config config/deepspeed_z3.json \
|
||||
--dataset_text_field prompt \
|
||||
--reward_model_path model/reward_model \
|
||||
--sft_model_path model/policy_model
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
import os
|
||||
import torch
|
||||
from datasets import load_from_disk
|
||||
from modeling_dolly import DollyForSequenceClassification
|
||||
from modeling_dolly import DollyForCausalLM
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoTokenizer,
|
||||
TrainingArguments,
|
||||
TrainerCallback
|
||||
)
|
||||
from trl import GRPOConfig, GRPOTrainer
|
||||
from transformers.trainer_utils import get_last_checkpoint
|
||||
|
||||
class ResumeTrainingCallback(TrainerCallback):
|
||||
def on_train_begin(self, args, state, control, **kwargs):
|
||||
if state.resume_from_checkpoint is None:
|
||||
checkpoint = get_last_checkpoint(args.output_dir)
|
||||
if checkpoint is not None:
|
||||
control.should_resume = True
|
||||
|
||||
class RewardModel:
|
||||
def __init__(self, model_path, device="cuda"):
|
||||
self.device = device
|
||||
self.model = DollyForSequenceClassification.from_pretrained(model_path).to(device)
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
self.model.eval()
|
||||
|
||||
def __call__(self, completions, **kwargs):
|
||||
inputs = self.tokenizer(
|
||||
completions,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=2048,
|
||||
return_tensors="pt"
|
||||
).to(self.device)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**inputs)
|
||||
return outputs.logits.squeeze(-1).cpu().tolist()
|
||||
|
||||
def main():
|
||||
training_args = GRPOConfig(
|
||||
output_dir="./grpo_output",
|
||||
per_device_train_batch_size=32,
|
||||
per_device_eval_batch_size=8,
|
||||
num_train_epochs=3,
|
||||
gradient_accumulation_steps=2,
|
||||
learning_rate=1e-5,
|
||||
logging_steps=10,
|
||||
save_steps=500,
|
||||
eval_steps=200,
|
||||
warmup_ratio=0.1,
|
||||
weight_decay=0.01,
|
||||
fp16=True,
|
||||
deepspeed="configs/deepspeed_z3.json",
|
||||
report_to="tensorboard",
|
||||
remove_unused_columns=False,
|
||||
)
|
||||
|
||||
checkpoint_dir = get_last_checkpoint(training_args.output_dir)
|
||||
if checkpoint_dir is not None:
|
||||
training_args.resume_from_checkpoint = checkpoint_dir
|
||||
print(f"恢复训练自检查点: {checkpoint_dir}")
|
||||
|
||||
policy_model_path = os.path.join("model", "qwen_local")
|
||||
tokenizer = AutoTokenizer.from_pretrained(policy_model_path)
|
||||
model = DollyForCausalLM.from_pretrained(
|
||||
policy_model_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="auto"
|
||||
)
|
||||
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
reward_model_path = "./reward_model_output"
|
||||
reward_model = RewardModel(reward_model_path)
|
||||
|
||||
dataset_path = os.path.join("datas", "my_grpo_data")
|
||||
if not os.path.exists(dataset_path):
|
||||
raise ValueError(f"数据集路径不存在: {dataset_path}")
|
||||
|
||||
dataset = load_from_disk(dataset_path)
|
||||
|
||||
trainer = GRPOTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
args=training_args,
|
||||
reward_funcs=[reward_model],
|
||||
train_dataset=dataset,
|
||||
)
|
||||
|
||||
trainer.add_callback(ResumeTrainingCallback())
|
||||
|
||||
trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
|
||||
|
||||
trainer.save_model(os.path.join(training_args.output_dir, "final_model"))
|
||||
|
||||
print("GRPO训练完成!")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
"""
|
||||
torchrun --nproc_per_node=8 \
|
||||
train_grpo.py \
|
||||
--output_dir ./grpo_output \
|
||||
--per_device_train_batch_size 16 \
|
||||
--num_train_epochs 3 \
|
||||
--gradient_accumulation_steps 2 \
|
||||
--learning_rate 1e-5 \
|
||||
--logging_steps 10 \
|
||||
--save_steps 500 \
|
||||
--eval_steps 200 \
|
||||
--fp16
|
||||
"""
|
||||
@@ -0,0 +1,125 @@
|
||||
import os
|
||||
import warnings
|
||||
import torch
|
||||
from datasets import load_from_disk
|
||||
from modeling_dolly import DollyForSequenceClassification
|
||||
from transformers import (
|
||||
AutoModelForSequenceClassification,
|
||||
AutoTokenizer,
|
||||
HfArgumentParser,
|
||||
TrainerCallback,
|
||||
TrainingArguments
|
||||
)
|
||||
from trl import (
|
||||
ModelConfig,
|
||||
RewardConfig,
|
||||
RewardTrainer,
|
||||
ScriptArguments,
|
||||
get_kbit_device_map,
|
||||
get_peft_config,
|
||||
get_quantization_config,
|
||||
setup_chat_format,
|
||||
)
|
||||
from transformers.trainer_utils import get_last_checkpoint
|
||||
|
||||
class ResumeTrainingCallback(TrainerCallback):
|
||||
def on_train_begin(self, args, state, control, **kwargs):
|
||||
if state.resume_from_checkpoint is None:
|
||||
checkpoint = get_last_checkpoint(args.output_dir)
|
||||
if checkpoint is not None:
|
||||
control.should_resume = True
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = HfArgumentParser((ScriptArguments, RewardConfig, ModelConfig))
|
||||
script_args, training_args, model_args = parser.parse_args_into_dataclasses()
|
||||
training_args.gradient_checkpointing_kwargs = dict(use_reentrant=False)
|
||||
|
||||
if script_args.deepspeed_config:
|
||||
training_args.deepspeed = script_args.deepspeed_config
|
||||
|
||||
checkpoint_dir = get_last_checkpoint(training_args.output_dir)
|
||||
if checkpoint_dir is not None:
|
||||
training_args.resume_from_checkpoint = checkpoint_dir
|
||||
print(f"恢复训练自检查点: {checkpoint_dir}")
|
||||
torch_dtype = (
|
||||
model_args.torch_dtype if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype)
|
||||
)
|
||||
quantization_config = get_quantization_config(model_args)
|
||||
|
||||
model_kwargs = dict(
|
||||
revision=model_args.model_revision,
|
||||
device_map=get_kbit_device_map() if quantization_config is not None else None,
|
||||
quantization_config=quantization_config,
|
||||
use_cache=False if training_args.gradient_checkpointing else True,
|
||||
torch_dtype=torch_dtype,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
os.path.join("model", model_args.model_name_or_path),
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
use_fast=True
|
||||
)
|
||||
model = DollyForSequenceClassification.from_pretrained(
|
||||
os.path.join("model", model_args.model_name_or_path),
|
||||
num_labels=1,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
**model_kwargs
|
||||
)
|
||||
|
||||
model.config.pad_token_id = tokenizer.pad_token_id
|
||||
if tokenizer.chat_template is None:
|
||||
model, tokenizer = setup_chat_format(model, tokenizer)
|
||||
|
||||
if model_args.use_peft and model_args.lora_task_type != "SEQ_CLS":
|
||||
warnings.warn(
|
||||
"You are using a `task_type` that is different than `SEQ_CLS` for PEFT. This will lead to silent bugs"
|
||||
" Make sure to pass --lora_task_type SEQ_CLS when using this script with PEFT.",
|
||||
UserWarning,
|
||||
)
|
||||
|
||||
dataset_path = os.path.join("datas", script_args.dataset_name)
|
||||
if not os.path.exists(dataset_path):
|
||||
raise ValueError(f"数据集路径不存在: {dataset_path}")
|
||||
|
||||
dataset = load_from_disk(dataset_path)
|
||||
trainer = RewardTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
args=training_args,
|
||||
train_dataset=dataset[script_args.dataset_train_split],
|
||||
eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None,
|
||||
peft_config=get_peft_config(model_args) if model_args.use_peft else None,
|
||||
)
|
||||
trainer.add_callback(ResumeTrainingCallback())
|
||||
trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
|
||||
if not model_args.use_peft:
|
||||
trainer.save_model(training_args.output_dir)
|
||||
else:
|
||||
trainer.model.save_pretrained(training_args.output_dir)
|
||||
|
||||
if training_args.eval_strategy != "no":
|
||||
metrics = trainer.evaluate()
|
||||
trainer.log_metrics("eval", metrics)
|
||||
trainer.save_metrics("eval", metrics)
|
||||
tokenizer.save_pretrained(training_args.output_dir)
|
||||
|
||||
|
||||
"""
|
||||
torchrun --nproc_per_node=8 \
|
||||
reward_training.py \
|
||||
--model_name_or_path qwen_local \
|
||||
--dataset_name my_reward_data \
|
||||
--output_dir ./reward_model_output \
|
||||
--per_device_train_batch_size 16 \
|
||||
--num_train_epochs 3 \
|
||||
--gradient_accumulation_steps 2 \
|
||||
--learning_rate 2e-5 \
|
||||
--eval_strategy steps \
|
||||
--eval_steps 200 \
|
||||
--save_strategy steps \
|
||||
--save_steps 500 \
|
||||
--max_length 2048 \
|
||||
--deepspeed_config configs/deepspeed_z3.json \
|
||||
--bf16 \
|
||||
--report_to tensorboard
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user