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,11 +28,22 @@ 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()
+    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)
+
+
+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)
@@ -91,6 +102,17 @@ def train_epoch(epoch, wandb):
 
         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()
 
 
@@ -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)