Skip to content

GitLab

Commit 06b52e5b authored by silencealiang's avatar silencealiang
Browse files

采集prof时kernel显示完整名称以及显示shape

parent 60ffce35
Pipeline #2218 passed with stage
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
megatron/training/training.py
View file @ 06b52e5b
...@@ -1424,7 +1424,7 @@ def train(forward_step_func, model, optimizer, opt_param_scheduler, ...@@ -1424,7 +1424,7 @@ def train(forward_step_func, model, optimizer, opt_param_scheduler,
group_by_stack_n=5).table(sort_by="self_cuda_time_total", group_by_stack_n=5).table(sort_by="self_cuda_time_total",
row_limit=-1, row_limit=-1,
max_src_column_width=100, max_src_column_width=100,
max_name_column_width=250, max_name_column_width=280,
max_shapes_column_width=200)) max_shapes_column_width=200))
p.export_chrome_trace("{path}/trace_rank{rank}_step{step}.json".format( p.export_chrome_trace("{path}/trace_rank{rank}_step{step}.json".format(
...@@ -1440,7 +1440,7 @@ def train(forward_step_func, model, optimizer, opt_param_scheduler, ...@@ -1440,7 +1440,7 @@ def train(forward_step_func, model, optimizer, opt_param_scheduler,
warmup=1 if args.profile_step_start > 0 else 0, warmup=1 if args.profile_step_start > 0 else 0,
active=args.profile_step_end-args.profile_step_start, active=args.profile_step_end-args.profile_step_start,
repeat=1), repeat=1),
#record_shapes=True, record_shapes=True,
#on_trace_ready=torch.profiler.tensorboard_trace_handler('./torch_prof_data')) #on_trace_ready=torch.profiler.tensorboard_trace_handler('./torch_prof_data'))
on_trace_ready=trace_handler) on_trace_ready=trace_handler)
prof.start() prof.start()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment