diff --git a/egomimic/pl_utils/ema.py b/egomimic/pl_utils/ema.py new file mode 100644 index 0000000000..4637262e28 --- /dev/null +++ b/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"]) diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index c98d1d5f86..2790dc7051 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/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() diff --git a/egomimic/utils/checkpoint_utils.py b/egomimic/utils/checkpoint_utils.py index 9d70bbe739..ed9940f7f1 100644 --- a/egomimic/utils/checkpoint_utils.py +++ b/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) diff --git a/tests/unit/test_ema.py b/tests/unit/test_ema.py new file mode 100644 index 0000000000..bc95fb4921 --- /dev/null +++ b/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))