refactor : replace iteration with consumed_samples

- Replace context.iteration with consumed_samples (global sample count)
- Add optimizer_step property derived from consumed_samples
- Checkpoint meta.json stores consumed_samples, drops iteration
- CLI --start_batch renamed to --start_samples (per-rank samples)
- Checkpoint dir naming: epoch_X_step_Y instead of epoch_X_iter_Y
- Metric log entries use step and consumed_samples fields
- Backward compat removed (old iteration checkpoints unsupported)

astrai/config/train_config.py

@@ -47,14 +47,18 @@ class TrainConfig(BaseConfig):
 
     # checkpoint setting
     start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."})
-    start_batch: int = field(
+    start_samples: int = field(
-        default=0, metadata={"help": "Start batch iteration for training."}
+        default=0,
+        metadata={
+            "help": "Start samples count (per rank). Superseded by checkpoint consumed_samples."
+        },
     )
     ckpt_dir: str = field(
         default="./checkpoint", metadata={"help": "Checkpoint directory."}
     )
     ckpt_interval: int = field(
-        default=5000, metadata={"help": "Number of iterations between checkpoints."}
+        default=5000,
+        metadata={"help": "Number of optimizer steps between checkpoints."},
     )
 
     # lora setting

astrai/serialization.py

@@ -136,7 +136,7 @@ def load_state_dict(path: Union[str, Path], broadcast: bool = False) -> dict:
 class Checkpoint:
     state_dict: Dict[str, Any] = field(default_factory=dict)
     epoch: int = 0
-    iteration: int = 0
+    consumed_samples: int = 0
     extra: Dict[str, Any] = field(default_factory=dict)
     meta: Dict[str, Any] = field(default_factory=dict)
     config: Dict[str, Any] = field(default_factory=dict)
@@ -150,7 +150,7 @@ class Checkpoint:
 
         meta = {
             "epoch": self.epoch,
-            "iteration": self.iteration,
+            "consumed_samples": self.consumed_samples,
             "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
             **self.meta,
         }
@@ -176,7 +176,7 @@ class Checkpoint:
         return cls(
             state_dict=state_dict,
             epoch=meta.get("epoch", 0),
-            iteration=meta.get("iteration", 0),
+            consumed_samples=meta.get("consumed_samples", 0),
             extra=extra,
             config=config,
         )

astrai/trainer/train_callback.py

@@ -139,34 +139,35 @@ class CheckpointCallback(TrainCallback):
         self.interval = interval
         self.weight_only = weight_only
         self.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra
-        self.last_ckpt_iter = 0
+        self.last_ckpt_step = 0
 
     def _save_checkpoint(self, context: TrainContext):
         state_dict = context.executor.unwrap_model(context.model)
-        self.last_ckpt_iter = context.iteration
+        self.last_ckpt_step = context.optimizer_step
 
         if get_rank() == 0:
             save_path = os.path.join(
-                self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
+                self.save_dir,
+                f"epoch_{context.epoch}_step_{context.optimizer_step}",
             )
             extra = self.save_extra_fn(context)
             meta = context.config.to_dict()
             context.checkpoint = Checkpoint(
                 state_dict=state_dict,
                 epoch=context.epoch,
-                iteration=context.iteration,
+                consumed_samples=context.consumed_samples,
+                config=context.model_config,
                 extra=extra,
                 meta=meta,
-                config=context.model_config,
             )
             context.checkpoint.save(save_path)
 
     def on_batch_end(self, context: TrainContext):
-        if context.iteration - self.last_ckpt_iter >= self.interval:
+        if context.optimizer_step - self.last_ckpt_step >= self.interval:
             self._save_checkpoint(context)
 
     def on_train_end(self, context: TrainContext):
-        if context.iteration != self.last_ckpt_iter:
+        if context.optimizer_step != self.last_ckpt_step:
             self._save_checkpoint(context)
 
     def on_error(self, context: TrainContext):
@@ -232,7 +233,7 @@ class MetricLoggerCallback(TrainCallback):
         log_interval: int = 1,
         metrics: List[str] = None,
     ):
-        self.last_log_iter = 0
+        self.last_log_flush_step = 0
         self._last_val_loss = None
         self._last_log_step = 0
         self.save_interval = save_interval
@@ -264,31 +265,30 @@ class MetricLoggerCallback(TrainCallback):
             "type": event_type,
             "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
             "epoch": context.epoch,
-            "step": context.iteration // context.config.grad_accum_steps,
+            "step": context.optimizer_step,
-            "iter": context.iteration,
+            "consumed_samples": context.consumed_samples,
             **extra,
         }
         self.log_cache.append(entry)
 
     @only_on_rank(0)
-    def _flush(self, epoch, iter):
+    def _flush(self, epoch, consumed):
-        log_file = self.log_dir / f"epoch_{epoch}_iter_{iter}_metric.jsonl"
+        log_file = self.log_dir / f"epoch_{epoch}_consumed_{consumed}_metric.jsonl"
         log_file.parent.mkdir(parents=True, exist_ok=True)
         with open(log_file, "w") as f:
             for log in self.log_cache:
                 f.write(json.dumps(log) + "\n")
 
     def on_batch_end(self, context):
-        if context.iteration - self.last_log_iter >= self.save_interval:
+        if context.optimizer_step - self.last_log_flush_step >= self.save_interval:
-            self._flush(context.epoch, context.iteration)
+            self._flush(context.epoch, context.optimizer_step)
-            self.last_log_iter = context.iteration
+            self.last_log_flush_step = context.optimizer_step
 
     def on_optimizer_step(self, context):
-        step = context.iteration // context.config.grad_accum_steps
+        if context.optimizer_step - self._last_log_step >= self.log_interval:
-        if step - self._last_log_step >= self.log_interval:
             step_metrics = [m for m in self.metrics if m != "val_loss"]
             self._append("step", context, **self._metrics(context, step_metrics))
-            self._last_log_step = step
+            self._last_log_step = context.optimizer_step
         if context.val_loss is not None and context.val_loss != self._last_val_loss:
             self._append("validation", context, val_loss=context.val_loss)
             self._last_val_loss = context.val_loss
@@ -297,11 +297,11 @@ class MetricLoggerCallback(TrainCallback):
         self._append("epoch", context)
 
     def on_train_end(self, context):
-        if context.iteration != self.last_log_iter:
+        if context.optimizer_step != self.last_log_flush_step:
-            self._flush(context.epoch, context.iteration)
+            self._flush(context.epoch, context.optimizer_step)
 
     def on_error(self, context):
-        self._flush(context.epoch, context.iteration)
+        self._flush(context.epoch, context.optimizer_step)
 
 
 @CallbackFactory.register("validation")
@@ -328,9 +328,9 @@ class ValidationCallback(TrainCallback):
         context.val_loss = avg_loss
         context.model.train()
 
-        step_count = context.iteration // context.config.grad_accum_steps
         logger.info(
-            f"Epoch {context.epoch + 1}, Step {step_count}, Val Loss: {avg_loss:.4f}"
+            f"Epoch {context.epoch + 1}, Step {context.optimizer_step}, "
+            f"Val Loss: {avg_loss:.4f}"
         )
 
     def on_optimizer_step(self, context: TrainContext):
@@ -339,6 +339,5 @@ class ValidationCallback(TrainCallback):
         cfg = context.config
         if cfg.val_step <= 0:
             return
-        step_count = context.iteration // cfg.grad_accum_steps
+        if context.optimizer_step % cfg.val_step == 0:
-        if step_count % cfg.val_step == 0:
             self._run_validation(context)

astrai/trainer/train_context.py

@@ -29,7 +29,7 @@ class TrainContext:
     executor: BaseExecutor = field(default=None)
 
     epoch: int = field(default=0)
-    iteration: int = field(default=0)
+    consumed_samples: int = field(default=0)
     loss: float = field(default=0.0)
     grad_norm: Optional[float] = field(default=None)
     val_dataloader: Optional[DataLoader] = field(default=None)
@@ -39,6 +39,14 @@ class TrainContext:
     rank: int = field(default=0)
     kwargs: Dict[str, Any] = field(default_factory=dict)
 
+    @property
+    def optimizer_step(self) -> int:
+        return self.consumed_samples // (
+            self.config.batch_per_device
+            * self.world_size
+            * self.config.grad_accum_steps
+        )
+
 
 class TrainContextBuilder:
     def __init__(
@@ -90,7 +98,10 @@ class TrainContextBuilder:
                 if checkpoint.config:
                     context.model_config = checkpoint.config
                 context.epoch = checkpoint.epoch or cfg.start_epoch
-                context.iteration = checkpoint.iteration or cfg.start_batch
+                if checkpoint.consumed_samples > 0:
+                    context.consumed_samples = checkpoint.consumed_samples
+                else:
+                    context.consumed_samples = cfg.start_samples * context.world_size
                 context.checkpoint = checkpoint
 
         if cfg.lora is not None:
@@ -116,7 +127,7 @@ class TrainContextBuilder:
                 cfg.dataset, [n_train, n_val], generator=generator
             )
 
-        sampler_offset = context.iteration * cfg.batch_per_device
+        sampler_offset = context.consumed_samples // context.world_size
         sampler = ResumableDistributedSampler(
             data_source=train_dataset,
             start_epoch=context.epoch,

astrai/trainer/trainer.py

@@ -74,7 +74,9 @@ class Trainer:
                         context.loss = loss.item()
                         stand_loss = loss / executor.grad_accum_steps
                         executor.backward(stand_loss)
-                        context.iteration += 1
+                        context.consumed_samples += (
+                            context.config.batch_per_device * context.world_size
+                        )
                         self._call_callbacks("on_batch_end", context)
 
                         if executor.sync_gradients:

scripts/tools/train.py

@@ -175,7 +175,10 @@ def parse_args() -> argparse.Namespace:
         "--start_epoch", type=int, default=0, help="Start epoch for training."
     )
     parser.add_argument(
-        "--start_batch", type=int, default=0, help="Start batch for training."
+        "--start_samples",
+        type=int,
+        default=0,
+        help="Start samples (per rank) for training.",
     )
 
     parser.add_argument(
@@ -317,7 +320,7 @@ def train(
     n_epoch: int,
     batch_per_device: int,
     start_epoch: int,
-    start_batch: int,
+    start_samples: int,
     grad_accum_steps: int,
     warmup_ratio: float,
     ckpt_interval: int,
@@ -444,7 +447,7 @@ def train(
         n_epoch=n_epoch,
         batch_per_device=batch_per_device,
         start_epoch=start_epoch,
-        start_batch=start_batch,
+        start_samples=start_samples,
         ckpt_interval=ckpt_interval,
         grad_accum_steps=grad_accum_steps,
         max_grad_norm=max_grad_norm,

tests/conftest.py

@@ -75,7 +75,7 @@ class MultiTurnDataset(Dataset):
 
 
 class EarlyStoppingDataset(Dataset):
-    """Dataset that triggers early stopping after a specified number of iterations."""
+    """Dataset that triggers early stopping after consuming a specified number of samples."""
 
     def __init__(self, length=10, stop_after=5):
         self.length = length

tests/data/test_checkpoint.py

@@ -25,7 +25,9 @@ def test_single_process():
 
         scheduler.step()
 
-    checkpoint = Checkpoint(state_dict=model.state_dict(), epoch=3, iteration=30)
+    checkpoint = Checkpoint(
+        state_dict=model.state_dict(), epoch=3, consumed_samples=120
+    )
 
     with tempfile.TemporaryDirectory() as tmpdir:
         checkpoint.save(tmpdir)
@@ -33,7 +35,7 @@ def test_single_process():
         loaded_checkpoint = Checkpoint.load(tmpdir)
 
         assert loaded_checkpoint.epoch == 3
-        assert loaded_checkpoint.iteration == 30
+        assert loaded_checkpoint.consumed_samples == 120
 
 
 def test_checkpoint_with_extra():
@@ -46,7 +48,10 @@ def test_checkpoint_with_extra():
         "scheduler": {"last_epoch": 5},
     }
     checkpoint = Checkpoint(
-        state_dict=model.state_dict(), epoch=1, iteration=10, extra=extra
+        state_dict=model.state_dict(),
+        epoch=1,
+        consumed_samples=40,
+        extra=extra,
     )
 
     with tempfile.TemporaryDirectory() as tmpdir:
@@ -77,7 +82,7 @@ def simple_training():
     checkpoint = Checkpoint(
         state_dict=model.state_dict(),
         epoch=2,
-        iteration=10,
+        consumed_samples=40,
     )
 
     rank = get_rank()

tests/trainer/conftest.py

@@ -52,7 +52,7 @@ def create_train_config(
         batch_per_device: Batch size per device (default: 2)
         grad_accum_steps: Gradient accumulation steps (default: 1)
         max_grad_norm: Maximum gradient norm for clipping (default: 1.0)
-        ckpt_interval: Checkpoint save interval in iterations (default: 5)
+        ckpt_interval: Checkpoint save interval in optimizer steps (default: 5)
         random_seed: Random seed for reproducibility (default: 42)
         **kwargs: Additional arguments passed to TrainConfig
 

tests/trainer/test_early_stopping.py

@@ -44,14 +44,14 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
         pass
 
     # Resume from latest checkpoint
-    load_dir = os.path.join(base_test_env["test_dir"], "epoch_0_iter_2")
+    load_dir = os.path.join(base_test_env["test_dir"], "epoch_0_step_1")
     trainer = Trainer(train_config)
     trainer.train(resume_dir=load_dir)
 
-    # Verify checkpoint was saved at expected iteration
+    # Verify checkpoint was saved at expected step
-    load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_iter_10")
+    load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_step_5")
     import json
 
     with open(os.path.join(load_dir, "meta.json")) as f:
         meta = json.load(f)
-    assert meta["iteration"] == 10
+    assert meta["consumed_samples"] == 20