Consolidate torch_checkpointing HF final saves

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()

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: