fix: proper checkpoint restore weights-only vs full training state

CHANGELOG.md

@@ -70,6 +70,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Fix `WeatherDataset.__len__` off-by-one in analysis mode (was undercounting by 1 sample), include `num_past_forcing_steps` in the forecast-mode minimum-horizon check, validate forcing-side forecast horizon when forcing is present, and use `min(n_state, n_forcing)` when both are present in analysis mode; raise `IndexError` for out-of-range indices in `WeatherDataset.__getitem__` (with Python-style negative indexing support) [\#312](https://github.com/mllam/neural-lam/pull/312) @kshirajahere
 
+- Fix `--load` to truly restore weights only by default: instead of overriding the optimizer state inside `on_load_checkpoint` (which left epoch / scheduler / callbacks inherited from the checkpoint), `train_model.py` now reconstructs the model via `ForecasterModule.load_from_checkpoint` and calls `trainer.fit(model)` with no `ckpt_path`, so all non-weight training state starts fresh. Also rename the CLI flag from `--restore_opt` to `--load_training_state` to reflect that it controls more than just the optimizer (epoch / scheduler / callbacks too), and assert it requires `--load` [\#240](https://github.com/mllam/neural-lam/pull/240) @Mani212005
+
 ### Maintenance
 
 - Group the existing Neural-LAM citation papers in the README under a `### Core Neural-LAM Publications` subheading for clearer structure [\#633](https://github.com/mllam/neural-lam/pull/633) @HetaviM29

neural_lam/models/module.py

@@ -822,7 +822,3 @@ def on_load_checkpoint(self, checkpoint):
                 )
                 loaded_state_dict[new_key] = loaded_state_dict[old_key]
                 del loaded_state_dict[old_key]
-
-        if not self.restore_opt:
-            opt = self.configure_optimizers()
-            checkpoint["optimizer_states"] = [opt.state_dict()]

neural_lam/train_model.py

@@ -127,9 +127,13 @@ def main(input_args=None):
         help="Path to load model parameters from",
     )
     runtime_group.add_argument(
-        "--restore_opt",
+        "--load_training_state",
         action="store_true",
-        help="If optimizer state should be restored with model",
+        help=(
+            "If the full training state (optimizer, epoch, LR scheduler, "
+            "callbacks) should be restored from the checkpoint passed to "
+            "--load. Without this flag, --load is weights-only."
+        ),
     )
 
     # Model architecture
@@ -371,6 +375,12 @@ def main(input_args=None):
             "Evaluation (--eval) without --load: no checkpoint will be loaded.",
         )
 
+    if args.load_training_state and not args.load:
+        parser.error(
+            "--load_training_state requires --load <checkpoint_path> to "
+            "specify which checkpoint to restore training state from."
+        )
+
     # Get an (actual) random run id as a unique identifier
     random_run_id = random.randint(0, 9999)
 
@@ -452,7 +462,7 @@ def main(input_args=None):
         datastore=datastore,
         loss=args.loss,
         lr=args.lr,
-        restore_opt=args.restore_opt,
+        restore_opt=args.load_training_state,
         n_example_pred=args.n_example_pred,
         create_gif=args.create_gif,
         val_steps_to_log=args.val_steps_to_log,
@@ -510,7 +520,19 @@ def main(input_args=None):
             datamodule=data_module,
             ckpt_path=args.load,
         )
+    elif args.load and not args.load_training_state:
+        # Weights-only restore: load model weights from the checkpoint but
+        # start a fresh training run (epoch/scheduler/optimizer/callbacks
+        # all reset). Passing `ckpt_path=args.load` to `trainer.fit` would
+        # also restore those, which is `--load_training_state` territory.
+        model = load_forecaster_module_from_checkpoint(
+            args.load, config=config, datastore=datastore
+        )
+        trainer.fit(model=model, datamodule=data_module)
     else:
+        # Either no checkpoint (`args.load is None`) or full resume
+        # (`--load_training_state`); in both cases Lightning handles the
+        # state.
         trainer.fit(model=model, datamodule=data_module, ckpt_path=args.load)
 
 

tests/test_cli.py

@@ -134,3 +134,21 @@ def test_unsupported_logger_raises_value_error():
 
     with pytest.raises(ValueError, match="Unsupported logger type"):
         setup_training_logger(datastore, args, run_name="my-run")
+
+
+def test_load_training_state_requires_load():
+    """--load_training_state without --load must fail at the CLI argument
+    check, before reaching `load_config_and_datastore`. The assertion lives
+    in train_model.main and protects against the user expecting to resume
+    training state without supplying a checkpoint."""
+    with patch("neural_lam.train_model.load_config_and_datastore") as mock_load:
+        with pytest.raises(SystemExit, match="2"):
+            neural_lam.train_model.main.__wrapped__(
+                [
+                    "--config_path",
+                    "dummy.yaml",
+                    "--load_training_state",
+                ]
+            )
+
+    mock_load.assert_not_called()

tests/test_train_model_warnings.py

@@ -24,6 +24,7 @@ def test_eval_without_load_warning(eval_val, load_val, expect_warning):
     mock_args.val_steps_to_log = []
     mock_args.var_leads_metrics_watch = "{}"
     mock_args.ar_steps_eval = 10
+    mock_args.load_training_state = False
 
     with patch(
         "neural_lam.train_model.ArgumentParser.parse_args",
@@ -55,6 +56,7 @@ def test_create_gif_forwarded_to_forecaster_module():
     mock_args.create_gif = True
     mock_args.devices = ["auto"]
     mock_args.model = "graph_lam"
+    mock_args.load_training_state = False
 
     captured_kwargs = {}