Route saves through the torch_checkpointing backend by ivy-zhou · Pull Request #4188 · pytorch/torchtitan

ivy-zhou · 2026-08-18T06:45:42Z

Summary:
TorchCheckpointingManager has been config-only since #4058: it builds its
backend and then raises NotImplementedError from every operation. This
implements the save path against that backend, leaving load for a later change.

Saves follow the same cadence contract as the DCP manager -- interval, first
step, and last step -- and hand the flattened state dict to the backend, which
returns a future tracked in save_future and awaited by maybe_wait_for_saving.
On a step where no save is due, the manager prewarms the backend's staging
buffers once, so the first real save does not also pay for pinned-memory
allocation.

The last step is special-cased. An async saver can outlive the process, so the
final write retires the async manager and issues a synchronous save through a
fresh one built by _with_sync_save. Load-only runs use the same synchronous
config with the barrier stripped, because no save is expected and a barrier
would block against ranks that never reach it. Conversely, a multi-rank run
whose writer has no barrier is rejected up front rather than deadlocking or
producing a torn checkpoint.

Backend storage is a defaulted __init__ parameter rather than a Config
field. Configurable.Config is Tyro-parsed, and a backend storage object is
not a command-line surface; callers needing remote storage pass it
programmatically. An explicit value is also pushed into the backend config, so
saves and loads use it rather than only this manager's own path probes.

Retention reuses the shared purge_thread worker from checkpointer/base.py,
draining through the backend Storage abstraction rather than the filesystem
directly, and matches the same step-(\d+) names the DCP manager purges.

last_save_in_hf is rejected in Config.__post_init__; HF export lands in a
later change and would otherwise fail deep in the backend.

This is a port of an internal change onto the current OSS layout, which moved
underneath it: components/torch_checkpointing_manager.py is now
components/checkpointer/torch_checkpointing.py, CheckpointManagerConfig is
BaseCheckpointManager.Config, purge_worker is purge_thread, and
LRSchedulersContainer comes from components.optimizer. The internal version
also carried its own if not self.enable guard in every public method; those
are gone, since the base class now owns that check and dispatches to the
_save / _wait_for_saving / _maybe_wait_for_staging / _close hooks.

_should_save and _create_checkpoint_id would otherwise be reimplemented
here, identically to the DCP manager's copies -- both depend only on config
fields BaseCheckpointManager.Config already declares, not on how a backend
reads or writes bytes. They move to the base instead, and the DCP manager's
copies are deleted; the base keeps the DCP signature for
_create_checkpoint_id, including its optional folder argument. Retention,
step discovery, state selection, and the last-step payload stay per-manager,
because each reaches storage differently and unifying them means introducing a
storage abstraction against a converged code path.

Test Plan:
pytest tests/unit_tests/test_torch_checkpointing.py: 17 passed (6 existing
plus 11 new).

The new tests cover save cadence and future tracking, prewarm running exactly
once before the first scheduled save, load-only selecting a synchronous
barrier-free backend, load-only never constructing a barrier at all, the
multi-rank no-barrier rejection, staging waits going through the backend lock,
save waits using the configured timeout, close() still draining the purge
thread and closing the backend when the save future raises, purge name matching
against step-N / tmp_step-N / step-N.partial, purge not depending on
checkpoint metadata, and the last step building a synchronous manager with a
model-only payload.

pytest tests/unit_tests/test_checkpoint.py tests/unit_tests/test_state_dict_keys.py torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py alongside the
above: 61 passed, 4 subtests passed.

Also verified:

Every backend API this depends on exists in torch_checkpointing 0.1.0:
SyncCheckpointSaverConfig, LocalFileSystemStorageConfig, Storage,
Config.with_sync_save, prewarm_staging, and lock.
TorchCheckpointingManager.__abstractmethods__ is empty, so the base
contract is fully implemented.
ufmt and flake8 --config=.flake8 clean on both changed files.

Stack created with Sapling. Best reviewed with ReviewStack.

meetv18 · 2026-08-19T17:38:19Z

+    for load-only runs, where no save is expected and the barrier would block
+    against ranks that never save.
+    """
+    writer_config = copy.deepcopy(config.save.writer_config)


why do we need this?

errrr I think because we might replace just one config on it, for configuring the sync save to not use a barrier?

meetv18 · 2026-08-19T17:43:04Z

+            manager.save(
+                self._create_checkpoint_id(curr_step),
+                _stateful_to_state_dict(states),
+            )


curious: why don't we just use the async manager that we already build and just do wait_for_upload()?

Why construct a sync specific manager?

This is the same flow as in [DCPv1] (https://github.com/pytorch/torchtitan/blob/main/torchtitan/components/checkpointer/dcp.py?fbclid=IwY2xjawTy5tNwZG9mA2V4dG4DYWVtAjExAGJyaWQRMWNNMlFuU3ZFZHpMWE1xd2lzcnRjBmFwcF9pZAEwAAEewUA7ZrzF3VhfEg8vxClY6Y8eBOuTfiaQWVtLdqfrjclJiTaKAwOUq7Kjp4Y_aem_XXVru1bj-SIS9gb7Ty3Uuw#L792) where the last step saves synchronously.

Agree it would be cleaner with just 1 saver, but this specific last save:

change dtypes

change FQNs

calls consolidation, which starts a new PG to have a varrier

So it ends up we are just using the same sync save flow here. I think both these things can be fixed (e.g. consolidation in the barrier should probably just use the existing barrier constructed in TC async save) but it requires some refactoring, and the metadata implications of the dtype and the FQN changes are difficult to reason through. I think we spoke about this before, and this was the compromise we came to, happy to revisit it in the future though!

…4183) Pyrefly reports three errors in this file, all of them real: _save was declared "-> None" while the base declares "-> bool", and it discarded super()._save()'s result. BaseCheckpointManager.save returns whatever _save returns, so save() handed back None for this manager. Nothing consumes it today -- torchft/trainer.py ignores the result -- but the contract was broken and the next caller to check it would have been surprised. A replica that skips the full save now reports False; the per-replica dataloader checkpoint is a side channel, not the checkpoint this value describes. _wait_for_saving dereferenced save_future without narrowing it. The base's maybe_wait_for_saving guarantees it is set before dispatching here, which the comment already said, so this just asserts what the comment claims. _ft_save assigned dcp_save's "Future | AsyncSaveResponse | None" straight into save_future, typed "Future | None". AsyncMode.ASYNC always yields a plain Future, so assert that, matching how the DCP manager narrows the same call in its own ASYNC branch. Only the first of the three is new: it arrived with the disabled-guard refactor (#4173), which renamed save to _save and made the base's return type load bearing. The other two predate it. Test Plan: python3 -m pyrefly check torchtitan/experiments/torchft/checkpoint.py -> 0 errors (was 3) python3 -m pytest torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py -q -> 2 passed Adds a test covering both branches of the participating_rank guard, so the return value is pinned rather than left to the type checker. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/pytorch/torchtitan/pull/4183). * #4191 * #4190 * #4189 * #4188 * #4197 * #4187 * #4186 * #4185 * #4184 * __->__ #4183

Summary: `components/checkpoint.py` is the last of the three re-export shims left behind when the checkpointer was grouped into a package, after the lr_scheduler shim in #4172 and checkpoint_utils in the preceding change. With it gone there are no compatibility shims left under `torchtitan/components/`. This one is a pure module rename at the callsite. `checkpointer/__init__.py` already re-exports exactly the same eight symbols the shim forwarded -- `AsyncMode`, `CheckpointManager`, `ModelWrapper`, and the `MODEL` / `OPTIMIZER` / `LR_SCHEDULER` / `DATALOADER` / `TRAIN_STATE` key constants -- so every importer changes only the module it names, with the imported names and their grouping untouched. Verified by parsing each importer and checking every imported symbol against the package's `__all__` before touching anything; nothing referenced a symbol the package does not expose, and no callsite used the plain `import torchtitan.components.checkpoint` form. Thirty modules are updated, spanning the checkpoint-conversion scripts, the forge, torchft, graph_trainer and rl experiments, and the unit tests. A thirty-first file, `experiments/rl/__init__.py`, carries the import inside its module docstring as a usage example; that is updated too, so the documented path matches the working one. This is an import-path change only; no runtime behavior changes. Test Plan: Full `pytest tests/unit_tests` (excluding `test_rope.py`, which cannot be collected without the optional `fla` package): 622 passed, 18 failed. The 18 are the same set that fails on unmodified main in this environment -- missing optional dependencies (`transformers`, `fla`) and environment-specific kernel/compile failures (helion rope, inductor lora). Also verified: - No `components.checkpoint` references remain anywhere in the repo, across all file types, and no `Compatibility imports` shim remains under `torchtitan/components/`. - `import torchtitan.components.checkpoint` now raises `ModuleNotFoundError`. - All eight symbols import cleanly from `torchtitan.components.checkpointer`. - All 31 changed files byte-compile, and the affected non-test modules (`trainer`, `forge.engine`, `torchft.checkpoint`, `torchft.optimizer`, both checkpoint-conversion scripts) import cleanly. `experiments.rl` fails only on the absent optional `vllm` package, unrelated to this change. - `ufmt` and `flake8 --config=.flake8` clean on all 31 files. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/pytorch/torchtitan/pull/4184). * #4240 * #4230 * #4191 * #4190 * #4189 * #4188 * #4197 * #4187 * #4186 * #4185 * __->__ #4184

…4185) Summary: `components/checkpoint_utils.py` became a re-export shim when the optimizer and checkpointer components were grouped into packages. Unlike the lr_scheduler shim removed in #4172, this one forwarded to two different destinations at once, which is what made it worth deleting rather than keeping: reading an import of `checkpoint_utils` told you nothing about whether the symbol was optimizer plumbing or checkpointer plumbing. Route each of the four importers to the module that actually defines the symbol. `canonical_fqn` lives in `checkpointer/utils.py`; `init_optim_state`, `get_flat_optim_state_dict`, and `load_flat_optim_state_dict` live in `optimizer/utils.py`. This also settles the naming objection fegin raised on #4140, that `canonical_fqn` does not belong under an optimizer-shaped name -- its importer in the rl trainer now names the checkpointer package directly. The three state-dict helpers are imported from `optimizer.utils` rather than re-exported through `optimizer/__init__.py`. They are low-level DCP plumbing with two callers between them, not part of the package's public surface, which stays `OptimizersContainer`, `LRSchedulersContainer`, `ParamGroupConfig`, and `default_adamw`. This is an import-path change only; no runtime behavior changes. Test Plan: `pytest tests/unit_tests/test_state_dict_keys.py tests/unit_tests/test_checkpoint.py tests/unit_tests/test_lr_scheduler.py tests/unit_tests/test_optimizer_param_groups.py tests/unit_tests/test_torch_checkpointing.py`: 77 passed, 4 subtests passed. `test_legacy_checkpoint_utils_imports`, which asserted the shim's symbols were identical to the submodule's, is dropped -- it cannot outlive the shim. `test_legacy_checkpoint_utils_can_be_imported_first` is kept but retargeted, as `test_state_dict_helpers_can_be_imported_first`. It guards a real property rather than the shim: the `optimizer` and `checkpointer` package `__init__` files import from each other, so importing either leaf `utils` module first in a fresh interpreter must not close an import cycle. It now subtests both leaf modules instead of the single shim. Also verified: - No `checkpoint_utils` references remain anywhere in the repo, across all file types, not only Python. - `import torchtitan.components.checkpoint_utils` now raises `ModuleNotFoundError`. - `ufmt` and `flake8 --config=.flake8` clean on the four changed files. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/pytorch/torchtitan/pull/4185). * #4240 * #4230 * #4191 * #4190 * #4189 * #4188 * #4197 * #4187 * #4186 * __->__ #4185 * #4184

…4183) Pyrefly reports three errors in this file, all of them real: _save was declared "-> None" while the base declares "-> bool", and it discarded super()._save()'s result. BaseCheckpointManager.save returns whatever _save returns, so save() handed back None for this manager. Nothing consumes it today -- torchft/trainer.py ignores the result -- but the contract was broken and the next caller to check it would have been surprised. A replica that skips the full save now reports False; the per-replica dataloader checkpoint is a side channel, not the checkpoint this value describes. _wait_for_saving dereferenced save_future without narrowing it. The base's maybe_wait_for_saving guarantees it is set before dispatching here, which the comment already said, so this just asserts what the comment claims. _ft_save assigned dcp_save's "Future | AsyncSaveResponse | None" straight into save_future, typed "Future | None". AsyncMode.ASYNC always yields a plain Future, so assert that, matching how the DCP manager narrows the same call in its own ASYNC branch. Only the first of the three is new: it arrived with the disabled-guard refactor (#4173), which renamed save to _save and made the base's return type load bearing. The other two predate it. Test Plan: python3 -m pyrefly check torchtitan/experiments/torchft/checkpoint.py -> 0 errors (was 3) python3 -m pytest torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py -q -> 2 passed Adds a test covering both branches of the participating_rank guard, so the return value is pinned rather than left to the type checker. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/pytorch/torchtitan/pull/4183). * #4191 * #4190 * #4189 * #4188 * #4197 * #4187 * #4186 * #4185 * #4184 * __->__ #4183

Summary: `TorchCheckpointingManager` has been config-only since #4058: it builds its backend and then raises `NotImplementedError` from every operation. This implements the save path against that backend, leaving load for a later change. Saves follow the same cadence contract as the DCP manager -- interval, first step, and last step -- and hand the flattened state dict to the backend, which returns a future tracked in `save_future` and awaited by `maybe_wait_for_saving`. On a step where no save is due, the manager prewarms the backend's staging buffers once, so the first real save does not also pay for pinned-memory allocation. The last step is special-cased. An async saver can outlive the process, so the final write retires the async manager and issues a synchronous save through a fresh one built by `_with_sync_save`. Load-only runs use the same synchronous config with the barrier stripped, because no save is expected and a barrier would block against ranks that never reach it. Conversely, a multi-rank run whose writer has no barrier is rejected up front rather than deadlocking or producing a torn checkpoint. Backend storage is a defaulted `__init__` parameter rather than a `Config` field. `Configurable.Config` is Tyro-parsed, and a backend storage object is not a command-line surface; callers needing remote storage pass it programmatically. An explicit value is also pushed into the backend config, so saves and loads use it rather than only this manager's own path probes. Retention reuses the shared `purge_thread` worker from `checkpointer/base.py`, draining through the backend `Storage` abstraction rather than the filesystem directly, and matches the same `step-(\d+)` names the DCP manager purges. `last_save_in_hf` is rejected in `Config.__post_init__`; HF export lands in a later change and would otherwise fail deep in the backend. This is a port of an internal change onto the current OSS layout, which moved underneath it: `components/torch_checkpointing_manager.py` is now `components/checkpointer/torch_checkpointing.py`, `CheckpointManagerConfig` is `BaseCheckpointManager.Config`, `purge_worker` is `purge_thread`, and `LRSchedulersContainer` comes from `components.optimizer`. The internal version also carried its own `if not self.enable` guard in every public method; those are gone, since the base class now owns that check and dispatches to the `_save` / `_wait_for_saving` / `_maybe_wait_for_staging` / `_close` hooks. `_should_save` and `_create_checkpoint_id` would otherwise be reimplemented here, identically to the DCP manager's copies -- both depend only on config fields `BaseCheckpointManager.Config` already declares, not on how a backend reads or writes bytes. They move to the base instead, and the DCP manager's copies are deleted; the base keeps the DCP signature for `_create_checkpoint_id`, including its optional `folder` argument. Retention, step discovery, state selection, and the last-step payload stay per-manager, because each reaches storage differently and unifying them means introducing a storage abstraction against a converged code path. Test Plan: `pytest tests/unit_tests/test_torch_checkpointing.py`: 17 passed (6 existing plus 11 new). The new tests cover save cadence and future tracking, prewarm running exactly once before the first scheduled save, load-only selecting a synchronous barrier-free backend, load-only never constructing a barrier at all, the multi-rank no-barrier rejection, staging waits going through the backend lock, save waits using the configured timeout, `close()` still draining the purge thread and closing the backend when the save future raises, purge name matching against `step-N` / `tmp_step-N` / `step-N.partial`, purge not depending on checkpoint metadata, and the last step building a synchronous manager with a model-only payload. `pytest tests/unit_tests/test_checkpoint.py tests/unit_tests/test_state_dict_keys.py torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py` alongside the above: 61 passed, 4 subtests passed. Also verified: - Every backend API this depends on exists in `torch_checkpointing` 0.1.0: `SyncCheckpointSaverConfig`, `LocalFileSystemStorageConfig`, `Storage`, `Config.with_sync_save`, `prewarm_staging`, and `lock`. - `TorchCheckpointingManager.__abstractmethods__` is empty, so the base contract is fully implemented. - `ufmt` and `flake8 --config=.flake8` clean on both changed files.

pytorch-bot Bot added the ciflow/8gpu label Aug 18, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 18, 2026

ivy-zhou force-pushed the pr4188 branch from 84983c3 to d7f9fb7 Compare August 18, 2026 14:55

ivy-zhou mentioned this pull request Aug 18, 2026

Teach retention about staged and abandoned checkpoints #4197

Open

ivy-zhou force-pushed the pr4187 branch from 6d6f4fb to 61ca324 Compare August 18, 2026 14:55

ivy-zhou changed the base branch from pr4187 to pr4197 August 18, 2026 14:55

ivy-zhou marked this pull request as ready for review August 18, 2026 15:00

ivy-zhou requested review from fegin, tianyu-l, wconstab and wwwjn as code owners August 18, 2026 15:00

ivy-zhou force-pushed the pr4188 branch from d7f9fb7 to 79c1047 Compare August 18, 2026 17:01

meetv18 reviewed Aug 19, 2026

View reviewed changes

ivy-zhou force-pushed the pr4197 branch from 6fefae8 to bde300c Compare August 19, 2026 19:35

ivy-zhou force-pushed the pr4188 branch from 79c1047 to 5d56807 Compare August 19, 2026 19:35

ivy-zhou mentioned this pull request Aug 19, 2026

[models] Handle replicated MoE placements in HF export #4230

Merged

meetv18 approved these changes Aug 19, 2026

View reviewed changes

ivy-zhou force-pushed the pr4197 branch from bde300c to 159bf1e Compare August 19, 2026 21:16

ivy-zhou force-pushed the pr4188 branch from 5d56807 to fe42ba3 Compare August 19, 2026 21:16

ivy-zhou mentioned this pull request Aug 20, 2026

[float8] Use stateless grouped MM for MoE training #4240

Merged

ivy-zhou force-pushed the pr4188 branch from fe42ba3 to d77362a Compare August 20, 2026 04:40

ivy-zhou force-pushed the pr4197 branch from 159bf1e to 3ca8228 Compare August 20, 2026 04:40

Neilblaze approved these changes Aug 21, 2026

View reviewed changes

ivy-zhou force-pushed the pr4197 branch from 3ca8228 to 18939d3 Compare August 22, 2026 02:07

ivy-zhou force-pushed the pr4188 branch from d77362a to 70dac0f Compare August 22, 2026 02:07

This was referenced Aug 22, 2026

[checkpoint] Log native DCP async save durations #4279

Open

[checkpoint] Centralize checkpoint trace spans #4278

Open

[checkpoint] Load Hugging Face checkpoints natively #4277

Open

Navigation Menu

Uh oh!

Route saves through the torch_checkpointing backend - #4188

Route saves through the torch_checkpointing backend#4188
ivy-zhou wants to merge 1 commit into
pr4197from
pr4188

ivy-zhou commented Aug 18, 2026 •

edited

Loading

Uh oh!

meetv18 Aug 19, 2026

Uh oh!

ivy-zhou Aug 19, 2026

Uh oh!

meetv18 Aug 19, 2026

Uh oh!

ivy-zhou Aug 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

ivy-zhou commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

meetv18 Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

ivy-zhou Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

meetv18 Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

ivy-zhou Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ivy-zhou commented Aug 18, 2026 •

edited

Loading