Repository: huggingface/optimum

Issue:
Issue when running transformers from GPU / CPU
Following [this unit test](https://github.com/huggingface/optimum/blob/81b85fa656730b3e0fb7b792781ff072d3f702cb/tests/onnxruntime/test_modeling_ort.py#L458-L463) I was able to successfully run the text-generation with the dummy model, however when I run it with the actual GPT2 model I get an inconsistency of GPU / CPU resources - namely I run:
```python
model_id = "gpt2"
onnx_model = ORTModelForCausalLM.from_pretrained(model_id, from_transformers=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)

pp = pipeline("text-generation", model=onnx_model, tokenizer=tokenizer)
text = "My Name is Philipp and i live"
outputs = pp(text)
```

However I get the following GPU/CPU error:

```
File ~/miniconda3/envs/optimum-mlserver/lib/python3.8/site-packages/transformers/generation_utils.py:587, in GenerationMixin._expand_inputs_for_generation(input_ids, expand_size, is_encoder_decoder, attention_mask, encoder_outputs, **model_kwargs)
    584     model_kwargs["token_type_ids"] = token_type_ids.index_select(0, expanded_return_idx)
    586 if attention_mask is not None:
--> 587     model_kwargs["attention_mask"] = attention_mask.index_select(0, expanded_return_idx)
    589 if is_encoder_decoder:
    590     if encoder_outputs is None:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select)

```

When I change the pipeline unfortunately I get an error where the ORT model does not support the to function (and then if that is circumvented there are other issues like the ORT function does not support config param, etc. Namely if I upate the pipeline to `pp = pipeline("text-generation", model=onnx_model, tokenizer=tokenizer, device=0)` - then the error is:
```
File ~/miniconda3/envs/optimum-mlserver/lib/python3.8/site-packages/transformers/pipelines/base.py:757, in Pipeline.__init__(self, model, tokenizer, feature_extractor, modelcard, framework, task, args_parser, device, binary_output, **kwargs)
    755 # Special handling
    756 if self.framework == "pt" and self.device.type == "cuda":
--> 757     self.model = self.model.to(self.device)
    759 # Update config with task specific parameters
    760 task_specific_params = self.model.config.task_specific_params

AttributeError: 'ORTModelForCausalLM' object has no attribute 'to'
```
For completeness I am getting the notice The model 'ORTModelForCausalLM' is not supported for text-generation.  - is that expected or does that suggest my transformers/hugging-face library is also not compatible with the latest live branch of optimum? I am using the following dependencies:
```
optimum (from main)
huggingface==0.0.1
huggingface-hub==0.5.1
transformers==4.18.0
```


Discussion:

Resolving patch:
diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py
--- a/optimum/onnxruntime/modeling_ort.py
+++ b/optimum/onnxruntime/modeling_ort.py
@@ -687,3 +687,26 @@ def forward(
         return CausalLMOutputWithCrossAttentions(
             logits=torch.from_numpy(outputs[self.model_outputs["logits"]]),
         )
+
+    # Adapted from https://github.com/huggingface/transformers/blob/99289c08a1b16a805dd4ee46de029e9fd23cba3d/src/transformers/generation_utils.py#L490
+    def _prepare_attention_mask_for_generation(
+        self,
+        inputs: torch.Tensor,
+        pad_token_id: int,
+        eos_token_id: int,
+    ) -> torch.LongTensor:
+        """
+        Overrides the base method of `GenerationMixin` to ensure input IDs and
+        attention mask are on the same device.
+        """
+        is_input_ids = len(inputs.shape) == 2 and inputs.dtype in [torch.int, torch.long]
+        is_pad_token_in_inputs = (pad_token_id is not None) and (pad_token_id in inputs)
+        is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or (
+            (eos_token_id is not None) and (pad_token_id != eos_token_id)
+        )
+        # Check if input is input_ids and padded -> only then is attention_mask defined
+        if is_input_ids and is_pad_token_in_inputs and is_pad_token_not_equal_to_eos_token_id:
+            return inputs.ne(pad_token_id).long()
+        else:
+            # Ensure attention mask is on the same device as the input IDs
+            return torch.ones(inputs.shape[:2], dtype=torch.long, device=inputs.device)


Test patch:
diff --git a/tests/onnxruntime/test_modeling_ort.py b/tests/onnxruntime/test_modeling_ort.py
--- a/tests/onnxruntime/test_modeling_ort.py
+++ b/tests/onnxruntime/test_modeling_ort.py
@@ -393,6 +393,7 @@ def test_pipeline(self, *args, **kwargs):
 class ORTModelForCausalLMIntergrationTest(unittest.TestCase):
     SUPPORTED_ARCHITECTURES_WITH_MODEL_ID = {
         "gpt2": "hf-internal-testing/tiny-random-gpt2",
+        "distilgpt2": "distilgpt2",
     }
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
@@ -436,6 +437,21 @@ def test_generate_utils(self, *args, **kwargs):
         self.assertTrue(isinstance(res[0], str))
         self.assertTrue(len(res[0]) > len(text))
 
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    def test_generate_utils_with_input_ids(self, *args, **kwargs):
+        model_arch, model_id = args
+        model = ORTModelForCausalLM.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        text = "This is a sample output"
+        tokens = tokenizer(
+            text,
+            return_tensors="pt",
+        )
+        outputs = model.generate(input_ids=tokens["input_ids"])
+        res = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        self.assertTrue(isinstance(res[0], str))
+        self.assertTrue(len(res[0]) > len(text))
+
     @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
     def test_compare_to_transformers(self, *args, **kwargs):
         model_arch, model_id = args