From 7c1a76f33669b6cfd854a65ec80843dd3fb39bea Mon Sep 17 00:00:00 2001 From: PlagueKind Date: Sun, 7 Jun 2026 22:09:32 +0200 Subject: [PATCH] Fix text encoder offload bug when caching embeddings (#868) --- toolkit/memory_management/manager.py | 87 +++++++++++++++++++++++++--- toolkit/unloader.py | 24 ++++++-- 2 files changed, 100 insertions(+), 11 deletions(-) diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index a43d3f3..0785f06 100644 --- a/toolkit/memory_management/manager.py +++ b/toolkit/memory_management/manager.py @@ -1,5 +1,5 @@ import torch -from .manager_modules import LinearLayerMemoryManager, ConvLayerMemoryManager +from .manager_modules import LinearLayerMemoryManager, ConvLayerMemoryManager, _DEVICE_STATE import random LINEAR_MODULES = [ @@ -67,9 +67,9 @@ class MemoryManager: @classmethod def attach( - cls, - module: torch.nn.Module, - device: torch.device, + cls, + module: torch.nn.Module, + device: torch.device, offload_percent: float = 1.0, ignore_modules: list[torch.nn.Module] = [] ): @@ -86,7 +86,7 @@ class MemoryManager: # add ignore modules to unmanaged list for im in ignore_modules: module._memory_manager.unmanaged_modules.append(im) - + # count ignore modules as processed modules_processed = [x for x in ignore_modules] # attach to all modules @@ -113,7 +113,7 @@ class MemoryManager: ara = child_module.ara_lora_ref() if ara not in modules_processed: MemoryManager.attach( - ara, + ara, device, ) modules_processed.append(child_module) @@ -138,7 +138,7 @@ class MemoryManager: ara = child_module.ara_lora_ref() if ara not in modules_processed: MemoryManager.attach( - ara, + ara, device, ) modules_processed.append(ara) @@ -151,3 +151,76 @@ class MemoryManager: module._memory_manager.unmanaged_modules.append(child_module) else: continue + + @classmethod + def detach(cls, module: torch.nn.Module): + """ + Reverse of attach(). Moves unmanaged modules back to CPU, restores the + original .to() and forward methods on all child layers, unpins CPU weight + tensors, and clears the global CUDA device state. + + Call this before unloading/replacing a module that had attach() applied. + """ + if not hasattr(module, "_memory_manager"): + return + + for unmanaged in module._memory_manager.unmanaged_modules: + try: + if isinstance(unmanaged, torch.nn.Parameter): + unmanaged.data = unmanaged.data.to('cpu') + else: + unmanaged.to('cpu') + except Exception: + pass + + 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 + + for param_name in ("weight", "bias"): + param = getattr(child, param_name, None) + if param is None or not isinstance(param, torch.nn.Parameter): + continue + try: + if param.data.is_pinned(): + object.__setattr__( + child, + param_name, + torch.nn.Parameter( + param.data.clone(), + requires_grad=param.requires_grad, + ), + ) + except Exception: + pass + + 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] + + torch.cuda.empty_cache() diff --git a/toolkit/unloader.py b/toolkit/unloader.py index 31f75e1..7b1670d 100644 --- a/toolkit/unloader.py +++ b/toolkit/unloader.py @@ -1,5 +1,7 @@ +import gc import torch from toolkit.basic import flush +from toolkit.memory_management import MemoryManager from typing import TYPE_CHECKING @@ -24,15 +26,21 @@ class FakeTextEncoder(torch.nn.Module): @property def device(self): return self._device - + @property def dtype(self): return self._dtype - + def to(self, *args, **kwargs): 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 @@ -45,13 +53,15 @@ def unload_text_encoder(model: "BaseModel"): # 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) te = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype) text_encoder_list.append(te) - pipe.text_encoder.to('cpu') pipe.text_encoder = te 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) @@ -59,6 +69,12 @@ def unload_text_encoder(model: "BaseModel"): model.text_encoder = text_encoder_list else: # only has a single text encoder - model.text_encoder = FakeTextEncoder(device=model.device_torch, dtype=model.torch_dtype) + _detach_and_cpu(model.text_encoder) + model.text_encoder = FakeTextEncoder( + device=model.device_torch, + dtype=model.torch_dtype + ) + torch.cuda.empty_cache() + gc.collect() flush()