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

megatron/training/training.py

@@ -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", 
                                                                row_limit=-1, 
                                                                max_src_column_width=100,
-                                                               max_name_column_width=250, 
+                                                               max_name_column_width=280, 
                                                                max_shapes_column_width=200))
                 
             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,
             warmup=1 if args.profile_step_start > 0 else 0,
             active=args.profile_step_end-args.profile_step_start,
             repeat=1),
-        #record_shapes=True, 
+        record_shapes=True, 
         #on_trace_ready=torch.profiler.tensorboard_trace_handler('./torch_prof_data'))
         on_trace_ready=trace_handler)
         prof.start()