From 4a3c349d3305a496a52a09e2ad4b4e646d9c4d55 Mon Sep 17 00:00:00 2001 From: "hanluzhi0429@gmail.com" Date: Sun, 22 Jun 2025 22:57:51 +0800 Subject: [PATCH] =?UTF-8?q?[2025.06.22]=20add=20reward=20model=E3=80=81PPO?= =?UTF-8?q?=E3=80=81GRPO=20train=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 6 +- src/step4_cold_start.py | 37 +++++--- src/step5.1_PPO.py | 152 ++++++++++++++++++++++++++++++++ src/step5.2_GRPO.py | 117 ++++++++++++++++++++++++ src/step5_train_reward_model.py | 125 ++++++++++++++++++++++++++ 5 files changed, 423 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index ba95ac4..d91a55b 100755 --- a/README.md +++ b/README.md @@ -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训练
查看更多 +#### 2025-06-19 +- 📝 pretoken corpus and update cold start code. #### 2025-06-05 - 📝 上传tokenizer的训练权重（1.56字符/token），调整tokenizer的格式和训练方式，和qwen2的tokenizer风格保持一致. diff --git a/src/step4_cold_start.py b/src/step4_cold_start.py index bf15f70..e193e61 100644 --- a/src/step4_cold_start.py +++ b/src/step4_cold_start.py @@ -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)) \ No newline at end of file + 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 \ No newline at end of file diff --git a/src/step5.1_PPO.py b/src/step5.1_PPO.py index e69de29..3b1accc 100644 --- a/src/step5.1_PPO.py +++ b/src/step5.1_PPO.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 + """ diff --git a/src/step5.2_GRPO.py b/src/step5.2_GRPO.py index e69de29..707f5ae 100644 --- a/src/step5.2_GRPO.py +++ b/src/step5.2_GRPO.py @@ -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 +""" \ No newline at end of file diff --git a/src/step5_train_reward_model.py b/src/step5_train_reward_model.py index e69de29..bdd7464 100644 --- a/src/step5_train_reward_model.py +++ b/src/step5_train_reward_model.py @@ -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 +"""