diff --git a/docs/50-example-calc-x.md b/docs/50-example-calc-x.md index ed97f66c7..dfcac2d77 100644 --- a/docs/50-example-calc-x.md +++ b/docs/50-example-calc-x.md @@ -71,3 +71,42 @@ bash run_minikube.sh `run_minikube.sh` starts `agl-server` and `agl-controller`, and writes their logs under `/tmp/`. The script also starts a new local Minikube single-node K8s cluster, and the agent runs in this cluster as Kubernetes Jobs. When `run_minikube.sh` exits, it automatically cleans up the server, controller, and Minikube it started. Minikube needs at least 64 GB of memory; otherwise, it may be killed due to insufficient memory. + +## Evaluating a Trained Checkpoint + +Training writes checkpoints under `trainer.default_local_dir`, which for this example is +`checkpoints/agentlightning/calc_x/`, one folder per `global_step_N`. To score one of them on +a held-out split without training anything further, pass `--val-only` and point `--resume` at +the checkpoint: + +```bash +source .venv/bin/activate +cd examples/calc_x +python train_calc_agent.py \ + --val-only \ + --resume checkpoints/agentlightning/calc_x/global_step_128 \ + --val-file data/test.parquet +``` + +`--resume` sets `trainer.resume_mode` and `trainer.resume_from_path`, which is how `verl` +restores a checkpoint. `--model` keeps its earlier meaning, the base model, because the +gateway requires its own model name and `actor_rollout_ref.model.path` to be identical +([API gateway configuration](25-api-gateway-configuration.md)); aiming either of them at a +checkpoint makes the proxy answer "model not found". + +`--val-only` then runs a single validation pass over `--val-file` and exits without taking an +optimizer step. The metrics are printed as `Initial validation metrics: {...}` and logged to +W&B under the `calc_x_eval` experiment. The one to read is `val/reward`: it is the mean reward +over the split, and the Calc-X reward is 1.0 for an answer matching the reference and 0.0 +otherwise, so it doubles as the accuracy. + +Under local mode the same flags go through `run_local.sh`, which forwards its arguments: +`bash run_local.sh --val-only --resume checkpoints/agentlightning/calc_x/global_step_128`. + +Two details worth knowing: + +- `--val-only` still wants a non-empty train split, because `run_ppo` sizes its sampler from + one. The example hands the held-out split in as both, so make sure it has at least + `data.train_batch_size` rows (32 by default). +- The checkpoint is only read, never written, so the same split can be scored again later + and compared across runs. diff --git a/examples/calc_x/train_calc_agent.py b/examples/calc_x/train_calc_agent.py index 8891fb65c..5147a5a8b 100644 --- a/examples/calc_x/train_calc_agent.py +++ b/examples/calc_x/train_calc_agent.py @@ -10,11 +10,15 @@ from omegaconf import OmegaConf -def verl_default_config() -> dict[str, Any]: +def verl_default_config(val_only: bool = False) -> dict[str, Any]: """VERL config overrides for Calc-X training. These are merged on top of Agent Lightning's base config (agentlightning/verl/config.yaml → verl/trainer/config/ppo_trainer.yaml). + + With ``val_only`` the run scores a saved checkpoint instead of training one. + VERL only takes that exit from the validation pass it runs before the first + step, so ``val_before_train`` has to be on as well. """ example_dir = Path(__file__).resolve().parent return { @@ -68,7 +72,8 @@ def verl_default_config() -> dict[str, Any]: }, "trainer": { "n_gpus_per_node": 1, - "val_before_train": False, + "val_before_train": val_only, + "val_only": val_only, "critic_warmup": 0, "logger": ["console", "wandb"], "project_name": "agentlightning", @@ -107,6 +112,8 @@ def build_config( agl_key: str | None = None, run_name: str | None = None, async_mode: bool = False, + val_only: bool = False, + resume_path: str | None = None, config_overrides: Sequence[str] = (), ) -> Any: """Build the full OmegaConf config by merging base + overrides. @@ -123,7 +130,16 @@ def build_config( with initialize_config_dir(config_dir=config_dir, version_base=None): base_cfg = compose(config_name="config") - overrides = verl_default_config() + overrides = verl_default_config(val_only=val_only) + + if resume_path: + # A checkpoint is restored through trainer.resume_from_path, not by pointing + # actor_rollout_ref.model.path at it. That key is the base model, and the gateway + # requires its own model name to match it exactly, so aiming it at a checkpoint + # makes the proxy answer "model not found" + # (docs/25-api-gateway-configuration.md). + overrides["trainer"]["resume_mode"] = "resume_path" + overrides["trainer"]["resume_from_path"] = resume_path if model: overrides["actor_rollout_ref"]["model"]["path"] = model @@ -138,6 +154,8 @@ def build_config( overrides["data"]["train_batch_size"] * 2 ) name_parts.append("async") + if val_only: + name_parts.append("eval") if run_name: name_parts.append(run_name) overrides["trainer"]["experiment_name"] = "_".join(name_parts) @@ -146,9 +164,30 @@ def build_config( cli_override_conf = OmegaConf.from_dotlist(list(config_overrides)) OmegaConf.set_struct(base_cfg, False) config = OmegaConf.merge(base_cfg, override_conf, cli_override_conf) + + # The CLI overrides merge last, so any of them can flip these back. In this mode the + # held-out split is also what run_ppo sizes its sampler from, so a val_before_train of + # false starts optimizing on it: the run would train on the data it is supposed to be + # scored on, with nothing in the logs saying so. Re-assert after the merge rather than + # trusting the keys to survive it. + if val_only: + config.trainer.val_only = True + config.trainer.val_before_train = True + config.trainer.resume_mode = "resume_path" + if resume_path: + config.trainer.resume_from_path = resume_path + return config +def load_split(path: str) -> Sequence[Any]: + """Read one parquet split into a list of rows.""" + return cast( + Sequence[Any], + HuggingFaceDataset.from_parquet(path).to_list(), # type: ignore + ) + + def train( *, train_file: str, @@ -163,15 +202,8 @@ def train( """Load datasets, build config, and launch VERL training via Agent Lightning.""" from agentlightning.verl.entrypoint import run_ppo - # Load datasets. - train_dataset: Sequence[Any] = cast( - Sequence[Any], - HuggingFaceDataset.from_parquet(train_file).to_list(), # type: ignore - ) - val_dataset: Sequence[Any] = cast( - Sequence[Any], - HuggingFaceDataset.from_parquet(val_file).to_list(), # type: ignore - ) + train_dataset = load_split(train_file) + val_dataset = load_split(val_file) print(f"Train dataset: {len(train_dataset)} samples") print(f"Val dataset: {len(val_dataset)} samples") @@ -193,6 +225,50 @@ def train( run_ppo(config, train_dataset=train_dataset, val_dataset=val_dataset) +def evaluate( + *, + checkpoint: str, + data_file: str, + model: str | None = None, + agl_base_url: str | None = None, + agl_key: str | None = None, + run_name: str | None = None, + config_overrides: Sequence[str] = (), +) -> None: + """Score a saved checkpoint on a held-out split, without training it. + + The checkpoint goes through ``trainer.resume_from_path``; ``--model`` stays what it + already was, the base model both the trainer and the gateway are configured with. + + VERL runs one validation pass and returns. ``run_ppo`` still wants a non-empty train + split, so the held-out split is handed in as both: in this mode it only sizes the + sampler and no optimizer step is taken. That also means the split has to be at least + ``data.train_batch_size`` rows long (32 by default), or the sampler has nothing to + build a batch from. + """ + dataset = load_split(data_file) + print(f"Held-out dataset: {len(dataset)} samples") + + config = build_config( + model=model, + resume_path=checkpoint, + agl_base_url=agl_base_url, + agl_key=agl_key, + run_name=run_name, + val_only=True, + config_overrides=config_overrides, + ) + + from pprint import pprint + + print("\n=== VERL Config ===") + pprint(OmegaConf.to_container(config, resolve=True)) + + from agentlightning.verl.entrypoint import run_ppo + + run_ppo(config, train_dataset=dataset, val_dataset=dataset) + + def main() -> None: parser = argparse.ArgumentParser( description="Train Calc-X agent with VERL on Agent Lightning.", @@ -239,8 +315,35 @@ def main() -> None: action="store_true", help="Enable async rollout with async batch size set to 2x train_batch_size", ) + parser.add_argument( + "--resume", + dest="resume_path", + type=str, + default=None, + help="Path to a saved checkpoint (the run's global_step_N folder) to score", + ) + parser.add_argument( + "--val-only", + dest="val_only", + action="store_true", + help="Score --resume on --val-file once, without training", + ) args, config_overrides = parser.parse_known_args() + if args.val_only: + if not args.resume_path: + parser.error("--val-only needs --resume to point at a saved checkpoint") + evaluate( + checkpoint=args.resume_path, + data_file=args.val_file, + model=args.model, + agl_base_url=args.agl_base_url, + agl_key=args.agl_key, + run_name=args.run_name, + config_overrides=config_overrides, + ) + return + train( train_file=args.train_file, val_file=args.val_file, diff --git a/tests/examples/test_calc_x_val_only.py b/tests/examples/test_calc_x_val_only.py new file mode 100644 index 000000000..612a14284 --- /dev/null +++ b/tests/examples/test_calc_x_val_only.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""The Calc-X example has to be able to score a checkpoint without training it. + +`fit()` only reaches its `val_only` exit from the validation pass it runs before +the first step, so turning one knob without the other silently trains instead of +evaluating. These tests pin that pair down, along with the two ways the mode can +be undone from outside: the checkpoint has to travel through `resume_from_path` +rather than the base model path the gateway keys on, and the CLI overrides merge +last, so they must not be able to switch the validation pass back off. +""" + +from __future__ import annotations + +import pytest + +from examples.calc_x.train_calc_agent import build_config, verl_default_config + +CHECKPOINT = "/runs/calc_x/global_step_128" + + +def test_training_still_skips_the_pre_training_validation() -> None: + trainer = verl_default_config()["trainer"] + assert trainer["val_only"] is False + assert trainer["val_before_train"] is False + + +def test_val_only_switches_on_the_validation_pass() -> None: + trainer = verl_default_config(val_only=True)["trainer"] + assert trainer["val_only"] is True + assert trainer["val_before_train"] is True + + +def test_evaluation_restores_the_checkpoint_through_the_resume_path() -> None: + # Only the trainer's own exit matters here; Hydra needs verl to compose. + pytest.importorskip("verl") + + config = build_config(resume_path=CHECKPOINT, val_only=True) + assert config.trainer.val_only is True + assert config.trainer.val_before_train is True + assert config.trainer.resume_mode == "resume_path" + assert config.trainer.resume_from_path == CHECKPOINT + assert config.trainer.experiment_name == "calc_x_eval" + + +def test_run_name_is_still_appended_after_the_evaluation_marker() -> None: + pytest.importorskip("verl") + + config = build_config(val_only=True, run_name="step128") + assert config.trainer.experiment_name == "calc_x_eval_step128" + + +def test_the_checkpoint_does_not_replace_the_base_model() -> None: + # The gateway contract is default_proxy.model_name == actor_rollout_ref.model.path + # (docs/25-api-gateway-configuration.md), so aiming either at a checkpoint makes the + # proxy answer "model not found" before the validation pass gets anywhere. + pytest.importorskip("verl") + + config = build_config( + model="Qwen/Qwen2.5-1.5B-Instruct", + resume_path=CHECKPOINT, + val_only=True, + ) + assert config.actor_rollout_ref.model.path == "Qwen/Qwen2.5-1.5B-Instruct" + assert config.trainer.resume_from_path == CHECKPOINT + + +def test_cli_overrides_cannot_turn_evaluation_back_into_training() -> None: + # The overrides merge last, and in this mode the held-out split is also what run_ppo + # sizes its sampler from, so a val_before_train that got through would optimise on the + # data being scored, silently. + pytest.importorskip("verl") + + config = build_config( + resume_path=CHECKPOINT, + val_only=True, + config_overrides=["trainer.val_before_train=false", "trainer.val_only=false"], + ) + assert config.trainer.val_only is True + assert config.trainer.val_before_train is True + assert config.trainer.resume_mode == "resume_path" + assert config.trainer.resume_from_path == CHECKPOINT