Log WandB tables on main process (#1951)

huggingface · Aug 20, 2024 · dc4cfab · dc4cfab
1 parent 66d3a82
commit dc4cfab
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 15 deletions.
diff --git a/trl/trainer/online_dpo_trainer.py b/trl/trainer/online_dpo_trainer.py
@@ -568,13 +568,14 @@ def generate_completions(self, sampling: bool = False):
                 if sampling:
                     break
         df = pd.DataFrame(table)
-        if self.accelerator.process_index == 0:
+
+        if self.accelerator.is_main_process:
             print_rich_table(df.iloc[0 : 0 + 5])
-        if "wandb" in args.report_to:
-            import wandb
+            if "wandb" in args.report_to:
+                import wandb
 
-            if wandb.run is not None:
-                wandb.log({"completions": wandb.Table(dataframe=df)})
+                if wandb.run is not None:
+                    wandb.log({"completions": wandb.Table(dataframe=df)})
 
     @wraps(Trainer.push_to_hub)
     def push_to_hub(

diff --git a/trl/trainer/ppov2_trainer.py b/trl/trainer/ppov2_trainer.py
@@ -604,13 +604,14 @@ def generate_completions(self, sampling: bool = False):
                 if sampling:
                     break
         df = pd.DataFrame(table)
-        if self.accelerator.process_index == 0:
+
+        if self.accelerator.is_main_process:
             print_rich_table(df.iloc[0 : 0 + 5])
-        if "wandb" in args.report_to:
-            import wandb
+            if "wandb" in args.report_to:
+                import wandb
 
-            if wandb.run is not None:
-                wandb.log({"completions": wandb.Table(dataframe=df)})
+                if wandb.run is not None:
+                    wandb.log({"completions": wandb.Table(dataframe=df)})
 
     @wraps(Trainer.push_to_hub)
     def push_to_hub(

diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py
@@ -502,13 +502,14 @@ def generate_completions(self, sampling: bool = False):
                 if sampling:
                     break
         df = pd.DataFrame(table)
-        if self.accelerator.process_index == 0:
+
+        if self.accelerator.is_main_process:
             print_rich_table(df.iloc[0 : 0 + 5])
-        if "wandb" in args.report_to:
-            import wandb
+            if "wandb" in args.report_to:
+                import wandb
 
-            if wandb.run is not None:
-                wandb.log({"completions": wandb.Table(dataframe=df)})
+                if wandb.run is not None:
+                    wandb.log({"completions": wandb.Table(dataframe=df)})
 
     @wraps(Trainer.push_to_hub)
     def push_to_hub(