Repository: huggingface/transformers
Pull request: Make benchmark-v2 to be device agnostic, to support more torch built-in devices like xpu

@remi-or , pls help review, thx very much.

Changed files:
filename: benchmark_v2/framework/benchmark_config.py
patch: @@ -7,7 +7,9 @@
 
 from transformers.generation.configuration_utils import CompileConfig
 from transformers.utils.import_utils import is_flash_attn_2_available, is_kernels_available
+from transformers.utils import is_torch_accelerator_available
 
+import torch
 
 KERNELIZATION_AVAILABLE = False
 try:
@@ -87,7 +89,7 @@ def __init__(
         self.kernelize = kernelize
         # Constant parameters
         self.dtype = "torch.bfloat16"
-        self.device = "cuda"
+        self.device = torch.accelerator.current_accelerator().type if is_torch_accelerator_available() else "cuda"
 
         self.check_validity(skip_validity_check)
         self.name = name if name is not None else self.infer_name()

filename: benchmark_v2/framework/benchmark_runner.py
patch: @@ -20,8 +20,10 @@
     AutoTokenizer,
     GenerationConfig,
     GenerationMixin,
+    is_torch_xpu_available,
 )
 from transformers.generation.streamers import BaseStreamer
+from transformers.utils import is_torch_accelerator_available
 
 from .benchmark_config import BenchmarkConfig
 from .data_classes import BenchmarkMetadata, BenchmarkResult, GPURawMetrics, pretty_print_dict
@@ -97,10 +99,13 @@ def flush_memory(flush_compile: bool = True) -> None:
             if hasattr(torch._inductor.codecache, "TritonFuture"):
                 if hasattr(torch._inductor.codecache.TritonFuture, "_compile_cache"):
                     torch._inductor.codecache.TritonFuture._compile_cache.clear()
-    # Clear CUDA cache
+    # Clear device cache
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.synchronize()
+    elif is_torch_xpu_available():
+        torch.xpu.empty_cache()
+        torch.xpu.synchronize()
     gc.collect()
 
 
@@ -156,6 +161,8 @@ def __init__(
         self._setup_for = ""
         # Attributes that are reset for each run
         self.model: GenerationMixin | None = None
+        self.device_type = torch.accelerator.current_accelerator().type if is_torch_accelerator_available() else "cuda"
+        self.torch_accelerator_module = getattr(torch, self.device_type, torch.cuda)
 
     def cleanup(self) -> None:
         del self.model
@@ -288,14 +295,20 @@ def time_generate(
         e2e_latency = wall_time_1 - wall_time_0
         timestamps = torch.tensor(timestamps).sub(wall_time_0).tolist()
         self.logger.info(
-            f"Time generate done in {e2e_latency:.2f} seconds. Memory usage: {torch.cuda.memory_allocated() / 1024**2:.2f} MB"
+            f"Time generate done in {e2e_latency:.2f} seconds. Memory usage: {self.torch_accelerator_module.memory_allocated() / 1024**2:.2f} MB"
         )
         return e2e_latency, timestamps, shape_and_decoded_output, gpu_metrics
 
     def profile_generate(self, num_tokens_to_profile: int, config_name: str) -> None:
         """Profile the latency of a call to model.generate() with the given (inputs) and (max_new_tokens)."""
+        activities = [torch.profiler.ProfilerActivity.CPU]
+        if self.device_type == "cuda":
+            activities.append(torch.profiler.ProfilerActivity.CUDA)
+        elif self.device_type == "xpu":
+            activities.append(torch.profiler.ProfilerActivity.XPU)
+
         profiler = torch.profiler.profile(
-            activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
+            activities=activities,
             record_shapes=True,
         )
         with profiler as prof:

filename: benchmark_v2/framework/hardware_metrics.py
patch: @@ -12,12 +12,16 @@
 import psutil
 import torch
 
+from transformers.utils import is_torch_accelerator_available
+
 
 # Data class to hold the hardware information
 def get_device_name_and_memory_total() -> tuple[str, float]:
     """Returns the name and memory total of GPU 0."""
-    device_name = torch.cuda.get_device_properties(0).name
-    device_memory_total = torch.cuda.get_device_properties(0).total_memory / 1024**3
+    device_type = torch.accelerator.current_accelerator().type if is_torch_accelerator_available() else "cuda"
+    torch_accelerator_module = getattr(torch, device_type, torch.cuda)
+    device_name = torch_accelerator_module.get_device_properties(0).name
+    device_memory_total = torch_accelerator_module.get_device_properties(0).total_memory / 1024**3
     return device_name, device_memory_total
 
 
@@ -61,6 +65,34 @@ def get_amd_gpu_stats() -> tuple[int, float]:
     gpu_stats.sort(key=lambda x: x[1], reverse=True)
     return int(gpu_stats[0][1]), float(gpu_stats[0][2]) / 1024**3
 
+def get_intel_xpu_stats() -> tuple[int, float]:
+    """Returns the utilization and memory used of an Intel XPU"""
+    # xpu-smi outputs CSV format: Timestamp, DeviceId, GPU Memory Utilization (%), GPU Memory Used (MiB)
+    xpu_smi_output = subprocess.check_output(["xpu-smi", "dump", "-m", "5,18", "-n", "1"])
+    lines = xpu_smi_output.decode("utf-8").strip().split("\n")
+
+    # Parse all data lines (skip header) and collect stats from all cards
+    xpu_stats = []
+    for line in lines[1:]:
+        data_line = line.split(",")
+        if len(data_line) < 4:
+            continue
+        device_id = data_line[1].strip()
+        utilization_str = data_line[2].strip()
+        memory_used_str = data_line[3].strip()
+        if utilization_str != "N/A" and memory_used_str != "N/A":
+            utilization = int(float(utilization_str))
+            memory_used_mib = float(memory_used_str)
+            xpu_stats.append((device_id, utilization, memory_used_mib))
+
+    if not xpu_stats:
+        return 0, 0.0
+ 
+    # Sort by utilization (descending) and pick the highest
+    xpu_stats.sort(key=lambda x: x[1], reverse=True)
+    device_id, utilization, memory_used_mib = xpu_stats[0]
+    memory_used_gb = memory_used_mib / 1024
+    return utilization, memory_used_gb
 
 def get_nvidia_gpu_stats() -> tuple[int, float]:
     """Returns the utilization and memory used of an NVIDIA GPU, both in percent"""
@@ -80,6 +112,8 @@ def __init__(self) -> None:
             self.get_utilization_and_memory_used = get_amd_gpu_stats
         elif "nvidia" in self.device_name.lower():
             self.get_utilization_and_memory_used = get_nvidia_gpu_stats
+        elif "intel" in self.device_name.lower() or device_type == "xpu":
+            self.get_utilization_and_memory_used = get_intel_xpu_stats
         else:
             raise RuntimeError(f"Unsupported GPU: {self.device_name}")
 
@@ -126,9 +160,11 @@ def __init__(self, sample_interval_sec: float = 0.1, logger: Logger | None = Non
         self.sample_interval_sec = sample_interval_sec
         self.logger = logger if logger is not None else logging.getLogger(__name__)
 
-        self.num_available_gpus = torch.cuda.device_count()
+        device_type = torch.accelerator.current_accelerator().type if is_torch_accelerator_available() else "cuda"
+        torch_accelerator_module = getattr(torch, device_type, torch.cuda)
+        self.num_available_gpus = torch_accelerator_module.device_count()
         if self.num_available_gpus == 0:
-            raise RuntimeError("No GPUs detected by torch.cuda.device_count().")
+            raise RuntimeError(f"No GPUs detected by torch.{device_type}.device_count().")
         self.gpu_stats_getter = GPUStatsCollector()
 
     def start(self):