Repository: huggingface/accelerate

Issue:
Can't copy model after dispatch_model called
### System Info

```Shell
- `Accelerate` version: 0.20.3
- Platform: Linux-5.15.0-70-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- Numpy version: 1.22.2
- PyTorch version (GPU?): 2.1.0a0+4136153 (True)
- PyTorch XPU available: False
```


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [X] My own task or dataset (give details below)

### Reproduction

Reproduce issue with script:
```python
import torch
from torch import nn
import accelerate
import copy

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1,10)
     
    def forward(self, x):
        x = self.fc1(x)

        print(f'In forward pass of model with id: {self.id}')

        return x
    

inp = torch.rand((1,1))
model = MyModel()
model.id = 1
model = accelerate.dispatch_model(model, device_map={'fc1': 'cpu'})

print(f'Output from first model:')
model(inp) # prints id 1 as expected

model_copy = copy.deepcopy(model)
model_copy.id = 2
print(f'Output from copied model:')
model_copy(inp) # incorrectly prints id of 1, when id of 2 was expected
```

The output of this script:
```
Output from first model:
In forward pass of model with id: 1
Output from copied model:
In forward pass of model with id: 1
```



### Expected behavior

Expected output:
```
Output from first model:
In forward pass of model with id: 1
Output from copied model:
In forward pass of model with id: 2
```

When a model is copied, the Accelerate hooks copied over to the new model incorrectly reference the original model. The source of the issue is in `add_hook_to_module` in `hooks.py`
https://github.com/huggingface/accelerate/blob/058a3546ea588abf57de3705409c250cc52423e7/src/accelerate/hooks.py#L158-L168
The issue is that the hooks overwrite the forward method of the model, but do not bind that forward method to the specific class instance. Instead, the `new_forward` method references `module` which is defined outside of the new function. Because of this, when calling the copied version of the model, it references the version of `new_forward` that calls the forward method of the original model (rather than the copied model as expected), leading to the incorrect ID being printed out.

One proposed solution is to the replace the above code block with:
```python
    @functools.wraps(old_forward)
    def new_forward(self, *args, **kwargs):
        args, kwargs = self._hf_hook.pre_forward(self, *args, **kwargs)
        if self._hf_hook.no_grad:
            with torch.no_grad():
                output = self._old_forward(*args, **kwargs)
        else:
            output = self._old_forward(*args, **kwargs)
        return self._hf_hook.post_forward(self, output)

    module.forward = new_forward.__get__(module) # second param to __get__ is optional
```

In this modified version, we bind the forward function to the specific class instance, which means that when we call the forward on the copied model, it now properly calls the forward method of the copied model and outputs the correct ID.

New output (matching expected):
```
Output from first model:
In forward pass of model with id: 1
Output from copied model:
In forward pass of model with id: 2
```

The general motivation for this issue is that `torch.compile` functionality relies on `copy.deepcopy` to copy a model with fake tensor model enabled:
https://github.com/pytorch/pytorch/blob/bb6b157458a34d8b3499932035381fdb12683703/torch/_dynamo/utils.py#L915-L917

Currently, due to the model forward not operating correctly when copied, the `torch.compile` Dynamo tracing fails to properly trace the model. The tracing goes through the original model's forward function with real tensor values from the original model and not the model that is copied with fake tensors, leading to an exception.

Addressing this issue would help step towards being able to use Accelerate to distribute a model and then apply `torch.compile` functionality to that distributed model. I'm also eager to hear if there are alternative solutions to this issue!


Discussion:
Thanks for opening the detailed issue. I'm not too sure the fix is correct: it would bind any future of that class with this new forward method which is specific to this model if I understand correctly. However a user might create a smaller/bigger version of the model in the same script which would require a different forward and those two models wouldn't be able to coexist.

If I'm mistaken and we can create different instances of the model class without any problems, then a PR with this fix would be welcome!
> Reproduce issue with script:

I ran the exact same script on the latest accelerate and I get the expected results:

```
Output from first model:
In forward pass of model with id: 1
Output from copied model:
In forward pass of model with id: 2
```
@BenjaminBossan Thanks for the message. You are correct, I updated my accelerate version to the latest commit and am seeing the correct output as you mentioned. However, the issue appears to actually be with my testing script. Here is an updated script that reproduces my original issue:
```python
import torch
from torch import nn
import accelerate
import copy

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1,10)
        self.fc2 = nn.Linear(10,1)
     
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)

        print(f'In forward pass of model with id: {self.id}')

        return x
    

inp = torch.rand((1,1))
model = MyModel()
model.id = 1
model = accelerate.dispatch_model(model, device_map={'fc1': 'cpu', 'fc2': 'cuda:0'})

print(f'Output from first model:')
model(inp) # prints id 1 as expected

model_copy = copy.deepcopy(model)
model_copy.id = 2
print(f'Output from copied model:')
model_copy(inp) # incorrectly prints id of 1, when id of 2 was expected
```

On the latest accelerate commit `4d13e4e474def51a67ca23cc7356926d45cd63bc` (accelerate `0.22.0.dev0`), I'm getting the incorrect output with this corrected script:
```
Output from first model:
In forward pass of model with id: 1
Output from copied model:
In forward pass of model with id: 1
```

The reason my original testing script was incorrect was that hooks were never getting attached since I only had one device specified in my device map, which is specifically checked for here (this check was added since I last tested my original script):
https://github.com/huggingface/accelerate/blob/4d13e4e474def51a67ca23cc7356926d45cd63bc/src/accelerate/big_modeling.py#L341

Let me know if you are able to reproduce the issue with this updated script. Thanks! 
@sgugger Thanks for the comment! My understanding is that the `module.forward = new_forward.__get__(module)` takes `new_forward` which is an unbound function and binds it to a class instance (in this case `module`). In doing so it produces a bound method for the specific instance but doesn't modify the future behavior of other instances of `MyModel` (the type of `module`). This bounded version of `new_forward` is then assigned as the `forward` function only to the specific model instance and not generally to the underlying `MyModel` class.

To demonstrate experimentally verify this, consider the following script:
```python
import torch
from torch import nn
import accelerate
import copy

class MyModel(nn.Module):
    def __init__(self, model_size):
        super().__init__()
        self.fc1 = nn.Linear(1,model_size)
        self.fc2 = nn.Linear(model_size,1)
        self.model_size = model_size
     
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)

        assert self.model_size == self.fc1.out_features and self.model_size == self.fc2.in_features # sanity check
        print(f'In forward pass of model with id: {self.id} and model size: {self.model_size}\n')

        return x
    

inp = torch.rand((1,1))

model = MyModel(10)
model.id = 1
model = accelerate.dispatch_model(model, device_map={'fc1': 'cpu', 'fc2': 'cuda:0'})
print(f'Output from first model (expected id: 1, model_size: 10):')
model(inp)

model_copy = copy.deepcopy(model)
model_copy.id = 2
print(f'Output from copied model (expected id: 2, model_size: 10):')
model_copy(inp)

diff_size_model = MyModel(20)
diff_size_model.id = 3
diff_size_model = accelerate.dispatch_model(diff_size_model, device_map={'fc1': 'cpu', 'fc2': 'cuda:0'})
print(f'Output from diff_size_model (expected id: 3, model_size: 20):')
diff_size_model(inp)

diff_size_model_copy = copy.deepcopy(diff_size_model)
diff_size_model_copy.id = 4
print(f'Output from copied diff_size_model (expected id: 4, model_size: 20):')
diff_size_model_copy(inp)

# For sanity verify our copy of the original model still produces correct output even though another copy was made on a different model
print(f'sanity: Output from copied model (expected id: 2, model_size: 10):')
model_copy(inp)
```

In this script, I demonstrate copying both a smaller and larger version of the same model. With my fix not included, the output is as follows:
```
Output from first model (expected id: 1, model_size: 10):
In forward pass of model with id: 1 and model size: 10

Output from copied model (expected id: 2, model_size: 10):
In forward pass of model with id: 1 and model size: 10

Output from diff_size_model (expected id: 3, model_size: 20):
In forward pass of model with id: 3 and model size: 20

Output from copied diff_size_model (expected id: 4, model_size: 20):
In forward pass of model with id: 3 and model size: 20

sanity: Output from copied model (expected id: 2, model_size: 10):
In forward pass of model with id: 1 and model size: 10
```

This is the incorrect output due to the issue I discussed in my initial thread. With my fix, the copy behavior now works and we still are able to have future models of varying sizes that are not impacted by this change.

The output with the proposed fix:
```
Output from first model (expected id: 1, model_size: 10):
In forward pass of model with id: 1 and model size: 10

Output from copied model (expected id: 2, model_size: 10):
In forward pass of model with id: 2 and model size: 10

Output from diff_size_model (expected id: 3, model_size: 20):
In forward pass of model with id: 3 and model size: 20

Output from copied diff_size_model (expected id: 4, model_size: 20):
In forward pass of model with id: 4 and model size: 20

sanity: Output from copied model (expected id: 2, model_size: 10):
In forward pass of model with id: 2 and model size: 10
```

After changing the line you indicated as a fix in the hooks module, I am personally getting an error on your code sample:

```
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
/tmp/ipykernel_376368/703981786.py in <module>
     27 model = accelerate.dispatch_model(model, device_map={'fc1': 'cpu', 'fc2': 'cuda:0'})
     28 print(f'Output from first model (expected id: 1, model_size: 10):')
---> 29 model(inp)
     30 
     31 model_copy = copy.deepcopy(model)

~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

~/git/accelerate/src/accelerate/hooks.py in new_forward(*args, **kwargs)
    158     @functools.wraps(old_forward)
    159     def new_forward(*args, **kwargs):
--> 160         args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
    161         if module._hf_hook.no_grad:
    162             with torch.no_grad():

~/git/accelerate/src/accelerate/hooks.py in pre_forward(self, module, *args, **kwargs)
    288                 )
    289 
--> 290         return send_to_device(args, self.execution_device), send_to_device(
    291             kwargs, self.execution_device, skip_keys=self.skip_keys
    292         )

~/git/accelerate/src/accelerate/utils/operations.py in send_to_device(tensor, device, non_blocking, skip_keys)
    149     """
    150     if isinstance(tensor, (tuple, list)):
--> 151         return honor_type(
    152             tensor, (send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys) for t in tensor)
    153         )

~/git/accelerate/src/accelerate/utils/operations.py in honor_type(obj, generator)
     81         return type(obj)(*list(generator))
     82     else:
---> 83         return type(obj)(generator)
     84 
     85 

~/git/accelerate/src/accelerate/utils/operations.py in <genexpr>(.0)
    150     if isinstance(tensor, (tuple, list)):
    151         return honor_type(
--> 152             tensor, (send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys) for t in tensor)
    153         )
    154     elif isinstance(tensor, Mapping):

~/git/accelerate/src/accelerate/utils/operations.py in send_to_device(tensor, device, non_blocking, skip_keys)
    165     elif hasattr(tensor, "to"):
    166         try:
--> 167             return tensor.to(device, non_blocking=non_blocking)
    168         except TypeError:  # .to() doesn't accept non_blocking as kwarg
    169             return tensor.to(device)

~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py in to(self, *args, **kwargs)
   1143             return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
   1144 
-> 1145         return self._apply(convert)
   1146 
   1147     def register_full_backward_pre_hook(

~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py in _apply(self, fn)
    795     def _apply(self, fn):
    796         for module in self.children():
--> 797             module._apply(fn)
    798 
    799         def compute_should_use_set_data(tensor, tensor_applied):

~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py in _apply(self, fn)
    818             # `with torch.no_grad():`
    819             with torch.no_grad():
--> 820                 param_applied = fn(param)
    821             should_use_set_data = compute_should_use_set_data(param, param_applied)
    822             if should_use_set_data:

~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py in convert(t)
   1141                 return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
   1142                             non_blocking, memory_format=convert_to_format)
-> 1143             return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
   1144 
   1145         return self._apply(convert)

NotImplementedError: Cannot copy out of meta tensor; no data!
```
@sgugger My proposed changes include that entire code block I listed (including changes in the `new_forward` function itself). From the stack trace it appears you are still using the unmodified version of `new_forward`. Could you try with the entire code block listed?
Indeed, sorry about that! Looks like the changes don't break anything on a first try. Would you like to open a PR so we can test more extensively?
@sgugger Thanks for testing it out! I will try to open a PR this week and will update here when I have done so.
@austinapatel @sgugger  `dispatch_model` make the model non-serializable (https://github.com/huggingface/accelerate/issues/1921). Excuse me, does fixing the issue mentioned in this issue also mean that the model can be serialized?

Resolving patch:
diff --git a/src/accelerate/hooks.py b/src/accelerate/hooks.py
--- a/src/accelerate/hooks.py
+++ b/src/accelerate/hooks.py
@@ -155,17 +155,17 @@ def add_hook_to_module(module: nn.Module, hook: ModelHook, append: bool = False)
     module = hook.init_hook(module)
     module._hf_hook = hook
 
-    @functools.wraps(old_forward)
-    def new_forward(*args, **kwargs):
+    def new_forward(module, *args, **kwargs):
         args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
         if module._hf_hook.no_grad:
             with torch.no_grad():
-                output = old_forward(*args, **kwargs)
+                output = module._old_forward(*args, **kwargs)
         else:
-            output = old_forward(*args, **kwargs)
+            output = module._old_forward(*args, **kwargs)
         return module._hf_hook.post_forward(module, output)
 
-    module.forward = new_forward
+    module.forward = functools.update_wrapper(functools.partial(new_forward, module), old_forward)
+
     return module
 
 


Test patch:
diff --git a/tests/test_big_modeling.py b/tests/test_big_modeling.py
--- a/tests/test_big_modeling.py
+++ b/tests/test_big_modeling.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import copy
 import os
 import unittest
 from tempfile import TemporaryDirectory
@@ -45,6 +45,18 @@ def forward(self, x):
         return self.linear2(self.batchnorm(self.linear1(x)))
 
 
+class ModelForTestCopy(nn.Module):
+    def __init__(self, id: int):
+        super().__init__()
+        self.id = id
+        self.linear1 = nn.Linear(3, 4)
+        self.batchnorm = nn.BatchNorm1d(4)
+        self.linear2 = nn.Linear(4, 5)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x))), self.id
+
+
 class ModelForTestTiedWeights(nn.Module):
     def __init__(self):
         super().__init__()
@@ -325,6 +337,25 @@ def test_dispatch_model_multi_gpu(self):
             output = model(x)
             self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
 
+    @require_cuda
+    def test_dispatch_model_copy(self):
+        original_model = ModelForTestCopy(id=1)
+        device_map = {"linear1": 0, "batchnorm": "cpu", "linear2": 0}
+
+        x = torch.randn(2, 3)
+        expected, original_output_id = original_model(x)
+
+        dispatch_model(original_model, device_map)
+
+        copied_model = copy.deepcopy(original_model)
+        copied_model.id = 2
+        output, copied_output_id = copied_model(x)
+
+        self.assertEqual(original_model.id, original_output_id)
+        self.assertEqual(copied_model.id, copied_output_id)
+        self.assertFalse(copied_model.linear1.forward is original_model.linear1.forward)
+        self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
     @require_cuda
     def test_dispatch_model_move_offloaded_model(self):
         model = ModelForTest()