Fix TorchFTCheckpointManager's contract with BaseCheckpointManager

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:

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