Support for ROCm 6.1

ml-energy · Jun 6, 2024 · 7d41581 · 7d41581
1 parent 9d2ba6f
commit 7d41581
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/zeus/device/gpu/__init__.py b/zeus/device/gpu/__init__.py
@@ -5,7 +5,7 @@
 
 !!! Important
     In theory, any NVIDIA GPU would be supported.
-    On the other hand, for AMD GPUs, we currently only support ROCm 6.0 and later.
+    On the other hand, for AMD GPUs, we currently only support ROCm 6.1 and later.
 
 ## Getting handles to GPUs
 

diff --git a/zeus/device/gpu/amd.py b/zeus/device/gpu/amd.py
@@ -69,7 +69,10 @@ def __init__(self, gpu_index: int) -> None:
         """Initialize the GPU object."""
         super().__init__(gpu_index)
         self._get_handle()
-        self._supportsGetTotalEnergyConsumption = None
+        # XXX(Jae-Won): Right now, the energy API's unit is broken (either the
+        # `power` field or the `counter_resolution` field). Before that, we're
+        # disabling the energy API.
+        self._supportsGetTotalEnergyConsumption = False
 
     _exception_map = {
         1: gpu_common.ZeusGPUInvalidArgError,  # amdsmi.amdsmi_wrapper.AMDSMI_STATUS_INVAL
@@ -258,7 +261,7 @@ class AMDGPUs(gpu_common.GPUs):
     """AMD GPU Manager object, containing individual AMDGPU objects, abstracting amdsmi calls and handling related exceptions.
 
     !!! Important
-        Currently only ROCM 6.0 is supported.
+        Currently only ROCm >= 6.1 is supported.
 
     `HIP_VISIBLE_DEVICES` environment variable is respected if set.
     For example, if there are 4 GPUs on the node and `HIP_VISIBLE_DEVICES=0,2`,