diff --git a/tests/unit_tests/test_torch_checkpointing.py b/tests/unit_tests/test_torch_checkpointing.py index dc46cd9024..835aeda5dc 100644 --- a/tests/unit_tests/test_torch_checkpointing.py +++ b/tests/unit_tests/test_torch_checkpointing.py @@ -53,6 +53,8 @@ class _BackendManager: def __init__(self) -> None: self.closed = False self.lock_calls = 0 + self.load_calls = [] + self.load_result = None self.prewarm_calls = [] self.save_calls = [] self.save_result = Future() @@ -68,6 +70,10 @@ def lock(self): self.lock_calls += 1 return nullcontext() + def load(self, checkpoint_id, into=None, **kwargs): + self.load_calls.append((checkpoint_id, into, kwargs)) + return self.load_result if self.load_result is not None else into + def close(self) -> None: self.closed = True @@ -100,6 +106,10 @@ def _build_manager( *, backend_config: BackendCheckpointManager.Config | None = None, storage_config=None, + base_folder: str = "/tmp", + model_parts=None, + optimizers=None, + states=None, ) -> tuple[TorchCheckpointingManager, _BackendManager]: if backend_config is None: backend_config = _default_backend_config() @@ -118,12 +128,12 @@ def _build_manager( ): manager = config.build( dataloader=None, - model_parts=[nn.Linear(2, 2)], - optimizers=_Stateful("optimizer"), + model_parts=model_parts or [nn.Linear(2, 2)], + optimizers=optimizers or _Stateful("optimizer"), lr_schedulers=_Stateful("scheduler"), - states={"train_state": _Stateful("train")}, + states=states or {"train_state": _Stateful("train")}, sd_adapter=None, - base_folder="/tmp", + base_folder=base_folder, storage_config=storage_config, ) return manager, backend_manager @@ -808,3 +818,169 @@ def test_hf_final_save_converts_and_consolidates_before_commit( storage_config=storage_config, ) manager.close() + + def test_native_load_restores_model_and_optimizer(self) -> None: + with tempfile.TemporaryDirectory() as base_folder: + checkpoint_id = os.path.join(base_folder, "checkpoint", "step-5") + os.makedirs(checkpoint_id) + with open(os.path.join(checkpoint_id, "metadata.pkl"), "wb"): + pass + model = nn.Linear(2, 2, bias=False) + optimizer = _Stateful("optimizer") + config = TorchCheckpointingManager.Config( + enable=True, + folder="checkpoint", + keep_latest_k=0, + initial_load_model_only=False, + load_only=True, + ) + manager, backend_manager = self._build_manager( + config, + base_folder=base_folder, + model_parts=[model], + optimizers=optimizer, + ) + expected_weight = torch.full_like(model.weight, 3) + backend_manager.load_result = { + "weight": expected_weight, + OPTIMIZER: {"value": "restored"}, + } + + self.assertTrue(manager.load(step=5)) + + torch.testing.assert_close(model.weight, expected_weight) + self.assertEqual("restored", optimizer.value) + self.assertEqual(checkpoint_id, backend_manager.load_calls[0][0]) + self.assertEqual(set(manager.states), set(backend_manager.load_calls[0][1])) + self.assertEqual({"strict": True}, backend_manager.load_calls[0][2]) + manager.close() + + def test_native_load_requires_every_requested_key(self) -> None: + # The backend skips absent keys by default, which would leave those + # parameters at their initialized values and quietly resume from a model + # that is not the one that was saved. + with tempfile.TemporaryDirectory() as base_folder: + checkpoint_id = os.path.join(base_folder, "checkpoint", "step-5") + os.makedirs(checkpoint_id) + with open(os.path.join(checkpoint_id, "metadata.pkl"), "wb"): + pass + config = TorchCheckpointingManager.Config( + enable=True, + folder="checkpoint", + keep_latest_k=0, + initial_load_model_only=False, + load_only=True, + ) + manager, backend_manager = self._build_manager( + config, + base_folder=base_folder, + ) + backend_manager.load = mock.Mock( + side_effect=RuntimeError( + f"Checkpoint at {checkpoint_id} is missing keys: " + f"['{MODEL}::weight']" + ) + ) + + with self.assertRaisesRegex(RuntimeError, "is missing keys"): + manager.load(step=5) + + self.assertIs(True, backend_manager.load.call_args.kwargs["strict"]) + manager.close() + + def test_load_latest_step_zero_loads_only_model_state(self) -> None: + with tempfile.TemporaryDirectory() as base_folder: + checkpoint_id = os.path.join(base_folder, "checkpoint", "step-0") + os.makedirs(checkpoint_id) + with open(os.path.join(checkpoint_id, "metadata.pkl"), "wb"): + pass + config = TorchCheckpointingManager.Config( + enable=True, + folder="checkpoint", + keep_latest_k=0, + initial_load_model_only=False, + load_only=True, + ) + manager, backend_manager = self._build_manager( + config, + base_folder=base_folder, + ) + + self.assertTrue(manager.load()) + + self.assertEqual({MODEL}, set(backend_manager.load_calls[0][1])) + manager.close() + + def test_load_latest_uses_configured_storage(self) -> None: + # The backend Storage the adapter wraps: step-7 is a directory holding a + # metadata.pkl file. + storage = mock.Mock() + storage.ls.return_value = ["step-7"] + storage.exists.return_value = True + storage.isdir.side_effect = lambda path: not str(path).endswith("metadata.pkl") + storage_config = mock.Mock() + storage_config.create_storage.return_value = storage + config = TorchCheckpointingManager.Config( + enable=True, + folder="checkpoint", + keep_latest_k=0, + initial_load_model_only=False, + load_only=True, + ) + backend_config = _default_backend_config() + backend_config.storage_config = storage_config + manager, backend_manager = self._build_manager( + config, + backend_config=backend_config, + base_folder="/custom", + ) + + self.assertTrue(manager.load()) + + self.assertEqual( + "/custom/checkpoint/step-7", + backend_manager.load_calls[0][0], + ) + manager.close() + + def test_load_latest_ignores_incomplete_checkpoint_directories(self) -> None: + with tempfile.TemporaryDirectory() as base_folder: + checkpoint_folder = os.path.join(base_folder, "checkpoint") + for step in (2, 5): + checkpoint_id = os.path.join(checkpoint_folder, f"step-{step}") + os.makedirs(checkpoint_id) + with open(os.path.join(checkpoint_id, "metadata.pkl"), "wb"): + pass + incomplete_checkpoint_id = os.path.join(checkpoint_folder, "step-8") + os.makedirs(incomplete_checkpoint_id) + # An interrupted save leaves its metadata behind, so the temporary + # directory looks complete. Resuming from it would build the id + # "step-9", which does not exist. + # Neither is a checkpoint id we could reconstruct: "tmp_step-9" is + # an interrupted save whose metadata already landed, so it looks + # complete, and "step-99.partial" is not a name we write at all. + # Matching either one loosely resolves to a step that does not exist. + for decoy in ("tmp_step-9", "step-99.partial"): + decoy_id = os.path.join(checkpoint_folder, decoy) + os.makedirs(decoy_id) + with open(os.path.join(decoy_id, "metadata.pkl"), "wb"): + pass + config = TorchCheckpointingManager.Config( + enable=True, + folder="checkpoint", + keep_latest_k=0, + initial_load_model_only=False, + load_only=True, + ) + manager, backend_manager = self._build_manager( + config, + base_folder=base_folder, + ) + + self.assertTrue(manager.load()) + + self.assertEqual( + os.path.join(checkpoint_folder, "step-5"), + backend_manager.load_calls[0][0], + ) + manager.close() diff --git a/torchtitan/components/checkpointer/torch_checkpointing.py b/torchtitan/components/checkpointer/torch_checkpointing.py index 2df2065767..07ab30779f 100644 --- a/torchtitan/components/checkpointer/torch_checkpointing.py +++ b/torchtitan/components/checkpointer/torch_checkpointing.py @@ -12,7 +12,8 @@ import queue import re import threading -from collections.abc import Callable +import time +from collections.abc import Callable, Mapping from concurrent.futures import Future from dataclasses import dataclass, replace from pathlib import Path @@ -22,6 +23,7 @@ import torch.distributed as dist import torch.nn as nn from torch.distributed.checkpoint.state_dict_saver import _stateful_to_state_dict +from torch.distributed.checkpoint.stateful import Stateful from torch_checkpointing.barriers import TCPStoreBarrierConfig from torch_checkpointing.checkpoint_layout import LayoutInfo, SafetensorsSerialization from torch_checkpointing.checkpoint_manager import ( @@ -177,6 +179,36 @@ def _item_specs() -> dict[str, ItemSpec]: } +def _restore_state_dict( + states: Mapping[str, Any], + state_dict: Mapping[str, Any], +) -> None: + """Write loaded values back into the live state objects. + + The backend returns a plain dict, so ``Stateful`` entries need an explicit + ``load_state_dict``; plain dicts are updated in place so callers holding a + reference observe the loaded values. + """ + missing = object() + model_state_dict = state_dict.get(MODEL, state_dict) + + for key, target in states.items(): + value = model_state_dict if key == MODEL else state_dict.get(key, missing) + if value is missing: + continue + if isinstance(target, Stateful): + target.load_state_dict(value) + elif isinstance(target, dict) and isinstance(value, Mapping): + if target is not value: + target.clear() + target.update(value) + elif target is not value: + raise TypeError( + f"Cannot restore non-Stateful checkpoint state {key!r} of type " + f"{type(target).__name__}" + ) + + def _default_backend_config() -> BackendCheckpointManager.Config: barrier_timeout_sec = _DEFAULT_BARRIER_TIMEOUT_SEC save_config = AsyncCheckpointSaverConfig( @@ -345,11 +377,69 @@ def __init__( def __del__(self) -> None: self.close() - # Load routing lands in a later change. + @sl.log_trace_span("checkpoint_load") + @torch.no_grad() def _load(self, step: int = -1) -> bool: - raise NotImplementedError( - "TorchCheckpointingManager does not implement load() yet." + has_checkpoint_folder = self._storage.isdir(self.folder) + load_step = -1 + if has_checkpoint_folder: + load_step = self._find_load_step() if step == -1 else step + if step != -1 and not has_checkpoint_folder: + raise FileNotFoundError( + f"--checkpoint.load_step={step} not found because " + f"checkpoint.folder {self.folder} does not exist" + ) + + if load_step == -1: + if self.initial_load_in_hf: + raise ValueError( + "TorchCheckpointingManager does not yet support loading " + "Hugging Face checkpoints." + ) + if not self.initial_load_path: + logger.info("No checkpoint was provided, this is a fresh start.") + return False + checkpoint_id = self.initial_load_path + model_only = self.initial_load_model_only + if not self._storage.isdir(checkpoint_id): + raise ValueError( + f"Checkpoint.initial_load_path is invalid: {checkpoint_id}" + ) + else: + step = load_step + # Step 0 is a seed checkpoint, which holds model state only. + model_only = step == 0 + checkpoint_id = self._create_checkpoint_id(step) + if not self._storage.isdir(checkpoint_id): + raise FileNotFoundError( + f"--checkpoint.load_step={step} not found at {checkpoint_id}" + ) + + if not self._is_valid_checkpoint(checkpoint_id): + raise ValueError( + f"Checkpoint {checkpoint_id!r} is not a native " + "torch_checkpointing checkpoint." + ) + logger.info("Loading the checkpoint from %s.", checkpoint_id) + begin = time.monotonic() + states = self._states_to_load(model_only) + # strict: the backend defaults to skipping anything the checkpoint does + # not carry, which would silently leave parameters at their initialized + # values and resume from a model that is not the one that was saved. + # exclude_from_loading is applied by _states_to_load, so anything still + # in `states` here is genuinely required. + loaded = self._manager.load( + checkpoint_id, + into=_stateful_to_state_dict(states), + strict=True, + ) + _restore_state_dict(states, loaded) + GarbageCollection.collect("GC collection for checkpoint loading.") + logger.info( + "Finished loading the checkpoint in %.2f seconds.", + time.monotonic() - begin, ) + return True @sl.log_trace_span("checkpoint_save") @torch.no_grad() @@ -520,3 +610,16 @@ def _save_last_step(self, curr_step: int) -> None: def _should_prewarm(self) -> bool: return self.enable and not self._prewarmed and not self.load_only + + def _states_to_load(self, model_only: bool) -> dict[str, Any]: + if model_only: + return {MODEL: self.states[MODEL]} + + for exclude_key in self.exclude_from_loading: + if exclude_key not in self.states: + raise ValueError(f"{exclude_key} not found in state_dict.") + return { + key: value + for key, value in self.states.items() + if key not in self.exclude_from_loading + }