github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge 437eae89d6 into cc312c1cc6

Jiyang 2026-09-19 09:03:44 +08:00 committed by GitHub
parent cc312c1cc6 437eae89d6
commit 3064fdb90e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 29 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
trainer/train_agent.py
View File

@ -379,7 +379,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=1, help="训练轮数") parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=2, help="批次大小") parser.add_argument("--batch_size", type=int, default=2, help="批次大小")
parser.add_argument("--learning_rate", type=float, default=3e-7, help="学习率") parser.add_argument("--learning_rate", type=float, default=3e-7, help="学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"), help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型 bfloat16/float16") parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型 bfloat16/float16")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
@ -422,9 +422,9 @@ if __name__ == "__main__":
max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe)) max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
wandb = None wandb = None
if args.use_wandb and is_main_process(): if args.use_wandb and is_main_process():

8
trainer/train_distillation.py
View File

@ -151,7 +151,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=6, help="训练轮数") parser.add_argument("--epochs", type=int, default=6, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=32, help="batch size") parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--learning_rate", type=float, default=5e-6, help="初始学习率") parser.add_argument("--learning_rate", type=float, default=5e-6, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"), help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
@ -188,9 +188,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -210,7 +210,7 @@ if __name__ == "__main__":
Logger(f'教师模型总参数量：{sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f} M') Logger(f'教师模型总参数量：{sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f} M')
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) scaler = torch.amp.GradScaler(device_type, enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========

8
trainer/train_dpo.py
View File

@ -135,7 +135,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=1, 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("--batch_size", type=int, default=4, help="batch size")
parser.add_argument("--learning_rate", type=float, default=4e-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("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"), help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
@ -166,9 +166,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -190,7 +190,7 @@ if __name__ == "__main__":
train_ds = DPODataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_ds = DPODataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) scaler = torch.amp.GradScaler(device_type, enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========

8
trainer/train_full_sft.py
View File

@ -88,7 +88,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=2, help="训练轮数") parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=16, help="batch size") parser.add_argument("--batch_size", type=int, default=16, help="batch size")
parser.add_argument("--learning_rate", type=float, default=1e-5, help="初始学习率") parser.add_argument("--learning_rate", type=float, default=1e-5, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"), help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
@ -119,9 +119,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -136,7 +136,7 @@ if __name__ == "__main__":
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) 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_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) scaler = torch.amp.GradScaler(device_type, enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========

6
trainer/train_grpo.py
View File

@ -212,7 +212,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=1, help="训练轮数") parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=2, help="batch size") parser.add_argument("--batch_size", type=int, default=2, help="batch size")
parser.add_argument("--learning_rate", type=float, default=3e-7, help="初始学习率") parser.add_argument("--learning_rate", type=float, default=3e-7, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"), help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
@ -257,9 +257,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None

8
trainer/train_lora.py
View File

@ -82,7 +82,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=10, help="训练轮数") parser.add_argument("--epochs", type=int, default=10, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=32, help="batch size") parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率") parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"), help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
@ -112,9 +112,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -148,7 +148,7 @@ if __name__ == "__main__":
# ========== 6. 定义数据和优化器 ========== # ========== 6. 定义数据和优化器 ==========
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) scaler = torch.amp.GradScaler(device_type, enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(lora_params, lr=args.learning_rate) optimizer = optim.AdamW(lora_params, lr=args.learning_rate)
# ========== 7. 从ckp恢复状态 ========== # ========== 7. 从ckp恢复状态 ==========

6
trainer/train_ppo.py
View File

@ -315,7 +315,7 @@ if __name__ == "__main__":
parser.add_argument("--batch_size", type=int, default=2, help="batch size") parser.add_argument("--batch_size", type=int, default=2, help="batch size")
parser.add_argument("--learning_rate", type=float, default=3e-7, help="Actor学习率") parser.add_argument("--learning_rate", type=float, default=3e-7, help="Actor学习率")
parser.add_argument("--critic_learning_rate", type=float, default=5e-7, help="Critic学习率") parser.add_argument("--critic_learning_rate", type=float, default=5e-7, help="Critic学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"), help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
@ -364,9 +364,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None

8
trainer/train_pretrain.py
View File

@ -87,7 +87,7 @@ if __name__ == "__main__":
parser.add_argument("--epochs", type=int, default=2, help="训练轮数") parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=32, help="batch size") parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率") parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"), help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数") parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
@ -118,9 +118,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else args.device
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.autocast(device_type=device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -135,7 +135,7 @@ if __name__ == "__main__":
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) 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_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) scaler = torch.amp.GradScaler(device_type, enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========