File tree Expand file tree Collapse file tree 2 files changed +8
-4
lines changed
Expand file tree Collapse file tree 2 files changed +8
-4
lines changed Original file line number Diff line number Diff line change @@ -400,8 +400,10 @@ def profile(self, is_start: bool = True):
400400 self .profiler .start ()
401401 else :
402402 self .profiler .stop ()
403- print (self .profiler .key_averages ().table (
404- sort_by = "self_cuda_time_total" ))
403+ # only print profiler results on rank 0
404+ if self .local_rank == 0 :
405+ print (self .profiler .key_averages ().table (
406+ sort_by = "self_cuda_time_total" ))
405407
406408 def execute_dummy_batch (self ) -> None :
407409 self .model_runner ._dummy_run (1 )
Original file line number Diff line number Diff line change @@ -128,8 +128,10 @@ def stop_profile(self):
128128 if self .profiler is None :
129129 raise RuntimeError ("Profiler is not enabled." )
130130 self .profiler .stop ()
131- print (
132- self .profiler .key_averages ().table (sort_by = "self_cuda_time_total" ))
131+ # only print profiler results on rank 0
132+ if self .local_rank == 0 :
133+ print (self .profiler .key_averages ().table (
134+ sort_by = "self_cuda_time_total" ))
133135
134136 def sleep (self , level : int = 1 ) -> None :
135137 free_bytes_before_sleep = torch .cuda .mem_get_info ()[0 ]
You can’t perform that action at this time.
0 commit comments