Repository: huggingface/accelerate

Issue:
`big_modeling.init_on_device` not working with WrapperTensor class like `FakeTensor`
### System Info

```Shell
- `Accelerate` version: 0.21.0
- Platform: Linux-5.19.0-45-generic-x86_64-with-glibc2.17
- Python version: 3.8.16
- Numpy version: 1.21.2
- PyTorch version (GPU?): 2.1.0.dev20230725+cu118 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- System RAM: 62.70 GB
- GPU type: NVIDIA GeForce RTX 3070
- `Accelerate` default config:
        Not found
```


### 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

My repro:

```
import torch
from torch._subclasses.fake_tensor import FakeTensorMode
from accelerate.big_modeling import init_on_device

fake_mode = FakeTensorMode()
with fake_mode:
    with init_on_device(torch.device('cuda')):
        net = torch.nn.Linear(3,4)
```

### Expected behavior

As torch introduced WrapperTensor class and TensorDispatchMode, it is possible that user may want to create custom tensors/parameters while loading from transformers/diffusers. Ideally the above code should work with any WrapperTensor class. However, `big_modeling.init_on_device` assume that everything in `WrapperTensor.__dict__` should be keyword argument of `WrapperTensor.new()`, which may not be the case for `WrapperTensor` classes.

Output:
```
Traceback (most recent call last):
  File "ex.py", line 16, in <module>
    net = torch.nn.Linear(3,4)
  File "python3.8/site-packages/torch/nn/modules/linear.py", line 96, in __init__
    self.weight = Parameter(torch.empty((out_features, in_features), **factory_kwargs))
  File "python3.8/site-packages/torch/nn/modules/module.py", line 1712, in __setattr__
    self.register_parameter(name, value)
  File "python3.8/site-packages/accelerate/big_modeling.py", line 105, in register_empty_parameter
    module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
TypeError: __new__() got an unexpected keyword argument 'fake_device'
```


Discussion:
I don't fully understand the use case of combining this with `init_on_device`. Isn't this new PyTorch API the way to directly init a model on devices (e.g. it's another way to do what Accelerate does with `init_on_device`)?

We can certainly look at a fix, but want to fully understand the use case first. (cc @SunMarc )
@sgugger Thanks for the response! My use case is that, with `TensorDispatchMode`, users can now run arbitrary pytorch/transformers code under the context and ensure that all tensors could be wrapped and dispatched to the user-specified behavior. E.g. user can have a `LoggingTensor` which does exactly the same as `torch.Tensor`, on the same device etc., as well as logging all tensor operations with a server. This Wrapper tensor class could have a different class init/new signature so `init_on_device` won't work. So `init_on_device` is more like `with torch.device()` in my opinion. And my use case looks like:

```
with CustomTensorDispatchMode():  # -> create custom tensors instead of torch.Tensor
  model = transformers.some_model.from_pretrained(...)  # -> this could call `init_on_device`
```

The other way to think about this is that now user can define custom tensor class with pytorch, the assumption in `init_on_device`, which is all items in `tensor.__dict__` would be a keyword argument to the tensor's contructor, no longer holds. 
> The other way to think about this is that now user can define custom tensor class with pytorch, the assumption in init_on_device, which is all items in tensor.__dict__ would be a keyword argument to the tensor's contructor, no longer holds.

Yes but at the same time, we don't have an API from PyTorch to move a tensor from the meta device without recreating a Tensor, so it's a tad bit hard to reconcile 😅 
If what we want to fix is `init_on_device`, what about changing it to a wrapper of `with torch.device(...)`? That way we don't need to re-create tensors as parameters. 

```
import torch
from torch._subclasses.fake_tensor import FakeTensorMode

with torch.device('cuda'):
  net = torch.nn.Linear(3,4)

net.weight.device  # -> device(type='cuda', index=0)

with FakeTensorMode():
   with torch.device('cuda'):
     net = torch.nn.Linear(3,4)

net.weight  # -> Parameter(FakeTensor(..., device='cuda:0', size=(4, 3), requires_grad=True))
net.weight.device # -> device(type='cuda', index=0)
```

It's probably not BC though.



That means dropping support for all older versions of PyTorch, so not something we can really do on our side.

Resolving patch:
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -34,6 +34,7 @@
     find_tied_parameters,
     get_balanced_memory,
     infer_auto_device_map,
+    is_torch_version,
     load_checkpoint_in_model,
     offload_state_dict,
     retie_parameters,
@@ -93,6 +94,12 @@ def init_on_device(device: torch.device, include_buffers: bool = False):
         tst = nn.Liner(100, 100)  # on `cuda` device
     ```
     """
+    # TODO(shingjan): remove the torch version check once older versions are deprecated
+    if is_torch_version(">=", "2.0") and include_buffers:
+        with device:
+            yield
+        return
+
     old_register_parameter = nn.Module.register_parameter
     if include_buffers:
         old_register_buffer = nn.Module.register_buffer


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
@@ -31,7 +31,7 @@
 )
 from accelerate.hooks import remove_hook_from_submodules
 from accelerate.test_utils import require_bnb, require_cuda, require_mps, require_multi_gpu, slow
-from accelerate.utils import offload_state_dict
+from accelerate.utils import is_torch_version, offload_state_dict
 
 
 class ModelForTest(nn.Module):
@@ -106,8 +106,14 @@ def test_init_empty_weights(self):
         self.assertEqual(module.running_mean.device, torch.device("cpu"))
 
         # Use with include_buffers=True
+        register_parameter_func = nn.Module.register_parameter
+        register_buffer_func = nn.Module.register_buffer
         with init_empty_weights(include_buffers=True):
             module = nn.BatchNorm1d(4)
+            # nn.Module.register_parameter/buffer shouldn't be changed with torch >= 2.0
+            if is_torch_version(">=", "2.0"):
+                self.assertEqual(register_parameter_func, nn.Module.register_parameter)
+                self.assertEqual(register_buffer_func, nn.Module.register_buffer)
         self.assertEqual(module.weight.device, torch.device("meta"))
         self.assertEqual(module.running_mean.device, torch.device("meta"))