From 8f7e07b8efbc638c08c7c4f6c3892f94aa6e15c6 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Tue, 28 Oct 2025 23:30:10 +0800 Subject: [PATCH 1/3] [feat] update trainer --- trainer/train_distill_reason.py | 2 +- trainer/train_distillation.py | 2 +- trainer/train_dpo.py | 2 +- trainer/train_full_sft.py | 2 +- trainer/train_grpo.py | 2 +- trainer/train_lora.py | 2 +- trainer/train_ppo.py | 2 +- trainer/train_pretrain.py | 2 +- trainer/train_spo.py | 2 +- 9 files changed, 9 insertions(+), 9 deletions(-) diff --git a/trainer/train_distill_reason.py b/trainer/train_distill_reason.py index f1e3526..3c882cd 100644 --- a/trainer/train_distill_reason.py +++ b/trainer/train_distill_reason.py @@ -15,7 +15,7 @@ from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import SFTDataset -from trainer.trainer_utils import * +from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler warnings.filterwarnings('ignore') diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index f1de5d2..848bf52 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -16,7 +16,7 @@ from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import SFTDataset -from trainer.trainer_utils import * +from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler warnings.filterwarnings('ignore') diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index b4e7b37..97716a5 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -16,7 +16,7 @@ from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import DPODataset -from trainer.trainer_utils import * +from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler warnings.filterwarnings('ignore') diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index 09fa941..6702159 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -15,7 +15,7 @@ from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import SFTDataset -from trainer.trainer_utils import * +from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler warnings.filterwarnings('ignore') diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 22a4e26..13727fe 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -19,7 +19,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR from transformers import AutoModel from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset -from trainer.trainer_utils import * +from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler warnings.filterwarnings('ignore') diff --git a/trainer/train_lora.py b/trainer/train_lora.py index b6fc2b0..1473d60 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -16,7 +16,7 @@ from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import SFTDataset from model.model_lora import save_lora, apply_lora -from trainer.trainer_utils import * +from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler warnings.filterwarnings('ignore') diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 27a82b7..51b0813 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -20,7 +20,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR from transformers import AutoModel from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset -from trainer.trainer_utils import * +from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler warnings.filterwarnings('ignore') diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 18cc445..87c79a7 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -15,7 +15,7 @@ from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import PretrainDataset -from trainer.trainer_utils import * +from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler warnings.filterwarnings('ignore') diff --git a/trainer/train_spo.py b/trainer/train_spo.py index 74dc72c..64c4e9f 100755 --- a/trainer/train_spo.py +++ b/trainer/train_spo.py @@ -19,7 +19,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR from transformers import AutoModel from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset -from trainer.trainer_utils import * +from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler warnings.filterwarnings('ignore') From acd592519308cc020dcde2e53b7b79aec0a5b7ca Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Wed, 29 Oct 2025 00:52:37 +0800 Subject: [PATCH 2/3] [feat] update trainer --- trainer/train_dpo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index 97716a5..fc2eb82 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -119,9 +119,9 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind DPO (Direct Preference Optimization)") parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录") parser.add_argument('--save_weight', default='dpo', type=str, help="保存权重的前缀名") - parser.add_argument("--epochs", type=int, default=2, help="训练轮数") + parser.add_argument("--epochs", type=int, default=1, help="训练轮数") parser.add_argument("--batch_size", type=int, default=4, help="batch size") - parser.add_argument("--learning_rate", type=float, default=5e-8, help="初始学习率（建议<=5e-8避免遗忘）") + parser.add_argument("--learning_rate", type=float, default=4e-8, help="初始学习率（建议<=5e-8避免遗忘）") parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数") From 1713c24114b6ef802da59005c6e1b4e5773a8c39 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Wed, 29 Oct 2025 10:36:28 +0800 Subject: [PATCH 3/3] [fix] model device --- trainer/train_distill_reason.py | 2 +- trainer/train_distillation.py | 4 ++-- trainer/train_dpo.py | 4 ++-- trainer/train_full_sft.py | 2 +- trainer/train_grpo.py | 22 +++++++-------------- trainer/train_lora.py | 2 +- trainer/train_ppo.py | 35 ++++++++++++++------------------- trainer/train_pretrain.py | 2 +- trainer/train_spo.py | 22 +++++++-------------- 9 files changed, 37 insertions(+), 58 deletions(-) diff --git a/trainer/train_distill_reason.py b/trainer/train_distill_reason.py index 3c882cd..a14a3aa 100644 --- a/trainer/train_distill_reason.py +++ b/trainer/train_distill_reason.py @@ -140,7 +140,7 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 定义模型、数据、优化器 ========== - model, tokenizer = init_model(lm_config, args.from_weight) + model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index 848bf52..2ef0c17 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -184,9 +184,9 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 定义学生和教师模型 ========== - model, tokenizer = init_model(lm_config_student, args.from_student_weight) + model, tokenizer = init_model(lm_config_student, args.from_student_weight, device=args.device) Logger(f'学生模型总参数量：{sum(p.numel() for p in model.parameters()) / 1e6:.3f} M') - teacher_model, _ = init_model(lm_config_teacher, args.from_teacher_weight) + teacher_model, _ = init_model(lm_config_teacher, args.from_teacher_weight, device=args.device) teacher_model.eval() teacher_model.requires_grad_(False) Logger(f'教师模型总参数量：{sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f} M') diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index fc2eb82..9dfe63b 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -166,10 +166,10 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 定义模型和参考模型 ========== - model, tokenizer = init_model(lm_config, args.from_weight) + model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) Logger(f'策略模型总参数量：{sum(p.numel() for p in model.parameters()) / 1e6:.3f} M') # 初始化参考模型（ref_model冻结） - ref_model, _ = init_model(lm_config, args.from_weight) + ref_model, _ = init_model(lm_config, args.from_weight, device=args.device) ref_model.eval() ref_model.requires_grad_(False) Logger(f'参考模型总参数量：{sum(p.numel() for p in ref_model.parameters()) / 1e6:.3f} M') diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index 6702159..9001dab 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -128,7 +128,7 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 定义模型、数据、优化器 ========== - model, tokenizer = init_model(lm_config, args.from_weight) + model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 13727fe..afebc7c 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -19,7 +19,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR from transformers import AutoModel from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset -from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler +from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model warnings.filterwarnings('ignore') @@ -240,25 +240,17 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 初始化模型和数据 ========== - tokenizer = AutoTokenizer.from_pretrained('../model/') - moe_suffix = '_moe' if lm_config.use_moe else '' base_weight = "reason" if args.reasoning == 1 else "full_sft" - ckp = f'{args.save_dir}/{base_weight}_{lm_config.hidden_size}{moe_suffix}.pth' - state_dict = torch.load(ckp, map_location=args.device) # Policy模型 - model = MiniMindForCausalLM(lm_config) - model.load_state_dict(state_dict, strict=False) - model = model.to(args.device) - Logger(f'Policy模型总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} M') + model, tokenizer = init_model(lm_config, base_weight, device=args.device) # Reference模型 - ref_model = MiniMindForCausalLM(lm_config) - ref_model.load_state_dict(state_dict, strict=False) - ref_model.eval().requires_grad_(False) - ref_model = ref_model.to(args.device) + ref_model, _ = init_model(lm_config, base_weight, device=args.device) + ref_model = ref_model.eval().requires_grad_(False) # Reward模型 reward_model = AutoModel.from_pretrained( - args.reward_model_path, device_map="cuda", torch_dtype=torch.float16, trust_remote_code=True - ).to(args.device).eval().requires_grad_(False) + args.reward_model_path, torch_dtype=torch.float16, trust_remote_code=True + ) + reward_model = reward_model.to(args.device).eval().requires_grad_(False) reward_tokenizer = AutoTokenizer.from_pretrained(args.reward_model_path, trust_remote_code=True) # 数据和优化器 train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len) diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 1473d60..a88a1af 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -123,7 +123,7 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 定义模型、应用LoRA、冻结非LoRA参数 ========== - model, tokenizer = init_model(lm_config, args.from_weight) + model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) apply_lora(model) # 统计参数 diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 51b0813..5b49aa3 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -20,7 +20,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR from transformers import AutoModel from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset -from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler +from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model warnings.filterwarnings('ignore') @@ -290,33 +290,28 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 初始化模型和数据 ========== - tokenizer = AutoTokenizer.from_pretrained('../model/', padding_side='left') - moe_suffix = '_moe' if lm_config.use_moe else '' base_weight = "reason" if args.reasoning == 1 else "full_sft" + # Actor模型 + actor_model, tokenizer = init_model(lm_config, base_weight, device=args.device) + tokenizer.padding_side = 'left' # PPO需要左侧padding + # Old Actor模型 + old_actor_model, _ = init_model(lm_config, base_weight, device=args.device) + old_actor_model = old_actor_model.eval().requires_grad_(False) + # Reference模型 + ref_model, _ = init_model(lm_config, base_weight, device=args.device) + ref_model = ref_model.eval().requires_grad_(False) + # Critic模型 + moe_suffix = '_moe' if lm_config.use_moe else '' ckp = f'{args.save_dir}/{base_weight}_{lm_config.hidden_size}{moe_suffix}.pth' state_dict = torch.load(ckp, map_location=args.device) - # Actor模型 - actor_model = MiniMindForCausalLM(lm_config) - actor_model.load_state_dict(state_dict, strict=False) - actor_model = actor_model.to(args.device) - Logger(f'Actor模型总参数量：{sum(p.numel() for p in actor_model.parameters() if p.requires_grad) / 1e6:.3f} M') - # Old Actor模型 - old_actor_model = MiniMindForCausalLM(lm_config) - old_actor_model.load_state_dict(state_dict, strict=False) - old_actor_model = old_actor_model.eval().requires_grad_(False).to(args.device) - # Reference模型 - ref_model = MiniMindForCausalLM(lm_config) - ref_model.load_state_dict(state_dict, strict=False) - ref_model = ref_model.eval().requires_grad_(False).to(args.device) - # Critic模型 critic_model = CriticModel(lm_config) critic_model.load_state_dict(state_dict, strict=False) critic_model = critic_model.to(args.device) - Logger(f'Critic模型总参数量：{sum(p.numel() for p in critic_model.parameters() if p.requires_grad) / 1e6:.3f} M') # Reward模型 reward_model = AutoModel.from_pretrained( - args.reward_model_path, device_map="cuda", torch_dtype=torch.float32, trust_remote_code=True - ).to(args.device).eval().requires_grad_(False) + args.reward_model_path, torch_dtype=torch.float16, trust_remote_code=True + ) + reward_model = reward_model.to(args.device).eval().requires_grad_(False) reward_tokenizer = AutoTokenizer.from_pretrained(args.reward_model_path, trust_remote_code=True) # 数据和优化器 train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=(args.max_seq_len + args.max_gen_len)) diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 87c79a7..83705dd 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -127,7 +127,7 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 定义模型、数据、优化器 ========== - model, tokenizer = init_model(lm_config, args.from_weight) + model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) diff --git a/trainer/train_spo.py b/trainer/train_spo.py index 64c4e9f..34a7451 100755 --- a/trainer/train_spo.py +++ b/trainer/train_spo.py @@ -19,7 +19,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR from transformers import AutoModel from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset -from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler +from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model warnings.filterwarnings('ignore') @@ -287,25 +287,17 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 初始化模型（Policy, Ref, Reward）和Value Tracker、数据 ========== - tokenizer = AutoTokenizer.from_pretrained('../model/') - moe_suffix = '_moe' if lm_config.use_moe else '' base_weight = "reason" if args.reasoning == 1 else "full_sft" - ckp = f'{args.save_dir}/{base_weight}_{lm_config.hidden_size}{moe_suffix}.pth' - state_dict = torch.load(ckp, map_location=args.device) # Policy模型 - model = MiniMindForCausalLM(lm_config) - model.load_state_dict(state_dict, strict=False) - model = model.to(args.device) - Logger(f'Policy模型总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} M') + model, tokenizer = init_model(lm_config, base_weight, device=args.device) # Reference模型 - ref_model = MiniMindForCausalLM(lm_config) - ref_model.load_state_dict(state_dict, strict=False) - ref_model.eval().requires_grad_(False) - ref_model = ref_model.to(args.device) + ref_model, _ = init_model(lm_config, base_weight, device=args.device) + ref_model = ref_model.eval().requires_grad_(False) # Reward模型 reward_model = AutoModel.from_pretrained( - args.reward_model_path, device_map="cuda", torch_dtype=torch.float16, trust_remote_code=True - ).to(args.device).eval().requires_grad_(False) + args.reward_model_path, torch_dtype=torch.float16, trust_remote_code=True + ) + reward_model = reward_model.to(args.device).eval().requires_grad_(False) reward_tokenizer = AutoTokenizer.from_pretrained(args.reward_model_path, trust_remote_code=True) # Value Tracker value_tracker = AutoAdaptiveValueTracker(rho_mode='kl', rho_const=0.9, D_half=0.06, clip_lower=0.5, clip_upper=0.96)