Merge branch '25-fix-multi-gpu-online-evaluation' into 'main'

Resolve "Fix multi-GPU online evaluation"

Closes #25

See merge request optlm/protein_tune_rl!23

Author: landajuela

protein_tune_rl/collator/dro_collator.py

@@ -169,6 +169,7 @@ def __call__(self, batch):
 
         if self.eval:
             return {
+                "__row_idx__": batch["__row_idx__"],
                 "input_ids": tokenized_masked_prompts_with_completions,
                 "prompts": tokenized_masked_prompts,
                 "labels": input_mask,
@@ -178,6 +179,7 @@ def __call__(self, batch):
             }
 
         return {
+            "__row_idx__": batch["__row_idx__"],
             "input_ids": tokenized_masked_prompts_with_completions,
             "prompts": tokenized_masked_prompts,
             "labels": input_mask,

protein_tune_rl/dataset/dro_dataset.py

@@ -9,6 +9,7 @@ def __init__(self, data_directory, chain, region, reward):
     def __getitem__(self, idx):
 
         return {
+            "__row_idx__": int(idx),
             "prompts": self.data[self.chain].iloc[idx],
             "completions": self.data[self.region].iloc[idx],
             "rewards": float(self.data[self.reward].iloc[idx]),
@@ -23,6 +24,7 @@ def __init__(self, data_directory, chain, region):
     def __getitem__(self, idx):
 
         return {
+            "__row_idx__": int(idx),
             "prompts": self.data[self.chain].iloc[idx],
             "completions": self.data[self.region].iloc[idx],
             "LC": self.data.LC.iloc[idx],

protein_tune_rl/dataset/infilling_dataset.py

@@ -14,6 +14,7 @@ def __len__(self):
     def __getitem__(self, idx):
 
         return {
+            "__row_idx__": int(idx),
             "prompts": self.data[self.chain].iloc[idx],
             "region": self.data[self.region].iloc[idx],
             "LC": self.data.LC.iloc[idx],

protein_tune_rl/protein_evaluator/iglm_evaluator.py

@@ -246,8 +246,11 @@ def run_with_ground_truth(self, output_dir=None):
             'generated_sequences': [],
             'heavy_chains': [],
             'light_chains': [],
+            '__row_idx__': [],
         }
 
+        self._log_dataset_info()
+
         for batch_number, batch in enumerate(iter(self.dataloader)):
             self.policy.eval()
 
@@ -272,6 +275,19 @@ def run_with_ground_truth(self, output_dir=None):
         eval_df = self._create_evaluation_dataframe(results)
         return gather_dataframes(eval_df, device=self.device)
 
+    def _log_dataset_info(self):
+        dataloader = self.dataloader
+        ddp_enabled = dist.is_available() and dist.is_initialized()
+        world_size = dist.get_world_size() if ddp_enabled else 1
+        sampler = getattr(dataloader, "sampler", None)
+        samples_per_rank = len(sampler) if sampler else len(dataloader.dataset)
+        batches_per_rank = len(dataloader)
+        logger.info(
+            f"Eval: world_size={world_size}, batch_size=1, "
+            f"per_rank={samples_per_rank} samples/{batches_per_rank} batches, "
+            f"global_batches_per_epoch={batches_per_rank * world_size}"
+        )
+
     def _generate_sequences_if_needed(self, tokenized_batch):
         """Generate sequences if any metric requires generated sequences."""
         if not any(self.metric_use_generated):
@@ -405,6 +421,7 @@ def _collect_sample_results(
             + "[MASK]"
             + tokenized_batch["seq_post_mask"][0]
         )
+        results['__row_idx__'].append(int(tokenized_batch["__row_idx__"][0]))
 
     def _create_evaluation_dataframe(self, results):
         """Create DataFrame from collected results."""
@@ -413,6 +430,7 @@ def _create_evaluation_dataframe(self, results):
         eval_df['HC'] = results['heavy_chains']
         eval_df['LC'] = results['light_chains']
         eval_df['prompts'] = results['prompts']
+        eval_df['__row_idx__'] = results['__row_idx__']
 
         for idx, metric in enumerate(self.config['metric']):
             eval_df[str(metric['name'])] = [

protein_tune_rl/protein_trainer/dro_trainer.py

@@ -161,6 +161,7 @@ def run_evaluation(self, output_dir, current_step):
             eval_df = self.evaluator.run_with_ground_truth()
 
         if dist.get_rank() == 0 and eval_df is not None:
+            eval_df = eval_df.sort_values("__row_idx__").reset_index(drop=True)
             eval_df.to_csv(
                 f"{output_dir}/evaluation_results_step_{current_step}.csv",
                 index=False,
@@ -173,67 +174,95 @@ def run(self, output_dir):
         """Run the DRO Trainer for the specified number of optimization steps."""
         log_df = pd.DataFrame()
 
-        logger.info(
-            f"Breaking down the training dataset into {len(self.dataloader)} batches."
-        )
+        self._log_dataset_info()
 
         current_step = 0
         while current_step < self.total_optimization_steps:
             for batch_number, batch in enumerate(iter(self.dataloader)):
-                self.value.train()
-                self.policy.train()
+                current_step = self._train_step(batch, current_step, batch_number)
+                self._log_step(log_df, output_dir, current_step, batch_number)
+                dist.barrier()
 
-                self.policy_optimizer.zero_grad()
-                self.value_optimizer.zero_grad()
+                if self._should_checkpoint(current_step):
+                    self._maybe_save_models(output_dir, current_step)
+                    self._maybe_run_evaluation(output_dir, current_step)
 
-                tokenized_batch = self.collator(batch)
+                if current_step >= self.total_optimization_steps:
+                    break
 
-                policy_loss, value_loss = self.model_optimizer.calculate_loss(
-                    tokenized_batch
-                )
+        self._final_save(output_dir)
+        return log_df
 
-                value_loss.backward()
-                policy_loss.backward()
+    def _log_dataset_info(self):
+        dl = self.dataloader
+        world = (
+            dist.get_world_size()
+            if dist.is_available() and dist.is_initialized()
+            else 1
+        )
+        sampler = getattr(dl, "sampler", None)
 
-                self.policy_optimizer.step()
-                self.value_optimizer.step()
+        per_rank_samples = len(sampler) if sampler is not None else len(dl.dataset)
+        per_rank_batches = len(dl)
 
-                current_step += 1
+        logger.info(
+            f"Per-rank: {per_rank_samples} samples → {per_rank_batches} batches "
+            f"(batch size={dl.batch_size}, drop_last={dl.drop_last}); "
+            f"Global: world_size={world}, effective batch size={dl.batch_size * world}, "
+            f"batches/epoch={per_rank_batches * world}."
+        )
 
-                logger.info(
-                    f"Step {current_step}, Batch {batch_number + 1}: "
-                    f"Policy Loss: {policy_loss.item():.4f}, "
-                    f"Value Loss: {value_loss.item():.4f}"
-                )
+    def _train_step(self, batch, current_step, batch_number):
+        """Perform a single training step on the provided batch."""
+        self.value.train()
+        self.policy.train()
 
-                if dist.get_rank() == 0:
-                    step_log_df = pd.DataFrame.from_dict(
-                        {
-                            "step": [current_step],
-                            "policy_loss": [policy_loss.item()],
-                            "value_loss": [value_loss.item()],
-                        }
-                    )
+        self.policy_optimizer.zero_grad()
+        self.value_optimizer.zero_grad()
 
-                    log_df = pd.concat([log_df, step_log_df])
-                    log_df.to_csv(f"{output_dir}/dro_trainer_log.csv", index=False)
-                dist.barrier()
+        tokenized_batch = self.collator(batch)
 
-                if (current_step % self.check_point_freq == 0) and (current_step > 0):
+        policy_loss, value_loss = self.model_optimizer.calculate_loss(tokenized_batch)
 
-                    if self.config["trainer"].get("save_models", True):
-                        if dist.get_rank() == 0:
-                            self.save_models(output_dir, current_step)
-                        dist.barrier()
+        value_loss.backward()
+        policy_loss.backward()
 
-                    # Run online evaluation if configured
-                    if self.config["trainer"].get("evaluate_during_training", False):
-                        self.run_evaluation(output_dir, current_step)
+        self.policy_optimizer.step()
+        self.value_optimizer.step()
 
-                if current_step >= self.total_optimization_steps:
-                    break
+        logger.info(
+            f"Step {current_step + 1}, Batch {batch_number + 1}: Policy Loss: {policy_loss.item():.4f}, Value Loss: {value_loss.item():.4f}"
+        )
 
-        # Final save after training completes
-        self.policy.module.save(output_dir / "models/final")
+        self._last_policy_loss = policy_loss
+        self._last_value_loss = value_loss
 
-        return log_df
+        return current_step + 1
+
+    def _log_step(self, log_df, output_dir, current_step, batch_number):
+        if dist.get_rank() == 0:
+            step_log_df = pd.DataFrame.from_dict(
+                {
+                    "step": [current_step],
+                    "policy_loss": [self._last_policy_loss.item()],
+                    "value_loss": [self._last_value_loss.item()],
+                }
+            )
+            log_df = pd.concat([log_df, step_log_df])
+            log_df.to_csv(f"{output_dir}/dro_trainer_log.csv", index=False)
+
+    def _should_checkpoint(self, current_step):
+        return (current_step % self.check_point_freq == 0) and (current_step > 0)
+
+    def _maybe_save_models(self, output_dir, current_step):
+        if self.config["trainer"].get("save_models", True):
+            if dist.get_rank() == 0:
+                self.save_models(output_dir, current_step)
+            dist.barrier()
+
+    def _maybe_run_evaluation(self, output_dir, current_step):
+        if self.config["trainer"].get("evaluate_during_training", False):
+            self.run_evaluation(output_dir, current_step)
+
+    def _final_save(self, output_dir):
+        self.policy.module.save(output_dir / "models/final")