From e007822450fe3278ab8c5104d4cfdc6338074089 Mon Sep 17 00:00:00 2001 From: Ivy Zhou Date: Fri, 21 Aug 2026 18:30:59 -0700 Subject: [PATCH] Consolidate torch_checkpointing HF final saves 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. --- tests/unit_tests/test_torch_checkpointing.py | 156 ++++++++++++++++++ .../checkpointer/torch_checkpointing.py | 69 +++++++- 2 files changed, 216 insertions(+), 9 deletions(-) diff --git a/tests/unit_tests/test_torch_checkpointing.py b/tests/unit_tests/test_torch_checkpointing.py index 76db114c99..dc46cd9024 100644 --- a/tests/unit_tests/test_torch_checkpointing.py +++ b/tests/unit_tests/test_torch_checkpointing.py @@ -7,7 +7,9 @@ import dataclasses import json import logging +import os import queue +import tempfile import unittest from concurrent.futures import Future from contextlib import nullcontext @@ -20,15 +22,19 @@ import torchtitan.components.checkpointer.torch_checkpointing as manager_module from torch.distributed.checkpoint.stateful import Stateful from torch_checkpointing.barriers import TCPStoreBarrierConfig +from torch_checkpointing.checkpoint_layout import SafetensorsSerialization from torch_checkpointing.checkpoint_manager import ( CheckpointManager as BackendCheckpointManager, ) +from torch_checkpointing.checkpoint_writer import CheckpointWriterConfig from torch_checkpointing.config import ( AsyncCheckpointSaverConfig, SyncCheckpointSaverConfig, ) from torch_checkpointing.default_resharder import DefaultResharder from torch_checkpointing.logging_utils import checkpoint_logging_context +from torch_checkpointing.schema import ItemSpec +from torch_checkpointing.storage.filesystem import LocalFileSystemStorageConfig from torchtitan.components.checkpointer import ( BaseCheckpointManager, CheckpointManager, @@ -77,6 +83,16 @@ def load_state_dict(self, state_dict) -> None: self.value = state_dict["value"] +class _StateDictAdapter: + def __init__(self) -> None: + self.fqn_to_index_mapping = {"hf_weight": 1} + self.to_hf_calls = [] + + def to_hf(self, state_dict): + self.to_hf_calls.append(state_dict) + return {"hf_weight": state_dict["weight"]} + + class TorchCheckpointingManagerTest(unittest.TestCase): def _build_manager( self, @@ -571,6 +587,59 @@ def test_save_stamps_the_step_on_backend_events(self) -> None: backend_manager.save_result.set_result(None) manager.close() + def test_hf_consolidation_uses_the_path_the_backend_supplies(self) -> None: + """Drive a real backend save and check what pre_finalize_callback gets. + + Every other test here mocks the backend, so they cannot catch the + callback's path contract changing underneath us -- which it has. This + asserts against the installed torch_checkpointing: whatever directory + the writer names, that is where the shards are, so the callback must + consolidate from it verbatim. + """ + received: list[str] = [] + with tempfile.TemporaryDirectory() as root: + checkpoint_id = os.path.join(root, "step-1", "sharded") + config = BackendCheckpointManager.Config( + default=ItemSpec(requires_copy=False), + save=SyncCheckpointSaverConfig( + writer_config=CheckpointWriterConfig(barrier_config=None) + ), + # O_DIRECT alignment support varies across CI filesystems and + # is unrelated to the callback-path contract under test. + storage_config=LocalFileSystemStorageConfig(use_direct_io=False), + pre_finalize_callback=lambda path, _logger: received.append(path), + ) + manager = config.build() + try: + manager.save(checkpoint_id, {MODEL: torch.ones(2)}) + finally: + manager.close() + + self.assertEqual(1, len(received)) + self.assertTrue( + os.listdir(received[0]), + f"callback was handed {received[0]!r}, which holds no shards", + ) + + def test_a_finished_hf_export_is_a_valid_checkpoint(self) -> None: + # A final HF export keeps the backend's metadata in its nested "sharded" + # directory and the consolidated files at the root. Recognising only the + # backend metadata would mark a finished export abandoned, and the next + # run's pre-save retention deletes abandoned directories outright. + manager = TorchCheckpointingManager.__new__(TorchCheckpointingManager) + manager._storage = mock.Mock(spec=CheckpointStorage) + + for marker in ("metadata.pkl", "model.safetensors.index.json"): + with self.subTest(marker=marker): + manager._storage.isfile.side_effect = ( + lambda path, marker=marker: path.endswith(marker) + ) + self.assertTrue(manager._is_valid_checkpoint("/tmp/checkpoint/step-5")) + + manager._storage.isfile.side_effect = None + manager._storage.isfile.return_value = False + self.assertFalse(manager._is_valid_checkpoint("/tmp/checkpoint/step-5")) + def test_subprocess_logging_initializes_and_delegates(self) -> None: calls = [] init_fn = mock.Mock(side_effect=lambda *_args: calls.append("existing")) @@ -652,3 +721,90 @@ def test_async_manager_composes_subprocess_logging_initializer(self) -> None: manager._manager_config.subprocess_init_args, ) manager.close() + + @mock.patch.object( + manager_module, + "consolidate_hf_safetensors_checkpoint", + create=True, + ) + def test_hf_final_save_converts_and_consolidates_before_commit( + self, + consolidate, + ) -> None: + adapter = _StateDictAdapter() + config = TorchCheckpointingManager.Config( + enable=True, + keep_latest_k=0, + initial_load_model_only=False, + last_save_model_only=True, + last_save_in_hf=True, + ) + storage_config = mock.Mock() + storage_config.create_storage.return_value = mock.Mock() + backend_config = _default_backend_config() + backend_config.storage_config = storage_config + backend_manager = _BackendManager() + sync_manager = _BackendManager() + sync_manager.save_result = None + with ( + mock.patch.object( + manager_module, + "_default_backend_config", + return_value=backend_config, + ), + mock.patch.object( + BackendCheckpointManager.Config, + "build", + autospec=True, + side_effect=[backend_manager, sync_manager], + ) as build, + ): + manager = config.build( + dataloader=None, + model_parts=[nn.Linear(2, 2)], + optimizers=_Stateful("optimizer"), + lr_schedulers=_Stateful("scheduler"), + states={"train_state": _Stateful("train")}, + sd_adapter=adapter, + base_folder="/tmp", + ) + self.assertTrue(manager.save(curr_step=5, last_step=True)) + + sync_config = build.call_args_list[1].args[0] + self.assertEqual( + "/tmp/checkpoint/step-5/sharded", + sync_manager.save_calls[0][0], + ) + checkpoint = sync_manager.save_calls[0][1] + self.assertEqual({MODEL}, set(checkpoint)) + self.assertEqual({"hf_weight"}, set(checkpoint[MODEL])) + torch.testing.assert_close( + checkpoint[MODEL]["hf_weight"], + manager.states[MODEL].state_dict()["weight"], + ) + model_spec = sync_config.items[MODEL] + self.assertIsInstance( + model_spec.layout.serialization_format, + SafetensorsSerialization, + ) + self.assertEqual(f"{MODEL}_{{rank}}.safetensors", model_spec.layout.file_path) + + # The backend hands the callback the directory the shards were actually + # written to -- its staging directory when a barrier is configured. Feed + # that in and assert it is consolidated as given, with no derivation. + self.assertIsNotNone(sync_config.save.writer_config.barrier_config) + # Built from the writer's public config rather than the backend's + # private _temp_dir_path helper, so a rename upstream cannot break + # collection of this module the way CheckpointWriter.TMP_PREFIX did. + save_path = Path(sync_manager.save_calls[0][0]) + prefix = sync_config.save.writer_config.temp_dir_prefix + staged = save_path.parent / f"{prefix}{save_path.name}" + sync_config.pre_finalize_callback(str(staged), mock.Mock()) + consolidate.assert_called_once_with( + "/tmp/checkpoint/step-5/tmp_sharded", + output_dir="/tmp/checkpoint/step-5", + item_key=MODEL, + fqn_to_index_mapping=adapter.fqn_to_index_mapping, + storage_config=storage_config, + ) + manager.close() diff --git a/torchtitan/components/checkpointer/torch_checkpointing.py b/torchtitan/components/checkpointer/torch_checkpointing.py index 22f4609fc7..2df2065767 100644 --- a/torchtitan/components/checkpointer/torch_checkpointing.py +++ b/torchtitan/components/checkpointer/torch_checkpointing.py @@ -23,6 +23,7 @@ import torch.nn as nn from torch.distributed.checkpoint.state_dict_saver import _stateful_to_state_dict from torch_checkpointing.barriers import TCPStoreBarrierConfig +from torch_checkpointing.checkpoint_layout import LayoutInfo, SafetensorsSerialization from torch_checkpointing.checkpoint_manager import ( CheckpointManager as BackendCheckpointManager, ) @@ -35,6 +36,7 @@ from torch_checkpointing.distributed_metadata import ( METADATA_FILE_NAME as TORCH_CHECKPOINTING_METADATA_FILE_NAME, ) +from torch_checkpointing.hf.consolidation import consolidate_hf_safetensors_checkpoint from torch_checkpointing.logging_utils import checkpoint_logging_context from torch_checkpointing.schema import ItemSpec from torch_checkpointing.staging import CheckpointStagerConfig @@ -66,6 +68,10 @@ _DEFAULT_BARRIER_INIT_TIMEOUT_SEC = 60 _DEFAULT_BARRIER_TIMEOUT_SEC = 600 +# Index the HF consolidation writes at the root of a final export; the +# backend names it after the checkpoint item it consolidated. +_HF_INDEX_FILE_NAME = f"{MODEL}.safetensors.index.json" + # Logger the backend emits its checkpoint events and metrics on. _BACKEND_LOGGER_NAME = "torch_checkpointing" @@ -229,12 +235,7 @@ class TorchCheckpointingManager(BaseCheckpointManager): @dataclass(kw_only=True, slots=True) class Config(BaseCheckpointManager.Config): - def __post_init__(self) -> None: - BaseCheckpointManager.Config.__post_init__(self) - if self.last_save_in_hf: - raise ValueError( - "TorchCheckpointingManager does not support last_save_in_hf yet." - ) + pass def __init__( self, @@ -391,9 +392,15 @@ def _parse_step(self, filename: str) -> tuple[int, bool] | None: return int(match.group("step")), bool(match.group("tmp")) def _is_valid_checkpoint(self, checkpoint_id: str) -> bool: + # Either shape this manager publishes. A resumable checkpoint has the + # backend's metadata at its root. A final HF export does not: its + # backend metadata sits in the nested "sharded" directory the shards + # were written to, and the root holds the consolidated HF files. Probing + # only for the former would classify a finished export as abandoned and + # let the next run's retention delete it. return self._storage.isfile( filesystem.join(checkpoint_id, TORCH_CHECKPOINTING_METADATA_FILE_NAME) - ) + ) or self._storage.isfile(filesystem.join(checkpoint_id, _HF_INDEX_FILE_NAME)) def _maybe_wait_for_staging(self) -> None: # Acquiring the backend lock is what blocks until staging for the last @@ -456,11 +463,55 @@ def _save_last_step(self, curr_step: int) -> None: # The final save must land before the process exits, so retire the async # manager and write synchronously through a fresh one. + checkpoint_id = self._create_checkpoint_id(curr_step) self._manager.close() - manager = _with_sync_save(self._manager_config).build() + manager_config = _with_sync_save(self._manager_config) + input_checkpoint_id = checkpoint_id + if self.last_save_in_hf: + assert self.sd_adapter is not None + states = {MODEL: self.sd_adapter.to_hf(states[MODEL])} + # Ranks write safetensors shards into a nested directory; the + # pre-finalize callback consolidates them up into checkpoint_id, so + # the published checkpoint is HF-layout rather than sharded. + input_checkpoint_id = filesystem.join(checkpoint_id, "sharded") + item_specs = dict(manager_config.items) + model_spec = item_specs.get( + MODEL, + ItemSpec(requires_copy=True, required=False), + ) + item_specs[MODEL] = replace( + model_spec, + layout=LayoutInfo( + f"{MODEL}_{{rank}}.safetensors", + SafetensorsSerialization(), + ), + ) + fqn_to_index_mapping = self.sd_adapter.fqn_to_index_mapping + hf_storage_config = ( + manager_config.storage_config + or LocalFileSystemStorageConfig(use_direct_io=False) + ) + manager_config = replace( + manager_config, + items=item_specs, + # The backend hands the callback the directory the shards were + # actually written to -- its staging directory when a write + # barrier is configured, the final path otherwise -- so + # consolidate from that path as given, deriving nothing from it. + pre_finalize_callback=lambda staged, _event_logger: ( + consolidate_hf_safetensors_checkpoint( + staged, + output_dir=checkpoint_id, + item_key=MODEL, + fqn_to_index_mapping=fqn_to_index_mapping, + storage_config=hf_storage_config, + ) + ), + ) + manager = manager_config.build() try: manager.save( - self._create_checkpoint_id(curr_step), + input_checkpoint_id, _stateful_to_state_dict(states), ) finally: