[chore] refactor profiler utils

examples/language/gpt/gemini/commons/performance_evaluator.py

@@ -0,0 +1 @@
+../../../performance_evaluator.py

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

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

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: