From 03cde49eba0580ed17f9ae2250832fd8af4ed756 Mon Sep 17 00:00:00 2001 From: chuanqiw Date: Mon, 22 Jul 2024 07:16:26 -0700 Subject: [PATCH] Align xpu models batch size with A100 (#2378) Summary: To align xpu batchsize for dynamobenchmark torchbench suite Pull Request resolved: https://github.com/pytorch/benchmark/pull/2378 Reviewed By: aaronenyeshi Differential Revision: D59961717 Pulled By: xuzhao9 fbshipit-source-id: c926d6d14d8b979284aa465738132c17425dafe9 --- torchbenchmark/util/model.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/torchbenchmark/util/model.py b/torchbenchmark/util/model.py index 1c37b9fdb2..83c41fe79f 100644 --- a/torchbenchmark/util/model.py +++ b/torchbenchmark/util/model.py @@ -31,7 +31,7 @@ ) from torchbenchmark.util.input import input_cast, ModelInputDescriptor -SPECIAL_DEVICE_MAPPING = {"AMD Instinct MI210": "NVIDIA A100-SXM4-40GB"} +SPECIAL_DEVICE_MAPPING = {"AMD Instinct MI210": "NVIDIA A100-SXM4-40GB", "Intel(R) Data Center GPU Max 1100": "NVIDIA A100-SXM4-40GB", "Intel(R) Data Center GPU Max 1550": "NVIDIA A100-SXM4-40GB"} class PostInitProcessor(type): @@ -211,9 +211,7 @@ def _determine_dynamic_num_batches( return 1 def _get_batch_size_from_metadata(self) -> Optional[str]: - if self.device != "cuda": - current_device_name = str(self.device) - else: + if self.device == "cuda": current_device_name = ( torch.cuda.get_device_name() if torch.cuda.get_device_name() @@ -221,6 +219,16 @@ def _get_batch_size_from_metadata(self) -> Optional[str]: ) if current_device_name in SPECIAL_DEVICE_MAPPING: current_device_name = SPECIAL_DEVICE_MAPPING[current_device_name] + elif self.device == "xpu": + current_device_name = ( + torch.xpu.get_device_name() + if torch.xpu.get_device_name() + else "UNKNOWN" + ) + if current_device_name in SPECIAL_DEVICE_MAPPING: + current_device_name = SPECIAL_DEVICE_MAPPING[current_device_name] + else: + current_device_name = str(self.device) # use the device suggestion on CUDA inference tests, key should be either eval_batch_size or train_batch_size device_batch_size_key = f"{self.test}_batch_size"