diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 0fa9f9145fe..4cf95173038 100644 --- a/miles/backends/megatron_utils/model.py +++ b/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: diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index c1e092ecd51..105150e00d3 100644 --- a/miles/ray/rollout.py +++ b/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,6 +59,8 @@ 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 @@ -65,6 +68,10 @@ def __init__(self, args, pg, wandb_run_id): 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: diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index c31f3e9c3b8..a7554c91dc6 100644 --- a/miles/utils/arguments.py +++ b/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. diff --git a/miles/utils/metric_checker.py b/miles/utils/metric_checker.py new file mode 100644 index 00000000000..9ecbe8af224 --- /dev/null +++ b/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) diff --git a/train.py b/train.py index f4d1b92a671..c82ebdeef81 100644 --- a/train.py +++ b/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() diff --git a/train_async.py b/train_async.py index 8a9a021fe03..f972c0794bb 100644 --- a/train_async.py +++ b/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()