remove eval_at_k

Author: chenyushuo

tests/trainer/trainer_test.py

@@ -84,14 +84,11 @@ def test_trainer(self):
         self.config.buffer.explorer_input.taskset.task_selector = TaskSelectorConfig(
             selector_type="shuffle", seed=42
         )
-        self.config.buffer.explorer_input.eval_tasksets.append(
-            get_unittest_dataset_config("countdown", "test")
-        )
-        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]
+        eval_tasksets = self.config.buffer.explorer_input.eval_tasksets
+        eval_tasksets.append(get_unittest_dataset_config("countdown", "test"))
+        eval_tasksets.append(get_unittest_dataset_config("copy_countdown", "test"))
+        eval_tasksets[0].repeat_times = 4
+        eval_tasksets[1].repeat_times = 4
         self.config.trainer.save_interval = 4
         self.config.check_and_update()
         _trainer_config = self.config.trainer.trainer_config
@@ -149,13 +146,12 @@ 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"]:
-            for eval_taskset, taskset_name in zip(eval_tasksets, ["countdown", "copy_countdown"]):
+            for taskset_name in ["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 k in [2, 4]:
                         for stats in ["mean", "std"]:
                             metric_name = f"{prefix}/{taskset_name}/score/{eval_stats}@{k}/{stats}"
                             metric_steps = parser.metric_steps(metric_name)
@@ -973,7 +969,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.buffer.explorer_input.eval_tasksets[0].repeat_times = 8
         self.config.model.model_path = get_model_path()
         self.config.algorithm.algorithm_type = "grpo"
         self.config.algorithm.advantage_fn = "grpo"
@@ -1021,12 +1017,11 @@ 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)
             for eval_stats in ["mean", "best", "worst"]:
-                for k in eval_taskset.eval_at_k:
+                for k in [2, 4, 8]:
                     for stats in ["mean", "std"]:
                         metric_name = f"{prefix}/gsm8k/accuracy/{eval_stats}@{k}/{stats}"
                         metric_steps = parser.metric_steps(metric_name)

trinity/buffer/schema/formatter.py

@@ -65,7 +65,6 @@ 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, it will be set by `algorithm.repeat_times` or `max(buffer.explorer_input.eval_tasksets[i].eval_at_k)`
+    # ! DO NOT SET, it will be set by `algorithm.repeat_times` or `max(buffer.explorer_input.eval_tasksets[i].repeat_times)`
     n: int = 1
 
 
@@ -196,7 +196,6 @@ 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
@@ -238,7 +237,6 @@ 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"
@@ -253,8 +251,8 @@ class TasksetConfig:
 
     # ! DO NOT SET, automatically load from checkpoint
     index: int = 0
-    # ! 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 in trainer_input, automatically set from `algorithm.repeat_times`
+    repeat_times: int = 1
     # ! DO NOT SET, automatically set based on train/eval
     is_eval: bool = False
     # ! DO NOT SET, automatically set from buffer.batch_size
@@ -270,7 +268,6 @@ 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,
@@ -910,7 +907,6 @@ 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
@@ -931,9 +927,6 @@ def _check_explorer_input(self) -> None:
             dataset.batch_size = self.buffer.batch_size
             if not dataset.name:
                 dataset.name = f"eval_taskset_{idx}"
-            if len(dataset.eval_at_k) == 0:
-                dataset.eval_at_k = [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)

trinity/common/workflows/workflow.py

@@ -30,7 +30,6 @@ 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/scheduler.py

@@ -31,15 +31,12 @@ class TaskWrapper:
     results: List[Tuple[Status, List[Experience]]] = field(default_factory=list)
 
 
-def calculate_task_level_metrics(
-    metrics: List[Dict], is_eval: bool, eval_at_k: List[int]
-) -> Dict[str, float]:
+def calculate_task_level_metrics(metrics: List[Dict], is_eval: bool) -> Dict[str, float]:
     """Calculate task level metrics (mean) from multiple runs of the same task.
 
     Args:
         metrics (`List[Dict]`): A list of metric dictionaries from multiple runs of the same task.
         is_eval (`bool`): Whether this is an evaluation task.
-        eval_at_k (`List[int]`): A list of k values to evaluate at.
 
     Returns:
         `Dict[str, float]`: A dictionary of aggregated metrics, where each metric is averaged over all runs.
@@ -54,9 +51,13 @@ def calculate_task_level_metrics(
     if is_eval:
         result = {}
         for key, values in aggregated_metrics.items():
-            for k in eval_at_k:
-                if k > len(values):
-                    continue
+            k_list = []
+            k = 2
+            while k < len(values):
+                k_list.append(k)
+                k *= 2
+            k_list.append(len(values))
+            for k in k_list:
                 result[f"{key}/mean@{k}"] = sum(values[:k]) / k
                 result[f"{key}/best@{k}"] = max(values[:k])
                 result[f"{key}/worst@{k}"] = min(values[:k])
@@ -356,11 +357,7 @@ def task_done_callback(self, async_task: asyncio.Task):
                 # calculate task level metrics
                 task_status = Status(
                     ok=all_success,
-                    metrics=[
-                        calculate_task_level_metrics(
-                            task_metrics, task.task.is_eval, task.task.eval_at_k
-                        )
-                    ],
+                    metrics=[calculate_task_level_metrics(task_metrics, task.task.is_eval)],
                 )
                 self.completed_tasks[task.batch_id].appendleft((task_status, task_experiences))
                 self.logger.debug(f"Task completed (batch_id {task.batch_id}).")