Skip to content

Commit 7b3c40b

Browse files
yaoyu-33yaoyu-33
authored
[ckpt] fix: select newest global recovery root (#5090)
Signed-off-by: Yu Yao <yaoyu.094@gmail.com>
1 parent 3173b86 commit 7b3c40b

2 files changed

Lines changed: 150 additions & 16 deletions

File tree

‎src/megatron/bridge/training/checkpointing.py‎

Lines changed: 39 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -149,17 +149,29 @@
149149
# ============================================================================
150150

151151

152+
def _get_global_non_persistent_checkpoint_dirs(
153+
load_dir: str | None,
154+
ckpt_cfg: CheckpointConfig,
155+
) -> tuple[str, ...]:
156+
"""Return the configured global non-persistent checkpoint sources."""
157+
if ckpt_cfg.non_persistent_global_ckpt_dir:
158+
return (ckpt_cfg.non_persistent_global_ckpt_dir,)
159+
160+
checkpoint_dirs = []
161+
for checkpoint_root in (getattr(ckpt_cfg, "save", None), load_dir):
162+
if isinstance(checkpoint_root, str) and checkpoint_root:
163+
checkpoint_dir = os.path.join(checkpoint_root, _NON_PERSISTENT_CKPT_SUBDIR)
164+
if checkpoint_dir not in checkpoint_dirs:
165+
checkpoint_dirs.append(checkpoint_dir)
166+
return tuple(checkpoint_dirs)
167+
168+
152169
def _has_global_non_persistent_checkpoint(load_dir: str | None, ckpt_cfg: CheckpointConfig) -> bool:
153170
"""Return whether the configured global non-persistent checkpoint source exists."""
154171
if ckpt_cfg.non_persistent_ckpt_type != "global":
155172
return False
156173

157-
non_persistent_global_dir = (
158-
ckpt_cfg.non_persistent_global_ckpt_dir
159-
if ckpt_cfg.non_persistent_global_ckpt_dir or load_dir is None
160-
else os.path.join(load_dir, _NON_PERSISTENT_CKPT_SUBDIR)
161-
)
162-
return checkpoint_exists(non_persistent_global_dir)
174+
return any(checkpoint_exists(path) for path in _get_global_non_persistent_checkpoint_dirs(load_dir, ckpt_cfg))
163175

164176

165177
def set_checkpoint_version(value: float) -> None:
@@ -3524,14 +3536,18 @@ def _load_base_checkpoint(
35243536
raise NotImplementedError(f"Checkpoint format {ckpt_format} not supported")
35253537

35263538
# Try to load non-persistent checkpoint first
3527-
non_persistent_global_dir = (
3528-
ckpt_cfg.non_persistent_global_ckpt_dir
3529-
if ckpt_cfg.non_persistent_global_ckpt_dir or load_dir is None
3530-
else os.path.join(load_dir, _NON_PERSISTENT_CKPT_SUBDIR)
3531-
)
3532-
non_persistent_iteration = _get_non_persistent_iteration(
3533-
non_persistent_global_dir, ckpt_cfg.non_persistent_ckpt_type, checkpointing_context
3534-
)
3539+
non_persistent_global_dir = ""
3540+
non_persistent_iteration = -1
3541+
if ckpt_cfg.non_persistent_ckpt_type == "global":
3542+
for candidate_dir in _get_global_non_persistent_checkpoint_dirs(load_dir, ckpt_cfg):
3543+
candidate_iteration = _get_non_persistent_iteration(candidate_dir, "global", checkpointing_context)
3544+
if candidate_iteration > non_persistent_iteration:
3545+
non_persistent_global_dir = candidate_dir
3546+
non_persistent_iteration = candidate_iteration
3547+
else:
3548+
non_persistent_iteration = _get_non_persistent_iteration(
3549+
non_persistent_global_dir, ckpt_cfg.non_persistent_ckpt_type, checkpointing_context
3550+
)
35353551

35363552
tracker_filename = "because load directory is not defined"
35373553
if load_dir is not None:
@@ -3541,8 +3557,15 @@ def _load_base_checkpoint(
35413557

35423558
if non_persistent_iteration != -1: # there is a non-persistent checkpoint
35433559
if non_persistent_iteration >= iteration:
3544-
# Non-persistent (global and local) state is written under checkpoint.save, already the root.
3545-
_record_dataloader_state_dir(checkpointing_context, ckpt_cfg.save)
3560+
dataloader_checkpoint_root = ckpt_cfg.save
3561+
if (
3562+
ckpt_cfg.non_persistent_ckpt_type == "global"
3563+
and not ckpt_cfg.non_persistent_global_ckpt_dir
3564+
and load_dir is not None
3565+
and non_persistent_global_dir == os.path.join(load_dir, _NON_PERSISTENT_CKPT_SUBDIR)
3566+
):
3567+
dataloader_checkpoint_root = load_dir
3568+
_record_dataloader_state_dir(checkpointing_context, dataloader_checkpoint_root)
35463569
return _load_non_persistent_base_checkpoint(
35473570
non_persistent_global_dir,
35483571
ckpt_cfg,

‎tests/unit_tests/training/test_checkpointing.py‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
_extract_megatron_lm_args_from_state_dict,
3636
_get_checkpoint_format,
3737
_get_non_persistent_iteration,
38+
_has_global_non_persistent_checkpoint,
3839
_load_base_checkpoint,
3940
_load_checkpoint_from_path,
4041
_load_hf_pretrained_checkpoint,
@@ -1931,6 +1932,116 @@ def mock_pg_collection(self):
19311932
mock_pg.tp.size.return_value = 1
19321933
return mock_pg
19331934

1935+
def test_global_non_persistent_checkpoint_is_found_under_distinct_save_dir(self, tmp_path):
1936+
"""An unchanged restart discovers the non-persistent checkpoint written under save."""
1937+
load_dir = tmp_path / "input"
1938+
save_dir = tmp_path / "output"
1939+
non_persistent_dir = save_dir / "non_persistent"
1940+
non_persistent_dir.mkdir(parents=True)
1941+
(non_persistent_dir / "latest_train_state.pt").touch()
1942+
1943+
ckpt_cfg = Mock(spec=CheckpointConfig)
1944+
ckpt_cfg.save = str(save_dir)
1945+
ckpt_cfg.non_persistent_ckpt_type = "global"
1946+
ckpt_cfg.non_persistent_global_ckpt_dir = None
1947+
1948+
assert _has_global_non_persistent_checkpoint(str(load_dir), ckpt_cfg)
1949+
1950+
@patch("megatron.bridge.training.checkpointing._load_global_dist_base_checkpoint")
1951+
@patch("megatron.bridge.training.checkpointing._get_checkpoint_format", return_value="torch_dist")
1952+
@patch("megatron.bridge.training.checkpointing._load_non_persistent_base_checkpoint")
1953+
@patch("megatron.bridge.training.checkpointing._resolve_checkpoint_iteration", return_value=(100, False))
1954+
def test_newer_global_non_persistent_checkpoint_is_loaded_from_distinct_save_dir(
1955+
self,
1956+
_mock_resolve,
1957+
mock_load_non_persistent,
1958+
_mock_get_format,
1959+
mock_load_persistent,
1960+
tmp_path,
1961+
base_config,
1962+
mock_pg_collection,
1963+
):
1964+
"""A restart prefers the newer recovery point written under checkpoint.save."""
1965+
load_dir = tmp_path / "input"
1966+
save_dir = tmp_path / "output"
1967+
non_persistent_dir = save_dir / "non_persistent"
1968+
non_persistent_dir.mkdir(parents=True)
1969+
torch.save(
1970+
TrainState(step=200).state_dict(),
1971+
get_checkpoint_train_state_filename(str(non_persistent_dir), prefix="latest"),
1972+
)
1973+
1974+
base_config.save = str(save_dir)
1975+
base_config.ckpt_format = "torch_dist"
1976+
base_config.non_persistent_ckpt_type = "global"
1977+
base_config.non_persistent_global_ckpt_dir = None
1978+
expected = ({"model": "newer"}, str(non_persistent_dir / "iter_0000200"), False, CheckpointType.GLOBAL)
1979+
mock_load_non_persistent.return_value = expected
1980+
1981+
checkpointing_context = {}
1982+
result = _load_base_checkpoint(
1983+
str(load_dir),
1984+
base_config,
1985+
checkpointing_context=checkpointing_context,
1986+
pg_collection=mock_pg_collection,
1987+
)
1988+
1989+
assert result == expected
1990+
assert mock_load_non_persistent.call_args.args[0] == str(non_persistent_dir)
1991+
assert mock_load_non_persistent.call_args.args[4] == 200
1992+
assert checkpointing_context["dataloader_state_dir"] == str(save_dir / DATALOADER_STATE_SUBDIR)
1993+
mock_load_persistent.assert_not_called()
1994+
1995+
@patch("megatron.bridge.training.checkpointing._load_non_persistent_base_checkpoint")
1996+
@patch("megatron.bridge.training.checkpointing._resolve_checkpoint_iteration", return_value=(50, False))
1997+
def test_newer_global_non_persistent_checkpoint_under_load_remains_available(
1998+
self,
1999+
_mock_resolve,
2000+
mock_load_non_persistent,
2001+
tmp_path,
2002+
base_config,
2003+
mock_pg_collection,
2004+
):
2005+
"""A distinct empty save root does not hide a newer recovery point under load."""
2006+
load_dir = tmp_path / "input"
2007+
save_dir = tmp_path / "output"
2008+
load_non_persistent_dir = load_dir / "non_persistent"
2009+
save_non_persistent_dir = save_dir / "non_persistent"
2010+
load_non_persistent_dir.mkdir(parents=True)
2011+
save_non_persistent_dir.mkdir(parents=True)
2012+
torch.save(
2013+
TrainState(step=200).state_dict(),
2014+
get_checkpoint_train_state_filename(str(load_non_persistent_dir), prefix="latest"),
2015+
)
2016+
torch.save(
2017+
TrainState(step=100).state_dict(),
2018+
get_checkpoint_train_state_filename(str(save_non_persistent_dir), prefix="latest"),
2019+
)
2020+
2021+
base_config.save = str(save_dir)
2022+
base_config.non_persistent_ckpt_type = "global"
2023+
base_config.non_persistent_global_ckpt_dir = None
2024+
expected = (
2025+
{"model": "newer"},
2026+
str(load_non_persistent_dir / "iter_0000200"),
2027+
False,
2028+
CheckpointType.GLOBAL,
2029+
)
2030+
mock_load_non_persistent.return_value = expected
2031+
2032+
checkpointing_context = {}
2033+
result = _load_base_checkpoint(
2034+
str(load_dir),
2035+
base_config,
2036+
checkpointing_context=checkpointing_context,
2037+
pg_collection=mock_pg_collection,
2038+
)
2039+
2040+
assert result == expected
2041+
assert mock_load_non_persistent.call_args.args[0] == str(load_non_persistent_dir)
2042+
assert mock_load_non_persistent.call_args.args[4] == 200
2043+
assert checkpointing_context["dataloader_state_dir"] == str(load_dir / DATALOADER_STATE_SUBDIR)
2044+
19342045
@patch("megatron.bridge.training.checkpointing._get_non_persistent_iteration")
19352046
@patch("megatron.bridge.training.checkpointing.file_exists")
19362047
def test_load_base_checkpoint_no_checkpoint(

0 commit comments

Comments
 (0)