Fix scale_warmup: apply SSR to warmup boundary in *WithWarmupScheduler classes

[AnjaniKAgr]

What does this PR do?

When scale_warmup=True, the warmup ramp is scaled by the scale-schedule ratio (via self.warmup_scheduler(state, ssr)), but the warmup-end boundary and the post-warmup decay origin are computed from the unscaled t_warmup:

t_warmup = _convert_time(self.t_warmup, state)   # ssr never applied
...
if state.timestamp < t_warmup:                   # boundary uses unscaled value
    if self.scale_warmup:
        return self.warmup_scheduler(state, ssr) # ramp IS scaled

The two halves of the schedule then disagree about where warmup ends, for absolute-unit warmups (ba/ep/tok/sp):

• ssr > 1: the LR multiplier jumps discontinuously mid-ramp. With t_warmup='500ba' and ssr=2.0, the ramp should span [0, 1000)ba, but the multiplier jumps from 0.5 to 1.0 at the unscaled 500ba boundary.
• ssr < 1: the LR freezes at its peak between ssr * t_warmup and t_warmup, and the decay phase is computed against the unscaled warmup end. With ssr=0.02 and t_warmup='10000ba' (max_duration 1000ep), the LR sits at 1.0 for 9,800 batches — 49% of the entire run — before decay begins.

This contradicts the documented contract ("SSR also scales the warmup period").

Fix: apply ssr to the boundary computation when scale_warmup=True, in all four affected classes (MultiStepWithWarmupScheduler, LinearWithWarmupScheduler, CosineAnnealingWithWarmupScheduler, PolynomialWithWarmupScheduler; ConstantWithWarmupScheduler delegates to Linear and is fixed transitively). The four edits are line-identical:

t_warmup = _convert_time(self.t_warmup, state, ssr=ssr if self.scale_warmup else 1.0)

The scaled t_warmup then flows consistently into both the boundary check and the decay fraction (t - t_warmup) / (t_max - t_warmup), so ramp endpoint and decay origin agree with no discontinuity.

A companion commit corrects the post-warmup tau_w formula in these schedulers' docstrings to match the implemented (t - t_warmup) / (t_max - t_warmup).

Behavior notes

• scale_warmup=False (default): unchanged — the ternary passes ssr=1.0.
• ssr == 1.0: unchanged.
• dur-unit warmups: unchanged — _convert_time ignores ssr for TimeUnit.DURATION since max_duration is pre-scaled by the Trainer. The bug only manifests for absolute units.
• All pre-existing test vectors pass unchanged.

Why existing tests didn't catch this

The existing scale_warmup=True vectors only sample points inside the scaled warmup and one point deep into decay — never at or past the boundary where the two halves disagree. Additionally, at the suite's MAX_DURATION='1000ep', the ssr<1 error near the boundary is ~1e-4 per batch, below the abs=1e-3 tolerance.

Tests

Added a regression case to test_scheduler_init: LinearWithWarmupScheduler(t_warmup='500ba', scale_warmup=True) at ssr=2.0, probing the mid-ramp point at the unscaled boundary. The error there is 0.5 in LR multiplier — 500x the test tolerance.

Verified locally: on unfixed main the regression case fails with Obtained: 1.0, Expected: 0.5 ± 1e-3 (1 failed, 1051 passed); with the fix the full tests/optim/test_scheduler.py suite passes (1052 passed). Happy to extend the regression coverage to the other three classes if desired — the call sites are identical.