Repository: pytorch
Issue: For PyTorch Nightly, failure when changing MPS device to CPU after PYTORCH_ENABLE_MPS_FALLBACK occurs.
Labels: module: crash,triaged,module: mps

### 🐛 Describe the bug

When trying to generate text with a GPT-2 from the transformers library, I get this error:

NotImplementedError: The operator 'aten::cumsum.out' is not current implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.

So I activated the environment variable (I set it in the terminal, because it didn't work with the version in the following code), but afterwards another error occurs (I posted it after the code used). I need to mention that if I only use CPU from the start, the generation works without problems.

```python
# Sample code to reproduce the problem

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import os
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = "1"

device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
tokenizer = AutoTokenizer.from_pretrained('readerbench/RoGPT2-base')
model = AutoModelForCausalLM.from_pretrained('readerbench/RoGPT2-base').to(device)

inputs = tokenizer('Salut priete', return_tensors='pt').to(device)
generation = model.generate(inputs['input_ids'], max_length = len(inputs['input_ids'][0]) + 10, no_repeat_ngram_size=2, num_beams=5, early_stopping=True, num_return_sequences=1)
```

```
The error message you got, with the full traceback.

Traceback (most recent call last):
  File "/Users/alexandrudima/home/Research/test.py", line 15, in <module>
    generation = model.generate(inputs['input_ids'], max_length = len(inputs['input_ids'][0]) + 10, no_repeat_ngram_size=2, num_beams=5, early_stopping=True, num_return_sequences=1)[0][len(inputs['input_ids'][0]):]
  File "/opt/homebrew/Caskroom/miniforge/base/envs/ml/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/ml/lib/python3.9/site-packages/transformers/generation_utils.py", line 1386, in generate
    return self.beam_search(
  File "/opt/homebrew/Caskroom/miniforge/base/envs/ml/lib/python3.9/site-packages/transformers/generation_utils.py", line 2232, in beam_search
    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/ml/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 1016, in prepare_inputs_for_generation
    position_ids = attention_mask.long().cumsum(-1) - 1
NotImplementedError: The operator 'aten::cumsum.out' is not current implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.
(ml) alexandrudima@Alex-MacBook Research % PYTORCH_ENABLE_MPS_FALLBACK=1 python test.py
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
tensor([[23640,   344,  3205]], device='mps:0') tensor([23640,   344,  3205], device='mps:0')
The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.
/opt/homebrew/Caskroom/miniforge/base/envs/ml/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py:1016: UserWarning: The operator 'aten::cumsum.out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/_temp/anaconda/conda-bld/pytorch_1661929883516/work/aten/src/ATen/mps/MPSFallback.mm:11.)
  position_ids = attention_mask.long().cumsum(-1) - 1
Traceback (most recent call last):
  File "/Users/alexandrudima/home/Research/test.py", line 15, in <module>
    generation = model.generate(inputs['input_ids'], max_length = len(inputs['input_ids'][0]) + 10, no_repeat_ngram_size=2, num_beams=5, early_stopping=True, num_return_sequences=1)[0][len(inputs['input_ids'][0]):]
  File "/opt/homebrew/Caskroom/miniforge/base/envs/ml/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/ml/lib/python3.9/site-packages/transformers/generation_utils.py", line 1386, in generate
    return self.beam_search(
  File "/opt/homebrew/Caskroom/miniforge/base/envs/ml/lib/python3.9/site-packages/transformers/generation_utils.py", line 2253, in beam_search
    next_token_scores_processed = logits_processor(input_ids, next_token_scores)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/ml/lib/python3.9/site-packages/transformers/generation_logits_process.py", line 92, in __call__
    scores = processor(input_ids, scores)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/ml/lib/python3.9/site-packages/transformers/generation_logits_process.py", line 333, in __call__
    scores[i, banned_tokens] = -float("inf")
RuntimeError: dst_.nbytes() >= dst_byte_offset INTERNAL ASSERT FAILED at "/Users/runner/work/_temp/anaconda/conda-bld/pytorch_1661929883516/work/aten/src/ATen/native/mps/operations/Copy.mm":184, please report a bug to PyTorch. 
```


### Versions

Collecting environment information...
PyTorch version: 1.13.0.dev20220831
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.5.1 (arm64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2.5)
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 17:00:33)  [Clang 13.0.1 ] (64-bit runtime)
Python platform: macOS-12.5.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.22.4
[pip3] pytorch-lightning==1.5.8
[pip3] torch==1.13.0.dev20220831
[pip3] torchaudio==0.13.0.dev20220831
[pip3] torchmetrics==0.9.3
[pip3] torchvision==0.14.0.dev20220831
[conda] numpy                     1.22.4           py39h7df2422_0    conda-forge
[conda] pytorch                   1.13.0.dev20220831         py3.9_0    pytorch-nightly
[conda] pytorch-lightning         1.5.8              pyhd8ed1ab_0    conda-forge
[conda] torchaudio                0.13.0.dev20220831        py39_cpu    pytorch-nightly
[conda] torchmetrics              0.9.3              pyhd8ed1ab_0    conda-forge
[conda] torchvision               0.14.0.dev20220831        py39_cpu    pytorch-nightly

cc @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen