From bb175fda915a3d84cc28080b5fe8e8bc37d7da83 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Wed, 15 Jul 2026 08:52:15 +0800
Subject: [PATCH] fix: resume optimizer LR, step display, and consumed_samples
 alignment

---
 astrai/trainer/train_callback.py | 2 +-
 astrai/trainer/train_context.py  | 7 ++++++-
 scripts/tools/train.py           | 1 +
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py
index 828194f..4abf1f9 100644
--- a/astrai/trainer/train_callback.py
+++ b/astrai/trainer/train_callback.py
@@ -210,7 +210,7 @@ class ProgressBarCallback(TrainCallback):
     @only_on_rank(0)
     def on_optimizer_step(self, context: TrainContext):
         postfix = {
-            "step": context.optimizer_step,
+            "step": f"{context.optimizer_step:d}",
             "loss": f"{context.loss:.4f}",
             "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
         }
diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py
index 933ac4c..8f36896 100644
--- a/astrai/trainer/train_context.py
+++ b/astrai/trainer/train_context.py
@@ -99,7 +99,12 @@ class TrainContextBuilder:
                     context.model_config = checkpoint.config
                 context.epoch = checkpoint.epoch or cfg.start_epoch
                 if checkpoint.consumed_samples > 0:
-                    context.consumed_samples = checkpoint.consumed_samples
+                    per_step = (
+                        cfg.batch_per_device * context.world_size * cfg.grad_accum_steps
+                    )
+                    context.consumed_samples = (
+                        checkpoint.consumed_samples // per_step
+                    ) * per_step
                 else:
                     context.consumed_samples = cfg.start_samples * context.world_size
                 context.checkpoint = checkpoint
diff --git a/scripts/tools/train.py b/scripts/tools/train.py
index b80c048..af3a70f 100644
--- a/scripts/tools/train.py
+++ b/scripts/tools/train.py
@@ -90,6 +90,7 @@ class MuonMix(optim.Optimizer):
     def load_state_dict(self, state_dict: Dict[str, Any]):
         self.muon.load_state_dict(state_dict["muon"])
         self.adamw.load_state_dict(state_dict["adamw"])
+        self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]
 
 
 def parse_args() -> argparse.Namespace: