[ramanathan831]

When LossSpikeRollback runs after optimizer.step() and scheduler.step(), changing only base_lrs leaves the next optimizer update at the old rate. This PR immediately rescales current optimizer-group rates by the change in the guard's scale, preserving the schedule's current warmup/decay multiplier and gradual recovery.

After adjusting the optimizer groups, each underlying scheduler refreshes _last_lr from its own optimizer. This keeps get_last_lr() and serialized scheduler state consistent with the rate consumed by the next update, including multiple optimizers with different parameter-group layouts. It does not call an extra scheduler.step() or advance scheduler counters.

The callback remains disabled by default. The backoff/recovery formulas, zero-scale protection, no-scheduler path, and checkpoint/resume policy are unchanged. This fixes the original finding on #291 and the reporting/state review.

Validation:

• 21 focused tests passed, including actual CUDA optimizer updates using the native generator and reasoner containers on H100. Coverage includes constant/decaying schedules, immediate consumption, repeated backoff/recovery, heterogeneous optimizer groups, reordered schedulers, unchanged counters, zero scale, and disabled-callback equivalence to an unguarded run.
• Model/optimizer/scheduler save-load tests check rate consistency immediately after backoff and recovery, then compare the next update with the uninterrupted optimizer/scheduler trajectory.
• The workflow-selected H100 L0 suite passed: 734 passed, 3 skipped. All three skips are existing hardware/backend capability guards. The separate attention check had 1 pass and 1 capability skip because this torch build lacks FlashAttention-4 sparse backward.
• Changed-file Ruff lint/format, both full pre-commit configurations, Pyrefly (zero errors), and git diff --check passed. AST comparison confirms that the follow-up only adds the cache refresh; other callback behavior and defaults are unchanged.

These are small optimizer/callback regressions and L0 checks, not a new full-model training or DCP resume qualification. The save-load tests do not claim persistence of the guard's separate rolling history or change the native containers' existing resume policies.

[lfengad]

Review of 45be0ce: the immediate LR backoff fix is reasonable, but there is one remaining LR reporting/state consistency issue.

In _scale_lr, lines 292–302, the scheduler bases and optimizer-group rates are updated, but the scheduler's cached _last_lr is not. Consequently, the rate used by the next optimizer update differs from scheduler.get_last_lr().

Minimal CPU reproduction using this PR's callback:

parameter = torch.nn.Parameter(torch.ones(1))
optimizer = torch.optim.AdamW([parameter], lr=0.01)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: 1.0)
optimizer.step()
scheduler.step()

guard = LossSpikeRollback(enabled=True)
guard._scale_lr(optimizer, scheduler, 0.5)

assert optimizer.param_groups[0]["lr"] == 0.004
assert scheduler.get_last_lr()[0] == 0.01
assert scheduler.state_dict()["_last_lr"][0] == 0.01

This has an actual observability consumer: the existing LR logger reads scheduler.get_last_lr()[0] before the optimizer update. After backoff/recovery it can therefore report a different rate from the one actually consumed. The serialized scheduler state is inconsistent with the optimizer rate as well; this is not a claim that every resume path necessarily uses that stale value.

Could you please synchronize the affected scheduler's reported/cached rates after changing optimizer groups, and add tests asserting agreement between optimizer LR, get_last_lr(), and serialized _last_lr after backoff and recovery, including multiple groups/optimizers? Please avoid an extra scheduler.step(), which would advance the schedule.

Validation: the 17 existing regression cases passed when executing the actual functions in an isolated CPU harness (package-level logging/base-class imports isolated); the reproduction above still shows the mismatch. No full training or GPU rerun was performed. The behavioral fix should remain limited to explicitly enabled rollback runs.

[lfengad]

Concrete implementation request following the LR-state review:

After applying the optimizer-group LR ratio, please refresh each underlying scheduler's cached/reporting state from its own optimizer groups, for example:

for sched in schedulers:
    sched._last_lr = [
        group["lr"] for group in sched.optimizer.param_groups
    ]

Place this after the group-rate adjustment in the scheduler branch of _scale_lr. Keep the existing backoff/recovery formulas and zero-scale protection; do not call an extra scheduler.step(), increment scheduler counters, or change callback defaults. Please cover multiple optimizers/groups and ensure optimizer rates, get_last_lr(), and serialized _last_lr agree.

This is a narrowly scoped consistency fix, not a request to change unrelated training or scheduler behavior.

[ramanathan831]

Addressed the LR reporting/state review and concrete follow-up in 8986409.

After the optimizer-group adjustment, _scale_lr() now refreshes each underlying scheduler's _last_lr from that scheduler's own optimizer groups. This makes the optimizer rate, get_last_lr() and serialized _last_lr agree at the next logging/update boundary. There is no extra scheduler.step() or counter increment. The backoff/recovery formulas, zero-scale protection, callback defaults and no-scheduler path are unchanged. See loss_spike_rollback.py:303.

The regressions cover:

• Backoff and healthy recovery in the trainer's actual hook order, including the rate consumed by the next optimizer step.
• Two optimizers with different rates and group counts, different per-group decay functions, and reversed scheduler order, so copying one optimizer's rates to another scheduler cannot pass. Both last_epoch and _step_count remain unchanged during rescaling.
• Zero-scale recovery and disabled-callback equivalence to an unguarded run, including optimizer state and model parameters.
• Actual CUDA AdamW updates with the native generator and reasoner optimizer/scheduler containers.
• Model/optimizer/scheduler serialization immediately after backoff and recovery, with matching reported rates and the same subsequent optimizer/scheduler update after loading.

See loss_spike_rollback_test.py. The new assertions reproduce eight failures against the reviewed implementation. With the four-line production fix, all 21 focused cases pass. The workflow-selected H100 L0 suite passed: 734 passed, 3 skipped. All three skips are existing hardware/backend capability guards. The separate attention check had 1 pass and 1 capability skip because this torch build lacks FlashAttention-4 sparse backward. Ruff/format, both pre-commit configurations, changed-file Pyrefly and diff checks pass.

This addresses the reporting/state inconsistency without claiming that every resume path previously consumed _last_lr. The container-specific resume policies and the guard's rolling-history persistence are outside this change; no full-model training or DCP resume run is claimed.