From a9a38531503c0f97bb730401ac9784ca60f85dca Mon Sep 17 00:00:00 2001 From: Minsoo Patrick Kang Date: Sun, 25 Feb 2024 00:16:35 -0800 Subject: [PATCH] profiler --- analog/analog.py | 2 +- analog/analysis/influence_function.py | 2 +- analog/logging/log_saver.py | 2 +- analog/logging/logger.py | 2 +- analog/{timer => monitor_util}/__init__.py | 1 + analog/monitor_util/profiler.py | 26 +++++++++++++++++++ analog/{timer => monitor_util}/timer.py | 8 +++--- .../mnist_influence/compute_influences.py | 2 +- 8 files changed, 37 insertions(+), 8 deletions(-) rename analog/{timer => monitor_util}/__init__.py (51%) create mode 100644 analog/monitor_util/profiler.py rename analog/{timer => monitor_util}/timer.py (95%) diff --git a/analog/analog.py b/analog/analog.py index d74ea089..cc4c0b7a 100644 --- a/analog/analog.py +++ b/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, diff --git a/analog/analysis/influence_function.py b/analog/analysis/influence_function.py index 24d54f27..2834cbc1 100644 --- a/analog/analysis/influence_function.py +++ b/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 diff --git a/analog/logging/log_saver.py b/analog/logging/log_saver.py index e3326597..e95d5ced 100644 --- a/analog/logging/log_saver.py +++ b/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 diff --git a/analog/logging/logger.py b/analog/logging/logger.py index edff8987..a3d6b3eb 100644 --- a/analog/logging/logger.py +++ b/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 diff --git a/analog/timer/__init__.py b/analog/monitor_util/__init__.py similarity index 51% rename from analog/timer/__init__.py rename to analog/monitor_util/__init__.py index 6187990c..45761224 100644 --- a/analog/timer/__init__.py +++ b/analog/monitor_util/__init__.py @@ -1 +1,2 @@ from .timer import FunctionTimer, Timer +from .profiler import memory_profiler diff --git a/analog/monitor_util/profiler.py b/analog/monitor_util/profiler.py new file mode 100644 index 00000000..647fe440 --- /dev/null +++ b/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 diff --git a/analog/timer/timer.py b/analog/monitor_util/timer.py similarity index 95% rename from analog/timer/timer.py rename to analog/monitor_util/timer.py index 923255b6..fed99c20 100644 --- a/analog/timer/timer.py +++ b/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) diff --git a/examples/mnist_influence/compute_influences.py b/examples/mnist_influence/compute_influences.py index 10120a6b..bf974cff 100644 --- a/examples/mnist_influence/compute_influences.py +++ b/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")