Add torch_checkpointing save config plumbing (#4058)

.github/workflows/lint.yaml

@@ -58,6 +58,12 @@ jobs:
         run: |
           python -m pip install -r requirements.txt -r requirements-dev.txt
           python -m pip install --force-reinstall --pre --index-url https://download.pytorch.org/whl/nightly/cu130 torch
+          # pyrefly type-checks torchtitan/components/checkpointer/torch_checkpointing.py,
+          # so the library has to be importable here. --no-deps, and after the
+          # nightly install above: torch_checkpointing pins torch>=2.6.0, which
+          # would otherwise pull a stable wheel over the nightly. Mirrors the
+          # install in .ci/docker/common/install_conda.sh.
+          python -m pip install --no-deps "git+https://github.com/meta-pytorch/torch_checkpointing.git@main"
           pre-commit install-hooks
 
       - name: Get changed files

pyproject.toml

@@ -13,6 +13,9 @@ dependencies = [
     # Stateful Dataloader
     "torchdata>=0.8.0",
 
+    # Checkpointing
+    "torch_checkpointing>=0.1.0",
+
     # Hugging Face integrations
     "datasets>=3.6.0,<4.8.0",
     "tokenizers>=0.15.0",
@@ -41,9 +44,6 @@ dev = [
     "expecttest", # test_tokenizer
     "pyrefly==0.45.1",
 ]
-torch-checkpointing = [
-    "torch_checkpointing>=0.1.0",
-]
 remat = [
     "torch_remat>=0.2.0",
 ]

tests/unit_tests/test_torch_checkpointing.py

@@ -0,0 +1,129 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import dataclasses
+import json
+import unittest
+from unittest import mock
+
+import torch.nn as nn
+from torch_checkpointing.barriers import TCPStoreBarrierConfig
+from torch_checkpointing.checkpoint_manager import (
+    CheckpointManager as BackendCheckpointManager,
+)
+from torch_checkpointing.config import AsyncCheckpointSaverConfig
+from torch_checkpointing.dtensor_resharder import DTensorResharder
+
+from torchtitan.components.checkpointer import (
+    BaseCheckpointManager,
+    CheckpointManager,
+    MODEL,
+    OPTIMIZER,
+)
+from torchtitan.components.checkpointer.torch_checkpointing import (
+    _default_backend_config,
+    DEFAULT_TORCH_CHECKPOINTING_BARRIER_TCPSTORE_PORT,
+    TorchCheckpointingManager,
+)
+
+
+class _BackendManager:
+    def __init__(self) -> None:
+        self.closed = False
+
+    def close(self) -> None:
+        self.closed = True
+
+
+class TorchCheckpointingManagerTest(unittest.TestCase):
+    def _build_manager(
+        self, config: TorchCheckpointingManager.Config
+    ) -> tuple[TorchCheckpointingManager, _BackendManager]:
+        backend_manager = _BackendManager()
+        with mock.patch.object(
+            BackendCheckpointManager.Config,
+            "build",
+            return_value=backend_manager,
+        ):
+            manager = config.build(
+                dataloader=None,
+                model_parts=[nn.Linear(2, 2)],
+                optimizers=object(),
+                lr_schedulers=object(),
+                states={},
+                sd_adapter=None,
+                base_folder="/tmp",
+            )
+        return manager, backend_manager
+
+    def test_config_builds_independent_manager(self) -> None:
+        config = TorchCheckpointingManager.Config(
+            enable=True,
+            keep_latest_k=0,
+            initial_load_model_only=False,
+        )
+
+        manager, backend_manager = self._build_manager(config)
+
+        self.assertIsInstance(manager, TorchCheckpointingManager)
+        self.assertNotIsInstance(manager, CheckpointManager)
+        manager.close()
+        self.assertTrue(backend_manager.closed)
+
+    def test_config_adds_no_checkpoint_options(self) -> None:
+        field_names = {
+            field.name for field in dataclasses.fields(TorchCheckpointingManager.Config)
+        }
+        base_field_names = {
+            field.name for field in dataclasses.fields(BaseCheckpointManager.Config)
+        }
+
+        self.assertEqual(field_names, base_field_names)
+
+    def test_config_to_dict_is_json_serializable(self) -> None:
+        config_dict = TorchCheckpointingManager.Config().to_dict()
+
+        json.dumps(config_dict)
+        self.assertNotIn("checkpoint_manager", config_dict)
+
+    def test_disabled_manager_lifecycle_is_noop(self) -> None:
+        config = TorchCheckpointingManager.Config(
+            enable=False,
+            initial_load_model_only=False,
+        )
+        manager, _ = self._build_manager(config)
+
+        self.assertFalse(manager.load())
+        self.assertFalse(manager.save(curr_step=1))
+        self.assertIsNone(manager.maybe_wait_for_staging())
+        manager.close()
+
+    @mock.patch.dict("os.environ", {"MASTER_ADDR": "checkpoint-host"})
+    def test_default_backend_configuration_owns_schema_and_barrier(self) -> None:
+        backend_config = _default_backend_config()
+
+        self.assertIsInstance(backend_config.save, AsyncCheckpointSaverConfig)
+        self.assertTrue(backend_config.save.staging_config.use_pinned_memory)
+        self.assertEqual(set(backend_config.items), {MODEL, OPTIMIZER})
+        for spec in backend_config.items.values():
+            self.assertTrue(spec.requires_copy)
+            self.assertFalse(spec.required)
+            self.assertIsInstance(spec.resharder, DTensorResharder)
+        self.assertIsNotNone(backend_config.default)
+        self.assertFalse(backend_config.default.requires_copy)
+        barrier_config = backend_config.save.writer_config.barrier_config
+        self.assertIsInstance(barrier_config, TCPStoreBarrierConfig)
+        self.assertEqual(barrier_config.master_address, "checkpoint-host")
+        self.assertEqual(
+            barrier_config.tcpstore_port,
+            DEFAULT_TORCH_CHECKPOINTING_BARRIER_TCPSTORE_PORT,
+        )
+
+    def test_legacy_config_has_no_backend_selector(self) -> None:
+        config = CheckpointManager.Config()
+
+        self.assertFalse(hasattr(config, "save_backend"))
+        self.assertFalse(hasattr(config, "load_backend"))

torchtitan/components/checkpointer/torch_checkpointing.py

@@ -0,0 +1,165 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Any
+
+import torch.nn as nn
+from torch_checkpointing.barriers import TCPStoreBarrierConfig
+from torch_checkpointing.checkpoint_manager import (
+    CheckpointManager as BackendCheckpointManager,
+)
+from torch_checkpointing.checkpoint_writer import CheckpointWriterConfig
+from torch_checkpointing.config import AsyncCheckpointSaverConfig
+from torch_checkpointing.dtensor_resharder import DTensorResharder
+from torch_checkpointing.schema import ItemSpec
+from torch_checkpointing.staging import CheckpointStagerConfig
+from torchtitan.components.dataloader import BaseDataLoader
+from torchtitan.components.optimizer import LRSchedulersContainer, OptimizersContainer
+from torchtitan.config import TORCH_DTYPE_MAP
+from torchtitan.protocols.state_dict_adapter import BaseStateDictAdapter
+from torchtitan.tools import filesystem
+
+from .base import (
+    BaseCheckpointManager,
+    DATALOADER,
+    LR_SCHEDULER,
+    MODEL,
+    ModelWrapper,
+    OPTIMIZER,
+)
+
+DEFAULT_TORCH_CHECKPOINTING_BARRIER_TCPSTORE_PORT = 43001
+_DEFAULT_BARRIER_INIT_TIMEOUT_SEC = 60
+_DEFAULT_BARRIER_TIMEOUT_SEC = 600
+
+
+def _item_specs() -> dict[str, ItemSpec]:
+    resharder = DTensorResharder()
+    return {
+        MODEL: ItemSpec(
+            requires_copy=True,
+            resharder=resharder,
+            required=False,
+        ),
+        OPTIMIZER: ItemSpec(
+            requires_copy=True,
+            resharder=resharder,
+            required=False,
+        ),
+    }
+
+
+def _default_backend_config() -> BackendCheckpointManager.Config:
+    barrier_timeout_sec = _DEFAULT_BARRIER_TIMEOUT_SEC
+    save_config = AsyncCheckpointSaverConfig(
+        writer_config=CheckpointWriterConfig(
+            checkpoint_write_barrier_timeout_sec=barrier_timeout_sec,
+            barrier_config=TCPStoreBarrierConfig(
+                master_address=os.environ.get("MASTER_ADDR", "localhost"),
+                tcpstore_port=DEFAULT_TORCH_CHECKPOINTING_BARRIER_TCPSTORE_PORT,
+                timeout_barrier_init_sec=_DEFAULT_BARRIER_INIT_TIMEOUT_SEC,
+                use_checkpoint_barrier_tcpstore_libuv=True,
+            ),
+        ),
+        staging_config=CheckpointStagerConfig(use_pinned_memory=True),
+        wait_timeout_secs=barrier_timeout_sec,
+    )
+    return BackendCheckpointManager.Config(
+        items=_item_specs(),
+        default=ItemSpec(requires_copy=False),
+        save=save_config,
+    )
+
+
+class TorchCheckpointingManager(BaseCheckpointManager):
+    """TorchTitan checkpoint manager backed by ``torch_checkpointing``."""
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(BaseCheckpointManager.Config):
+        pass
+
+    def __init__(
+        self,
+        config: Config,
+        *,
+        dataloader: BaseDataLoader | None,
+        model_parts: list[nn.Module],
+        optimizers: OptimizersContainer,
+        lr_schedulers: LRSchedulersContainer,
+        states: dict[str, Any],
+        sd_adapter: BaseStateDictAdapter | None,
+        base_folder: str = "",
+    ) -> None:
+        self.enable = config.enable
+        if not self.enable:
+            return
+
+        self.folder = filesystem.join(base_folder, config.folder)
+        self.interval = config.interval
+        self.states = states
+        self.states.update(
+            {
+                MODEL: ModelWrapper(model_parts),
+                OPTIMIZER: optimizers,
+                DATALOADER: dataloader,
+                LR_SCHEDULER: lr_schedulers,
+            }
+        )
+
+        self.load_only = config.load_only
+        self.exclude_from_loading = config.exclude_from_loading
+        self.initial_load_path = config.initial_load_path
+        self.initial_load_model_only = config.initial_load_model_only
+        self.initial_load_in_hf = config.initial_load_in_hf
+        self.initial_load_in_hf_quantized = config.initial_load_in_hf_quantized
+        self.enable_first_step_checkpoint = config.enable_first_step_checkpoint
+        self.last_save_model_only = config.last_save_model_only
+        self.last_save_in_hf = config.last_save_in_hf
+        self.export_dtype = TORCH_DTYPE_MAP[config.export_dtype]
+        self.keep_latest_k = config.keep_latest_k
+        self.sd_adapter = sd_adapter
+        if self.last_save_in_hf and self.sd_adapter is None:
+            raise ValueError(
+                "checkpoint.last_save_in_hf is True, but sd_adapter is not provided."
+            )
+
+        self._manager = _default_backend_config().build()
+
+    def __del__(self) -> None:
+        self.close()
+
+    # Save and load routing land in later changes; this one only plumbs config.
+    # The methods are stubbed rather than omitted because BaseCheckpointManager
+    # declares them abstract, so a partial implementation cannot be instantiated.
+
+    def load(self, step: int = -1) -> bool:
+        if not self.enable:
+            return False
+        raise NotImplementedError(
+            "TorchCheckpointingManager does not implement load() yet."
+        )
+
+    def save(self, curr_step: int, last_step: bool = False) -> bool:
+        if not self.enable:
+            return False
+        raise NotImplementedError(
+            "TorchCheckpointingManager does not implement save() yet."
+        )
+
+    def maybe_wait_for_staging(self) -> None:
+        if not self.enable:
+            return
+        raise NotImplementedError(
+            "TorchCheckpointingManager does not implement maybe_wait_for_staging() yet."
+        )
+
+    def close(self) -> None:
+        if hasattr(self, "_manager"):
+            self._manager.close()