[Accelerate] Fix offloaded_dispatch, implement disable_offloading (#355)

* fix offloaded_dispatch, implement disable_offloading

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

* update params

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

* small speedup

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

---------

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

Author: kylesayrs

src/compressed_tensors/utils/offload.py

@@ -31,9 +31,10 @@
 import warnings
 from functools import wraps
 from operator import attrgetter
-from typing import Any, Callable, Dict, Iterable, Literal, Optional, Union
+from typing import Any, Callable, Dict, Iterable, Literal, Optional, Tuple, Union
 
 import torch
+from compressed_tensors.utils import patch_attr
 
 
 try:
@@ -83,6 +84,7 @@
     "register_offload_module",
     "delete_offload_module",
     "offloaded_dispatch",
+    "disable_offloading",
 ]
 
 
@@ -214,7 +216,7 @@ def register_offload_parameter(
 def update_offload_parameter(
     module: torch.nn.Module,
     name: str,
-    data: Optional[torch.Tensor],
+    data: torch.Tensor,
     offload_device: Optional[Union[torch.device, Literal["disk"]]] = None,
 ):
     """
@@ -227,15 +229,15 @@ def update_offload_parameter(
     :param offload_device: device on which weight will be offloaded to. If None is
         provided, then infer device from parameters on module
     """
-    param = getattr(module, name)
+    param: torch.nn.Parameter = getattr(module, name)
     if param.data.shape != data.shape:
         warnings.warn(
             f"Shape of parameter being updated {param.data.shape} does not match shape "
             f"of update data {data.shape}"
         )
 
     # copy data into onloaded parameter if applicable
-    if param.device != torch.device("meta"):
+    if param.device != torch.device("meta") and data is not param.data:
         param.data.copy_(data)
 
     # update offload dict
@@ -501,7 +503,9 @@ def offloaded_dispatch(
         raise NotImplementedError("Disk offloading is not currently supported")
 
     # create weights map
-    weights_map = OffloadedWeightsLoader(state_dict=module.state_dict(), device="cpu")
+    state_dict = module.state_dict()
+    state_dict = {key: val.to(offload_device) for key, val in state_dict.items()}
+    weights_map = OffloadedWeightsLoader(state_dict=state_dict, device=offload_device)
 
     # create tied params map
     tied_params = find_tied_parameters(module)
@@ -522,6 +526,36 @@ def offloaded_dispatch(
     return module
 
 
+@contextlib.contextmanager
+def disable_offloading():
+    """
+    Keep modules onloaded and disable offloading until this context exits.
+    Affects modules which have been hooked with accelerate's `AlignDevicesHook`
+    """
+    original_pre_forward = AlignDevicesHook.pre_forward
+    onloaded_modules: Dict[torch.nn.Module, Tuple[AlignDevicesHook, bool]] = dict()
+
+    # onload once and disable any future onloading/offloading steps
+    def keep_onload_pre_forward(self: AlignDevicesHook, module, *args, **kwargs):
+        ret = original_pre_forward(self, module, *args, **kwargs)
+        if module not in onloaded_modules:
+            onloaded_modules[module] = (self, self.offload)
+            self.offload = False
+        return ret
+
+    # use the patched pre_forward function within the context
+    with patch_attr(AlignDevicesHook, "pre_forward", keep_onload_pre_forward):
+        yield
+
+    # manually offload all modules that were onloaded
+    # update any parameters which may have changed
+    for module, (hook, offload) in onloaded_modules.items():
+        hook.offload = offload
+        for name, param in module.named_parameters():
+            update_offload_parameter(module, name, param.data)
+        hook.post_forward(module, None)
+
+
 """ Upstreamed Functions """
 
 

tests/test_utils/test_offload.py

@@ -19,6 +19,7 @@
     delete_offload_module,
     delete_offload_parameter,
     disable_hf_hook,
+    disable_offloading,
     get_execution_device,
     has_offloaded_params,
     offloaded_dispatch,
@@ -397,29 +398,37 @@ def test_delete_offload_module(exec_device):
 
 @requires_gpu
 @requires_accelerate()
-@pytest.mark.parametrize("exec_device", [torch.device("cpu"), torch.device("cuda")])
-def test_offloaded_dispatch(exec_device):
+@pytest.mark.parametrize(
+    "exec_device,offload_device",
+    [
+        (torch.device("cpu"), torch.device("cpu")),
+        (torch.device("cpu"), torch.device("cuda:0")),
+        (torch.device("cuda:0"), torch.device("cpu")),
+        (torch.device("cuda:0"), torch.device("cuda:0")),
+    ],
+)
+def test_offloaded_dispatch(exec_device, offload_device):
     # single module
-    module = torch.nn.Linear(1, 2)
-    module = offloaded_dispatch(module, exec_device)
+    module = torch.nn.Linear(1, 2, device=offload_device)
+    module = offloaded_dispatch(module, exec_device, offload_device)
     assert has_offloaded_params(module)
     assert module._hf_hook.offload
     assert module.weight.device == torch.device("meta")
-    assert "weight" in module._hf_hook.weights_map
+    assert module._hf_hook.weights_map["weight"].device == offload_device
     assert module._hf_hook.tied_params_map is not None
 
     # can run
     module(torch.empty(1, device=exec_device))
 
     # model
     model = ExampleModel()
-    model = offloaded_dispatch(model, exec_device)
+    model = offloaded_dispatch(model, exec_device, offload_device)
     assert not has_offloaded_params(model)
 
     assert has_offloaded_params(model.linear)
     assert model.linear._hf_hook.offload
     assert model.linear.weight.device == torch.device("meta")
-    assert "weight" in model.linear._hf_hook.weights_map
+    assert model.linear._hf_hook.weights_map["weight"].device == offload_device
     assert model.linear._hf_hook.tied_params_map is not None
 
     # can run
@@ -429,4 +438,43 @@ def test_offloaded_dispatch(exec_device):
     parameter = torch.nn.Parameter(torch.tensor(1.0))
     register_offload_parameter(module, "new_param", parameter)
     assert module.new_param.device == torch.device("meta")
-    assert module._hf_hook.weights_map["new_param"].device == torch.device("cpu")
+    assert module._hf_hook.weights_map["new_param"].device == offload_device
+
+
+@requires_gpu
+@requires_accelerate()
+@pytest.mark.parametrize(
+    "exec_device,offload_device",
+    [
+        (torch.device("cpu"), torch.device("cpu")),
+        (torch.device("cpu"), torch.device("cuda:0")),
+        (torch.device("cuda:0"), torch.device("cpu")),
+        (torch.device("cuda:0"), torch.device("cuda:0")),
+    ],
+)
+def test_disable_offloading(exec_device, offload_device):
+    module = torch.nn.Linear(1, 2, device=exec_device)
+
+    # non-offloaded modules are unaffected
+    with disable_offloading():
+        output = module(torch.empty(1, device=exec_device))
+        assert module.weight.device == exec_device
+        assert output.device == exec_device
+
+    # offloaded modules stay on device until context exit
+    offloaded_dispatch(module, exec_device, offload_device)
+    assert module.weight.device == torch.device("meta")
+    assert module._hf_hook.weights_map["weight"].device == offload_device
+
+    with disable_offloading():
+        assert module.weight.device == torch.device("meta")
+        output = module(torch.empty(1, device=exec_device))
+        assert module.weight.device == exec_device
+        assert output.device == exec_device
+
+        output = module(torch.empty(1, device=exec_device))
+        assert module.weight.device == exec_device
+        assert output.device == exec_device
+
+    assert module.weight.device == torch.device("meta")
+    assert module._hf_hook.weights_map["weight"].device == offload_device