Add DeMemWM optimizer LR groups and warmup

algorithms/dememwm/df_base.py

@@ -60,14 +60,17 @@ class DiffusionForcingBase(BasePytorchAlgo):
         return optimizer_dynamics
 
     def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
-        # update params
+        self._apply_optimizer_lrs(optimizer, self.trainer.global_step)
         optimizer.step(closure=optimizer_closure)
-
-        # manually warm up lr without a scheduler
-        if self.trainer.global_step < self.cfg.warmup_steps:
-            lr_scale = min(1.0, float(self.trainer.global_step + 1) / self.cfg.warmup_steps)
-            for pg in optimizer.param_groups:
-                pg["lr"] = lr_scale * self.cfg.lr
+        self._apply_optimizer_lrs(optimizer, self.trainer.global_step + 1)
+
+    def _apply_optimizer_lrs(self, optimizer, step: int) -> None:
+        warmup_steps = int(getattr(self.cfg, "warmup_steps", 0) or 0)
+        for pg in optimizer.param_groups:
+            warmup_start_step = int(pg.get("warmup_start_step", 0) or 0)
+            warmup_progress = max(0, int(step) - warmup_start_step + 1)
+            lr_scale = 1.0 if warmup_steps <= 0 else min(1.0, float(warmup_progress) / warmup_steps)
+            pg["lr"] = float(pg.get("target_lr", self.cfg.lr)) * lr_scale
 
     def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
         xs, conditions, masks = self._preprocess_batch(batch)

algorithms/dememwm/df_video.py

@@ -103,6 +103,60 @@ def _dememwm_optimizer_parameters(diffusion_model, vae, trainability):
         yield from vae.parameters()
 
 
+def _dememwm_group_target_lr(group_name: str, trainability, default_lr, global_step: int) -> float:
+    lr_cfg = _cfg_get(trainability, "lr", {})
+    if group_name == "memory_modules":
+        return float(_cfg_get(lr_cfg, "memory_modules", default_lr))
+    if group_name == "base_dit":
+        if not _dememwm_full_dit_active(trainability, global_step):
+            return 0.0
+        return float(_cfg_get(lr_cfg, "base_dit", default_lr))
+    return float(default_lr)
+
+
+def _dememwm_group_warmup_start_step(group_name: str, trainability) -> int:
+    if group_name != "base_dit":
+        return 0
+    start_step = _cfg_get(trainability, "full_dit_start_step", 0)
+    return 0 if start_step is None else int(start_step)
+
+
+def _dememwm_optimizer_parameter_groups(diffusion_model, vae, trainability, default_lr, global_step: int):
+    include_full_dit = bool(_cfg_get(trainability, "train_full_dit", False))
+    grouped = {"memory_modules": [], "base_dit": []}
+    for name, param in diffusion_model.named_parameters():
+        if _is_dememwm_memory_trainable_parameter(name, trainability):
+            grouped["memory_modules"].append(param)
+        elif include_full_dit:
+            grouped["base_dit"].append(param)
+
+    param_groups = []
+    for group_name, params in grouped.items():
+        if params:
+            target_lr = _dememwm_group_target_lr(group_name, trainability, default_lr, global_step)
+            param_groups.append({
+                "params": params,
+                "lr": target_lr,
+                "target_lr": target_lr,
+                "warmup_start_step": _dememwm_group_warmup_start_step(group_name, trainability),
+                "name": group_name,
+            })
+
+    if vae is not None and not bool(_cfg_get(trainability, "freeze_vae", True)):
+        target_lr = float(_cfg_get(_cfg_get(trainability, "lr", {}), "vae", default_lr))
+        param_groups.append({"params": tuple(vae.parameters()), "lr": target_lr, "target_lr": target_lr, "name": "vae"})
+
+    return param_groups
+
+
+def _apply_dememwm_optimizer_group_lrs(optimizer, trainability, default_lr, global_step: int) -> None:
+    for param_group in optimizer.param_groups:
+        group_name = param_group.get("name")
+        if group_name in {"memory_modules", "base_dit"}:
+            param_group["target_lr"] = _dememwm_group_target_lr(group_name, trainability, default_lr, global_step)
+            param_group["warmup_start_step"] = _dememwm_group_warmup_start_step(group_name, trainability)
+
+
 def _derive_memory_condition_length(cfg) -> int:
     memory_cfg = _cfg_get(cfg, "memory_selection")
     if memory_cfg is None:
@@ -491,15 +545,24 @@ class DeMemWMMinecraft(DiffusionForcingBase):
     def configure_optimizers(self):
         trainability = _trainability_cfg(self.cfg)
         self._apply_trainability()
-        params = tuple(_dememwm_optimizer_parameters(self.diffusion_model, getattr(self, "vae", None), trainability))
-        if not params:
+        param_groups = _dememwm_optimizer_parameter_groups(
+            self.diffusion_model,
+            getattr(self, "vae", None),
+            trainability,
+            self.cfg.lr,
+            self._global_step_for_trainability(),
+        )
+        if not param_groups:
             raise ValueError("DeMemWM trainability selected no optimizer parameters")
         return torch.optim.AdamW(
-            params, lr=self.cfg.lr, weight_decay=self.cfg.weight_decay, betas=self.cfg.optimizer_beta
+            param_groups, weight_decay=self.cfg.weight_decay, betas=self.cfg.optimizer_beta
         )
 
     def on_train_batch_start(self, batch, batch_idx, dataloader_idx=0) -> None:
         self._apply_trainability()
+        trainability = _trainability_cfg(self.cfg)
+        for optimizer in getattr(getattr(self, "trainer", None), "optimizers", []) or []:
+            _apply_dememwm_optimizer_group_lrs(optimizer, trainability, self.cfg.lr, self._global_step_for_trainability())
 
     def _generate_noise_levels(self, xs: torch.Tensor, masks = None) -> torch.Tensor:
         """

configurations/algorithm/dememwm_base.yaml

@@ -23,6 +23,9 @@ trainability:
   geometry_projections: true
   train_full_dit: false
   full_dit_start_step: null
+  lr:
+    memory_modules: 8.0e-5
+    base_dit: 2.0e-5
 
 n_tokens: ${dataset.n_frames}
 action_cond_dim: ${dataset.action_cond_dim}

tests/test_dememwm_latent_dataset.py

@@ -383,8 +383,8 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
             def reweight_loss(self, loss, weight=None):
                 raise AssertionError("target mask path should reduce the latent dict loss directly")
 
-            def log(self, name, value):
-                self.logged.append((name, value))
+            def log(self, name, value, **kwargs):
+                self.logged.append((name, value, kwargs))
 
         batch = {
             "latents": torch.arange(5, dtype=torch.float32).view(1, 5, 1, 1, 1),
@@ -492,9 +492,12 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
     def test_trainability_controls_memory_groups_full_dit_ramp_and_vae_freeze(self):
         from algorithms.dememwm.df_video import (
             _apply_dememwm_trainability,
+            _apply_dememwm_optimizer_group_lrs,
+            _dememwm_optimizer_parameter_groups,
             _dememwm_optimizer_parameters,
             _trainability_cfg,
         )
+        from algorithms.dememwm.df_base import DiffusionForcingBase
 
         class TinyReferenceAttention(nn.Module):
             def __init__(self):
@@ -520,7 +523,15 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
 
         diffusion = TinyDiffusion()
         vae = nn.Linear(2, 2, bias=False)
-        cfg = OmegaConf.create({"trainability": {"freeze_vae": True, "train_full_dit": True, "full_dit_start_step": 10, "geometry_projections": False}})
+        cfg = OmegaConf.create({
+            "trainability": {
+                "freeze_vae": True,
+                "train_full_dit": True,
+                "full_dit_start_step": 10,
+                "geometry_projections": False,
+                "lr": {"memory_modules": 8.0e-5, "base_dit": 2.0e-5},
+            }
+        })
         trainability = _trainability_cfg(cfg)
 
         _apply_dememwm_trainability(diffusion, vae, trainability, global_step=0)
@@ -536,6 +547,34 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
 
         opt_ids = {id(param) for param in _dememwm_optimizer_parameters(diffusion, vae, trainability)}
         self.assertEqual(opt_ids, {id(param) for param in diffusion.parameters()})
+        param_groups = _dememwm_optimizer_parameter_groups(diffusion, vae, trainability, default_lr=2.0e-5, global_step=0)
+        lr_by_name = {group["name"]: group["target_lr"] for group in param_groups}
+        warmup_start_by_name = {group["name"]: group.get("warmup_start_step", 0) for group in param_groups}
+        self.assertEqual(set(lr_by_name), {"memory_modules", "base_dit"})
+        self.assertAlmostEqual(lr_by_name["memory_modules"], 8.0e-5)
+        self.assertEqual(lr_by_name["base_dit"], 0.0)
+        self.assertEqual(warmup_start_by_name["memory_modules"], 0)
+        self.assertEqual(warmup_start_by_name["base_dit"], 10)
+
+        optimizer = torch.optim.AdamW(param_groups)
+        lr_owner = object.__new__(DiffusionForcingBase)
+        lr_owner.cfg = OmegaConf.create({"lr": 2.0e-5, "warmup_steps": 10})
+        lr_owner._apply_optimizer_lrs(optimizer, step=0)
+        lr_by_name = {group["name"]: group["lr"] for group in optimizer.param_groups}
+        self.assertAlmostEqual(lr_by_name["memory_modules"], 8.0e-6)
+        self.assertEqual(lr_by_name["base_dit"], 0.0)
+
+        _apply_dememwm_optimizer_group_lrs(optimizer, trainability, default_lr=2.0e-5, global_step=10)
+        lr_owner._apply_optimizer_lrs(optimizer, step=10)
+        lr_by_name = {group["name"]: group["lr"] for group in optimizer.param_groups}
+        self.assertAlmostEqual(lr_by_name["memory_modules"], 8.0e-5)
+        self.assertAlmostEqual(lr_by_name["base_dit"], 2.0e-6)
+
+        _apply_dememwm_optimizer_group_lrs(optimizer, trainability, default_lr=2.0e-5, global_step=19)
+        lr_owner._apply_optimizer_lrs(optimizer, step=19)
+        lr_by_name = {group["name"]: group["lr"] for group in optimizer.param_groups}
+        self.assertAlmostEqual(lr_by_name["memory_modules"], 8.0e-5)
+        self.assertAlmostEqual(lr_by_name["base_dit"], 2.0e-5)
 
         _apply_dememwm_trainability(diffusion, vae, trainability, global_step=10)
         self.assertTrue(all(param.requires_grad for param in diffusion.parameters()))