From 98f0c3932bb653142efbcd183421b6413194db59 Mon Sep 17 00:00:00 2001 From: Ivy Zhou Date: Sat, 22 Aug 2026 03:01:32 -0700 Subject: [PATCH] Route native loads through torch_checkpointing Summary: `TorchCheckpointingManager` could save but not load: `_load` raised `NotImplementedError`, so a run configured for this backend could write checkpoints and never resume from them. Implement the native load path, completing the round trip. Resolution mirrors the DCP manager. With no explicit step, the latest valid step in the checkpoint folder is used; with an explicit step, a missing folder or missing checkpoint is an error rather than a silent fresh start. When no step is found, `initial_load_path` is used if configured, and otherwise the run starts fresh. Step 0 is treated as a seed checkpoint and loads model state only. Step discovery only accepts directories that actually contain the backend's `metadata.pkl`. A crashed save leaves a `step-N` directory with no metadata, and without this check that partial directory would win the `max()` and be loaded as the resume point. Directory listing, existence checks, and the load itself all go through the configured `Storage`, so remote checkpoint folders work the same as local ones. The backend returns a plain dict rather than writing through the live objects, so `_restore_state_dict` puts values back: `Stateful` entries get `load_state_dict`, plain dicts are updated in place so callers holding a reference see the result, and anything else that is not already the same object raises rather than being silently skipped. Keys in `exclude_from_loading` are dropped before the load, and an excluded key that is not in the state dict is an error, matching the DCP manager. Hugging Face loads are explicitly rejected for now; they land separately. Ported from an internal change. Besides the path and naming translation, the internal version defines `load()` with its own `if not self.enable` guard; here the body moves to the `_load` hook and the guard is dropped, since `BaseCheckpointManager` owns it. Test Plan: `pytest tests/unit_tests/test_torch_checkpointing.py`: 24 passed. Four new tests: an explicit-step load restores both model weights and optimizer state from the backend payload and passes the full state dict as the load target; a latest-step load of `step-0` requests model state only; discovery and load go through a configured non-local `Storage` rather than the filesystem; and a folder containing valid `step-2` and `step-5`, a metadata-less `step-8`, and a `tmp_step-9` resolves to `step-5`, so an interrupted save is not mistaken for the newest checkpoint. `pytest tests/unit_tests/test_checkpoint.py tests/unit_tests/observability/ torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py` alongside the above: 147 passed, 2 subtests passed. Also verified: - `TorchCheckpointingManager.__abstractmethods__` is empty; with `_load` implemented the manager now satisfies the full `BaseCheckpointManager` contract. - `METADATA_FILE_NAME` and `CheckpointManager.load(checkpoint_id, into=...)` exist in `torch_checkpointing` 0.1.0. - `ufmt` and `flake8 --config=.flake8` clean on both changed files. --- tests/unit_tests/test_torch_checkpointing.py | 184 +++++++++++++++++- .../checkpointer/torch_checkpointing.py | 111 ++++++++++- 2 files changed, 287 insertions(+), 8 deletions(-) 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 + }