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