From 6d3e45fb10f6ffb64ef34b617584a2382776dd3e Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 14 Jul 2026 17:26:29 -0700 Subject: [PATCH] Require effective GRPO updates --- docs/opencode-grpo.md | 6 + docs/training-pipeline.md | 20 +- pipelines/benchflow-task-posttrain/README.md | 13 +- .../configs/qwen3.5-9b-data-agent-full.toml | 5 +- .../benchflow_pipeline/config.py | 5 +- .../posttrainarena/benchflow_pipeline/grpo.py | 241 +++++++++++++++- .../benchflow_pipeline/pipeline.py | 133 ++++++++- .../benchflow_pipeline/publishing.py | 1 + .../tests/test_config.py | 5 +- .../tests/test_grpo.py | 188 ++++++++++++- .../tests/test_pipeline.py | 258 +++++++++++++++++- .../tests/test_publishing.py | 2 + 12 files changed, 846 insertions(+), 31 deletions(-) diff --git a/docs/opencode-grpo.md b/docs/opencode-grpo.md index 170be81..3a3b453 100644 --- a/docs/opencode-grpo.md +++ b/docs/opencode-grpo.md @@ -103,6 +103,12 @@ TRL GRPOTrainer -> GRPO policy update ``` +Production uses eight OpenCode rollouts per task group, while low-cost smoke +recipes may use two. `training_diagnostics.json` records the exact TRL recipe, +reward variance for every group, LoRA-B update statistics, and trainer log +history. The production recipe rejects an all-zero-variance run instead of +publishing a no-op adapter. + The BenchFlow LiteLLM proxy is invoked with `BENCHFLOW_CAPTURE_TOKEN_LOGPROBS=1`. It requests sampled-token logprobs from the chat-completions endpoint and preserves them in diff --git a/docs/training-pipeline.md b/docs/training-pipeline.md index d47754f..6d0b092 100644 --- a/docs/training-pipeline.md +++ b/docs/training-pipeline.md @@ -247,8 +247,11 @@ evaluation health artifacts. Use a new run name when changing a recipe or task list. Incomplete strict teacher collection reuses finished `attempt-*` directories and continues with only tasks that still lack an eligible rollout. Resume may -increase `[teacher].max_attempts` or `[runtime].max_completion_length`; all -dataset, model, reward, and optimizer semantics remain immutable. +increase `[teacher].max_attempts`, `[runtime].max_completion_length`, +`[runtime].num_generations`, and `[grpo].generation_batch_size`, or tighten +`[grpo].require_reward_variance`. Changing the GRPO recipe invalidates GRPO and +downstream evaluation artifacts while preserving the verified SFT checkpoint; +all dataset, model, reward, and optimizer semantics remain immutable. The snapshot boundary also rejects byte-equivalent task packages under different task IDs after normalizing the package's declared task name. This is @@ -274,6 +277,10 @@ stall before the first tool call. OpenCode title/summary helpers are not seeded. GRPO rollout requests opt into sampled-token logprobs and do not receive a forced seed, preserving within-group reward variance without making pass-rate comparisons depend on uncontrolled random decoding. +The production Qwen3.5 recipe uses eight generations per task, matching TRL's +official default instead of the two-generation smoke setting. It records +per-group reward ranges and LoRA-B update statistics, and fails before +publishing a GRPO checkpoint when every complete group has zero reward variance. The evaluator itself has a real SkillsBench + Daytona canary with score `1.0`, complete provider telemetry, and healthy `results.jsonl` and @@ -350,10 +357,11 @@ The Qwen3.5 recipe requires separate physical devices for the trainer and TRL vLLM worker. Two H100 80 GB GPUs are the initial canary topology. Exact memory and runtime depend on completion length, generation count, sandbox latency, and task trajectory length. Run the 1x1 canary before the full 2,238-task run. -The checked-in GRPO recipe trains one same-prompt pair at a time: two -generations, generation batch two, and no gradient accumulation. This keeps -long OpenCode trajectories within the 80 GB trainer GPU without reducing the -one-epoch task coverage. The bootstrap also enables PyTorch expandable CUDA +The checked-in production GRPO recipe trains one same-prompt group at a time: +eight generations, generation batch eight, and no gradient accumulation. The +trainer still recomputes policy logprobs with a per-device microbatch of one to +fit long OpenCode trajectories on the 80 GB trainer GPU. The bootstrap also +enables PyTorch expandable CUDA segments to reduce allocator fragmentation, and the GRPO trainer clears cached CUDA allocations between steps. diff --git a/pipelines/benchflow-task-posttrain/README.md b/pipelines/benchflow-task-posttrain/README.md index 9e4d95d..3a08eb3 100644 --- a/pipelines/benchflow-task-posttrain/README.md +++ b/pipelines/benchflow-task-posttrain/README.md @@ -156,8 +156,10 @@ can be reused. If strict teacher coverage is incomplete, resume reuses completed attempts and continues only missing tasks. The retry budget may be increased without changing the rest of the persisted run plan. -An interrupted GRPO stage also permits increasing only the aggregate completion -budget, then restarts cleanly from the saved SFT checkpoint. +An interrupted GRPO stage also permits increasing the aggregate completion +budget, generation count, and generation batch, or enabling strict reward +variance checks. It then restarts GRPO and downstream evaluation cleanly from +the saved SFT checkpoint. The public OpenCode endpoint is `posttrainarena-train model-bridge`, which forwards to the TRL server at `TRL_VLLM_SERVER_BASE_URL`. The pipeline @@ -182,7 +184,8 @@ runs//reports/score.json Important fields include `baseline_score`, `sft_score`, `grpo_gate_score`, `score_after_posttrain`, `delta_score`, `grpo_planned`, `grpo_ran`, exact task IDs, dataset revisions, BenchFlow commit, `grpo_run_policy`, and the recorded -stage commands. The report also records the SFT and GRPO adapter and merged +stage commands. The report also records `grpo_effective_update`, the compact +reward-variance/update summary, and the SFT and GRPO adapter and merged checkpoint paths. A dry-run may set `grpo_planned` while leaving `grpo_ran` false. @@ -191,6 +194,10 @@ false. The Qwen3.5 full recipe sets `grpo.run_policy = "always"` because GRPO is part of the fixed organizer recipe. The engine still supports `on_reward` for low-cost experiments that should skip a constant-zero reward distribution. +That full recipe samples eight generations per task group, matching TRL's +official default, and requires at least one group with nonzero verifier-reward +variance before it publishes a GRPO adapter. Two-generation settings remain in +the canary and smoke recipes for cost control. Do not use held-out eval tasks to tune this gate. Production recipes should use separate training, gate/development, and final evaluation lists. diff --git a/pipelines/benchflow-task-posttrain/configs/qwen3.5-9b-data-agent-full.toml b/pipelines/benchflow-task-posttrain/configs/qwen3.5-9b-data-agent-full.toml index e7daacd..d33a237 100644 --- a/pipelines/benchflow-task-posttrain/configs/qwen3.5-9b-data-agent-full.toml +++ b/pipelines/benchflow-task-posttrain/configs/qwen3.5-9b-data-agent-full.toml @@ -18,7 +18,7 @@ task_list = "../task-lists/data-agent-eval-all.txt" sandbox = "docker" sandbox_user = "agent" max_completion_length = 40960 -num_generations = 2 +num_generations = 8 [harness] agent = "opencode" @@ -75,8 +75,9 @@ lora_r = 16 lora_alpha = 32 lora_dropout = 0.05 log_completions = false -generation_batch_size = 2 +generation_batch_size = 8 rollout_attempts = 2 +require_reward_variance = true vllm_server_base_url_env = "TRL_VLLM_SERVER_BASE_URL" [tracking] diff --git a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/config.py b/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/config.py index f199649..e3b4290 100644 --- a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/config.py +++ b/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/config.py @@ -68,7 +68,7 @@ class RuntimeConfig: sandbox: str | None = None sandbox_user: str | None = "agent" max_completion_length: int = 2048 - num_generations: int = 2 + num_generations: int = 8 @dataclass(frozen=True) @@ -140,6 +140,7 @@ class GrpoConfig: log_completions: bool = False generation_batch_size: int | None = None rollout_attempts: int = 2 + require_reward_variance: bool = False vllm_server_base_url_env: str = "TRL_VLLM_SERVER_BASE_URL" @@ -266,6 +267,8 @@ def validate(self) -> None: ) if not _is_positive_int(self.grpo.rollout_attempts): errors.append("grpo.rollout_attempts must be positive") + if not isinstance(self.grpo.require_reward_variance, bool): + errors.append("grpo.require_reward_variance must be boolean") if ( not isinstance(self.grpo.vllm_server_base_url_env, str) or not self.grpo.vllm_server_base_url_env.strip() diff --git a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/grpo.py b/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/grpo.py index ceebe85..4713247 100644 --- a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/grpo.py +++ b/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/grpo.py @@ -9,11 +9,12 @@ import gc from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass +from importlib.metadata import PackageNotFoundError, version from pathlib import Path from collections.abc import Callable, Mapping from typing import Any, Sequence -from .config import PipelineConfig +from .config import BENCHFLOW_COMMIT, PipelineConfig from .io import CommandRunner, supported_kwargs, write_json from .model_bridge import normalize_tool_call_arguments from .opencode import ServedModelRole, evaluate, served_model @@ -22,6 +23,184 @@ TASK_HANDLE_PREFIX = "benchflow-task://" +def effective_generation_batch_size(config: PipelineConfig) -> int: + return ( + config.grpo.generation_batch_size + or config.runtime.num_generations * config.harness.concurrency + ) + + +def _implementation_sha256() -> str: + digest = hashlib.sha256() + module_dir = Path(__file__).resolve().parent + for name in ("config.py", "grpo.py", "model_bridge.py", "opencode.py"): + path = module_dir / name + digest.update(name.encode()) + digest.update(b"\0") + digest.update(path.read_bytes()) + return digest.hexdigest() + + +def _distribution_version(name: str) -> str: + try: + return version(name) + except PackageNotFoundError as exc: + raise RuntimeError( + f"GRPO training dependency is not installed: {name}" + ) from exc + + +def grpo_training_recipe(config: PipelineConfig) -> dict[str, Any]: + return { + "trl_version": _distribution_version("trl"), + "peft_version": _distribution_version("peft"), + "transformers_version": _distribution_version("transformers"), + "torch_version": _distribution_version("torch"), + "benchflow_commit": BENCHFLOW_COMMIT, + "implementation_sha256": _implementation_sha256(), + "num_generations": config.runtime.num_generations, + "max_completion_length": config.runtime.max_completion_length, + "generation_batch_size": effective_generation_batch_size(config), + "num_train_epochs": config.grpo.num_train_epochs, + "max_steps": config.grpo.max_steps, + "learning_rate": config.grpo.learning_rate, + "gradient_accumulation_steps": config.grpo.gradient_accumulation_steps, + "gradient_checkpointing": config.grpo.gradient_checkpointing, + "lora_r": config.grpo.lora_r, + "lora_alpha": config.grpo.lora_alpha, + "lora_dropout": config.grpo.lora_dropout, + "rollout_attempts": config.grpo.rollout_attempts, + "require_reward_variance": config.grpo.require_reward_variance, + "bf16": True, + "per_device_train_batch_size": 1, + "loss_type": "dapo", + "scale_rewards": "group", + "target_modules": "all-linear", + "lora_bias": "none", + "task_type": "CAUSAL_LM", + "torch_empty_cache_steps": 1, + "use_vllm": True, + "vllm_mode": "server", + "vllm_importance_sampling_correction": True, + } + + +def reward_group_diagnostics( + records: Sequence[Mapping[str, Any]], + *, + num_generations: int, +) -> dict[str, Any]: + groups: dict[int, dict[str, Any]] = {} + for record in records: + group_index = record.get("group_index") + global_step = record.get("global_step") + task_id = record.get("task_id") + reward = record.get("reward") + if ( + not isinstance(group_index, int) + or isinstance(group_index, bool) + or not isinstance(global_step, int) + or isinstance(global_step, bool) + or not isinstance(task_id, str) + or not isinstance(reward, int | float) + or isinstance(reward, bool) + or not math.isfinite(float(reward)) + ): + raise RuntimeError("GRPO rollout record is invalid") + group = groups.setdefault( + group_index, + {"global_steps": set(), "task_ids": set(), "rewards": []}, + ) + group["global_steps"].add(global_step) + group["task_ids"].add(task_id) + group["rewards"].append(float(reward)) + + group_rows = [] + complete_group_count = 0 + nonzero_variance_group_count = 0 + for group_index, group in sorted(groups.items()): + global_steps = group["global_steps"] + task_ids = group["task_ids"] + rewards = group["rewards"] + reward_range = max(rewards) - min(rewards) + consistent = len(global_steps) == 1 and len(task_ids) == 1 + complete = consistent and len(rewards) == num_generations + has_variance = complete and not math.isclose( + reward_range, + 0.0, + rel_tol=0.0, + abs_tol=1e-12, + ) + complete_group_count += int(complete) + nonzero_variance_group_count += int(has_variance) + group_rows.append( + { + "group_index": group_index, + "global_step": next(iter(global_steps)) + if len(global_steps) == 1 + else None, + "task_id": next(iter(task_ids)) if len(task_ids) == 1 else None, + "count": len(rewards), + "min_reward": min(rewards), + "max_reward": max(rewards), + "reward_range": reward_range, + "consistent": consistent, + "complete": complete, + "has_variance": has_variance, + } + ) + zero_variance_group_count = complete_group_count - nonzero_variance_group_count + return { + "num_generations": num_generations, + "rollout_count": len(records), + "group_count": len(group_rows), + "complete_group_count": complete_group_count, + "incomplete_group_count": len(group_rows) - complete_group_count, + "nonzero_variance_group_count": nonzero_variance_group_count, + "zero_variance_group_count": zero_variance_group_count, + "zero_variance_fraction": ( + zero_variance_group_count / complete_group_count + if complete_group_count + else None + ), + "groups": group_rows, + } + + +def lora_b_update_diagnostics(model: Any) -> dict[str, Any]: + named_parameters = getattr(model, "named_parameters", None) + if not callable(named_parameters): + return { + "available": False, + "tensor_count": 0, + "nonzero_tensor_count": 0, + "nonfinite_tensor_count": 0, + "max_abs": 0.0, + } + tensor_count = 0 + nonzero_tensor_count = 0 + nonfinite_tensor_count = 0 + max_abs = 0.0 + for name, parameter in named_parameters(): + if "lora_B" not in name: + continue + tensor_count += 1 + tensor = parameter.detach() + tensor_max = float(tensor.float().abs().max().item()) if tensor.numel() else 0.0 + if not math.isfinite(tensor_max): + nonfinite_tensor_count += 1 + continue + max_abs = max(max_abs, tensor_max) + nonzero_tensor_count += int(tensor_max > 0.0) + return { + "available": tensor_count > 0, + "tensor_count": tensor_count, + "nonzero_tensor_count": nonzero_tensor_count, + "nonfinite_tensor_count": nonfinite_tensor_count, + "max_abs": max_abs, + } + + def _directory_sha256(path: Path) -> str: files = sorted( item @@ -790,6 +969,7 @@ def _collect_one( payload=payload, attempt_root=attempt_root, attempt=attempt, + rollout_index=rollout_index, task_id=task_id, tokenizer=tokenizer, global_step=global_step, @@ -816,6 +996,7 @@ def _materialize_rollout( payload: dict[str, Any], attempt_root: Path, attempt: int, + rollout_index: int, task_id: str, tokenizer: Any, global_step: int, @@ -858,6 +1039,8 @@ def _materialize_rollout( trace_resolver=self._resolve_bridge_trace, ) record = { + "rollout_index": rollout_index, + "group_index": rollout_index // self.config.runtime.num_generations, "task_id": task_id, "reward": float(reward), "rollout_dir": str(rollout_dir), @@ -915,10 +1098,7 @@ def train_grpo( label="TRL vLLM server endpoint", ) processing_class = _load_tokenizer(config, model) - generation_batch_size = ( - config.grpo.generation_batch_size - or config.runtime.num_generations * config.harness.concurrency - ) + generation_batch_size = effective_generation_batch_size(config) values = { "output_dir": str(adapter_dir), "run_name": run_name, @@ -935,6 +1115,8 @@ def train_grpo( "vllm_mode": "server", "vllm_server_base_url": vllm_server_base_url, "vllm_importance_sampling_correction": True, + "loss_type": "dapo", + "scale_rewards": "group", "logging_steps": 1, "per_device_train_batch_size": 1, "gradient_accumulation_steps": config.grpo.gradient_accumulation_steps, @@ -970,6 +1152,51 @@ def train_grpo( result = trainer.train() finally: _close_weight_communicator(getattr(trainer, "vllm_generation", None)) + reward_diagnostics = reward_group_diagnostics( + collector.records, + num_generations=config.runtime.num_generations, + ) + update_diagnostics = lora_b_update_diagnostics(trainer.model) + training_log = list( + getattr(getattr(trainer, "state", None), "log_history", []) or [] + ) + jobs_dir.mkdir(parents=True, exist_ok=True) + write_json( + jobs_dir / "training_diagnostics.json", + { + "training_recipe": grpo_training_recipe(config), + "reward_groups": reward_diagnostics, + "lora_b_update": update_diagnostics, + "metrics": result.metrics, + "training_log": training_log, + }, + ) + if config.grpo.require_reward_variance: + train_loss = result.metrics.get("train_loss") + if ( + not isinstance(train_loss, int | float) + or isinstance(train_loss, bool) + or not math.isfinite(float(train_loss)) + ): + raise RuntimeError("GRPO training did not produce a finite train_loss") + if reward_diagnostics["incomplete_group_count"]: + raise RuntimeError( + "GRPO produced incomplete reward groups; inspect " + f"{jobs_dir / 'training_diagnostics.json'}" + ) + if not reward_diagnostics["nonzero_variance_group_count"]: + raise RuntimeError( + "GRPO produced zero within-group reward variance; increase " + "runtime.num_generations or improve reward shaping" + ) + if not update_diagnostics["available"]: + raise RuntimeError("GRPO could not inspect LoRA-B update tensors") + if update_diagnostics["nonfinite_tensor_count"]: + raise RuntimeError("GRPO produced non-finite LoRA-B update tensors") + if not update_diagnostics["nonzero_tensor_count"]: + raise RuntimeError( + "GRPO reward variance was nonzero but every LoRA-B tensor remained zero" + ) adapter_dir.mkdir(parents=True, exist_ok=True) trainer.save_model(str(adapter_dir)) processing_class = getattr(trainer, "processing_class", None) @@ -1021,6 +1248,10 @@ def train_grpo( ), "max_steps": config.grpo.max_steps, "generation_batch_size": generation_batch_size, + "training_recipe": grpo_training_recipe(config), + "reward_group_diagnostics": reward_diagnostics, + "lora_b_update_diagnostics": update_diagnostics, + "training_log": training_log, "quantization": None, "metrics": result.metrics, "jobs_dir": str(jobs_dir), diff --git a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/pipeline.py b/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/pipeline.py index a987d07..2dd8af9 100644 --- a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/pipeline.py +++ b/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/pipeline.py @@ -45,6 +45,8 @@ def _resume_plan_compatible(existing: dict[str, Any], current: dict[str, Any]) - for section, field in ( ("teacher", "max_attempts"), ("runtime", "max_completion_length"), + ("runtime", "num_generations"), + ("grpo", "generation_batch_size"), ): existing_section = existing.get(section) current_section = normalized_current.get(section) @@ -54,6 +56,33 @@ def _resume_plan_compatible(existing: dict[str, Any], current: dict[str, Any]) - return False existing_value = existing_section.get(field) current_value = current_section.get(field) + if existing_value == current_value: + continue + if ( + section == "grpo" + and field == "generation_batch_size" + and existing_value is None + and isinstance(current_value, int) + and not isinstance(current_value, bool) + ): + existing_runtime = existing.get("runtime") + existing_harness = existing.get("harness") + if not isinstance(existing_runtime, dict) or not isinstance( + existing_harness, dict + ): + return False + existing_generations = existing_runtime.get("num_generations") + existing_concurrency = existing_harness.get("concurrency") + if ( + not isinstance(existing_generations, int) + or isinstance(existing_generations, bool) + or not isinstance(existing_concurrency, int) + or isinstance(existing_concurrency, bool) + or current_value < existing_generations * existing_concurrency + ): + return False + current_section[field] = None + continue if ( not isinstance(existing_value, int) or isinstance(existing_value, bool) @@ -63,6 +92,22 @@ def _resume_plan_compatible(existing: dict[str, Any], current: dict[str, Any]) - ): return False current_section[field] = existing_value + existing_grpo = existing.get("grpo") + current_grpo = normalized_current.get("grpo") + if not isinstance(existing_grpo, dict) or not isinstance(current_grpo, dict): + return False + existing_variance = existing_grpo.get("require_reward_variance", False) + current_variance = current_grpo.get("require_reward_variance", False) + if not isinstance(existing_variance, bool) or not isinstance( + current_variance, bool + ): + return False + if existing_variance and not current_variance: + return False + if "require_reward_variance" in existing_grpo: + current_grpo["require_reward_variance"] = existing_variance + else: + current_grpo.pop("require_reward_variance", None) return existing == normalized_current @@ -926,10 +971,6 @@ def _train_sft(self, output_model: str) -> None: ) ): return - if self.resume: - for path in (self.layout.sft_adapter, Path(output_model)): - if path.exists(): - shutil.rmtree(path) if self.dry_run: self.runner.commands.append( { @@ -940,6 +981,40 @@ def _train_sft(self, output_model: str) -> None: } ) return + if self.resume: + for path in ( + self.layout.jobs / "sft", + self.layout.jobs / "grpo-gate", + self.layout.jobs / "grpo-train", + self.layout.jobs / "posttrain", + self.layout.sft_adapter, + Path(output_model), + self.layout.grpo_adapter, + self.layout.grpo_merged, + ): + if path.exists(): + shutil.rmtree(path) + for artifact in ( + self.layout.results / "sft_endpoint_sync.json", + self.layout.results / "sft_eval.json", + self.layout.results / "sft_eval_health.json", + self.layout.results / "sft_eval_task_manifest.json", + self.layout.results / "sft_eval_run_config.json", + self.layout.results / "grpo_gate_eval.json", + self.layout.results / "grpo_gate_eval_health.json", + self.layout.results / "grpo_gate_eval_task_manifest.json", + self.layout.results / "grpo_gate_eval_run_config.json", + self.layout.results / "grpo_endpoint_sync.json", + self.layout.results / "posttrain_eval.json", + self.layout.results / "posttrain_eval_health.json", + self.layout.results / "posttrain_eval_task_manifest.json", + self.layout.results / "posttrain_eval_run_config.json", + self.layout.reports / "EVAL_LIFT.md", + self.layout.reports / "eval_lift.json", + self.layout.reports / "SCORE.md", + self.layout.reports / "score.json", + ): + artifact.unlink(missing_ok=True) from .sft import train_sft train_sft( @@ -985,10 +1060,6 @@ def _train_grpo(self, *, input_model: str, output_model: str) -> None: ): return jobs_dir = self.layout.jobs / "grpo-train" - if self.resume: - for path in (jobs_dir, self.layout.grpo_adapter, Path(output_model)): - if path.exists(): - shutil.rmtree(path) if self.dry_run: self.runner.commands.append( { @@ -1000,6 +1071,27 @@ def _train_grpo(self, *, input_model: str, output_model: str) -> None: } ) return + if self.resume: + for path in ( + jobs_dir, + self.layout.jobs / "posttrain", + self.layout.grpo_adapter, + Path(output_model), + ): + if path.exists(): + shutil.rmtree(path) + for artifact in ( + self.layout.results / "grpo_endpoint_sync.json", + self.layout.results / "posttrain_eval.json", + self.layout.results / "posttrain_eval_health.json", + self.layout.results / "posttrain_eval_task_manifest.json", + self.layout.results / "posttrain_eval_run_config.json", + self.layout.reports / "EVAL_LIFT.md", + self.layout.reports / "eval_lift.json", + self.layout.reports / "SCORE.md", + self.layout.reports / "score.json", + ): + artifact.unlink(missing_ok=True) from .grpo import train_grpo train_grpo( @@ -1020,6 +1112,8 @@ def _grpo_checkpoint_is_current( input_model: str, output_model: Path, ) -> bool: + from .grpo import grpo_training_recipe + revision = ( self.config.model_revision if input_model == self.config.model else None ) @@ -1029,6 +1123,7 @@ def _grpo_checkpoint_is_current( metrics.get("mode") == "grpo" and metrics.get("model") == input_model and metrics.get("task_ids") == self.train_task_ids + and metrics.get("training_recipe") == grpo_training_recipe(self.config) and metrics.get("adapter_dir") == str(self.layout.grpo_adapter) and metrics.get("merged_model_dir") == str(output_model) and metrics.get("base_checkpoint_sha256") @@ -1122,6 +1217,25 @@ def _write_score( if baseline_score is None or final_score is None else final_score - baseline_score ) + grpo_training = None + grpo_effective_update = None + grpo_metrics_path = self.layout.grpo_merged / "train_metrics.json" + if grpo_ran and grpo_metrics_path.is_file(): + grpo_metrics = load_json(grpo_metrics_path) + reward_groups = dict(grpo_metrics.get("reward_group_diagnostics") or {}) + reward_groups.pop("groups", None) + lora_b_update = dict(grpo_metrics.get("lora_b_update_diagnostics") or {}) + grpo_training = { + "training_recipe": grpo_metrics.get("training_recipe"), + "metrics": grpo_metrics.get("metrics"), + "reward_groups": reward_groups, + "lora_b_update": lora_b_update, + } + grpo_effective_update = bool( + reward_groups.get("nonzero_variance_group_count") + and lora_b_update.get("available") + and lora_b_update.get("nonzero_tensor_count") + ) summary = { "schema_version": 1, "run_name": self.run_name, @@ -1139,6 +1253,8 @@ def _write_score( "grpo_run_policy": self.config.grpo.run_policy, "grpo_planned": grpo_planned, "grpo_ran": grpo_ran, + "grpo_effective_update": grpo_effective_update, + "grpo_training": grpo_training, "checkpoints": { "sft_adapter": ( str(self.layout.sft_adapter) if self.config.sft.enabled else None @@ -1182,6 +1298,7 @@ def _write_score( f"- Delta: `{delta}`", f"- GRPO planned: `{grpo_planned}`", f"- GRPO ran: `{grpo_ran}`", + f"- GRPO effective update: `{grpo_effective_update}`", f"- Training tasks: `{len(self.train_task_ids)}`", f"- Eval tasks: `{len(self.eval_task_ids)}`", "", diff --git a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/publishing.py b/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/publishing.py index 334f4d2..d371121 100644 --- a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/publishing.py +++ b/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/publishing.py @@ -73,6 +73,7 @@ def build_run_record( benchmark_summary.get("benchmarks") if benchmark_summary else [] ), "grpo_ran": score.get("grpo_ran"), + "grpo_effective_update": score.get("grpo_effective_update"), "train_task_count": len(score.get("train_task_ids", [])), "eval_task_count": len(score.get("eval_task_ids", [])), "artifact_url": artifact_url, diff --git a/pipelines/benchflow-task-posttrain/tests/test_config.py b/pipelines/benchflow-task-posttrain/tests/test_config.py index f993363..b73fbda 100644 --- a/pipelines/benchflow-task-posttrain/tests/test_config.py +++ b/pipelines/benchflow-task-posttrain/tests/test_config.py @@ -66,6 +66,7 @@ def test_qwen35_full_recipe_uses_all_tasks_and_one_epoch_lora() -> None: assert config.teacher.require_all_tasks is True assert config.teacher.min_verified == 2238 assert config.runtime.max_completion_length == 40960 + assert config.runtime.num_generations == 8 assert config.sft.num_train_epochs == 1.0 assert config.sft.max_steps is None assert config.sft.lora_r == 16 @@ -77,7 +78,8 @@ def test_qwen35_full_recipe_uses_all_tasks_and_one_epoch_lora() -> None: assert config.grpo.lora_alpha == 32 assert config.grpo.log_completions is False assert config.grpo.gradient_accumulation_steps == 1 - assert config.grpo.generation_batch_size == 2 + assert config.grpo.generation_batch_size == 8 + assert config.grpo.require_reward_variance is True assert config.tracking.report_to == "none" assert config.evaluation.sync_base_to_vllm is True @@ -164,6 +166,7 @@ def test_config_rejects_unknown_grpo_run_policy() -> None: ("lora_dropout", 1.0, "grpo.lora_dropout"), ("log_completions", "yes", "grpo.log_completions"), ("generation_batch_size", 3, "grpo.generation_batch_size"), + ("require_reward_variance", "yes", "grpo.require_reward_variance"), ( "vllm_server_base_url_env", "", diff --git a/pipelines/benchflow-task-posttrain/tests/test_grpo.py b/pipelines/benchflow-task-posttrain/tests/test_grpo.py index 1de5c60..d3583a6 100644 --- a/pipelines/benchflow-task-posttrain/tests/test_grpo.py +++ b/pipelines/benchflow-task-posttrain/tests/test_grpo.py @@ -18,6 +18,8 @@ _chat_prompt_ids, attest_served_policy, build_grpo_rows, + grpo_training_recipe, + reward_group_diagnostics, sync_checkpoint_to_vllm, sync_model_to_vllm, sync_reference_to_vllm, @@ -387,6 +389,53 @@ def test_verifier_reward_uses_rollout_metadata() -> None: verifier_reward(["a"], rollout_reward=None) +def test_reward_group_diagnostics_detects_mixed_and_constant_groups() -> None: + records = [ + { + "group_index": 0, + "global_step": 0, + "task_id": "task-a", + "reward": reward, + } + for reward in (0.0, 1.0, 1.0, 1.0) + ] + [ + { + "group_index": 1, + "global_step": 4, + "task_id": "task-b", + "reward": 1.0, + } + for _ in range(4) + ] + + diagnostics = reward_group_diagnostics(records, num_generations=4) + + assert diagnostics["rollout_count"] == 8 + assert diagnostics["complete_group_count"] == 2 + assert diagnostics["nonzero_variance_group_count"] == 1 + assert diagnostics["zero_variance_group_count"] == 1 + assert diagnostics["zero_variance_fraction"] == 0.5 + + +def test_reward_group_diagnostics_rejects_invalid_records() -> None: + with pytest.raises(RuntimeError, match="rollout record is invalid"): + reward_group_diagnostics( + [ + { + "group_index": 0, + "global_step": 0, + "task_id": "task-a", + "reward": float("nan"), + } + ], + num_generations=2, + ) + + diagnostics = reward_group_diagnostics([], num_generations=2) + assert diagnostics["group_count"] == 0 + assert diagnostics["zero_variance_fraction"] is None + + def test_collector_runs_one_opencode_rollout_per_prompt( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -742,7 +791,16 @@ def __init__(self, **kwargs): self.values = kwargs class FakeModel: - pass + def named_parameters(self): + import torch + + if captured.get("hide_lora"): + return + if captured.get("force_nonfinite_lora"): + value = float("inf") + else: + value = 0.0 if captured.get("force_zero_lora") else 1.0 + yield "layer.lora_B.default.weight", torch.tensor([value]) class FakeMerged: def save_pretrained(self, path, **kwargs): @@ -775,9 +833,33 @@ def __init__(self, **kwargs): self.model = FakeModel() self.processing_class = FakeProcessor() self.vllm_generation = SimpleNamespace(vllm_client=FakeVllmClient()) + self.state = SimpleNamespace(log_history=[{"loss": 0.5, "grad_norm": 1.0}]) + self.args = kwargs["args"].values + self.rollout_func = kwargs["rollout_func"] def train(self): - return SimpleNamespace(metrics={"loss": 0.5}) + rewards = [1.0] * self.args["num_generations"] + if not captured.get("force_zero_variance"): + rewards[0] = 0.0 + if captured.get("force_incomplete_group"): + rewards.pop() + self.rollout_func.records = [ + { + "rollout_index": index, + "group_index": 0, + "global_step": 0, + "task_id": "task-a", + "reward": reward, + } + for index, reward in enumerate(rewards) + ] + return SimpleNamespace( + metrics={ + "train_loss": ( + float("inf") if captured.get("force_nonfinite_loss") else 0.5 + ) + } + ) def save_model(self, path): captured["model_path"] = path @@ -842,8 +924,11 @@ def save_model(self, path): assert args["vllm_mode"] == "server" assert args["vllm_server_base_url"] == "http://127.0.0.1:8000" assert args["vllm_importance_sampling_correction"] is True + assert args["loss_type"] == "dapo" + assert args["scale_rewards"] == "group" assert args["log_completions"] is False assert args["generation_batch_size"] == 2 + assert args["num_generations"] == 2 assert args["max_steps"] == config.grpo.max_steps assert "num_train_epochs" not in args assert args["gradient_checkpointing"] is True @@ -861,6 +946,15 @@ def save_model(self, path): assert payload["base_checkpoint_sha256"] assert payload["adapter_sha256"] assert payload["merged_model_sha256"] + assert payload["training_recipe"] == grpo_training_recipe(config) + assert payload["reward_group_diagnostics"]["nonzero_variance_group_count"] == 1 + assert payload["lora_b_update_diagnostics"]["nonzero_tensor_count"] == 1 + assert payload["training_log"] == [{"loss": 0.5, "grad_norm": 1.0}] + diagnostics = json.loads((tmp_path / "jobs/training_diagnostics.json").read_text()) + assert diagnostics["training_recipe"] == grpo_training_recipe(config) + assert diagnostics["reward_groups"]["groups"][0]["has_variance"] is True + assert diagnostics["lora_b_update"]["nonzero_tensor_count"] == 1 + assert diagnostics["metrics"] == {"train_loss": 0.5} dependency = json.loads((adapter_dir / "adapter_dependency.json").read_text()) assert dependency["base_checkpoint"] == config.model assert dependency["published_base_sibling"] == "../sft-merged" @@ -881,10 +975,96 @@ def save_model(self, path): assert full_args["num_train_epochs"] == 1.0 assert "max_steps" not in full_args - assert full_args["generation_batch_size"] == 2 + assert full_args["generation_batch_size"] == 8 + assert full_args["num_generations"] == 8 assert full_payload["num_train_epochs"] == 1.0 assert full_payload["max_steps"] is None - assert full_payload["generation_batch_size"] == 2 + assert full_payload["generation_batch_size"] == 8 + assert full_payload["reward_group_diagnostics"]["nonzero_variance_group_count"] == 1 + + captured["force_zero_variance"] = True + with pytest.raises(RuntimeError, match="zero within-group reward variance"): + train_grpo( + config=full_config, + model=full_config.model, + tasks_dir=tasks_dir, + task_ids=["task-a"], + jobs_dir=tmp_path / "zero-variance-jobs", + adapter_dir=tmp_path / "zero-variance-adapter", + output_dir=tmp_path / "zero-variance-checkpoint", + run_name="zero-variance-run", + ) + assert (tmp_path / "zero-variance-jobs/training_diagnostics.json").is_file() + + captured["force_zero_variance"] = False + captured["force_zero_lora"] = True + with pytest.raises(RuntimeError, match="every LoRA-B tensor remained zero"): + train_grpo( + config=full_config, + model=full_config.model, + tasks_dir=tasks_dir, + task_ids=["task-a"], + jobs_dir=tmp_path / "zero-update-jobs", + adapter_dir=tmp_path / "zero-update-adapter", + output_dir=tmp_path / "zero-update-checkpoint", + run_name="zero-update-run", + ) + + captured["force_zero_lora"] = False + captured["force_incomplete_group"] = True + with pytest.raises(RuntimeError, match="incomplete reward groups"): + train_grpo( + config=full_config, + model=full_config.model, + tasks_dir=tasks_dir, + task_ids=["task-a"], + jobs_dir=tmp_path / "incomplete-group-jobs", + adapter_dir=tmp_path / "incomplete-group-adapter", + output_dir=tmp_path / "incomplete-group-checkpoint", + run_name="incomplete-group-run", + ) + + captured["force_incomplete_group"] = False + captured["hide_lora"] = True + with pytest.raises(RuntimeError, match="could not inspect LoRA-B"): + train_grpo( + config=full_config, + model=full_config.model, + tasks_dir=tasks_dir, + task_ids=["task-a"], + jobs_dir=tmp_path / "missing-update-jobs", + adapter_dir=tmp_path / "missing-update-adapter", + output_dir=tmp_path / "missing-update-checkpoint", + run_name="missing-update-run", + ) + + captured["hide_lora"] = False + captured["force_nonfinite_lora"] = True + with pytest.raises(RuntimeError, match="non-finite LoRA-B"): + train_grpo( + config=full_config, + model=full_config.model, + tasks_dir=tasks_dir, + task_ids=["task-a"], + jobs_dir=tmp_path / "nonfinite-update-jobs", + adapter_dir=tmp_path / "nonfinite-update-adapter", + output_dir=tmp_path / "nonfinite-update-checkpoint", + run_name="nonfinite-update-run", + ) + + captured["force_nonfinite_lora"] = False + captured["force_nonfinite_loss"] = True + with pytest.raises(RuntimeError, match="finite train_loss"): + train_grpo( + config=full_config, + model=full_config.model, + tasks_dir=tasks_dir, + task_ids=["task-a"], + jobs_dir=tmp_path / "nonfinite-loss-jobs", + adapter_dir=tmp_path / "nonfinite-loss-adapter", + output_dir=tmp_path / "nonfinite-loss-checkpoint", + run_name="nonfinite-loss-run", + ) def test_sync_model_to_vllm_closes_weight_communicator( diff --git a/pipelines/benchflow-task-posttrain/tests/test_pipeline.py b/pipelines/benchflow-task-posttrain/tests/test_pipeline.py index 06c6896..2161289 100644 --- a/pipelines/benchflow-task-posttrain/tests/test_pipeline.py +++ b/pipelines/benchflow-task-posttrain/tests/test_pipeline.py @@ -8,6 +8,7 @@ from posttrainarena.benchflow_pipeline.config import load_config from posttrainarena.benchflow_pipeline.io import directory_sha256 +from posttrainarena.benchflow_pipeline.grpo import grpo_training_recipe from posttrainarena.benchflow_pipeline.pipeline import ( Pipeline, _sha256, @@ -73,6 +74,8 @@ def test_dry_run_writes_score_schema_without_heavy_dependencies(tmp_path: Path) assert saved["schema_version"] == 1 assert saved["grpo_planned"] is True assert saved["grpo_ran"] is False + assert saved["grpo_effective_update"] is None + assert saved["grpo_training"] is None assert saved["harness"]["agent"] == "opencode" assert saved["teacher"]["require_all_tasks"] is True assert saved["sft"]["lora_r"] == config.sft.lora_r @@ -121,6 +124,80 @@ def test_dry_run_writes_score_schema_without_heavy_dependencies(tmp_path: Path) ) +def test_score_compacts_grpo_training_diagnostics(tmp_path: Path) -> None: + config = load_config(ROOT / "configs/qwen3-4b-data-agent-smoke.toml") + config = replace(config, output_root=tmp_path) + pipeline = Pipeline(config, run_name="score-diagnostics", dry_run=False) + pipeline.layout.grpo_merged.mkdir(parents=True) + (pipeline.layout.grpo_merged / "train_metrics.json").write_text( + json.dumps( + { + "training_recipe": grpo_training_recipe(config), + "metrics": {"train_loss": 0.25}, + "reward_group_diagnostics": { + "nonzero_variance_group_count": 1, + "zero_variance_group_count": 0, + "groups": [{"task_id": "task-a"}], + }, + "lora_b_update_diagnostics": { + "available": True, + "nonzero_tensor_count": 1, + }, + } + ) + ) + + pipeline._write_score( + baseline_score=0.0, + sft_score=0.5, + grpo_gate_score=0.5, + final_score=1.0, + final_model=str(pipeline.layout.grpo_merged), + grpo_planned=True, + grpo_ran=True, + ) + + saved = json.loads((pipeline.layout.reports / "score.json").read_text()) + assert saved["grpo_effective_update"] is True + assert saved["grpo_training"]["metrics"] == {"train_loss": 0.25} + assert "groups" not in saved["grpo_training"]["reward_groups"] + + saved_metrics = json.loads( + (pipeline.layout.grpo_merged / "train_metrics.json").read_text() + ) + saved_metrics["lora_b_update_diagnostics"] = { + "available": False, + "nonzero_tensor_count": 0, + } + (pipeline.layout.grpo_merged / "train_metrics.json").write_text( + json.dumps(saved_metrics) + ) + pipeline._write_score( + baseline_score=0.0, + sft_score=0.5, + grpo_gate_score=0.5, + final_score=1.0, + final_model=str(pipeline.layout.grpo_merged), + grpo_planned=True, + grpo_ran=True, + ) + saved = json.loads((pipeline.layout.reports / "score.json").read_text()) + assert saved["grpo_effective_update"] is False + + pipeline._write_score( + baseline_score=0.0, + sft_score=0.5, + grpo_gate_score=0.5, + final_score=0.5, + final_model=str(pipeline.layout.sft_merged), + grpo_planned=False, + grpo_ran=False, + ) + saved = json.loads((pipeline.layout.reports / "score.json").read_text()) + assert saved["grpo_effective_update"] is None + assert saved["grpo_training"] is None + + def test_pipeline_rejects_train_eval_overlap(tmp_path: Path) -> None: config = load_config(ROOT / "configs/qwen3-4b-data-agent-smoke.toml") overlap = tmp_path / "eval.txt" @@ -418,17 +495,83 @@ def test_grpo_run_policy_can_force_zero_reward_training(tmp_path: Path) -> None: assert forced._should_run_grpo(0.0) is True +def test_resumed_sft_restarts_downstream_artifacts( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = load_config(ROOT / "configs/qwen3-4b-data-agent-smoke.toml") + config = replace(config, output_root=tmp_path) + pipeline = Pipeline( + config, + run_name="resume-sft", + dry_run=False, + resume=True, + ) + monkeypatch.setattr( + "posttrainarena.benchflow_pipeline.sft.train_sft", + lambda **_: None, + ) + stale_paths = [ + pipeline.layout.sft_adapter / "stale.json", + pipeline.layout.sft_merged / "stale.json", + pipeline.layout.grpo_adapter / "stale.json", + pipeline.layout.grpo_merged / "stale.json", + pipeline.layout.jobs / "sft" / "stale.json", + pipeline.layout.jobs / "grpo-train" / "stale.json", + pipeline.layout.jobs / "posttrain" / "stale.json", + ] + for path in stale_paths: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("{}") + stale_metrics = pipeline.layout.results / "posttrain_eval.json" + stale_metrics.parent.mkdir(parents=True) + stale_metrics.write_text("{}") + stale_score = pipeline.layout.reports / "score.json" + stale_score.parent.mkdir(parents=True) + stale_score.write_text("{}") + + pipeline._train_sft(str(pipeline.layout.sft_merged)) + + assert all(not path.exists() for path in stale_paths) + assert not stale_metrics.exists() + assert not stale_score.exists() + + +def test_resumed_sft_dry_run_does_not_delete_artifacts(tmp_path: Path) -> None: + config = load_config(ROOT / "configs/qwen3-4b-data-agent-smoke.toml") + config = replace(config, output_root=tmp_path) + pipeline = Pipeline( + config, + run_name="resume-sft-dry-run", + dry_run=True, + resume=True, + ) + stale = pipeline.layout.sft_adapter / "stale.json" + stale.parent.mkdir(parents=True) + stale.write_text("{}") + + pipeline._train_sft(str(pipeline.layout.sft_merged)) + + assert stale.is_file() + assert pipeline.runner.commands[-1]["name"] == "train_sft" + + def test_resumed_grpo_restarts_stage_instead_of_reusing_rollouts( tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, ) -> None: config = load_config(ROOT / "configs/qwen3-4b-data-agent-smoke.toml") config = replace(config, output_root=tmp_path) pipeline = Pipeline( config, run_name="resume-grpo", - dry_run=True, + dry_run=False, resume=True, ) + monkeypatch.setattr( + "posttrainarena.benchflow_pipeline.grpo.train_grpo", + lambda **_: None, + ) stale = pipeline.layout.jobs / "grpo-train" / "stale.json" stale.parent.mkdir(parents=True) stale.write_text("{}") @@ -438,6 +581,15 @@ def test_resumed_grpo_restarts_stage_instead_of_reusing_rollouts( stale_merged = pipeline.layout.grpo_merged / "stale.json" stale_merged.parent.mkdir(parents=True) stale_merged.write_text("{}") + stale_posttrain = pipeline.layout.jobs / "posttrain" / "stale.json" + stale_posttrain.parent.mkdir(parents=True) + stale_posttrain.write_text("{}") + stale_posttrain_metrics = pipeline.layout.results / "posttrain_eval.json" + stale_posttrain_metrics.parent.mkdir(parents=True) + stale_posttrain_metrics.write_text("{}") + stale_score = pipeline.layout.reports / "score.json" + stale_score.parent.mkdir(parents=True) + stale_score.write_text("{}") pipeline._train_grpo( input_model=config.model, @@ -447,6 +599,30 @@ def test_resumed_grpo_restarts_stage_instead_of_reusing_rollouts( assert not stale.exists() assert not stale_adapter.exists() assert not stale_merged.exists() + assert not stale_posttrain.exists() + assert not stale_posttrain_metrics.exists() + assert not stale_score.exists() + + +def test_resumed_grpo_dry_run_does_not_delete_artifacts(tmp_path: Path) -> None: + config = load_config(ROOT / "configs/qwen3-4b-data-agent-smoke.toml") + config = replace(config, output_root=tmp_path) + pipeline = Pipeline( + config, + run_name="resume-grpo-dry-run", + dry_run=True, + resume=True, + ) + stale = pipeline.layout.jobs / "grpo-train" / "stale.json" + stale.parent.mkdir(parents=True) + stale.write_text("{}") + + pipeline._train_grpo( + input_model=config.model, + output_model=str(pipeline.layout.grpo_merged), + ) + + assert stale.is_file() assert pipeline.runner.commands[-1]["resume_policy"] == "restart-stage" @@ -623,6 +799,7 @@ def test_resume_checkpoint_digests_detect_tampering(tmp_path: Path) -> None: "mode": "grpo", "model": str(pipeline.layout.sft_merged), "task_ids": pipeline.train_task_ids, + "training_recipe": grpo_training_recipe(config), "adapter_dir": str(pipeline.layout.grpo_adapter), "merged_model_dir": str(pipeline.layout.grpo_merged), "base_checkpoint_sha256": directory_sha256(pipeline.layout.sft_merged), @@ -636,6 +813,20 @@ def test_resume_checkpoint_digests_detect_tampering(tmp_path: Path) -> None: input_model=str(pipeline.layout.sft_merged), output_model=pipeline.layout.grpo_merged, ) + for changed_config in ( + replace(config, runtime=replace(config.runtime, num_generations=4)), + replace(config, grpo=replace(config.grpo, generation_batch_size=4)), + ): + changed_pipeline = Pipeline( + changed_config, + run_name="checkpoint-digests", + dry_run=True, + ) + assert not changed_pipeline._grpo_checkpoint_is_current( + grpo_metrics_path, + input_model=str(pipeline.layout.sft_merged), + output_model=pipeline.layout.grpo_merged, + ) (pipeline.layout.grpo_merged / "model.safetensors").write_bytes(b"tampered") assert not pipeline._grpo_checkpoint_is_current( @@ -700,6 +891,71 @@ def test_resume_allows_increased_grpo_completion_budget(tmp_path: Path) -> None: ) +def test_resume_allows_stricter_grpo_sampling_recipe(tmp_path: Path) -> None: + config = load_config(ROOT / "configs/qwen3-4b-data-agent-smoke.toml") + config = replace(config, output_root=tmp_path) + original = Pipeline(config, run_name="resume-grpo-sampling", dry_run=True) + original._prepare_run_plan() + plan_path = original.layout.reports / "plan.json" + original_plan = json.loads(plan_path.read_text()) + original_plan["grpo"].pop("require_reward_variance") + plan_path.write_text(json.dumps(original_plan)) + changed = Pipeline( + replace( + config, + runtime=replace(config.runtime, num_generations=8), + grpo=replace( + config.grpo, + generation_batch_size=8, + require_reward_variance=True, + ), + ), + run_name="resume-grpo-sampling", + dry_run=True, + resume=True, + ) + + changed._prepare_run_plan() + + updated = json.loads((changed.layout.reports / "plan.json").read_text()) + assert updated["runtime"]["num_generations"] == 8 + assert updated["grpo"]["generation_batch_size"] == 8 + assert updated["grpo"]["require_reward_variance"] is True + + +def test_resume_rejects_weaker_grpo_sampling_recipe(tmp_path: Path) -> None: + config = load_config(ROOT / "configs/qwen3-4b-data-agent-smoke.toml") + config = replace( + config, + output_root=tmp_path, + runtime=replace(config.runtime, num_generations=8), + grpo=replace( + config.grpo, + generation_batch_size=8, + require_reward_variance=True, + ), + ) + original = Pipeline(config, run_name="resume-grpo-weaker", dry_run=True) + original._prepare_run_plan() + changed = Pipeline( + replace( + config, + runtime=replace(config.runtime, num_generations=2), + grpo=replace( + config.grpo, + generation_batch_size=2, + require_reward_variance=False, + ), + ), + run_name="resume-grpo-weaker", + dry_run=True, + resume=True, + ) + + with pytest.raises(RuntimeError, match="Changed fields: grpo, runtime"): + changed._prepare_run_plan() + + def test_resume_rejects_other_teacher_plan_changes(tmp_path: Path) -> None: config = load_config(ROOT / "configs/qwen3-4b-data-agent-smoke.toml") config = replace(config, output_root=tmp_path) diff --git a/pipelines/benchflow-task-posttrain/tests/test_publishing.py b/pipelines/benchflow-task-posttrain/tests/test_publishing.py index 5572146..f6bf6dc 100644 --- a/pipelines/benchflow-task-posttrain/tests/test_publishing.py +++ b/pipelines/benchflow-task-posttrain/tests/test_publishing.py @@ -27,6 +27,7 @@ def test_build_run_record_reads_score_and_removes_contact_data(tmp_path: Path) - "score_after_posttrain": 0.25, "delta_score": 0.15, "grpo_ran": True, + "grpo_effective_update": True, "train_task_ids": ["a", "b"], "eval_task_ids": ["c"], } @@ -42,6 +43,7 @@ def test_build_run_record_reads_score_and_removes_contact_data(tmp_path: Path) - ) assert record["delta_score"] == 0.15 + assert record["grpo_effective_update"] is True assert record["train_task_count"] == 2 assert record["eval_task_count"] == 1 assert "contact_email" not in record