Save log every step

Author: landajuela

protein_tune_rl/protein_trainer/dpo_trainer.py

@@ -156,11 +156,14 @@ def run(self, output_dir):
                 batch = self.collator(raw_batch)
                 # Perform one training step
                 current_step = self._train_step(batch, current_step, batch_idx)
+                # Log the step results
+                log_df = self._log_step(log_df, output_dir, current_step)
 
                 # Checkpointing and evaluation
                 if self._should_checkpoint(current_step, self.check_point_freq):
                     # Log the step results
-                    log_df = self._log_step(log_df, output_dir, current_step)
+                    if dist.get_rank() == 0:
+                        log_df.to_csv(f"{output_dir}/dpo_trainer_log.csv", index=False)
                     dist.barrier()
                     # Save model checkpoints
                     if dist.get_rank() == 0 and self.config["trainer"].get(
@@ -217,5 +220,4 @@ def _log_step(self, log_df, output_dir, current_step):
                 "avg_margin": self._last_avg_margin,
             }
             log_df = pd.concat([log_df, pd.DataFrame([step_data])], ignore_index=True)
-            log_df.to_csv(f"{output_dir}/dpo_trainer_log.csv", index=False)
         return log_df

protein_tune_rl/protein_trainer/dro_trainer.py

@@ -181,12 +181,12 @@ def run(self, output_dir):
 
                 # Perform one training step
                 current_step = self._train_step(batch, current_step, batch_number)
+                # Log the step results
+                log_df = self._log_step(log_df, output_dir, current_step, batch_number)
 
                 if self._should_checkpoint(current_step, self.check_point_freq):
-                    # Log the step results
-                    log_df = self._log_step(
-                        log_df, output_dir, current_step, batch_number
-                    )
+                    if dist.get_rank() == 0:
+                        log_df.to_csv(f"{output_dir}/dro_trainer_log.csv", index=False)
                     dist.barrier()
                     self._maybe_save_models(output_dir, current_step)
                     self._maybe_run_evaluation(output_dir, current_step)
@@ -234,7 +234,6 @@ def _log_step(self, log_df, output_dir, current_step, batch_number):
                 }
             )
             log_df = pd.concat([log_df, step_log_df])
-            log_df.to_csv(f"{output_dir}/dro_trainer_log.csv", index=False)
         return log_df
 
     def _maybe_save_models(self, output_dir, current_step):