Reworked freeing memory manager for removing text encoder completly when not needed.

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

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