From 196ff1181ebd7dcb6c14d592627ed81c8b434a69 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Thu, 21 Sep 2023 15:20:31 -0400
Subject: [PATCH] skip the gpu memory checks if the device is set to 'auto'
 (#609)

* skip the gpu memory checks if the device is set to 'auto'

* skip gpu mem logging if cpu too

* don't worry about log_gpu_memory_usage since it calls another annotated fn

* rename decorator internal
---
 src/axolotl/utils/bench.py | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/axolotl/utils/bench.py b/src/axolotl/utils/bench.py
index b460b2ba7c..685be526f0 100644
--- a/src/axolotl/utils/bench.py
+++ b/src/axolotl/utils/bench.py
@@ -1,14 +1,40 @@
 """Benchmarking and measurement utilities"""
+import functools
 
 import pynvml
 import torch
 from pynvml.nvml import NVMLError
 
 
+def check_cuda_device(default_value):
+    """
+    wraps a function and returns the default value instead of running the
+    wrapped function if cuda isn't available or the device is auto
+    :param default_value:
+    :return:
+    """
+
+    def deco(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            device = kwargs.get("device", args[0] if args else None)
+
+            if not torch.cuda.is_available() or device == "auto" or device == "cpu":
+                return default_value
+
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return deco
+
+
+@check_cuda_device(0.0)
 def gpu_memory_usage(device=0):
     return torch.cuda.memory_allocated(device) / 1024.0**3
 
 
+@check_cuda_device((0.0, 0.0, 0.0))
 def gpu_memory_usage_all(device=0):
     usage = torch.cuda.memory_allocated(device) / 1024.0**3
     reserved = torch.cuda.memory_reserved(device) / 1024.0**3
@@ -16,6 +42,7 @@ def gpu_memory_usage_all(device=0):
     return usage, reserved - usage, max(0, smi - reserved)
 
 
+@check_cuda_device(0.0)
 def gpu_memory_usage_smi(device=0):
     if isinstance(device, torch.device):
         device = device.index
@@ -31,9 +58,6 @@ def gpu_memory_usage_smi(device=0):
 
 
 def log_gpu_memory_usage(log, msg, device):
-    if not torch.cuda.is_available() or device == "auto":
-        return (0, 0, 0)
-
     usage, cache, misc = gpu_memory_usage_all(device)
     extras = []
     if cache > 0: