Support checking accuracy correctness programmatically

miles/backends/megatron_utils/model.py

@@ -486,7 +486,7 @@ def train(rollout_id, model, optimizer, opt_param_scheduler, data_iterator, num_
                 tb = _TensorboardAdapter(args)
                 tb.log(data=log_dict, step=accumulated_step_id)
 
-            if args.ci_test:
+            if args.ci_test and not args.ci_disable_kl_checker:
                 if step_id == 0 and "train/ppo_kl" in log_dict and "train/pg_clipfrac" in log_dict:
                     assert log_dict["train/ppo_kl"] == 0.0 and log_dict["train/pg_clipfrac"] == 0.0, f"{loss_dict=}"
                 if accumulated_step_id == 0 and "train/kl_loss" in log_dict:

miles/ray/rollout.py

@@ -14,6 +14,7 @@
 from miles.backends.sglang_utils.sglang_engine import SGLangEngine
 from miles.ray.rollout_data_source import RolloutDataSourceWithBuffer
 from miles.utils.http_utils import find_available_port, get_host_info, init_http_client
+from miles.utils.metric_checker import MetricChecker
 from miles.utils.misc import load_function
 from miles.utils.ray_utils import Box
 from miles.utils.types import Sample
@@ -58,13 +59,19 @@ def __init__(self, args, pg, wandb_run_id):
         self.rollout_engines = self.all_rollout_engines[:: self.nodes_per_engine]
         self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote()
 
+        self._metric_checker = MetricChecker.maybe_create(args)
+
         # fault tolerance
         self._health_monitor_thread = None
         self._health_monitor_stop_event = None
         self._health_check_interval = args.rollout_health_check_interval
         self._health_check_timeout = args.rollout_health_check_timeout
         self._health_check_first_wait = args.rollout_health_check_first_wait
 
+    def dispose(self):
+        if self._metric_checker is not None:
+            self._metric_checker.dispose()
+
     def get_rollout_engines_and_lock(self):
         return self.rollout_engines, self.rollout_engine_lock, self.num_new_engines
 
@@ -94,7 +101,9 @@ def eval(self, rollout_id):
             return
         # TODO: add fault tolerance to eval
         data = self.eval_generate_rollout(self.args, rollout_id, self.data_source, evaluation=True)
-        _log_eval_rollout_data(rollout_id, self.args, data)
+        metrics = _log_eval_rollout_data(rollout_id, self.args, data)
+        if self._metric_checker is not None:
+            self._metric_checker.on_eval(metrics)
 
     def save(self, rollout_id):
         self.data_source.save(rollout_id)
@@ -475,6 +484,8 @@ def _log_eval_rollout_data(rollout_id, args, data):
             ),
         )
 
+    return log_dict
+
 
 def _log_rollout_data(rollout_id, args, samples, rollout_time):
     if args.load_debug_rollout_data:

miles/utils/arguments.py

@@ -933,6 +933,20 @@ def add_ci_arguments(parser):
                 "--ci-test",
                 action="store_true",
             )
+            parser.add_argument(
+                "--ci-disable-kl-checker",
+                action="store_true",
+            )
+            parser.add_argument(
+                "--ci-metric-checker-key",
+                type=str,
+                default=None,
+            )
+            parser.add_argument(
+                "--ci-metric-checker-threshold",
+                type=float,
+                default=None,
+            )
             return parser
 
         # Add custom arguments in front to prevent overwritten some miles arguments.

miles/utils/metric_checker.py

@@ -0,0 +1,26 @@
+from typing import Dict
+
+
+class MetricChecker:
+    @staticmethod
+    def maybe_create(args):
+        if args.ci_test and (args.ci_metric_checker_key is not None):
+            return MetricChecker(args)
+        return None
+
+    def __init__(self, args):
+        self.args = args
+        self._exists_check_success = False
+
+    def on_eval(self, metrics: Dict[str, float]):
+        actual_value = metrics.get(self.args.ci_metric_checker_key)
+        assert actual_value is not None, f"{metrics=} {self.args.ci_metric_checker_key=}"
+
+        check_success = actual_value >= self.args.ci_metric_checker_threshold
+        print(f"[MetricChecker] {check_success=} {actual_value=} {self.args.ci_metric_checker_threshold=}")
+
+        self._exists_check_success |= check_success
+
+    def dispose(self):
+        assert self._exists_check_success, "[MetricChecker] accuracy check failed"
+        print(f"[MetricChecker] pass dispose check", flush=True)

train.py

@@ -84,6 +84,8 @@ def train(args):
         ):
             ray.get(rollout_manager.eval.remote(rollout_id))
 
+    ray.get(rollout_manager.dispose.remote())
+
 
 if __name__ == "__main__":
     args = parse_args()

train_async.py

@@ -63,6 +63,8 @@ def train(args):
         ):
             ray.get(rollout_manager.eval.remote(rollout_id))
 
+    ray.get(rollout_manager.dispose.remote())
+
 
 if __name__ == "__main__":
     args = parse_args()