From ecb2bfbe5e1fc0577fcb490658eb72de57247da6 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 2 Sep 2026 09:22:39 -0600 Subject: [PATCH] Fixed issue where some unmanaged modules could be lefs on cpu --- toolkit/memory_management/manager.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index 3a89d16..6fd195b 100644 --- a/toolkit/memory_management/manager.py +++ b/toolkit/memory_management/manager.py @@ -155,8 +155,14 @@ class MemoryManager: except Exception: pass - # count ignore modules as processed + # count ignore modules as processed — including their whole subtree: an + # ignored module stays resident, so none of its children may get a + # layer manager (a managed child would pin its weight back to cpu + # behind the parent's back) modules_processed = [x for x in ignore_modules] + for im in ignore_modules: + if isinstance(im, torch.nn.Module): + modules_processed.extend(im.modules()) # weights tied to an embedding (lm_head <-> embed_tokens) must not be # managed: pinning the linear's weight to cpu strands the (unmanaged)