Commit the last step on world_process_zero in WandbCallback (#9805)

* Commit the last step on world_process_zero in WandbCallback

* Use the environment variable WANDB_LOG_MODEL as a default value in WandbCallback

src/transformers/integrations.py

@@ -516,6 +516,8 @@ class WandbCallback(TrainerCallback):
             else:
                 self._wandb = wandb
         self._initialized = False
+        # log outputs
+        self._log_model = os.getenv("WANDB_LOG_MODEL", "FALSE").upper() in ENV_VARS_TRUE_VALUES.union({"TRUE"})
 
     def setup(self, args, state, model, reinit, **kwargs):
         """
@@ -569,9 +571,6 @@ class WandbCallback(TrainerCallback):
                     model, log=os.getenv("WANDB_WATCH", "gradients"), log_freq=max(100, args.logging_steps)
                 )
 
-            # log outputs
-            self._log_model = os.getenv("WANDB_LOG_MODEL", "FALSE").upper() in ENV_VARS_TRUE_VALUES.union({"TRUE"})
-
     def on_train_begin(self, args, state, control, model=None, **kwargs):
         if self._wandb is None:
             return
@@ -583,6 +582,7 @@ class WandbCallback(TrainerCallback):
         if self._wandb is None:
             return
         # commit last step
+        if state.is_world_process_zero:
             self._wandb.log({})
         if self._log_model and self._initialized and state.is_world_process_zero:
             from .trainer import Trainer