add delete_offload_module

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/compressed_tensors/utils/offload.py

@@ -77,6 +77,7 @@
     "align_modules",
     "align_module_device",
     "register_offload_module",
+    "delete_offload_module",
     "force_cpu_offload",
 ]
 
@@ -398,7 +399,6 @@ def align_modules(
         yield
 
 
-@check_accelerate(fallback=None)
 def register_offload_module(base: torch.nn.Module, name: str, module: torch.nn.Module):
     """
     Register a submodule with offloading if the parent module is offloaded
@@ -459,6 +459,20 @@ def register_offload_module(base: torch.nn.Module, name: str, module: torch.nn.M
     # online way, assume that all pointers are shared. This comes at no runtime cost
 
 
+def delete_offload_module(base: torch.nn.Module, name: str):
+    """
+    Delete a submodule from a model which may contain offloading
+    :param base: parent module to delete submodule from
+    :param name: name of submodule on parent
+    """
+    module: torch.nn.Module = getattr(base, name)
+
+    for param_name, _ in list(module.named_parameters()):
+        delete_offload_parameter(module, param_name)
+
+    delattr(base, name)
+
+
 @check_accelerate(fallback="error")
 def force_cpu_offload(
     module: torch.nn.Module, execution_device: torch.device

tests/test_utils/test_offload.py

@@ -16,6 +16,7 @@
 from compressed_tensors.utils import (
     align_module_device,
     align_modules,
+    delete_offload_module,
     delete_offload_parameter,
     disable_hf_hook,
     force_cpu_offload,
@@ -368,6 +369,32 @@ def test_register_offload_module(exec_device):
     child(torch.empty(2, device=exec_device))
 
 
+@requires_gpu
+@requires_accelerate()
+@pytest.mark.parametrize("exec_device", [torch.device("cpu"), torch.device("cuda")])
+def test_delete_offload_module(exec_device):
+    # no offloading
+    model = ExampleModel()
+    child = torch.nn.Linear(2, 3)
+    register_offload_module(model, "child", child)
+    register_offload_module(model.linear, "child", child)
+    delete_offload_module(model, "child")
+    delete_offload_module(model.linear, "child")
+    assert not child in model.children()
+    assert not child in model.linear.children()
+
+    # with offloading
+    model = ExampleModel()
+    child = torch.nn.Linear(2, 3)
+    force_cpu_offload(model, exec_device)
+    register_offload_module(model, "child", child)
+    register_offload_module(model.linear, "child", child)
+    delete_offload_module(model, "child")
+    delete_offload_module(model.linear, "child")
+    assert not child in model.children()
+    assert not child in model.linear.children()
+
+
 @requires_gpu
 @requires_accelerate()
 @pytest.mark.parametrize("exec_device", [torch.device("cpu"), torch.device("cuda")])