Change progress logging to once across all nodes

huggingface · Jan 9, 2024 · 5b67725 · 5b67725
1 parent 976189a
commit 5b67725
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/src/transformers/trainer_callback.py b/src/transformers/trainer_callback.py
@@ -489,42 +489,42 @@ def __init__(self):
         self.prediction_bar = None
 
     def on_train_begin(self, args, state, control, **kwargs):
-        if state.is_local_process_zero:
+        if state.is_world_process_zero:
             self.training_bar = tqdm(total=state.max_steps, dynamic_ncols=True)
         self.current_step = 0
 
     def on_step_end(self, args, state, control, **kwargs):
-        if state.is_local_process_zero:
+        if state.is_world_process_zero:
             self.training_bar.update(state.global_step - self.current_step)
             self.current_step = state.global_step
 
     def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):
-        if state.is_local_process_zero and has_length(eval_dataloader):
+        if state.is_world_process_zero and has_length(eval_dataloader):
             if self.prediction_bar is None:
                 self.prediction_bar = tqdm(
                     total=len(eval_dataloader), leave=self.training_bar is None, dynamic_ncols=True
                 )
             self.prediction_bar.update(1)
 
     def on_evaluate(self, args, state, control, **kwargs):
-        if state.is_local_process_zero:
+        if state.is_world_process_zero:
             if self.prediction_bar is not None:
                 self.prediction_bar.close()
             self.prediction_bar = None
 
     def on_predict(self, args, state, control, **kwargs):
-        if state.is_local_process_zero:
+        if state.is_world_process_zero:
             if self.prediction_bar is not None:
                 self.prediction_bar.close()
             self.prediction_bar = None
 
     def on_log(self, args, state, control, logs=None, **kwargs):
-        if state.is_local_process_zero and self.training_bar is not None:
+        if state.is_world_process_zero and self.training_bar is not None:
             _ = logs.pop("total_flos", None)
             self.training_bar.write(str(logs))
 
     def on_train_end(self, args, state, control, **kwargs):
-        if state.is_local_process_zero:
+        if state.is_world_process_zero:
             self.training_bar.close()
             self.training_bar = None