Resolved device mismatch error when exporting Hugging Face's BERTModel to TorchScript
Author: lsrock1
Pull request: yes

# What does this PR do?


Fixes #29205 


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@ArthurZucker @fxmarty 
<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
@lsrock1 Thank you for the PR. Regarding the issue https://github.com/huggingface/transformers/issues/29205, is it not possible to trace the model on the same device it is intended to be used?

Overall, doing `torch.arange(key_length, dtype=torch.long).to("cuda")` is more expensive than `torch.arange(key_length, dtype=torch.long, device="cuda")` so I think this downside outweights the benefit from this change. We could probably have a flag `is_jit_tracing` passed around (we would maybe not want to call `torch.jit.is_tracing()` at each layer) so that your code path would be used, but it looks to be complexity added for a small benefit.

To be fair, this case of CPU tracing -> GPU usage (or reverse) is not tested in our CI and I am not sure whether it is something we want to support, I'm personally not in favor of it.

Related: https://discuss.pytorch.org/t/jit-tracing-fails-with-basic-embeddingbag-when-changing-devices/174719/2
https://github.com/pytorch/pytorch/issues/43134
https://discuss.huggingface.co/t/pytorch-trace-on-cpu-and-use-on-gpu/181
https://github.com/pytorch/pytorch/issues/31141

@ArthurZucker @fxmarty  After hearing your opinion, there seem to be two drawbacks in this PR. Firstly, there is a significant overhead in creating and switching devices. Secondly, checking whether tracing is in progress adds additional time. Therefore, the current implementation doesn't seem appropriate.

The reason why resolving this issue is necessary is not only for simply transitioning from CPU to GPU, but also because it's impossible to switch between devices, such as from device:1 to device:2, in a multi-GPU environment. Even if there are four GPUs and you want to deploy the model on each device, it's impossible to move the traced model due to enforced device constraints.

Anyway, thanks for your answering!

Do you mean that tracing on device 1, serializing, and loading on device 2 is impossible as well? For your use case, do you necessarily need to serialize the traced model, and not e.g. trace on the fly on device?

@fxmarty 
While the model's weights can be moved to device:2 using map_location, tensors created with a fixed device during tracing cannot be switched to a different device. Consequently, the model becomes unusable.

@lsrock1 I see. It is indeed not ideal.

cc @IlyasMoutawwakil  now @fxmarty is gone 

Hi @lsrock1 did you try something like:
```python
@torch.jit.script_if_tracing
def get_device(tensor: torch.Tensor):
    """
    Get the tensor's device.
    """
    return tensor.device
```
And then use this scripted method instead for initializing tensors with `torch.constructor(..., device=get_device(tensor))`.
If I understand correctly, if the method is scripted, the value of the device won't be hardcoded.