From 0f56ba1459eacae6d1748b00df63b4af9e6f988f Mon Sep 17 00:00:00 2001 From: Ivy Zhou Date: Mon, 17 Aug 2026 23:09:37 -0700 Subject: [PATCH] Fix TorchFTCheckpointManager's contract with BaseCheckpointManager Pyrefly reports three errors in this file, all of them real: _save was declared "-> None" while the base declares "-> bool", and it discarded super()._save()'s result. BaseCheckpointManager.save returns whatever _save returns, so save() handed back None for this manager. Nothing consumes it today -- torchft/trainer.py ignores the result -- but the contract was broken and the next caller to check it would have been surprised. A replica that skips the full save now reports False; the per-replica dataloader checkpoint is a side channel, not the checkpoint this value describes. _wait_for_saving dereferenced save_future without narrowing it. The base's maybe_wait_for_saving guarantees it is set before dispatching here, which the comment already said, so this just asserts what the comment claims. _ft_save assigned dcp_save's "Future | AsyncSaveResponse | None" straight into save_future, typed "Future | None". AsyncMode.ASYNC always yields a plain Future, so assert that, matching how the DCP manager narrows the same call in its own ASYNC branch. Only the first of the three is new: it arrived with the disabled-guard refactor (#4173), which renamed save to _save and made the base's return type load bearing. The other two predate it. Test Plan: python3 -m pyrefly check torchtitan/experiments/torchft/checkpoint.py -> 0 errors (was 3) python3 -m pytest torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py -q -> 2 passed Adds a test covering both branches of the participating_rank guard, so the return value is pinned rather than left to the type checker. --- torchtitan/experiments/torchft/checkpoint.py | 20 ++++++-- .../torchft/tests/test_torchft_checkpoint.py | 47 +++++++++++++++++++ 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/torchtitan/experiments/torchft/checkpoint.py b/torchtitan/experiments/torchft/checkpoint.py index 915c055a35..877c792073 100644 --- a/torchtitan/experiments/torchft/checkpoint.py +++ b/torchtitan/experiments/torchft/checkpoint.py @@ -15,6 +15,7 @@ from __future__ import annotations import time +from concurrent.futures import Future from dataclasses import dataclass from typing import Any, cast @@ -140,7 +141,7 @@ def load_state_dict(state_dict): self.pg = cast(dist.ProcessGroup, dist.new_group(backend="gloo")) @torch.no_grad() - def _save(self, curr_step: int, last_step: bool = False) -> None: + def _save(self, curr_step: int, last_step: bool = False) -> bool: # FT dataloader checkpoint is saved every step (not gated by interval) # to minimize data replay on replica failure. if self.enable_ft_dataloader_checkpoints: @@ -151,14 +152,18 @@ def _save(self, curr_step: int, last_step: bool = False) -> None: # pyrefly: ignore [missing-attribute] and self.ft_manager.participating_rank() == 0 ): - super()._save(curr_step, last_step) - elif self.enable_ft_dataloader_checkpoints: + return super()._save(curr_step, last_step) + if self.enable_ft_dataloader_checkpoints: assert self.ft_manager is not None logger.info( "Replica %d doesn't save checkpoint.", # pyrefly: ignore [missing-attribute] self.ft_manager.participating_rank(), ) + # The per-replica dataloader checkpoint above is a side channel, not the + # checkpoint this return value describes, so a replica that skipped the + # full save reports False. + return False @torch.no_grad() def _load(self, step: int = -1) -> bool: @@ -177,7 +182,8 @@ def _wait_for_saving(self) -> None: # so save_future can exist even when self.async_mode is DISABLED. The DCP # manager would incorrectly raise in that case, so we override to handle # it. BaseCheckpointManager.maybe_wait_for_saving has already checked - # that save_future is set. + # that save_future is set; assert to narrow it for the type checker. + assert self.save_future is not None self.save_future.result() # ASYNC_WITH_PINNED_MEM: the stager manages the future's lifecycle; # all other modes (ASYNC, DISABLED with FT) should clear the future. @@ -199,9 +205,13 @@ def _ft_save(self, step: int) -> None: begin = time.monotonic() self.maybe_wait_for_saving() checkpoint_id = self._create_checkpoint_id(step, folder=self._ft_folder()) - self.save_future = self.dcp_save( + result = self.dcp_save( self.ft_states, checkpoint_id=checkpoint_id, async_mode=AsyncMode.ASYNC ) + # AsyncMode.ASYNC always yields a plain Future; the AsyncSaveResponse and + # None arms of dcp_save's return type belong to the other modes. + assert isinstance(result, Future) + self.save_future = result logger.info(f"Staging torchft checkpoint took {time.monotonic() - begin} secs.") def _ft_load(self) -> None: diff --git a/torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py b/torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py index cad45e445f..d7c8c35960 100644 --- a/torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py +++ b/torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py @@ -163,6 +163,53 @@ def test_torchft_async_save_calls_maybe_wait_for_saving( manager.close() + def _manager(self, participating_rank: int) -> TorchFTCheckpointManager: + config = TorchFTCheckpointManager.Config( + enable=True, + async_mode="disabled", + folder=self.test_folder, + interval=1, + keep_latest_k=0, + last_save_model_only=False, + export_dtype="float32", + exclude_from_loading=[], + initial_load_path=None, + initial_load_model_only=False, + enable_ft_dataloader_checkpoints=True, + ) + return TorchFTCheckpointManager( + config, + dataloader=self.data_loader, + model_parts=self.model_parts, + optimizers=self.optimizers, + lr_schedulers=self.lr_schedulers, + states=self.states, + sd_adapter=None, + base_folder=self.test_folder, + ft_manager=DummyFTManager( + enabled=True, participating_rank=participating_rank + ), + ) + + @mock.patch("torch.cuda.Stream") + @mock.patch.object(dist_checkpoint, "async_save", side_effect=fake_async_save) + def test_save_returns_whether_the_full_checkpoint_was_written( + self, + mock_async_save, + mock_cuda_stream, + ): + # BaseCheckpointManager.save returns _save's result, so this override has + # to report a bool. The per-replica dataloader checkpoint is a side + # channel and does not count as writing the checkpoint. + with mock.patch.object(dist_checkpoint, "save"): + participating = self._manager(participating_rank=0) + self.assertIs(True, participating.save(curr_step=5)) + participating.close() + + bystander = self._manager(participating_rank=1) + self.assertIs(False, bystander.save(curr_step=5)) + bystander.close() + if __name__ == "__main__": unittest.main()