Add files via upload

增加接续训练功能

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)

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):
+def train_epoch(epoch, start_step, wandb):
-    loss_fct = nn.CrossEntropyLoss(reduction='none')
     start_time = time.time()
-    for step, (X, Y, loss_mask) in enumerate(train_loader):
+    if epoch == start_epoch:
-        X = X.to(args.device)
+        train_loader_iter = iter(train_loader)
-        Y = Y.to(args.device)
+        for _ in range(start_step + 1):
-        loss_mask = loss_mask.to(args.device)
+            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:
+def train_step(epoch, step, X, Y, loss_mask, wandb, start_time):
-            res = model(X)
+    X = X.to(args.device)
-            loss = loss_fct(
+    Y = Y.to(args.device)
-                res.logits.view(-1, res.logits.size(-1)),
+    loss_mask = loss_mask.to(args.device)
-                Y.view(-1)
-            ).view(Y.size())
-            loss = (loss * loss_mask).sum() / loss_mask.sum()
-            loss += res.aux_loss
-            loss = loss / args.accumulation_steps
 
-        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:
+    with ctx:
-            scaler.unscale_(optimizer)
+        res = model(X)
-            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        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.scale(loss).backward()
-            scaler.update()
 
-            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:
+        scaler.step(optimizer)
-            spend_time = time.time() - start_time
+        scaler.update()
-            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 (wandb is not None) and (not ddp or dist.get_rank() == 0):
+        optimizer.zero_grad(set_to_none=True)
-                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})
 
-        if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
+    if step % args.log_interval == 0:
-            model.eval()
+        spend_time = time.time() - start_time
-            moe_path = '_moe' if lm_config.use_moe else ''
+        Logger(
-            ckp = f'{args.save_dir}/pretrain_{lm_config.hidden_size}{moe_path}.pth'
+            '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):
+        if (wandb is not None) and (not ddp or dist.get_rank() == 0):
-                state_dict = model.module.state_dict()
+            wandb.log({"loss": loss.item() * args.accumulation_steps,
-            else:
+                       "lr": optimizer.param_groups[-1]['lr'],
-                state_dict = model.state_dict()
+                       "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()}  # 半精度保存
+    if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
-            torch.save(state_dict, ckp)
+        model.eval()
-            model.train()
+        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")
+    parser.add_argument("--out_dir", type=str, default="../out")  # 输出目录
-    # 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。
+    parser.add_argument("--epochs", type=int, default=8)          # 训练轮次
-    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=8)     # 物理批次大小
-    parser.add_argument("--batch_size", type=int, default=32)
     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("--dtype", type=str, default="bfloat16")  # 混合精度类型
-    parser.add_argument("--use_wandb", action="store_true")
+    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("--num_workers", type=int, default=1)     # 数据加载线程数
-    parser.add_argument("--ddp", action="store_true")
+    parser.add_argument("--ddp", action="store_true")             # 是否启用DDP
-    parser.add_argument("--accumulation_steps", type=int, default=8)
+    parser.add_argument("--accumulation_steps", type=int, default=4)  # 梯度累积步数
-    parser.add_argument("--grad_clip", type=float, default=1.0)
+    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("--log_interval", type=int, default=100)
+    parser.add_argument("--save_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('--local_rank', type=int, default=-1)
+    # 模型架构参数
-    parser.add_argument('--hidden_size', default=512, type=int)
+    parser.add_argument('--hidden_size', default=512, type=int)   # 隐藏层维度
-    parser.add_argument('--num_hidden_layers', default=8, 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('--max_seq_len', default=4096, type=int)   # 序列长度
-    parser.add_argument('--use_moe', default=False, type=bool)
+    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)