support parallel reward function

CarperAI · Dec 7, 2023 · 2cf8af5 · 2cf8af5
1 parent 6838580
commit 2cf8af5
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/trlx/trainer/accelerate_base_trainer.py b/trlx/trainer/accelerate_base_trainer.py
@@ -421,9 +421,11 @@ def evaluate(self):  # noqa: C901
 
                 if self.accelerator.is_main_process:
                     columns = ["prompt", "output"]
-                    columns_data = [str_prompts, str_outputs]
-                    if not self.config.train.reward_only_in_main_process:
-                        columns_data = self.accelerator.gather_for_metrics(columns_data)
+
+                # gather should be invoked in every process, not just the main process
+                columns_data = [str_prompts, str_outputs]
+                if not self.config.train.reward_only_in_main_process:
+                    columns_data = self.accelerator.gather_for_metrics(columns_data)
 
                 metadata, *xs = all_metadata
                 for k in metadata:
@@ -447,9 +449,11 @@ def evaluate(self):  # noqa: C901
                     else:
                         rewards = torch.tensor(rewards, dtype=float)
 
+                    # gather should be invoked in every process, not just the main process
+                    if not self.config.train.reward_only_in_main_process:
+                        rewards = self.accelerator.gather(rewards)
+
                     if self.accelerator.is_main_process:
-                        if not self.config.train.reward_only_in_main_process:
-                            rewards = self.accelerator.gather(rewards)
                         mean_reward = rewards.mean().item()
 
                         columns.append("reward")