admin/LLM-from-scratch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[2025.06.22] add reward model、PPO、GRPO train code

Browse Source
This commit is contained in:
hanluzhi0429@gmail.com
2025-06-22 22:57:51 +08:00
parent 2a0810012b
commit 4a3c349d33
5 changed files with 423 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+4 -2
View File
@@ -4,11 +4,13 @@ LLM-Dolly is a custom LLM built from scratch
🗓️ 更新日志
#### 2025-06-19
- 📝 pretoken corpus and update cold start code.
#### 2025-06-22
- 📝 增加奖励模型训练、PPO训练、GRPO训练
<details>
<summary>查看更多</summary>
#### 2025-06-19
- 📝 pretoken corpus and update cold start code.
#### 2025-06-05
- 📝 上传tokenizer的训练权重（1.56字符/token），调整tokenizer的格式和训练方式，和qwen2的tokenizer风格保持一致.
src/step4_cold_start.py
+25 -12
View File
@@ -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
src/step5.1_PPO.py
+152
View File
@@ -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
"""
src/step5.2_GRPO.py
+117
View File
@@ -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
"""
src/step5_train_reward_model.py
+125
View File
@@ -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
"""