From 154720ba6ee10631added5486eabd1134d480c9e Mon Sep 17 00:00:00 2001 From: hxwang Date: Tue, 28 May 2024 12:41:42 +0000 Subject: [PATCH] [chore] refactor profiler utils --- .../gpt/gemini/commons/performance_evaluator.py | 1 + examples/language/gpt/gemini/commons/utils.py | 16 ---------------- examples/language/gpt/gemini/train_gpt_demo.py | 3 ++- examples/language/performance_evaluator.py | 4 ++-- 4 files changed, 5 insertions(+), 19 deletions(-) create mode 120000 examples/language/gpt/gemini/commons/performance_evaluator.py diff --git a/examples/language/gpt/gemini/commons/performance_evaluator.py b/examples/language/gpt/gemini/commons/performance_evaluator.py new file mode 120000 index 000000000..152602774 --- /dev/null +++ b/examples/language/gpt/gemini/commons/performance_evaluator.py @@ -0,0 +1 @@ +../../../performance_evaluator.py \ No newline at end of file diff --git a/examples/language/gpt/gemini/commons/utils.py b/examples/language/gpt/gemini/commons/utils.py index 03054c0a2..ba80cc4a6 100644 --- a/examples/language/gpt/gemini/commons/utils.py +++ b/examples/language/gpt/gemini/commons/utils.py @@ -1,8 +1,6 @@ import time -from contextlib import nullcontext import torch -from torch.profiler import ProfilerActivity, profile, schedule, tensorboard_trace_handler class DummyProfiler: @@ -24,20 +22,6 @@ def get_tflops(model_numel, batch_size, seq_len, step_time): return model_numel * batch_size * seq_len * 8 / 1e12 / (step_time + 1e-12) -def get_profile_context(enable_flag, warmup_steps, active_steps, save_dir): - if enable_flag: - return profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - schedule=schedule(wait=0, warmup=warmup_steps, active=active_steps), - on_trace_ready=tensorboard_trace_handler(save_dir), - # record_shapes=True, - # profile_memory=True, - with_stack=True, - ) - else: - return nullcontext(DummyProfiler()) - - def get_time_stamp(): cur_time = time.strftime("%d-%H:%M", time.localtime()) return cur_time diff --git a/examples/language/gpt/gemini/train_gpt_demo.py b/examples/language/gpt/gemini/train_gpt_demo.py index 4911ff124..cb5d2c32c 100644 --- a/examples/language/gpt/gemini/train_gpt_demo.py +++ b/examples/language/gpt/gemini/train_gpt_demo.py @@ -8,7 +8,8 @@ import psutil import torch import torch.nn as nn from commons.model_zoo import model_builder -from commons.utils import get_data, get_profile_context, get_tflops, get_time_stamp +from commons.performance_evaluator import get_profile_context +from commons.utils import get_data, get_tflops, get_time_stamp from packaging import version import colossalai diff --git a/examples/language/performance_evaluator.py b/examples/language/performance_evaluator.py index 99df8f1da..6b8daf37d 100644 --- a/examples/language/performance_evaluator.py +++ b/examples/language/performance_evaluator.py @@ -47,8 +47,8 @@ def get_profile_context(enable_flag, warmup_steps, active_steps, save_dir): activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], schedule=schedule(wait=0, warmup=warmup_steps, active=active_steps), on_trace_ready=tensorboard_trace_handler(save_dir), - # record_shapes=True, - # profile_memory=True, + record_shapes=True, + profile_memory=True, with_stack=True, ) else: