implement recursive case

src/compressed_tensors/utils/offload.py

@@ -249,15 +249,23 @@ def delete_offload_parameter(module: torch.nn.Module, name: str):
 @check_accelerate(fallback=contextlib.nullcontext())
 @contextlib.contextmanager
 def disable_hf_hook(module: torch.nn.Module, recurse: bool = False):
-    offloaded = has_offloaded_params(module)
-    if offloaded:
-        hook = module._hf_hook
-        remove_hook_from_module(module, recurse=recurse)
+    hooks = {}
+    def collect_hooks(module):
+        nonlocal hooks
+        if hasattr(module, "_hf_hook"):
+            hooks[module] = module._hf_hook
+            remove_hook_from_module(module)
+
+        for submodule in module.children():
+            print(submodule)
+            collect_hooks(submodule)
+
+    collect_hooks(module)
 
     yield
 
-    if offloaded:
-        add_hook_to_module(module, hook)
+    for submodule, hook in hooks.items():
+        add_hook_to_module(submodule, hook)
 
 
 """ Upstreamed Functions """

tests/test_utils/test_offload.py

@@ -161,3 +161,23 @@ def custom_forward():
 
     assert hasattr(module, "_hf_hook")
     assert module._old_forward == custom_forward
+
+
+@requires_accelerate()
+def test_disable_hf_hook_model_recurse():
+    from accelerate.hooks import attach_align_device_hook
+
+    module0 = ExampleModule()
+    module1 = ExampleModule()
+    module2 = ExampleModule()
+    model = torch.nn.Sequential(module0, torch.nn.Sequential(module1, module2))
+    attach_align_device_hook(model, offload=True, weights_map=model.state_dict())
+
+    with disable_hf_hook(model):
+        assert not hasattr(module0, "_hf_hook")
+        assert not hasattr(module1, "_hf_hook")
+        assert not hasattr(module2, "_hf_hook")
+
+    assert hasattr(module0, "_hf_hook")
+    assert hasattr(module1, "_hf_hook")
+    assert hasattr(module2, "_hf_hook")