ud

Author: LovelyBuggies

comlrl/trainers/actor_critic/ac_base.py

@@ -368,18 +368,22 @@ def evaluate(self) -> Dict[str, float]:
         turn_groups: Dict[int, List[Any]] = {}
         seen = 0
 
-        with torch.no_grad():
-            for batch in dataloader:
-                for item in batch:
-                    rollouts = self._collect_rollouts(item)
-                    for sample in rollouts:
-                        t_idx = int(sample.metadata.get("turn_idx", 0))
-                        turn_groups.setdefault(t_idx, []).append(sample)
-                    seen += 1
+        self._in_eval = True
+        try:
+            with torch.no_grad():
+                for batch in dataloader:
+                    for item in batch:
+                        rollouts = self._collect_rollouts(item)
+                        for sample in rollouts:
+                            t_idx = int(sample.metadata.get("turn_idx", 0))
+                            turn_groups.setdefault(t_idx, []).append(sample)
+                        seen += 1
+                        if seen >= num_samples:
+                            break
                     if seen >= num_samples:
                         break
-                if seen >= num_samples:
-                    break
+        finally:
+            self._in_eval = False
 
         eval_log: Dict[str, float] = {}
         for turn_idx, samples in sorted(turn_groups.items()):

comlrl/trainers/actor_critic/iac.py

@@ -465,7 +465,9 @@ def _generate_rollout(
             "do_sample": True,
             "temperature": self.args.temperature,
             "top_p": self.args.top_p,
-            "num_return_sequences": num_ret,
+            "num_return_sequences": (
+                1 if bool(getattr(self, "_in_eval", False)) else num_ret
+            ),
             "num_beams": 1,
         }
         if self.args.top_k is not None:
@@ -558,7 +560,11 @@ def _collect_rollouts(self, item: Dict[str, Any]) -> List[RolloutSample]:
         if num_turns > 1:
             return self._collect_rollouts_multi_turn(item, num_turns)
 
-        num_ret = int(getattr(self.args, "num_generations", 1))
+        num_ret = (
+            1
+            if bool(getattr(self, "_in_eval", False))
+            else int(getattr(self.args, "num_generations", 1))
+        )
         turn_prompts = [
             self._resolve_turn_prompt(item, agent_idx)
             for agent_idx in range(self.args.num_agents)

comlrl/trainers/actor_critic/maac.py

@@ -435,7 +435,11 @@ def _generate(self, agent_model, prompt: str, agent_idx: int) -> Dict[str, Any]:
         prompt_attention_mask = encoded_prompt["attention_mask"]
         prompt_len = prompt_input_ids.size(1)
 
-        num_ret = int(self.args.num_generations)
+        num_ret = (
+            1
+            if bool(getattr(self, "_in_eval", False))
+            else int(self.args.num_generations)
+        )
         generation_kwargs: Dict[str, Any] = {
             "input_ids": prompt_input_ids,
             "attention_mask": prompt_attention_mask,

comlrl/trainers/reinforce/magrpo.py

@@ -495,18 +495,21 @@ def get_eval_dataloader(self) -> Optional[DataLoader]:
             num_workers=0,
         )
 
-    def evaluate(self, num_eval_samples: int = 4) -> Dict[str, float]:
+    def evaluate(self, num_eval_samples: Optional[int] = None) -> Dict[str, float]:
         """
         Unified evaluation that supports both single-turn and multi-turn.
 
         Args:
-            num_eval_samples: Number of samples to evaluate
+            num_eval_samples: Number of samples to evaluate. Defaults to args.eval_num_samples.
 
         Returns:
             Dictionary containing evaluation metrics
         """
         if self.eval_dataset is None:
             return {}
+        if num_eval_samples is None:
+            num_eval_samples = int(getattr(self.args, "eval_num_samples", 4))
+        num_eval_samples = int(num_eval_samples)
 
         # Storage for completions across turns for all agents
         all_agent_completions_turns = [[] for _ in range(self.num_agents)]

docs/content/docs/dev/changelog.md

@@ -10,7 +10,7 @@ weight: 3
 - Remove the redundant sampling hyperparameters in algorithms.
 - Allow multi-gpu training with MP.
 
-## Version 1.3.6
+## Latest Changes
 
 - Fixed critical bug of loading heterogeneous models and reform the model loading logics.
 - Polish the docs.

docs/content/docs/user-guide/training-parallelization.md

@@ -32,3 +32,7 @@ CUDA_VISIBLE_DEVICES=0,1,2,3 python train_iac.py
     iac.agent_devices='["cuda:0","cuda:1"]'
     iac.critic_devices='["cuda:2","cuda:3"]'
 ```
+
+{{% hint note %}}
+Note that when `parallel_training=mp`, even if the same models with same sampling are used on the same seed, the training is not deterministic due to the non-deterministic GPU scheduling and aggregation on CPU.
+{{% /hint %}}

tests/test_distributed_metrics.py

@@ -57,3 +57,35 @@ def _fake_log(metrics, step):  # noqa: ARG001
     trainer._log_metrics({"loss": 1.0})
     trainer._log_metrics({"reward": 2.0})
     assert called["log"] == [{"loss": 1.0}, {"reward": 2.0}]
+
+
+def test_evaluate_calls_collect_rollouts_with_eval_flag():
+    class _EvalDummyTrainer(ActorCriticTrainerBase):
+        def __init__(self):
+            self.eval_dataset = [{"id": 0}, {"id": 1}]
+            self.args = SimpleNamespace(eval_batch_size=1, eval_num_samples=1)
+            self.wandb_initialized = False
+            self.env_step = 0
+            self.dist_env = _ctx()
+            self.verbose = False
+            self._collect_calls = 0
+            self.eval_flags = []
+
+        def _collect_rollouts(self, item):  # noqa: ARG002
+            self._collect_calls += 1
+            self.eval_flags.append(bool(getattr(self, "_in_eval", False)))
+            return [
+                SimpleNamespace(
+                    metadata={},
+                    reward=torch.tensor([1.0]),
+                    returns=torch.tensor([1.5]),
+                    old_value=torch.tensor([0.5]),
+                )
+            ]
+
+    trainer = _EvalDummyTrainer()
+    metrics = trainer.evaluate()
+    assert trainer._collect_calls == 1
+    assert trainer.eval_flags == [True]
+    assert trainer._in_eval is False
+    assert "eval/turn_1/reward_mean" in metrics