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

feat: enable MPS across training entrypoints

mengxin 2026-09-06 02:41:42 +08:00
parent 3d4d313617
commit fce15c0bf4
9 changed files with 142 additions and 103 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
eval_llm.py
View File

@ -6,7 +6,7 @@ import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
from model.model_lora import * from model.model_lora import *
from trainer.trainer_utils import setup_seed, get_model_params from trainer.trainer_utils import setup_seed, get_model_params, get_default_device
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
def init_model(args): def init_model(args):
@ -45,7 +45,7 @@ def main():
parser.add_argument('--open_thinking', default=0, type=int, help="是否开启自适应思考（0=否，1=是）") parser.add_argument('--open_thinking', default=0, type=int, help="是否开启自适应思考（0=否，1=是）")
parser.add_argument('--historys', default=0, type=int, help="携带历史对话轮数（需为偶数，0表示不携带历史）") parser.add_argument('--historys', default=0, type=int, help="携带历史对话轮数（需为偶数，0表示不携带历史）")
parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）") parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）")
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备") parser.add_argument('--device', default=get_default_device(), type=str, help="运行设备 (cuda:0/mps/cpu)")
args = parser.parse_args() args = parser.parse_args()
prompts = [ prompts = [
@ -91,4 +91,4 @@ def main():
print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n') print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n')
if __name__ == "__main__": if __name__ == "__main__":
main() main()

39
trainer/train_agent.py
View File

@ -16,7 +16,6 @@ import warnings
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
import torch.distributed as dist import torch.distributed as dist
from contextlib import nullcontext
from torch import optim from torch import optim
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler from torch.utils.data import DataLoader, DistributedSampler
@ -24,7 +23,11 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
from transformers import AutoTokenizer from transformers import AutoTokenizer
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
from dataset.lm_dataset import AgentRLDataset from dataset.lm_dataset import AgentRLDataset
from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint,
init_distributed_mode, setup_seed,
SkipBatchSampler, init_model, LMForRewardModel,
get_default_device, get_autocast_context,
get_grad_scaler, should_pin_memory)
from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
@ -330,11 +333,13 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
policy_loss = (((per_token_loss * completion_mask).sum(dim=1)[valid_rows] / token_counts[valid_rows].clamp(min=1)).mean() policy_loss = (((per_token_loss * completion_mask).sum(dim=1)[valid_rows] / token_counts[valid_rows].clamp(min=1)).mean()
if valid_rows.any() else per_token_loss.sum() * 0.0) if valid_rows.any() else per_token_loss.sum() * 0.0)
loss = (policy_loss + aux_loss) / args.accumulation_steps loss = (policy_loss + aux_loss) / args.accumulation_steps
loss.backward() scaler.scale(loss).backward()
if step % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) if args.grad_clip > 0:
optimizer.step(); scheduler.step(); optimizer.zero_grad() scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer); scaler.update(); scheduler.step(); optimizer.zero_grad()
if step % args.log_interval == 0 or step == iters: if step % args.log_interval == 0 or step == iters:
pl = loss.item() * args.accumulation_steps pl = loss.item() * args.accumulation_steps
@ -357,7 +362,8 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
state_dict = raw_model.state_dict() state_dict = raw_model.state_dict()
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp) torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler) epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints',
scheduler=scheduler, scaler=scaler)
model.train() model.train()
del state_dict del state_dict
@ -367,8 +373,10 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask
if last_step > start_step and last_step % args.accumulation_steps != 0: if last_step > start_step and last_step % args.accumulation_steps != 0:
if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) if args.grad_clip > 0:
optimizer.step(); scheduler.step(); optimizer.zero_grad() scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer); scaler.update(); scheduler.step(); optimizer.zero_grad()
if __name__ == "__main__": if __name__ == "__main__":
@ -378,8 +386,8 @@ 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=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型 bfloat16/float16") parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "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="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
@ -412,7 +420,7 @@ if __name__ == "__main__":
parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_agent", help="SGLang共享存储路径") parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_agent", help="SGLang共享存储路径")
args = parser.parse_args() args = parser.parse_args()
local_rank = init_distributed_mode() local_rank = init_distributed_mode(args.device)
if dist.is_initialized(): args.device = f"cuda:{local_rank}" if dist.is_initialized(): args.device = f"cuda:{local_rank}"
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
@ -421,9 +429,8 @@ 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" autocast_ctx = get_autocast_context(args.device, args.dtype)
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 scaler = get_grad_scaler(args.device, args.dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
wandb = None wandb = None
if args.use_wandb and is_main_process(): if args.use_wandb and is_main_process():
@ -464,6 +471,7 @@ if __name__ == "__main__":
model.load_state_dict(ckp_data['model']) model.load_state_dict(ckp_data['model'])
optimizer.load_state_dict(ckp_data['optimizer']) optimizer.load_state_dict(ckp_data['optimizer'])
scheduler.load_state_dict(ckp_data['scheduler']) scheduler.load_state_dict(ckp_data['scheduler'])
if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler'])
start_epoch = ckp_data['epoch'] start_epoch = ckp_data['epoch']
start_step = ckp_data.get('step', 0) start_step = ckp_data.get('step', 0)
@ -480,7 +488,8 @@ if __name__ == "__main__":
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
skip = start_step if (epoch == start_epoch and start_step > 0) else 0 skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) 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, collate_fn=collate_fn) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
pin_memory=should_pin_memory(args.device), collate_fn=collate_fn)
if skip > 0: if skip > 0:
Logger(f'Epoch [{epoch+1}/{args.epochs}]: skip {start_step} steps') Logger(f'Epoch [{epoch+1}/{args.epochs}]: skip {start_step} steps')
rl_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang")) rl_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))

24
trainer/train_distillation.py
View File

@ -11,13 +11,16 @@ import warnings
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
import torch.distributed as dist import torch.distributed as dist
from contextlib import nullcontext
from torch import optim from torch import optim
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler from torch.utils.data import DataLoader, DistributedSampler
from model.model_minimind import MiniMindConfig from model.model_minimind import MiniMindConfig
from dataset.lm_dataset import SFTDataset from dataset.lm_dataset import SFTDataset
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
init_distributed_mode, setup_seed, init_model,
SkipBatchSampler, get_default_device,
get_autocast_context, get_grad_scaler,
should_pin_memory)
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
@ -151,8 +154,8 @@ 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=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "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="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
@ -177,7 +180,7 @@ if __name__ == "__main__":
args = parser.parse_args() args = parser.parse_args()
# ========== 1. 初始化环境和随机种子 ========== # ========== 1. 初始化环境和随机种子 ==========
local_rank = init_distributed_mode() local_rank = init_distributed_mode(args.device)
if dist.is_initialized(): args.device = f"cuda:{local_rank}" if dist.is_initialized(): args.device = f"cuda:{local_rank}"
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
@ -188,9 +191,8 @@ 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" autocast_ctx = get_autocast_context(args.device, args.dtype)
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 scaler = get_grad_scaler(args.device, args.dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -210,7 +212,6 @@ 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'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========
@ -235,7 +236,8 @@ if __name__ == "__main__":
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
skip = start_step if (epoch == start_epoch and start_step > 0) else 0 skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) 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) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
pin_memory=should_pin_memory(args.device))
if skip > 0: if skip > 0:
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
train_epoch(epoch, loader, len(loader) + skip, teacher_model, lm_config_student, start_step, wandb, args.alpha, args.temperature) train_epoch(epoch, loader, len(loader) + skip, teacher_model, lm_config_student, start_step, wandb, args.alpha, args.temperature)
@ -245,4 +247,4 @@ if __name__ == "__main__":
# ========== 9. 清理分布进程 ========== # ========== 9. 清理分布进程 ==========
if dist.is_initialized(): if dist.is_initialized():
dist.barrier() dist.barrier()
dist.destroy_process_group() dist.destroy_process_group()

24
trainer/train_dpo.py
View File

@ -11,13 +11,16 @@ import warnings
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
import torch.distributed as dist import torch.distributed as dist
from contextlib import nullcontext
from torch import optim from torch import optim
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler from torch.utils.data import DataLoader, DistributedSampler
from model.model_minimind import MiniMindConfig from model.model_minimind import MiniMindConfig
from dataset.lm_dataset import DPODataset from dataset.lm_dataset import DPODataset
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
init_distributed_mode, setup_seed, init_model,
SkipBatchSampler, get_default_device,
get_autocast_context, get_grad_scaler,
should_pin_memory)
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
@ -135,8 +138,8 @@ 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=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "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="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
@ -156,7 +159,7 @@ if __name__ == "__main__":
args = parser.parse_args() args = parser.parse_args()
# ========== 1. 初始化环境和随机种子 ========== # ========== 1. 初始化环境和随机种子 ==========
local_rank = init_distributed_mode() local_rank = init_distributed_mode(args.device)
if dist.is_initialized(): args.device = f"cuda:{local_rank}" if dist.is_initialized(): args.device = f"cuda:{local_rank}"
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
@ -166,9 +169,8 @@ 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" autocast_ctx = get_autocast_context(args.device, args.dtype)
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 scaler = get_grad_scaler(args.device, args.dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -190,7 +192,6 @@ 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'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========
@ -215,7 +216,8 @@ if __name__ == "__main__":
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
skip = start_step if (epoch == start_epoch and start_step > 0) else 0 skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) 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) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
pin_memory=should_pin_memory(args.device))
if skip > 0: if skip > 0:
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
train_epoch(epoch, loader, len(loader) + skip, ref_model, lm_config, start_step, wandb, args.beta) train_epoch(epoch, loader, len(loader) + skip, ref_model, lm_config, start_step, wandb, args.beta)
@ -225,4 +227,4 @@ if __name__ == "__main__":
# ========== 9. 清理分布进程 ========== # ========== 9. 清理分布进程 ==========
if dist.is_initialized(): if dist.is_initialized():
dist.barrier() dist.barrier()
dist.destroy_process_group() dist.destroy_process_group()

24
trainer/train_full_sft.py
View File

@ -10,13 +10,16 @@ import time
import warnings import warnings
import torch import torch
import torch.distributed as dist import torch.distributed as dist
from contextlib import nullcontext
from torch import optim, nn from torch import optim, nn
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler from torch.utils.data import DataLoader, DistributedSampler
from model.model_minimind import MiniMindConfig from model.model_minimind import MiniMindConfig
from dataset.lm_dataset import SFTDataset from dataset.lm_dataset import SFTDataset
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
init_distributed_mode, setup_seed, init_model,
SkipBatchSampler, get_default_device,
get_autocast_context, get_grad_scaler,
should_pin_memory)
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
@ -88,8 +91,8 @@ 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=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "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="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
@ -109,7 +112,7 @@ if __name__ == "__main__":
args = parser.parse_args() args = parser.parse_args()
# ========== 1. 初始化环境和随机种子 ========== # ========== 1. 初始化环境和随机种子 ==========
local_rank = init_distributed_mode() local_rank = init_distributed_mode(args.device)
if dist.is_initialized(): args.device = f"cuda:{local_rank}" if dist.is_initialized(): args.device = f"cuda:{local_rank}"
setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0)) setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0))
@ -119,9 +122,8 @@ 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" autocast_ctx = get_autocast_context(args.device, args.dtype)
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 scaler = get_grad_scaler(args.device, args.dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -136,7 +138,6 @@ 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'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========
@ -161,7 +162,8 @@ if __name__ == "__main__":
setup_seed(args.seed + 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 skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) 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) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
pin_memory=should_pin_memory(args.device))
if skip > 0: if skip > 0:
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb) train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
@ -171,4 +173,4 @@ if __name__ == "__main__":
# ========== 9. 清理分布进程 ========== # ========== 9. 清理分布进程 ==========
if dist.is_initialized(): if dist.is_initialized():
dist.barrier() dist.barrier()
dist.destroy_process_group() dist.destroy_process_group()

37
trainer/train_grpo.py
View File

@ -14,7 +14,6 @@ import torch
import torch.nn.functional as F import torch.nn.functional as F
import torch.distributed as dist import torch.distributed as dist
from transformers import AutoTokenizer from transformers import AutoTokenizer
from contextlib import nullcontext
from torch import optim from torch import optim
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler from torch.utils.data import DataLoader, DistributedSampler
@ -22,7 +21,11 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
from transformers import AutoModel from transformers import AutoModel
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
from dataset.lm_dataset import RLAIFDataset from dataset.lm_dataset import RLAIFDataset
from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint,
init_distributed_mode, setup_seed,
SkipBatchSampler, init_model, LMForRewardModel,
get_default_device, get_autocast_context,
get_grad_scaler, should_pin_memory)
from trainer.rollout_engine import create_rollout_engine from trainer.rollout_engine import create_rollout_engine
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
@ -142,12 +145,14 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl) per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp(min=1)).mean() policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp(min=1)).mean()
loss = (policy_loss + aux_loss) / args.accumulation_steps # scalar loss = (policy_loss + aux_loss) / args.accumulation_steps # scalar
loss.backward() scaler.scale(loss).backward()
if step % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
if args.grad_clip > 0: if args.grad_clip > 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
optimizer.step() scaler.step(optimizer)
scaler.update()
scheduler.step() scheduler.step()
optimizer.zero_grad() optimizer.zero_grad()
@ -186,7 +191,8 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
state_dict = raw_model.state_dict() state_dict = raw_model.state_dict()
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp) torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler) epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints',
scheduler=scheduler, scaler=scaler)
model.train() model.train()
del state_dict del state_dict
@ -197,8 +203,10 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
if step > start_step and step % args.accumulation_steps != 0: if step > start_step and step % args.accumulation_steps != 0:
if args.grad_clip > 0: if args.grad_clip > 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
optimizer.step() scaler.step(optimizer)
scaler.update()
scheduler.step() scheduler.step()
optimizer.zero_grad() optimizer.zero_grad()
@ -210,8 +218,8 @@ 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=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "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="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
@ -244,7 +252,7 @@ if __name__ == "__main__":
args = parser.parse_args() args = parser.parse_args()
# ========== 1. 初始化环境和随机种子 ========== # ========== 1. 初始化环境和随机种子 ==========
local_rank = init_distributed_mode() local_rank = init_distributed_mode(args.device)
if dist.is_initialized(): args.device = f"cuda:{local_rank}" if dist.is_initialized(): args.device = f"cuda:{local_rank}"
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
@ -255,9 +263,8 @@ 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" autocast_ctx = get_autocast_context(args.device, args.dtype)
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 scaler = get_grad_scaler(args.device, args.dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -303,6 +310,7 @@ if __name__ == "__main__":
model.load_state_dict(ckp_data['model']) model.load_state_dict(ckp_data['model'])
optimizer.load_state_dict(ckp_data['optimizer']) optimizer.load_state_dict(ckp_data['optimizer'])
scheduler.load_state_dict(ckp_data['scheduler']) scheduler.load_state_dict(ckp_data['scheduler'])
if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler'])
start_epoch = ckp_data['epoch'] start_epoch = ckp_data['epoch']
start_step = ckp_data.get('step', 0) start_step = ckp_data.get('step', 0)
@ -321,7 +329,8 @@ if __name__ == "__main__":
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
skip = start_step if (epoch == start_epoch and start_step > 0) else 0 skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) 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) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
pin_memory=should_pin_memory(args.device))
if skip > 0: if skip > 0:
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
grpo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang")) grpo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
@ -331,4 +340,4 @@ if __name__ == "__main__":
# ========== 9. 清理分布进程 ========== # ========== 9. 清理分布进程 ==========
if dist.is_initialized(): if dist.is_initialized():
dist.barrier() dist.barrier()
dist.destroy_process_group() dist.destroy_process_group()

24
trainer/train_lora.py
View File

@ -10,14 +10,17 @@ import time
import warnings import warnings
import torch import torch
import torch.distributed as dist import torch.distributed as dist
from contextlib import nullcontext
from torch import optim, nn from torch import optim, nn
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler from torch.utils.data import DataLoader, DistributedSampler
from model.model_minimind import MiniMindConfig from model.model_minimind import MiniMindConfig
from dataset.lm_dataset import SFTDataset from dataset.lm_dataset import SFTDataset
from model.model_lora import save_lora, apply_lora from model.model_lora import save_lora, apply_lora
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
init_distributed_mode, setup_seed, init_model,
SkipBatchSampler, get_default_device,
get_autocast_context, get_grad_scaler,
should_pin_memory)
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
@ -82,8 +85,8 @@ 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=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "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="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
@ -102,7 +105,7 @@ if __name__ == "__main__":
args = parser.parse_args() args = parser.parse_args()
# ========== 1. 初始化环境和随机种子 ========== # ========== 1. 初始化环境和随机种子 ==========
local_rank = init_distributed_mode() local_rank = init_distributed_mode(args.device)
if dist.is_initialized(): args.device = f"cuda:{local_rank}" if dist.is_initialized(): args.device = f"cuda:{local_rank}"
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
@ -112,9 +115,8 @@ 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" autocast_ctx = get_autocast_context(args.device, args.dtype)
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 scaler = get_grad_scaler(args.device, args.dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -148,7 +150,6 @@ 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'))
optimizer = optim.AdamW(lora_params, lr=args.learning_rate) optimizer = optim.AdamW(lora_params, lr=args.learning_rate)
# ========== 7. 从ckp恢复状态 ========== # ========== 7. 从ckp恢复状态 ==========
@ -173,7 +174,8 @@ if __name__ == "__main__":
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
skip = start_step if (epoch == start_epoch and start_step > 0) else 0 skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) 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) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
pin_memory=should_pin_memory(args.device))
if skip > 0: if skip > 0:
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
train_epoch(epoch, loader, len(loader) + skip, lora_params, start_step, wandb) train_epoch(epoch, loader, len(loader) + skip, lora_params, start_step, wandb)
@ -183,4 +185,4 @@ if __name__ == "__main__":
# ========== 10. 清理分布进程 ========== # ========== 10. 清理分布进程 ==========
if dist.is_initialized(): if dist.is_initialized():
dist.barrier() dist.barrier()
dist.destroy_process_group() dist.destroy_process_group()

43
trainer/train_ppo.py
View File

@ -12,7 +12,6 @@ import warnings
import torch import torch
import torch.distributed as dist import torch.distributed as dist
import torch.nn.functional as F import torch.nn.functional as F
from contextlib import nullcontext
from torch import optim, nn from torch import optim, nn
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler from torch.utils.data import DataLoader, DistributedSampler
@ -20,7 +19,11 @@ from torch.nn.utils import clip_grad_norm_
from torch.optim.lr_scheduler import CosineAnnealingLR from torch.optim.lr_scheduler import CosineAnnealingLR
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
from dataset.lm_dataset import RLAIFDataset from dataset.lm_dataset import RLAIFDataset
from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint,
init_distributed_mode, setup_seed,
SkipBatchSampler, init_model, LMForRewardModel,
get_default_device, get_autocast_context,
get_grad_scaler, should_pin_memory)
from trainer.rollout_engine import create_rollout_engine from trainer.rollout_engine import create_rollout_engine
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
@ -224,7 +227,7 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
else: else:
loss = (policy_loss + args.vf_coef * value_loss + aux_loss) / args.accumulation_steps loss = (policy_loss + args.vf_coef * value_loss + aux_loss) / args.accumulation_steps
loss.backward() scaler.scale(loss).backward()
policy_loss_sum += policy_loss.item() policy_loss_sum += policy_loss.item()
value_loss_sum += value_loss.item() value_loss_sum += value_loss.item()
@ -237,20 +240,26 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
grad_accum_step += 1 grad_accum_step += 1
if grad_accum_step % args.accumulation_steps == 0: if grad_accum_step % args.accumulation_steps == 0:
scaler.unscale_(actor_optimizer)
scaler.unscale_(critic_optimizer)
clip_grad_norm_(actor_model.parameters(), args.grad_clip) clip_grad_norm_(actor_model.parameters(), args.grad_clip)
clip_grad_norm_(critic_model.parameters(), args.grad_clip) clip_grad_norm_(critic_model.parameters(), args.grad_clip)
actor_optimizer.step() scaler.step(actor_optimizer)
critic_optimizer.step() scaler.step(critic_optimizer)
scaler.update()
actor_scheduler.step() actor_scheduler.step()
critic_scheduler.step() critic_scheduler.step()
actor_optimizer.zero_grad() actor_optimizer.zero_grad()
critic_optimizer.zero_grad() critic_optimizer.zero_grad()
if grad_accum_step % args.accumulation_steps != 0: if grad_accum_step % args.accumulation_steps != 0:
scaler.unscale_(actor_optimizer)
scaler.unscale_(critic_optimizer)
clip_grad_norm_(actor_model.parameters(), args.grad_clip) clip_grad_norm_(actor_model.parameters(), args.grad_clip)
clip_grad_norm_(critic_model.parameters(), args.grad_clip) clip_grad_norm_(critic_model.parameters(), args.grad_clip)
actor_optimizer.step() scaler.step(actor_optimizer)
critic_optimizer.step() scaler.step(critic_optimizer)
scaler.update()
actor_scheduler.step() actor_scheduler.step()
critic_scheduler.step() critic_scheduler.step()
actor_optimizer.zero_grad() actor_optimizer.zero_grad()
@ -297,7 +306,8 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
lm_checkpoint(lm_config, weight=args.save_weight, model=actor_model, optimizer=actor_optimizer, lm_checkpoint(lm_config, weight=args.save_weight, model=actor_model, optimizer=actor_optimizer,
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints',
scheduler=actor_scheduler, critic_model=critic_model, scheduler=actor_scheduler, critic_model=critic_model,
critic_optimizer=critic_optimizer, critic_scheduler=critic_scheduler) critic_optimizer=critic_optimizer, critic_scheduler=critic_scheduler,
scaler=scaler)
actor_model.train() actor_model.train()
del actor_state del actor_state
@ -314,8 +324,8 @@ 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=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "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="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
@ -353,7 +363,7 @@ if __name__ == "__main__":
args = parser.parse_args() args = parser.parse_args()
# ========== 1. 初始化环境和随机种子 ========== # ========== 1. 初始化环境和随机种子 ==========
local_rank = init_distributed_mode() local_rank = init_distributed_mode(args.device)
if dist.is_initialized(): args.device = f"cuda:{local_rank}" if dist.is_initialized(): args.device = f"cuda:{local_rank}"
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
@ -363,9 +373,8 @@ 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" autocast_ctx = get_autocast_context(args.device, args.dtype)
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 scaler = get_grad_scaler(args.device, args.dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -419,6 +428,7 @@ if __name__ == "__main__":
critic_optimizer.load_state_dict(ckp_data['critic_optimizer']) critic_optimizer.load_state_dict(ckp_data['critic_optimizer'])
actor_scheduler.load_state_dict(ckp_data['scheduler']) actor_scheduler.load_state_dict(ckp_data['scheduler'])
critic_scheduler.load_state_dict(ckp_data['critic_scheduler']) critic_scheduler.load_state_dict(ckp_data['critic_scheduler'])
if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler'])
start_epoch = ckp_data['epoch'] start_epoch = ckp_data['epoch']
start_step = ckp_data.get('step', 0) start_step = ckp_data.get('step', 0)
@ -438,7 +448,8 @@ if __name__ == "__main__":
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
skip = start_step if (epoch == start_epoch and start_step > 0) else 0 skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) 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) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
pin_memory=should_pin_memory(args.device))
if skip > 0: if skip > 0:
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
ppo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step, wandb) ppo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step, wandb)
@ -448,4 +459,4 @@ if __name__ == "__main__":
# ========== 9. 清理分布进程 ========== # ========== 9. 清理分布进程 ==========
if dist.is_initialized(): if dist.is_initialized():
dist.barrier() dist.barrier()
dist.destroy_process_group() dist.destroy_process_group()

24
trainer/train_pretrain.py
View File

@ -10,13 +10,16 @@ import time
import warnings import warnings
import torch import torch
import torch.distributed as dist import torch.distributed as dist
from contextlib import nullcontext
from torch import optim, nn from torch import optim, nn
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler from torch.utils.data import DataLoader, DistributedSampler
from model.model_minimind import MiniMindConfig from model.model_minimind import MiniMindConfig
from dataset.lm_dataset import PretrainDataset from dataset.lm_dataset import PretrainDataset
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint,
init_distributed_mode, setup_seed, init_model,
SkipBatchSampler, get_default_device,
get_autocast_context, get_grad_scaler,
should_pin_memory)
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
@ -87,8 +90,8 @@ 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=get_default_device(), help="训练设备 (cuda:0/mps/cpu)")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "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="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
@ -108,7 +111,7 @@ if __name__ == "__main__":
args = parser.parse_args() args = parser.parse_args()
# ========== 1. 初始化环境和随机种子 ========== # ========== 1. 初始化环境和随机种子 ==========
local_rank = init_distributed_mode() local_rank = init_distributed_mode(args.device)
if dist.is_initialized(): args.device = f"cuda:{local_rank}" if dist.is_initialized(): args.device = f"cuda:{local_rank}"
setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0)) setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0))
@ -118,9 +121,8 @@ 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" autocast_ctx = get_autocast_context(args.device, args.dtype)
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 scaler = get_grad_scaler(args.device, args.dtype)
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -135,7 +137,6 @@ 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'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========
@ -160,7 +161,8 @@ if __name__ == "__main__":
setup_seed(args.seed + 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 skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) 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) loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers,
pin_memory=should_pin_memory(args.device))
if skip > 0: if skip > 0:
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb) train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
@ -170,4 +172,4 @@ if __name__ == "__main__":
# ========== 9. 清理分布进程 ========== # ========== 9. 清理分布进程 ==========
if dist.is_initialized(): if dist.is_initialized():
dist.barrier() dist.barrier()
dist.destroy_process_group() dist.destroy_process_group()