Route native loads through torch_checkpointing

[ivy-zhou]

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.

---

Stack created with Sapling. Best reviewed with ReviewStack.

• [checkpoint] Log native DCP async save durations #4279
• [checkpoint] Centralize checkpoint trace spans #4278
• [checkpoint] Load Hugging Face checkpoints natively #4277
• -> Route native loads through torch_checkpointing #4191
• Consolidate torch_checkpointing HF final saves #4190
• Bridge checkpoint metrics into structured logging #4189
• Route saves through the torch_checkpointing backend #4188
• Teach retention about staged and abandoned checkpoints #4197
• Move checkpoint discovery and retention into BaseCheckpointManager #4187

[meetv18]

is this logic commonly shared with dcpv1?