[BUG] Learning rate not overrided when set `--override-opt_param-scheduler`

[TissueC]

Describe the bug
When setting override-opt_param-scheduler (but still load optimizer and load rng) and setting new learning rate scheduler (including max lr, min lr, decay style, etc.), the learning rate still persists its original scheduler.

A related issue could be #963

To Reproduce

1. Set max_lr as 6e-4 and constant learning rate.
2. Training some steps and save checkpoint.
3. Load the checkpoint (including optimizer params) and override the scheduler (e.g. cosine 6e-4 to 6e-5).
4. Then the bug shows: the learning rate is still constantly 6e-4.

Expected behavior
The learning rate scheduler should be overrided.

Environment (please complete the following information):

• Megatron-LM commit ID [9bcd417]

Proposed fix
The following code (https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/optimizer_param_scheduler.py#L121) leads to reloading the optimizer's learning rate when setting override-opt_param-scheduler

def get_lr(self, param_group: dict) -> float:
        """Learning rate decay functions from:
        https://openreview.net/pdf?id=BJYwwY9ll pg. 4

        Args:
            param_group (dict): parameter group from the optimizer.
        """

        max_lr = param_group.get('max_lr', self.max_lr)
        min_lr = param_group.get('min_lr', self.min_lr)

A possible solution could be

def get_lr(self, param_group: dict) -> float:
        """Learning rate decay functions from:
        https://openreview.net/pdf?id=BJYwwY9ll pg. 4

        Args:
            param_group (dict): parameter group from the optimizer.
        """

        max_lr = self.max_lr
        min_lr = self.min_lr

[puneeshkhanna]

@TissueC - Even I hit this issue. Proposed solution may not work when there are decoupled lrs. I have a fix. Will test and raise PR soon for Nvidia folks to review and merge.

[TissueC]

@puneeshkhanna Great! I'm looking forward to it. Actually the proposed solution was just reverting the commit of decouple lr, thus it indeed could not work when there are decoupled lrs.

[puneeshkhanna]

@TissueC - In training.py , you can do below till I get time to raise a proper PR.
Please import _update_min_and_max_lr_in_param_groups

    if args.load is not None or args.pretrained_checkpoint is not None:
        one_logger and one_logger.log_metrics({
            'load_checkpoint_start_time': one_logger_utils.get_timestamp_in_ms()
        })
        timers('load-checkpoint', log_level=0).start(barrier=True)
        args.iteration, args.num_floating_point_operations_so_far = load_checkpoint(
            model, optimizer, opt_param_scheduler)
        timers('load-checkpoint').stop(barrier=True)
        timers.log(['load-checkpoint'])
        one_logger and one_logger.log_metrics({
            'load_checkpoint_finish_time': one_logger_utils.get_timestamp_in_ms(),
            'load_checkpoint_time': timers('load-checkpoint').active_time()
        })
    else:
        args.iteration = 0
        args.num_floating_point_operations_so_far = 0

    if args.override_opt_param_scheduler:
        # Call below again so that the load checkpoint values are overriden again by the input values
        optimizer.param_groups = _update_min_and_max_lr_in_param_groups(
            optimizer.param_groups,
            lr=config.lr,
            min_lr=config.min_lr,
            decoupled_lr=config.decoupled_lr,
            decoupled_min_lr=config.decoupled_min_lr,
        )

[github-actions]

Marking as stale. No activity in 60 days.

[github-actions]

Marking as stale. No activity in 60 days.

[github-actions]

This issue was closed because it has been inactive for 7 days since being marked as stale.