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: