Skip to content

Consolidate torch_checkpointing HF final saves - #4190

Open
ivy-zhou wants to merge 1 commit into
pr4189from
pr4190
Open

Consolidate torch_checkpointing HF final saves#4190
ivy-zhou wants to merge 1 commit into
pr4189from
pr4190

Conversation

@ivy-zhou

@ivy-zhou ivy-zhou commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Summary:
TorchCheckpointingManager rejected last_save_in_hf outright, because the
backend writes one shard per rank and Hugging Face consumers expect a single
consolidated safetensors checkpoint. Implement that consolidation and lift the
rejection.

On a final save with last_save_in_hf, the model state is converted through
the state dict adapter's to_hf, and the model item is given a safetensors
layout so each rank writes model_{rank}.safetensors. Those shards go into a
nested sharded/ directory rather than the checkpoint root, and the backend's
pre_finalize_callback merges them up into the checkpoint directory before the
checkpoint is published. Consolidating pre-finalize rather than afterwards
means a reader never observes a checkpoint directory in the sharded
intermediate state.

The callback reuses the manager's configured storage, falling back to local
filesystem with direct IO disabled, so consolidation reads and writes through
the same storage abstraction as the save itself rather than assuming a POSIX
path.

Only the final save is affected. Periodic saves are unchanged and remain native
format, matching the DCP manager, where last_save_in_hf is likewise a
final-save-only option.

Ported from an internal change onto the current OSS layout. The Config
validation this removes was added by the preceding commit in this stack, which
introduced the rejection as a deliberate placeholder.

Test Plan:
pytest tests/unit_tests/test_torch_checkpointing.py: 20 passed.

The new test drives a full final save with last_save_in_hf=True and asserts
the whole chain: the payload is routed through the adapter's to_hf, the
resulting tensor matches the model's own weight, shards are written to
step-5/sharded rather than step-5, the model item carries a
SafetensorsSerialization layout with the model_{rank}.safetensors file
pattern, and invoking the registered pre_finalize_callback calls
consolidate_hf_safetensors_checkpoint with the checkpoint root as
output_dir, the adapter's fqn_to_index_mapping, and the manager's own
storage config.

pytest tests/unit_tests/test_checkpoint.py tests/unit_tests/observability/
alongside the above: 142 passed, 2 subtests passed.

Also verified:

  • LayoutInfo, SafetensorsSerialization,
    consolidate_hf_safetensors_checkpoint, EventLogger, StorageConfig,
    Config.pre_finalize_callback, and ItemSpec.layout all 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.

@meetv18

meetv18 commented Aug 19, 2026

Copy link
Copy Markdown

do we have consolidation numbers for DCPv1 v/s v2? I would suggest to add to the test plan of this diff the following:

  1. Bitwise identical HF checkpoints between DCPv1 and DCPv2
  2. Latency regression checks for E2E HF checkpoint writing (incl. consolidation).

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
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
Summary:
`TorchCheckpointingManager` rejected `last_save_in_hf` outright, because the
backend writes one shard per rank and Hugging Face consumers expect a single
consolidated safetensors checkpoint. Implement that consolidation and lift the
rejection.

On a final save with `last_save_in_hf`, the model state is converted through
the state dict adapter's `to_hf`, and the model item is given a safetensors
layout so each rank writes `model_{rank}.safetensors`. Those shards go into a
nested `sharded/` directory rather than the checkpoint root, and the backend's
`pre_finalize_callback` merges them up into the checkpoint directory before the
checkpoint is published. Consolidating pre-finalize rather than afterwards
means a reader never observes a checkpoint directory in the sharded
intermediate state.

The callback reuses the manager's configured storage, falling back to local
filesystem with direct IO disabled, so consolidation reads and writes through
the same storage abstraction as the save itself rather than assuming a POSIX
path.

Only the final save is affected. Periodic saves are unchanged and remain native
format, matching the DCP manager, where `last_save_in_hf` is likewise a
final-save-only option.

Ported from an internal change onto the current OSS layout. The `Config`
validation this removes was added by the preceding commit in this stack, which
introduced the rejection as a deliberate placeholder.

Test Plan:
`pytest tests/unit_tests/test_torch_checkpointing.py`: 20 passed.

The new test drives a full final save with `last_save_in_hf=True` and asserts
the whole chain: the payload is routed through the adapter's `to_hf`, the
resulting tensor matches the model's own weight, shards are written to
`step-5/sharded` rather than `step-5`, the model item carries a
`SafetensorsSerialization` layout with the `model_{rank}.safetensors` file
pattern, and invoking the registered `pre_finalize_callback` calls
`consolidate_hf_safetensors_checkpoint` with the checkpoint root as
`output_dir`, the adapter's `fqn_to_index_mapping`, and the manager's own
storage config.

`pytest tests/unit_tests/test_checkpoint.py tests/unit_tests/observability/`
alongside the above: 142 passed, 2 subtests passed.

Also verified:
- `LayoutInfo`, `SafetensorsSerialization`,
  `consolidate_hf_safetensors_checkpoint`, `EventLogger`, `StorageConfig`,
  `Config.pre_finalize_callback`, and `ItemSpec.layout` all exist in
  `torch_checkpointing` 0.1.0.
- `ufmt` and `flake8 --config=.flake8` clean on both changed files.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/8gpu CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants