Wonder if _update_learning_rates is properly implemented

wgchang · April 19, 2023, 1:49am

Lightning-AI/lightning/blob/21ae19c69ff2bc96812ca12a32cf57077407cb97/src/lightning/pytorch/loops/training_epoch_loop.py#L325


      
          """
          trainer = self.trainer
          
          
if not trainer.lr_scheduler_configs or not trainer.lightning_module.automatic_optimization:
              return
          
          
for config in trainer.lr_scheduler_configs:
              if update_plateau_schedulers ^ config.reduce_on_plateau:
                  continue
          
          
    current_idx = self.batch_idx if interval == "step" else trainer.current_epoch
              current_idx += 1  # account for both batch and epoch starts from 0
              # Take step if call to update_learning_rates matches the interval key and
              # the current step modulo the schedulers frequency is zero
              if config.interval == interval and current_idx % config.frequency == 0:
                  monitor_val = None
                  if config.reduce_on_plateau:
                      monitor_key = config.monitor
                      assert monitor_key is not None
                      monitor_val = self._get_monitor_value(monitor_key)
                      if monitor_val is None:

When the interval is “step”, the lr scheduler updates the learning rate based on batch_idx. Shouldn’t this part be updated based on global_step? If gradient accumulation takes place at a specific batch_idx, lt seems that lr_scheduler is expected to update the learning rate in the wrong way.