pl_model: RDT-style weight EMA, validated in place of raw weights

egomimic/pl_utils/ema.py

@@ -0,0 +1,91 @@
+"""Exponential moving average of a model's trainable weights."""
+
+from __future__ import annotations
+
+import torch
+
+
+class WeightEMA:
+    """EMA shadow of ``named_params`` with RDT's decay warmup.
+
+    Decay after ``n`` updates is ``1 - (1 + (n - 1) / inv_gamma) ** -power``,
+    capped at ``max_decay`` (0 for the first two updates), as in RDT's
+    ``models/ema_model.py``. With power 0.75 it passes 0.999 near 10k updates
+    and reaches 0.9999 near 215k, so the average window grows with the run.
+    """
+
+    def __init__(
+        self,
+        named_params,
+        power: float = 0.75,
+        max_decay: float = 0.9999,
+        inv_gamma: float = 1.0,
+    ):
+        named_params = list(named_params)
+        self.names = [name for name, _ in named_params]
+        self.params = [param for _, param in named_params]
+        self.shadow = [param.detach().float().clone() for param in self.params]
+        self.power = power
+        self.max_decay = max_decay
+        self.inv_gamma = inv_gamma
+        self.num_updates = 0
+        self._backup = None
+
+    def decay(self) -> float:
+        step = self.num_updates - 1
+        if step <= 0:
+            return 0.0
+        return min(1 - (1 + step / self.inv_gamma) ** -self.power, self.max_decay)
+
+    @torch.no_grad()
+    def update(self) -> float:
+        decay = self.decay()
+        current = [p.detach().to(s.dtype) for p, s in zip(self.params, self.shadow)]
+        torch._foreach_lerp_(self.shadow, current, 1.0 - decay)
+        self.num_updates += 1
+        return decay
+
+    @property
+    def swapped(self) -> bool:
+        return self._backup is not None
+
+    @torch.no_grad()
+    def swap_in(self) -> None:
+        """Load the EMA weights into the model, keeping the raw ones aside."""
+        if self.swapped:
+            return
+        # On CPU: a second GPU copy of a 1B model can OOM validation.
+        self._backup = [p.detach().to("cpu", copy=True) for p in self.params]
+        for param, shadow in zip(self.params, self.shadow):
+            param.copy_(shadow)
+
+    @torch.no_grad()
+    def swap_out(self) -> None:
+        if not self.swapped:
+            return
+        for param, raw in zip(self.params, self._backup):
+            param.copy_(raw)
+        self._backup = None
+
+    def raw_by_id(self) -> dict:
+        """``id(param) -> raw weights`` while swapped in, else empty."""
+        if not self.swapped:
+            return {}
+        return {id(p): raw for p, raw in zip(self.params, self._backup)}
+
+    def state_dict(self) -> dict:
+        return {
+            "shadow": dict(zip(self.names, self.shadow)),
+            "num_updates": self.num_updates,
+        }
+
+    @torch.no_grad()
+    def load_state_dict(self, state: dict) -> None:
+        missing = [name for name in self.names if name not in state["shadow"]]
+        if missing:
+            raise KeyError(
+                f"EMA state is missing {len(missing)} weight(s), e.g. {missing[:5]}"
+            )
+        for name, shadow in zip(self.names, self.shadow):
+            shadow.copy_(state["shadow"][name])
+        self.num_updates = int(state["num_updates"])

egomimic/pl_utils/pl_model.py

@@ -10,6 +10,7 @@
 from omegaconf import DictConfig, OmegaConf
 
 import egomimic.utils.tensor_utils as TensorUtils
+from egomimic.pl_utils.ema import WeightEMA
 from egomimic.pl_utils.pl_data_utils import head_of_loader, video_loader_name
 from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset
 
@@ -88,6 +89,20 @@ def __init__(
         # MultiDataModuleWrapper.val_loader_names(); None = [valid, train_viz].
         self.val_loader_names = None
 
+        ema_cfg = (
+            self._as_config(config_tree).model.get("ema")
+            if config_tree is not None
+            else None
+        )
+        self._ema_cfg = (
+            OmegaConf.to_container(ema_cfg)
+            if ema_cfg is not None and ema_cfg.get("enabled", True)
+            else None
+        )
+        self.ema = None
+        self._ema_loaded = None
+        self._ema_last_step = None
+
     @staticmethod
     def _as_config(cfg):
         if cfg is None:
@@ -228,6 +243,48 @@ def on_after_backward(self):
         for k, v in info.items():
             self.log("Train/" + k, v, on_step=False, on_epoch=True, sync_dist=True)
 
+    def _ensure_ema(self) -> None:
+        if self._ema_cfg is None or self.ema is not None:
+            return
+        opts = {k: v for k, v in self._ema_cfg.items() if k != "enabled"}
+        self.ema = WeightEMA(
+            ((n, p) for n, p in self.named_parameters() if p.requires_grad), **opts
+        )
+        if self._ema_loaded is not None:
+            self.ema.load_state_dict(self._ema_loaded)
+            self._ema_loaded = None
+
+    def on_train_batch_end(self, outputs, batch, batch_idx):
+        # Once per optimizer step, not per batch, so gradient accumulation is fine.
+        if self._ema_cfg is None or self.global_step == self._ema_last_step:
+            return
+        self._ema_last_step = self.global_step
+        self._ensure_ema()
+        decay = self.ema.update()
+        self.log("Train/ema_decay", decay, on_step=False, on_epoch=True)
+
+    def on_save_checkpoint(self, checkpoint):
+        if self.ema is None:
+            if self._ema_loaded is not None:
+                checkpoint["ema"] = self._ema_loaded
+            return
+        checkpoint["ema"] = self.ema.state_dict()
+        # ModelCheckpoint can save mid-validation, while the EMA weights are
+        # swapped in; the state_dict must still hold the raw ones.
+        raw = self.ema.raw_by_id()
+        if raw:
+            for key, value in self.state_dict(keep_vars=True).items():
+                if id(value) in raw:
+                    checkpoint["state_dict"][key] = raw[id(value)]
+
+    def on_load_checkpoint(self, checkpoint):
+        if self._ema_cfg is None or "ema" not in checkpoint:
+            return
+        if self.ema is not None:
+            self.ema.load_state_dict(checkpoint["ema"])
+        else:
+            self._ema_loaded = checkpoint["ema"]
+
     def _val_heads(self) -> dict:
         return {
             "valid": self.evaluator,
@@ -244,6 +301,11 @@ def on_validation_start(self):
         if callable(reset_eval_rng):
             reset_eval_rng()
 
+        # Validate the EMA weights once there are any, trained or loaded.
+        if self.ema is not None or self._ema_loaded is not None:
+            self._ensure_ema()
+            self.ema.swap_in()
+
         heads = [h for h in self._val_heads().values() if h is not None]
         if not heads:
             return
@@ -307,6 +369,8 @@ def validation_step(self, batch, batch_idx, dataloader_idx=0):
 
     def on_validation_end(self):
         print(f"[ON_VALIDATION_END] rank={self.global_rank}", flush=True)
+        if self.ema is not None:
+            self.ema.swap_out()
         for head in self._val_heads().values():
             if head is not None:
                 head.on_validation_end()

egomimic/utils/checkpoint_utils.py

@@ -39,4 +39,9 @@ def load_checkpoint_weights(
             f"the model needs, e.g. {result.missing_keys[:5]}; the model would run "
             "on random weights. Is this checkpoint from the same model config?"
         )
+    # Lightning's restore calls this hook; a bare state_dict load does not, and
+    # ModelWrapper takes its EMA weights from it.
+    on_load = getattr(model, "on_load_checkpoint", None)
+    if callable(on_load):
+        on_load(checkpoint)
     log.info("Loaded weights from %s", ckpt_path)

tests/unit/test_ema.py

@@ -0,0 +1,114 @@
+import torch
+import torch.nn as nn
+from omegaconf import OmegaConf
+
+from egomimic.pl_utils.ema import WeightEMA
+from egomimic.pl_utils.pl_model import ModelWrapper
+
+
+def _ema(module, **kw):
+    return WeightEMA(module.named_parameters(), **kw)
+
+
+def test_decay_matches_rdt_schedule():
+    ema = _ema(nn.Linear(2, 2))
+    decays = []
+    for _ in range(5):
+        decays.append(ema.decay())
+        ema.num_updates += 1
+    assert decays[:2] == [0.0, 0.0]
+    assert decays[2] == 1 - 2**-0.75
+    ema.num_updates = 10**7
+    assert ema.decay() == 0.9999
+
+
+def test_update_tracks_weights():
+    lin = nn.Linear(2, 2)
+    ema = _ema(lin)
+    ema.update()
+    with torch.no_grad():
+        lin.weight.fill_(1.0)
+    ema.update()
+    assert torch.equal(ema.shadow[0], lin.weight.detach())
+    with torch.no_grad():
+        lin.weight.fill_(3.0)
+    decay = ema.update()
+    assert torch.allclose(ema.shadow[0], torch.full((2, 2), decay + 3 * (1 - decay)))
+
+
+def test_swap_round_trip():
+    lin = nn.Linear(2, 2)
+    raw = lin.weight.detach().clone()
+    ema = _ema(lin)
+    ema.shadow[0].fill_(7.0)
+    ema.swap_in()
+    assert torch.equal(lin.weight.detach(), torch.full((2, 2), 7.0))
+    assert torch.equal(ema.raw_by_id()[id(lin.weight)], raw)
+    ema.swap_out()
+    assert torch.equal(lin.weight.detach(), raw)
+    assert not ema.swapped
+
+
+def test_state_round_trip():
+    src = _ema(nn.Linear(2, 2))
+    src.shadow[0].fill_(5.0)
+    src.num_updates = 42
+    dst = _ema(nn.Linear(2, 2))
+    dst.load_state_dict(src.state_dict())
+    assert torch.equal(dst.shadow[0], src.shadow[0])
+    assert dst.num_updates == 42
+
+
+class _Algo:
+    def __init__(self):
+        self.nets = nn.ModuleDict({"policy": nn.Linear(2, 2)})
+
+
+def _wrapper(monkeypatch, ema_cfg):
+    monkeypatch.setattr(ModelWrapper, "_instantiate_model", lambda *a, **k: _Algo())
+    model = {"robomimic_model": {}}
+    if ema_cfg is not None:
+        model["ema"] = ema_cfg
+    return ModelWrapper(
+        config_tree=OmegaConf.create({"model": model}), norm_stats_state=None
+    )
+
+
+def test_ema_is_off_unless_asked(monkeypatch):
+    assert _wrapper(monkeypatch, None)._ema_cfg is None
+    assert _wrapper(monkeypatch, {"enabled": False})._ema_cfg is None
+
+
+def test_checkpoint_saved_mid_validation_keeps_raw_weights(monkeypatch):
+    wrapper = _wrapper(monkeypatch, {"enabled": True})
+    wrapper._ensure_ema()
+    raw = wrapper.nets["policy"].weight.detach().clone()
+    wrapper.ema.shadow[0].fill_(9.0)
+    wrapper.ema.swap_in()
+    checkpoint = {"state_dict": wrapper.state_dict()}
+    wrapper.on_save_checkpoint(checkpoint)
+    keys = [k for k in checkpoint["state_dict"] if k.endswith("policy.weight")]
+    assert keys
+    for key in keys:
+        assert torch.equal(checkpoint["state_dict"][key], raw)
+    assert torch.equal(
+        next(iter(checkpoint["ema"]["shadow"].values())), torch.full((2, 2), 9.0)
+    )
+
+
+def test_loaded_ema_survives_until_first_use(monkeypatch):
+    src = _wrapper(monkeypatch, {"enabled": True})
+    src._ensure_ema()
+    src.ema.shadow[0].fill_(4.0)
+    src.ema.num_updates = 11
+    checkpoint = {"state_dict": src.state_dict()}
+    src.on_save_checkpoint(checkpoint)
+
+    dst = _wrapper(monkeypatch, {"enabled": True})
+    dst.on_load_checkpoint(checkpoint)
+    resaved = {"state_dict": dst.state_dict()}
+    dst.on_save_checkpoint(resaved)
+    assert resaved["ema"]["num_updates"] == 11
+    dst._ensure_ema()
+    assert dst.ema.num_updates == 11
+    assert torch.equal(dst.ema.shadow[0], torch.full((2, 2), 4.0))