sync from github

src/backend/hflm_with_measurement.py

@@ -37,7 +37,7 @@ from lm_eval.models.utils import (
     stop_sequences_criteria,
 )
 from lm_eval.models.huggingface import HFLM
-from src.utils import get_gpu_number, get_gpu_details, get_peak_bw, transfer_precision2bytes, get_peak_flops
+from src.utils import get_gpu_details, get_peak_bw, transfer_precision2bytes, get_peak_flops
 from src.submission.check_validity import get_model_size
 from src.envs import API
 
@@ -73,6 +73,18 @@ class HFLMWithMeasurement(HFLM):
         self.pretrained = kwargs.get("pretrained", None)
         self.revision = kwargs.get("revision", None)
         self.precision = kwargs.get("dtype", None)
+        self.num_gpus = None
+
+    def _detect_num_gpus_used(self):
+        if self.num_gpus is not None:
+            return self.num_gpus
+        gpus = []
+        for p in self.model.parameters():
+            if p.device.type == "cuda":
+                gpus.append(p.device.index)
+                
+        self.num_gpus = len(set(gpus))
+        return self.num_gpus
 
     def _loglikelihood_tokens(
         self,
@@ -352,7 +364,8 @@ class HFLMWithMeasurement(HFLM):
                 else:
                     continue
         print(f"linear_count: {linear_count}")
-        print(f"element_wise_mul: {element_wise_mul}")     
+        print(f"element_wise_mul: {element_wise_mul}")
+        print(f"GPU usage: {self._detect_num_gpus_used()}")
 
         stopping_criteria = stop_sequences_criteria(
             self.tokenizer, stop, context.shape[1], context.shape[0]
@@ -423,7 +436,7 @@ class HFLMWithMeasurement(HFLM):
         per_token_kv_size = 2 * n_layers * d_model * precision_bytes
         
         peak_bw_single = get_peak_bw(get_gpu_details())
-        peak_bw = peak_bw_single * get_gpu_number()
+        peak_bw = peak_bw_single * self._detect_num_gpus_used()
         
         context_prefill_size = context_length
         kv_size = context_prefill_size * per_token_kv_size + (output_length - 1) * per_token_kv_size / 2
@@ -441,7 +454,7 @@ class HFLMWithMeasurement(HFLM):
         avg_context_length = context_length + (output_length - 1) / 2
         flops_per_token = 2 * model_size + ((linear_count + element_wise_mul) * n_layers * avg_context_length * d_model) + 4 * d_model + 2 * d_model * n_vocab
         peak_flops_single = get_peak_flops(get_gpu_details(), self.precision)
-        peak_flops = peak_flops_single * get_gpu_number()
+        peak_flops = peak_flops_single * self._detect_num_gpus_used()
         
         ## TODO only support llama-type decoder only models and moe models of switch transformer and mixtrial
         mfu = token_per_sec * flops_per_token / peak_flops

src/utils.py

@@ -174,43 +174,6 @@ def analyze_gpu_stats(stats_list):
 
     return avg_stats
 
-def get_gpu_number():
-    visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', None)
-    if visible_devices is not None:
-        gpu_indices = visible_devices.split(',')
-    else:
-        # Query all GPU indices if CUDA_VISIBLE_DEVICES is not set
-        result = subprocess.run(['nvidia-smi', '--query-gpu=index', '--format=csv,noheader'], capture_output=True, text=True)
-        if result.returncode != 0:
-            print("Failed to query GPU indices.")
-            return []
-        gpu_indices = result.stdout.strip().split('\n')
-    # print(f"gpu_indices: {gpu_indices}")
-    gpu_stats = []
-
-    gpu_info_pattern = re.compile(r'(\d+)C\s+P\d+\s+(\d+)W / \d+W\s+\|\s+(\d+)MiB / \d+MiB\s+\|\s+(\d+)%')
-
-    for index in gpu_indices:
-        result = subprocess.run(['nvidia-smi', '-i', index], capture_output=True, text=True)
-        output = result.stdout.strip()
-        lines = output.split("\n")
-        for line in lines:
-            match = gpu_info_pattern.search(line)
-            gpu_info = {}
-            if match:
-                temp, power_usage, mem_usage, gpu_util = map(int, match.groups())
-                gpu_info.update({
-                    GPU_TEMP: temp,
-                    GPU_Power: power_usage,
-                    GPU_Mem: round(mem_usage / 1024, 2),
-                    GPU_Util: gpu_util
-                })
-
-            if len(gpu_info) >= 4:
-                gpu_stats.append(gpu_info)
-                
-    return len(gpu_stats)
-
 def get_gpu_details():
     gpus = GPUtil.getGPUs()
     gpu = gpus[0]