add outer lr scheduler

Author: samsja

open_diloco/hivemind_diloco.py

@@ -334,6 +334,7 @@ def __init__(
         inner_optimizer: OptimizerFactory,
         params: Optional[Union[Parameters, ParamGroups]] = None,
         scheduler: Optional[SchedulerFactory] = None,
+        outer_scheduler: Optional[SchedulerFactory] = None,
         averager_opts: Optional[dict] = None,
         grad_compression: CompressionBase = NoCompression(),
         tracker_opts: Optional[dict] = None,
@@ -365,7 +366,7 @@ def __init__(
         # since we have two optimizers, we need to persist the params to a list
         self.num_inner_steps = num_inner_steps
 
-        for opt_or_scheduler in [outer_optimizer, scheduler]:
+        for opt_or_scheduler in [outer_optimizer, scheduler, outer_scheduler]:
             if not (callable(opt_or_scheduler) or opt_or_scheduler is None):
                 raise TypeError("You need to pass inner and outer optimizer as well as scheduler as callable")
 
@@ -405,6 +406,8 @@ def __init__(
         )
         self.diloco_grad_averager = self._make_gradient_averager(compression=grad_compression)
 
+        self.outer_scheduler = outer_scheduler(self.state_averager.optimizer)
+
     def _check_kwargs(self, kwargs) -> None:
         """DiLoCo Optimizer only support a subset of Hivemind Optimizer kwargs.
         This function raise an error if some kwargs are not supported"""
@@ -555,6 +558,9 @@ def step(
         if self.tracker.ready_to_update_epoch:
             self._update_global_epoch()
 
+        if self.outer_scheduler is not None:
+            self.outer_scheduler.step()
+
         return loss
 
     def _compute_schema_hash(self) -> int:

open_diloco/train_fsdp.py

@@ -92,6 +92,8 @@ class HvConfig(BaseConfig):
     galaxy_size: int
     fail_rank_drop: bool = False  # fail if we lose a diloco worker
     warmup_outerstep: int = 10
+    outer_lr_min: float = 0.3
+    outer_scheduler: bool = False
 
     @model_validator(mode="before")
     def cast_str_to_list(cls, values: dict[str, Any]) -> dict[str, Any]:
@@ -179,17 +181,12 @@ def _get_cosine_schedule_with_warmup_lr_lambda(
     *,
     num_warmup_steps: int,
     num_training_steps: int,
-    num_inner_steps: int,
-    warmup_outerstep: int | None,
     num_cycles: float,
     min_lr_rate: float = 0.0,
 ):
     if current_step < num_warmup_steps:
         return float(current_step) / float(max(1, num_warmup_steps))
 
-    if warmup_outerstep is not None and current_step % num_inner_steps < warmup_outerstep:
-        return 0
-
     progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
     factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))
     factor = factor * (1 - min_lr_rate) + min_lr_rate
@@ -201,13 +198,36 @@ def get_cosine_schedule_with_warmup(optimizer, config: Config):
         _get_cosine_schedule_with_warmup_lr_lambda,
         num_warmup_steps=config.warmup_steps,
         num_training_steps=config.total_steps,
-        num_inner_steps=config.hv.local_steps,
-        warmup_outerstep=config.hv.warmup_outerstep,
         num_cycles=0.5,
     )
     return LambdaLR(optimizer, lambda_lr, -1)
 
 
+def _get_lr_outer(
+    current_step: int,
+    *,
+    num_warmup_steps: int,
+    num_training_steps: int,
+    min_lr_rate: float = 0.0,
+):
+    if current_step < num_warmup_steps:
+        return 1
+
+    progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+    factor = 0.5 * (1.0 + math.cos(math.pi * 2.0 * progress))
+    factor = factor * (1 - min_lr_rate) + min_lr_rate
+    return max(0, factor)
+
+
+def get_lr_outer(optimizer, config: Config):
+    lambda_lr = partial(
+        _get_lr_outer,
+        num_warmup_steps=config.warmup_steps,
+        num_training_steps=config.total_steps,
+    )
+    return LambdaLR(optimizer, lambda_lr, -1)
+
+
 def train(config: Config):
     sharding_strategy = get_sharding_strategy(config.sharding_strategy)
     local_rank = int(os.environ["LOCAL_RANK"])
@@ -298,6 +318,9 @@ def scheduler_fn(opt):
             config=config,
         )
 
+    def outer_scheduler_fn(opt):
+        return get_lr_outer(opt, config=config)
+
     if config.hv is not None:
         if config.ckpt.resume:
             # We need to load with a fake optimizer to set the model parameters correctly before initializing the DiLoCoOptimizer
@@ -323,6 +346,7 @@ def scheduler_fn(opt):
             outer_optimizer=outer_optimizer,
             inner_optimizer=inner_optimizer,
             scheduler=None,
+            outer_scheduler=outer_scheduler_fn if config.hv.outer_scheduler else None,
             params=model.parameters(),
             delay_optimizer_step=False,
             delay_grad_averaging=False,
@@ -438,6 +462,7 @@ def scheduler_fn(opt):
             scaler.update()
 
             scheduler.step()
+
             optimizer.zero_grad()
 
             if logging_activations_steps: