From 5ab18bec487c5dea7b279f0fccc80ddfa45c2c1d Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Mon, 27 Jul 2026 23:54:52 +0800
Subject: [PATCH] fix: correct epoch computation on resume to avoid redoing
 whole epoch

---
 astrai/trainer/train_context.py | 28 ++++++++++++++++------------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py
index 136fa63..3771ee5 100644
--- a/astrai/trainer/train_context.py
+++ b/astrai/trainer/train_context.py
@@ -101,18 +101,13 @@ class TrainContextBuilder:
                 if checkpoint.config:
                     model_config = checkpoint.config
                 if self._resume:
-                    preloaded_epoch = checkpoint.epoch or cfg.start_epoch
-                    if checkpoint.consumed_samples > 0:
-                        per_step = (
-                            cfg.batch_per_device
-                            * get_world_size()
-                            * cfg.grad_accum_steps
-                        )
-                        preloaded_consumed = (
-                            checkpoint.consumed_samples // per_step
-                        ) * per_step
-                    else:
-                        preloaded_consumed = cfg.start_samples * get_world_size()
+                    preloaded_epoch = checkpoint.epoch
+                    per_step = (
+                        cfg.batch_per_device * get_world_size() * cfg.grad_accum_steps
+                    )
+                    preloaded_consumed = (
+                        checkpoint.consumed_samples // per_step
+                    ) * per_step
                     preloaded_checkpoint = checkpoint
 
         if not model_config and hasattr(cfg.model_fn(), "config"):
@@ -162,6 +157,15 @@ class TrainContextBuilder:
             )
 
         sampler_offset = context.consumed_samples // context.world_size
+
+        if self._resume and sampler_offset > 0:
+            offset = context.world_size - 1
+            num_samples_per_replica = (
+                len(train_dataset) + offset
+            ) // context.world_size
+            if num_samples_per_replica > 0:
+                context.epoch = sampler_offset // num_samples_per_replica
+
         sampler = RDSampler(
             data_source=train_dataset,
             start_epoch=context.epoch,