From 2acfecf4845f6a1ec2bea5cb84cbe0b60c5f1a81 Mon Sep 17 00:00:00 2001 From: Pian Pawakapan Date: Mon, 17 Aug 2026 15:32:29 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .github/workflows/lint.yaml | 6 + pyproject.toml | 6 +- tests/unit_tests/test_torch_checkpointing.py | 129 ++++++++++++++ .../checkpointer/torch_checkpointing.py | 165 ++++++++++++++++++ 4 files changed, 303 insertions(+), 3 deletions(-) create mode 100644 tests/unit_tests/test_torch_checkpointing.py create mode 100644 torchtitan/components/checkpointer/torch_checkpointing.py diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml index 6a72fff95e..f65b84e3ad 100644 --- a/.github/workflows/lint.yaml +++ b/.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 diff --git a/pyproject.toml b/pyproject.toml index 5074eef478..c20565e68b 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/tests/unit_tests/test_torch_checkpointing.py b/tests/unit_tests/test_torch_checkpointing.py new file mode 100644 index 0000000000..91fbcdbdec --- /dev/null +++ b/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")) diff --git a/torchtitan/components/checkpointer/torch_checkpointing.py b/torchtitan/components/checkpointer/torch_checkpointing.py new file mode 100644 index 0000000000..8743bed38a --- /dev/null +++ b/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()