Conversation
This was referenced Aug 18, 2026
meta-cla
Bot
added
the
CLA Signed
This was referenced Aug 18, 2026
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
added a commit
that referenced
this pull request
Aug 19, 2026
…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
meetv18
reviewed
Aug 19, 2026
| if match is None: | ||
| return None | ||
| return int(match.group("step")) | ||
| return int(match.group("step")), bool(match.group("tmp")) |
There was a problem hiding this comment.
why? tmp would mean the checkpoint isn't complete yet?
Contributor
Author
There was a problem hiding this comment.
It's just returning the step from the file name and whether or not the checkpoint file is a temporary one.
meetv18
approved these changes
Aug 19, 2026
ivy-zhou
added a commit
that referenced
this pull request
Aug 20, 2026
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
ivy-zhou
added a commit
that referenced
this pull request
Aug 20, 2026
…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
acisseJZhong
pushed a commit
that referenced
this pull request
Aug 21, 2026
…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
Retention counts every step-N directory by name. That is right for a manager
that writes each checkpoint straight to its published name, but not for one
whose backend stages a save under a temporary name and renames it on success.
Two things go wrong there. A directory left behind by a save that never got
renamed holds nothing anyone can resume from, yet it occupies one of the k
slots and evicts a checkpoint that can be resumed from. And the slot arithmetic
assumes the new checkpoint is already on disk, which is false for a manager
that purges before issuing its save.
Both follow from one fact -- whether a save is being written right now -- so
that is the single argument the caller passes. It is a caller's property, not
the method's: the base does not implement _save, so it has no business knowing
when a subclass purges, and keeping the value at the call site means moving the
call cannot silently change retention.
_parse_step widens to report whether a name is a staging directory. The DCP
manager answers False always and passes the default, so its behavior is
unchanged in every respect.
Abandoned directories are deleted synchronously rather than queued. The reason
they are safe to remove is that no save is being written, and that holds only
at the instant of the check -- a queued delete could land after a retry at the
same step recreated the directory it names.
Test Plan:
python3 -m pytest tests/unit_tests/test_checkpoint.py \
tests/unit_tests/test_torch_checkpointing.py \
torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py -q
-> 53 + 26 + 2 passed
Tests cover an incomplete directory failing to evict a valid checkpoint, an
in-flight directory not being mistaken for an abandoned one, and abandoned
directories bypassing the purge queue.
ufmt + flake8 clean.
This was referenced Aug 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Retention counts every step-N directory by name. That is right for a manager
that writes each checkpoint straight to its published name, but not for one
whose backend stages a save under a temporary name and renames it on success.
Two things go wrong there. A directory left behind by a save that never got
renamed holds nothing anyone can resume from, yet it occupies one of the k
slots and evicts a checkpoint that can be resumed from. And the slot arithmetic
assumes the new checkpoint is already on disk, which is false for a manager
that purges before issuing its save.
Both follow from one fact -- whether a save is being written right now -- so
that is the single argument the caller passes. It is a caller's property, not
the method's: the base does not implement _save, so it has no business knowing
when a subclass purges, and keeping the value at the call site means moving the
call cannot silently change retention.
_parse_step widens to report whether a name is a staging directory. The DCP
manager answers False always and passes the default, so its behavior is
unchanged in every respect.
Abandoned directories are deleted synchronously rather than queued. The reason
they are safe to remove is that no save is being written, and that holds only
at the instant of the check -- a queued delete could land after a retry at the
same step recreated the directory it names.
Test Plan:
python3 -m pytest tests/unit_tests/test_checkpoint.py
tests/unit_tests/test_torch_checkpointing.py
torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py -q
-> 53 + 26 + 2 passed
Tests cover an incomplete directory failing to evict a valid checkpoint, an
in-flight directory not being mistaken for an abandoned one, and abandoned
directories bypassing the purge queue.
ufmt + flake8 clean.
Stack created with Sapling. Best reviewed with ReviewStack.