diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index e2a9a2f..cc6e47d 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -30,10 +30,12 @@ def get_lr(current_step, total_steps, lr): return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps)) -def train_epoch(epoch, wandb): +def train_epoch(epoch, wandb, start_step=0): loss_fct = nn.CrossEntropyLoss(reduction='none') start_time = time.time() for step, (X, Y, loss_mask) in enumerate(train_loader): + if step < start_step: + continue X = X.to(args.device) Y = Y.to(args.device) loss_mask = loss_mask.to(args.device) @@ -90,6 +92,15 @@ def train_epoch(epoch, wandb): state_dict = model.state_dict() state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存 torch.save(state_dict, ckp) + + # 保存优化器状态、当前epoch和step + checkpoint = { + 'epoch': epoch, + 'step': step + 1, + 'optimizer_state_dict': optimizer.state_dict(), + 'scaler_state_dict': scaler.state_dict() + } + torch.save(checkpoint, f'{args.save_dir}/checkpoint.pth') model.train() @@ -141,6 +152,7 @@ if __name__ == "__main__": parser.add_argument('--max_seq_len', default=512, type=int) parser.add_argument('--use_moe', default=False, type=bool) parser.add_argument("--data_path", type=str, default="../dataset/sft_mini_512.jsonl") + parser.add_argument("--resume", action="store_true", help="Resume training from checkpoint") args = parser.parse_args() @@ -198,5 +210,20 @@ if __name__ == "__main__": model = DistributedDataParallel(model, device_ids=[ddp_local_rank]) iter_per_epoch = len(train_loader) - for epoch in range(args.epochs): - train_epoch(epoch, wandb) + + start_epoch = 0 + start_step = 0 + if args.resume: + checkpoint_path = f'{args.save_dir}/checkpoint.pth' + if os.path.exists(checkpoint_path): + checkpoint = torch.load(checkpoint_path) + start_epoch = checkpoint['epoch'] + start_step = checkpoint['step'] + optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + scaler.load_state_dict(checkpoint['scaler_state_dict']) + Logger(f"Resuming training from epoch {start_epoch + 1}, step {start_step}") + else: + Logger("Checkpoint not found, starting training from scratch.") + + for epoch in range(start_epoch, args.epochs): + train_epoch(epoch, wandb, start_step if epoch == start_epoch else 0) \ No newline at end of file diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index f0cfc58..a614933 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -28,70 +28,92 @@ def Logger(content): def get_lr(current_step, total_steps, lr): return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps)) +loss_fct = nn.CrossEntropyLoss(reduction='none') -def train_epoch(epoch, wandb): - loss_fct = nn.CrossEntropyLoss(reduction='none') +def train_epoch(epoch, start_step, wandb): start_time = time.time() - for step, (X, Y, loss_mask) in enumerate(train_loader): - X = X.to(args.device) - Y = Y.to(args.device) - loss_mask = loss_mask.to(args.device) + if epoch == start_epoch: + train_loader_iter = iter(train_loader) + for _ in range(start_step + 1): + next(train_loader_iter) + for step, (X, Y, loss_mask) in enumerate(train_loader_iter, start=start_step + 1): + train_step(epoch, step, X, Y, loss_mask, wandb, start_time) + else: + for step, (X, Y, loss_mask) in enumerate(train_loader): + train_step(epoch, step, X, Y, loss_mask, wandb, start_time) - lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate) - for param_group in optimizer.param_groups: - param_group['lr'] = lr - with ctx: - res = model(X) - loss = loss_fct( - res.logits.view(-1, res.logits.size(-1)), - Y.view(-1) - ).view(Y.size()) - loss = (loss * loss_mask).sum() / loss_mask.sum() - loss += res.aux_loss - loss = loss / args.accumulation_steps +def train_step(epoch, step, X, Y, loss_mask, wandb, start_time): + X = X.to(args.device) + Y = Y.to(args.device) + loss_mask = loss_mask.to(args.device) - scaler.scale(loss).backward() + lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate) + for param_group in optimizer.param_groups: + param_group['lr'] = lr - if (step + 1) % args.accumulation_steps == 0: - scaler.unscale_(optimizer) - torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + with ctx: + res = model(X) + loss = loss_fct( + res.logits.view(-1, res.logits.size(-1)), + Y.view(-1) + ).view(Y.size()) + loss = (loss * loss_mask).sum() / loss_mask.sum() + loss += res.aux_loss + loss = loss / args.accumulation_steps - scaler.step(optimizer) - scaler.update() + scaler.scale(loss).backward() - optimizer.zero_grad(set_to_none=True) + if (step + 1) % args.accumulation_steps == 0: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - if step % args.log_interval == 0: - spend_time = time.time() - start_time - Logger( - 'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format( - epoch + 1, - args.epochs, - step, - iter_per_epoch, - loss.item() * args.accumulation_steps, - optimizer.param_groups[-1]['lr'], - spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60)) + scaler.step(optimizer) + scaler.update() - if (wandb is not None) and (not ddp or dist.get_rank() == 0): - wandb.log({"loss": loss.item() * args.accumulation_steps, - "lr": optimizer.param_groups[-1]['lr'], - "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60}) + optimizer.zero_grad(set_to_none=True) - if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0): - model.eval() - moe_path = '_moe' if lm_config.use_moe else '' - ckp = f'{args.save_dir}/pretrain_{lm_config.hidden_size}{moe_path}.pth' + if step % args.log_interval == 0: + spend_time = time.time() - start_time + Logger( + 'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format( + epoch + 1, + args.epochs, + step, + iter_per_epoch, + loss.item() * args.accumulation_steps, + optimizer.param_groups[-1]['lr'], + spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60)) - if isinstance(model, torch.nn.parallel.DistributedDataParallel): - state_dict = model.module.state_dict() - else: - state_dict = model.state_dict() + if (wandb is not None) and (not ddp or dist.get_rank() == 0): + wandb.log({"loss": loss.item() * args.accumulation_steps, + "lr": optimizer.param_groups[-1]['lr'], + "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60}) - state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存 - torch.save(state_dict, ckp) - model.train() + if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0): + model.eval() + moe_path = '_moe' if lm_config.use_moe else '' + ckp = f'{args.save_dir}/pretrain_{lm_config.hidden_size}{moe_path}.pth' + + if isinstance(model, torch.nn.parallel.DistributedDataParallel): + state_dict = model.module.state_dict() + else: + state_dict = model.state_dict() + + state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存 + torch.save(state_dict, ckp) + + # 保存训练状态 + training_state = { + 'epoch': epoch, + 'step': step, + 'optimizer_state_dict': optimizer.state_dict(), + 'scaler_state_dict': scaler.state_dict() if scaler is not None else None + } + training_state_path = f'{args.save_dir}/training_state.pth' + torch.save(training_state, training_state_path) + + model.train() def init_model(lm_config): @@ -115,28 +137,28 @@ def init_distributed_mode(): # torchrun --nproc_per_node 2 1-pretrain.py if __name__ == "__main__": + # 命令行参数解析 parser = argparse.ArgumentParser(description="MiniMind Pretraining") - parser.add_argument("--out_dir", type=str, default="../out") - # 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。 - parser.add_argument("--epochs", type=int, default=1) - parser.add_argument("--batch_size", type=int, default=32) + parser.add_argument("--out_dir", type=str, default="../out") # 输出目录 + parser.add_argument("--epochs", type=int, default=8) # 训练轮次 + parser.add_argument("--batch_size", type=int, default=8) # 物理批次大小 parser.add_argument("--learning_rate", type=float, default=5e-4) parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") - parser.add_argument("--dtype", type=str, default="bfloat16") - parser.add_argument("--use_wandb", action="store_true") + parser.add_argument("--dtype", type=str, default="bfloat16") # 混合精度类型 + parser.add_argument("--use_wandb", action="store_true") # 是否使用WandB parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain") - parser.add_argument("--num_workers", type=int, default=1) - parser.add_argument("--ddp", action="store_true") - parser.add_argument("--accumulation_steps", type=int, default=8) - parser.add_argument("--grad_clip", type=float, default=1.0) - parser.add_argument("--warmup_iters", type=int, default=0) - parser.add_argument("--log_interval", type=int, default=100) - parser.add_argument("--save_interval", type=int, default=100) - parser.add_argument('--local_rank', type=int, default=-1) - parser.add_argument('--hidden_size', default=512, type=int) - parser.add_argument('--num_hidden_layers', default=8, type=int) - parser.add_argument('--max_seq_len', default=512, type=int) - parser.add_argument('--use_moe', default=False, type=bool) + parser.add_argument("--num_workers", type=int, default=1) # 数据加载线程数 + parser.add_argument("--ddp", action="store_true") # 是否启用DDP + parser.add_argument("--accumulation_steps", type=int, default=4) # 梯度累积步数 + parser.add_argument("--grad_clip", type=float, default=1.0) # 梯度裁剪阈值 + parser.add_argument("--log_interval", type=int, default=100) # 日志间隔（步） + parser.add_argument("--save_interval", type=int, default=100) # 保存间隔（步） + parser.add_argument('--local_rank', type=int, default=-1) # DDP自动传入参数 + # 模型架构参数 + parser.add_argument('--hidden_size', default=512, type=int) # 隐藏层维度 + parser.add_argument('--num_hidden_layers', default=8, type=int) # 层数 + parser.add_argument('--max_seq_len', default=4096, type=int) # 序列长度 + parser.add_argument('--use_moe', default=False, type=bool) # 是否使用MoE parser.add_argument("--data_path", type=str, default="../dataset/pretrain_hq.jsonl") args = parser.parse_args() @@ -194,5 +216,19 @@ if __name__ == "__main__": model = DistributedDataParallel(model, device_ids=[ddp_local_rank]) iter_per_epoch = len(train_loader) - for epoch in range(args.epochs): - train_epoch(epoch, wandb) + + # 检查是否存在保存的训练状态文件 + training_state_path = f'{args.save_dir}/training_state.pth' + if os.path.exists(training_state_path): + training_state = torch.load(training_state_path) + start_epoch = training_state['epoch'] + start_step = training_state['step'] + optimizer.load_state_dict(training_state['optimizer_state_dict']) + if scaler is not None and training_state['scaler_state_dict'] is not None: + scaler.load_state_dict(training_state['scaler_state_dict']) + else: + start_epoch = 0 + start_step = 0 + + for epoch in range(start_epoch, args.epochs): + train_epoch(epoch, start_step if epoch == start_epoch else 0, wandb) \ No newline at end of file