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

fix: recover DeepSpeed resume step from checkpoint tag

pull/791/head
baihaochen 2026-06-04 14:20:21 +08:00
parent e9c93951cc
commit 4dcda45bfa
1 changed files with 24 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
trainer/trainer_utils.py
View File

@ -5,6 +5,7 @@ import os
import sys import sys
__package__ = "trainer" __package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import re
import json import json
import random import random
import math import math
@ -224,6 +225,16 @@ def _candidate_deepspeed_tags(lm_config, weight, save_dir="../checkpoints"):
return dedup return dedup
def _parse_deepspeed_tag_metadata(tag):
match = re.search(r"_epoch(\d+)_step(\d+)$", tag)
if not match:
return {}
return {
"epoch": int(match.group(1)),
"step": int(match.group(2)),
}
def _prune_old_deepspeed_step_tags(lm_config, weight, keep_tag, save_dir="../checkpoints"): def _prune_old_deepspeed_step_tags(lm_config, weight, keep_tag, save_dir="../checkpoints"):
if not os.path.exists(save_dir): if not os.path.exists(save_dir):
return return
@ -294,6 +305,7 @@ def save_deepspeed_checkpoint(model_engine, lm_config, weight, epoch=0, step=0,
def load_deepspeed_checkpoint(model_engine, lm_config, weight, save_dir="../checkpoints"): def load_deepspeed_checkpoint(model_engine, lm_config, weight, save_dir="../checkpoints"):
errors = [] errors = []
loaded_tag = None
for tag in _candidate_deepspeed_tags(lm_config, weight, save_dir): for tag in _candidate_deepspeed_tags(lm_config, weight, save_dir):
try: try:
load_path, client_state = model_engine.load_checkpoint(save_dir, tag=tag) load_path, client_state = model_engine.load_checkpoint(save_dir, tag=tag)
@ -303,18 +315,27 @@ def load_deepspeed_checkpoint(model_engine, lm_config, weight, save_dir="../chec
if load_path is None: if load_path is None:
errors.append(f"{tag}: load_path=None") errors.append(f"{tag}: load_path=None")
continue continue
loaded_tag = tag
break break
else: else:
if is_main_process(): if is_main_process():
Logger("DeepSpeed checkpoint not loaded" + (f": {'; '.join(errors)}" if errors else "")) Logger("DeepSpeed checkpoint not loaded" + (f": {'; '.join(errors)}" if errors else ""))
return None return None
saved_ws = client_state.get("world_size", 1) if client_state else 1 client_state = client_state or {}
tag_state = _parse_deepspeed_tag_metadata(loaded_tag or "")
for key, value in tag_state.items():
client_state.setdefault(key, value)
client_state.setdefault("world_size", 1)
saved_ws = client_state.get("world_size", 1)
current_ws = dist.get_world_size() if dist.is_initialized() else 1 current_ws = dist.get_world_size() if dist.is_initialized() else 1
if saved_ws != current_ws and is_main_process(): if saved_ws != current_ws and is_main_process():
Logger(f"DeepSpeed checkpoint world_size changed: {saved_ws} -> {current_ws}") Logger(f"DeepSpeed checkpoint world_size changed: {saved_ws} -> {current_ws}")
if is_main_process(): if is_main_process():
Logger(f"DeepSpeed checkpoint loaded: {load_path}") Logger(
return client_state or {} f"DeepSpeed checkpoint loaded: {load_path}, "
f"resume epoch={client_state.get('epoch')}, step={client_state.get('step')}"
)
return client_state
def build_lm_config(hidden_size=768, num_hidden_layers=8, use_moe=False, attention_type="gqa", kv_lora_rank=128, **kwargs): def build_lm_config(hidden_size=768, num_hidden_layers=8, use_moe=False, attention_type="gqa", kv_lora_rank=128, **kwargs):