diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index cd9f763..b080aef 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -101,6 +101,7 @@ if __name__ == "__main__": parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构(0=否,1=是)") parser.add_argument('--num_experts', default=4, type=int, help="MoE专家数量(仅use_moe=1时生效)") parser.add_argument('--num_experts_per_tok', default=1, type=int, help="每个token路由到的专家数(仅use_moe=1时生效)") + parser.add_argument('--seed', default=42, type=int, help="随机种子(DDP下每个rank为seed+rank,每轮为seed+epoch)") parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl", help="训练数据路径") parser.add_argument('--from_weight', default='pretrain', type=str, help="基于哪个权重训练,为none则不基于任何权重训练") parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训(0=否,1=是)") @@ -112,7 +113,7 @@ if __name__ == "__main__": # ========== 1. 初始化环境和随机种子 ========== local_rank = init_distributed_mode() if dist.is_initialized(): args.device = f"cuda:{local_rank}" - setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) + setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0)) # ========== 2. 配置目录、模型参数、检查ckp ========== os.makedirs(args.save_dir, exist_ok=True) @@ -160,7 +161,7 @@ if __name__ == "__main__": # ========== 8. 开始训练 ========== for epoch in range(start_epoch, args.epochs): train_sampler and train_sampler.set_epoch(epoch) - setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() + setup_seed(args.seed + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 1994296..42384c8 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -100,6 +100,7 @@ if __name__ == "__main__": parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构(0=否,1=是)") parser.add_argument('--num_experts', default=4, type=int, help="MoE专家数量(仅use_moe=1时生效)") parser.add_argument('--num_experts_per_tok', default=1, type=int, help="每个token路由到的专家数(仅use_moe=1时生效)") + parser.add_argument('--seed', default=42, type=int, help="随机种子(DDP下每个rank为seed+rank,每轮为seed+epoch)") parser.add_argument("--data_path", type=str, default="../dataset/pretrain_t2t_mini.jsonl", help="预训练数据路径") parser.add_argument('--from_weight', default='none', type=str, help="基于哪个权重训练,为none则从头开始") parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训(0=否,1=是)") @@ -111,7 +112,7 @@ if __name__ == "__main__": # ========== 1. 初始化环境和随机种子 ========== local_rank = init_distributed_mode() if dist.is_initialized(): args.device = f"cuda:{local_rank}" - setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) + setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0)) # ========== 2. 配置目录、模型参数、检查ckp ========== os.makedirs(args.save_dir, exist_ok=True) @@ -159,7 +160,7 @@ if __name__ == "__main__": # ========== 8. 开始训练 ========== for epoch in range(start_epoch, args.epochs): train_sampler and train_sampler.set_epoch(epoch) - setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() + setup_seed(args.seed + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)