Fix text encoder offload bug when caching embeddings (#868)

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()

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()