fix: harden DeepSpeed SFT resume

trainer/train_full_sft_deepspeed.py

@@ -27,6 +27,7 @@ from trainer.trainer_utils import (
     add_model_profile_args,
     apply_model_profile,
     build_lm_config,
+    build_training_signature,
     assert_no_training_output_collision,
     init_model,
     is_main_process,
@@ -34,16 +35,42 @@ from trainer.trainer_utils import (
     load_deepspeed_checkpoint,
     load_deepspeed_config,
     log_training_setup,
+    normalize_resume_position,
     reduce_metrics,
     resolve_attention_type,
     save_deepspeed_checkpoint,
     save_model_weights,
     setup_seed,
+    validate_training_signature,
 )
 
 warnings.filterwarnings("ignore")
 
 
+def validate_sft_supervision(train_ds, max_samples=32):
+    checked = min(len(train_ds), max_samples)
+    supervised_tokens = 0
+    total_tokens = 0
+    valid_samples = 0
+    for index in range(checked):
+        _, labels = train_ds[index]
+        current_supervised = int(labels.ne(-100).sum().item())
+        supervised_tokens += current_supervised
+        total_tokens += int(labels.numel())
+        valid_samples += int(current_supervised > 0)
+    if valid_samples == 0:
+        raise RuntimeError(
+            "SFT data validation failed: sampled labels are all -100. "
+            "Check conversations roles and tokenizer chat-template markers."
+        )
+    if is_main_process():
+        ratio = supervised_tokens / max(total_tokens, 1)
+        Logger(
+            f"SFT supervision check: samples={checked}, valid_samples={valid_samples}, "
+            f"supervised_tokens={supervised_tokens}, supervised_ratio={ratio:.4f}"
+        )
+
+
 def init_deepspeed_runtime(args, deepspeed):
     deepspeed.init_distributed(dist_backend="nccl")
     local_rank = int(os.environ.get("LOCAL_RANK", args.local_rank if args.local_rank >= 0 else 0))
@@ -112,6 +139,7 @@ def train_epoch(args, lm_config, model_engine, loader, iters, start_step=0, wand
                 step=step,
                 wandb=wandb,
                 save_dir="../checkpoints",
+                training_signature=args.training_signature,
             )
             if is_main_process():
                 ckp = save_model_weights(lm_config, model_engine, save_dir=args.save_dir, weight=args.save_weight)
@@ -150,6 +178,7 @@ def parse_args():
     parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl")
     parser.add_argument("--from_weight", default="pretrain", type=str)
     parser.add_argument("--from_resume", default=0, type=int, choices=[0, 1])
+    parser.add_argument("--allow_resume_mismatch", default=0, type=int, choices=[0, 1])
     parser.add_argument("--overwrite_output", default=0, type=int, choices=[0, 1])
     parser.add_argument("--use_wandb", action="store_true")
     parser.add_argument("--wandb_project", type=str, default="SearchLM-Full-SFT-DS")
@@ -198,9 +227,11 @@ if __name__ == "__main__":
         Logger("torch.compile enabled before DeepSpeed initialize")
 
     train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
+    validate_sft_supervision(train_ds)
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
     loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler)
     iters = len(loader_for_count)
+    args.training_signature = build_training_signature(args, args.data_path, len(train_ds), iters)
     total_optimizer_steps = max(1, math.ceil(iters / args.accumulation_steps) * args.epochs)
     ds_config = load_deepspeed_config(args, total_optimizer_steps)
 
@@ -217,6 +248,12 @@ if __name__ == "__main__":
             lm_config,
             weight=args.save_weight,
             save_dir="../checkpoints",
+            required=True,
+        )
+        validate_training_signature(
+            ckp_data.get("training_signature"),
+            args.training_signature,
+            allow_mismatch=bool(args.allow_resume_mismatch),
         )
 
     wandb = None
@@ -249,6 +286,11 @@ if __name__ == "__main__":
 
     start_epoch = ckp_data.get("epoch", 0) if ckp_data else 0
     start_step = ckp_data.get("step", 0) if ckp_data else 0
+    start_epoch, start_step = normalize_resume_position(start_epoch, start_step, iters)
+    if args.from_resume == 1 and is_main_process():
+        Logger(f"SFT resume position: epoch={start_epoch}, step={start_step}, iters_per_epoch={iters}")
+    if start_epoch >= args.epochs and is_main_process():
+        Logger(f"Training already complete: resume epoch={start_epoch}, target epochs={args.epochs}")
 
     for epoch in range(start_epoch, args.epochs):
         args.current_epoch = epoch

trainer/train_pretrain_deepspeed.py

@@ -27,6 +27,7 @@ from trainer.trainer_utils import (
     add_model_profile_args,
     apply_model_profile,
     build_lm_config,
+    build_training_signature,
     assert_no_training_output_collision,
     init_model,
     is_main_process,
@@ -34,11 +35,13 @@ from trainer.trainer_utils import (
     load_deepspeed_checkpoint,
     load_deepspeed_config,
     log_training_setup,
+    normalize_resume_position,
     reduce_metrics,
     resolve_attention_type,
     save_deepspeed_checkpoint,
     save_model_weights,
     setup_seed,
+    validate_training_signature,
 )
 
 warnings.filterwarnings("ignore")
@@ -112,6 +115,7 @@ def train_epoch(args, lm_config, model_engine, loader, iters, start_step=0, wand
                 step=step,
                 wandb=wandb,
                 save_dir="../checkpoints",
+                training_signature=args.training_signature,
             )
             if is_main_process():
                 ckp = save_model_weights(lm_config, model_engine, save_dir=args.save_dir, weight=args.save_weight)
@@ -150,6 +154,7 @@ def parse_args():
     parser.add_argument("--data_path", type=str, default="../dataset/pretrain_t2t_mini.jsonl")
     parser.add_argument("--from_weight", default="none", type=str)
     parser.add_argument("--from_resume", default=0, type=int, choices=[0, 1])
+    parser.add_argument("--allow_resume_mismatch", default=0, type=int, choices=[0, 1])
     parser.add_argument("--overwrite_output", default=0, type=int, choices=[0, 1])
     parser.add_argument("--use_wandb", action="store_true")
     parser.add_argument("--wandb_project", type=str, default="SearchLM-Pretrain-DS")
@@ -201,6 +206,7 @@ if __name__ == "__main__":
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
     loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler)
     iters = len(loader_for_count)
+    args.training_signature = build_training_signature(args, args.data_path, len(train_ds), iters)
     total_optimizer_steps = max(1, math.ceil(iters / args.accumulation_steps) * args.epochs)
     ds_config = load_deepspeed_config(args, total_optimizer_steps)
 
@@ -217,6 +223,12 @@ if __name__ == "__main__":
             lm_config,
             weight=args.save_weight,
             save_dir="../checkpoints",
+            required=True,
+        )
+        validate_training_signature(
+            ckp_data.get("training_signature"),
+            args.training_signature,
+            allow_mismatch=bool(args.allow_resume_mismatch),
         )
 
     wandb = None
@@ -249,6 +261,11 @@ if __name__ == "__main__":
 
     start_epoch = ckp_data.get("epoch", 0) if ckp_data else 0
     start_step = ckp_data.get("step", 0) if ckp_data else 0
+    start_epoch, start_step = normalize_resume_position(start_epoch, start_step, iters)
+    if args.from_resume == 1 and is_main_process():
+        Logger(f"Pretrain resume position: epoch={start_epoch}, step={start_step}, iters_per_epoch={iters}")
+    if start_epoch >= args.epochs and is_main_process():
+        Logger(f"Training already complete: resume epoch={start_epoch}, target epochs={args.epochs}")
 
     for epoch in range(start_epoch, args.epochs):
         args.current_epoch = epoch

trainer/trainer_utils.py

@@ -154,6 +154,74 @@ def get_cuda_peak_memory_gb():
     return torch.cuda.max_memory_allocated(torch.cuda.current_device()) / (1024 ** 3)
 
 
+def build_training_signature(args, data_path, dataset_len, iters):
+    """记录恢复训练必须保持一致的关键参数。"""
+    resolved_data_path = os.path.realpath(data_path)
+    data_stat = os.stat(resolved_data_path)
+    return {
+        "data_path": resolved_data_path,
+        "data_size": data_stat.st_size,
+        "data_mtime_ns": data_stat.st_mtime_ns,
+        "dataset_len": int(dataset_len),
+        "iters_per_epoch": int(iters),
+        "epochs": int(args.epochs),
+        "batch_size": int(args.batch_size),
+        "accumulation_steps": int(args.accumulation_steps),
+        "max_seq_len": int(args.max_seq_len),
+        "learning_rate": float(args.learning_rate),
+        "dtype": getattr(args, "dtype", "float16"),
+        "world_size": dist.get_world_size() if dist.is_initialized() else 1,
+        "model_profile": getattr(args, "model_profile", ""),
+        "hidden_size": int(args.hidden_size),
+        "num_hidden_layers": int(args.num_hidden_layers),
+        "num_attention_heads": int(args.num_attention_heads),
+        "num_key_value_heads": int(args.num_key_value_heads),
+        "intermediate_size": getattr(args, "intermediate_size", None),
+        "attention_type": resolve_attention_type(args),
+        "use_moe": int(getattr(args, "use_moe", 0)),
+        "kv_lora_rank": int(getattr(args, "kv_lora_rank", 0)),
+        "q_lora_rank": int(getattr(args, "q_lora_rank", 0)),
+        "rope_dim": getattr(args, "rope_dim", None),
+    }
+
+
+def validate_training_signature(saved, current, allow_mismatch=False):
+    """阻止参数变化后伪装成严格断点续训。旧 checkpoint 没有签名时保持兼容。"""
+    if not saved:
+        if is_main_process():
+            Logger("Checkpoint has no training signature; resume compatibility cannot be fully verified")
+        return
+    mismatches = []
+    for key, current_value in current.items():
+        saved_value = saved.get(key)
+        if saved_value != current_value:
+            mismatches.append(f"{key}: checkpoint={saved_value!r}, current={current_value!r}")
+    if not mismatches:
+        return
+    message = "Resume configuration mismatch: " + "; ".join(mismatches)
+    if allow_mismatch:
+        if is_main_process():
+            Logger("WARNING: " + message)
+        return
+    raise RuntimeError(message + ". Use the original settings, or pass --allow_resume_mismatch 1 only if discontinuity is acceptable.")
+
+
+def normalize_resume_position(epoch, step, iters):
+    """若 checkpoint 位于 epoch 末尾，直接从下一 epoch 开始。"""
+    epoch = int(epoch or 0)
+    step = int(step or 0)
+    if iters <= 0:
+        return epoch, step
+    if step > iters:
+        raise RuntimeError(
+            f"Checkpoint step {step} exceeds current iters_per_epoch {iters}; "
+            "the dataset, GPU count, or batch size likely changed."
+        )
+    if step == iters:
+        return epoch + 1, 0
+    return epoch, step
+
+
 def reduce_metrics(metrics, average=True):
     """聚合所有 rank 的标量指标，避免日志只反映 rank0 局部 batch。"""
     if not dist.is_initialized():
@@ -304,7 +372,9 @@ def save_deepspeed_checkpoint(model_engine, lm_config, weight, epoch=0, step=0,
     return tag
 
 
-def load_deepspeed_checkpoint(model_engine, lm_config, weight, save_dir="../checkpoints"):
+def load_deepspeed_checkpoint(model_engine, lm_config, weight, save_dir="../checkpoints", required=False):
+    # 这些 checkpoint 由本项目生成，包含 optimizer/scheduler 等非纯权重状态。
+    os.environ.setdefault("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1")
     errors = []
     loaded_tag = None
     for tag in _candidate_deepspeed_tags(lm_config, weight, save_dir):
@@ -321,8 +391,11 @@ def load_deepspeed_checkpoint(model_engine, lm_config, weight, save_dir="../chec
         break
     #这个的else的for else 意思是如果for循环没有被break打断就走else分支
     else:
+        message = "DeepSpeed checkpoint not loaded" + (f": {'; '.join(errors)}" if errors else "")
+        if required:
+            raise RuntimeError(message)
         if is_main_process():
-            Logger("DeepSpeed checkpoint not loaded" + (f": {'; '.join(errors)}" if errors else ""))
+            Logger(message)
         return None
     client_state = client_state or {}
     tag_state = _parse_deepspeed_tag_metadata(loaded_tag or "")