Fix LR scheduler when warmup exceeds cycle length (#3086)

## Summary
- clamp warmup steps to cycle length before constructing per-cycle
schedule phases
- clamp decay steps to remaining cycle budget and fall back to constant
schedule when no decay steps remain
- add regression coverage for a long-warmup/short-run case to prevent
end-of-cycle LR jumps

## Testing
- `uv run --group test pytest tests/test_optimizer_config.py -q` (run in
`lib/levanter`)
- `uv run --with ruff ruff check src/levanter/optim/config.py
tests/test_optimizer_config.py` (run in `lib/levanter`)

Author: dlwh

lib/levanter/src/levanter/optim/config.py

@@ -263,50 +263,56 @@ def lr_scheduler(self, num_train_steps, override_lr=None):
                 warmup_steps = _convert_frac_or_steps(self.warmup, cycle_steps)
             else:
                 warmup_steps = _convert_frac_or_steps(self.rewarmup, cycle_steps)
+            warmup_steps = min(warmup_steps, cycle_steps)
 
             if warmup_steps != 0:
                 warmup = optax.linear_schedule(previous_end, learning_rate, warmup_steps)
                 schedules.append(warmup)
                 boundaries.append(start + warmup_steps)
 
-            lr_decay_steps = (
-                _convert_frac_or_steps(self.decay, cycle_steps)
-                if self.decay is not None
-                else cycle_steps - warmup_steps
+            max_decay_steps = max(cycle_steps - warmup_steps, 0)
+            requested_decay_steps = (
+                _convert_frac_or_steps(self.decay, cycle_steps) if self.decay is not None else max_decay_steps
             )
+            lr_decay_steps = min(max(requested_decay_steps, 0), max_decay_steps)
             stable_steps = cycle_steps - warmup_steps - lr_decay_steps
 
-            if stable_steps != 0:
+            if stable_steps > 0:
                 stable = optax.constant_schedule(learning_rate)
                 schedules.append(stable)
                 boundaries.append(start + warmup_steps + stable_steps)
 
-            if isinstance(self.lr_schedule, str):
-                match self.lr_schedule:
-                    case "constant":
-                        schedule = optax.constant_schedule(learning_rate)
-                    case "cosine":
-                        schedule = optax.cosine_decay_schedule(learning_rate, lr_decay_steps, self.min_lr_ratio)
-                    case "linear":
-                        schedule = optax.linear_schedule(learning_rate, min_lr, lr_decay_steps)
-                    case "inv_sqrt":
-                        schedule = _inv_sqrt_decay_schedule(learning_rate, min_lr, warmup_steps, 10000)
-                    case "inv":
-                        schedule = _inv_decay_schedule(learning_rate, min_lr, lr_decay_steps)
-                    case _:
-                        raise ValueError(f"Unknown lr_schedule: {self.lr_schedule}")
-            elif isinstance(self.lr_schedule, LrSchedule):
-                schedule = self.lr_schedule.build(
-                    LrScheduleContext(
-                        warmup_steps=warmup_steps,
-                        decay_steps=lr_decay_steps,
-                        learning_rate=learning_rate,
-                        min_lr_ratio=self.min_lr_ratio,
-                        min_lr=min_lr,
+            if lr_decay_steps > 0:
+                if isinstance(self.lr_schedule, str):
+                    match self.lr_schedule:
+                        case "constant":
+                            schedule = optax.constant_schedule(learning_rate)
+                        case "cosine":
+                            schedule = optax.cosine_decay_schedule(learning_rate, lr_decay_steps, self.min_lr_ratio)
+                        case "linear":
+                            schedule = optax.linear_schedule(learning_rate, min_lr, lr_decay_steps)
+                        case "inv_sqrt":
+                            schedule = _inv_sqrt_decay_schedule(learning_rate, min_lr, warmup_steps, 10000)
+                        case "inv":
+                            schedule = _inv_decay_schedule(learning_rate, min_lr, lr_decay_steps)
+                        case _:
+                            raise ValueError(f"Unknown lr_schedule: {self.lr_schedule}")
+                elif isinstance(self.lr_schedule, LrSchedule):
+                    schedule = self.lr_schedule.build(
+                        LrScheduleContext(
+                            warmup_steps=warmup_steps,
+                            decay_steps=lr_decay_steps,
+                            learning_rate=learning_rate,
+                            min_lr_ratio=self.min_lr_ratio,
+                            min_lr=min_lr,
+                        )
+                    )
+                else:
+                    raise ValueError(
+                        f"lr_schedule must be a string or an instance of LrSchedule, got {self.lr_schedule}"
                     )
-                )
             else:
-                raise ValueError(f"lr_schedule must be a string or an instance of LrSchedule, got {self.lr_schedule}")
+                schedule = optax.constant_schedule(learning_rate)
 
             previous_end = schedule(lr_decay_steps)
 

lib/levanter/tests/test_optimizer_config.py

@@ -246,3 +246,20 @@ def test_wsds_schedule_with_cycle_points():
     assert np.isclose(sched_fn(701), 1e-3)
     assert np.isclose(sched_fn(969), 1e-3)
     assert sched_fn(971) < 1e-3
+
+
+def test_warmup_longer_than_run_does_not_jump():
+    optimizer = AdamConfig(
+        learning_rate=3e-3,
+        weight_decay=0.0,
+        warmup=1000,
+        decay=0.2,
+        min_lr_ratio=0.1,
+        lr_schedule="cosine",
+    )
+
+    sched_fn = optimizer.lr_scheduler(200)
+
+    assert np.isclose(sched_fn(160), 0.0024, atol=1e-6)
+    assert sched_fn(161) > sched_fn(160)
+    assert np.isclose(sched_fn(200), 3e-3, atol=1e-6)