flatten optimizerwrapper in checkpoint

test_runner.py

@@ -372,7 +372,7 @@ def build_test_list():
                     "--checkpoint.enable_checkpoint",
                     "--experimental.pipeline_parallel_degree 2",
                     "--training.enable_cpu_offload True",
-                    "--training.enable_optimizer_in_backward True",
+                    "--optimizer.backward True",
                 ],
             ],
             "Enable CPU Offload with PP",

torchtitan/checkpoint.py

@@ -103,9 +103,15 @@ def __init__(
         self,
         model: Union[nn.Module, List[nn.Module]],
         optim: Union[torch.optim.Optimizer, List[torch.optim.Optimizer]],
+        optim_in_bwd: bool = False,
     ) -> None:
         self.model = [model] if isinstance(model, nn.Module) else model
-        self.optim = [optim] if isinstance(optim, torch.optim.Optimizer) else optim
+        if not optim_in_bwd:
+            self.optim = [optim] if isinstance(optim, torch.optim.Optimizer) else optim
+        else:
+            self.optim = [
+                sub_optim for optim_group in optim for sub_optim in optim_group
+            ]
 
     def state_dict(self) -> Dict[str, Any]:
         func = functools.partial(
@@ -123,30 +129,6 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         list(map(func, self.model, self.optim))
 
 
-class OptimizerInBackwardWrapper(OptimizerWrapper):
-    def state_dict(self) -> Dict[str, Any]:
-        func = functools.partial(
-            get_optimizer_state_dict,
-            options=StateDictOptions(flatten_optimizer_state_dict=True),
-        )
-        state_dict = {}
-        for optim in self.optim:
-            for sub_opt in optim:
-                for sd in map(func, self.model, (sub_opt,)):
-                    state_dict.update(sd)
-        return state_dict
-
-    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        func = functools.partial(
-            set_optimizer_state_dict,
-            optim_state_dict=state_dict,
-            options=StateDictOptions(flatten_optimizer_state_dict=True),
-        )
-        for optim in self.optim:
-            for sub_opt in optim:
-                list(map(func, self.model, (sub_opt,)))
-
-
 class Terminate:
     pass
 
@@ -240,16 +222,10 @@ def __init__(
         self.states.update(
             {
                 "model": ModelWrapper(model_parts),
-                "optimizer": (
-                    OptimizerWrapper(
-                        model_parts,
-                        optimizers,
-                    )
-                    if not job_config.optimizer.backward
-                    else OptimizerInBackwardWrapper(
-                        model_parts,
-                        optimizers,
-                    )
+                "optimizer": OptimizerWrapper(
+                    model_parts,
+                    optimizers,
+                    job_config.optimizer.backward,
                 ),
                 "dataloader": dataloader,
             }

torchtitan/optimizer.py

@@ -170,10 +170,7 @@ def step(self):
                 schedulers.step()
 
     class SchedulersInBackwardContainer(SchedulersContainer):
-        """Util for calling step on multiple learning rate schedulers needed for virtual pipeline stages"""
-
-        def __init__(self, schedulers):
-            self.schedulers = schedulers
+        """Util for calling step on multiple learning rate schedulers when optimizers are in backward"""
 
         def step(self):
             for schedulers in self.schedulers: