Add DeMemWM curriculum training support

configurations/experiment/exp_video.yaml

@@ -9,6 +9,17 @@ training:
   batch_size: 4
   max_epochs: -1
   max_steps: 2000005
+  curriculum:
+    enabled: true
+    stages:
+      - name: no_updown
+        until_step: 240000
+        dataset:
+          wo_updown: true
+      - name: with_updown
+        until_step: ${experiment.training.max_steps}
+        dataset:
+          wo_updown: false
   checkpointing:
     every_n_train_steps: 2500
   optim:

experiments/exp_base.py

@@ -21,7 +21,7 @@ from lightning.pytorch.utilities.types import TRAIN_DATALOADERS
 from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint
 from pytorch_lightning.utilities import rank_zero_info
 
-from omegaconf import DictConfig
+from omegaconf import DictConfig, open_dict
 
 from utils.print_utils import cyan
 from utils.distributed_utils import is_rank_zero
@@ -395,6 +395,108 @@ class BaseLightningExperiment(BaseExperiment):
         callbacks = []
         if self.logger:
             callbacks.append(LearningRateMonitor("step", True))
+        if "checkpointing" in self.cfg.training:
+            callbacks.append(
+                ModelCheckpoint(
+                    pathlib.Path(hydra.core.hydra_config.HydraConfig.get()["runtime"]["output_dir"]) / "checkpoints",
+                    filename="epoch{epoch}_step{step}",
+                    auto_insert_metric_name=False,
+                    **self.cfg.training.checkpointing,
+                )
+            )
+        return callbacks
+
+    def _build_training_trainer(self, max_steps: int):
+        return pl.Trainer(
+            accelerator="auto",
+            devices="auto",
+            strategy=DDPStrategy(find_unused_parameters=True) if torch.cuda.device_count() > 1 else "auto",
+            logger=self.logger or False,
+            callbacks=self._build_trainer_callbacks(),
+            gradient_clip_val=self.cfg.training.optim.gradient_clip_val or 0.0,
+            val_check_interval=self.cfg.validation.val_every_n_step if self.cfg.validation.val_every_n_step else None,
+            limit_val_batches=self.cfg.validation.limit_batch,
+            check_val_every_n_epoch=self.cfg.validation.val_every_n_epoch if not self.cfg.validation.val_every_n_step else None,
+            accumulate_grad_batches=self.cfg.training.optim.accumulate_grad_batches or 1,
+            precision=self.cfg.training.precision or 32,
+            detect_anomaly=False,
+            num_sanity_val_steps=int(self.cfg.debug) if self.cfg.debug else 0,
+            max_epochs=self.cfg.training.max_epochs,
+            max_steps=max_steps,
+            max_time=self.cfg.training.max_time
+        )
+
+    def _curriculum_stages(self):
+        curriculum = self.cfg.training.get("curriculum", None)
+        if curriculum is None or not bool(curriculum.get("enabled", True)):
+            return None
+        stages = curriculum.get("stages", None)
+        return stages if stages is not None and len(stages) > 0 else None
+
+    def _merge_curriculum_updates(self, target_cfg: DictConfig, updates: DictConfig) -> None:
+        with open_dict(target_cfg):
+            for key, value in updates.items():
+                if key in target_cfg and isinstance(target_cfg[key], DictConfig) and isinstance(value, DictConfig):
+                    self._merge_curriculum_updates(target_cfg[key], value)
+                else:
+                    target_cfg[key] = value
+
+    def _apply_curriculum_stage(self, stage: DictConfig, stage_idx: int, until_step: int) -> None:
+        for cfg_name in ("dataset", "algorithm"):
+            updates = stage.get(cfg_name, None)
+            if updates is not None:
+                self._merge_curriculum_updates(self.root_cfg[cfg_name], updates)
+        if is_rank_zero:
+            print(cyan("Curriculum stage:"), f"{stage.get('name', stage_idx + 1)} until step {until_step}")
+
+    def _curriculum_until_step(self, stage: DictConfig, stage_idx: int) -> int:
+        if "until_step" not in stage:
+            raise ValueError(f"experiment.training.curriculum.stages[{stage_idx}] must define until_step")
+        return int(stage.until_step)
+
+    def _curriculum_checkpoint_path(self, stage: DictConfig, stage_idx: int, until_step: int) -> pathlib.Path:
+        checkpoint_dir = pathlib.Path(hydra.core.hydra_config.HydraConfig.get()["runtime"]["output_dir"]) / "checkpoints"
+        stage_name = str(stage.get("name", f"stage_{stage_idx + 1}"))
+        safe_name = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in stage_name)
+        return checkpoint_dir / f"curriculum_{stage_idx + 1:02d}_{safe_name}_step{until_step}.ckpt"
+
+    def _prepare_training_weights(self):
+        if self.auto_resuming:
+            return self.ckpt_path
+        if self.customized_load:
+            self._load_custom_checkpoint_for_task()
+
+        if self.zero_init_gate:
+            for name, para in self.algo.diffusion_model.named_parameters():
+                if 'r_adaLN_modulation' in name:
+                    para.requires_grad_(False)
+                    para[2*1024:3*1024] = 0
+                    para[5*1024:6*1024] = 0
+                    para.requires_grad_(True)
+
+        if self.only_tune_memory:
+            for name, para in self.algo.diffusion_model.named_parameters():
+                para.requires_grad_(False)
+                if 'r_' in name or 'pose_embedder' in name or 'pose_cond_mlp' in name or 'lora_' in name:
+                    para.requires_grad_(True)
+
+        return None if self.customized_load else self.ckpt_path
+
+    def _fit_training_curriculum(self, stages, ckpt_path) -> None:
+        for stage_idx, stage in enumerate(stages):
+            until_step = self._curriculum_until_step(stage, stage_idx)
+            self._apply_curriculum_stage(stage, stage_idx, until_step)
+            trainer = self._build_training_trainer(max_steps=until_step)
+            trainer.fit(
+                self.algo,
+                train_dataloaders=self._build_training_loader(),
+                val_dataloaders=self._build_validation_loader(),
+                ckpt_path=ckpt_path,
+            )
+            if stage_idx + 1 < len(stages):
+                ckpt_path = self._curriculum_checkpoint_path(stage, stage_idx, until_step)
+                ckpt_path.parent.mkdir(parents=True, exist_ok=True)
+                trainer.save_checkpoint(ckpt_path)
 
     def _build_training_loader(self) -> Optional[Union[TRAIN_DATALOADERS, pl.LightningDataModule]]:
         train_dataset = self._build_dataset("training")
@@ -453,84 +555,20 @@ class BaseLightningExperiment(BaseExperiment):
         if self.cfg.training.compile:
             self.algo = torch.compile(self.algo)
 
-        callbacks = []
-        if self.logger:
-            callbacks.append(LearningRateMonitor("step", True))
-        if "checkpointing" in self.cfg.training:
-            callbacks.append(
-                ModelCheckpoint(
-                    pathlib.Path(hydra.core.hydra_config.HydraConfig.get()["runtime"]["output_dir"]) / "checkpoints",
-                    filename="epoch{epoch}_step{step}",
-                    auto_insert_metric_name=False,
-                    **self.cfg.training.checkpointing,
-                )
-            )
-
-        trainer = pl.Trainer(
-            accelerator="auto",
-            devices="auto",
-            strategy=DDPStrategy(find_unused_parameters=True) if torch.cuda.device_count() > 1 else "auto",
-            logger=self.logger or False,
-            callbacks=callbacks,
-            gradient_clip_val=self.cfg.training.optim.gradient_clip_val or 0.0,
-            val_check_interval=self.cfg.validation.val_every_n_step if self.cfg.validation.val_every_n_step else None,
-            limit_val_batches=self.cfg.validation.limit_batch,
-            check_val_every_n_epoch=self.cfg.validation.val_every_n_epoch if not self.cfg.validation.val_every_n_step else None,
-            accumulate_grad_batches=self.cfg.training.optim.accumulate_grad_batches or 1,
-            precision=self.cfg.training.precision or 32,
-            detect_anomaly=False,
-            num_sanity_val_steps=int(self.cfg.debug) if self.cfg.debug else 0,
-            max_epochs=self.cfg.training.max_epochs,
-            max_steps=self.cfg.training.max_steps,
-            max_time=self.cfg.training.max_time
+        ckpt_path = self._prepare_training_weights()
+        curriculum_stages = self._curriculum_stages()
+        if curriculum_stages is not None:
+            self._fit_training_curriculum(curriculum_stages, ckpt_path)
+            return
+
+        trainer = self._build_training_trainer(max_steps=self.cfg.training.max_steps)
+        trainer.fit(
+            self.algo,
+            train_dataloaders=self._build_training_loader(),
+            val_dataloaders=self._build_validation_loader(),
+            ckpt_path=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():
-                    if 'r_adaLN_modulation' in name:
-                        para.requires_grad_(False)
-                        para[2*1024:3*1024] = 0
-                        para[5*1024:6*1024] = 0
-                        para.requires_grad_(True)
-
-            if self.only_tune_memory:
-                for name, para in self.algo.diffusion_model.named_parameters():
-                    para.requires_grad_(False)
-                    if 'r_' in name or 'pose_embedder' in name or 'pose_cond_mlp' in name or 'lora_' in name:
-                        para.requires_grad_(True)
-                    
-            trainer.fit(
-                self.algo,
-                train_dataloaders=self._build_training_loader(),
-                val_dataloaders=self._build_validation_loader(),
-                ckpt_path=None,
-            )
-        else:
-
-            if self.only_tune_memory:
-                for name, para in self.algo.diffusion_model.named_parameters():
-                    para.requires_grad_(False)
-                    if 'r_' in name or 'pose_embedder' in name or 'pose_cond_mlp' in name or 'lora_' in name:
-                        para.requires_grad_(True)
-            
-            trainer.fit(
-                self.algo,
-                train_dataloaders=self._build_training_loader(),
-                val_dataloaders=self._build_validation_loader(),
-                ckpt_path=self.ckpt_path,
-            )
-
     def validation(self) -> None:
         """
         All validation happens here

main.py

@@ -29,7 +29,7 @@ WANDB_RUN_ID_FILE = ".wandb_run_id"
 
 
 def _process_rank() -> int:
-    return int(os.environ.get("RANK", 0))
+    return int(os.environ.get("RANK", os.environ.get("SLURM_PROCID", 0)))
 
 
 def _checkpoint_step(path: Path) -> int:

tests/test_resume_checkpoint_logic.py

@@ -1,4 +1,7 @@
+import importlib
 import unittest
+import os
+from pathlib import Path
 from unittest.mock import patch
 
 from omegaconf import OmegaConf
@@ -21,12 +24,14 @@ class DummyExperiment(exp_base.BaseLightningExperiment):
 
 
 class DummyTrainer:
+    init_calls = []
     fit_calls = []
     validate_calls = []
     test_calls = []
+    save_checkpoint_calls = []
 
     def __init__(self, *args, **kwargs):
-        pass
+        self.init_calls.append({"args": args, "kwargs": kwargs})
 
     def fit(self, *args, **kwargs):
         self.fit_calls.append({"args": args, "kwargs": kwargs})
@@ -37,6 +42,9 @@ class DummyTrainer:
     def test(self, *args, **kwargs):
         self.test_calls.append({"args": args, "kwargs": kwargs})
 
+    def save_checkpoint(self, path):
+        self.save_checkpoint_calls.append(path)
+
 
 def _root_cfg(auto_resuming=True, algorithm_name="dummy", zero_init_gate=False, only_tune_memory=False):
     return OmegaConf.create(
@@ -120,6 +128,62 @@ class ResumeCheckpointLogicTests(unittest.TestCase):
 
         self.assertEqual(DummyTrainer.fit_calls[0]["kwargs"]["ckpt_path"], "/tmp/last.ckpt")
 
+    def test_curriculum_updates_stage_config_and_checkpoint_handoff(self):
+        DummyTrainer.init_calls = []
+        DummyTrainer.fit_calls = []
+        DummyTrainer.save_checkpoint_calls = []
+        cfg = _root_cfg(auto_resuming=True)
+        cfg.experiment.training.curriculum = {
+            "enabled": True,
+            "stages": [
+                {
+                    "name": "near",
+                    "until_step": 3,
+                    "dataset": {"wo_updown": True, "memory_selection": {"pose_similarity_radius": 2.0}},
+                    "algorithm": {"memory_selection": {"pose_similarity_radius": 2.0}},
+                },
+                {
+                    "name": "full",
+                    "until_step": 5,
+                    "dataset": {"wo_updown": False, "memory_selection": {"pose_similarity_radius": 8.0}},
+                    "algorithm": {"memory_selection": {"pose_similarity_radius": 8.0}},
+                },
+            ],
+        }
+        experiment = DummyExperiment(cfg, logger=None, ckpt_path="/tmp/last.ckpt")
+        experiment.algo = object()
+        seen_dataset = []
+
+        def build_training_loader():
+            seen_dataset.append(
+                (
+                    bool(experiment.root_cfg.dataset.wo_updown),
+                    float(experiment.root_cfg.dataset.memory_selection.pose_similarity_radius),
+                )
+            )
+            return None
+
+        experiment._build_training_loader = build_training_loader
+
+        with patch.object(exp_base.pl, "Trainer", DummyTrainer), patch.object(
+            DummyExperiment,
+            "_curriculum_checkpoint_path",
+            return_value=Path("/tmp/curriculum_stage.ckpt"),
+        ):
+            experiment.training()
+
+        self.assertEqual(seen_dataset, [(True, 2.0), (False, 8.0)])
+        self.assertEqual(
+            [call["kwargs"]["ckpt_path"] for call in DummyTrainer.fit_calls],
+            ["/tmp/last.ckpt", Path("/tmp/curriculum_stage.ckpt")],
+        )
+        self.assertEqual(
+            [call["kwargs"]["max_steps"] for call in DummyTrainer.init_calls],
+            [3, 5],
+        )
+        self.assertEqual(DummyTrainer.save_checkpoint_calls, [Path("/tmp/curriculum_stage.ckpt")])
+        self.assertEqual(float(experiment.root_cfg.algorithm.memory_selection.pose_similarity_radius), 8.0)
+
     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")
@@ -142,6 +206,22 @@ class ResumeCheckpointLogicTests(unittest.TestCase):
         load_mock.assert_called_once_with(algo=experiment.algo, checkpoint_path="/tmp/trained.ckpt")
         self.assertIsNone(DummyTrainer.test_calls[0]["kwargs"]["ckpt_path"])
 
+    def test_pre_lightning_rank_uses_slurm_procid(self):
+        import main
+        import utils.distributed_utils as distributed_utils
+
+        old_env = os.environ.copy()
+        try:
+            os.environ.pop("RANK", None)
+            os.environ["SLURM_PROCID"] = "1"
+            self.assertEqual(main._process_rank(), 1)
+            reloaded = importlib.reload(distributed_utils)
+            self.assertFalse(reloaded.is_rank_zero)
+        finally:
+            os.environ.clear()
+            os.environ.update(old_env)
+            importlib.reload(distributed_utils)
+
 
 if __name__ == "__main__":
     unittest.main()

utils/distributed_utils.py

@@ -1,3 +1,8 @@
 import os
 
-is_rank_zero = int(os.environ.get("RANK", 0)) == 0
+
+def process_rank() -> int:
+    return int(os.environ.get("RANK", os.environ.get("SLURM_PROCID", 0)))
+
+
+is_rank_zero = process_rank() == 0