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):
400
400
self .profiler .start ()
401
401
else :
402
402
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" ))
405
407
406
408
def execute_dummy_batch (self ) -> None :
407
409
self .model_runner ._dummy_run (1 )
Original file line number Diff line number Diff line change @@ -128,8 +128,10 @@ def stop_profile(self):
128
128
if self .profiler is None :
129
129
raise RuntimeError ("Profiler is not enabled." )
130
130
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" ))
133
135
134
136
def sleep (self , level : int = 1 ) -> None :
135
137
free_bytes_before_sleep = torch .cuda .mem_get_info ()[0 ]
You can’t perform that action at this time.
0 commit comments