Fix DeMemWM checkpoint load priority

configurations/training.yaml

@@ -12,7 +12,7 @@ wandb:
   project: worldmem # wandb project name; if not provided, defaults to root folder name [fixme]
   mode: online # set wandb logging to online, offline or dryrun
 
-resume: null # wandb run id to resume logging and loading checkpoint from
+resume: null # wandb run id to resume logging
 load: null # wandb run id containing checkpoint or a path to a checkpoint file
 auto_resume: true # automatically resume training from output_dir/checkpoints when available
 resume_ckpt_path: null # explicit full Lightning checkpoint path for deterministic training resume

experiments/exp_base.py

@@ -65,6 +65,8 @@ _DEMEMWM_FRAME_MEMORY_MARKERS = (
     "r_attn_anchor",
     "r_attn_dynamic",
     "r_attn_revisit",
+    "r_adaLN_modulation",
+    "r_mlp",
     "frame_memory",
     "memory_adapter",
 )
@@ -308,6 +310,19 @@ class BaseExperiment(ABC):
         self.diffusion_model_path = getattr(root_cfg, "diffusion_model_path", None)
         self.vae_path = getattr(root_cfg, "vae_path", None)
         self.pose_predictor_path = getattr(root_cfg, "pose_predictor_path", None)
+        self.auto_resuming = getattr(root_cfg, "_auto_resuming", False)
+
+    def _load_custom_checkpoint_for_task(self) -> None:
+        if self.ckpt_path:
+            load_custom_checkpoint(algo=self.algo, checkpoint_path=self.ckpt_path)
+        elif self.seperate_load:
+            if 'oasis500m' in self.diffusion_model_path:
+                load_custom_checkpoint(algo=self.algo.diffusion_model.model, checkpoint_path=self.diffusion_model_path)
+            else:
+                load_custom_checkpoint(algo=self.algo.diffusion_model, checkpoint_path=self.diffusion_model_path)
+            load_custom_checkpoint(algo=self.algo.vae, checkpoint_path=self.vae_path)
+        else:
+            load_custom_checkpoint(algo=self.algo, checkpoint_path=self.ckpt_path)
 
     def _build_algo(self):
         """
@@ -466,15 +481,15 @@ class BaseLightningExperiment(BaseExperiment):
         )
 
 
-        if self.customized_load:
-            if self.seperate_load:
-                if 'oasis500m' in self.diffusion_model_path:
-                    load_custom_checkpoint(algo=self.algo.diffusion_model.model,checkpoint_path=self.diffusion_model_path)
-                else:
-                    load_custom_checkpoint(algo=self.algo.diffusion_model,checkpoint_path=self.diffusion_model_path)
-                load_custom_checkpoint(algo=self.algo.vae,checkpoint_path=self.vae_path)
-            else:
-                load_custom_checkpoint(algo=self.algo,checkpoint_path=self.ckpt_path)
+        if self.auto_resuming:
+            trainer.fit(
+                self.algo,
+                train_dataloaders=self._build_training_loader(),
+                val_dataloaders=self._build_validation_loader(),
+                ckpt_path=self.ckpt_path,
+            )
+        elif self.customized_load:
+            self._load_custom_checkpoint_for_task()
 
             if self.zero_init_gate:
                 for name, para in self.algo.diffusion_model.named_parameters():
@@ -536,14 +551,7 @@ class BaseLightningExperiment(BaseExperiment):
         )
 
         if self.customized_load:
-            if self.seperate_load:
-                if 'oasis500m' in self.diffusion_model_path:
-                    load_custom_checkpoint(algo=self.algo.diffusion_model.model,checkpoint_path=self.diffusion_model_path)
-                else:
-                    load_custom_checkpoint(algo=self.algo.diffusion_model,checkpoint_path=self.diffusion_model_path)
-                load_custom_checkpoint(algo=self.algo.vae,checkpoint_path=self.vae_path)
-            else:
-                load_custom_checkpoint(algo=self.algo,checkpoint_path=self.ckpt_path)
+            self._load_custom_checkpoint_for_task()
 
             if self.zero_init_gate:
                 for name, para in self.algo.diffusion_model.named_parameters():
@@ -590,14 +598,7 @@ class BaseLightningExperiment(BaseExperiment):
         )
 
         if self.customized_load:
-            if self.seperate_load:
-                if 'oasis500m' in self.diffusion_model_path:
-                    load_custom_checkpoint(algo=self.algo.diffusion_model.model,checkpoint_path=self.diffusion_model_path)
-                else:
-                    load_custom_checkpoint(algo=self.algo.diffusion_model,checkpoint_path=self.diffusion_model_path)
-                load_custom_checkpoint(algo=self.algo.vae,checkpoint_path=self.vae_path)
-            else:
-                load_custom_checkpoint(algo=self.algo,checkpoint_path=self.ckpt_path)
+            self._load_custom_checkpoint_for_task()
 
             if self.zero_init_gate:
                 for name, para in self.algo.diffusion_model.named_parameters():

tests/test_checkpoint_loading.py

@@ -18,6 +18,7 @@ class DummyFrameMemoryBlock(nn.Module):
     def __init__(self, hidden_size=2):
         super().__init__()
         self.r_adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size))
+        self.r_mlp = nn.Linear(hidden_size, hidden_size)
         self.r_attn_anchor = nn.Linear(hidden_size, hidden_size, bias=False)
         self.r_attn_dynamic = nn.Linear(hidden_size, hidden_size, bias=False)
         self.r_attn_revisit = nn.Linear(hidden_size, hidden_size, bias=False)
@@ -32,8 +33,15 @@ class DummyDeMemWM(nn.Module):
 
 def _ones_state(module, skip_frame_memory=False):
     state = {}
+    memory_markers = (
+        "r_attn_anchor",
+        "r_attn_dynamic",
+        "r_attn_revisit",
+        "r_adaLN_modulation",
+        "r_mlp",
+    )
     for key, value in module.state_dict().items():
-        if skip_frame_memory and any(marker in key for marker in ("r_attn_anchor", "r_attn_dynamic", "r_attn_revisit")):
+        if skip_frame_memory and any(marker in key for marker in memory_markers):
             continue
         state[key] = torch.ones_like(value)
     return state
@@ -107,7 +115,7 @@ class CheckpointLoadingTests(unittest.TestCase):
         target = DummyDeMemWM()
         state = _ones_state(target, skip_frame_memory=True)
         for key, value in target.state_dict().items():
-            if "r_attn_anchor" in key:
+            if "r_attn_anchor" in key or "r_adaLN_modulation" in key or "r_mlp" in key:
                 state[key] = torch.ones_like(value)
 
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -116,6 +124,32 @@ class CheckpointLoadingTests(unittest.TestCase):
             with self.assertRaisesRegex(RuntimeError, "Incomplete DeMemWM checkpoint.*r_attn_dynamic"):
                 load_custom_checkpoint(target, ckpt_path)
 
+    def test_incomplete_dememwm_checkpoint_missing_adaln_mlp_fails(self):
+        target = DummyDeMemWM()
+        state = _ones_state(target, skip_frame_memory=True)
+        for key, value in target.state_dict().items():
+            if "r_attn_" in key:
+                state[key] = torch.ones_like(value)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            ckpt_path = Path(tmpdir) / "incomplete_dememwm_missing_adaln_mlp.ckpt"
+            torch.save({"state_dict": state}, ckpt_path)
+            with self.assertRaisesRegex(RuntimeError, "Incomplete DeMemWM checkpoint.*r_adaLN_modulation"):
+                load_custom_checkpoint(target, ckpt_path)
+
+    def test_incomplete_dememwm_checkpoint_missing_r_mlp_fails(self):
+        target = DummyDeMemWM()
+        state = _ones_state(target, skip_frame_memory=True)
+        for key, value in target.state_dict().items():
+            if "r_attn_" in key or "r_adaLN_modulation" in key:
+                state[key] = torch.ones_like(value)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            ckpt_path = Path(tmpdir) / "incomplete_dememwm_missing_r_mlp.ckpt"
+            torch.save({"state_dict": state}, ckpt_path)
+            with self.assertRaisesRegex(RuntimeError, "Incomplete DeMemWM checkpoint.*r_mlp"):
+                load_custom_checkpoint(target, ckpt_path)
+
     def test_directory_load_skips_unreadable_latest_checkpoint(self):
         target = nn.Linear(2, 2)
         state = {key: torch.full_like(value, 3.0) for key, value in target.state_dict().items()}

tests/test_resume_checkpoint_logic.py

@@ -0,0 +1,128 @@
+import unittest
+from unittest.mock import patch
+
+from omegaconf import OmegaConf
+
+from experiments import exp_base
+
+
+class DummyExperiment(exp_base.BaseLightningExperiment):
+    compatible_algorithms = {}
+    compatible_datasets = {}
+
+    def _build_training_loader(self):
+        return None
+
+    def _build_validation_loader(self):
+        return None
+
+    def _build_test_loader(self):
+        return None
+
+
+class DummyTrainer:
+    fit_calls = []
+    validate_calls = []
+    test_calls = []
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def fit(self, *args, **kwargs):
+        self.fit_calls.append({"args": args, "kwargs": kwargs})
+
+    def validate(self, *args, **kwargs):
+        self.validate_calls.append({"args": args, "kwargs": kwargs})
+
+    def test(self, *args, **kwargs):
+        self.test_calls.append({"args": args, "kwargs": kwargs})
+
+
+def _root_cfg(auto_resuming=True):
+    return OmegaConf.create(
+        {
+            "debug": False,
+            "_auto_resuming": auto_resuming,
+            "customized_load": True,
+            "seperate_load": True,
+            "diffusion_model_path": "oasis500m.safetensors",
+            "vae_path": "vit-l-20.safetensors",
+            "algorithm": {"_name": "dummy"},
+            "dataset": {"_name": "dummy"},
+            "experiment": {
+                "debug": False,
+                "num_nodes": 1,
+                "tasks": ["training"],
+                "training": {
+                    "compile": False,
+                    "precision": 32,
+                    "batch_size": 1,
+                    "max_epochs": 1,
+                    "max_steps": 1,
+                    "max_time": None,
+                    "data": {"shuffle": False, "num_workers": 0},
+                    "optim": {"gradient_clip_val": 0, "accumulate_grad_batches": 1},
+                },
+                "validation": {
+                    "compile": False,
+                    "precision": 32,
+                    "inference_mode": True,
+                    "val_every_n_step": None,
+                    "val_every_n_epoch": None,
+                    "limit_batch": 0,
+                    "batch_size": 1,
+                    "data": {"shuffle": False, "num_workers": 0},
+                },
+                "test": {
+                    "compile": False,
+                    "precision": 32,
+                    "inference_mode": True,
+                    "limit_batch": 0,
+                    "batch_size": 1,
+                    "data": {"shuffle": False, "num_workers": 0},
+                },
+            },
+        }
+    )
+
+
+class ResumeCheckpointLogicTests(unittest.TestCase):
+    def test_auto_resume_takes_priority_over_custom_separate_load(self):
+        DummyTrainer.fit_calls = []
+        experiment = DummyExperiment(_root_cfg(auto_resuming=True), logger=None, ckpt_path="/tmp/last.ckpt")
+        experiment.algo = object()
+
+        with patch.object(exp_base.pl, "Trainer", DummyTrainer), patch.object(
+            exp_base,
+            "load_custom_checkpoint",
+            side_effect=AssertionError("custom load should not run during auto-resume"),
+        ):
+            experiment.training()
+
+        self.assertEqual(DummyTrainer.fit_calls[0]["kwargs"]["ckpt_path"], "/tmp/last.ckpt")
+
+    def test_validation_load_checkpoint_takes_priority_over_separate_load(self):
+        DummyTrainer.validate_calls = []
+        experiment = DummyExperiment(_root_cfg(auto_resuming=False), logger=None, ckpt_path="/tmp/trained.ckpt")
+        experiment.algo = object()
+
+        with patch.object(exp_base.pl, "Trainer", DummyTrainer), patch.object(exp_base, "load_custom_checkpoint") as load_mock:
+            experiment.validation()
+
+        load_mock.assert_called_once_with(algo=experiment.algo, checkpoint_path="/tmp/trained.ckpt")
+        self.assertIsNone(DummyTrainer.validate_calls[0]["kwargs"]["ckpt_path"])
+
+    def test_test_load_checkpoint_takes_priority_over_separate_load(self):
+        DummyTrainer.test_calls = []
+        experiment = DummyExperiment(_root_cfg(auto_resuming=False), logger=None, ckpt_path="/tmp/trained.ckpt")
+        experiment.algo = object()
+
+        with patch.object(exp_base.pl, "Trainer", DummyTrainer), patch.object(exp_base, "load_custom_checkpoint") as load_mock:
+            experiment.test()
+
+        load_mock.assert_called_once_with(algo=experiment.algo, checkpoint_path="/tmp/trained.ckpt")
+        self.assertIsNone(DummyTrainer.test_calls[0]["kwargs"]["ckpt_path"])
+
+
+if __name__ == "__main__":
+    unittest.main()

utils/wandb_utils.py

@@ -40,22 +40,6 @@ class SpaceEfficientWandbLogger(WandbLogger):
         expiration_days: Optional[int] = 5,
         **kwargs: Any,
     ) -> None:
-        super().__init__(
-            name=name,
-            save_dir=save_dir,
-            version=version,
-            offline=False,
-            dir=dir,
-            id=id,
-            anonymous=anonymous,
-            project=project,
-            log_model=log_model,
-            experiment=experiment,
-            prefix=prefix,
-            checkpoint_name=checkpoint_name,
-            **kwargs,
-        )
-
         super().__init__(
             name=name,
             save_dir=save_dir,