diff --git a/cosmos_framework/callbacks/loss_spike_rollback.py b/cosmos_framework/callbacks/loss_spike_rollback.py index a17cf905..65c9a943 100644 --- a/cosmos_framework/callbacks/loss_spike_rollback.py +++ b/cosmos_framework/callbacks/loss_spike_rollback.py @@ -261,14 +261,17 @@ def _restore(self, model: "ImaginaireModel", optimizer: torch.optim.Optimizer) - return True def _scale_lr(self, optimizer: torch.optim.Optimizer, scheduler: Any, factor: float) -> None: - """Back the learning rate off by rescaling the scheduler's ``base_lrs``. + """Rescale both the next update's rate and the scheduler's base rates. Writing ``param_group["lr"]`` directly would not survive: ``LambdaLR`` recomputes ``lr = base_lrs[i] * lr_lambda(step)`` on every ``scheduler.step()``, so a direct write is discarded on the very next iteration. ``base_lrs`` is the input to that - product and therefore the only durable place to apply the backoff. + product and therefore the durable place to apply the backoff. The trainer + has already stepped the scheduler before this hook, so the current rates + must also change before the next optimizer update. """ schedulers = _real_schedulers(scheduler) + previous_scale = self._lr_scale if self._base_lrs is None: if schedulers: self._base_lrs = [list(sched.base_lrs) for sched in schedulers] @@ -289,6 +292,18 @@ def _scale_lr(self, optimizer: torch.optim.Optimizer, scheduler: Any, factor: fl if schedulers: for sched, bases in zip(schedulers, self._base_lrs): sched.base_lrs = [base * self._lr_scale for base in bases] + # Preserve the scheduler's current warmup/decay multiplier. Changing + # only base_lrs takes effect after the next optimizer update, too late + # to protect that update. The ratio also handles gradual recovery and + # repeated backoffs without applying the cumulative scale twice. + if self._lr_scale != previous_scale: + for opt in _real_optimizers(optimizer): + for group in opt.param_groups: + group["lr"] = group["lr"] * self._lr_scale / previous_scale + # Keep logging and checkpoint state aligned with the next update's + # rate without advancing the schedule or its counters. + for sched in schedulers: + sched._last_lr = [group["lr"] for group in sched.optimizer.param_groups] return for opt, bases in zip(_real_optimizers(optimizer), self._base_lrs): for group, base in zip(opt.param_groups, bases): diff --git a/cosmos_framework/callbacks/loss_spike_rollback_test.py b/cosmos_framework/callbacks/loss_spike_rollback_test.py index e3de7861..b7e84659 100644 --- a/cosmos_framework/callbacks/loss_spike_rollback_test.py +++ b/cosmos_framework/callbacks/loss_spike_rollback_test.py @@ -6,12 +6,18 @@ injected gradient rather than by waiting for training to misbehave. """ +from io import BytesIO +from types import SimpleNamespace + +import pytest import torch from torch import nn from torch.optim.lr_scheduler import LambdaLR from cosmos_framework.callbacks.loss_spike_rollback import LossSpikeRollback +pytestmark = pytest.mark.L0 + def _harness(**kwargs): torch.manual_seed(0) @@ -28,8 +34,8 @@ def _step(model, optimizer, scheduler, callback, grad_scale): parameter.grad = torch.full_like(parameter, grad_scale) callback.on_before_optimizer_step(model, optimizer, scheduler, None) optimizer.step() - callback.on_before_zero_grad(model, optimizer, scheduler) scheduler.step() + callback.on_before_zero_grad(model, optimizer, scheduler) optimizer.zero_grad(set_to_none=True) @@ -38,8 +44,8 @@ def _step_at(model, optimizer, scheduler, callback, grad_scale, iteration): parameter.grad = torch.full_like(parameter, grad_scale) callback.on_before_optimizer_step(model, optimizer, scheduler, None, iteration=iteration) optimizer.step() - callback.on_before_zero_grad(model, optimizer, scheduler, iteration=iteration) scheduler.step() + callback.on_before_zero_grad(model, optimizer, scheduler, iteration=iteration) optimizer.zero_grad(set_to_none=True) @@ -47,6 +53,12 @@ def _spike_at(model, optimizer, scheduler, callback, iteration): _step_at(model, optimizer, scheduler, callback, 10.0, iteration) +def _assert_reported_lrs(scheduler): + rates = [group["lr"] for group in scheduler.optimizer.param_groups] + assert scheduler.get_last_lr() == pytest.approx(rates) + assert scheduler.state_dict()["_last_lr"] == pytest.approx(rates) + + def test_no_rollback_on_steady_gradients(): model, optimizer, scheduler, callback = _harness() for _ in range(20): @@ -90,6 +102,159 @@ def test_lr_backoff_survives_scheduler_step(): assert optimizer.param_groups[0]["lr"] < 0.01, "backoff was undone by scheduler.step()" +@pytest.mark.parametrize("decay", [1.0, 0.95]) +def test_lr_backoff_applies_to_the_next_update(decay): + """Use the trainer order: optimizer, scheduler, rollback, then zero_grad.""" + model, optimizer, _, callback = _harness() + scheduler = LambdaLR(optimizer, lr_lambda=lambda step: decay**step) + for _ in range(20): + _step(model, optimizer, scheduler, callback, 0.01) + _step(model, optimizer, scheduler, callback, 10.0) + + assert callback.rollbacks == 1 + expected = 0.01 * decay**scheduler.last_epoch * 0.8 * 0.5 + assert optimizer.param_groups[0]["lr"] == pytest.approx(expected) + _assert_reported_lrs(scheduler) + used_rates = [] + hook = optimizer.register_step_pre_hook(lambda opt, args, kwargs: used_rates.append(opt.param_groups[0]["lr"])) + try: + _step(model, optimizer, scheduler, callback, 0.01) + finally: + hook.remove() + assert used_rates == pytest.approx([expected]) + # Healthy recovery updates the already-computed rate, not just base_lrs. + assert optimizer.param_groups[0]["lr"] == pytest.approx(0.01 * decay**scheduler.last_epoch * callback._lr_scale) + _assert_reported_lrs(scheduler) + + +@pytest.mark.parametrize("with_scheduler", [True, False]) +def test_scaling_preserves_multiple_optimizers_groups_and_schedule(with_scheduler): + model = nn.Sequential(nn.Linear(3, 4), nn.Linear(4, 2)) + optimizers = [ + torch.optim.AdamW([{"params": [model[0].weight], "lr": 0.02}, {"params": [model[0].bias], "lr": 0.005}]), + torch.optim.AdamW(model[1].parameters(), lr=0.007), + ] + schedulers = [ + LambdaLR(optimizers[0], lr_lambda=[lambda step: 0.9**step, lambda step: 0.8**step]), + LambdaLR(optimizers[1], lr_lambda=lambda step: 0.7**step), + ] + # Reporting must follow each scheduler's own optimizer, not container order. + scheduler = SimpleNamespace(schedulers=list(reversed(schedulers))) if with_scheduler else None + callback = LossSpikeRollback(enabled=True) + for opt, sched in zip(optimizers, schedulers): + opt.step() + if with_scheduler: + sched.step() + initial = [[group["lr"] for group in opt.param_groups] for opt in optimizers] + counters = [(sched.last_epoch, sched._step_count) for sched in schedulers] + for factor in (0.5, 1.02, 0.5, 1.02): + callback._scale_lr(optimizers, scheduler, factor) + for opt, rates in zip(optimizers, initial): + assert [group["lr"] for group in opt.param_groups] == pytest.approx( + [rate * callback._lr_scale for rate in rates] + ) + assert [(sched.last_epoch, sched._step_count) for sched in schedulers] == counters + if with_scheduler: + for sched in schedulers: + _assert_reported_lrs(sched) + if with_scheduler: + next_rates = [[0.02 * 0.9**2, 0.005 * 0.8**2], [0.007 * 0.7**2]] + for opt, sched, rates in zip(optimizers, schedulers, next_rates): + opt.step() + sched.step() + assert [group["lr"] for group in opt.param_groups] == pytest.approx( + [rate * callback._lr_scale for rate in rates] + ) + _assert_reported_lrs(sched) + + +@pytest.mark.parametrize("recover", [False, True]) +def test_optimizer_and_scheduler_checkpoint_rates_agree(recover): + model, optimizer, scheduler, callback = _harness() + for _ in range(20): + _step(model, optimizer, scheduler, callback, 0.01) + _step(model, optimizer, scheduler, callback, 10.0) + if recover: + _step(model, optimizer, scheduler, callback, 0.01) + _assert_reported_lrs(scheduler) + + checkpoint = BytesIO() + torch.save( + {"model": model.state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict()}, + checkpoint, + ) + checkpoint.seek(0) + saved = torch.load(checkpoint, weights_only=True) + restored_model, restored_optimizer, restored_scheduler, _ = _harness() + restored_model.load_state_dict(saved["model"]) + restored_optimizer.load_state_dict(saved["optimizer"]) + restored_scheduler.load_state_dict(saved["scheduler"]) + assert restored_scheduler.state_dict() == scheduler.state_dict() + _assert_reported_lrs(restored_scheduler) + + # Check the next optimizer/scheduler update, without claiming to persist the + # guard's separate rolling history or change its existing resume policy. + for net, opt, sched in ((model, optimizer, scheduler), (restored_model, restored_optimizer, restored_scheduler)): + for parameter in net.parameters(): + parameter.grad = torch.full_like(parameter, 0.01) + opt.step() + sched.step() + _assert_reported_lrs(sched) + assert restored_scheduler.state_dict() == scheduler.state_dict() + for actual, expected in zip(restored_model.parameters(), model.parameters()): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.GPU +@pytest.mark.parametrize("kind", ["generator", "reasoner"]) +def test_native_scheduler_container_reports_backoff_and_recovery(kind): + model = nn.Sequential() + model.add_module("net", nn.Sequential(nn.Linear(8, 8), nn.Linear(8, 8))) + model.cuda() + if kind == "generator": + from cosmos_framework.utils.generator.optimizer import LRSchedulersContainer + from cosmos_framework.utils.generator.optimizer import OptimizersContainer as GeneratorOptimizersContainer + + optimizer = GeneratorOptimizersContainer(model, "adamw", fused=True, lr=0.01, lr_multipliers={"bias": 0.25}) + scheduler = LRSchedulersContainer(optimizer, "constant") + else: + from cosmos_framework.utils.reasoner.optimizer import OptimizersContainer as ReasonerOptimizersContainer + from cosmos_framework.utils.reasoner.optimizer import SchedulersContainer + + optimizer = ReasonerOptimizersContainer( + list(model.net), + {"lr": 0.01, "weight_decay": 0.0, "fused": True}, + "adamw", + {"first": 1.0, "second": 0.25}, + ["first", "second"], + ) + scheduler = SchedulersContainer(optimizer, lr_lambda=lambda step: 0.95**step) + assert [len(opt.param_groups) for opt in optimizer] == ([2] if kind == "generator" else [1, 1]) + callback = LossSpikeRollback(enabled=True, min_observations=5, window=10, rollback_depth=2) + for _ in range(20): + _step(model, optimizer, scheduler, callback, 0.01) + for gradient, expected_scale in ((10.0, 0.4), (0.01, 0.4 * 1.02)): + _step(model, optimizer, scheduler, callback, gradient) + assert callback.rollbacks == 1 + assert callback._lr_scale == pytest.approx(expected_scale) + rates = [[group["lr"] for group in sched.optimizer.param_groups] for sched in scheduler.schedulers] + for sched in scheduler.schedulers: + _assert_reported_lrs(sched) + assert scheduler.get_last_lr() == ([lr for group in rates for lr in group] if kind == "generator" else rates) + assert scheduler.state_dict()["_last_lr"] == rates[0] + + +def test_zero_scale_remains_zero_during_recovery(): + _, optimizer, scheduler, callback = _harness(lr_min_scale=0.0, lr_backoff=0.0) + callback._scale_lr(optimizer, scheduler, 0.5) + _assert_reported_lrs(scheduler) + counters = (scheduler.last_epoch, scheduler._step_count) + callback._scale_lr(optimizer, scheduler, 1.02) + assert optimizer.param_groups[0]["lr"] == 0.0 + _assert_reported_lrs(scheduler) + assert (scheduler.last_epoch, scheduler._step_count) == counters + + def test_stands_down_after_sustained_divergence(): """A run that has truly diverged must not be frozen by an unyielding guard.""" model, optimizer, scheduler, callback = _harness(max_consecutive=3) @@ -102,11 +267,20 @@ def test_stands_down_after_sustained_divergence(): def test_disabled_is_inert(): model, optimizer, scheduler, callback = _harness(enabled=False) - for _ in range(20): - _step(model, optimizer, scheduler, callback, 0.01) - _step(model, optimizer, scheduler, callback, 10.0) + control_model, control_optimizer, control_scheduler, _ = _harness() + for gradient in [0.01] * 20 + [10.0]: + _step(model, optimizer, scheduler, callback, gradient) + for parameter in control_model.parameters(): + parameter.grad = torch.full_like(parameter, gradient) + control_optimizer.step() + control_scheduler.step() + control_optimizer.zero_grad(set_to_none=True) assert callback.rollbacks == 0 assert not callback._ring, "disabled guard should not pay the snapshot memory cost" + assert scheduler.state_dict() == control_scheduler.state_dict() + torch.testing.assert_close(optimizer.state_dict(), control_optimizer.state_dict(), rtol=0, atol=0) + for actual, expected in zip(model.parameters(), control_model.parameters()): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) def test_baseline_does_not_chase_a_deteriorating_run():