feat: add --resume flag to decouple weight loading from training resumption

- Add --resume bool flag to train.py CLI
- --param_path always loads weights only by default
- --resume restores epoch, consumed_samples, optimizer & scheduler
- Checkpoint.load() now preserves full meta dict
- Update test_early_stopping to use new param_path/resume API

astrai/serialization/checkpoint.py

@@ -2,7 +2,6 @@
 
 import io
 import json
-import os
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -178,6 +177,7 @@ class Checkpoint:
             epoch=meta.get("epoch", 0),
             consumed_samples=meta.get("consumed_samples", 0),
             extra=extra,
+            meta=meta,
             config=config,
         )
 

astrai/trainer/train_context.py

@@ -54,10 +54,12 @@ class TrainContextBuilder:
         config: TrainConfig,
     ):
         self.config = config
-        self._resume_dir: Optional[str] = None
+        self._param_path: Optional[str] = None
+        self._resume: bool = False
 
-    def with_resume_dir(self, resume_dir: Optional[str]) -> Self:
+    def with_param_path(self, param_path: Optional[str], resume: bool = False) -> Self:
-        self._resume_dir = resume_dir
+        self._param_path = param_path
+        self._resume = resume
         return self
 
     def build(self) -> TrainContext:
@@ -74,8 +76,8 @@ class TrainContextBuilder:
         model = model.to(device=device)
 
         model_config = {}
-        if self._resume_dir:
+        if self._param_path:
-            config_path = Path(self._resume_dir) / "config.json"
+            config_path = Path(self._param_path) / "config.json"
             if config_path.exists():
                 model_config = load_json(config_path)
 
@@ -91,23 +93,29 @@ class TrainContextBuilder:
             executor=executor,
         )
 
-        if self._resume_dir:
+        if self._param_path:
-            checkpoint = Checkpoint.load_any(self._resume_dir)
+            checkpoint = Checkpoint.load_any(self._param_path)
             if checkpoint is not None:
                 model.load_state_dict(checkpoint.state_dict, strict=False)
                 if checkpoint.config:
                     context.model_config = checkpoint.config
-                context.epoch = checkpoint.epoch or cfg.start_epoch
+
-                if checkpoint.consumed_samples > 0:
+                if self._resume:
-                    per_step = (
+                    context.epoch = checkpoint.epoch or cfg.start_epoch
-                        cfg.batch_per_device * context.world_size * cfg.grad_accum_steps
+                    if checkpoint.consumed_samples > 0:
-                    )
+                        per_step = (
-                    context.consumed_samples = (
+                            cfg.batch_per_device
-                        checkpoint.consumed_samples // per_step
+                            * context.world_size
-                    ) * per_step
+                            * cfg.grad_accum_steps
-                else:
+                        )
-                    context.consumed_samples = cfg.start_samples * context.world_size
+                        context.consumed_samples = (
-                context.checkpoint = checkpoint
+                            checkpoint.consumed_samples // per_step
+                        ) * per_step
+                    else:
+                        context.consumed_samples = (
+                            cfg.start_samples * context.world_size
+                        )
+                    context.checkpoint = checkpoint
 
         if cfg.lora is not None:
             inject_lora(

astrai/trainer/trainer.py

@@ -52,9 +52,11 @@ class Trainer:
             if method:
                 method(context)
 
-    def _trainer_loop(self, resume_dir: Optional[str] = None):
+    def _trainer_loop(self, param_path: Optional[str] = None, resume: bool = False):
         context = (
-            TrainContextBuilder(self.train_config).with_resume_dir(resume_dir).build()
+            TrainContextBuilder(self.train_config)
+            .with_param_path(param_path, resume=resume)
+            .build()
         )
         executor = context.executor
         self._call_callbacks("on_train_begin", context)
@@ -95,7 +97,7 @@ class Trainer:
         finally:
             self._call_callbacks("on_train_end", context)
 
-    def train(self, resume_dir: Optional[str] = None):
+    def train(self, param_path: Optional[str] = None, resume: bool = False):
         cfg = self.train_config
         spawn_parallel_fn(
             self._trainer_loop,
@@ -105,5 +107,6 @@ class Trainer:
             master_port=cfg.master_port,
             device_type=cfg.device_type,
             start_method=cfg.start_method,
-            resume_dir=resume_dir,
+            param_path=param_path,
+            resume=resume,
         )

scripts/tools/train.py

@@ -116,6 +116,13 @@ def parse_args() -> argparse.Namespace:
         required=True,
         help="Path to the model parameters or resume checkpoint.",
     )
+    parser.add_argument(
+        "--resume",
+        action="store_true",
+        default=False,
+        help="Resume training from checkpoint at --param_path "
+        "(restore epoch, consumed_samples, optimizer & scheduler state).",
+    )
 
     parser.add_argument(
         "--n_epoch", type=int, default=1, help="Number of epochs to train."
@@ -385,6 +392,7 @@ def train(
     train_type: str,
     param_path: str,
     data_root_path: str,
+    resume: bool,
     n_epoch: int,
     batch_per_device: int,
     start_epoch: int,
@@ -531,7 +539,7 @@ def train(
     )
 
     trainer = Trainer(train_config)
-    trainer.train(resume_dir=param_path)
+    trainer.train(param_path=param_path, resume=resume)
 
 
 if __name__ == "__main__":

tests/trainer/test_early_stopping.py

@@ -46,7 +46,7 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
     # Resume from latest checkpoint
     load_dir = os.path.join(base_test_env["test_dir"], "epoch_0_step_1")
     trainer = Trainer(train_config)
-    trainer.train(resume_dir=load_dir)
+    trainer.train(param_path=load_dir, resume=True)
 
     # Verify checkpoint was saved at expected step
     load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_step_5")