Repository: huggingface/optimum

Issue:
Tried to load a model to both cores of a two-core GPUs, but only one core (device:0) is used. 
I used Optimum to convert a finetuned `xlm-roberta-base` model to an ONNX model by following the tutorial in [this link](https://huggingface.co/docs/optimum/quicktour). I have a two core GPU. When I tried to load the original finetuned model, I was able to load it to both cores of my GPU. But for some reason, when I try it with the ONNX model, the model only gets loaded into one of the cores. And because of that, that core's memory gets maxed out. 

The code that deals with model loading is below, and any suggestion and advice would be much appreciated! 

```
def worker(i, sents, model_path):
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = ORTModelForSequenceClassification.from_pretrained(model_path)
    if i % 2 == 0:        
        pipe = pipeline("text-classification", 
                        model=model
                        tokenizer=tokenizer, 
                        padding="max_length", 
                        max_length=96, 
                        truncation=True, 
                        device=0, 
                        batch_size=64)
    else:
        pipe = pipeline("text-classification", 
                        model=model, 
                        tokenizer=tokenizer, 
                        padding="max_length", 
                        max_length=96, 
                        truncation=True, 
                        device=1, 
                        batch_size=64)
    scores = pipe(sents)
    return scores
```




Discussion:
Thank you! This a limitation from https://github.com/huggingface/optimum/blob/d15b47f6f0ab750b1203463f98869a8bcbb99c77/optimum/onnxruntime/modeling_ort.py#L141-L153 . Let me fix this shortly.

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
@@ -49,7 +49,7 @@
 from huggingface_hub import HfApi, hf_hub_download
 
 from ..modeling_base import FROM_PRETRAINED_START_DOCSTRING, OptimizedModel
-from .utils import ONNX_WEIGHTS_NAME, get_device_for_provider, get_provider_for_device
+from .utils import ONNX_WEIGHTS_NAME, get_device_for_provider, get_provider_for_device, parse_device
 
 
 logger = logging.getLogger(__name__)
@@ -138,18 +138,25 @@ def device(self) -> torch.device:
     def device(self, value: torch.device):
         self._device = value
 
-    def to(self, device: torch.device):
+    def to(self, device: Union[torch.device, str, int]):
         """
         Changes the ONNX Runtime provider according to the device.
+
+        Arguments:
+            device (`torch.device` or `str` or `int`):
+                Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, a positive will run
+                the model on the associated CUDA device id. You can pass native `torch.device` or a `str` too.
+
+        Returns:
+            `ORTModel`: the model placed on the requested device.
         """
-        # convert string device input (ie. "cuda") to torch.device
-        if type(device) == str:
-            device = torch.device(device)
+        device, provider_options = parse_device(device)
 
         self.device = device
         provider = get_provider_for_device(self.device)
-        self.model.set_providers([provider])
+        self.model.set_providers([provider], provider_options=[provider_options])
         self.providers = self.model.get_providers()
+
         return self
 
     def forward(self, *args, **kwargs):
@@ -230,6 +237,9 @@ def from_pretrained(
             possible providers. Defaults to `CPUExecutionProvider`.
         session_options (`onnxruntime.SessionOptions`, *optional*),:
             ONNX Runtime session options to use for loading the model. Defaults to `None`.
+        provider_options (`Dict`, **optional**):
+            Provider option dictionaries corresponding to the provider used. See available options
+            for each provider: https://onnxruntime.ai/docs/api/c/group___global.html . Defaults to `None`.
 
         Returns:
             `ORTModel`: The loaded ORTModel model.
diff --git a/optimum/onnxruntime/modeling_seq2seq.py b/optimum/onnxruntime/modeling_seq2seq.py
--- a/optimum/onnxruntime/modeling_seq2seq.py
+++ b/optimum/onnxruntime/modeling_seq2seq.py
@@ -40,6 +40,7 @@
     _is_gpu_available,
     get_device_for_provider,
     get_provider_for_device,
+    parse_device,
 )
 
 
@@ -474,23 +475,28 @@ def _from_transformers(
         kwargs["config"] = model.config.__dict__
         return cls._from_pretrained(save_dir, **kwargs)
 
-    def to(self, device):
+    def to(self, device: Union[torch.device, str, int]):
         """
         Changes the ONNX Runtime provider according to the device.
+
+        Arguments:
+            device (`torch.device` or `str` or `int`):
+                Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, a positive will run
+                the model on the associated CUDA device id. You can pass native `torch.device` or a `str` too.
+
+        Returns:
+            `ORTModel`: the model placed on the requested device.
         """
-        # convert string device input (ie. "cuda") to torch.device
-        if type(device) == str:
-            device = torch.device(device)
+        device, provider_options = parse_device(device)
 
-        self.device = device
         provider = get_provider_for_device(device)
         self.encoder._device = device
-        self.encoder.session.set_providers([provider])
+        self.encoder.session.set_providers([provider], provider_options=[provider_options])
         self.decoder._device = device
-        self.decoder.session.set_providers([provider])
+        self.decoder.session.set_providers([provider], provider_options=[provider_options])
         if self.decoder_with_past is not None:
             self.decoder_with_past._device = device
-            self.decoder_with_past.session.set_providers([provider])
+            self.decoder_with_past.session.set_providers([provider], provider_options=[provider_options])
         return self
 
 
diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py
--- a/optimum/onnxruntime/utils.py
+++ b/optimum/onnxruntime/utils.py
@@ -12,6 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from enum import Enum
+from typing import Dict, Tuple, Union
 
 import torch
 from transformers.onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast
@@ -168,6 +169,21 @@ def get_provider_for_device(device: torch.device) -> str:
     return "CUDAExecutionProvider" if device.type.lower() == "cuda" else "CPUExecutionProvider"
 
 
+def parse_device(device: Union[torch.device, str, int]) -> Tuple[torch.device, Dict]:
+    """Get the relevant torch.device from the passed device, and if relevant the provider options (e.g. to set the GPU id)."""
+
+    if device == -1:
+        device = torch.device("cpu")
+    else:
+        device = torch._C._nn._parse_to(device)[0]
+
+    provider_options = {}
+    if device.type == "cuda" and device.index:
+        provider_options["device_id"] = device.index
+
+    return device, provider_options
+
+
 class ORTQuantizableOperator(Enum):
     # Common ops
     Gather = "Gather"


Test patch:
diff --git a/tests/benchmark/test_transformers_optimum_examples_parity.py b/tests/benchmark/test_transformers_optimum_examples_parity.py
--- a/tests/benchmark/test_transformers_optimum_examples_parity.py
+++ b/tests/benchmark/test_transformers_optimum_examples_parity.py
@@ -271,6 +271,9 @@ def test_question_answering_parity(self):
             optimum_results["exact_match"], benchmark_results["evaluation"]["others"]["optimized"]["exact_match"]
         )
 
+    @unittest.skip(
+        "failing related to shuffle issue https://github.com/huggingface/datasets/issues/5145 , skip for now"
+    )
     def test_image_classification_parity(self):
         # dummy test until this question is solved
         # https://discuss.huggingface.co/t/why-use-val-transforms-function-in-image-classification-example-instead-of-feature-extractor/19976
diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
--- a/tests/onnxruntime/test_modeling.py
+++ b/tests/onnxruntime/test_modeling.py
@@ -34,7 +34,7 @@
     set_seed,
 )
 from transformers.onnx.utils import get_preprocessor
-from transformers.testing_utils import require_torch_gpu
+from transformers.testing_utils import get_gpu_count, require_torch_gpu
 
 import onnxruntime
 import requests
@@ -253,6 +253,20 @@ def test_passing_provider_options(self):
         )
         self.assertEqual(model.model.get_provider_options()["CUDAExecutionProvider"]["do_copy_in_default_stream"], "0")
 
+    @unittest.skipIf(get_gpu_count() <= 1, "this test requires multi-gpu")
+    def test_model_on_gpu_id(self):
+        model = ORTModel.from_pretrained(self.ONNX_MODEL_ID)
+        model.to(torch.device("cuda:1"))
+        self.assertEqual(model.model.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1")
+
+        model = ORTModel.from_pretrained(self.ONNX_MODEL_ID)
+        model.to(1)
+        self.assertEqual(model.model.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1")
+
+        model = ORTModel.from_pretrained(self.ONNX_MODEL_ID)
+        model.to("cuda:1")
+        self.assertEqual(model.model.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1")
+
     @require_torch_gpu
     def test_passing_provider_options_seq2seq(self):
         model = ORTModelForSeq2SeqLM.from_pretrained(self.ONNX_SEQ2SEQ_MODEL_ID, provider="CUDAExecutionProvider")
@@ -316,6 +330,32 @@ def test_seq2seq_model_on_gpu(self):
         self.assertEqual(model.decoder_with_past.session.get_providers()[0], "CUDAExecutionProvider")
         self.assertListEqual(model.providers, ["CUDAExecutionProvider", "CPUExecutionProvider"])
 
+    @unittest.skipIf(get_gpu_count() <= 1, "this test requires multi-gpu")
+    def test_seq2seq_model_on_gpu_id(self):
+        model = ORTModelForSeq2SeqLM.from_pretrained(self.ONNX_SEQ2SEQ_MODEL_ID, use_cache=True)
+        model.to(torch.device("cuda:1"))
+        self.assertEqual(model.encoder.session.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1")
+        self.assertEqual(model.decoder.session.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1")
+        self.assertEqual(
+            model.decoder_with_past.session.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1"
+        )
+
+        model = ORTModelForSeq2SeqLM.from_pretrained(self.ONNX_SEQ2SEQ_MODEL_ID, use_cache=True)
+        model.to(1)
+        self.assertEqual(model.encoder.session.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1")
+        self.assertEqual(model.decoder.session.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1")
+        self.assertEqual(
+            model.decoder_with_past.session.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1"
+        )
+
+        model = ORTModelForSeq2SeqLM.from_pretrained(self.ONNX_SEQ2SEQ_MODEL_ID, use_cache=True)
+        model.to("cuda:1")
+        self.assertEqual(model.encoder.session.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1")
+        self.assertEqual(model.decoder.session.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1")
+        self.assertEqual(
+            model.decoder_with_past.session.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1"
+        )
+
     # test string device input for to()
     @require_torch_gpu
     def test_seq2seq_model_on_gpu_str(self):