profiler

analog/analog.py

@@ -17,7 +17,7 @@
 from analog.lora import LoRAHandler
 from analog.lora.utils import is_lora
 from analog.state import AnaLogState
-from analog.timer.timer import DeviceFunctionTimer
+from analog.monitor_util.timer import DeviceFunctionTimer
 from analog.utils import (
     get_logger,
     get_rank,

analog/analysis/influence_function.py

@@ -5,7 +5,7 @@
 from einops import einsum, rearrange, reduce
 from analog.config import InfluenceConfig
 from analog.state import AnaLogState
-from analog.timer.timer import DeviceFunctionTimer
+from analog.monitor_util.timer import DeviceFunctionTimer
 from analog.utils import get_logger, nested_dict
 from analog.analysis.utils import synchronize_device
 

analog/logging/log_saver.py

@@ -1,7 +1,7 @@
 from concurrent.futures import ThreadPoolExecutor
 import torch
 
-from analog.timer.timer import HostFunctionTimer, DeviceFunctionTimer
+from analog.monitor_util.timer import DeviceFunctionTimer
 from analog.utils import nested_dict, to_numpy, get_rank
 from analog.logging.mmap import MemoryMapHandler
 

analog/logging/logger.py

@@ -10,7 +10,7 @@
 from analog.logging.option import LogOption
 from analog.logging.log_saver import LogSaver
 from analog.logging.utils import compute_per_sample_gradient
-from analog.timer.timer import DeviceFunctionTimer
+from analog.monitor_util.timer import DeviceFunctionTimer
 from analog.utils import get_logger
 
 

analog/timer/__init__.py → analog/monitor_util/__init__.py

@@ -1 +1,2 @@
 from .timer import FunctionTimer, Timer
+from .profiler import memory_profiler

analog/monitor_util/profiler.py

@@ -0,0 +1,26 @@
+import torch
+import functools
+from torch.profiler import profile, ProfilerActivity
+
+
+def memory_profiler(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        activities = [ProfilerActivity.CPU]
+        if device.type == "cuda":
+            activities.append(ProfilerActivity.CUDA)
+
+        with profile(activities=activities, profile_memory=True) as prof:
+            result = func(*args, **kwargs)
+
+        print(
+            prof.key_averages().table(
+                sort_by="self_cuda_memory_usage"
+                if device.type == "cuda"
+                else "self_cpu_memory_usage"
+            )
+        )
+        return result
+
+    return wrapper

analog/timer/timer.py → analog/monitor_util/timer.py

@@ -119,7 +119,9 @@ class DeviceFunctionTimer(FunctionTimer):
     if torch.cuda.is_available():
         host_timer = False
     else:
-        logging.warning("CUDA is not set, setting the timer is set to host timer.")
+        logging.warning(
+            "CUDA is not set, setting the monitor_util is set to host monitor_util."
+        )
         host_timer = True
 
 
@@ -135,13 +137,13 @@ def __init__(self):
     def start_timer(self, name, host_timer=False):
         if host_timer:
             if name in self.timers["cpu"]:
-                logging.warning(f"timer for {name} already exist")
+                logging.warning(f"monitor_util for {name} already exist")
                 return
             start_time = time.time()
             self.timers["cpu"][name] = [start_time]
         else:
             if name in self.timers["gpu"]:
-                logging.warning(f"timer for {name} already exist")
+                logging.warning(f"monitor_util for {name} already exist")
                 return
             self.is_synchronized = False
             start_event = torch.cuda.Event(enable_timing=True)

examples/mnist_influence/compute_influences.py

@@ -11,7 +11,7 @@
     construct_mlp,
 )
 
-from analog.timer import FunctionTimer
+from analog.monitor_util import FunctionTimer
 
 parser = argparse.ArgumentParser("MNIST Influence Analysis")
 parser.add_argument("--data", type=str, default="mnist", help="mnist or fmnist")