From 7fea0b581bd6c1d95e6e04402948464a472e5040 Mon Sep 17 00:00:00 2001
From: CHOU Yuhong <217777185+ChouYuhong@users.noreply.github.com>
Date: Sun, 13 Jul 2025 21:57:24 +0800
Subject: [PATCH] Update checkpoint.py

Dear Lingua Team. Thank you for your Lingua codebase project! Here is a small PR to fix an issue with the learning rate on resume. I saw that after loading a checkpoint, the first step would use the wrong LR, so I added a quick sync to pull the correct value from the scheduler into the optimizer right away.
---
 lingua/checkpoint.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/lingua/checkpoint.py b/lingua/checkpoint.py
index 058a33e8..503c85a6 100644
--- a/lingua/checkpoint.py
+++ b/lingua/checkpoint.py
@@ -290,6 +290,18 @@ def load(
         )
         dcp.load(state_dict, checkpoint_id=path)
         logger.info("Model and optim reloaded")
+
+        current_lrs = train_state.scheduler.get_last_lr()
+        for i, param_group in enumerate(optimizer.param_groups):
+            if i < len(correct_lrs):
+                param_group['lr'] = current_lrs[i]
+            else:
+                logger.warning(
+                    f"Optimizer has more param groups ({len(optimizer.param_groups)}) "
+                    f"than the scheduler has LRs ({len(current_lrs)}). "
+                    f"Param group {i} was not updated."
+                )
+        logger.info("Set the lr to the Right value")
     
     @classmethod
     def instantiate_and_make_dir(cls, args: CheckpointArgs):