github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[chore] refactor profiler utils

pull/5749/head
hxwang 2024-05-28 12:41:42 +00:00
parent 8547562884
commit 154720ba6e
4 changed files with 5 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
examples/language/gpt/gemini/commons/performance_evaluator.py Symbolic link
View File

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

16
examples/language/gpt/gemini/commons/utils.py
View File

@ -1,8 +1,6 @@
import time import time
from contextlib import nullcontext
import torch import torch
from torch.profiler import ProfilerActivity, profile, schedule, tensorboard_trace_handler
class DummyProfiler: 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) 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(): def get_time_stamp():
cur_time = time.strftime("%d-%H:%M", time.localtime()) cur_time = time.strftime("%d-%H:%M", time.localtime())
return cur_time return cur_time

3
examples/language/gpt/gemini/train_gpt_demo.py
View File

@ -8,7 +8,8 @@ import psutil
import torch import torch
import torch.nn as nn import torch.nn as nn
from commons.model_zoo import model_builder 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 from packaging import version
import colossalai import colossalai

4
examples/language/performance_evaluator.py
View File

@ -47,8 +47,8 @@ def get_profile_context(enable_flag, warmup_steps, active_steps, save_dir):
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
schedule=schedule(wait=0, warmup=warmup_steps, active=active_steps), schedule=schedule(wait=0, warmup=warmup_steps, active=active_steps),
on_trace_ready=tensorboard_trace_handler(save_dir), on_trace_ready=tensorboard_trace_handler(save_dir),
# record_shapes=True, record_shapes=True,
# profile_memory=True, profile_memory=True,
with_stack=True, with_stack=True,
) )
else: else: