github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: harden DeepSpeed SFT resume

pull/791/head
baihaochen 2026-06-11 21:45:31 +08:00
parent eaf3a8cbec
commit 5012b232ff
3 changed files with 134 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
trainer/train_full_sft_deepspeed.py
View File

@ -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

17
trainer/train_pretrain_deepspeed.py
View File

@ -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

77
trainer/trainer_utils.py
View File

@ -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 "")