adjust dcp api using

pytorch · Dec 13, 2024 · d913398 · d913398
1 parent b0590d4
commit d913398
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/torchtune/training/checkpointing/_checkpoint_client.py b/torchtune/training/checkpointing/_checkpoint_client.py
@@ -188,7 +188,7 @@ def _save_checkpoint_sync(
             # To prevent GPU memory from spiking during checkpoint save,
             # we consolidate the full model and optim state dicts on CPU for rank 0
             model_state_dict = training.gather_cpu_state_dict(
-                model.state_dict(),
+                model,
                 self._is_rank_zero,
                 device=self._device,
             )
@@ -208,6 +208,7 @@ def _save_checkpoint_sync(
             if no_dist:
                 if not self._optimizer_in_bwd:
                     optim_state_dict = training.get_full_optimizer_state_dict(
+                        model,
                         optimizer,
                         self._is_rank_zero,
                         device=self._device,
@@ -217,7 +218,7 @@ def _save_checkpoint_sync(
                         optim_state_dict[
                             param
                         ] = training.get_full_optimizer_state_dict(
-                            opt, self._is_rank_zero, device=self._device
+                            model, opt, self._is_rank_zero, device=self._device
                         )
             else:
                 optim_state_dict = optimizer.state_dict()