Cosine LR scheduler can increase LR after resuming from backup (missing progress clamp)

[dxqb]

Summary

With the plain COSINE LR scheduler, resuming from a backup can make the learning rate drop far below where it stopped and then increase over subsequent steps, instead of continuing to decay. A user reported this after resuming a run from an epoch-50 backup (the yellow/orange curve sits low and rises after the resume point).

Root cause: the schedule is rebuilt on resume and can change length

The LR scheduler state is not saved with the backup. On resume it is reconstructed from scratch (GenericTrainer → create.create_lr_scheduler) from global_step (restored from meta.json) plus the current config/dataset:

scheduler_steps = epochs * approximate_epoch_length / grad_accum - warmup
last_epoch      = global_step / grad_accum - 1

Nothing ties scheduler_steps to the value the original run used. The step count global_step is restored correctly, but the denominator scheduler_steps is recomputed — so if epochs or approximate_epoch_length differs from the original run, the whole cosine curve is silently rescaled.

Reading it off the reporter's chart:

• Original run: at step 43 212 (epoch 50) the LR is ~5.5e-7 ≈ factor 0.5, so progress ≈ 0.5 and the original scheduler_steps ≈ 86 000. It was in the middle of the decay — nothing unusual.
• Resumed run: at the same step the LR is ~1.3e-7 ≈ factor 0.13 and rising, so progress ≈ 1.2 and the rebuilt scheduler_steps ≈ 36 000.

The same global_step now maps to a much later point on a shorter curve. In the reported case the reporter had cleared the latent cache and restarted before returning to the backup, which re-scans the dataset and can change approximate_epoch_length (and any change to the configured epochs does the same). This is why it is hard to reproduce: it only appears when the rebuilt schedule differs from the original.

Why the LR then rises (secondary)

Once the rebuilt schedule is shorter than the resumed step (progress > 1), the plain cosine climbs back up, because the cosine is periodic and lr_lambda_cosine does not bound its input (modules/util/lr_scheduler_util.py):

progress = float(current_step) / float(scheduler_steps)   # can exceed 1
cos_val = 0.5 * (1.0 + math.cos(progress * math.pi))       # rises again for progress in (1, 2)
factor = max(0.0, cos_val)                                  # never negative here, so this doesn't help

0.5·(1 + cos(progress·π)) reaches its trough at progress = 1 and rises back toward 1 for progress ∈ (1, 2). The other decay schedulers hold past the end instead (lr_lambda_linear via max(0.0, …), lr_lambda_rex via an else branch; the restart variants intentionally keep the periodicity but bound it with min(current_step, scheduler_steps - 1)).

Suggested fixes

The real fix is to keep scheduler_steps stable across a resume — e.g. persist the scheduler (or at least the original total_steps/scheduler_steps) in the backup and reuse it, so the restored global_step lands at the same point on the same curve it did during the original run. This restores the correct LR (~factor 0.5 in the example), not just the direction.

As a smaller, defensive guard, lr_lambda_cosine could clamp its input to match its siblings:

progress = min(1.0, float(current_step) / float(scheduler_steps))

Note this only stops the rising pathology — it would pin the resumed LR at min_factor rather than restore the correct mid-decay value, so it is a symptom guard, not a full fix.

---

Drafted by Claude