Illegal memory access when using Trainer API on GPU with PyTorch 2.0's Inductor backend
Author: Lokiiiiii
Pull request: no

### System Info

### System Information

- `transformers` version: 4.26.1
- Platform: Linux-5.15.0-1028-aws-x86_64-with-glibc2.10
- Python version: 3.8.16
- Huggingface_hub version: 0.12.1
- PyTorch version (GPU?): 2.0.0a0+git8693604 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes (A10G/A100)
- Using distributed or parallel set-up in script?: No

### More information

1. Training is successful in eager mode with batch size of 128
2. Training is successful in dynamo + eager mode with batch size of 128
3. Training is only able to succeed with dynamo + inductor with batch size of 2
4. [Dynamo benchmarks](https://github.com/pytorch/pytorch/blob/master/benchmarks/dynamo/huggingface.py) which use the same HF models without Trainer API are able to succeed.


### Error Signature 

| 1677525140847 | Traceback (most recent call last):   File "./run_mlm.py", line 694, in <module>                                                                                                                                                                                                                                                                                                                                                                                 |
| 1677525140847 | main()   File "./run_mlm.py", line 635, in main                                                                                                                                                                                                                                                                                                                                                                                                                 |
| 1677525140847 | train_result = trainer.train(resume_from_checkpoint=checkpoint)   File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 1631, in train                                                                                                                                                                                                                                                                                                    |
| 1677525140847 | return inner_training_loop(   File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 1898, in _inner_training_loop                                                                                                                                                                                                                                                                                                                         |
| 1677525140848 | tr_loss_step = self.training_step(model, inputs)   File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 2640, in training_step                                                                                                                                                                                                                                                                                                           |
| 1677525140848 | loss = self.compute_loss(model, inputs)   File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 2672, in compute_loss                                                                                                                                                                                                                                                                                                                     |
| 1677525140848 | outputs = model(**inputs)   File "/opt/conda/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 82, in __call__     return self.dynamo_ctx(self._orig_mod.__call__)(*args, **kwargs)   File "/opt/conda/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 215, in _fn                                                                                                                                                                  |
| 1677525140848 | return fn(*args, **kwargs)   File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl                                                                                                                                                                                                                                                                                                                                 |
| 1677525140848 | return forward_call(*args, **kwargs)   File "/opt/conda/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 1324, in forward                                                                                                                                                                                                                                                                                                           |
| 1677525140848 | @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))   File "/opt/conda/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 215, in _fn                                                                                                                                                                                                                                                               |
| 1677525140848 | return fn(*args, **kwargs)   File "/opt/conda/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 2816, in forward                                                                                                                                                                                                                                                                                                                              |
| 1677525140848 | return compiled_fn(full_args)   File "/opt/conda/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 1222, in g                                                                                                                                                                                                                                                                                                                                 |
| 1677525140848 | return f(*args)   File "/opt/conda/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 2383, in debug_compiled_function                                                                                                                                                                                                                                                                                                                         |
| 1677525140848 | return compiled_function(*args)   File "/opt/conda/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 1895, in runtime_wrapper                                                                                                                                                                                                                                                                                                                 |
| 1677525140848 | all_outs = call_func_with_args(   File "/opt/conda/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 1256, in call_func_with_args                                                                                                                                                                                                                                                                                                             |
| 1677525140848 | out = normalize_as_list(f(*args))   File "/opt/conda/lib/python3.8/site-packages/torch/autograd/function.py", line 506, in apply                                                                                                                                                                                                                                                                                                                                |
| 1677525140848 | return super().apply(*args, **kwargs)  # type: ignore[misc]   File "/opt/conda/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 2148, in forward                                                                                                                                                                                                                                                                                             |
| 1677525140848 | fw_outs = call_func_with_args(   File "/opt/conda/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 1247, in call_func_with_args                                                                                                                                                                                                                                                                                                              |
| 1677525140848 | out = normalize_as_list(f(args))   File "/opt/conda/lib/python3.8/site-packages/torch/_inductor/compile_fx.py", line 248, in run                                                                                                                                                                                                                                                                                                                                |
| 1677525140848 | return model(new_inputs)   File "/tmp/torchinductor_root/rj/crjch2m3bp6tuhd3s6n2apgbibxay4o6o5jlrfbwsfiokrv2rkep.py", line 4483, in call                                                                                                                                                                                                                                                                                                                        |
| 1677525140848 | triton__49.run(primals_208, buf513, buf514, 128, 128, grid=grid(128), stream=stream0)   File "/opt/conda/lib/python3.8/site-packages/torch/_inductor/triton_ops/autotune.py", line 184, in run                                                                                                                                                                                                                                                                  |
| 1677525140848 | self.autotune_to_one_config(*args, grid=grid)   File "/opt/conda/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper                                                                                                                                                                                                                                                                                                                 |
| 1677525140848 | r = func(*args, **kwargs)   File "/opt/conda/lib/python3.8/site-packages/torch/_inductor/triton_ops/autotune.py", line 171, in autotune_to_one_config     timings = {   File "/opt/conda/lib/python3.8/site-packages/torch/_inductor/triton_ops/autotune.py", line 172, in <dictcomp>                                                                                                                                                                           |
| 1677525140848 | launcher: self.bench(launcher, *cloned_args, **kwargs)   File "/opt/conda/lib/python3.8/site-packages/torch/_inductor/triton_ops/autotune.py", line 153, in bench     return do_bench(kernel_call, rep=40, fast_flush=True)   File "/opt/conda/lib/python3.8/site-packages/triton/testing.py", line 144, in do_bench                                                                                                                                            |
| 1677525140848 | torch.cuda.synchronize()   File "/opt/conda/lib/python3.8/site-packages/torch/cuda/__init__.py", line 711, in synchronize                                                                                                                                                                                                                                                                                                                                       |
| 1677525140848 | return torch._C._cuda_synchronize()                                                                                                                                                                                                                                                                                                                                                                                                                             |
| 1677525140848 | RuntimeError: CUDA error: an illegal memory access was encountered                                                                                                                                                                                                                           

### Who can help?

@sgugger

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Training Command: ['python', './run_mlm.py', '--model_name_or_path', 'bert-base-uncased', '--output_dir', '/opt/ml/model', '--fp16', '--dataloader_drop_last', '--dataset_config_name', 'wikitext-2-raw-v1', '--dataset_name', 'wikitext', '--do_train', '--evaluation_strategy', 'no', '--logging_strategy', 'epoch', '--max_seq_length', '128', '--num_train_epochs', '50', '--overwrite_output_dir', '--per_device_train_batch_size', '128', '--save_strategy', 'no', '--torch_compile_backend', 'inductor']

### Expected behavior

No exceptions when using Inductor backend with the trainer API.


Comments:
Pretty much the same as #21826 , inductor backend is not yet supported.

@ArthurZucker  Can you elaborate on this ?

Do you mean the Trainer class is using a piece of code that is unsupported by inductor backend ?
Is this something that you will wait on PyTorch to fix or are you amenable to workarounds in the Trainer class for the inductor backend ?

Is there already a root cause that boils down to a few pieces of code ?


Hi @Lokiiiiii We haven't investigated the issue yet, to make sure if it comes from a bug in PyTorch or in Transformers. Stay tuned!

I can reproduce this issue and have something smaller:
```py
import torch

import transformers
from transformers import AutoModelForMaskedLM, AutoTokenizer, DataCollatorForLanguageModeling

def main():
    model_name = "bert-base-uncased"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForMaskedLM.from_pretrained(model_name)
    
    texts = ["This is a text for the example."] * 16
    tokenized_texts = tokenizer(texts, padding="max_length", truncation=True, max_length=128, return_tensors="pt")      
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer)
    batch = tokenized_texts
    batch["input_ids"], batch["labels"] = data_collator.torch_mask_tokens(batch["input_ids"])

    model = torch.compile(model, backend="inductor")
    model.to("cuda")

    batch = {k: v.to("cuda") for k, v in batch.items()}
    outputs = model(**batch)
    loss = outputs.loss
    loss.backward()


if __name__ == "__main__":
    main()
```
Reaching out to the PyTorch folks as it is raised in the model forward, so not something in Transformers at first glance.

See https://github.com/pytorch/pytorch/issues/95794

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.