add ckpt for outer schedler

Author: samsja

open_diloco/ckpt_utils.py

@@ -40,6 +40,7 @@ def save_checkpoint(
     model: torch.nn.Module,
     optimizer: torch.optim.Optimizer,
     scheduler: torch.optim.lr_scheduler.LambdaLR,
+    outer_scheduler: torch.optim.lr_scheduler.LambdaLR | None = None,
     outer_optimizer: torch.optim.Optimizer | None = None,
     scaler: torch.cuda.amp.GradScaler | None = None,
     loss: float | None = None,
@@ -81,6 +82,8 @@ def save_checkpoint(
 
     # 2. Save global states
     global_state_dict = {"scheduler": scheduler.state_dict(), "loss": loss if loss is not None else 0}
+    if outer_scheduler is not None:
+        global_state_dict["outer_scheduler"] = outer_scheduler.state_dict()
     if outer_optimizer is not None:
         global_state_dict["outer_optimizer"] = outer_optimizer.state_dict()
     if scaler is not None:
@@ -95,6 +98,7 @@ def load_checkpoint(
     model: torch.nn.Module,
     optimizer: torch.optim.Optimizer,
     scheduler: torch.optim.lr_scheduler.LambdaLR | None = None,
+    outer_scheduler: torch.optim.lr_scheduler.LambdaLR | None = None,
     outer_optimizer: torch.optim.Optimizer | None = None,
     scaler: torch.cuda.amp.GradScaler | None = None,
     data_loader: StatefulDataLoader | None = None,
@@ -139,8 +143,13 @@ def load_checkpoint(
     if scheduler is not None:
         scheduler.load_state_dict(global_state_dict["scheduler"])
         optimizer.param_groups[0]["lr"] = scheduler.get_last_lr()[0]
+
     if outer_optimizer is not None:
         outer_optimizer.load_state_dict(global_state_dict["outer_optimizer"])
+        if outer_scheduler is not None:
+            outer_scheduler.load_state_dict(global_state_dict["outer_scheduler"])
+            outer_optimizer.param_groups[0]["lr"] = outer_scheduler.get_last_lr()[0]
+
     if scaler is not None:
         scaler.load_state_dict(global_state_dict["scaler"])
     return global_state_dict["loss"]

open_diloco/train_fsdp.py

@@ -539,6 +539,7 @@ def outer_scheduler_fn(opt):
                             model=model,
                             optimizer=optimizer.inner_optimizer,
                             scheduler=scheduler,
+                            outer_scheduler=optimizer.outer_scheduler,
                             outer_optimizer=optimizer.state_averager.optimizer,
                             loss=loss_batch.item(),
                             scaler=scaler,