load/store outer optimizer state dict (#277)

Summary:
We don't restore outer optimizer state currently which can lead to bumps in loss because of high learning rate from a new replica. So save the outer optimizer state in the diloco specific state dict.

Pull Request resolved: #277

Reviewed By: d4l3k

Differential Revision: D83512078

fbshipit-source-id: 07c3ca7f4830f2115c3a4586d93c6d0883a38660

Author: tushar00jain

torchft/_test/diloco_trainer.py

@@ -227,10 +227,6 @@ def load_state_dict(self, state_dict: Dict[str, Dict[str, object]]) -> None:
         self.model.to(self.device)
 
         self.inner_optimizer.load_state_dict(state_dict["inner_optim"])
-        for i, optimizer in enumerate(self.outer_optimizers):
-            optimizer.load_state_dict(
-                cast(dict[str, torch.Tensor], state_dict[f"outer_optim"][f"{i}"])
-            )
 
     def state_dict(self) -> Dict[str, Dict[str, object]]:
         """
@@ -244,10 +240,6 @@ def state_dict(self) -> Dict[str, Dict[str, object]]:
         return {
             "model": self.model.state_dict(),
             "inner_optim": self.inner_optimizer.state_dict(),
-            "outer_optim": {
-                f"{i}": optimizer.state_dict()
-                for i, optimizer in enumerate(self.outer_optimizers)
-            },
         }
 
     def train_loop(self) -> dict[str, Any]:

torchft/diloco_regression_test.py

@@ -221,12 +221,18 @@ def train_loop(self) -> Dict[str, Any]:
 
                     for i in range(self.n_fragments):
                         value = cast(
-                            dict[str, torch.Tensor],
+                            dict[str, dict[str, torch.Tensor]],
                             user_state_dict[f"StreamingDiLoCoFragment_{i}"],
                         )
                         parameter_history["global_parameter_history"][local_step][
                             f"layers.{i}.weight"
-                        ] = value["weight"].data.clone().detach().cpu().tolist()
+                        ] = (
+                            value["original_parameters"]["weight"]
+                            .data.clone()
+                            .detach()
+                            .cpu()
+                            .tolist()
+                        )
 
                     manager_steps.add(manager_curr_step)
 

torchft/local_sgd.py

@@ -259,16 +259,21 @@ def register_state_dict_fn(self) -> None:
         fragment_key = f"StreamingDiLoCoFragment_{self._fragment_id}"
 
         # Define load function for this fragment
-        def load_fn(state_dict: Dict[str, torch.Tensor]) -> None:
-            for name, param in state_dict.items():
+        def load_fn(state_dict: Dict[str, Dict[str, torch.Tensor]]) -> None:
+            for name, param in state_dict["original_parameters"].items():
                 if name in self.original_parameters:
                     self.original_parameters[name].copy_(param)
 
+            self._outer_optimizer.load_state_dict(state_dict["outer_optimizer"])
+
         # Define save function for this fragment
-        def save_fn() -> Dict[str, torch.Tensor]:
+        def save_fn() -> Dict[str, Dict[str, torch.Tensor]]:
             return {
-                name: extract_local_tensor(param)
-                for name, param in self.original_parameters.items()
+                "outer_optimizer": self._outer_optimizer.state_dict(),
+                "original_parameters": {
+                    name: extract_local_tensor(param)
+                    for name, param in self.original_parameters.items()
+                },
             }
 
         # Register the functions with the manager

torchft/local_sgd_integ_test.py

@@ -140,29 +140,29 @@ def assert_equal_global_state(
     for step in rep0.keys():
         for i in range(n_fragments):
             torch.testing.assert_close(
-                rep1[step]["user"][f"StreamingDiLoCoFragment_{i}"],
-                rep0[step]["user"][f"StreamingDiLoCoFragment_{i}"],
+                rep1[step]["user"][f"StreamingDiLoCoFragment_{i}"][
+                    "original_parameters"
+                ],
+                rep0[step]["user"][f"StreamingDiLoCoFragment_{i}"][
+                    "original_parameters"
+                ],
                 check_device=False,
                 msg=f"{step=} {i=}",
             )
-        # Check all outer optimizers
-        for i in range(
-            len(
-                cast(
-                    dict[str, dict[str, torch.Tensor]],
-                    rep0[step]["user"]["default"]["outer_optim"],
-                ).keys()
-            )
-        ):
+            # Check all outer optimizers
             torch.testing.assert_close(
                 cast(
                     dict[str, dict[str, torch.Tensor]],
-                    rep1[step]["user"]["default"]["outer_optim"],
-                )[f"{i}"],
+                    rep1[step]["user"][f"StreamingDiLoCoFragment_{i}"][
+                        "outer_optimizer"
+                    ],
+                ),
                 cast(
                     dict[str, dict[str, torch.Tensor]],
-                    rep0[step]["user"]["default"]["outer_optim"],
-                )[f"{i}"],
+                    rep0[step]["user"][f"StreamingDiLoCoFragment_{i}"][
+                        "outer_optimizer"
+                    ],
+                ),
                 check_device=False,
             )