fix: make DeepSpeed resume checkpoints robust

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`

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

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)
+    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:
-        if is_main_process():
+            errors.append(f"{tag}: {exc}")
-            Logger(f"DeepSpeed checkpoint not loaded: {exc}")
+            continue
-        return None
         if load_path is None:
+            errors.append(f"{tag}: load_path=None")
+            continue
+        break
+    else:
+        if is_main_process():
+            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