Changes required to detect AMD GPUs

f98d9bac53ef98401c7cd47ddbe32ea5f653db79 · 2025-03-05 22:49:29 -0500 · DeftDawg

Files touched

M exo/topology/device_capabilities.py
M setup.py

Diff

commit f98d9bac53ef98401c7cd47ddbe32ea5f653db79
Author: DeftDawg <deftdawg@gmail.com>
Date:   Wed Mar 5 22:49:29 2025 -0500

    Changes required to detect AMD GPUs
---
 exo/topology/device_capabilities.py | 15 ++++++---------
 setup.py                            |  1 +
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/exo/topology/device_capabilities.py b/exo/topology/device_capabilities.py
index ae28e757..e60003f2 100644
--- a/exo/topology/device_capabilities.py
+++ b/exo/topology/device_capabilities.py
@@ -198,22 +198,19 @@ async def linux_device_capabilities() -> DeviceCapabilities:
       flops=CHIP_FLOPS.get(gpu_name, DeviceFlops(fp32=0, fp16=0, int8=0)),
     )
   elif Device.DEFAULT == "AMD":
-    # For AMD GPUs, pyrsmi is the way (Official python package for rocm-smi)
-    from pyrsmi import rocml
+    import pyamdgpuinfo
 
-    rocml.smi_initialize()
-    gpu_name = rocml.smi_get_device_name(0).upper()
-    gpu_memory_info = rocml.smi_get_device_memory_total(0)
+    gpu_raw_info = pyamdgpuinfo.get_gpu(0)
+    gpu_name = gpu_raw_info.name
+    gpu_memory_info = gpu_raw_info.memory_info["vram_size"]
 
     if DEBUG >= 2: print(f"AMD device {gpu_name=} {gpu_memory_info=}")
 
-    rocml.smi_shutdown()
-
     return DeviceCapabilities(
-      model="Linux Box ({gpu_name})",
+      model="Linux Box (" + gpu_name + ")",
       chip=gpu_name,
       memory=gpu_memory_info // 2**20,
-      flops=DeviceFlops(fp32=0, fp16=0, int8=0),
+      flops=CHIP_FLOPS.get(gpu_name, DeviceFlops(fp32=0, fp16=0, int8=0)),
     )
 
   else:
diff --git a/setup.py b/setup.py
index e9431ca4..98b8357e 100644
--- a/setup.py
+++ b/setup.py
@@ -20,6 +20,7 @@ install_requires = [
   "prometheus-client==0.20.0",
   "protobuf==5.28.1",
   "psutil==6.0.0",
+  "pyamdgpuinfo==2.1.6;platform_system=='Linux'",
   "pydantic==2.9.2",
   "requests==2.32.3",
   "rich==13.7.1",

← 013d2573 remove dead links in README · back to Exo · upgrade grpcio and grpcio-tools to 1.71.0 2857975b →