From 682b27c6eec4d302bc3300f47a44812bc3550fe4 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 9 Aug 2026 07:08:24 -0600 Subject: [PATCH] Reworked freeing memory manager for removing text encoder completly when not needed. --- toolkit/memory_management/manager.py | 83 ++++++++++++++++++++++++++++ toolkit/unloader.py | 27 +++------ 2 files changed, 92 insertions(+), 18 deletions(-) diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index c048e940..7e8267a8 100644 --- a/toolkit/memory_management/manager.py +++ b/toolkit/memory_management/manager.py @@ -1,3 +1,6 @@ +import ctypes +import gc + import torch from .manager_modules import ( LinearLayerMemoryManager, @@ -252,3 +255,83 @@ class MemoryManager: del _DEVICE_STATE[key] torch.cuda.empty_cache() + + @classmethod + def free(cls, module: torch.nn.Module): + """ + Detach memory management (if attached) and destroy the module's weights + by moving them to the meta device. + + Unlike detach(), nothing is staged back to CPU first: to('meta') frees + each storage from wherever it currently lives, so no transient host + allocation is made for data that is about to be discarded, and pinned + tensors are freed without the clone that unpinning requires. Freed + pinned storages land in torch's caching host allocator, not the OS; + call release_cached_memory() afterward to get the RSS back. + """ + if hasattr(module, "_memory_manager"): + if hasattr(module, "_mm_to"): + module.to = module._mm_to + del module._mm_to + + del module._memory_manager + + for child in module.modules(): + lmm = getattr(child, "_layer_memory_manager", None) + if lmm is None: + continue + + original_forward = getattr(lmm, "_original_forward", None) + if original_forward is not None: + if hasattr(child, "ara_lora_ref"): + ara = child.ara_lora_ref() + if ara is not None: + ara.org_forward = original_forward + else: + child.forward = original_forward + + del child._layer_memory_manager + if hasattr(child, "_memory_management_device"): + del child._memory_management_device + if hasattr(child, "_is_memory_managed"): + del child._is_memory_managed + + keys_to_delete = [ + dev for dev in _DEVICE_STATE + if isinstance(dev, torch.device) and dev.type == "cuda" + ] + for key in keys_to_delete: + del _DEVICE_STATE[key] + + # bypass any overridden/nopped-out .to() so the storages are actually freed + torch.nn.Module.to(module, "meta") + torch.cuda.empty_cache() + + @classmethod + def release_cached_memory(cls): + """ + Return freed memory to the OS. Freed pinned-host storages sit in + torch's caching host allocator and freed pageable memory sits in + glibc's arenas; neither shows up as reclaimed RSS without an + explicit flush. Call after free()ing a large module. + """ + gc.collect() + torch.cuda.empty_cache() + + # torch's pinned-host cache; private API, name varies by torch version + for fn_name in ("_accelerator_emptyHostCache", "_host_emptyCache"): + fn = getattr(torch._C, fn_name, None) + if fn is not None: + try: + fn() + break + except Exception: + pass + + # glibc keeps freed arenas mapped; CDLL(None) resolves malloc_trim in + # the running process where glibc is present and fails cleanly on + # macOS/musl + try: + ctypes.CDLL(None).malloc_trim(0) + except Exception: + pass diff --git a/toolkit/unloader.py b/toolkit/unloader.py index 7b1670dc..3364755f 100644 --- a/toolkit/unloader.py +++ b/toolkit/unloader.py @@ -1,6 +1,4 @@ -import gc import torch -from toolkit.basic import flush from toolkit.memory_management import MemoryManager from typing import TYPE_CHECKING @@ -35,12 +33,6 @@ class FakeTextEncoder(torch.nn.Module): return self -def _detach_and_cpu(te: torch.nn.Module): - MemoryManager.detach(te) - # bypass any nopped-out .to() override and force an actual CPU move - torch.nn.Module.to(te, 'cpu') - - def unload_text_encoder(model: "BaseModel"): # unload the text encoder in a way that will work with all models and will not throw errors # we need to make it appear as a text encoder module without actually having one so all @@ -52,8 +44,8 @@ def unload_text_encoder(model: "BaseModel"): pipe = model.pipeline # the pipeline stores text encoders like text_encoder, text_encoder_2, text_encoder_3, etc. - if hasattr(pipe, "text_encoder"): - _detach_and_cpu(pipe.text_encoder) + if getattr(pipe, "text_encoder", None) is not None: + MemoryManager.free(pipe.text_encoder) te = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype) text_encoder_list.append(te) pipe.text_encoder = te @@ -61,20 +53,19 @@ def unload_text_encoder(model: "BaseModel"): i = 2 while hasattr(pipe, f"text_encoder_{i}"): real_te = getattr(pipe, f"text_encoder_{i}") - _detach_and_cpu(real_te) - te = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype) - text_encoder_list.append(te) - setattr(pipe, f"text_encoder_{i}", te) + if real_te is not None: + MemoryManager.free(real_te) + te = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype) + text_encoder_list.append(te) + setattr(pipe, f"text_encoder_{i}", te) i += 1 model.text_encoder = text_encoder_list else: # only has a single text encoder - _detach_and_cpu(model.text_encoder) + MemoryManager.free(model.text_encoder) model.text_encoder = FakeTextEncoder( device=model.device_torch, dtype=model.torch_dtype ) - torch.cuda.empty_cache() - gc.collect() - flush() + MemoryManager.release_cached_memory()