Fix repeat times in evaluation

Author: chenyushuo

tests/trainer/trainer_test.py

@@ -90,6 +90,8 @@ def test_trainer(self):
         self.config.buffer.explorer_input.eval_tasksets.append(
             get_unittest_dataset_config("copy_countdown", "test")
         )
+        self.config.buffer.explorer_input.eval_tasksets[0].eval_at_k = [1, 2]
+        self.config.buffer.explorer_input.eval_tasksets[1].eval_at_k = [3, 4]
         self.config.trainer.save_interval = 4
         self.config.check_and_update()
         _trainer_config = self.config.trainer.trainer_config
@@ -147,15 +149,17 @@ def test_trainer(self):
         self.config.check_and_update()
         bench(self.config)
         parser = TensorBoardParser(os.path.join(self.config.monitor.cache_dir, "tensorboard"))
+        eval_tasksets = self.config.buffer.explorer_input.eval_tasksets
         for prefix in ["eval", "bench"]:
-            countdown_metrics = parser.metric_list(f"{prefix}/countdown")
-            copy_countdown_metrics = parser.metric_list(f"{prefix}/copy_countdown")
-            self.assertTrue(len(countdown_metrics) > 0)
-            self.assertTrue(len(copy_countdown_metrics) > 0)
-            countdown_metric_steps = parser.metric_steps(countdown_metrics[0])
-            countdown_copy_metric_steps = parser.metric_steps(copy_countdown_metrics[0])
-            self.assertEqual([0, 4, 8], countdown_metric_steps)
-            self.assertEqual([0, 4, 8], countdown_copy_metric_steps)
+            for eval_taskset, taskset_name in zip(eval_tasksets, ["countdown", "copy_countdown"]):
+                metrics = parser.metric_list(f"{prefix}/{taskset_name}")
+                self.assertTrue(len(metrics) > 0)
+                for eval_stats in ["mean", "best", "worst"]:
+                    for k in eval_taskset.eval_at_k:
+                        for stats in ["mean", "std"]:
+                            metric_name = f"{prefix}/{taskset_name}/score/{eval_stats}@{k}/{stats}"
+                            metric_steps = parser.metric_steps(metric_name)
+                            self.assertEqual(metric_steps, [0, 4, 8])
 
     def tearDown(self):
         # remove dir only when the test passed
@@ -969,6 +973,7 @@ def test_trainer(self):
         self.config.buffer.explorer_input.eval_tasksets.append(
             get_unittest_dataset_config("gsm8k", "test")
         )
+        self.config.buffer.explorer_input.eval_tasksets[0].eval_at_k = [1, 2, 4]
         self.config.model.model_path = get_model_path()
         self.config.algorithm.algorithm_type = "grpo"
         self.config.algorithm.advantage_fn = "grpo"
@@ -1016,11 +1021,16 @@ def test_trainer(self):
         self.config.check_and_update()
         bench(self.config)
         parser = TensorBoardParser(os.path.join(self.config.monitor.cache_dir, "tensorboard"))
+        eval_taskset = self.config.buffer.explorer_input.eval_tasksets[0]
         for prefix in ["eval", "bench"]:
             gsm8k_metrics = parser.metric_list(f"{prefix}/gsm8k")
             self.assertTrue(len(gsm8k_metrics) > 0)
-            gsm8k_metric_steps = parser.metric_steps(gsm8k_metrics[0])
-            self.assertEqual([0, 2], gsm8k_metric_steps)
+            for eval_stats in ["mean", "best", "worst"]:
+                for k in eval_taskset.eval_at_k:
+                    for stats in ["mean", "std"]:
+                        metric_name = f"{prefix}/gsm8k/accuracy/{eval_stats}@{k}/{stats}"
+                        metric_steps = parser.metric_steps(metric_name)
+                        self.assertEqual(metric_steps, [0, 2])
 
     def tearDown(self):
         shutil.rmtree(self.config.checkpoint_job_dir)

trinity/algorithm/sample_strategy/mix_sample_strategy.py

@@ -114,12 +114,12 @@ def default_args(cls) -> Dict:
 
     def state_dict(self) -> dict:
         return {
-            "usal_buffer": self.usual_exp_buffer.state_dict(),
+            "usual_buffer": self.usual_exp_buffer.state_dict(),
             "expert_buffer": self.expert_exp_buffer.state_dict(),
         }
 
     def load_state_dict(self, state_dict: dict) -> None:
-        if state_dict.get("usal_buffer", None):
-            self.usual_exp_buffer.load_state_dict(state_dict["usal_buffer"])
+        if state_dict.get("usual_buffer", None):
+            self.usual_exp_buffer.load_state_dict(state_dict["usual_buffer"])
         if state_dict.get("expert_buffer", None):
             self.expert_exp_buffer.load_state_dict(state_dict["expert_buffer"])

trinity/buffer/reader/file_reader.py

@@ -176,7 +176,7 @@ def __init__(self, config: StorageConfig):
             total_epochs=self.config.total_epochs if not self.config.is_eval else 1,
             offset=self.config.index,
             drop_last=not self.config.is_eval,
-            total_steps=self.config.total_steps,
+            total_steps=self.config.total_steps if not self.config.is_eval else None,
             enable_progress_bar=self.config.enable_progress_bar,
         )
         self.formatter = FORMATTER.get("task")(config)

trinity/buffer/schema/formatter.py

@@ -65,6 +65,7 @@ def format(self, sample: Dict) -> Task:
             workflow_args=self.config.workflow_args,
             reward_fn_args=self.config.reward_fn_args,
             is_eval=self.config.is_eval,
+            eval_at_k=self.config.eval_at_k,
             raw_task=sample,
         )
 

trinity/common/config.py

@@ -84,7 +84,7 @@ class GenerationConfig:
     logprobs: Optional[int] = None  # 0  # vLLM return `logprobs + 1` elements
     max_tokens: Optional[int] = None  # if None, use model.max_response_tokens
     # repeat each task for `n` times
-    # ! DO NOT SET in `buffer.explorer_input.taskset.rollout_args`
+    # ! DO NOT SET, it will be set by `algorithm.repeat_times` or `max(buffer.explorer_input.eval_tasksets[i].eval_at_k)`
     n: int = 1
 
 
@@ -196,6 +196,7 @@ class StorageConfig:
     workflow_args: dict = field(default_factory=dict)
     reward_fn_args: dict = field(default_factory=dict)
     task_selector: TaskSelectorConfig = field(default_factory=TaskSelectorConfig)
+    eval_at_k: List[int] = field(default_factory=lambda: [1])
 
     # enable progress bar (tqdm) for _HFBatchReader
     enable_progress_bar: Optional[bool] = False
@@ -237,6 +238,7 @@ class TasksetConfig:
     workflow_args: dict = field(default_factory=dict)
     reward_fn_args: dict = field(default_factory=dict)
     task_selector: TaskSelectorConfig = field(default_factory=TaskSelectorConfig)
+    eval_at_k: List[int] = field(default_factory=lambda: [1])
 
     # used for StorageType.FILE
     split: str = "train"
@@ -251,8 +253,8 @@ class TasksetConfig:
 
     # ! DO NOT SET, automatically load from checkpoint
     index: int = 0
-    # ! DO NOT SET, automatically set from algorithm.repeat_times
-    repeat_times: int = 1
+    # ! DO NOT SET, automatically set from `algorithm.repeat_times` or `max(buffer.explorer_input.eval_tasksets[i].eval_at_k)`
+    repeat_times: Optional[int] = None
     # ! DO NOT SET, automatically set based on train/eval
     is_eval: bool = False
     # ! DO NOT SET, automatically set from buffer.batch_size
@@ -268,6 +270,7 @@ def to_storage_config(self) -> StorageConfig:
             storage_type=self.storage_type,
             path=self.path,
             task_selector=self.task_selector,
+            eval_at_k=self.eval_at_k,
             repeat_times=self.repeat_times,
             split=self.split,
             subset_name=self.subset_name,
@@ -907,6 +910,7 @@ def _check_explorer_input(self) -> None:
                     "`buffer.explorer_input.taskset.repeat_times` is set to `algorithm.repeat_times`"
                     f" (={self.algorithm.repeat_times})."
                 )
+            taskset.eval_at_k = []
             taskset.total_epochs = self.buffer.total_epochs
             taskset.total_steps = self.buffer.total_steps
             taskset.batch_size = self.buffer.batch_size
@@ -927,7 +931,8 @@ def _check_explorer_input(self) -> None:
             dataset.batch_size = self.buffer.batch_size
             if not dataset.name:
                 dataset.name = f"eval_taskset_{idx}"
-            set_if_none(dataset, "repeat_times", 1)
+            dataset.repeat_times = max(dataset.eval_at_k)
+
             # eval_workflow has higher priority than workflow in eval tasksets, so we set it first
             set_if_none(dataset, "default_workflow_type", explorer_input.default_eval_workflow_type)
             set_if_none(dataset, "default_workflow_type", explorer_input.default_workflow_type)

trinity/common/workflows/workflow.py

@@ -30,6 +30,7 @@ class Task(dict):
     workflow_args: dict = field(default_factory=dict)
     reward_fn_args: dict = field(default_factory=dict)
     is_eval: bool = False
+    eval_at_k: List[int] = field(default_factory=lambda: [1])
     reward_fn: Optional[Type[RewardFn]] = None
     raw_task: Optional[dict] = None  # The raw data sample
 

trinity/explorer/explorer.py

@@ -397,7 +397,9 @@ async def _finish_eval_step(self, step: Optional[int] = None, prefix: str = "eva
             metric[f"{prefix}/{eval_task_name}/finished_task_count"] = len(statuses)
             metric.update(
                 gather_metrics(
-                    [status.metrics[0] for status in statuses], f"{prefix}/{eval_task_name}"
+                    [status.metrics[0] for status in statuses],
+                    f"{prefix}/{eval_task_name}",
+                    output_stats=["mean", "std"],
                 )
             )
         if self.eval_start_time is not None:

trinity/explorer/scheduler.py

@@ -30,7 +30,9 @@ class TaskWrapper:
     results: List[Tuple[Status, List[Experience]]] = field(default_factory=list)
 
 
-def calculate_task_level_metrics(metrics: List[Dict]) -> Dict[str, float]:
+def calculate_task_level_metrics(
+    metrics: List[Dict], is_eval: bool, eval_at_k: List[int]
+) -> Dict[str, float]:
     """Calculate task level metrics from experiences."""
     if not metrics:
         return {}
@@ -39,7 +41,20 @@ def calculate_task_level_metrics(metrics: List[Dict]) -> Dict[str, float]:
         for key, value in m.items():
             if isinstance(value, (int, float)):
                 aggregated_metrics[key].append(value)
-    return {key: sum(values) / len(values) for key, values in aggregated_metrics.items() if values}
+    if is_eval:
+        result = {}
+        for key, values in aggregated_metrics.items():
+            for k in eval_at_k:
+                if k > len(values):
+                    continue
+                result[f"{key}/mean@{k}"] = sum(values[:k]) / k
+                result[f"{key}/best@{k}"] = max(values[:k]) / k
+                result[f"{key}/worst@{k}"] = min(values[:k]) / k
+        return result
+    else:
+        return {
+            key: sum(values) / len(values) for key, values in aggregated_metrics.items() if values
+        }
 
 
 class RunnerWrapper:
@@ -327,7 +342,12 @@ def task_done_callback(self, async_task: asyncio.Task):
                     if not s.ok:
                         all_success = False
                 task_status = Status(
-                    ok=all_success, metrics=[calculate_task_level_metrics(task_metrics)]
+                    ok=all_success,
+                    metrics=[
+                        calculate_task_level_metrics(
+                            task_metrics, task.task.is_eval, task.task.eval_at_k
+                        )
+                    ],
                 )
                 self.completed_tasks[task.batch_id].appendleft((task_status, task_experiences))
                 self.logger.debug(f"Task completed (batch_id {task.batch_id}).")

trinity/utils/monitor.py

@@ -25,18 +25,19 @@
 MONITOR = Registry("monitor")
 
 
-def gather_metrics(metric_list: List[Dict], prefix: str) -> Dict:
+def gather_metrics(
+    metric_list: List[Dict], prefix: str, output_stats: List[str] = ["mean", "max", "min"]
+) -> Dict:
     if not metric_list:
         return {}
     try:
         df = pd.DataFrame(metric_list)
         numeric_df = df.select_dtypes(include=[np.number])
-        stats_df = numeric_df.agg(["mean", "max", "min"])
+        stats_df = numeric_df.agg(output_stats)
         metric = {}
         for col in stats_df.columns:
-            metric[f"{prefix}/{col}/mean"] = stats_df.loc["mean", col].item()
-            metric[f"{prefix}/{col}/max"] = stats_df.loc["max", col].item()
-            metric[f"{prefix}/{col}/min"] = stats_df.loc["min", col].item()
+            for stats in output_stats:
+                metric[f"{prefix}/{col}/{stats}"] = stats_df.loc[stats, col].item()
         return metric
     except Exception as e:
         raise ValueError(f"Failed to gather metrics: {e}") from e