From 0425c4f6226252b877565eb24328c507f0157979 Mon Sep 17 00:00:00 2001 From: baihaochen <1291308544@qq.com> Date: Tue, 2 Jun 2026 20:41:38 +0800 Subject: [PATCH] fix: make DeepSpeed resume checkpoints robust --- docs/resume_experiment_plan.md | 24 ++++++++++++ docs/search_shortqa_project.md | 42 ++++++++++++++++++++ trainer/trainer_utils.py | 72 ++++++++++++++++++++++++++++++---- 3 files changed, 130 insertions(+), 8 deletions(-) diff --git a/docs/resume_experiment_plan.md b/docs/resume_experiment_plan.md index 617c149..b6bddb9 100644 --- a/docs/resume_experiment_plan.md +++ b/docs/resume_experiment_plan.md @@ -69,6 +69,30 @@ deepspeed --num_gpus=7 train_search_grpo_deepspeed.py \ --accumulation_steps 4 ``` +断点续训： + +```bash +cd trainer + +deepspeed --num_gpus=7 train_search_grpo_deepspeed.py \ + --model_profile SearchLM-300M \ + --data_path ../dataset/search_shortqa/search_shortqa_train.jsonl \ + --from_weight full_sft_search \ + --save_weight search_grpo \ + --num_generations 6 \ + --dtype float16 \ + --batch_size 1 \ + --accumulation_steps 4 \ + --save_interval 25 \ + --from_resume 1 +``` + +注意： + +- 恢复时 `--epochs` 是总 epoch，不是剩余 epoch。 +- 严格连续训练时不要改变 GPU 数量、batch、accumulation、数据文件和模型结构。 +- `../checkpoints/` 是完整训练状态；`../out/` 是推理权重。在线 GPU 环境要保证这两个目录都在持久磁盘上。 + 3090 不支持 bf16 的完整高效路径时，优先用 `float16`。如果 GRPO 显存吃紧，先降： - `--batch_size 1` diff --git a/docs/search_shortqa_project.md b/docs/search_shortqa_project.md index 3091745..25e894d 100644 --- a/docs/search_shortqa_project.md +++ b/docs/search_shortqa_project.md @@ -246,6 +246,48 @@ pip install deepspeed 默认使用 `trainer/ds_config_zero2.json`，采用 ZeRO-2 切分 optimizer state 和 gradient，适合 7×RTX3090 的 100M/300M/500M 训练链路。完整 DeepSpeed checkpoint 保存在 `../checkpoints/`，同时 rank0 会导出一份轻量 `.pth` 权重到 `../out/` 用于评测和推理。 +## 断点续训 + +DeepSpeed 训练脚本支持 `--from_resume 1`。恢复的是 `../checkpoints/` 下的完整训练状态，包括模型、optimizer、scheduler、DeepSpeed engine 和已训练到的 `epoch/step`；`../out/` 下的 `.pth` 只是轻量推理权重，不能单独用于连续训练。 + +第一次训练建议降低 `save_interval`，减少在线 GPU 中断后的损失： + +```bash +cd trainer + +deepspeed --num_gpus=7 train_pretrain_deepspeed.py \ + --model_profile SearchLM-300M \ + --data_path ../dataset/pretrain_t2t_mini.jsonl \ + --save_weight pretrain_searchlm \ + --dtype float16 \ + --save_interval 200 +``` + +中断后用同样的核心参数恢复： + +```bash +cd trainer + +deepspeed --num_gpus=7 train_pretrain_deepspeed.py \ + --model_profile SearchLM-300M \ + --data_path ../dataset/pretrain_t2t_mini.jsonl \ + --save_weight pretrain_searchlm \ + --dtype float16 \ + --save_interval 200 \ + --from_resume 1 +``` + +恢复时需要保持以下参数一致：`--model_profile`、模型结构相关参数、`--save_weight`、`--data_path`、`--max_seq_len`、GPU 数量、`--batch_size` 和 `--accumulation_steps`。`--epochs` 表示总目标 epoch 数，不是剩余 epoch 数。 + +恢复成功时日志会出现： + +```text +DeepSpeed checkpoint loaded: ... +Epoch [x/y]: 跳过前N个step，从step N+1开始 +``` + +为了避免中断时覆盖最后一个可用 checkpoint，DeepSpeed checkpoint 会保存成 step 级目录，例如 `search_grpo_1024_mla_epoch0_step50`，并用 `search_grpo_1024_mla_latest` 指针记录最后完整保存的 checkpoint。 + 推荐完整训练顺序： ```bash diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index ac6a633..21c4610 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -169,6 +169,55 @@ def get_deepspeed_tag(lm_config, weight): return f"{weight}_{lm_config.hidden_size}{get_model_suffix(lm_config)}" +def get_deepspeed_step_tag(lm_config, weight, epoch, step): + return f"{get_deepspeed_tag(lm_config, weight)}_epoch{epoch}_step{step}" + + +def get_deepspeed_latest_file(lm_config, weight, save_dir="../checkpoints"): + return os.path.join(save_dir, f"{get_deepspeed_tag(lm_config, weight)}_latest") + + +def _read_deepspeed_latest_tag(lm_config, weight, save_dir="../checkpoints"): + latest_file = get_deepspeed_latest_file(lm_config, weight, save_dir) + if not os.path.exists(latest_file): + return None + with open(latest_file, "r", encoding="utf-8") as f: + tag = f.read().strip() + return tag or None + + +def _write_deepspeed_latest_tag(lm_config, weight, tag, save_dir="../checkpoints"): + latest_file = get_deepspeed_latest_file(lm_config, weight, save_dir) + tmp_file = latest_file + ".tmp" + with open(tmp_file, "w", encoding="utf-8") as f: + f.write(tag) + os.replace(tmp_file, latest_file) + + +def _candidate_deepspeed_tags(lm_config, weight, save_dir="../checkpoints"): + base_tag = get_deepspeed_tag(lm_config, weight) + candidates = [] + latest_tag = _read_deepspeed_latest_tag(lm_config, weight, save_dir) + if latest_tag: + candidates.append(latest_tag) + candidates.append(base_tag) + if os.path.exists(save_dir): + step_tags = [] + prefix = f"{base_tag}_epoch" + for name in os.listdir(save_dir): + path = os.path.join(save_dir, name) + if name.startswith(prefix) and os.path.isdir(path): + step_tags.append((os.path.getmtime(path), name)) + candidates.extend(name for _, name in sorted(step_tags, reverse=True)) + dedup = [] + seen = set() + for tag in candidates: + if tag not in seen: + seen.add(tag) + dedup.append(tag) + return dedup + + def save_deepspeed_checkpoint(model_engine, lm_config, weight, epoch=0, step=0, wandb=None, save_dir="../checkpoints", **kwargs): """DeepSpeed checkpoint 必须所有 rank 都调用，不能只在 rank0 保存。""" os.makedirs(save_dir, exist_ok=True) @@ -184,22 +233,29 @@ def save_deepspeed_checkpoint(model_engine, lm_config, weight, epoch=0, step=0, "wandb_id": wandb_id, } client_state.update({k: v for k, v in kwargs.items() if v is not None}) - tag = get_deepspeed_tag(lm_config, weight) + tag = get_deepspeed_step_tag(lm_config, weight, epoch, step) model_engine.save_checkpoint(save_dir, tag=tag, client_state=client_state) if is_main_process(): + _write_deepspeed_latest_tag(lm_config, weight, tag, save_dir) Logger(f"DeepSpeed checkpoint saved: {save_dir}/{tag}") return tag def load_deepspeed_checkpoint(model_engine, lm_config, weight, save_dir="../checkpoints"): - tag = get_deepspeed_tag(lm_config, weight) - try: - load_path, client_state = model_engine.load_checkpoint(save_dir, tag=tag) - except Exception as exc: + errors = [] + for tag in _candidate_deepspeed_tags(lm_config, weight, save_dir): + try: + load_path, client_state = model_engine.load_checkpoint(save_dir, tag=tag) + except Exception as exc: + errors.append(f"{tag}: {exc}") + continue + if load_path is None: + errors.append(f"{tag}: load_path=None") + continue + break + else: if is_main_process(): - Logger(f"DeepSpeed checkpoint not loaded: {exc}") - return None - if load_path is None: + Logger("DeepSpeed checkpoint not loaded" + (f": {'; '.join(errors)}" if errors else "")) return None saved_ws = client_state.get("world_size", 1) if client_state else 1 current_ws = dist.get_world_size() if dist.is_initialized() else 1