From decb81c5f2e596579d3eb4e3963921c6e457a651 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Mon, 28 Sep 2026 18:36:31 +0800 Subject: [PATCH] =?UTF-8?q?fix(distill):=20engram-resilient=20resume=20?= =?UTF-8?q?=E2=80=94=20strict=3DFalse=20with=20an=20allowlist?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The engram run's crash-loop: a pre-engram checkpoint (or any older save) lacks _engram.* keys, and strict load_state_dict crashed every restart into a retry cascade. The module is zero-init (identity), so resuming with a fresh table is safe; the allowlist rejects every other missing/unexpected key as the real mismatch it would be. --- src/gpu/modal_distill.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index e81c1f8..148194a 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -1149,9 +1149,18 @@ def _usable(ck: Path) -> bool: key=lambda p: int(p.name.split("-")[1]), ) if ckpts: - student.load_state_dict( - torch.load(ckpts[-1] / "student.pt", map_location="cpu", weights_only=True) - ) + sd = torch.load(ckpts[-1] / "student.pt", map_location="cpu", weights_only=True) + missing, unexpected = student.load_state_dict(sd, strict=False) + # engram params may be absent in older checkpoints: the module + # is zero-init (identity), so a fresh table is a safe resume. + # Anything else missing or unexpected is a real mismatch. + bad_missing = [k for k in missing if not k.startswith("_engram.")] + if bad_missing or unexpected: + raise RuntimeError( + f"resume mismatch: missing={bad_missing} unexpected={list(unexpected)}" + ) + if missing: + print(f"[{spec_id}] resume with fresh engram: {missing}", flush=True) optimizer.load_state_dict( torch.load(ckpts[-1] / "optim.pt", map_location="cpu", weights_only=True) )