From 9c2bb23e7a56137b4491f5681b3da357ec2ef1aa Mon Sep 17 00:00:00 2001
From: Yan Barros <yanbarrosyan@gmail.com>
Date: Thu, 24 Sep 2026 12:29:16 -0300
Subject: [PATCH] Save LR scheduler and grad scaler state in checkpoints

Resuming from a checkpoint re-initialized the learning rate scheduler and
the AMP grad scaler, so a resumed run diverged from an uninterrupted one.
Store both state dicts and restore them in load_checkpoint.

Also rename the misspelled optimizer_state_dit key to optimizer_state_dict,
while still reading the old key so existing checkpoints keep loading.

Fixes #77
---
 tests/benchmark/test_trainer_checkpoint.py | 60 ++++++++++++++++++++++
 the_well/benchmark/trainer/training.py     | 37 +++++++++----
 2 files changed, 86 insertions(+), 11 deletions(-)
 create mode 100644 tests/benchmark/test_trainer_checkpoint.py

diff --git a/tests/benchmark/test_trainer_checkpoint.py b/tests/benchmark/test_trainer_checkpoint.py
new file mode 100644
index 00000000..21903f9c
--- /dev/null
+++ b/tests/benchmark/test_trainer_checkpoint.py
@@ -0,0 +1,60 @@
+import pytest
+import torch
+
+pytest.importorskip("wandb")
+from the_well.benchmark.trainer.training import Trainer
+
+
+def _make_trainer(lr_scheduler_steps: int = 0) -> Trainer:
+    # Bypass __init__, which requires a full datamodule, and set only what the
+    # checkpointing methods use.
+    trainer = object.__new__(Trainer)
+    trainer.model = torch.nn.Linear(3, 3)
+    trainer.optimizer = torch.optim.Adam(trainer.model.parameters(), lr=1.0)
+    trainer.lr_scheduler = torch.optim.lr_scheduler.StepLR(
+        trainer.optimizer, step_size=1, gamma=0.5
+    )
+    trainer.grad_scaler = torch.GradScaler("cpu", init_scale=1024.0)
+    trainer.best_val_loss = 0.25
+    for _ in range(lr_scheduler_steps):
+        trainer.optimizer.step()
+        trainer.lr_scheduler.step()
+    return trainer
+
+
+def test_checkpoint_restores_scheduler_and_grad_scaler(tmp_path):
+    trainer = _make_trainer(lr_scheduler_steps=3)
+    trainer.grad_scaler.scale(torch.ones(()))  # initializes the scale
+    trainer.grad_scaler.update(new_scale=64.0)
+    path = tmp_path / "recent.pt"
+    trainer.save_model(epoch=3, validation_loss=0.5, output_path=path)
+
+    resumed = _make_trainer()
+    resumed.load_checkpoint(path)
+
+    assert resumed.starting_epoch == 4
+    assert resumed.lr_scheduler.last_epoch == 3
+    assert resumed.lr_scheduler.get_last_lr() == trainer.lr_scheduler.get_last_lr()
+    assert resumed.optimizer.param_groups[0]["lr"] == pytest.approx(0.125)
+    assert resumed.grad_scaler.get_scale() == 64.0
+
+
+def test_load_checkpoint_with_legacy_optimizer_key(tmp_path):
+    trainer = _make_trainer(lr_scheduler_steps=2)
+    path = tmp_path / "legacy.pt"
+    torch.save(
+        {
+            "epoch": 2,
+            "model_state_dict": trainer.model.state_dict(),
+            "optimizer_state_dit": trainer.optimizer.state_dict(),
+            "validation_loss": 0.5,
+            "best_validation_loss": 0.25,
+        },
+        path,
+    )
+
+    resumed = _make_trainer()
+    resumed.load_checkpoint(path)
+
+    assert resumed.optimizer.param_groups[0]["lr"] == pytest.approx(0.25)
+    assert resumed.lr_scheduler.last_epoch == 0
diff --git a/the_well/benchmark/trainer/training.py b/the_well/benchmark/trainer/training.py
index 7d9e0ad7..497e6855 100755
--- a/the_well/benchmark/trainer/training.py
+++ b/the_well/benchmark/trainer/training.py
@@ -152,16 +152,17 @@ def __init__(
 
     def save_model(self, epoch: int, validation_loss: float, output_path: str):
         """Save the model checkpoint."""
-        torch.save(
-            {
-                "epoch": epoch,
-                "model_state_dict": self.model.state_dict(),
-                "optimizer_state_dit": self.optimizer.state_dict(),
-                "validation_loss": validation_loss,
-                "best_validation_loss": self.best_val_loss,
-            },
-            output_path,
-        )
+        checkpoint = {
+            "epoch": epoch,
+            "model_state_dict": self.model.state_dict(),
+            "optimizer_state_dict": self.optimizer.state_dict(),
+            "grad_scaler_state_dict": self.grad_scaler.state_dict(),
+            "validation_loss": validation_loss,
+            "best_validation_loss": self.best_val_loss,
+        }
+        if self.lr_scheduler is not None:
+            checkpoint["lr_scheduler_state_dict"] = self.lr_scheduler.state_dict()
+        torch.save(checkpoint, output_path)
 
     def load_checkpoint(self, checkpoint_path: str):
         """Load the model checkpoint."""
@@ -170,7 +171,21 @@ def load_checkpoint(self, checkpoint_path: str):
         if self.model is not None:
             self.model.load_state_dict(checkpoint["model_state_dict"])
         if self.optimizer is not None:
-            self.optimizer.load_state_dict(checkpoint["optimizer_state_dit"])
+            # Checkpoints written by older versions used a misspelled key
+            optimizer_state = checkpoint.get(
+                "optimizer_state_dict", checkpoint.get("optimizer_state_dit")
+            )
+            self.optimizer.load_state_dict(optimizer_state)
+        if self.lr_scheduler is not None:
+            if "lr_scheduler_state_dict" in checkpoint:
+                self.lr_scheduler.load_state_dict(checkpoint["lr_scheduler_state_dict"])
+            else:
+                logger.warning(
+                    "Checkpoint has no learning rate scheduler state; "
+                    "the scheduler restarts from its initial state."
+                )
+        if "grad_scaler_state_dict" in checkpoint:
+            self.grad_scaler.load_state_dict(checkpoint["grad_scaler_state_dict"])
         self.best_val_loss = checkpoint["best_validation_loss"]
         self.starting_val_loss = checkpoint["validation_loss"]
         self.starting_epoch = (