diff --git a/trainer/train_full_sft_deepspeed.py b/trainer/train_full_sft_deepspeed.py index 1740712..567a8d6 100644 --- a/trainer/train_full_sft_deepspeed.py +++ b/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 diff --git a/trainer/train_pretrain_deepspeed.py b/trainer/train_pretrain_deepspeed.py index 8603c14..1273d0b 100644 --- a/trainer/train_pretrain_deepspeed.py +++ b/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 diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index 60d2e70..12094a4 100644 --- a/trainer/trainer_utils.py +++ b/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 "")