lr scheduler cosine with restarts number of restart is half the specified value if gradient accumulation set to 2

[zhuliyi0]

As title. I set num_cycle to 10 but got 5.

[bghira]

this would be an upstream bug

[bghira]

in fact, it seems this is just how gradient accumulation steps work.

effectively most parts of the tool become halved when you accumulate the loss over 2 iterations.

The warmup phase happens for num_warmup_steps and the learning rate increases linearly from 0 to its initial value during this period.
After the warmup, the learning rate follows a cosine schedule with hard restarts for (num_training_steps - num_warmup_steps) steps.
The number of hard restarts during this period is determined by num_cycles.
If you use gradient accumulation with 2 steps, you would have two micro-batches per optimizer update.

Here's how it would affect the number of effective cycles:

Warmup Steps: The warmup steps will take twice as long in terms of real time because each optimizer update will only happen after two micro-batches. However, in terms of scheduler steps, it remains the same.

Cosine Schedule: Similarly, each cycle of the cosine schedule will take twice as long in real-time but will remain the same in terms of scheduler steps.

Effective Cycles: The effective number of cycles remains the same, which is num_cycles, but each cycle will take twice as many real iterations (forward and backward passes) to complete.

So, the effective number of cycles would still be 4. However, each cycle would last twice as long in terms of real-time iterations, making your training effectively longer.

here is the relevant snippet:

def get_cosine_with_hard_restarts_schedule_with_warmup(
    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1
):
    """
    Create a schedule with a learning rate that decreases following the values of the cosine function between the
    initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases
    linearly between 0 and the initial lr set in the optimizer.

    Args:
        optimizer ([`~torch.optim.Optimizer`]):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (`int`):
            The number of steps for the warmup phase.
        num_training_steps (`int`):
            The total number of training steps.
        num_cycles (`int`, *optional*, defaults to 1):
            The number of hard restarts to use.
        last_epoch (`int`, *optional*, defaults to -1):
            The index of the last epoch when resuming training.

    Return:
        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        if progress >= 1.0:
            return 0.0
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))
    return LambdaLR(optimizer, lr_lambda, last_epoch)