From 36f28d27daa9b0216c61901345a4e85fcef1c180 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 13:26:00 +0800 Subject: [PATCH 01/16] feat: add MemAgent multi-turn training example --- examples/mem_agent/README.md | 79 ++++++ examples/mem_agent/__init__.py | 2 + examples/mem_agent/compare_results.py | 83 ++++++ examples/mem_agent/config.yaml | 16 ++ examples/mem_agent/convert-to-hf.sh | 18 ++ examples/mem_agent/convert.py | 116 ++++++++ examples/mem_agent/eval_ruler_hqa.py | 259 ++++++++++++++++++ examples/mem_agent/metrics.py | 63 +++++ examples/mem_agent/prepare-data.sh | 33 +++ examples/mem_agent/prepare_data.py | 221 +++++++++++++++ examples/mem_agent/prompts.py | 77 ++++++ examples/mem_agent/reward.py | 57 ++++ examples/mem_agent/rollout.py | 237 ++++++++++++++++ examples/mem_agent/run-eval.sh | 70 +++++ examples/mem_agent/run-paired-eval.sh | 38 +++ examples/mem_agent/run-pipeline.sh | 34 +++ examples/mem_agent/run-qwen3-4B-train.sh | 131 +++++++++ relax/backends/megatron/actor.py | 36 ++- relax/core/controller.py | 19 +- relax/distributed/ray/rollout.py | 85 ++++-- relax/engine/rollout/sglang_rollout.py | 22 +- relax/utils/utils.py | 71 ++++- tests/examples/mem_agent/__init__.py | 1 + .../mem_agent/test_compare_results.py | 14 + tests/examples/mem_agent/test_convert.py | 91 ++++++ .../mem_agent/test_data_and_metrics.py | 65 +++++ tests/examples/mem_agent/test_eval.py | 47 ++++ .../mem_agent/test_mock_integration.py | 109 ++++++++ .../mem_agent/test_recipe_contract.py | 38 +++ tests/examples/mem_agent/test_reward.py | 41 +++ tests/examples/mem_agent/test_rollout.py | 210 ++++++++++++++ tests/utils/test_custom_sample_converter.py | 139 ++++++++++ 32 files changed, 2485 insertions(+), 37 deletions(-) create mode 100644 examples/mem_agent/README.md create mode 100644 examples/mem_agent/__init__.py create mode 100644 examples/mem_agent/compare_results.py create mode 100644 examples/mem_agent/config.yaml create mode 100755 examples/mem_agent/convert-to-hf.sh create mode 100644 examples/mem_agent/convert.py create mode 100644 examples/mem_agent/eval_ruler_hqa.py create mode 100644 examples/mem_agent/metrics.py create mode 100755 examples/mem_agent/prepare-data.sh create mode 100644 examples/mem_agent/prepare_data.py create mode 100644 examples/mem_agent/prompts.py create mode 100644 examples/mem_agent/reward.py create mode 100644 examples/mem_agent/rollout.py create mode 100755 examples/mem_agent/run-eval.sh create mode 100755 examples/mem_agent/run-paired-eval.sh create mode 100755 examples/mem_agent/run-pipeline.sh create mode 100755 examples/mem_agent/run-qwen3-4B-train.sh create mode 100644 tests/examples/mem_agent/__init__.py create mode 100644 tests/examples/mem_agent/test_compare_results.py create mode 100644 tests/examples/mem_agent/test_convert.py create mode 100644 tests/examples/mem_agent/test_data_and_metrics.py create mode 100644 tests/examples/mem_agent/test_eval.py create mode 100644 tests/examples/mem_agent/test_mock_integration.py create mode 100644 tests/examples/mem_agent/test_recipe_contract.py create mode 100644 tests/examples/mem_agent/test_reward.py create mode 100644 tests/examples/mem_agent/test_rollout.py create mode 100644 tests/utils/test_custom_sample_converter.py diff --git a/examples/mem_agent/README.md b/examples/mem_agent/README.md new file mode 100644 index 000000000..91f3e5ffb --- /dev/null +++ b/examples/mem_agent/README.md @@ -0,0 +1,79 @@ +# MemAgent on ReLax + +This example trains Qwen3-4B to update a bounded textual memory while reading a long document chunk by chunk. Every memory-update turn and the final-answer turn is saved as an independent training row. Only the final boxed answer receives a rule-based reward; GRPO normalization happens before the trajectory is expanded. + +The reproducibility contract is frozen to: + +- model: `Qwen/Qwen3-4B@1cfa9a7208912126459214e8b04321603b3df60c`; +- dataset: `BytedTsinghua-SIA/hotpotqa@27275ff4fee67ac0acb6478e405e7ac07efbdc1a`; +- chunk/memory/final limits: 2048/1024/256 tokens, at most 64 chunks; +- GRPO group size 8, split credit, LR `1e-6`, KL coefficient `0.001`; +- 100 rollout steps with checkpoints every 50 steps. + +## Prepare model and data + +Download the exact model revision to a local directory with your preferred Hugging Face client. Then prepare all frozen train/eval files and their SHA-256 manifest: + +```bash +DATA_DIR=/data/mem-agent bash examples/mem_agent/prepare-data.sh +``` + +`prepare-data.sh` downloads `hotpotqa_train_32k.parquet`, `hotpotqa_dev.parquet`, and `eval_50/200/800.json` at the pinned dataset revision. It writes converted JSONL files plus `artifact_manifest.json`. + +## Train + +```bash +MODEL_PATH=/data/models/Qwen3-4B \ +DATA_DIR=/data/mem-agent \ +SAVE_DIR=/data/checkpoints/mem-agent-relax \ +bash examples/mem_agent/run-qwen3-4B-train.sh +``` + +For the required two-step correctness smoke, add `NUM_ROLLOUT=2` and use the same command. A smoke run validates the pipeline but is not an effects result. + +The train-side SGLang context envelope is 8192 tokens because each request is an independent chunk turn (2K chunk + 1K memory + response), not the concatenated trajectory. This retains the frozen 9216-token per-GPU packing budget and sample-mean loss used for split credit. + +## Convert and evaluate + +```bash +MODEL_PATH=/data/models/Qwen3-4B \ +CHECKPOINT_DIR=/data/checkpoints/mem-agent-relax \ +CHECKPOINT_TAG=iter_0000099 \ +bash examples/mem_agent/convert-to-hf.sh + +MODEL_PATH=/data/checkpoints/mem-agent-relax-HF/iter_0000099 \ +TOKENIZER_PATH=/data/models/Qwen3-4B \ +DATA_DIR=/data/mem-agent \ +RESULTS_DIR=/data/results/mem-agent-relax \ +bash examples/mem_agent/run-eval.sh +``` + +The evaluator writes raw per-sample JSONL and a summary JSON for HotpotQA dev and RULER-HQA 50/200/800. Failed requests remain in the denominator with score zero. `boxed_em_pct` is the HotpotQA reward-compatible accuracy and `sub_em_pct` is the primary VIME-compatible RULER-HQA metric. Set `MODE=base` to run the single-context base baseline; its context truncation always preserves the question and answer instruction. + +`TOKENIZER_PATH` should point to the frozen base snapshot. `run-pipeline.sh` preserves it automatically before switching `MODEL_PATH` to the converted checkpoint. When `NUM_ROLLOUT=2` is used, the pipeline also selects `iter_0000001` automatically instead of the 100-step default `iter_0000099`. + +For the VIME reproduction tolerance, evaluate an official VIME checkpoint when one is available; otherwise use a checkpoint produced once from the fixed VIME recipe. The paired runner holds the tokenizer, data, prompts, sampling parameters, and evaluator constant. It evaluates RULER-HQA 50/200/800 by default; `LENGTHS` can freeze a smaller pre-agreed subset before either result is observed. It exits non-zero when any selected `sub_em_pct` differs by more than 3 percentage points and retains raw per-sample output for review: + +```bash +VIME_MODEL_PATH=/data/checkpoints/vime-hf \ +RELAX_MODEL_PATH=/data/checkpoints/mem-agent-relax-hf \ +TOKENIZER_PATH=/data/models/Qwen3-4B \ +DATA_DIR=/data/mem-agent \ +RESULTS_DIR=/data/results/vime-vs-relax \ +LENGTHS="50 200 800" \ +bash examples/mem_agent/run-paired-eval.sh +``` + +## One-command chain + +With the environment paths set, `run-pipeline.sh` executes data preparation, training, checkpoint conversion, and evaluation in order: + +```bash +MODEL_PATH=/data/models/Qwen3-4B \ +DATA_DIR=/data/mem-agent \ +SAVE_DIR=/data/checkpoints/mem-agent-relax \ +RESULTS_DIR=/data/results/mem-agent-relax \ +bash examples/mem_agent/run-pipeline.sh +``` + +GPU execution is intentionally not started by the CPU test suite. The caller remains responsible for starting the ReLax/Ray environment described by the repository deployment guide. diff --git a/examples/mem_agent/__init__.py b/examples/mem_agent/__init__.py new file mode 100644 index 000000000..d577a07bd --- /dev/null +++ b/examples/mem_agent/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""MemAgent recurrent-memory training and evaluation example.""" diff --git a/examples/mem_agent/compare_results.py b/examples/mem_agent/compare_results.py new file mode 100644 index 000000000..b3d2d85fb --- /dev/null +++ b/examples/mem_agent/compare_results.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Compare paired VIME and ReLax evaluation summaries.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +def compare_pair( + label: str, + vime_summary: dict[str, Any], + relax_summary: dict[str, Any], + metric: str = "sub_em_pct", + tolerance_pp: float = 3.0, +) -> dict[str, Any]: + """Build one auditable percentage-point comparison. + + Both summaries must come from the same evaluator/data recipe. The function + intentionally compares the reported percentage field directly: ``3.0`` + therefore means three percentage points, not a three-percent relative gap. + """ + if metric not in vime_summary or metric not in relax_summary: + raise KeyError(f"Metric {metric!r} must exist in both summaries.") + vime_value = float(vime_summary[metric]) + relax_value = float(relax_summary[metric]) + gap_pp = abs(relax_value - vime_value) + return { + "label": label, + "metric": metric, + "vime": vime_value, + "relax": relax_value, + "absolute_gap_pp": gap_pp, + "tolerance_pp": tolerance_pp, + "passed": gap_pp <= tolerance_pp, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--pair", + nargs=3, + action="append", + metavar=("LABEL", "VIME_SUMMARY", "RELAX_SUMMARY"), + required=True, + help="Repeat for every RULER-HQA length selected for acceptance.", + ) + parser.add_argument("--metric", default="sub_em_pct") + parser.add_argument("--tolerance-pp", type=float, default=3.0) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + comparisons = [] + for label, vime_path, relax_path in args.pair: + with Path(vime_path).open(encoding="utf-8") as source: + vime_summary = json.load(source) + with Path(relax_path).open(encoding="utf-8") as source: + relax_summary = json.load(source) + comparison = compare_pair(label, vime_summary, relax_summary, args.metric, args.tolerance_pp) + comparison["vime_summary"] = str(vime_path) + comparison["relax_summary"] = str(relax_path) + comparisons.append(comparison) + + report = { + "metric": args.metric, + "tolerance_pp": args.tolerance_pp, + "passed": all(item["passed"] for item in comparisons), + "comparisons": comparisons, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as destination: + json.dump(report, destination, ensure_ascii=False, indent=2) + destination.write("\n") + print(json.dumps(report, ensure_ascii=False, indent=2)) + if not report["passed"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/mem_agent/config.yaml b/examples/mem_agent/config.yaml new file mode 100644 index 000000000..8ead9f1b5 --- /dev/null +++ b/examples/mem_agent/config.yaml @@ -0,0 +1,16 @@ +mem_agent_chunk_tokens: 2048 +mem_agent_max_memory_tokens: 1024 +mem_agent_max_final_tokens: 256 +mem_agent_max_chunks: 64 +mem_agent_credit_assignment: split +mem_agent_strict_alignment: true +# One trajectory has at most 64 memory turns plus one final-answer turn. +# These fields reserve enough queue capacity, disable post-expansion GRPO +# regrouping, and tell the actor to consume every converted row in one step. +custom_train_sample_expansion_factor: 65 +custom_train_data_group_size: 1 +custom_train_expanded_batch: true +# TP=2 on eight actor GPUs gives four data-parallel consumers (CP stays 1). +mem_agent_train_rows_multiple: 4 +model_id: Qwen/Qwen3-4B +model_revision: 1cfa9a7208912126459214e8b04321603b3df60c diff --git a/examples/mem_agent/convert-to-hf.sh b/examples/mem_agent/convert-to-hf.sh new file mode 100755 index 000000000..5eb085ce7 --- /dev/null +++ b/examples/mem_agent/convert-to-hf.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +RELAX_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" +MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to the frozen base model directory.}" +CHECKPOINT_DIR="${CHECKPOINT_DIR:?Set CHECKPOINT_DIR to the ReLax checkpoint root.}" +CHECKPOINT_TAG="${CHECKPOINT_TAG:-iter_0000099}" +HF_OUTPUT_DIR="${HF_OUTPUT_DIR:-${CHECKPOINT_DIR}-HF/${CHECKPOINT_TAG}}" + +python3 "${RELAX_ROOT}/scripts/tools/convert_torch_dist_to_hf_bridge.py" \ + --input-dir "${CHECKPOINT_DIR}/${CHECKPOINT_TAG}" \ + --output-dir "${HF_OUTPUT_DIR}" \ + --origin-hf-dir "${MODEL_PATH}" + +echo "Converted checkpoint: ${HF_OUTPUT_DIR}" diff --git a/examples/mem_agent/convert.py b/examples/mem_agent/convert.py new file mode 100644 index 000000000..b6cfcfe33 --- /dev/null +++ b/examples/mem_agent/convert.py @@ -0,0 +1,116 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Expand MemAgent trajectories into independent ReLax training samples.""" + +from __future__ import annotations + +from typing import Any + +from relax.utils.types import Sample +from relax.utils.utils import dict_to_tensordict, post_process_rewards + + +def _get_turns(sample: Sample) -> list[dict[str, Any]]: + metadata = sample.train_metadata or {} + turns = metadata.get("mem_agent_turns", metadata.get("turns")) + if not isinstance(turns, list) or not turns: + raise ValueError(f"Sample index={sample.index} has no MemAgent turns.") + return turns + + +def _validate_turn(sample: Sample, turn: dict[str, Any]) -> None: + required = {"tokens", "response_length", "loss_mask", "rollout_log_probs"} + missing = required.difference(turn) + if missing: + raise ValueError(f"Sample index={sample.index} turn is missing fields: {sorted(missing)}") + response_length = int(turn["response_length"]) + if response_length <= 0 or len(turn["tokens"]) < response_length: + raise ValueError(f"Sample index={sample.index} has an invalid response_length.") + if len(turn["loss_mask"]) != response_length: + raise ValueError(f"Sample index={sample.index} has a misaligned loss_mask.") + if len(turn["rollout_log_probs"]) != response_length: + raise ValueError(f"Sample index={sample.index} has misaligned rollout_log_probs.") + + +def convert_samples(args: Any, samples: list[Sample]): + """Normalize trajectory rewards first, then expand every saved turn.""" + if not samples: + raise ValueError("MemAgent converter received an empty sample list.") + if not getattr(args, "mem_agent_strict_alignment", True): + raise ValueError("MemAgent training requires mem_agent_strict_alignment=true.") + for sample in samples: + if sample.status in (Sample.Status.ABORTED, Sample.Status.FAILED): + raise ValueError(f"Cannot train from sample index={sample.index} with status={sample.status.value}.") + + # Normalize while there is still exactly one row per trajectory. Expanding + # first would overweight long documents in the group statistics. + raw_rewards, advantages = post_process_rewards(args, samples) + credit_assignment = getattr(args, "mem_agent_credit_assignment", "split") + if credit_assignment not in ("split", "share"): + raise ValueError("mem_agent_credit_assignment must be either 'split' or 'share'.") + + # Every trajectory in one GRPO prompt group reads the same context, so it + # must have the same number of turns. Besides catching partial trajectories, + # this makes each expanded group n_samples_per_prompt * turn_count rows. + turn_counts_by_group: dict[int, set[int]] = {} + for sample in samples: + if sample.group_index is None: + raise ValueError("MemAgent samples require group_index.") + turn_counts_by_group.setdefault(sample.group_index, set()).add(len(_get_turns(sample))) + inconsistent_groups = { + group_index: counts for group_index, counts in turn_counts_by_group.items() if len(counts) != 1 + } + if inconsistent_groups: + raise ValueError(f"MemAgent prompt groups have inconsistent turn counts: {inconsistent_groups}") + + train_data: dict[str, list[Any]] = { + "tokens": [], + "response_lengths": [], + "loss_masks": [], + "rollout_log_probs": [], + "rewards": [], + "raw_reward": [], + "truncated": [], + "sample_indices": [], + "trajectory_indices": [], + "turn_indices": [], + "total_lengths": [], + } + + for trajectory_position, (sample, raw_reward, advantage) in enumerate( + zip(samples, raw_rewards, advantages, strict=True) + ): + turns = _get_turns(sample) + turn_credit = float(advantage) / len(turns) if credit_assignment == "split" else float(advantage) + # This expansion is deliberately lossless. Unlike the fixed VIME + # helper, no tail rows are trimmed to a global-batch multiple. + for fallback_turn_index, turn in enumerate(turns): + _validate_turn(sample, turn) + tokens = list(turn["tokens"]) + response_length = int(turn["response_length"]) + loss_mask = list(turn["loss_mask"]) + if sample.remove_sample: + loss_mask = [0] * response_length + train_data["tokens"].append(tokens) + train_data["response_lengths"].append(response_length) + train_data["loss_masks"].append(loss_mask) + train_data["rollout_log_probs"].append(list(turn["rollout_log_probs"])) + train_data["rewards"].append(turn_credit) + train_data["raw_reward"].append(float(raw_reward)) + train_data["truncated"].append(int(turn.get("finish_reason") == Sample.Status.TRUNCATED.value)) + train_data["sample_indices"].append(sample.index) + train_data["trajectory_indices"].append(trajectory_position) + train_data["turn_indices"].append(int(turn.get("turn_index", fallback_turn_index))) + train_data["total_lengths"].append(len(tokens)) + + required_multiple = int(getattr(args, "mem_agent_train_rows_multiple", 1)) + if required_multiple <= 0: + raise ValueError("mem_agent_train_rows_multiple must be positive.") + if len(train_data["tokens"]) % required_multiple: + raise ValueError( + f"Expanded MemAgent row count {len(train_data['tokens'])} is not divisible by " + f"mem_agent_train_rows_multiple={required_multiple}." + ) + + if getattr(args, "debug_train_only", False): + return train_data + return dict_to_tensordict(train_data, len(train_data["tokens"])) diff --git a/examples/mem_agent/eval_ruler_hqa.py b/examples/mem_agent/eval_ruler_hqa.py new file mode 100644 index 000000000..d472e1e97 --- /dev/null +++ b/examples/mem_agent/eval_ruler_hqa.py @@ -0,0 +1,259 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Evaluate recurrent MemAgent or a single-context baseline on HotpotQA/RULER- +HQA.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +from pathlib import Path +from typing import Any + +import aiohttp + +from examples.mem_agent.metrics import aggregate, exact_match, f1_score, sub_exact_match +from examples.mem_agent.prompts import ( + NO_MEMORY, + final_instruction, + memory_instruction, + strip_stop_tokens, + truncate_text_to_tokens, +) +from examples.mem_agent.reward import exact_match_any, extract_last_boxed + + +def load_data(path: Path) -> list[dict[str, Any]]: + if path.suffix == ".jsonl": + with path.open(encoding="utf-8") as source: + raw = [json.loads(line) for line in source if line.strip()] + else: + with path.open(encoding="utf-8") as source: + raw = json.load(source) + if isinstance(raw, dict): + raw = list(raw.values()) + + data = [] + for index, item in enumerate(raw): + if "input" in item: + normalized = dict(item) + normalized.setdefault("_id", index) + else: + metadata = item.get("metadata") or {} + normalized = { + "_id": index, + "input": item.get("prompt", metadata.get("question", "")), + "answers": metadata.get("ground_truth", [item.get("label", "")]), + "context": metadata.get("context", ""), + "num_docs": metadata.get("num_docs", 0), + } + if normalized.get("input") and normalized.get("context") and normalized.get("answers"): + data.append(normalized) + return data + + +async def _chat_once( + session: aiohttp.ClientSession, + base_url: str, + api_key: str, + model: str, + instruction: str, + temperature: float, + top_p: float, + max_tokens: int, +) -> str: + payload = { + "model": model, + "messages": [{"role": "user", "content": instruction}], + "temperature": temperature, + "top_p": top_p, + "max_tokens": max_tokens, + } + async with session.post( + f"{base_url.rstrip('/')}/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json=payload, + ) as response: + body = await response.text() + if response.status != 200: + raise RuntimeError(f"HTTP {response.status}: {body[:300]}") + result = json.loads(body) + return str(result["choices"][0]["message"]["content"]) + + +async def recurrent_infer( + item: dict[str, Any], + args: argparse.Namespace, + tokenizer: Any, + session: aiohttp.ClientSession, +) -> tuple[str, dict[str, Any]]: + question = str(item["input"]).strip() + context_ids = tokenizer.encode(str(item["context"]).strip(), add_special_tokens=False) + all_chunks = [ + context_ids[offset : offset + args.chunk_tokens] for offset in range(0, len(context_ids), args.chunk_tokens) + ] + chunks = all_chunks[: args.max_chunks] + memory = NO_MEMORY + memory_lengths = [] + for chunk_ids in chunks: + chunk = tokenizer.decode(chunk_ids, skip_special_tokens=True) + generated_memory = strip_stop_tokens( + await _chat_once( + session, + args.base_url, + args.api_key, + args.model, + memory_instruction(question, memory, chunk), + args.temperature, + args.top_p, + args.max_memory_tokens, + ) + ) + # Match training exactly: the next turn only sees the re-tokenized, + # bounded output from the immediately preceding memory update. + memory, memory_length = truncate_text_to_tokens(tokenizer, generated_memory, args.max_memory_tokens) + memory_lengths.append(memory_length) + answer = await _chat_once( + session, + args.base_url, + args.api_key, + args.model, + final_instruction(question, memory), + args.temperature, + args.top_p, + args.max_final_tokens, + ) + return answer, { + "num_chunks": len(chunks), + "context_truncated": len(all_chunks) > args.max_chunks, + "memory_token_lengths": memory_lengths, + } + + +async def base_infer( + item: dict[str, Any], + args: argparse.Namespace, + tokenizer: Any, + session: aiohttp.ClientSession, +) -> tuple[str, dict[str, Any]]: + suffix = f"\n\nQuestion: {item['input']}\nPlease answer the question and put the answer in \\boxed{{}}." + suffix_ids = tokenizer.encode(suffix, add_special_tokens=False) + if len(suffix_ids) > args.max_input_tokens: + raise ValueError("Question and answer instruction exceed max_input_tokens; context cannot be retained.") + context_ids = tokenizer.encode(str(item["context"]), add_special_tokens=False) + context_budget = args.max_input_tokens - len(suffix_ids) + retained_context_ids = context_ids[:context_budget] + instruction = tokenizer.decode(retained_context_ids, skip_special_tokens=True) + suffix + # Decoding then concatenating can change a BPE boundary by a token. Trim + # context only, never the question suffix, until the served prompt fits. + while len(tokenizer.encode(instruction, add_special_tokens=False)) > args.max_input_tokens: + retained_context_ids = retained_context_ids[:-1] + instruction = tokenizer.decode(retained_context_ids, skip_special_tokens=True) + suffix + context_truncated = len(retained_context_ids) < len(context_ids) + answer = await _chat_once( + session, + args.base_url, + args.api_key, + args.model, + instruction, + args.temperature, + args.top_p, + args.max_final_tokens, + ) + return answer, {"num_chunks": 1, "context_truncated": context_truncated, "memory_token_lengths": []} + + +async def run_evaluation( + data: list[dict[str, Any]], args: argparse.Namespace, tokenizer: Any +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + semaphore = asyncio.Semaphore(args.concurrency) + timeout = aiohttp.ClientTimeout(total=args.timeout) + + async with aiohttp.ClientSession(timeout=timeout) as session: + + async def evaluate_one(item: dict[str, Any]) -> dict[str, Any]: + try: + async with semaphore: + if args.mode == "recurrent": + response, diagnostics = await recurrent_infer(item, args, tokenizer, session) + else: + response, diagnostics = await base_infer(item, args, tokenizer, session) + answers = item["answers"] if isinstance(item["answers"], list) else [item["answers"]] + # RULER-HQA's VIME-compatible metrics score the first reference. + # boxed_em additionally mirrors the HotpotQA training reward and + # accepts any annotated answer. + ground_truth = str(answers[0]) + prediction = extract_last_boxed(response[-300:]) + return { + "_id": item["_id"], + "answer": ground_truth, + "answers": [str(answer) for answer in answers], + "pred": prediction, + "judge_f1": f1_score(prediction, ground_truth), + "judge_em": exact_match(prediction, ground_truth), + "judge_sub_em": sub_exact_match(prediction, ground_truth), + "judge_boxed_em": float(bool(prediction) and exact_match_any(prediction, answers)), + "response": response, + **diagnostics, + } + except Exception as exc: + return {"_id": item["_id"], "error": f"{type(exc).__name__}: {exc}"} + + records = await asyncio.gather(*(evaluate_one(item) for item in data)) + + summary = { + **aggregate(records), + "mode": args.mode, + "model": args.model, + "data_file": str(args.data_file), + "temperature": args.temperature, + "top_p": args.top_p, + "sampling_count": 1, + "chunk_tokens": args.chunk_tokens, + "max_memory_tokens": args.max_memory_tokens, + "max_final_tokens": args.max_final_tokens, + "max_chunks": args.max_chunks, + } + return records, summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-file", type=Path, required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--tokenizer", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--run-name", required=True) + parser.add_argument("--mode", choices=("recurrent", "base"), default="recurrent") + parser.add_argument("--base-url", default="http://127.0.0.1:8000/v1") + parser.add_argument("--api-key", default="EMPTY") + parser.add_argument("--temperature", type=float, default=0.7) + parser.add_argument("--top-p", type=float, default=0.95) + parser.add_argument("--chunk-tokens", type=int, default=2048) + parser.add_argument("--max-memory-tokens", type=int, default=1024) + parser.add_argument("--max-final-tokens", type=int, default=256) + parser.add_argument("--max-chunks", type=int, default=64) + parser.add_argument("--max-input-tokens", type=int, default=7936) + parser.add_argument("--concurrency", type=int, default=16) + parser.add_argument("--timeout", type=int, default=86400) + args = parser.parse_args() + + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + data = load_data(args.data_file) + records, summary = asyncio.run(run_evaluation(data, args, tokenizer)) + args.output_dir.mkdir(parents=True, exist_ok=True) + records_path = args.output_dir / f"{args.run_name}.jsonl" + summary_path = args.output_dir / f"{args.run_name}.summary.json" + with records_path.open("w", encoding="utf-8") as destination: + for record in records: + destination.write(json.dumps(record, ensure_ascii=False) + "\n") + with summary_path.open("w", encoding="utf-8") as destination: + json.dump(summary, destination, ensure_ascii=False, indent=2) + destination.write("\n") + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/examples/mem_agent/metrics.py b/examples/mem_agent/metrics.py new file mode 100644 index 000000000..b47bb633e --- /dev/null +++ b/examples/mem_agent/metrics.py @@ -0,0 +1,63 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""RULER-HQA metrics used by MemAgent evaluation.""" + +from __future__ import annotations + +from collections import Counter + +from examples.mem_agent.reward import normalize_answer + + +def exact_match(prediction: str, ground_truth: str) -> float: + prediction = normalize_answer(prediction) + ground_truth = normalize_answer(ground_truth) + if prediction in ("yes", "no", "noanswer") and prediction != ground_truth: + return 0.0 + if ground_truth in ("yes", "no", "noanswer") and prediction != ground_truth: + return 0.0 + return float(prediction == ground_truth) + + +def sub_exact_match(prediction: str, ground_truth: str) -> float: + prediction = normalize_answer(prediction) + ground_truth = normalize_answer(ground_truth) + if not prediction or not ground_truth: + return 0.0 + return float(ground_truth in prediction or prediction in ground_truth) + + +def f1_score(prediction: str, ground_truth: str) -> float: + prediction_tokens = normalize_answer(prediction).split() + ground_truth_tokens = normalize_answer(ground_truth).split() + if not prediction_tokens or not ground_truth_tokens: + return float(prediction_tokens == ground_truth_tokens) + overlap = Counter(prediction_tokens) & Counter(ground_truth_tokens) + same = sum(overlap.values()) + if same == 0: + return 0.0 + precision = same / len(prediction_tokens) + recall = same / len(ground_truth_tokens) + return 2 * precision * recall / (precision + recall) + + +def aggregate(records: list[dict]) -> dict[str, float | int]: + successful = [record for record in records if not record.get("error")] + # Evaluation failures are zero-score examples, not dropped observations. + # Keeping the original denominator prevents transient serving errors from + # making a run look artificially better. + total = len(records) + result: dict[str, float | int] = { + "total": total, + "successful": len(successful), + "errors": total - len(successful), + } + for output_key, record_key in ( + ("f1", "judge_f1"), + ("em", "judge_em"), + ("sub_em", "judge_sub_em"), + ("boxed_em", "judge_boxed_em"), + ): + value = sum(float(record.get(record_key, 0.0)) for record in successful) / total if total else 0.0 + result[output_key] = value + result[f"{output_key}_pct"] = value * 100 + return result diff --git a/examples/mem_agent/prepare-data.sh b/examples/mem_agent/prepare-data.sh new file mode 100755 index 000000000..fa9c0f9df --- /dev/null +++ b/examples/mem_agent/prepare-data.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +DATA_DIR="${DATA_DIR:?Set DATA_DIR to the MemAgent data output directory.}" +HF_CACHE_DIR="${HF_CACHE_DIR:-${DATA_DIR}/hf-cache}" +MANIFEST="${DATA_DIR}/artifact_manifest.json" +REPO_ID="BytedTsinghua-SIA/hotpotqa" +REVISION="27275ff4fee67ac0acb6478e405e7ac07efbdc1a" + +mkdir -p "${DATA_DIR}" + +convert_file() { + local source_file="$1" + local output_file="$2" + python3 "${SCRIPT_DIR}/prepare_data.py" \ + --hf-file "${source_file}" \ + --output "${DATA_DIR}/${output_file}" \ + --repo-id "${REPO_ID}" \ + --revision "${REVISION}" \ + --cache-dir "${HF_CACHE_DIR}" \ + --manifest "${MANIFEST}" +} + +convert_file hotpotqa_train_32k.parquet train.jsonl +convert_file hotpotqa_dev.parquet dev.jsonl +for length in 50 200 800; do + convert_file "eval_${length}.json" "eval_${length}.jsonl" +done + +echo "Prepared frozen MemAgent data under ${DATA_DIR}" diff --git a/examples/mem_agent/prepare_data.py b/examples/mem_agent/prepare_data.py new file mode 100644 index 000000000..200df86d9 --- /dev/null +++ b/examples/mem_agent/prepare_data.py @@ -0,0 +1,221 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Convert frozen HotpotQA/MemAgent files to ReLax JSONL input.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +from typing import Any, Iterable + + +DEFAULT_DATASET_ID = "BytedTsinghua-SIA/hotpotqa" +DEFAULT_DATASET_REVISION = "27275ff4fee67ac0acb6478e405e7ac07efbdc1a" + + +def _json_safe(value: Any) -> Any: + if hasattr(value, "tolist"): + return _json_safe(value.tolist()) + if hasattr(value, "item") and not isinstance(value, (str, bytes, dict, list, tuple)): + return value.item() + if isinstance(value, dict): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + return value + + +def _first_user_content(prompt: Any) -> str: + prompt = _json_safe(prompt) + if isinstance(prompt, str): + return prompt + if isinstance(prompt, list) and prompt: + first = prompt[0] + return str(first.get("content", "")) if isinstance(first, dict) else str(first) + return "" + + +def convert_row(row: dict[str, Any]) -> dict[str, Any] | None: + """Normalize either training parquet rows or RULER-HQA JSON rows.""" + row = _json_safe(row) + context = row.get("context", "") + if isinstance(context, list): + context = "\n\n".join(str(part) for part in context) + context = str(context).strip() + if not context: + return None + + if row.get("input"): + question = str(row["input"]).strip() + ground_truth = row.get("answers", []) + extra_info = row + else: + extra_info = row.get("extra_info") or {} + question = _first_user_content(row.get("prompt")) or str(extra_info.get("question", "")) + reward_model = row.get("reward_model") or {} + if isinstance(reward_model, str): + reward_model = json.loads(reward_model) + ground_truth = reward_model.get("ground_truth", row.get("ground_truth", [])) + + if isinstance(ground_truth, str): + ground_truth = [ground_truth] + ground_truth = [str(answer) for answer in _json_safe(ground_truth) if str(answer).strip()] + question = question.strip() + if not question or not ground_truth: + return None + + return { + "prompt": question, + "label": ground_truth[0], + "metadata": { + "question": question, + "context": context, + "ground_truth": ground_truth, + "num_docs": int(extra_info.get("num_docs", row.get("num_docs", 0)) or 0), + "data_source": str(row.get("data_source", "hotpotqa")), + }, + } + + +def read_rows(path: Path) -> Iterable[dict[str, Any]]: + suffix = path.suffix.lower() + if suffix == ".parquet": + try: + import pandas as pd + except ImportError as exc: + raise RuntimeError("Reading parquet requires pandas and pyarrow.") from exc + yield from pd.read_parquet(path).to_dict(orient="records") + return + if suffix == ".jsonl": + with path.open(encoding="utf-8") as source: + for line in source: + if line.strip(): + yield json.loads(line) + return + if suffix == ".json": + with path.open(encoding="utf-8") as source: + payload = json.load(source) + if isinstance(payload, list): + yield from payload + elif isinstance(payload, dict) and all(isinstance(value, dict) for value in payload.values()): + yield from payload.values() + else: + yield payload + return + raise ValueError(f"Unsupported input format: {path}") + + +def convert_file(input_path: Path, output_path: Path) -> tuple[int, int]: + output_path.parent.mkdir(parents=True, exist_ok=True) + written = 0 + skipped = 0 + with output_path.open("w", encoding="utf-8") as destination: + for row in read_rows(input_path): + converted = convert_row(row) + if converted is None: + skipped += 1 + continue + destination.write(json.dumps(converted, ensure_ascii=False) + "\n") + written += 1 + return written, skipped + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def update_manifest( + manifest_path: Path, + source_path: Path, + filename: str, + repo_id: str, + revision: str, + output_path: Path | None = None, + written: int | None = None, + skipped: int | None = None, +) -> None: + manifest = {"artifacts": []} + if manifest_path.exists(): + with manifest_path.open(encoding="utf-8") as source: + manifest = json.load(source) + artifact = { + "repo_id": repo_id, + "resolved_revision": revision, + "filename": filename, + "local_path": str(source_path.resolve()), + "size_bytes": source_path.stat().st_size, + "sha256": sha256_file(source_path), + } + # Hash the normalized JSONL as well as the downloaded source. This makes + # preprocessing changes reviewable instead of proving only the input blob. + if output_path is not None: + artifact["converted"] = { + "local_path": str(output_path.resolve()), + "size_bytes": output_path.stat().st_size, + "sha256": sha256_file(output_path), + "written_rows": written, + "skipped_rows": skipped, + } + artifacts = [item for item in manifest.get("artifacts", []) if item.get("filename") != filename] + artifacts.append(artifact) + manifest["artifacts"] = sorted(artifacts, key=lambda item: item["filename"]) + manifest_path.parent.mkdir(parents=True, exist_ok=True) + with manifest_path.open("w", encoding="utf-8") as destination: + json.dump(manifest, destination, ensure_ascii=False, indent=2) + destination.write("\n") + + +def download_hf_file(repo_id: str, revision: str, filename: str, cache_dir: Path | None) -> Path: + try: + from huggingface_hub import hf_hub_download + except ImportError as exc: + raise RuntimeError("Hugging Face download requires huggingface_hub.") from exc + return Path( + hf_hub_download( + repo_id=repo_id, + repo_type="dataset", + revision=revision, + filename=filename, + cache_dir=None if cache_dir is None else str(cache_dir), + ) + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--input", type=Path, help="Local parquet/JSON/JSONL file") + source.add_argument("--hf-file", help="Filename in the frozen Hugging Face dataset") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--repo-id", default=DEFAULT_DATASET_ID) + parser.add_argument("--revision", default=DEFAULT_DATASET_REVISION) + parser.add_argument("--cache-dir", type=Path) + parser.add_argument("--manifest", type=Path) + args = parser.parse_args() + + input_path = args.input + source_filename = input_path.name if input_path else args.hf_file + if input_path is None: + input_path = download_hf_file(args.repo_id, args.revision, args.hf_file, args.cache_dir) + written, skipped = convert_file(input_path, args.output) + if args.manifest: + update_manifest( + args.manifest, + input_path, + source_filename, + args.repo_id, + args.revision, + output_path=args.output, + written=written, + skipped=skipped, + ) + print(f"Written={written} skipped={skipped} output={args.output}") + + +if __name__ == "__main__": + main() diff --git a/examples/mem_agent/prompts.py b/examples/mem_agent/prompts.py new file mode 100644 index 000000000..d19bc784e --- /dev/null +++ b/examples/mem_agent/prompts.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Prompt templates shared by MemAgent rollout and evaluation.""" + +from __future__ import annotations + +from typing import Any + + +NO_MEMORY = "No previous memory" +STOP_TOKEN_STRINGS = ("<|im_end|>", "<|endoftext|>") + +MEMORY_TEMPLATE = """You are presented with a problem, a section of an article that may contain the answer to the problem, and a previous memory. Please read the provided section carefully and update the memory with the new information that helps to answer the problem. Be sure to retain all relevant details from the previous memory while adding any new, useful information. + + +{question} + + + +{memory} + + +
+{chunk} +
+ +Updated memory: +""" + +FINAL_TEMPLATE = """You are presented with a problem and a previous memory. Please answer the problem based on the previous memory and put the answer in \\boxed{{}}. + + +{question} + + + +{memory} + + +Your answer: +""" + + +def strip_stop_tokens(text: str) -> str: + """Remove model terminators retained by some serving backends.""" + for token in STOP_TOKEN_STRINGS: + text = text.replace(token, "") + return text.strip() + + +def truncate_text_to_tokens(tokenizer: Any, text: str, max_tokens: int) -> tuple[str, int]: + """Clamp generated memory after re-tokenization and report its length.""" + if max_tokens <= 0: + raise ValueError("max_tokens must be positive.") + token_ids = tokenizer.encode(text, add_special_tokens=False)[:max_tokens] + bounded_text = tokenizer.decode(token_ids, skip_special_tokens=True) + bounded_ids = tokenizer.encode(bounded_text, add_special_tokens=False) + # Decode/encode is not guaranteed to be token-id preserving for every + # tokenizer. Trim characters only in that rare case so the text inserted + # into the next prompt has a verifiable hard token bound. + while len(bounded_ids) > max_tokens and bounded_text: + bounded_text = bounded_text[:-1] + bounded_ids = tokenizer.encode(bounded_text, add_special_tokens=False) + return bounded_text, len(bounded_ids) + + +def memory_instruction(question: str, memory: str, chunk: str) -> str: + return MEMORY_TEMPLATE.format(question=question, memory=memory, chunk=chunk) + + +def final_instruction(question: str, memory: str) -> str: + return FINAL_TEMPLATE.format(question=question, memory=memory) + + +def render_chat_prompt(tokenizer: Any, instruction: str) -> str: + """Render one independent user turn with the model's chat template.""" + messages = [{"role": "user", "content": instruction}] + return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) diff --git a/examples/mem_agent/reward.py b/examples/mem_agent/reward.py new file mode 100644 index 000000000..87d90b46d --- /dev/null +++ b/examples/mem_agent/reward.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Rule-based final-answer reward for MemAgent.""" + +from __future__ import annotations + +import re +import string +from typing import Any + + +def extract_last_boxed(text: str) -> str: + """Extract the payload of the last balanced ``\\boxed{...}``.""" + start = max(text.rfind("\\boxed{"), text.rfind("\\fbox{")) + if start < 0: + return "" + left = text.find("{", start) + depth = 0 + for position in range(left, len(text)): + if text[position] == "{": + depth += 1 + elif text[position] == "}": + depth -= 1 + if depth == 0: + return text[left + 1 : position].strip() + return "" + + +def normalize_answer(answer: str) -> str: + """Apply the standard HotpotQA exact-match normalization.""" + answer = answer.lower() + answer = "".join(character for character in answer if character not in set(string.punctuation)) + answer = re.sub(r"\b(a|an|the)\b", " ", answer) + return " ".join(answer.split()) + + +def exact_match_any(prediction: str, ground_truths: list[str]) -> bool: + normalized_prediction = normalize_answer(prediction) + return any(normalized_prediction == normalize_answer(str(answer)) for answer in ground_truths) + + +async def reward_func(args: Any, sample: Any, **kwargs: Any) -> dict[str, Any]: + """Score only the final boxed answer; memory turns receive no direct + score.""" + del args, kwargs + metadata = sample.metadata if isinstance(sample.metadata, dict) else {} + final_output = str(metadata.get("final_output") or sample.response or "") + ground_truths = metadata.get("ground_truth") or ([] if sample.label is None else [sample.label]) + if isinstance(ground_truths, str): + ground_truths = [ground_truths] + prediction = extract_last_boxed(final_output[-300:]) + score = float(bool(prediction) and exact_match_any(prediction, list(ground_truths))) + return { + "score": score, + "pred": prediction, + "gt": str(ground_truths[0]) if ground_truths else "", + "diagnostic": "matched" if score else ("missing_boxed" if not prediction else "answer_mismatch"), + } diff --git a/examples/mem_agent/rollout.py b/examples/mem_agent/rollout.py new file mode 100644 index 000000000..f3aa11f2e --- /dev/null +++ b/examples/mem_agent/rollout.py @@ -0,0 +1,237 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Chunk-by-chunk recurrent memory rollout for MemAgent.""" + +from __future__ import annotations + +from argparse import Namespace +from collections.abc import Awaitable, Callable +from typing import Any + +from examples.mem_agent.prompts import ( + NO_MEMORY, + final_instruction, + memory_instruction, + render_chat_prompt, + strip_stop_tokens, + truncate_text_to_tokens, +) +from relax.utils.logging_utils import get_logger +from relax.utils.types import Sample + + +logger = get_logger(__name__) +TurnGenerator = Callable[[Namespace, Sample, dict[str, Any], bool], Awaitable[Sample]] + + +def chunk_context(tokenizer: Any, context: str, chunk_tokens: int, max_chunks: int) -> tuple[list[list[int]], bool]: + """Split context on tokenizer boundaries without overlap or token loss.""" + if chunk_tokens <= 0: + raise ValueError("mem_agent_chunk_tokens must be positive.") + if max_chunks <= 0: + raise ValueError("mem_agent_max_chunks must be positive.") + token_ids = tokenizer.encode(context, add_special_tokens=False) + chunks = [token_ids[offset : offset + chunk_tokens] for offset in range(0, len(token_ids), chunk_tokens)] + return chunks[:max_chunks], len(chunks) > max_chunks + + +def _question_from_sample(sample: Sample) -> str: + question = sample.metadata.get("question") if isinstance(sample.metadata, dict) else None + if question: + return str(question) + if isinstance(sample.prompt, str): + return sample.prompt + if sample.prompt: + return str(sample.prompt[-1].get("content", "")) + return "" + + +def _validate_turn(turn: dict[str, Any], require_log_probs: bool) -> None: + response_length = turn["response_length"] + if response_length <= 0: + raise ValueError("MemAgent turn returned no trainable response tokens.") + if len(turn["tokens"]) < response_length: + raise ValueError("MemAgent turn response_length exceeds total token count.") + if len(turn["loss_mask"]) != response_length: + raise ValueError("MemAgent turn loss_mask is not aligned with response tokens.") + log_probs = turn["rollout_log_probs"] + if require_log_probs and len(log_probs) != response_length: + raise ValueError("MemAgent turn rollout_log_probs are not aligned with response tokens.") + + +def _turn_record(turn_sample: Sample, turn_index: int, kind: str, evaluation: bool) -> dict[str, Any]: + record = { + "turn_index": turn_index, + "kind": kind, + "tokens": list(turn_sample.tokens), + "response_length": turn_sample.response_length, + "loss_mask": list(turn_sample.loss_mask or [1] * turn_sample.response_length), + "rollout_log_probs": list(turn_sample.rollout_log_probs or []), + "finish_reason": turn_sample.status.value, + } + _validate_turn(record, require_log_probs=not evaluation) + return record + + +async def _run_turn( + args: Namespace, + parent: Sample, + prompt: str, + sampling_params: dict[str, Any], + max_new_tokens: int, + evaluation: bool, + generator: TurnGenerator, +) -> Sample: + turn = Sample( + group_index=parent.group_index, + index=parent.index, + prompt=prompt, + metadata={}, + session_id=parent.session_id, + ) + params = {**sampling_params, "max_new_tokens": max_new_tokens} + return await generator(args, turn, params, evaluation) + + +async def generate_trajectory( + args: Namespace, + sample: Sample, + sampling_params: dict[str, Any], + tokenizer: Any, + generator: TurnGenerator | None = None, + evaluation: bool = False, +) -> Sample: + """Generate all independent memory turns and the final-answer turn.""" + sample.metadata = sample.metadata if isinstance(sample.metadata, dict) else {} + context = str(sample.metadata.get("context", "")) + question = _question_from_sample(sample).strip() + if not context or not question: + sample.status = Sample.Status.ABORTED + sample.rollout_log_probs = [] + return sample + + chunk_tokens = int(getattr(args, "mem_agent_chunk_tokens", 2048)) + max_memory_tokens = int(getattr(args, "mem_agent_max_memory_tokens", 1024)) + max_final_tokens = int(getattr(args, "mem_agent_max_final_tokens", 256)) + max_chunks = int(getattr(args, "mem_agent_max_chunks", 64)) + if not getattr(args, "mem_agent_strict_alignment", True): + raise ValueError("MemAgent training requires mem_agent_strict_alignment=true.") + if max_memory_tokens <= 0 or max_final_tokens <= 0: + raise ValueError("MemAgent memory and final response limits must be positive.") + + chunks, context_truncated = chunk_context(tokenizer, context, chunk_tokens, max_chunks) + if not chunks: + sample.status = Sample.Status.ABORTED + sample.rollout_log_probs = [] + return sample + if generator is None: + raise ValueError("A turn generator is required for a non-empty MemAgent trajectory.") + + turns: list[dict[str, Any]] = [] + memory = NO_MEMORY + memory_token_lengths: list[int] = [] + any_turn_truncated = False + + # Every chunk starts an independent conversation. Only the generated + # memory text survives; prior prompts and token history are not appended. + for chunk_ids in chunks: + chunk = tokenizer.decode(chunk_ids, skip_special_tokens=True) + prompt = render_chat_prompt(tokenizer, memory_instruction(question, memory, chunk)) + turn_sample = await _run_turn( + args, + sample, + prompt, + sampling_params, + max_memory_tokens, + evaluation, + generator, + ) + if turn_sample.status in (Sample.Status.ABORTED, Sample.Status.FAILED): + # ReLax retries ABORTED groups. Never return a partially populated + # FAILED trajectory that the transfer path could accidentally see. + sample.status = Sample.Status.ABORTED + sample.rollout_log_probs = [] + return sample + turns.append(_turn_record(turn_sample, len(turns), "memory", evaluation)) + # Overwrite on every turn, including an empty post-processed response; + # carrying the old value forward would silently change the recurrence. + memory, memory_length = truncate_text_to_tokens( + tokenizer, + strip_stop_tokens(turn_sample.response), + max_memory_tokens, + ) + memory_token_lengths.append(memory_length) + any_turn_truncated = any_turn_truncated or turn_sample.status == Sample.Status.TRUNCATED + + # The final request deliberately excludes context/chunks. This enforces + # the question + latest-memory information boundary from the task spec. + final_prompt = render_chat_prompt(tokenizer, final_instruction(question, memory)) + final_sample = await _run_turn( + args, + sample, + final_prompt, + sampling_params, + max_final_tokens, + evaluation, + generator, + ) + if final_sample.status in (Sample.Status.ABORTED, Sample.Status.FAILED): + sample.status = Sample.Status.ABORTED + sample.rollout_log_probs = [] + return sample + turns.append(_turn_record(final_sample, len(turns), "final", evaluation)) + + final_output = strip_stop_tokens(final_sample.response) + # Preserve a valid top-level Sample for logs and failure recovery. The + # custom converter trains from all turn records below, not this final view. + sample.prompt = final_prompt + sample.response = final_output + sample.tokens = list(final_sample.tokens) + sample.rollout_tokens = list(final_sample.rollout_tokens) + sample.response_length = final_sample.response_length + sample.loss_mask = list(final_sample.loss_mask or [1] * final_sample.response_length) + sample.rollout_log_probs = list(final_sample.rollout_log_probs or []) + sample.metadata.update( + { + "question": question, + "final_output": final_output, + "num_chunks": len(chunks), + "context_truncated": context_truncated, + "any_turn_truncated": any_turn_truncated or final_sample.status == Sample.Status.TRUNCATED, + "memory_token_lengths": memory_token_lengths, + "final_memory_tokens": memory_token_lengths[-1], + } + ) + sample.train_metadata = {"mem_agent_turns": turns} + sample.status = final_sample.status + return sample + + +async def generate( + args: Namespace, + sample: Sample, + sampling_params: dict[str, Any], + evaluation: bool = False, +) -> Sample: + """ReLax custom-generate entry point.""" + # Keep SGLang imports lazy so pure chunk/prompt tests can run in a + # lightweight CPU environment without initializing serving dependencies. + from relax.engine.rollout.sglang_rollout import GenerateState + from relax.engine.rollout.sglang_rollout import generate as sglang_generate + + tokenizer = GenerateState(args).tokenizer + try: + return await generate_trajectory( + args, + sample, + sampling_params, + tokenizer, + generator=sglang_generate, + evaluation=evaluation, + ) + except Exception as exc: + logger.error(f"MemAgent rollout failed for sample index={sample.index}: {exc}") + sample.response = "" + sample.rollout_log_probs = [] + sample.train_metadata = None + sample.status = Sample.Status.ABORTED + return sample diff --git a/examples/mem_agent/run-eval.sh b/examples/mem_agent/run-eval.sh new file mode 100755 index 000000000..83dcca0ea --- /dev/null +++ b/examples/mem_agent/run-eval.sh @@ -0,0 +1,70 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to the Hugging Face checkpoint to evaluate.}" +TOKENIZER_PATH="${TOKENIZER_PATH:-${MODEL_PATH}}" +DATA_DIR="${DATA_DIR:?Set DATA_DIR to the prepared MemAgent data directory.}" +RESULTS_DIR="${RESULTS_DIR:?Set RESULTS_DIR to the evaluation output directory.}" +RUN_NAME="${RUN_NAME:-$(basename "${MODEL_PATH}")}" +MODE="${MODE:-recurrent}" +LENGTHS="${LENGTHS:-50 200 800}" +SERVE_HOST="${SERVE_HOST:-127.0.0.1}" +SERVE_PORT="${SERVE_PORT:-8000}" +TP="${TP:-1}" +CONCURRENCY="${CONCURRENCY:-16}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-8192}" +GPU_MEMORY_UTIL="${GPU_MEMORY_UTIL:-0.85}" + +mkdir -p "${RESULTS_DIR}" +SERVER_LOG="${RESULTS_DIR}/${RUN_NAME}.server.log" + +vllm serve "${MODEL_PATH}" \ + --tensor-parallel-size "${TP}" \ + --host "${SERVE_HOST}" \ + --port "${SERVE_PORT}" \ + --max-model-len "${MAX_MODEL_LEN}" \ + --gpu-memory-utilization "${GPU_MEMORY_UTIL}" \ + --trust-remote-code \ + >"${SERVER_LOG}" 2>&1 & +SERVER_PID=$! +trap 'kill -TERM "${SERVER_PID}" 2>/dev/null || true; wait "${SERVER_PID}" 2>/dev/null || true' EXIT INT TERM + +for _ in $(seq 1 120); do + if ! kill -0 "${SERVER_PID}" 2>/dev/null; then + echo "vLLM server exited early; see ${SERVER_LOG}" >&2 + exit 1 + fi + if curl -fsS "http://${SERVE_HOST}:${SERVE_PORT}/v1/models" >/dev/null; then + break + fi + sleep 5 +done +curl -fsS "http://${SERVE_HOST}:${SERVE_PORT}/v1/models" >/dev/null + +run_eval() { + local data_file="$1" + local suffix="$2" + python3 "${SCRIPT_DIR}/eval_ruler_hqa.py" \ + --data-file "${data_file}" \ + --model "${MODEL_PATH}" \ + --tokenizer "${TOKENIZER_PATH}" \ + --output-dir "${RESULTS_DIR}" \ + --run-name "${RUN_NAME}-${suffix}" \ + --mode "${MODE}" \ + --base-url "http://${SERVE_HOST}:${SERVE_PORT}/v1" \ + --temperature 0.7 \ + --top-p 0.95 \ + --chunk-tokens 2048 \ + --max-memory-tokens 1024 \ + --max-final-tokens 256 \ + --max-chunks 64 \ + --concurrency "${CONCURRENCY}" +} + +run_eval "${DATA_DIR}/dev.jsonl" hotpotqa-dev +for length in ${LENGTHS}; do + run_eval "${DATA_DIR}/eval_${length}.jsonl" "ruler-hqa-${length}" +done diff --git a/examples/mem_agent/run-paired-eval.sh b/examples/mem_agent/run-paired-eval.sh new file mode 100755 index 000000000..2f4c14c23 --- /dev/null +++ b/examples/mem_agent/run-paired-eval.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +VIME_MODEL_PATH="${VIME_MODEL_PATH:?Set VIME_MODEL_PATH to the official or equivalently trained VIME checkpoint.}" +RELAX_MODEL_PATH="${RELAX_MODEL_PATH:?Set RELAX_MODEL_PATH to the converted ReLax checkpoint.}" +TOKENIZER_PATH="${TOKENIZER_PATH:?Set TOKENIZER_PATH to the frozen Qwen3-4B tokenizer snapshot.}" +DATA_DIR="${DATA_DIR:?Set DATA_DIR to the prepared frozen evaluation data.}" +RESULTS_DIR="${RESULTS_DIR:?Set RESULTS_DIR for raw and summary results.}" +LENGTHS="${LENGTHS:-50 200 800}" + +mkdir -p "${RESULTS_DIR}/vime" "${RESULTS_DIR}/relax" + +# Start and stop each checkpoint independently while preserving every other +# serving and evaluation variable. run-eval.sh writes both raw JSONL records +# and summary JSON, so the final comparison remains sample-auditable. +MODEL_PATH="${VIME_MODEL_PATH}" RUN_NAME=vime RESULTS_DIR="${RESULTS_DIR}/vime" \ + TOKENIZER_PATH="${TOKENIZER_PATH}" DATA_DIR="${DATA_DIR}" LENGTHS="${LENGTHS}" \ + bash "${SCRIPT_DIR}/run-eval.sh" +MODEL_PATH="${RELAX_MODEL_PATH}" RUN_NAME=relax RESULTS_DIR="${RESULTS_DIR}/relax" \ + TOKENIZER_PATH="${TOKENIZER_PATH}" DATA_DIR="${DATA_DIR}" LENGTHS="${LENGTHS}" \ + bash "${SCRIPT_DIR}/run-eval.sh" + +COMPARE_ARGS=() +for length in ${LENGTHS}; do + COMPARE_ARGS+=( + --pair "ruler-hqa-${length}" + "${RESULTS_DIR}/vime/vime-ruler-hqa-${length}.summary.json" + "${RESULTS_DIR}/relax/relax-ruler-hqa-${length}.summary.json" + ) +done +python3 "${SCRIPT_DIR}/compare_results.py" \ + "${COMPARE_ARGS[@]}" \ + --metric sub_em_pct \ + --tolerance-pp 3.0 \ + --output "${RESULTS_DIR}/vime-vs-relax.json" diff --git a/examples/mem_agent/run-pipeline.sh b/examples/mem_agent/run-pipeline.sh new file mode 100755 index 000000000..26f7caca1 --- /dev/null +++ b/examples/mem_agent/run-pipeline.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +DATA_DIR="${DATA_DIR:?Set DATA_DIR.}" +MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH.}" +SAVE_DIR="${SAVE_DIR:?Set SAVE_DIR.}" +RESULTS_DIR="${RESULTS_DIR:?Set RESULTS_DIR.}" +NUM_ROLLOUT="${NUM_ROLLOUT:-100}" +if ((NUM_ROLLOUT <= 0)); then + echo "NUM_ROLLOUT must be positive." >&2 + exit 1 +fi + +export DATA_DIR MODEL_PATH SAVE_DIR RESULTS_DIR NUM_ROLLOUT + +if [[ "${SKIP_PREPARE:-0}" != "1" ]]; then + bash "${SCRIPT_DIR}/prepare-data.sh" +fi +bash "${SCRIPT_DIR}/run-qwen3-4B-train.sh" + +export CHECKPOINT_DIR="${SAVE_DIR}" +# ReLax numbers checkpoints from zero, so a two-step pipeline produces +# iter_0000001 while the frozen 100-step recipe produces iter_0000099. +printf -v LAST_ROLLOUT_ID "%07d" "$((NUM_ROLLOUT - 1))" +export CHECKPOINT_TAG="${CHECKPOINT_TAG:-iter_${LAST_ROLLOUT_ID}}" +export HF_OUTPUT_DIR="${HF_OUTPUT_DIR:-${SAVE_DIR}-HF/${CHECKPOINT_TAG}}" +export TOKENIZER_PATH="${TOKENIZER_PATH:-${MODEL_PATH}}" +bash "${SCRIPT_DIR}/convert-to-hf.sh" + +export MODEL_PATH="${HF_OUTPUT_DIR}" +bash "${SCRIPT_DIR}/run-eval.sh" diff --git a/examples/mem_agent/run-qwen3-4B-train.sh b/examples/mem_agent/run-qwen3-4B-train.sh new file mode 100755 index 000000000..b1faab058 --- /dev/null +++ b/examples/mem_agent/run-qwen3-4B-train.sh @@ -0,0 +1,131 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail +set -x + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +RELAX_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" + +if [[ -z "${RELAX_ENTRYPOINT_MODE:-}" ]]; then + source "${RELAX_ROOT}/scripts/entrypoint/local.sh" +fi +source "${MODEL_CONFIG_DIR}/qwen3-4B.sh" + +MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to the frozen Qwen3-4B snapshot directory.}" +DATA_DIR="${DATA_DIR:?Set DATA_DIR to the prepared MemAgent data directory.}" +SAVE_DIR="${SAVE_DIR:?Set SAVE_DIR to the checkpoint output directory.}" +TRAIN_DATA="${TRAIN_DATA:-${DATA_DIR}/train.jsonl}" +NUM_ROLLOUT="${NUM_ROLLOUT:-100}" +SAVE_INTERVAL="${SAVE_INTERVAL:-50}" +ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-8}" +N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" +GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-64}" +RUN_NAME="${RUN_NAME:-mem-agent-qwen3-4b}" + +[[ -f "${TRAIN_DATA}" ]] || { echo "Missing training data: ${TRAIN_DATA}" >&2; exit 1; } +[[ -f "${MODEL_PATH}/config.json" ]] || { echo "Missing model config: ${MODEL_PATH}/config.json" >&2; exit 1; } +mkdir -p "${SAVE_DIR}" "${RELAX_ROOT}/logs" + +CKPT_ARGS=( + --hf-checkpoint "${MODEL_PATH}" + --ref-load "${MODEL_PATH}" + --megatron-to-hf-mode bridge + --save "${SAVE_DIR}" + --save-interval "${SAVE_INTERVAL}" + --max-actor-ckpt-to-keep 3 +) + +ROLLOUT_ARGS=( + --prompt-data "${TRAIN_DATA}" + --input-key prompt + --label-key label + --metadata-key metadata + --custom-generate-function-path examples.mem_agent.rollout.generate + --custom-rm-path examples.mem_agent.reward.reward_func + --custom-convert-samples-to-train-data-path examples.mem_agent.convert.convert_samples + --custom-config-path "${SCRIPT_DIR}/config.yaml" + --reward-key score + --num-rollout "${NUM_ROLLOUT}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" + --rollout-max-response-len 1024 + # Each request is an independent <=2K chunk + <=1K memory turn. An 8K + # engine limit covers the real request envelope while retaining VIME's 9K + # per-GPU packing budget and sample-mean loss semantics. + --rollout-max-context-len 8192 + --rollout-temperature 1.0 + --rollout-top-p 1.0 + --rollout-seed 42 + --rollout-shuffle + --global-batch-size "${GLOBAL_BATCH_SIZE}" + --balance-data +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.001 + --kl-loss-type low_var_kl + --entropy-coef 0.0 + --eps-clip 0.2 + --eps-clip-high 0.3 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 + --log-probs-max-tokens-per-gpu 32768 +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 2 + --sglang-mem-fraction-static 0.7 +) + +MISC_ARGS=( + --seed 1234 + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --skip-eval-before-train + --max-staleness 0 + --num-data-storage-units 1 + --colocate + --use-health-check +) + +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://127.0.0.1:8265" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"actor": [1, 8], "rollout": [1, 8]}' \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" \ + 2>&1 | tee "${RELAX_ROOT}/logs/${RUN_NAME}.log" diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 1c1355426..f84139f86 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -627,14 +627,28 @@ def train(self, rollout_id: int) -> None: if is_sft_mode(self.args): batch_size = self.args.global_batch_size // dp_size rollout_mini_local_sample_counts = None + elif getattr(self.args, "custom_train_expanded_batch", False): + batch_size = None + rollout_mini_local_sample_counts = None else: plan = build_rollout_minibatch_plan(self.args, dp_size) batch_size = plan.mini_local_sample_request * plan.num_rollout_minis rollout_mini_local_sample_counts = [ plan.mini_local_sample_request for _ in range(plan.num_rollout_minis) ] - rollout_data = get_debug_data(self.args, rollout_id, batch_size, dp_rank=mpu.get_data_parallel_rank()) + rollout_data = get_debug_data( + self.args, + rollout_id, + batch_size, + dp_rank=mpu.get_data_parallel_rank(), + dp_size=dp_size, + ) post_process_rollout_data(self.args, rollout_data) + if getattr(self.args, "custom_train_expanded_batch", False): + # Debug dumps must preserve the same one-rollout boundary as + # online expanded batches; otherwise get_data_iterator may + # reinterpret turn rows as several fixed global batches. + rollout_mini_local_sample_counts = [len(rollout_data["total_lengths"])] if rollout_mini_local_sample_counts is not None: if sum(rollout_mini_local_sample_counts) != len(rollout_data["total_lengths"]): raise RuntimeError( @@ -657,7 +671,17 @@ def train(self, rollout_id: int) -> None: num_rollout_minis = 1 else: dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) - if self.args.partial_rollout and self.args.use_dynamic_global_batch_size: + if getattr(self.args, "custom_train_expanded_batch", False): + # The rollout-ID keyed getter waits until this exact + # partition has been fully converted and transferred. + expanded_rows = ray.get(self.rollout_manager.get_train_batch_row_count.remote(rollout_id)) + if expanded_rows % dp_size: + raise RuntimeError( + f"Expanded train row count {expanded_rows} must be divisible by dp_size={dp_size}." + ) + batch_size = expanded_rows // dp_size + num_rollout_minis = 1 + elif self.args.partial_rollout and self.args.use_dynamic_global_batch_size: dynamic_size = ray.get(self.rollout_manager.get_dynamic_global_batch_size.remote()) batch_size = dynamic_size // dp_size num_rollout_minis = 1 @@ -1247,7 +1271,13 @@ def train_hybrid(self, rollout_id) -> None: # rollout mini windows. logger.info(f"start to get rollout_id: {rollout_id} data from debug rollout data for train_hybrid.") full_batch_size = plan.mini_local_sample_request * plan.num_rollout_minis - debug_data = get_debug_data(self.args, rollout_id, full_batch_size, dp_rank=mpu.get_data_parallel_rank()) + debug_data = get_debug_data( + self.args, + rollout_id, + full_batch_size, + dp_rank=mpu.get_data_parallel_rank(), + dp_size=dp_size, + ) post_process_rollout_data(self.args, debug_data) for sub_batch in self._split_rollout_batch(debug_data, plan.num_rollout_minis): if len(sub_batch["total_lengths"]) != batch_size: diff --git a/relax/core/controller.py b/relax/core/controller.py index 96bd1f72b..ccdf157d4 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -143,22 +143,33 @@ def __init__(self, config: Namespace, runtime_env: dict = None) -> None: logger.info("Global health check system disabled (use --use-health-check to enable)") def _initialize_data_system(self): + from relax.utils.utils import get_train_data_group_size, get_train_sample_expansion_factor + algo_key = resolve_sft_algo_key(self.config) batch_size_for_capacity = ( self.config.over_sampling_batch_size if self.config.partial_rollout and self.config.use_dynamic_global_batch_size else self.config.rollout_batch_size ) + # A custom converter can expand one trajectory into multiple training + # rows. This declared factor is only a storage upper bound; the actual + # row count is measured after conversion by RolloutManager. total_storage_size = ( - batch_size_for_capacity * (self.config.max_staleness + 1) * self.config.n_samples_per_prompt + batch_size_for_capacity + * (self.config.max_staleness + 1) + * self.config.n_samples_per_prompt + * get_train_sample_expansion_factor(self.config) ) + # MemAgent completes GRPO normalization before turn expansion, so its + # converted rows opt into sampler group size 1. + train_data_group_size = get_train_data_group_size(self.config) if getattr(self.config, "fully_async", False) and getattr(self.config, "use_dynamic_batch_size", False): # Fully-async + dynamic-batch path streams data per DP via token # budget; the controller-side sampler maintains per-DP buckets and # balances tokens at small-unit granularity. See # docs/zh/guide/fully-async-training.md. sampler = StreamingTokenBudgetSampler( - n_samples_per_prompt=self.config.n_samples_per_prompt, + n_samples_per_prompt=train_data_group_size, ) logger.info("Using StreamingTokenBudgetSampler (fully_async + dynamic batch)") elif algo_key == "sft" or getattr(self.config, "balance_data", False): @@ -166,12 +177,12 @@ def _initialize_data_system(self): # since the GRPO grouped sampler assumes n_samples_per_prompt > 1 rollouts. dp_size = compute_dp_size(self.config) sampler = SeqlenBalancedSampler( - n_samples_per_prompt=self.config.n_samples_per_prompt, + n_samples_per_prompt=train_data_group_size, dp_size=dp_size, ) logger.info(f"Using SeqlenBalancedSampler with dp_size={dp_size}") else: - sampler = GRPOGroupNSampler(n_samples_per_prompt=self.config.n_samples_per_prompt) + sampler = GRPOGroupNSampler(n_samples_per_prompt=train_data_group_size) tq_config = OmegaConf.create( { diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index 128fa5f4e..2eaea9dc7 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -794,6 +794,12 @@ def __init__(self, args, pg, data_source=None): self.pg = pg self.args = args self._dynamic_global_batch_size = None + # Expanded row counts are keyed by rollout ID because rollout and actor + # scheduling can overlap. A single "latest" value lets step N+1 reuse + # step N's count and can leave TransferQueue rows unread. + self._train_batch_row_counts: dict[int, int] = {} + self._train_batch_row_errors: dict[int, str] = {} + self._train_batch_row_events: dict[int, asyncio.Event] = {} init_tracking(args, primary=False) @@ -809,12 +815,6 @@ def __init__(self, args, pg, data_source=None): self.custom_reward_post_process_func = None if self.args.custom_reward_post_process_path is not None: self.custom_reward_post_process_func = load_function(self.args.custom_reward_post_process_path) - self.custom_convert_samples_to_train_data_func = None - if self.args.custom_convert_samples_to_train_data_path is not None: - self.custom_convert_samples_to_train_data_func = load_function( - self.args.custom_convert_samples_to_train_data_path - ) - if self.args.use_agentic_rollout: from relax.agentic.rollout import init_agentic_resident_pipeline @@ -999,6 +999,21 @@ def get_dynamic_global_batch_size(self): ) return self._dynamic_global_batch_size + async def get_train_batch_row_count(self, rollout_id: int) -> int: + """Wait for and return the exact rows transferred for one rollout. + + The getter may be scheduled before ``generate`` starts. Reusing the + same event object lets both call orders converge without polling or + falling back to a stale count from the previous step. + """ + event = self._train_batch_row_events.setdefault(rollout_id, asyncio.Event()) + await event.wait() + if error := self._train_batch_row_errors.get(rollout_id): + raise RuntimeError(f"Failed to determine train rows for rollout_id={rollout_id}: {error}") + if rollout_id not in self._train_batch_row_counts: + raise RuntimeError(f"Train row count is unavailable for rollout_id={rollout_id}.") + return self._train_batch_row_counts[rollout_id] + def get_num_rollout_per_epoch(self): assert self.args.rollout_global_dataset return ray.get(self.data_source.lengths.remote()) // self.args.rollout_batch_size @@ -1006,19 +1021,51 @@ def get_num_rollout_per_epoch(self): async def generate(self, rollout_id): self.rollout_id = rollout_id self.health_monitoring_resume() - if self.args.ci_test and self.args.use_fault_tolerance and rollout_id >= 2: - self._try_ci_fault_injection() - output = await asyncio.to_thread( - call_rollout_fn, - self.generate_rollout, - self.args, - rollout_id, - self.data_source, - self.data_system_client, - evaluation=False, - ) - if self.args.partial_rollout and self.args.use_dynamic_global_batch_size: - self._dynamic_global_batch_size = len(output.samples) * self.args.n_samples_per_prompt + row_event = self._train_batch_row_events.setdefault(rollout_id, asyncio.Event()) + row_event.clear() + self._train_batch_row_counts.pop(rollout_id, None) + self._train_batch_row_errors.pop(rollout_id, None) + # Counts are only coordination state; retain a short overlap window for + # late actor/critic consumers and bound the actor's memory usage. + for stale_rollout_id in [key for key in self._train_batch_row_events if key < rollout_id - 2]: + self._train_batch_row_events.pop(stale_rollout_id, None) + self._train_batch_row_counts.pop(stale_rollout_id, None) + self._train_batch_row_errors.pop(stale_rollout_id, None) + try: + if self.args.ci_test and self.args.use_fault_tolerance and rollout_id >= 2: + self._try_ci_fault_injection() + output = await asyncio.to_thread( + call_rollout_fn, + self.generate_rollout, + self.args, + rollout_id, + self.data_source, + self.data_system_client, + evaluation=False, + ) + if getattr(self.args, "custom_train_expanded_batch", False): + # The transfer helper reports the exact rows it actually put. + # Re-converting output.samples here is unsafe because custom + # converters may filter rows or be stateful. + metrics = output.metrics or {} + row_count = metrics.get("rollout/train_batch_row_count") + if row_count is None: + raise RuntimeError("Expanded rollout did not report its transferred train row count.") + row_count = int(row_count) + if row_count <= 0: + raise RuntimeError(f"Expanded rollout produced an invalid train row count: {row_count}.") + # This count controls TQ consumption only. global_batch_size + # remains trajectory-based for GRPO loss normalization. + self._train_batch_row_counts[rollout_id] = row_count + elif self.args.partial_rollout and self.args.use_dynamic_global_batch_size: + self._dynamic_global_batch_size = len(output.samples) * self.args.n_samples_per_prompt + except Exception as exc: + if getattr(self.args, "custom_train_expanded_batch", False): + self._train_batch_row_errors[rollout_id] = f"{type(exc).__name__}: {exc}" + raise + finally: + if getattr(self.args, "custom_train_expanded_batch", False): + row_event.set() async def eval(self, rollout_id): self.health_monitoring_resume() diff --git a/relax/engine/rollout/sglang_rollout.py b/relax/engine/rollout/sglang_rollout.py index e578f3d25..3059f0fa3 100644 --- a/relax/engine/rollout/sglang_rollout.py +++ b/relax/engine/rollout/sglang_rollout.py @@ -731,6 +731,10 @@ async def generate_rollout_async( do_print = True pbar = tqdm(total=target_data_size * args.n_samples_per_prompt, desc=f"Rollout {rollout_id} generation") transfer_tasks = [] + # Producer-side conversion is the source of truth for the number of train + # rows. MemAgent expands one trajectory into a variable number of turns, so + # the actor cannot derive this count from rollout_batch_size. + transferred_train_rows = 0 batch_to_transfer = [] aborted_samples = [] # Completed groups beyond target_data_size (over-sampling surplus). Carried back to @@ -947,7 +951,10 @@ def target_reached() -> bool: logger.info(f"Generator exhausted. Waiting for {len(transfer_tasks)} transfer tasks to complete...") # Wait for all transfer tasks to complete if transfer_tasks: - await asyncio.gather(*transfer_tasks) + transfer_results = await asyncio.gather(*transfer_tasks) + transferred_train_rows += sum( + row_count for destination_rollout_id, row_count in transfer_results if destination_rollout_id == rollout_id + ) pbar.close() # Stop SGLang profiling if enabled (no-op if num_steps was set — SGLang auto-stops) @@ -1020,7 +1027,10 @@ def target_reached() -> bool: for group in accepted: data.append(group) if accepted: - await transfer_batch_to_data_system(args, accepted, len(accepted), rollout_id, data_system_client) + _, accepted_rows = await transfer_batch_to_data_system( + args, accepted, len(accepted), rollout_id, data_system_client + ) + transferred_train_rows += accepted_rows logger.info(f"Transferred {len(accepted)} extra completed groups to training ") global CURRENT_ROLLOUT_BATCH @@ -1049,7 +1059,13 @@ def target_reached() -> bool: state.reset() - return RolloutFnTrainOutput(samples=data, metrics=metric_gatherer.collect()), aborted_samples + metrics = metric_gatherer.collect() + if getattr(args, "custom_train_expanded_batch", False): + # This private coordination metric is consumed by RolloutManager. It is + # emitted only for expanded converters so ordinary rollout metrics stay + # byte-for-byte compatible. + metrics["rollout/train_batch_row_count"] = transferred_train_rows + return RolloutFnTrainOutput(samples=data, metrics=metrics), aborted_samples EVAL_PROMPT_DATASET = {} diff --git a/relax/utils/utils.py b/relax/utils/utils.py index 87cfaf594..e3ad8ab7e 100644 --- a/relax/utils/utils.py +++ b/relax/utils/utils.py @@ -21,6 +21,23 @@ CURRENT_ROLLOUT_BATCH = [] +def get_train_sample_expansion_factor(args: Any) -> int: + """Return the declared upper bound for custom converter row expansion.""" + factor = int(getattr(args, "custom_train_sample_expansion_factor", 1)) + if factor <= 0: + raise ValueError("custom_train_sample_expansion_factor must be positive.") + return factor + + +def get_train_data_group_size(args: Any) -> int: + """Return TransferQueue's grouping unit for converted training rows.""" + default_group_size = getattr(args, "n_samples_per_prompt", 1) + group_size = int(getattr(args, "custom_train_data_group_size", default_group_size)) + if group_size <= 0: + raise ValueError("custom_train_data_group_size must be positive.") + return group_size + + def _extract_images_seqlens(multimodal_train_inputs) -> list[int]: """Extract per-image ViT token counts from multimodal_train_inputs. @@ -166,6 +183,15 @@ def convert_samples_to_train_data(args: Any, samples: list[Sample] | list[list[S return rollout_batch +def convert_samples_to_train_data_with_custom(args: Any, samples: list[Sample] | list[list[Sample]]): + """Use the configured sample converter while preserving the default + path.""" + converter = convert_samples_to_train_data + if custom_path := getattr(args, "custom_convert_samples_to_train_data_path", None): + converter = load_function(custom_path) + return converter(args, samples) + + def post_process_rewards(args: Any, samples: list[Sample] | list[list[Sample]]): """Post-process rewards and return (raw_rewards, possibly-normalized rewards). @@ -326,9 +352,10 @@ def post_process_env(args, env): if "env_vars" not in env or not isinstance(env["env_vars"], dict): env["env_vars"] = {} - # Dynamic-batch streaming ends via the producer's is_last signal, not a - # pre-allocated partition, so pre-allocate the minimum (1) and let it grow. - # The non-dynamic path still pre-allocates the exact count for its .all() check. + # TQ pre-allocation is a trajectory-level sizing hint, not the final row + # count. A custom converter may expand each trajectory by a data-dependent + # amount, so reserving the declared maximum would leave an unread tail. The + # controller separately raises the storage capacity for expanded rows. if getattr(args, "fully_async", False) and getattr(args, "use_dynamic_batch_size", False): env["env_vars"]["TQ_PRE_ALLOC_SAMPLE_NUM"] = str( args.rollout_batch_size * args.n_samples_per_prompt @@ -411,7 +438,13 @@ def merge_dict_list(dict_list): return merged -def get_debug_data(args, rollout_id: int, batch_size, dp_rank: int) -> Dict[str, Any]: +def get_debug_data( + args, + rollout_id: int, + batch_size: int | None, + dp_rank: int, + dp_size: int | None = None, +) -> Dict[str, Any]: """Fetch debug data for a given rollout_id from the data system. Parameters: @@ -446,7 +479,23 @@ def get_debug_data(args, rollout_id: int, batch_size, dp_rank: int) -> Dict[str, logger.info( f"Subsample loaded debug rollout data using {ratio=} and change num rows {original_num_rows} -> {len(data)}" ) - rollout_batch = convert_samples_to_train_data(args, data) + rollout_batch = convert_samples_to_train_data_with_custom(args, data) + + # Custom converters can expand the debug dump. The partial-rollout dynamic + # path has the same requirement: slice converted rows, not trajectories. + custom_expanded_batch = getattr(args, "custom_train_expanded_batch", False) + partial_dynamic_batch = getattr(args, "partial_rollout", False) and getattr( + args, "use_dynamic_global_batch_size", False + ) + if custom_expanded_batch or partial_dynamic_batch: + if dp_size is None or dp_size <= 0: + raise ValueError("dp_size is required for dynamic debug rollout replay.") + converted_size = len(rollout_batch["total_lengths"]) + if converted_size % dp_size: + raise ValueError(f"Converted debug row count {converted_size} must be divisible by dp_size={dp_size}.") + batch_size = converted_size // dp_size + if batch_size is None: + raise ValueError("batch_size is required for debug rollout replay.") for key in rollout_batch: rollout_batch[key] = rollout_batch[key][dp_rank * batch_size : (dp_rank + 1) * batch_size] @@ -460,7 +509,7 @@ async def transfer_batch_to_data_system( rollout_id: int, data_system_client: Any, is_last: bool = False, -) -> None: +) -> tuple[int, int]: """Helper function to transfer a batch of samples to the data system client. @@ -472,6 +521,11 @@ async def transfer_batch_to_data_system( is_last: Mark this as the final batch of the partition train_{rollout_id} so the data system can detect streaming end-of-stream without a preset global batch size. See the is_last bookkeeping in generate_rollout. + + Returns: + The destination rollout ID and the exact number of converted rows put + into its partition. The rollout manager uses this producer-side count + instead of running a second, potentially divergent conversion. """ try: # Guard against empty batch_samples @@ -479,7 +533,7 @@ async def transfer_batch_to_data_system( logger.warning( f"transfer_batch_to_data_system called with empty batch_samples for rollout_id={rollout_id}, batch_count={batch_count}" ) - return + return rollout_id, 0 batch_samples = sorted( batch_samples, key=lambda group: group[0][0].index if isinstance(group[0], list) else group[0].index ) @@ -488,7 +542,7 @@ async def transfer_batch_to_data_system( batch_samples = sum(batch_samples, []) global CURRENT_ROLLOUT_BATCH CURRENT_ROLLOUT_BATCH.extend(batch_samples) - rollout_batch = convert_samples_to_train_data(args, batch_samples) + rollout_batch = convert_samples_to_train_data_with_custom(args, batch_samples) logger.info(f"Prepared rollout batch {batch_count} with {rollout_batch.numel()} samples for transfer") logger.info(f"Transferring batch rollout_batch: {rollout_batch}") @@ -505,6 +559,7 @@ async def transfer_batch_to_data_system( ) logger.info(f"Batch {batch_count} transferred successfully for rollout_id: {rollout_id}") + return rollout_id, len(rollout_batch["total_lengths"]) except Exception as e: logger.error(f"Error transferring batch {batch_count}: {e}") raise diff --git a/tests/examples/mem_agent/__init__.py b/tests/examples/mem_agent/__init__.py new file mode 100644 index 000000000..9f3863608 --- /dev/null +++ b/tests/examples/mem_agent/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. diff --git a/tests/examples/mem_agent/test_compare_results.py b/tests/examples/mem_agent/test_compare_results.py new file mode 100644 index 000000000..65b38584b --- /dev/null +++ b/tests/examples/mem_agent/test_compare_results.py @@ -0,0 +1,14 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import pytest + +from examples.mem_agent.compare_results import compare_pair + + +def test_compare_pair_uses_absolute_percentage_points(): + passed = compare_pair("50", {"sub_em_pct": 41.0}, {"sub_em_pct": 43.9}) + failed = compare_pair("50", {"sub_em_pct": 41.0}, {"sub_em_pct": 44.1}) + + assert passed["absolute_gap_pp"] == pytest.approx(2.9) + assert passed["passed"] is True + assert failed["passed"] is False diff --git a/tests/examples/mem_agent/test_convert.py b/tests/examples/mem_agent/test_convert.py new file mode 100644 index 000000000..7251c6554 --- /dev/null +++ b/tests/examples/mem_agent/test_convert.py @@ -0,0 +1,91 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from examples.mem_agent.convert import convert_samples +from relax.utils.types import Sample + + +def _turn(turn_index: int) -> dict: + return { + "turn_index": turn_index, + "kind": "final" if turn_index == 2 else "memory", + "tokens": [100 + turn_index, 200 + turn_index, 201 + turn_index], + "response_length": 2, + "loss_mask": [1, 1], + "rollout_log_probs": [-0.1, -0.2], + "finish_reason": Sample.Status.COMPLETED.value, + } + + +def _sample(index: int, reward: float, turn_count: int) -> Sample: + return Sample( + index=index, + group_index=9, + reward={"score": reward}, + status=Sample.Status.COMPLETED, + train_metadata={"mem_agent_turns": [_turn(turn_index) for turn_index in range(turn_count)]}, + ) + + +def _args(credit_assignment="split", debug_train_only=True): + return SimpleNamespace( + reward_key="score", + custom_reward_post_process_path=None, + agentic_custom_advantage_path=None, + advantage_estimator="grpo", + rewards_normalization=True, + n_samples_per_prompt=2, + grpo_std_normalization=True, + mem_agent_credit_assignment=credit_assignment, + debug_train_only=debug_train_only, + ) + + +def test_converter_normalizes_before_expansion_and_keeps_all_turns(): + samples = [_sample(3, 0.0, 3), _sample(4, 1.0, 3)] + data = convert_samples(_args(), samples) + normalized = torch.tensor([0.0, 1.0]) + normalized = (normalized - normalized.mean()) / (normalized.std() + 1e-6) + + assert len(data["tokens"]) == 6 + assert data["sample_indices"] == [3, 3, 3, 4, 4, 4] + assert data["turn_indices"] == [0, 1, 2, 0, 1, 2] + assert data["rewards"][:3] == pytest.approx([normalized[0].item() / 3] * 3) + assert data["rewards"][3:] == pytest.approx([normalized[1].item() / 3] * 3) + assert data["raw_reward"] == [0.0, 0.0, 0.0, 1.0, 1.0, 1.0] + + +def test_converter_share_credit_and_tensordict_contract(): + samples = [_sample(3, 0.0, 2), _sample(4, 1.0, 2)] + data = convert_samples(_args(credit_assignment="share", debug_train_only=False), samples) + assert len(data["total_lengths"]) == 4 + assert data.batch_size[0] == 4 + assert data["rewards"][0].item() == pytest.approx(data["rewards"][1].item()) + + +def test_converter_rejects_misaligned_or_failed_trajectory(): + sample = _sample(3, 1.0, 1) + sample.train_metadata["mem_agent_turns"][0]["rollout_log_probs"] = [-0.1] + with pytest.raises(ValueError, match="misaligned rollout_log_probs"): + convert_samples(_args(), [sample, _sample(4, 0.0, 1)]) + + failed = _sample(5, 0.0, 1) + failed.status = Sample.Status.FAILED + with pytest.raises(ValueError, match="status=failed"): + convert_samples(_args(), [failed, _sample(6, 1.0, 1)]) + + +def test_converter_rejects_inconsistent_group_turn_counts_and_non_divisible_rows(): + with pytest.raises(ValueError, match="inconsistent turn counts"): + convert_samples(_args(), [_sample(3, 0.0, 1), _sample(4, 1.0, 2)]) + + args = _args() + args.mem_agent_train_rows_multiple = 4 + with pytest.raises(ValueError, match="is not divisible"): + convert_samples(args, [_sample(3, 0.0, 3), _sample(4, 1.0, 3)]) diff --git a/tests/examples/mem_agent/test_data_and_metrics.py b/tests/examples/mem_agent/test_data_and_metrics.py new file mode 100644 index 000000000..f0090a56f --- /dev/null +++ b/tests/examples/mem_agent/test_data_and_metrics.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import json + +import pytest + +from examples.mem_agent.metrics import aggregate, exact_match, f1_score, sub_exact_match +from examples.mem_agent.prepare_data import convert_file, convert_row, update_manifest + + +def test_convert_row_supports_training_and_ruler_formats(): + training = convert_row( + { + "prompt": [{"role": "user", "content": "Who?"}], + "context": "Document text", + "reward_model": {"ground_truth": ["Alice", "A. Alice"]}, + "extra_info": {"num_docs": 200}, + } + ) + ruler = convert_row({"input": "Where?", "answers": ["Paris"], "context": "Context", "num_docs": 50}) + + assert training["prompt"] == "Who?" + assert training["metadata"]["ground_truth"] == ["Alice", "A. Alice"] + assert training["metadata"]["num_docs"] == 200 + assert ruler["label"] == "Paris" + assert ruler["metadata"]["question"] == "Where?" + + +def test_convert_file_and_manifest_are_deterministic(tmp_path): + source = tmp_path / "eval.json" + source.write_text( + json.dumps([{"input": "Where?", "answers": ["Paris"], "context": "Context", "num_docs": 50}]), + encoding="utf-8", + ) + output = tmp_path / "eval.jsonl" + manifest = tmp_path / "artifact_manifest.json" + + assert convert_file(source, output) == (1, 0) + update_manifest(manifest, source, source.name, "dataset/id", "revision", output, written=1, skipped=0) + payload = json.loads(manifest.read_text(encoding="utf-8")) + assert payload["artifacts"][0]["filename"] == "eval.json" + assert len(payload["artifacts"][0]["sha256"]) == 64 + assert payload["artifacts"][0]["converted"]["written_rows"] == 1 + assert len(payload["artifacts"][0]["converted"]["sha256"]) == 64 + assert json.loads(output.read_text(encoding="utf-8"))["metadata"]["ground_truth"] == ["Paris"] + + +def test_ruler_metrics_match_expected_semantics(): + assert exact_match("The Eiffel Tower", "Eiffel Tower") == 1.0 + assert sub_exact_match("located in Paris France", "Paris") == 1.0 + assert f1_score("Paris France", "Paris") == 2 / 3 + summary = aggregate( + [ + {"judge_f1": 1.0, "judge_em": 1.0, "judge_sub_em": 1.0, "judge_boxed_em": 1.0}, + {"judge_f1": 0.0, "judge_em": 0.0, "judge_sub_em": 0.0, "judge_boxed_em": 0.0}, + {"error": "boom"}, + ] + ) + assert summary["total"] == 3 + assert summary["successful"] == 2 + assert summary["errors"] == 1 + assert summary["sub_em_pct"] == pytest.approx(100 / 3) + assert summary["boxed_em_pct"] == pytest.approx(100 / 3) diff --git a/tests/examples/mem_agent/test_eval.py b/tests/examples/mem_agent/test_eval.py new file mode 100644 index 000000000..aff2f3595 --- /dev/null +++ b/tests/examples/mem_agent/test_eval.py @@ -0,0 +1,47 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from examples.mem_agent.eval_ruler_hqa import base_infer + + +class CharacterTokenizer: + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [ord(character) for character in text] + + def decode(self, token_ids, skip_special_tokens=True): + del skip_special_tokens + return "".join(chr(token_id) for token_id in token_ids) + + +@pytest.mark.asyncio +async def test_base_infer_truncates_context_without_dropping_question(monkeypatch): + captured = {} + + async def fake_chat_once(session, base_url, api_key, model, instruction, temperature, top_p, max_tokens): + del session, base_url, api_key, model, temperature, top_p, max_tokens + captured["instruction"] = instruction + return r"\boxed{x}" + + monkeypatch.setattr("examples.mem_agent.eval_ruler_hqa._chat_once", fake_chat_once) + args = SimpleNamespace( + max_input_tokens=100, + base_url="http://unused", + api_key="EMPTY", + model="model", + temperature=0.0, + top_p=1.0, + max_final_tokens=16, + ) + _, diagnostics = await base_infer( + {"context": "c" * 500, "input": "Which answer?"}, args, CharacterTokenizer(), object() + ) + + assert diagnostics["context_truncated"] is True + assert "Question: Which answer?" in captured["instruction"] + assert len(captured["instruction"]) <= args.max_input_tokens diff --git a/tests/examples/mem_agent/test_mock_integration.py b/tests/examples/mem_agent/test_mock_integration.py new file mode 100644 index 000000000..eaf67ab51 --- /dev/null +++ b/tests/examples/mem_agent/test_mock_integration.py @@ -0,0 +1,109 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""CPU-only integration of rollout, reward, expansion, and queue transfer.""" + +from __future__ import annotations + +from collections import defaultdict +from types import SimpleNamespace + +import pytest + +from examples.mem_agent.reward import reward_func +from examples.mem_agent.rollout import generate_trajectory +from relax.utils.types import Sample +from relax.utils.utils import CURRENT_ROLLOUT_BATCH, transfer_batch_to_data_system + + +class FakeTokenizer: + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [ord(character) for character in text] + + def decode(self, token_ids, skip_special_tokens=True): + del skip_special_tokens + return "".join(chr(token_id) for token_id in token_ids) + + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True): + assert not tokenize and add_generation_prompt + return f"{messages[0]['content']}" + + +def _args(): + return SimpleNamespace( + mem_agent_chunk_tokens=3, + mem_agent_max_memory_tokens=16, + mem_agent_max_final_tokens=16, + mem_agent_max_chunks=8, + mem_agent_credit_assignment="split", + custom_convert_samples_to_train_data_path="examples.mem_agent.convert.convert_samples", + custom_reward_post_process_path=None, + agentic_custom_advantage_path=None, + reward_key="score", + advantage_estimator="grpo", + rewards_normalization=True, + grpo_std_normalization=True, + n_samples_per_prompt=2, + debug_train_only=False, + ) + + +@pytest.mark.asyncio +async def test_mock_mem_agent_pipeline_transfers_every_turn_once(): + tokenizer = FakeTokenizer() + turn_positions = defaultdict(int) + observed_prompts = defaultdict(list) + + async def fake_generate(args, turn, sampling_params, evaluation): + del args, sampling_params, evaluation + position = turn_positions[turn.index] + turn_positions[turn.index] += 1 + observed_prompts[turn.index].append(turn.prompt) + if position < 2: + response = f"M{position + 1}" + else: + response = r"\boxed{x}" if turn.index == 0 else r"\boxed{wrong}" + response_ids = tokenizer.encode(response) + turn.response = response + turn.tokens = tokenizer.encode(turn.prompt) + response_ids + turn.rollout_tokens = list(turn.tokens) + turn.response_length = len(response_ids) + turn.loss_mask = [1] * len(response_ids) + turn.rollout_log_probs = [-0.2] * len(response_ids) + turn.status = Sample.Status.COMPLETED + return turn + + args = _args() + samples = [] + for index in range(2): + sample = Sample( + index=index, + group_index=5, + prompt="Question?", + metadata={"context": "abcdef", "ground_truth": ["x"]}, + ) + sample = await generate_trajectory(args, sample, {}, tokenizer, generator=fake_generate) + sample.reward = await reward_func(args, sample) + samples.append(sample) + + class Client: + payload = None + + async def async_put(self, **kwargs): + self.payload = kwargs + + client = Client() + CURRENT_ROLLOUT_BATCH.clear() + destination_rollout_id, transferred_rows = await transfer_batch_to_data_system( + args, [samples], 1, 0, client, is_last=True + ) + train_data = client.payload["data"] + + assert destination_rollout_id == 0 + assert transferred_rows == 6 + assert len(train_data["tokens"]) == 6 + assert train_data["sample_indices"].tolist() == [0, 0, 0, 1, 1, 1] + assert train_data["turn_indices"].tolist() == [0, 1, 2, 0, 1, 2] + assert train_data["raw_reward"].tolist() == [1.0, 1.0, 1.0, 0.0, 0.0, 0.0] + assert len(client.payload["custom_meta"]) == 6 + assert all("
" not in prompts[-1] for prompts in observed_prompts.values()) + assert all("abc" not in prompts[-1] and "def" not in prompts[-1] for prompts in observed_prompts.values()) diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py new file mode 100644 index 000000000..33cfbccf6 --- /dev/null +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -0,0 +1,38 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Static checks for the frozen MemAgent launch and reproducibility recipe.""" + +from __future__ import annotations + +from pathlib import Path + +import yaml + + +ROOT = Path(__file__).resolve().parents[3] +EXAMPLE = ROOT / "examples" / "mem_agent" + + +def test_custom_config_freezes_model_memory_and_expanded_batch_contract(): + config = yaml.safe_load((EXAMPLE / "config.yaml").read_text(encoding="utf-8")) + + assert config["model_id"] == "Qwen/Qwen3-4B" + assert config["model_revision"] == "1cfa9a7208912126459214e8b04321603b3df60c" + assert config["mem_agent_chunk_tokens"] == 2048 + assert config["mem_agent_max_memory_tokens"] == 1024 + assert config["mem_agent_max_final_tokens"] == 256 + assert config["mem_agent_max_chunks"] == 64 + assert config["mem_agent_credit_assignment"] == "split" + assert config["custom_train_sample_expansion_factor"] == 65 + assert config["custom_train_data_group_size"] == 1 + assert config["custom_train_expanded_batch"] is True + + +def test_train_script_keeps_trajectory_loss_and_real_turn_context_envelope(): + script = (EXAMPLE / "run-qwen3-4B-train.sh").read_text(encoding="utf-8") + + assert '--global-batch-size "${GLOBAL_BATCH_SIZE}"' in script + assert "--rollout-max-context-len 8192" in script + assert "--max-tokens-per-gpu 9216" in script + assert "--custom-convert-samples-to-train-data-path examples.mem_agent.convert.convert_samples" in script + assert "--use-dynamic-global-batch-size" not in script + assert "--calculate-per-token-loss" not in script diff --git a/tests/examples/mem_agent/test_reward.py b/tests/examples/mem_agent/test_reward.py new file mode 100644 index 000000000..5e3fd4e08 --- /dev/null +++ b/tests/examples/mem_agent/test_reward.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from examples.mem_agent.reward import extract_last_boxed, normalize_answer, reward_func +from relax.utils.types import Sample + + +def test_extract_last_boxed_supports_nested_braces_and_uses_last_answer(): + assert extract_last_boxed(r"first \boxed{wrong}; final \boxed{New {York}}") == "New {York}" + assert extract_last_boxed(r"incomplete \boxed{answer") == "" + + +def test_normalize_answer_uses_hotpotqa_rules(): + assert normalize_answer("The, Eiffel Tower!") == "eiffel tower" + + +@pytest.mark.asyncio +async def test_reward_scores_only_final_output_against_all_ground_truths(): + sample = Sample( + response=r"ignored \boxed{wrong}", + label="wrong", + metadata={"final_output": r"answer: \boxed{The Eiffel Tower}", "ground_truth": ["Paris", "Eiffel Tower"]}, + train_metadata={"mem_agent_turns": [{"response": "unscored memory"}]}, + ) + result = await reward_func(SimpleNamespace(), sample) + assert result["score"] == 1.0 + assert result["pred"] == "The Eiffel Tower" + assert result["diagnostic"] == "matched" + + +@pytest.mark.asyncio +async def test_reward_reports_missing_boxed_as_zero(): + sample = Sample(response="plain answer", label="answer") + result = await reward_func(SimpleNamespace(), sample) + assert result["score"] == 0.0 + assert result["diagnostic"] == "missing_boxed" diff --git a/tests/examples/mem_agent/test_rollout.py b/tests/examples/mem_agent/test_rollout.py new file mode 100644 index 000000000..f62c3d225 --- /dev/null +++ b/tests/examples/mem_agent/test_rollout.py @@ -0,0 +1,210 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +from examples.mem_agent.prompts import NO_MEMORY, truncate_text_to_tokens +from examples.mem_agent.rollout import chunk_context, generate, generate_trajectory +from relax.utils.types import Sample + + +class FakeTokenizer: + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [ord(character) for character in text] + + def decode(self, token_ids, skip_special_tokens=True): + del skip_special_tokens + return "".join(chr(token_id) for token_id in token_ids) + + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True): + assert tokenize is False + assert add_generation_prompt is True + return f"{messages[0]['content']}" + + +def _args(): + return SimpleNamespace( + mem_agent_chunk_tokens=3, + mem_agent_max_memory_tokens=64, + mem_agent_max_final_tokens=32, + mem_agent_max_chunks=8, + ) + + +@pytest.mark.asyncio +async def test_generate_trajectory_overwrites_memory_and_expands_every_turn(): + tokenizer = FakeTokenizer() + responses = ["M1<|im_end|>", "M2", r"The answer is \boxed{x}."] + prompts = [] + max_new_tokens = [] + + async def fake_generate(args, turn, sampling_params, evaluation): + del args + prompts.append(turn.prompt) + max_new_tokens.append(sampling_params["max_new_tokens"]) + response = responses.pop(0) + prompt_ids = tokenizer.encode(turn.prompt) + response_ids = tokenizer.encode(response) + turn.tokens = prompt_ids + response_ids + turn.rollout_tokens = list(turn.tokens) + turn.response = response + turn.response_length = len(response_ids) + turn.loss_mask = [1] * len(response_ids) + turn.rollout_log_probs = [] if evaluation else [-0.25] * len(response_ids) + turn.status = Sample.Status.COMPLETED + return turn + + sample = Sample(index=11, group_index=2, prompt="Where?", metadata={"context": "abcdef"}, session_id="s") + result = await generate_trajectory( + _args(), sample, {"temperature": 1.0}, tokenizer, generator=fake_generate, evaluation=False + ) + + assert result.status == Sample.Status.COMPLETED + assert result.response == r"The answer is \boxed{x}." + assert result.metadata["num_chunks"] == 2 + assert result.metadata["context_truncated"] is False + assert result.metadata["memory_token_lengths"] == [2, 2] + assert len(result.train_metadata["mem_agent_turns"]) == 3 + assert NO_MEMORY in prompts[0] + assert "abc" in prompts[0] + assert "\nM1\n" in prompts[1] + assert "abc" not in prompts[1] + assert "def" in prompts[1] + assert "\nM2\n" in prompts[2] + assert "
" not in prompts[2] + assert "abc" not in prompts[2] and "def" not in prompts[2] + assert max_new_tokens == [64, 64, 32] + + for turn in result.train_metadata["mem_agent_turns"]: + assert len(turn["loss_mask"]) == turn["response_length"] + assert len(turn["rollout_log_probs"]) == turn["response_length"] + assert len(turn["tokens"]) >= turn["response_length"] + + +def test_chunk_context_preserves_boundaries_and_marks_truncation(): + tokenizer = FakeTokenizer() + chunks, truncated = chunk_context(tokenizer, "abcdefgh", chunk_tokens=3, max_chunks=2) + assert [tokenizer.decode(chunk) for chunk in chunks] == ["abc", "def"] + assert truncated is True + + +def test_truncate_text_to_tokens_retokenizes_to_the_hard_limit(): + text, length = truncate_text_to_tokens(FakeTokenizer(), "MEMORY", 3) + assert text == "MEM" + assert length == 3 + + class ExpandingTokenizer(FakeTokenizer): + def encode(self, text, add_special_tokens=False): + ids = super().encode(text, add_special_tokens) + return ids + ([0] if text.endswith("!") else []) + + text, length = truncate_text_to_tokens(ExpandingTokenizer(), "AB!", 3) + assert text == "AB" + assert length == 2 + + +@pytest.mark.asyncio +async def test_generate_trajectory_aborts_without_context(): + sample = Sample(index=0, prompt="question", metadata={}) + result = await generate_trajectory(_args(), sample, {}, FakeTokenizer()) + assert result.status == Sample.Status.ABORTED + assert result.train_metadata is None + + +@pytest.mark.asyncio +async def test_generate_trajectory_bounds_memory_before_next_turn(): + tokenizer = FakeTokenizer() + responses = ["MEMORY", r"\boxed{x}"] + prompts = [] + + async def fake_generate(args, turn, sampling_params, evaluation): + del args, evaluation + prompts.append(turn.prompt) + response = responses.pop(0) + response_ids = tokenizer.encode(response)[: sampling_params["max_new_tokens"]] + turn.response = tokenizer.decode(response_ids) + turn.tokens = tokenizer.encode(turn.prompt) + response_ids + turn.rollout_tokens = list(turn.tokens) + turn.response_length = len(response_ids) + turn.loss_mask = [1] * len(response_ids) + turn.rollout_log_probs = [-0.1] * len(response_ids) + turn.status = Sample.Status.TRUNCATED if response == "MEMORY" else Sample.Status.COMPLETED + return turn + + args = _args() + args.mem_agent_max_memory_tokens = 3 + sample = Sample(index=0, group_index=0, prompt="Q", metadata={"context": "abc"}) + result = await generate_trajectory(args, sample, {}, tokenizer, generator=fake_generate) + assert result.metadata["memory_token_lengths"] == [3] + assert "\nMEM\n" in prompts[-1] + assert "MEMORY" not in prompts[-1] + + +@pytest.mark.asyncio +async def test_generate_trajectory_empty_update_does_not_reuse_old_memory(): + tokenizer = FakeTokenizer() + responses = ["M1", "<|im_end|>", r"\boxed{x}"] + prompts = [] + + async def fake_generate(args, turn, sampling_params, evaluation): + del args, sampling_params, evaluation + prompts.append(turn.prompt) + response = responses.pop(0) + response_ids = tokenizer.encode(response) + turn.response = response + turn.tokens = tokenizer.encode(turn.prompt) + response_ids + turn.rollout_tokens = list(turn.tokens) + turn.response_length = len(response_ids) + turn.loss_mask = [1] * len(response_ids) + turn.rollout_log_probs = [-0.1] * len(response_ids) + turn.status = Sample.Status.COMPLETED + return turn + + sample = Sample(index=0, group_index=0, prompt="Q", metadata={"context": "abcdef"}) + result = await generate_trajectory(_args(), sample, {}, tokenizer, generator=fake_generate) + + assert result.metadata["memory_token_lengths"] == [2, 0] + assert "\n\n" in prompts[-1] + assert "\nM1\n" not in prompts[-1] + + +@pytest.mark.asyncio +async def test_public_generate_entry_uses_sglang_state_and_turn_generator(monkeypatch): + """Exercise the exact callable wired by --custom-generate-function-path.""" + tokenizer = FakeTokenizer() + responses = iter(["MEM", r"\boxed{x}"]) + + class FakeGenerateState: + def __init__(self, args): + del args + self.tokenizer = tokenizer + + async def fake_sglang_generate(args, turn, sampling_params, evaluation): + del args, sampling_params, evaluation + response = next(responses) + response_ids = tokenizer.encode(response) + turn.response = response + turn.tokens = tokenizer.encode(turn.prompt) + response_ids + turn.rollout_tokens = list(turn.tokens) + turn.response_length = len(response_ids) + turn.loss_mask = [1] * len(response_ids) + turn.rollout_log_probs = [-0.1] * len(response_ids) + turn.status = Sample.Status.COMPLETED + return turn + + fake_module = ModuleType("relax.engine.rollout.sglang_rollout") + fake_module.GenerateState = FakeGenerateState + fake_module.generate = fake_sglang_generate + monkeypatch.setitem(sys.modules, fake_module.__name__, fake_module) + + sample = Sample(index=3, group_index=0, prompt="Q", metadata={"context": "abc"}) + result = await generate(_args(), sample, {"temperature": 1.0}) + + assert result.status == Sample.Status.COMPLETED + assert result.response == r"\boxed{x}" + assert [turn["kind"] for turn in result.train_metadata["mem_agent_turns"]] == ["memory", "final"] diff --git a/tests/utils/test_custom_sample_converter.py b/tests/utils/test_custom_sample_converter.py new file mode 100644 index 000000000..0804e251d --- /dev/null +++ b/tests/utils/test_custom_sample_converter.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from examples.mem_agent.convert import convert_samples +from relax.utils.types import Sample +from relax.utils.utils import ( + CURRENT_ROLLOUT_BATCH, + convert_samples_to_train_data_with_custom, + get_debug_data, + get_train_data_group_size, + get_train_sample_expansion_factor, + transfer_batch_to_data_system, +) + + +def _sample() -> Sample: + turns = [] + for turn_index in range(3): + turns.append( + { + "turn_index": turn_index, + "kind": "final" if turn_index == 2 else "memory", + "tokens": [10 + turn_index, 20, 21], + "response_length": 2, + "loss_mask": [1, 1], + "rollout_log_probs": [-0.1, -0.2], + "finish_reason": Sample.Status.COMPLETED.value, + } + ) + return Sample( + index=7, + group_index=0, + reward={"score": 1.0}, + status=Sample.Status.COMPLETED, + train_metadata={"mem_agent_turns": turns}, + ) + + +def _args(**overrides): + values = { + "custom_convert_samples_to_train_data_path": "examples.mem_agent.convert.convert_samples", + "custom_reward_post_process_path": None, + "agentic_custom_advantage_path": None, + "reward_key": "score", + "advantage_estimator": "grpo", + "rewards_normalization": False, + "grpo_std_normalization": True, + "n_samples_per_prompt": 1, + "mem_agent_credit_assignment": "split", + "debug_train_only": False, + "load_debug_rollout_data_subsample": None, + "use_dynamic_global_batch_size": False, + "custom_train_expanded_batch": False, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_converter_dispatch_preserves_default_and_custom_paths(monkeypatch): + sentinel = object() + monkeypatch.setattr("relax.utils.utils.convert_samples_to_train_data", lambda args, samples: sentinel) + args = _args(custom_convert_samples_to_train_data_path=None) + assert convert_samples_to_train_data_with_custom(args, [_sample()]) is sentinel + + custom = object() + monkeypatch.setattr("relax.utils.utils.load_function", lambda path: lambda args, samples: custom) + args.custom_convert_samples_to_train_data_path = "package.converter" + assert convert_samples_to_train_data_with_custom(args, [_sample()]) is custom + + +def test_debug_replay_uses_custom_converter(tmp_path): + path_template = str(tmp_path / "rollout_{rollout_id}.pt") + torch.save({"samples": [_sample().to_dict()]}, path_template.format(rollout_id=3)) + args = _args(debug_train_only=True, load_debug_rollout_data=path_template) + replay = get_debug_data(args, rollout_id=3, batch_size=3, dp_rank=0) + assert replay["turn_indices"] == [0, 1, 2] + assert len(replay["tokens"]) == 3 + + +def test_dynamic_debug_replay_slices_expanded_rows_across_dp(tmp_path): + first = _sample() + second = _sample() + second.index = 8 + second.group_index = 1 + path_template = str(tmp_path / "rollout_{rollout_id}.pt") + torch.save({"samples": [first.to_dict(), second.to_dict()]}, path_template.format(rollout_id=4)) + args = _args( + debug_train_only=True, + load_debug_rollout_data=path_template, + custom_train_expanded_batch=True, + ) + replay = get_debug_data(args, rollout_id=4, batch_size=None, dp_rank=1, dp_size=2) + assert len(replay["tokens"]) == 3 + assert replay["sample_indices"] == [8, 8, 8] + + +@pytest.mark.asyncio +async def test_online_transfer_queues_every_expanded_turn(): + class Client: + payload = None + + async def async_put(self, **kwargs): + self.payload = kwargs + + client = Client() + CURRENT_ROLLOUT_BATCH.clear() + destination_rollout_id, transferred_rows = await transfer_batch_to_data_system( + _args(), [[_sample()]], 1, 5, client, is_last=True + ) + assert destination_rollout_id == 5 + assert transferred_rows == 3 + assert len(client.payload["data"]["total_lengths"]) == 3 + assert len(client.payload["custom_meta"]) == 3 + assert client.payload["partition_id"] == "train_5" + assert client.payload["is_last"] is True + + +def test_expansion_capacity_and_group_defaults_are_backward_compatible(): + default = SimpleNamespace(n_samples_per_prompt=8) + expanded = SimpleNamespace( + n_samples_per_prompt=8, + custom_train_sample_expansion_factor=65, + custom_train_data_group_size=1, + ) + assert get_train_sample_expansion_factor(default) == 1 + assert get_train_data_group_size(default) == 8 + assert get_train_sample_expansion_factor(expanded) == 65 + assert get_train_data_group_size(expanded) == 1 + + +def test_custom_converter_function_returns_all_rows_without_tail_trim(): + data = convert_samples(_args(debug_train_only=True), [_sample()]) + assert len(data["tokens"]) == 3 From 7d34bbb0ec1e4c5e7bf183286c98619ab50e1a6c Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 14:12:36 +0800 Subject: [PATCH 02/16] fix: harden MemAgent acceptance contracts --- examples/mem_agent/README.md | 5 +- examples/mem_agent/compare_results.py | 75 ++++++++++++++++++- examples/mem_agent/convert.py | 61 ++++++++++++--- examples/mem_agent/eval_ruler_hqa.py | 27 ++++++- examples/mem_agent/rollout.py | 22 ++++-- examples/mem_agent/run-eval.sh | 1 + examples/mem_agent/run-paired-eval.sh | 25 +++++-- relax/backends/megatron/actor.py | 9 +++ relax/distributed/ray/rollout.py | 46 +++--------- relax/distributed/ray/utils.py | 52 +++++++++++++ .../ray/test_train_batch_row_count_tracker.py | 46 ++++++++++++ .../mem_agent/test_compare_results.py | 34 ++++++++- tests/examples/mem_agent/test_convert.py | 32 +++++++- .../mem_agent/test_data_and_metrics.py | 49 +++++++++++- tests/examples/mem_agent/test_eval.py | 40 +++++++++- .../mem_agent/test_recipe_contract.py | 23 ++++++ tests/examples/mem_agent/test_rollout.py | 26 +++++++ tests/utils/test_custom_sample_converter.py | 2 + 18 files changed, 507 insertions(+), 68 deletions(-) create mode 100644 tests/distributed/ray/test_train_batch_row_count_tracker.py diff --git a/examples/mem_agent/README.md b/examples/mem_agent/README.md index 91f3e5ffb..b9de3beff 100644 --- a/examples/mem_agent/README.md +++ b/examples/mem_agent/README.md @@ -48,13 +48,14 @@ RESULTS_DIR=/data/results/mem-agent-relax \ bash examples/mem_agent/run-eval.sh ``` -The evaluator writes raw per-sample JSONL and a summary JSON for HotpotQA dev and RULER-HQA 50/200/800. Failed requests remain in the denominator with score zero. `boxed_em_pct` is the HotpotQA reward-compatible accuracy and `sub_em_pct` is the primary VIME-compatible RULER-HQA metric. Set `MODE=base` to run the single-context base baseline; its context truncation always preserves the question and answer instruction. +The evaluator writes raw per-sample JSONL and a summary JSON for HotpotQA dev and RULER-HQA 50/200/800. Failed requests keep their ground truth in the raw file and remain in the denominator with score zero. `boxed_em_pct` is the HotpotQA reward-compatible accuracy and `sub_em_pct` is the primary VIME-compatible RULER-HQA metric. Set `MODE=base` to run the optional single-context diagnostic; its context truncation always preserves the question and answer instruction. `TOKENIZER_PATH` should point to the frozen base snapshot. `run-pipeline.sh` preserves it automatically before switching `MODEL_PATH` to the converted checkpoint. When `NUM_ROLLOUT=2` is used, the pipeline also selects `iter_0000001` automatically instead of the 100-step default `iter_0000099`. -For the VIME reproduction tolerance, evaluate an official VIME checkpoint when one is available; otherwise use a checkpoint produced once from the fixed VIME recipe. The paired runner holds the tokenizer, data, prompts, sampling parameters, and evaluator constant. It evaluates RULER-HQA 50/200/800 by default; `LENGTHS` can freeze a smaller pre-agreed subset before either result is observed. It exits non-zero when any selected `sub_em_pct` differs by more than 3 percentage points and retains raw per-sample output for review: +For the VIME reproduction tolerance, evaluate an official VIME checkpoint when one is available; otherwise use a checkpoint produced once from the fixed VIME recipe. The acceptance runner holds the tokenizer, data, prompts, sampling parameters, recurrent inference mode, and evaluator constant across frozen base, VIME, and ReLax. It evaluates RULER-HQA 50/200/800 by default; `LENGTHS` can freeze a smaller pre-agreed subset before any result is observed. It exits non-zero unless every selected VIME/ReLax `sub_em_pct` gap is at most 3 percentage points, ReLax beats frozen base on every selected RULER-HQA `sub_em_pct`, and ReLax beats frozen base on HotpotQA `boxed_em_pct`. Raw per-sample files are retained for review: ```bash +BASE_MODEL_PATH=/data/models/Qwen3-4B \ VIME_MODEL_PATH=/data/checkpoints/vime-hf \ RELAX_MODEL_PATH=/data/checkpoints/mem-agent-relax-hf \ TOKENIZER_PATH=/data/models/Qwen3-4B \ diff --git a/examples/mem_agent/compare_results.py b/examples/mem_agent/compare_results.py index b3d2d85fb..2b476dd75 100644 --- a/examples/mem_agent/compare_results.py +++ b/examples/mem_agent/compare_results.py @@ -9,6 +9,35 @@ from typing import Any +COMPATIBILITY_FIELDS = ( + "data_file", + "mode", + "tokenizer", + "temperature", + "top_p", + "sampling_count", + "chunk_tokens", + "max_memory_tokens", + "max_final_tokens", + "max_chunks", + "max_input_tokens", + "server_max_model_len", + "total", +) + + +def validate_compatible_summaries(*summaries: dict[str, Any]) -> None: + """Reject a comparison when any controlled evaluation field differs.""" + if len(summaries) < 2: + raise ValueError("At least two summaries are required for compatibility validation.") + for field in COMPATIBILITY_FIELDS: + if any(field not in summary for summary in summaries): + raise KeyError(f"Compatibility field {field!r} must exist in every summary.") + values = [summary[field] for summary in summaries] + if any(value != values[0] for value in values[1:]): + raise ValueError(f"Evaluation summaries differ on controlled field {field!r}: {values}") + + def compare_pair( label: str, vime_summary: dict[str, Any], @@ -38,6 +67,28 @@ def compare_pair( } +def compare_baseline( + label: str, + base_summary: dict[str, Any], + relax_summary: dict[str, Any], + metric: str, +) -> dict[str, Any]: + """Require the trained ReLax checkpoint to strictly beat frozen base.""" + if metric not in base_summary or metric not in relax_summary: + raise KeyError(f"Metric {metric!r} must exist in both summaries.") + base_value = float(base_summary[metric]) + relax_value = float(relax_summary[metric]) + improvement_pp = relax_value - base_value + return { + "label": label, + "metric": metric, + "base": base_value, + "relax": relax_value, + "improvement_pp": improvement_pp, + "passed": improvement_pp > 0, + } + + def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( @@ -50,6 +101,14 @@ def main() -> None: ) parser.add_argument("--metric", default="sub_em_pct") parser.add_argument("--tolerance-pp", type=float, default=3.0) + parser.add_argument( + "--baseline-pair", + nargs=4, + action="append", + default=[], + metavar=("LABEL", "METRIC", "BASE_SUMMARY", "RELAX_SUMMARY"), + help="Require ReLax to strictly exceed frozen base; repeat for every required metric/dataset.", + ) parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() @@ -59,16 +118,30 @@ def main() -> None: vime_summary = json.load(source) with Path(relax_path).open(encoding="utf-8") as source: relax_summary = json.load(source) + validate_compatible_summaries(vime_summary, relax_summary) comparison = compare_pair(label, vime_summary, relax_summary, args.metric, args.tolerance_pp) comparison["vime_summary"] = str(vime_path) comparison["relax_summary"] = str(relax_path) comparisons.append(comparison) + baseline_comparisons = [] + for label, metric, base_path, relax_path in args.baseline_pair: + with Path(base_path).open(encoding="utf-8") as source: + base_summary = json.load(source) + with Path(relax_path).open(encoding="utf-8") as source: + relax_summary = json.load(source) + validate_compatible_summaries(base_summary, relax_summary) + comparison = compare_baseline(label, base_summary, relax_summary, metric) + comparison["base_summary"] = str(base_path) + comparison["relax_summary"] = str(relax_path) + baseline_comparisons.append(comparison) + report = { "metric": args.metric, "tolerance_pp": args.tolerance_pp, - "passed": all(item["passed"] for item in comparisons), + "passed": all(item["passed"] for item in comparisons + baseline_comparisons), "comparisons": comparisons, + "baseline_comparisons": baseline_comparisons, } args.output.parent.mkdir(parents=True, exist_ok=True) with args.output.open("w", encoding="utf-8") as destination: diff --git a/examples/mem_agent/convert.py b/examples/mem_agent/convert.py index b6cfcfe33..542d69121 100644 --- a/examples/mem_agent/convert.py +++ b/examples/mem_agent/convert.py @@ -5,10 +5,14 @@ from typing import Any +from relax.utils.logging_utils import get_logger from relax.utils.types import Sample from relax.utils.utils import dict_to_tensordict, post_process_rewards +logger = get_logger(__name__) + + def _get_turns(sample: Sample) -> list[dict[str, Any]]: metadata = sample.train_metadata or {} turns = metadata.get("mem_agent_turns", metadata.get("turns")) @@ -17,20 +21,46 @@ def _get_turns(sample: Sample) -> list[dict[str, Any]]: return turns -def _validate_turn(sample: Sample, turn: dict[str, Any]) -> None: - required = {"tokens", "response_length", "loss_mask", "rollout_log_probs"} +def _validate_turn(args: Any, sample: Sample, turn: dict[str, Any]) -> None: + required = {"kind", "turn_index", "tokens", "response_length", "loss_mask", "rollout_log_probs", "finish_reason"} missing = required.difference(turn) if missing: raise ValueError(f"Sample index={sample.index} turn is missing fields: {sorted(missing)}") + kind = turn["kind"] + if kind not in ("memory", "final"): + raise ValueError(f"Sample index={sample.index} has invalid turn kind={kind!r}.") + if turn["finish_reason"] not in (Sample.Status.COMPLETED.value, Sample.Status.TRUNCATED.value): + raise ValueError(f"Sample index={sample.index} has invalid finish_reason={turn['finish_reason']!r}.") response_length = int(turn["response_length"]) if response_length <= 0 or len(turn["tokens"]) < response_length: raise ValueError(f"Sample index={sample.index} has an invalid response_length.") + limit_name = "mem_agent_max_memory_tokens" if kind == "memory" else "mem_agent_max_final_tokens" + max_response_tokens = int(getattr(args, limit_name, 1024 if kind == "memory" else 256)) + if response_length > max_response_tokens: + raise ValueError( + f"Sample index={sample.index} {kind} response_length={response_length} exceeds {max_response_tokens}." + ) if len(turn["loss_mask"]) != response_length: raise ValueError(f"Sample index={sample.index} has a misaligned loss_mask.") if len(turn["rollout_log_probs"]) != response_length: raise ValueError(f"Sample index={sample.index} has misaligned rollout_log_probs.") +def _validate_trajectory(args: Any, sample: Sample, turns: list[dict[str, Any]]) -> None: + """Require ordered memory turns followed by exactly one final turn.""" + expected_indices = list(range(len(turns))) + actual_indices = [int(turn.get("turn_index", -1)) for turn in turns] + if actual_indices != expected_indices: + raise ValueError( + f"Sample index={sample.index} has non-contiguous turn indices: {actual_indices}, expected {expected_indices}." + ) + kinds = [turn.get("kind") for turn in turns] + if kinds[-1] != "final" or any(kind != "memory" for kind in kinds[:-1]): + raise ValueError(f"Sample index={sample.index} must contain memory turns followed by exactly one final turn.") + for turn in turns: + _validate_turn(args, sample, turn) + + def convert_samples(args: Any, samples: list[Sample]): """Normalize trajectory rewards first, then expand every saved turn.""" if not samples: @@ -41,9 +71,6 @@ def convert_samples(args: Any, samples: list[Sample]): if sample.status in (Sample.Status.ABORTED, Sample.Status.FAILED): raise ValueError(f"Cannot train from sample index={sample.index} with status={sample.status.value}.") - # Normalize while there is still exactly one row per trajectory. Expanding - # first would overweight long documents in the group statistics. - raw_rewards, advantages = post_process_rewards(args, samples) credit_assignment = getattr(args, "mem_agent_credit_assignment", "split") if credit_assignment not in ("split", "share"): raise ValueError("mem_agent_credit_assignment must be either 'split' or 'share'.") @@ -51,17 +78,26 @@ def convert_samples(args: Any, samples: list[Sample]): # Every trajectory in one GRPO prompt group reads the same context, so it # must have the same number of turns. Besides catching partial trajectories, # this makes each expanded group n_samples_per_prompt * turn_count rows. + validated_turns: list[list[dict[str, Any]]] = [] turn_counts_by_group: dict[int, set[int]] = {} for sample in samples: if sample.group_index is None: raise ValueError("MemAgent samples require group_index.") - turn_counts_by_group.setdefault(sample.group_index, set()).add(len(_get_turns(sample))) + turns = _get_turns(sample) + _validate_trajectory(args, sample, turns) + validated_turns.append(turns) + turn_counts_by_group.setdefault(sample.group_index, set()).add(len(turns)) inconsistent_groups = { group_index: counts for group_index, counts in turn_counts_by_group.items() if len(counts) != 1 } if inconsistent_groups: raise ValueError(f"MemAgent prompt groups have inconsistent turn counts: {inconsistent_groups}") + # Normalize only after the full trajectory contract is known to be valid, + # while there is still exactly one row per trajectory. Expanding first + # would overweight long documents in the group statistics. + raw_rewards, advantages = post_process_rewards(args, samples) + train_data: dict[str, list[Any]] = { "tokens": [], "response_lengths": [], @@ -76,15 +112,13 @@ def convert_samples(args: Any, samples: list[Sample]): "total_lengths": [], } - for trajectory_position, (sample, raw_reward, advantage) in enumerate( - zip(samples, raw_rewards, advantages, strict=True) + for trajectory_position, (sample, turns, raw_reward, advantage) in enumerate( + zip(samples, validated_turns, raw_rewards, advantages, strict=True) ): - turns = _get_turns(sample) turn_credit = float(advantage) / len(turns) if credit_assignment == "split" else float(advantage) # This expansion is deliberately lossless. Unlike the fixed VIME # helper, no tail rows are trimmed to a global-batch multiple. for fallback_turn_index, turn in enumerate(turns): - _validate_turn(sample, turn) tokens = list(turn["tokens"]) response_length = int(turn["response_length"]) loss_mask = list(turn["loss_mask"]) @@ -111,6 +145,13 @@ def convert_samples(args: Any, samples: list[Sample]): f"mem_agent_train_rows_multiple={required_multiple}." ) + # One concise count line makes the GPU smoke invariant reviewable: + # saved turns and emitted train rows must be identical, with no tail trim. + logger.info( + f"Expanded MemAgent trajectories={len(samples)} saved_turns={sum(len(turns) for turns in validated_turns)} " + f"train_rows={len(train_data['tokens'])}" + ) + if getattr(args, "debug_train_only", False): return train_data return dict_to_tensordict(train_data, len(train_data["tokens"])) diff --git a/examples/mem_agent/eval_ruler_hqa.py b/examples/mem_agent/eval_ruler_hqa.py index d472e1e97..ad40d7a53 100644 --- a/examples/mem_agent/eval_ruler_hqa.py +++ b/examples/mem_agent/eval_ruler_hqa.py @@ -172,22 +172,23 @@ async def run_evaluation( async with aiohttp.ClientSession(timeout=timeout) as session: async def evaluate_one(item: dict[str, Any]) -> dict[str, Any]: + answers = item["answers"] if isinstance(item["answers"], list) else [item["answers"]] + answers = [str(answer) for answer in answers] + ground_truth = answers[0] try: async with semaphore: if args.mode == "recurrent": response, diagnostics = await recurrent_infer(item, args, tokenizer, session) else: response, diagnostics = await base_infer(item, args, tokenizer, session) - answers = item["answers"] if isinstance(item["answers"], list) else [item["answers"]] # RULER-HQA's VIME-compatible metrics score the first reference. # boxed_em additionally mirrors the HotpotQA training reward and # accepts any annotated answer. - ground_truth = str(answers[0]) prediction = extract_last_boxed(response[-300:]) return { "_id": item["_id"], "answer": ground_truth, - "answers": [str(answer) for answer in answers], + "answers": answers, "pred": prediction, "judge_f1": f1_score(prediction, ground_truth), "judge_em": exact_match(prediction, ground_truth), @@ -197,7 +198,21 @@ async def evaluate_one(item: dict[str, Any]) -> dict[str, Any]: **diagnostics, } except Exception as exc: - return {"_id": item["_id"], "error": f"{type(exc).__name__}: {exc}"} + # Preserve the target and explicit zero scores in raw output. + # Reviewers can therefore audit every input row even when the + # serving request itself failed. + return { + "_id": item["_id"], + "answer": ground_truth, + "answers": answers, + "pred": "", + "judge_f1": 0.0, + "judge_em": 0.0, + "judge_sub_em": 0.0, + "judge_boxed_em": 0.0, + "response": "", + "error": f"{type(exc).__name__}: {exc}", + } records = await asyncio.gather(*(evaluate_one(item) for item in data)) @@ -205,6 +220,7 @@ async def evaluate_one(item: dict[str, Any]) -> dict[str, Any]: **aggregate(records), "mode": args.mode, "model": args.model, + "tokenizer": args.tokenizer, "data_file": str(args.data_file), "temperature": args.temperature, "top_p": args.top_p, @@ -213,6 +229,8 @@ async def evaluate_one(item: dict[str, Any]) -> dict[str, Any]: "max_memory_tokens": args.max_memory_tokens, "max_final_tokens": args.max_final_tokens, "max_chunks": args.max_chunks, + "max_input_tokens": args.max_input_tokens, + "server_max_model_len": args.server_max_model_len, } return records, summary @@ -234,6 +252,7 @@ def main() -> None: parser.add_argument("--max-final-tokens", type=int, default=256) parser.add_argument("--max-chunks", type=int, default=64) parser.add_argument("--max-input-tokens", type=int, default=7936) + parser.add_argument("--server-max-model-len", type=int, default=8192) parser.add_argument("--concurrency", type=int, default=16) parser.add_argument("--timeout", type=int, default=86400) args = parser.parse_args() diff --git a/examples/mem_agent/rollout.py b/examples/mem_agent/rollout.py index f3aa11f2e..5bcdc4673 100644 --- a/examples/mem_agent/rollout.py +++ b/examples/mem_agent/rollout.py @@ -45,10 +45,16 @@ def _question_from_sample(sample: Sample) -> str: return "" -def _validate_turn(turn: dict[str, Any], require_log_probs: bool) -> None: +def _validate_turn(turn: dict[str, Any], require_log_probs: bool, max_response_tokens: int) -> None: + if turn["finish_reason"] not in (Sample.Status.COMPLETED.value, Sample.Status.TRUNCATED.value): + raise ValueError(f"MemAgent turn has invalid finish status: {turn['finish_reason']}.") response_length = turn["response_length"] if response_length <= 0: raise ValueError("MemAgent turn returned no trainable response tokens.") + if response_length > max_response_tokens: + raise ValueError( + f"MemAgent turn returned {response_length} response tokens, exceeding its limit {max_response_tokens}." + ) if len(turn["tokens"]) < response_length: raise ValueError("MemAgent turn response_length exceeds total token count.") if len(turn["loss_mask"]) != response_length: @@ -58,7 +64,13 @@ def _validate_turn(turn: dict[str, Any], require_log_probs: bool) -> None: raise ValueError("MemAgent turn rollout_log_probs are not aligned with response tokens.") -def _turn_record(turn_sample: Sample, turn_index: int, kind: str, evaluation: bool) -> dict[str, Any]: +def _turn_record( + turn_sample: Sample, + turn_index: int, + kind: str, + evaluation: bool, + max_response_tokens: int, +) -> dict[str, Any]: record = { "turn_index": turn_index, "kind": kind, @@ -68,7 +80,7 @@ def _turn_record(turn_sample: Sample, turn_index: int, kind: str, evaluation: bo "rollout_log_probs": list(turn_sample.rollout_log_probs or []), "finish_reason": turn_sample.status.value, } - _validate_turn(record, require_log_probs=not evaluation) + _validate_turn(record, require_log_probs=not evaluation, max_response_tokens=max_response_tokens) return record @@ -151,7 +163,7 @@ async def generate_trajectory( sample.status = Sample.Status.ABORTED sample.rollout_log_probs = [] return sample - turns.append(_turn_record(turn_sample, len(turns), "memory", evaluation)) + turns.append(_turn_record(turn_sample, len(turns), "memory", evaluation, max_memory_tokens)) # Overwrite on every turn, including an empty post-processed response; # carrying the old value forward would silently change the recurrence. memory, memory_length = truncate_text_to_tokens( @@ -178,7 +190,7 @@ async def generate_trajectory( sample.status = Sample.Status.ABORTED sample.rollout_log_probs = [] return sample - turns.append(_turn_record(final_sample, len(turns), "final", evaluation)) + turns.append(_turn_record(final_sample, len(turns), "final", evaluation, max_final_tokens)) final_output = strip_stop_tokens(final_sample.response) # Preserve a valid top-level Sample for logs and failure recovery. The diff --git a/examples/mem_agent/run-eval.sh b/examples/mem_agent/run-eval.sh index 83dcca0ea..c5e075e57 100755 --- a/examples/mem_agent/run-eval.sh +++ b/examples/mem_agent/run-eval.sh @@ -61,6 +61,7 @@ run_eval() { --max-memory-tokens 1024 \ --max-final-tokens 256 \ --max-chunks 64 \ + --server-max-model-len "${MAX_MODEL_LEN}" \ --concurrency "${CONCURRENCY}" } diff --git a/examples/mem_agent/run-paired-eval.sh b/examples/mem_agent/run-paired-eval.sh index 2f4c14c23..530918797 100755 --- a/examples/mem_agent/run-paired-eval.sh +++ b/examples/mem_agent/run-paired-eval.sh @@ -4,6 +4,7 @@ set -euo pipefail SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +BASE_MODEL_PATH="${BASE_MODEL_PATH:?Set BASE_MODEL_PATH to the frozen Qwen3-4B snapshot.}" VIME_MODEL_PATH="${VIME_MODEL_PATH:?Set VIME_MODEL_PATH to the official or equivalently trained VIME checkpoint.}" RELAX_MODEL_PATH="${RELAX_MODEL_PATH:?Set RELAX_MODEL_PATH to the converted ReLax checkpoint.}" TOKENIZER_PATH="${TOKENIZER_PATH:?Set TOKENIZER_PATH to the frozen Qwen3-4B tokenizer snapshot.}" @@ -11,15 +12,19 @@ DATA_DIR="${DATA_DIR:?Set DATA_DIR to the prepared frozen evaluation data.}" RESULTS_DIR="${RESULTS_DIR:?Set RESULTS_DIR for raw and summary results.}" LENGTHS="${LENGTHS:-50 200 800}" -mkdir -p "${RESULTS_DIR}/vime" "${RESULTS_DIR}/relax" +mkdir -p "${RESULTS_DIR}/base" "${RESULTS_DIR}/vime" "${RESULTS_DIR}/relax" # Start and stop each checkpoint independently while preserving every other -# serving and evaluation variable. run-eval.sh writes both raw JSONL records -# and summary JSON, so the final comparison remains sample-auditable. -MODEL_PATH="${VIME_MODEL_PATH}" RUN_NAME=vime RESULTS_DIR="${RESULTS_DIR}/vime" \ +# serving and evaluation variable. Frozen base also uses recurrent mode: the +# policy weights are the only variable in the trained-vs-base acceptance. +# run-eval.sh retains raw JSONL so every aggregate remains sample-auditable. +MODEL_PATH="${BASE_MODEL_PATH}" RUN_NAME=base RESULTS_DIR="${RESULTS_DIR}/base" MODE=recurrent \ TOKENIZER_PATH="${TOKENIZER_PATH}" DATA_DIR="${DATA_DIR}" LENGTHS="${LENGTHS}" \ bash "${SCRIPT_DIR}/run-eval.sh" -MODEL_PATH="${RELAX_MODEL_PATH}" RUN_NAME=relax RESULTS_DIR="${RESULTS_DIR}/relax" \ +MODEL_PATH="${VIME_MODEL_PATH}" RUN_NAME=vime RESULTS_DIR="${RESULTS_DIR}/vime" MODE=recurrent \ + TOKENIZER_PATH="${TOKENIZER_PATH}" DATA_DIR="${DATA_DIR}" LENGTHS="${LENGTHS}" \ + bash "${SCRIPT_DIR}/run-eval.sh" +MODEL_PATH="${RELAX_MODEL_PATH}" RUN_NAME=relax RESULTS_DIR="${RESULTS_DIR}/relax" MODE=recurrent \ TOKENIZER_PATH="${TOKENIZER_PATH}" DATA_DIR="${DATA_DIR}" LENGTHS="${LENGTHS}" \ bash "${SCRIPT_DIR}/run-eval.sh" @@ -29,10 +34,18 @@ for length in ${LENGTHS}; do --pair "ruler-hqa-${length}" "${RESULTS_DIR}/vime/vime-ruler-hqa-${length}.summary.json" "${RESULTS_DIR}/relax/relax-ruler-hqa-${length}.summary.json" + --baseline-pair "ruler-hqa-${length}" sub_em_pct + "${RESULTS_DIR}/base/base-ruler-hqa-${length}.summary.json" + "${RESULTS_DIR}/relax/relax-ruler-hqa-${length}.summary.json" ) done +COMPARE_ARGS+=( + --baseline-pair hotpotqa-dev boxed_em_pct + "${RESULTS_DIR}/base/base-hotpotqa-dev.summary.json" + "${RESULTS_DIR}/relax/relax-hotpotqa-dev.summary.json" +) python3 "${SCRIPT_DIR}/compare_results.py" \ "${COMPARE_ARGS[@]}" \ --metric sub_em_pct \ --tolerance-pp 3.0 \ - --output "${RESULTS_DIR}/vime-vs-relax.json" + --output "${RESULTS_DIR}/acceptance-comparison.json" diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index f84139f86..3c2c35e44 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -681,6 +681,10 @@ def train(self, rollout_id: int) -> None: ) batch_size = expanded_rows // dp_size num_rollout_minis = 1 + logger.info( + f"Expanded train row contract rollout_id={rollout_id}: " + f"transferred_global_rows={expanded_rows}, dp_size={dp_size}, local_rows={batch_size}" + ) elif self.args.partial_rollout and self.args.use_dynamic_global_batch_size: dynamic_size = ray.get(self.rollout_manager.get_dynamic_global_batch_size.remote()) batch_size = dynamic_size // dp_size @@ -735,6 +739,11 @@ def train(self, rollout_id: int) -> None: f"batch_index={batch_index - 1}: expected {batch_size}, " f"got {len(rollout_data['total_lengths'])}." ) + if getattr(self.args, "custom_train_expanded_batch", False): + logger.info( + f"Consumed expanded train rows rollout_id={rollout_id}: " + f"local_rows={len(rollout_data['total_lengths'])}, expected_local_rows={batch_size}" + ) rollout_mini_batches.append(rollout_data) rollout_mini_batch_metas.append(batch_meta) rollout_mini_local_sample_counts.append(len(rollout_data["total_lengths"])) diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index 2eaea9dc7..c35d35548 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -58,7 +58,7 @@ from relax.utils.types import Sample from relax.utils.utils import get_ray_accelerator_kwargs -from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock +from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock, TrainBatchRowCountTracker logging.getLogger("httpx").setLevel(logging.WARNING) @@ -794,12 +794,9 @@ def __init__(self, args, pg, data_source=None): self.pg = pg self.args = args self._dynamic_global_batch_size = None - # Expanded row counts are keyed by rollout ID because rollout and actor - # scheduling can overlap. A single "latest" value lets step N+1 reuse - # step N's count and can leave TransferQueue rows unread. - self._train_batch_row_counts: dict[int, int] = {} - self._train_batch_row_errors: dict[int, str] = {} - self._train_batch_row_events: dict[int, asyncio.Event] = {} + # Actor and rollout scheduling can overlap. The tracker prevents step + # N+1 from reusing step N's expanded row count. + self._train_batch_row_tracker = TrainBatchRowCountTracker() init_tracking(args, primary=False) @@ -1006,13 +1003,7 @@ async def get_train_batch_row_count(self, rollout_id: int) -> int: same event object lets both call orders converge without polling or falling back to a stale count from the previous step. """ - event = self._train_batch_row_events.setdefault(rollout_id, asyncio.Event()) - await event.wait() - if error := self._train_batch_row_errors.get(rollout_id): - raise RuntimeError(f"Failed to determine train rows for rollout_id={rollout_id}: {error}") - if rollout_id not in self._train_batch_row_counts: - raise RuntimeError(f"Train row count is unavailable for rollout_id={rollout_id}.") - return self._train_batch_row_counts[rollout_id] + return await self._train_batch_row_tracker.wait(rollout_id) def get_num_rollout_per_epoch(self): assert self.args.rollout_global_dataset @@ -1021,16 +1012,9 @@ def get_num_rollout_per_epoch(self): async def generate(self, rollout_id): self.rollout_id = rollout_id self.health_monitoring_resume() - row_event = self._train_batch_row_events.setdefault(rollout_id, asyncio.Event()) - row_event.clear() - self._train_batch_row_counts.pop(rollout_id, None) - self._train_batch_row_errors.pop(rollout_id, None) - # Counts are only coordination state; retain a short overlap window for - # late actor/critic consumers and bound the actor's memory usage. - for stale_rollout_id in [key for key in self._train_batch_row_events if key < rollout_id - 2]: - self._train_batch_row_events.pop(stale_rollout_id, None) - self._train_batch_row_counts.pop(stale_rollout_id, None) - self._train_batch_row_errors.pop(stale_rollout_id, None) + expanded_batch = getattr(self.args, "custom_train_expanded_batch", False) + if expanded_batch: + self._train_batch_row_tracker.start(rollout_id) try: if self.args.ci_test and self.args.use_fault_tolerance and rollout_id >= 2: self._try_ci_fault_injection() @@ -1043,7 +1027,7 @@ async def generate(self, rollout_id): self.data_system_client, evaluation=False, ) - if getattr(self.args, "custom_train_expanded_batch", False): + if expanded_batch: # The transfer helper reports the exact rows it actually put. # Re-converting output.samples here is unsafe because custom # converters may filter rows or be stateful. @@ -1051,21 +1035,15 @@ async def generate(self, rollout_id): row_count = metrics.get("rollout/train_batch_row_count") if row_count is None: raise RuntimeError("Expanded rollout did not report its transferred train row count.") - row_count = int(row_count) - if row_count <= 0: - raise RuntimeError(f"Expanded rollout produced an invalid train row count: {row_count}.") # This count controls TQ consumption only. global_batch_size # remains trajectory-based for GRPO loss normalization. - self._train_batch_row_counts[rollout_id] = row_count + self._train_batch_row_tracker.complete(rollout_id, row_count) elif self.args.partial_rollout and self.args.use_dynamic_global_batch_size: self._dynamic_global_batch_size = len(output.samples) * self.args.n_samples_per_prompt except Exception as exc: - if getattr(self.args, "custom_train_expanded_batch", False): - self._train_batch_row_errors[rollout_id] = f"{type(exc).__name__}: {exc}" + if expanded_batch: + self._train_batch_row_tracker.fail(rollout_id, exc) raise - finally: - if getattr(self.args, "custom_train_expanded_batch", False): - row_event.set() async def eval(self, rollout_id): self.health_monitoring_resume() diff --git a/relax/distributed/ray/utils.py b/relax/distributed/ray/utils.py index 7f798fa73..75a5f5e93 100644 --- a/relax/distributed/ray/utils.py +++ b/relax/distributed/ray/utils.py @@ -1,4 +1,6 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. # Adapted from https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/trainer/ray/utils.py#L1 +import asyncio import os import ray @@ -26,6 +28,56 @@ ] +class TrainBatchRowCountTracker: + """Coordinate data-dependent train row counts across rollout steps. + + A custom converter can produce a different number of rows per rollout. + Actor training may ask for step N before rollout generation for N has + finished, so counts must be keyed by rollout ID and waiters must be woken + on both success and failure. + """ + + def __init__(self, retained_rollouts: int = 2): + self._retained_rollouts = retained_rollouts + self._counts: dict[int, int] = {} + self._errors: dict[int, str] = {} + self._events: dict[int, asyncio.Event] = {} + + def start(self, rollout_id: int) -> None: + """Reset one rollout without replacing an event held by a waiter.""" + event = self._events.setdefault(rollout_id, asyncio.Event()) + event.clear() + self._counts.pop(rollout_id, None) + self._errors.pop(rollout_id, None) + for stale_rollout_id in [key for key in self._events if key < rollout_id - self._retained_rollouts]: + self._events.pop(stale_rollout_id, None) + self._counts.pop(stale_rollout_id, None) + self._errors.pop(stale_rollout_id, None) + + def complete(self, rollout_id: int, row_count: int) -> None: + """Publish a successful producer-side count and wake consumers.""" + row_count = int(row_count) + if row_count <= 0: + raise ValueError(f"Expanded rollout produced an invalid train row count: {row_count}.") + self._counts[rollout_id] = row_count + self._events.setdefault(rollout_id, asyncio.Event()).set() + + def fail(self, rollout_id: int, error: BaseException) -> None: + """Publish a generation/conversion failure so waiters cannot hang.""" + self._errors[rollout_id] = f"{type(error).__name__}: {error}" + self._events.setdefault(rollout_id, asyncio.Event()).set() + + async def wait(self, rollout_id: int) -> int: + """Wait for the exact count associated with ``rollout_id``.""" + event = self._events.setdefault(rollout_id, asyncio.Event()) + await event.wait() + if error := self._errors.get(rollout_id): + raise RuntimeError(f"Failed to determine train rows for rollout_id={rollout_id}: {error}") + if rollout_id not in self._counts: + raise RuntimeError(f"Train row count is unavailable for rollout_id={rollout_id}.") + return self._counts[rollout_id] + + def ray_noset_visible_devices(env_vars=os.environ): return any(env_vars.get(env_var) for env_var in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST) diff --git a/tests/distributed/ray/test_train_batch_row_count_tracker.py b/tests/distributed/ray/test_train_batch_row_count_tracker.py new file mode 100644 index 000000000..cccb087db --- /dev/null +++ b/tests/distributed/ray/test_train_batch_row_count_tracker.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import asyncio + +import pytest + +from relax.distributed.ray.utils import TrainBatchRowCountTracker + + +@pytest.mark.asyncio +async def test_train_batch_row_tracker_waits_for_matching_rollout(): + tracker = TrainBatchRowCountTracker() + tracker.start(0) + tracker.complete(0, 24) + + # A completed previous step must not satisfy a waiter for the next step. + waiter = asyncio.create_task(tracker.wait(1)) + await asyncio.sleep(0) + assert waiter.done() is False + + # Starting after wait() must reuse the same event object. + tracker.start(1) + tracker.complete(1, 32) + assert await waiter == 32 + assert await tracker.wait(0) == 24 + + +@pytest.mark.asyncio +async def test_train_batch_row_tracker_propagates_failure_to_waiter(): + tracker = TrainBatchRowCountTracker() + tracker.start(7) + waiter = asyncio.create_task(tracker.wait(7)) + await asyncio.sleep(0) + tracker.fail(7, ValueError("converter failed")) + + with pytest.raises(RuntimeError, match="rollout_id=7.*converter failed"): + await waiter + + +def test_train_batch_row_tracker_rejects_non_positive_count(): + tracker = TrainBatchRowCountTracker() + tracker.start(3) + with pytest.raises(ValueError, match="invalid train row count"): + tracker.complete(3, 0) diff --git a/tests/examples/mem_agent/test_compare_results.py b/tests/examples/mem_agent/test_compare_results.py index 65b38584b..f2d4c129e 100644 --- a/tests/examples/mem_agent/test_compare_results.py +++ b/tests/examples/mem_agent/test_compare_results.py @@ -2,7 +2,7 @@ import pytest -from examples.mem_agent.compare_results import compare_pair +from examples.mem_agent.compare_results import compare_baseline, compare_pair, validate_compatible_summaries def test_compare_pair_uses_absolute_percentage_points(): @@ -12,3 +12,35 @@ def test_compare_pair_uses_absolute_percentage_points(): assert passed["absolute_gap_pp"] == pytest.approx(2.9) assert passed["passed"] is True assert failed["passed"] is False + + +def test_compare_baseline_requires_strict_improvement(): + improved = compare_baseline("50", {"sub_em_pct": 41.0}, {"sub_em_pct": 41.1}, "sub_em_pct") + tied = compare_baseline("50", {"sub_em_pct": 41.0}, {"sub_em_pct": 41.0}, "sub_em_pct") + + assert improved["improvement_pp"] == pytest.approx(0.1) + assert improved["passed"] is True + assert tied["passed"] is False + + +def test_comparator_rejects_control_variable_mismatch(): + common = { + "data_file": "eval_50.jsonl", + "mode": "recurrent", + "tokenizer": "frozen-tokenizer", + "temperature": 0.7, + "top_p": 0.95, + "sampling_count": 1, + "chunk_tokens": 2048, + "max_memory_tokens": 1024, + "max_final_tokens": 256, + "max_chunks": 64, + "max_input_tokens": 7936, + "server_max_model_len": 8192, + "total": 128, + } + mismatched = {**common, "chunk_tokens": 4096} + + validate_compatible_summaries(common, dict(common)) + with pytest.raises(ValueError, match="chunk_tokens"): + validate_compatible_summaries(common, mismatched) diff --git a/tests/examples/mem_agent/test_convert.py b/tests/examples/mem_agent/test_convert.py index 7251c6554..681bb7766 100644 --- a/tests/examples/mem_agent/test_convert.py +++ b/tests/examples/mem_agent/test_convert.py @@ -11,10 +11,10 @@ from relax.utils.types import Sample -def _turn(turn_index: int) -> dict: +def _turn(turn_index: int, kind: str) -> dict: return { "turn_index": turn_index, - "kind": "final" if turn_index == 2 else "memory", + "kind": kind, "tokens": [100 + turn_index, 200 + turn_index, 201 + turn_index], "response_length": 2, "loss_mask": [1, 1], @@ -29,7 +29,12 @@ def _sample(index: int, reward: float, turn_count: int) -> Sample: group_index=9, reward={"score": reward}, status=Sample.Status.COMPLETED, - train_metadata={"mem_agent_turns": [_turn(turn_index) for turn_index in range(turn_count)]}, + train_metadata={ + "mem_agent_turns": [ + _turn(turn_index, "final" if turn_index == turn_count - 1 else "memory") + for turn_index in range(turn_count) + ] + }, ) @@ -43,6 +48,8 @@ def _args(credit_assignment="split", debug_train_only=True): n_samples_per_prompt=2, grpo_std_normalization=True, mem_agent_credit_assignment=credit_assignment, + mem_agent_max_memory_tokens=1024, + mem_agent_max_final_tokens=256, debug_train_only=debug_train_only, ) @@ -80,6 +87,11 @@ def test_converter_rejects_misaligned_or_failed_trajectory(): with pytest.raises(ValueError, match="status=failed"): convert_samples(_args(), [failed, _sample(6, 1.0, 1)]) + aborted = _sample(7, 0.0, 1) + aborted.status = Sample.Status.ABORTED + with pytest.raises(ValueError, match="status=aborted"): + convert_samples(_args(), [aborted, _sample(8, 1.0, 1)]) + def test_converter_rejects_inconsistent_group_turn_counts_and_non_divisible_rows(): with pytest.raises(ValueError, match="inconsistent turn counts"): @@ -89,3 +101,17 @@ def test_converter_rejects_inconsistent_group_turn_counts_and_non_divisible_rows args.mem_agent_train_rows_multiple = 4 with pytest.raises(ValueError, match="is not divisible"): convert_samples(args, [_sample(3, 0.0, 3), _sample(4, 1.0, 3)]) + + +def test_converter_rejects_structurally_invalid_or_overlong_turns(): + malformed = _sample(3, 0.0, 2) + malformed.train_metadata["mem_agent_turns"][-1]["kind"] = "memory" + with pytest.raises(ValueError, match="exactly one final turn"): + convert_samples(_args(), [malformed, _sample(4, 1.0, 2)]) + + overlong = _sample(5, 0.0, 2) + overlong.train_metadata["mem_agent_turns"][0]["response_length"] = 2 + args = _args() + args.mem_agent_max_memory_tokens = 1 + with pytest.raises(ValueError, match="response_length=2 exceeds 1"): + convert_samples(args, [overlong, _sample(6, 1.0, 2)]) diff --git a/tests/examples/mem_agent/test_data_and_metrics.py b/tests/examples/mem_agent/test_data_and_metrics.py index f0090a56f..f6b4eb34f 100644 --- a/tests/examples/mem_agent/test_data_and_metrics.py +++ b/tests/examples/mem_agent/test_data_and_metrics.py @@ -3,11 +3,19 @@ from __future__ import annotations import json +import sys +from types import ModuleType, SimpleNamespace import pytest from examples.mem_agent.metrics import aggregate, exact_match, f1_score, sub_exact_match -from examples.mem_agent.prepare_data import convert_file, convert_row, update_manifest +from examples.mem_agent.prepare_data import ( + convert_file, + convert_row, + download_hf_file, + read_rows, + update_manifest, +) def test_convert_row_supports_training_and_ruler_formats(): @@ -47,6 +55,45 @@ def test_convert_file_and_manifest_are_deterministic(tmp_path): assert json.loads(output.read_text(encoding="utf-8"))["metadata"]["ground_truth"] == ["Paris"] +def test_parquet_reader_and_huggingface_download_are_pinned_and_normalizable(monkeypatch, tmp_path): + parquet_row = { + "prompt": [{"role": "user", "content": "Who?"}], + "context": "Document", + "reward_model": {"ground_truth": ["Alice"]}, + } + + class FakeFrame: + def to_dict(self, orient): + assert orient == "records" + return [parquet_row] + + monkeypatch.setitem( + sys.modules, + "pandas", + SimpleNamespace(read_parquet=lambda path: FakeFrame()), + ) + assert convert_row(list(read_rows(tmp_path / "train.parquet"))[0])["label"] == "Alice" + + captured = {} + fake_hub = ModuleType("huggingface_hub") + + def fake_download(**kwargs): + captured.update(kwargs) + return str(tmp_path / kwargs["filename"]) + + fake_hub.hf_hub_download = fake_download + monkeypatch.setitem(sys.modules, "huggingface_hub", fake_hub) + path = download_hf_file("dataset/id", "fixed-revision", "train.parquet", tmp_path / "cache") + assert path == tmp_path / "train.parquet" + assert captured == { + "repo_id": "dataset/id", + "repo_type": "dataset", + "revision": "fixed-revision", + "filename": "train.parquet", + "cache_dir": str(tmp_path / "cache"), + } + + def test_ruler_metrics_match_expected_semantics(): assert exact_match("The Eiffel Tower", "Eiffel Tower") == 1.0 assert sub_exact_match("located in Paris France", "Paris") == 1.0 diff --git a/tests/examples/mem_agent/test_eval.py b/tests/examples/mem_agent/test_eval.py index aff2f3595..d05f3f523 100644 --- a/tests/examples/mem_agent/test_eval.py +++ b/tests/examples/mem_agent/test_eval.py @@ -6,7 +6,7 @@ import pytest -from examples.mem_agent.eval_ruler_hqa import base_infer +from examples.mem_agent.eval_ruler_hqa import base_infer, run_evaluation class CharacterTokenizer: @@ -45,3 +45,41 @@ async def fake_chat_once(session, base_url, api_key, model, instruction, tempera assert diagnostics["context_truncated"] is True assert "Question: Which answer?" in captured["instruction"] assert len(captured["instruction"]) <= args.max_input_tokens + + +@pytest.mark.asyncio +async def test_evaluation_error_keeps_ground_truth_and_counts_as_zero(monkeypatch, tmp_path): + async def failed_infer(item, args, tokenizer, session): + del item, args, tokenizer, session + raise RuntimeError("server unavailable") + + monkeypatch.setattr("examples.mem_agent.eval_ruler_hqa.recurrent_infer", failed_infer) + args = SimpleNamespace( + mode="recurrent", + concurrency=1, + timeout=1, + model="model", + tokenizer="tokenizer", + data_file=tmp_path / "eval.jsonl", + temperature=0.7, + top_p=0.95, + chunk_tokens=2048, + max_memory_tokens=1024, + max_final_tokens=256, + max_chunks=64, + max_input_tokens=7936, + server_max_model_len=8192, + ) + records, summary = await run_evaluation( + [{"_id": "q1", "input": "Question", "context": "Context", "answers": ["A", "Alias"]}], + args, + CharacterTokenizer(), + ) + + assert records[0]["answers"] == ["A", "Alias"] + assert records[0]["pred"] == "" + assert records[0]["judge_boxed_em"] == 0.0 + assert "server unavailable" in records[0]["error"] + assert summary["total"] == 1 + assert summary["errors"] == 1 + assert summary["sub_em_pct"] == 0.0 diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index 33cfbccf6..cc5d531d0 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -36,3 +36,26 @@ def test_train_script_keeps_trajectory_loss_and_real_turn_context_envelope(): assert "--custom-convert-samples-to-train-data-path examples.mem_agent.convert.convert_samples" in script assert "--use-dynamic-global-batch-size" not in script assert "--calculate-per-token-loss" not in script + + +def test_pipeline_and_acceptance_scripts_cover_required_stages_and_metrics(): + pipeline = (EXAMPLE / "run-pipeline.sh").read_text(encoding="utf-8") + paired = (EXAMPLE / "run-paired-eval.sh").read_text(encoding="utf-8") + evaluator = (EXAMPLE / "run-eval.sh").read_text(encoding="utf-8") + + for stage in ("prepare-data.sh", "run-qwen3-4B-train.sh", "convert-to-hf.sh", "run-eval.sh"): + assert stage in pipeline + assert 'BASE_MODEL_PATH="${BASE_MODEL_PATH:?' in paired + assert paired.count("MODE=recurrent") == 3 + assert '--pair "ruler-hqa-${length}"' in paired + assert '--baseline-pair "ruler-hqa-${length}" sub_em_pct' in paired + assert "--baseline-pair hotpotqa-dev boxed_em_pct" in paired + assert 'LENGTHS="${LENGTHS:-50 200 800}"' in paired + for value in ( + "--temperature 0.7", + "--top-p 0.95", + "--chunk-tokens 2048", + "--max-memory-tokens 1024", + '--server-max-model-len "${MAX_MODEL_LEN}"', + ): + assert value in evaluator diff --git a/tests/examples/mem_agent/test_rollout.py b/tests/examples/mem_agent/test_rollout.py index f62c3d225..7939b14c0 100644 --- a/tests/examples/mem_agent/test_rollout.py +++ b/tests/examples/mem_agent/test_rollout.py @@ -45,6 +45,7 @@ async def test_generate_trajectory_overwrites_memory_and_expands_every_turn(): async def fake_generate(args, turn, sampling_params, evaluation): del args + assert turn.session_id == "s" prompts.append(turn.prompt) max_new_tokens.append(sampling_params["max_new_tokens"]) response = responses.pop(0) @@ -69,6 +70,7 @@ async def fake_generate(args, turn, sampling_params, evaluation): assert result.metadata["num_chunks"] == 2 assert result.metadata["context_truncated"] is False assert result.metadata["memory_token_lengths"] == [2, 2] + assert result.session_id == "s" assert len(result.train_metadata["mem_agent_turns"]) == 3 assert NO_MEMORY in prompts[0] assert "abc" in prompts[0] @@ -208,3 +210,27 @@ async def fake_sglang_generate(args, turn, sampling_params, evaluation): assert result.status == Sample.Status.COMPLETED assert result.response == r"\boxed{x}" assert [turn["kind"] for turn in result.train_metadata["mem_agent_turns"]] == ["memory", "final"] + + +@pytest.mark.asyncio +async def test_generate_trajectory_rejects_response_longer_than_turn_limit(): + tokenizer = FakeTokenizer() + + async def oversized_generate(args, turn, sampling_params, evaluation): + del args, sampling_params, evaluation + response_ids = tokenizer.encode("TOO-LONG") + turn.response = "TOO-LONG" + turn.tokens = tokenizer.encode(turn.prompt) + response_ids + turn.rollout_tokens = list(turn.tokens) + turn.response_length = len(response_ids) + turn.loss_mask = [1] * len(response_ids) + turn.rollout_log_probs = [-0.1] * len(response_ids) + turn.status = Sample.Status.COMPLETED + return turn + + args = _args() + args.mem_agent_max_memory_tokens = 3 + sample = Sample(index=0, group_index=0, prompt="Q", metadata={"context": "abc"}) + + with pytest.raises(ValueError, match="exceeding its limit 3"): + await generate_trajectory(args, sample, {}, tokenizer, generator=oversized_generate) diff --git a/tests/utils/test_custom_sample_converter.py b/tests/utils/test_custom_sample_converter.py index 0804e251d..7248ba7d8 100644 --- a/tests/utils/test_custom_sample_converter.py +++ b/tests/utils/test_custom_sample_converter.py @@ -53,6 +53,8 @@ def _args(**overrides): "grpo_std_normalization": True, "n_samples_per_prompt": 1, "mem_agent_credit_assignment": "split", + "mem_agent_max_memory_tokens": 1024, + "mem_agent_max_final_tokens": 256, "debug_train_only": False, "load_debug_rollout_data_subsample": None, "use_dynamic_global_batch_size": False, From 1d74271e101be47fbc8d311dc5b283240ececa32 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 14:41:31 +0800 Subject: [PATCH 03/16] fix: tighten MemAgent reproduction evidence --- examples/mem_agent/README.md | 2 +- examples/mem_agent/compare_results.py | 18 ++++++++++- examples/mem_agent/eval_ruler_hqa.py | 26 ++++++++++++++-- examples/mem_agent/prompts.py | 31 +++++++++++++++---- .../mem_agent/test_compare_results.py | 12 +++++++ tests/examples/mem_agent/test_eval.py | 6 +++- tests/examples/mem_agent/test_rollout.py | 14 ++++++++- 7 files changed, 97 insertions(+), 12 deletions(-) diff --git a/examples/mem_agent/README.md b/examples/mem_agent/README.md index b9de3beff..0209828ce 100644 --- a/examples/mem_agent/README.md +++ b/examples/mem_agent/README.md @@ -48,7 +48,7 @@ RESULTS_DIR=/data/results/mem-agent-relax \ bash examples/mem_agent/run-eval.sh ``` -The evaluator writes raw per-sample JSONL and a summary JSON for HotpotQA dev and RULER-HQA 50/200/800. Failed requests keep their ground truth in the raw file and remain in the denominator with score zero. `boxed_em_pct` is the HotpotQA reward-compatible accuracy and `sub_em_pct` is the primary VIME-compatible RULER-HQA metric. Set `MODE=base` to run the optional single-context diagnostic; its context truncation always preserves the question and answer instruction. +The evaluator writes raw per-sample JSONL and a summary JSON for HotpotQA dev and RULER-HQA 50/200/800. Failed requests keep their ground truth in the raw file and remain in the denominator with score zero. Formal comparison additionally rejects empty runs and any run with request errors. Each summary records the normalized input file SHA-256 and evaluator schema version, so equal paths with different bytes or incompatible evaluator revisions cannot be compared. `boxed_em_pct` is the HotpotQA reward-compatible accuracy and `sub_em_pct` is the primary VIME-compatible RULER-HQA metric. Set `MODE=base` to run the optional single-context diagnostic; its context truncation always preserves the question and answer instruction. `TOKENIZER_PATH` should point to the frozen base snapshot. `run-pipeline.sh` preserves it automatically before switching `MODEL_PATH` to the converted checkpoint. When `NUM_ROLLOUT=2` is used, the pipeline also selects `iter_0000001` automatically instead of the 100-step default `iter_0000099`. diff --git a/examples/mem_agent/compare_results.py b/examples/mem_agent/compare_results.py index 2b476dd75..67c5a004d 100644 --- a/examples/mem_agent/compare_results.py +++ b/examples/mem_agent/compare_results.py @@ -11,6 +11,8 @@ COMPATIBILITY_FIELDS = ( "data_file", + "data_sha256", + "evaluator_schema_version", "mode", "tokenizer", "temperature", @@ -24,12 +26,26 @@ "server_max_model_len", "total", ) +COMPLETENESS_FIELDS = ("successful", "errors") def validate_compatible_summaries(*summaries: dict[str, Any]) -> None: - """Reject a comparison when any controlled evaluation field differs.""" + """Reject incomplete runs or mismatched controlled evaluation fields.""" if len(summaries) < 2: raise ValueError("At least two summaries are required for compatibility validation.") + for summary in summaries: + for field in COMPLETENESS_FIELDS: + if field not in summary: + raise KeyError(f"Completeness field {field!r} must exist in every summary.") + total = int(summary.get("total", 0)) + successful = int(summary["successful"]) + errors = int(summary["errors"]) + # A request error is retained as a zero-score row for diagnosis, but + # an effects claim must come from a non-empty, fully completed run. + if total <= 0 or errors != 0 or successful != total: + raise ValueError( + f"Evaluation summary is incomplete: total={total}, successful={successful}, errors={errors}." + ) for field in COMPATIBILITY_FIELDS: if any(field not in summary for summary in summaries): raise KeyError(f"Compatibility field {field!r} must exist in every summary.") diff --git a/examples/mem_agent/eval_ruler_hqa.py b/examples/mem_agent/eval_ruler_hqa.py index ad40d7a53..2561f52cc 100644 --- a/examples/mem_agent/eval_ruler_hqa.py +++ b/examples/mem_agent/eval_ruler_hqa.py @@ -6,6 +6,7 @@ import argparse import asyncio +import hashlib import json from pathlib import Path from typing import Any @@ -23,6 +24,18 @@ from examples.mem_agent.reward import exact_match_any, extract_last_boxed +EVALUATOR_SCHEMA_VERSION = "mem-agent-vime-eval-v1" + + +def sha256_file(path: Path) -> str: + """Hash the exact normalized dataset consumed by one evaluation run.""" + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + def load_data(path: Path) -> list[dict[str, Any]]: if path.suffix == ".jsonl": with path.open(encoding="utf-8") as source: @@ -103,7 +116,7 @@ async def recurrent_infer( args.base_url, args.api_key, args.model, - memory_instruction(question, memory, chunk), + memory_instruction(question, memory, chunk, evaluation=True), args.temperature, args.top_p, args.max_memory_tokens, @@ -118,7 +131,7 @@ async def recurrent_infer( args.base_url, args.api_key, args.model, - final_instruction(question, memory), + final_instruction(question, memory, evaluation=True), args.temperature, args.top_p, args.max_final_tokens, @@ -166,6 +179,10 @@ async def base_infer( async def run_evaluation( data: list[dict[str, Any]], args: argparse.Namespace, tokenizer: Any ) -> tuple[list[dict[str, Any]], dict[str, Any]]: + # Capture the digest before any long-running inference. A later overwrite + # of the path must not make the summary claim that different bytes were + # evaluated. + data_sha256 = sha256_file(args.data_file) semaphore = asyncio.Semaphore(args.concurrency) timeout = aiohttp.ClientTimeout(total=args.timeout) @@ -222,6 +239,11 @@ async def evaluate_one(item: dict[str, Any]) -> dict[str, Any]: "model": args.model, "tokenizer": args.tokenizer, "data_file": str(args.data_file), + # Path equality alone is not proof that sequential checkpoint runs saw + # identical bytes. The comparator requires this digest and evaluator + # schema before it accepts a controlled comparison. + "data_sha256": data_sha256, + "evaluator_schema_version": EVALUATOR_SCHEMA_VERSION, "temperature": args.temperature, "top_p": args.top_p, "sampling_count": 1, diff --git a/examples/mem_agent/prompts.py b/examples/mem_agent/prompts.py index d19bc784e..b4ba11dcb 100644 --- a/examples/mem_agent/prompts.py +++ b/examples/mem_agent/prompts.py @@ -11,7 +11,7 @@ MEMORY_TEMPLATE = """You are presented with a problem, a section of an article that may contain the answer to the problem, and a previous memory. Please read the provided section carefully and update the memory with the new information that helps to answer the problem. Be sure to retain all relevant details from the previous memory while adding any new, useful information. - +{problem_tag_suffix} {question} @@ -28,7 +28,7 @@ FINAL_TEMPLATE = """You are presented with a problem and a previous memory. Please answer the problem based on the previous memory and put the answer in \\boxed{{}}. - +{problem_tag_suffix} {question} @@ -63,12 +63,31 @@ def truncate_text_to_tokens(tokenizer: Any, text: str, max_tokens: int) -> tuple return bounded_text, len(bounded_ids) -def memory_instruction(question: str, memory: str, chunk: str) -> str: - return MEMORY_TEMPLATE.format(question=question, memory=memory, chunk=chunk) +def _problem_tag_suffix(evaluation: bool) -> str: + """Preserve the one-character prompt difference in fixed VIME code. + VIME training has `` `` while its evaluator has ````. + Keeping the variants explicit makes both reproduction paths byte-aligned + instead of silently choosing one template for both. + """ + return "" if evaluation else " " -def final_instruction(question: str, memory: str) -> str: - return FINAL_TEMPLATE.format(question=question, memory=memory) + +def memory_instruction(question: str, memory: str, chunk: str, *, evaluation: bool = False) -> str: + return MEMORY_TEMPLATE.format( + question=question, + memory=memory, + chunk=chunk, + problem_tag_suffix=_problem_tag_suffix(evaluation), + ) + + +def final_instruction(question: str, memory: str, *, evaluation: bool = False) -> str: + return FINAL_TEMPLATE.format( + question=question, + memory=memory, + problem_tag_suffix=_problem_tag_suffix(evaluation), + ) def render_chat_prompt(tokenizer: Any, instruction: str) -> str: diff --git a/tests/examples/mem_agent/test_compare_results.py b/tests/examples/mem_agent/test_compare_results.py index f2d4c129e..44fa9b1ca 100644 --- a/tests/examples/mem_agent/test_compare_results.py +++ b/tests/examples/mem_agent/test_compare_results.py @@ -26,6 +26,8 @@ def test_compare_baseline_requires_strict_improvement(): def test_comparator_rejects_control_variable_mismatch(): common = { "data_file": "eval_50.jsonl", + "data_sha256": "a" * 64, + "evaluator_schema_version": "mem-agent-vime-eval-v1", "mode": "recurrent", "tokenizer": "frozen-tokenizer", "temperature": 0.7, @@ -38,9 +40,19 @@ def test_comparator_rejects_control_variable_mismatch(): "max_input_tokens": 7936, "server_max_model_len": 8192, "total": 128, + "successful": 128, + "errors": 0, } mismatched = {**common, "chunk_tokens": 4096} validate_compatible_summaries(common, dict(common)) with pytest.raises(ValueError, match="chunk_tokens"): validate_compatible_summaries(common, mismatched) + + changed_data = {**common, "data_sha256": "b" * 64} + with pytest.raises(ValueError, match="data_sha256"): + validate_compatible_summaries(common, changed_data) + + incomplete = {**common, "successful": 127, "errors": 1} + with pytest.raises(ValueError, match="incomplete"): + validate_compatible_summaries(common, incomplete) diff --git a/tests/examples/mem_agent/test_eval.py b/tests/examples/mem_agent/test_eval.py index d05f3f523..0056e50aa 100644 --- a/tests/examples/mem_agent/test_eval.py +++ b/tests/examples/mem_agent/test_eval.py @@ -54,13 +54,15 @@ async def failed_infer(item, args, tokenizer, session): raise RuntimeError("server unavailable") monkeypatch.setattr("examples.mem_agent.eval_ruler_hqa.recurrent_infer", failed_infer) + data_file = tmp_path / "eval.jsonl" + data_file.write_text('{"input":"Question","context":"Context","answers":["A"]}\n', encoding="utf-8") args = SimpleNamespace( mode="recurrent", concurrency=1, timeout=1, model="model", tokenizer="tokenizer", - data_file=tmp_path / "eval.jsonl", + data_file=data_file, temperature=0.7, top_p=0.95, chunk_tokens=2048, @@ -83,3 +85,5 @@ async def failed_infer(item, args, tokenizer, session): assert summary["total"] == 1 assert summary["errors"] == 1 assert summary["sub_em_pct"] == 0.0 + assert len(summary["data_sha256"]) == 64 + assert summary["evaluator_schema_version"] == "mem-agent-vime-eval-v1" diff --git a/tests/examples/mem_agent/test_rollout.py b/tests/examples/mem_agent/test_rollout.py index 7939b14c0..c13b74605 100644 --- a/tests/examples/mem_agent/test_rollout.py +++ b/tests/examples/mem_agent/test_rollout.py @@ -7,7 +7,7 @@ import pytest -from examples.mem_agent.prompts import NO_MEMORY, truncate_text_to_tokens +from examples.mem_agent.prompts import NO_MEMORY, final_instruction, memory_instruction, truncate_text_to_tokens from examples.mem_agent.rollout import chunk_context, generate, generate_trajectory from relax.utils.types import Sample @@ -95,6 +95,18 @@ def test_chunk_context_preserves_boundaries_and_marks_truncation(): assert truncated is True +def test_prompt_variants_match_fixed_vime_training_and_evaluation_templates(): + train_memory = memory_instruction("Q", "M", "C") + eval_memory = memory_instruction("Q", "M", "C", evaluation=True) + train_final = final_instruction("Q", "M") + eval_final = final_instruction("Q", "M", evaluation=True) + + assert " \nQ\n" in train_memory + assert "\nQ\n" in eval_memory + assert " \nQ\n" in train_final + assert "\nQ\n" in eval_final + + def test_truncate_text_to_tokens_retokenizes_to_the_hard_limit(): text, length = truncate_text_to_tokens(FakeTokenizer(), "MEMORY", 3) assert text == "MEM" From 048bb0182098416b227777ade66082b30505bda5 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 15:12:30 +0800 Subject: [PATCH 04/16] fix: complete MemAgent pre-GPU evidence --- examples/mem_agent/README.md | 4 ++- examples/mem_agent/reward.py | 5 ++++ examples/mem_agent/rollout.py | 7 +++-- examples/mem_agent/run-eval.sh | 3 ++ .../mem_agent/test_recipe_contract.py | 6 ++++ tests/examples/mem_agent/test_reward.py | 9 ++++++ tests/examples/mem_agent/test_rollout.py | 29 +++++++++++++++++++ 7 files changed, 60 insertions(+), 3 deletions(-) diff --git a/examples/mem_agent/README.md b/examples/mem_agent/README.md index 0209828ce..ff712f969 100644 --- a/examples/mem_agent/README.md +++ b/examples/mem_agent/README.md @@ -2,6 +2,8 @@ This example trains Qwen3-4B to update a bounded textual memory while reading a long document chunk by chunk. Every memory-update turn and the final-answer turn is saved as an independent training row. Only the final boxed answer receives a rule-based reward; GRPO normalization happens before the trajectory is expanded. +The training log reports the trajectory-level 0/1 outcome as `rollout/mem_agent_raw_reward/mean` on every rollout step. This diagnostic mirrors the primary `score` exactly but is not consumed by GRPO, making first/last-window reward comparisons auditable without changing optimization. + The reproducibility contract is frozen to: - model: `Qwen/Qwen3-4B@1cfa9a7208912126459214e8b04321603b3df60c`; @@ -48,7 +50,7 @@ RESULTS_DIR=/data/results/mem-agent-relax \ bash examples/mem_agent/run-eval.sh ``` -The evaluator writes raw per-sample JSONL and a summary JSON for HotpotQA dev and RULER-HQA 50/200/800. Failed requests keep their ground truth in the raw file and remain in the denominator with score zero. Formal comparison additionally rejects empty runs and any run with request errors. Each summary records the normalized input file SHA-256 and evaluator schema version, so equal paths with different bytes or incompatible evaluator revisions cannot be compared. `boxed_em_pct` is the HotpotQA reward-compatible accuracy and `sub_em_pct` is the primary VIME-compatible RULER-HQA metric. Set `MODE=base` to run the optional single-context diagnostic; its context truncation always preserves the question and answer instruction. +The evaluator writes raw per-sample JSONL and a summary JSON for HotpotQA dev and RULER-HQA 50/200/800. Its 64-chunk limit is the effective value of fixed VIME's official `run-eval.sh`: that script sources `_common.sh`, which exports `MEM_MAX_CHUNKS=64`, even though the Python evaluator alone has a 512 fallback. Failed requests keep their ground truth in the raw file and remain in the denominator with score zero. Formal comparison additionally rejects empty runs and any run with request errors. Each summary records the normalized input file SHA-256 and evaluator schema version, so equal paths with different bytes or incompatible evaluator revisions cannot be compared. `boxed_em_pct` is the HotpotQA reward-compatible accuracy and `sub_em_pct` is the primary VIME-compatible RULER-HQA metric. Set `MODE=base` to run the optional single-context diagnostic; its context truncation always preserves the question and answer instruction. `TOKENIZER_PATH` should point to the frozen base snapshot. `run-pipeline.sh` preserves it automatically before switching `MODEL_PATH` to the converted checkpoint. When `NUM_ROLLOUT=2` is used, the pipeline also selects `iter_0000001` automatically instead of the 100-step default `iter_0000099`. diff --git a/examples/mem_agent/reward.py b/examples/mem_agent/reward.py index 87d90b46d..f75790db8 100644 --- a/examples/mem_agent/reward.py +++ b/examples/mem_agent/reward.py @@ -51,6 +51,11 @@ async def reward_func(args: Any, sample: Any, **kwargs: Any) -> dict[str, Any]: score = float(bool(prediction) and exact_match_any(prediction, list(ground_truths))) return { "score": score, + # ReLax intentionally excludes the primary reward key from auxiliary + # metric aggregation. Mirror the same 0/1 value under a diagnostic key + # so every rollout step logs rollout/mem_agent_raw_reward/mean without + # changing the score used by GRPO. + "mem_agent_raw_reward": score, "pred": prediction, "gt": str(ground_truths[0]) if ground_truths else "", "diagnostic": "matched" if score else ("missing_boxed" if not prediction else "answer_mismatch"), diff --git a/examples/mem_agent/rollout.py b/examples/mem_agent/rollout.py index 5bcdc4673..4419eead3 100644 --- a/examples/mem_agent/rollout.py +++ b/examples/mem_agent/rollout.py @@ -147,7 +147,10 @@ async def generate_trajectory( # memory text survives; prior prompts and token history are not appended. for chunk_ids in chunks: chunk = tokenizer.decode(chunk_ids, skip_special_tokens=True) - prompt = render_chat_prompt(tokenizer, memory_instruction(question, memory, chunk)) + prompt = render_chat_prompt( + tokenizer, + memory_instruction(question, memory, chunk, evaluation=evaluation), + ) turn_sample = await _run_turn( args, sample, @@ -176,7 +179,7 @@ async def generate_trajectory( # The final request deliberately excludes context/chunks. This enforces # the question + latest-memory information boundary from the task spec. - final_prompt = render_chat_prompt(tokenizer, final_instruction(question, memory)) + final_prompt = render_chat_prompt(tokenizer, final_instruction(question, memory, evaluation=evaluation)) final_sample = await _run_turn( args, sample, diff --git a/examples/mem_agent/run-eval.sh b/examples/mem_agent/run-eval.sh index c5e075e57..8c7644a8b 100755 --- a/examples/mem_agent/run-eval.sh +++ b/examples/mem_agent/run-eval.sh @@ -47,6 +47,9 @@ curl -fsS "http://${SERVE_HOST}:${SERVE_PORT}/v1/models" >/dev/null run_eval() { local data_file="$1" local suffix="$2" + # Fixed VIME's Python evaluator has a 512 fallback, but its official + # run-eval.sh sources _common.sh, which exports MEM_MAX_CHUNKS=64. Pin the + # effective official-run value explicitly instead of relying on inheritance. python3 "${SCRIPT_DIR}/eval_ruler_hqa.py" \ --data-file "${data_file}" \ --model "${MODEL_PATH}" \ diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index cc5d531d0..4357716f5 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -38,6 +38,11 @@ def test_train_script_keeps_trajectory_loss_and_real_turn_context_envelope(): assert "--calculate-per-token-loss" not in script +def test_reward_exposes_a_step_level_raw_reward_metric(): + reward_source = (EXAMPLE / "reward.py").read_text(encoding="utf-8") + assert '"mem_agent_raw_reward": score' in reward_source + + def test_pipeline_and_acceptance_scripts_cover_required_stages_and_metrics(): pipeline = (EXAMPLE / "run-pipeline.sh").read_text(encoding="utf-8") paired = (EXAMPLE / "run-paired-eval.sh").read_text(encoding="utf-8") @@ -56,6 +61,7 @@ def test_pipeline_and_acceptance_scripts_cover_required_stages_and_metrics(): "--top-p 0.95", "--chunk-tokens 2048", "--max-memory-tokens 1024", + "--max-chunks 64", '--server-max-model-len "${MAX_MODEL_LEN}"', ): assert value in evaluator diff --git a/tests/examples/mem_agent/test_reward.py b/tests/examples/mem_agent/test_reward.py index 5e3fd4e08..37fbf9de8 100644 --- a/tests/examples/mem_agent/test_reward.py +++ b/tests/examples/mem_agent/test_reward.py @@ -7,6 +7,7 @@ import pytest from examples.mem_agent.reward import extract_last_boxed, normalize_answer, reward_func +from relax.utils.metrics.metric_utils import compute_rollout_explicit_reward_metrics from relax.utils.types import Sample @@ -29,9 +30,17 @@ async def test_reward_scores_only_final_output_against_all_ground_truths(): ) result = await reward_func(SimpleNamespace(), sample) assert result["score"] == 1.0 + assert result["mem_agent_raw_reward"] == 1.0 assert result["pred"] == "The Eiffel Tower" assert result["diagnostic"] == "matched" + sample.reward = result + metrics = compute_rollout_explicit_reward_metrics( + SimpleNamespace(reward_key="score", log_passrate=False, n_samples_per_prompt=1), + [sample], + ) + assert metrics["mem_agent_raw_reward/mean"] == 1.0 + @pytest.mark.asyncio async def test_reward_reports_missing_boxed_as_zero(): diff --git a/tests/examples/mem_agent/test_rollout.py b/tests/examples/mem_agent/test_rollout.py index c13b74605..d849a9fef 100644 --- a/tests/examples/mem_agent/test_rollout.py +++ b/tests/examples/mem_agent/test_rollout.py @@ -107,6 +107,35 @@ def test_prompt_variants_match_fixed_vime_training_and_evaluation_templates(): assert "\nQ\n" in eval_final +@pytest.mark.asyncio +async def test_generate_trajectory_uses_evaluation_prompt_variant(): + tokenizer = FakeTokenizer() + responses = iter(["MEM", r"\boxed{x}"]) + prompts = [] + + async def fake_generate(args, turn, sampling_params, evaluation): + del args, sampling_params + assert evaluation is True + prompts.append(turn.prompt) + response = next(responses) + response_ids = tokenizer.encode(response) + turn.response = response + turn.tokens = tokenizer.encode(turn.prompt) + response_ids + turn.rollout_tokens = list(turn.tokens) + turn.response_length = len(response_ids) + turn.loss_mask = [1] * len(response_ids) + turn.rollout_log_probs = [] + turn.status = Sample.Status.COMPLETED + return turn + + sample = Sample(index=0, group_index=0, prompt="Q", metadata={"context": "abc"}) + result = await generate_trajectory(_args(), sample, {}, tokenizer, generator=fake_generate, evaluation=True) + + assert result.status == Sample.Status.COMPLETED + assert all("\nQ\n" in prompt for prompt in prompts) + assert all(" \nQ\n" not in prompt for prompt in prompts) + + def test_truncate_text_to_tokens_retokenizes_to_the_hard_limit(): text, length = truncate_text_to_tokens(FakeTokenizer(), "MEMORY", 3) assert text == "MEM" From 92a15824011a680957694b4b21a7976d580731de Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 15:33:58 +0800 Subject: [PATCH 05/16] feat: summarize MemAgent training rewards --- examples/mem_agent/README.md | 2 +- examples/mem_agent/run-pipeline.sh | 11 +- examples/mem_agent/summarize_reward.py | 134 ++++++++++++++++++ .../mem_agent/test_recipe_contract.py | 10 +- .../mem_agent/test_summarize_reward.py | 46 ++++++ 5 files changed, 200 insertions(+), 3 deletions(-) create mode 100644 examples/mem_agent/summarize_reward.py create mode 100644 tests/examples/mem_agent/test_summarize_reward.py diff --git a/examples/mem_agent/README.md b/examples/mem_agent/README.md index ff712f969..a86032d75 100644 --- a/examples/mem_agent/README.md +++ b/examples/mem_agent/README.md @@ -2,7 +2,7 @@ This example trains Qwen3-4B to update a bounded textual memory while reading a long document chunk by chunk. Every memory-update turn and the final-answer turn is saved as an independent training row. Only the final boxed answer receives a rule-based reward; GRPO normalization happens before the trajectory is expanded. -The training log reports the trajectory-level 0/1 outcome as `rollout/mem_agent_raw_reward/mean` on every rollout step. This diagnostic mirrors the primary `score` exactly but is not consumed by GRPO, making first/last-window reward comparisons auditable without changing optimization. +The training log reports the trajectory-level 0/1 outcome as `rollout/mem_agent_raw_reward/mean` on every rollout step. This diagnostic mirrors the primary `score` exactly but is not consumed by GRPO. `run-pipeline.sh` then runs `summarize_reward.py` and writes `training-reward.summary.json`, containing every raw point, the first/last-window means, their delta, and the peak. It rejects a run whose rollout ids are incomplete instead of producing a partial trend. The reproducibility contract is frozen to: diff --git a/examples/mem_agent/run-pipeline.sh b/examples/mem_agent/run-pipeline.sh index 26f7caca1..af3bb2b68 100755 --- a/examples/mem_agent/run-pipeline.sh +++ b/examples/mem_agent/run-pipeline.sh @@ -9,18 +9,27 @@ MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH.}" SAVE_DIR="${SAVE_DIR:?Set SAVE_DIR.}" RESULTS_DIR="${RESULTS_DIR:?Set RESULTS_DIR.}" NUM_ROLLOUT="${NUM_ROLLOUT:-100}" +RUN_NAME="${RUN_NAME:-mem-agent-qwen3-4b}" if ((NUM_ROLLOUT <= 0)); then echo "NUM_ROLLOUT must be positive." >&2 exit 1 fi -export DATA_DIR MODEL_PATH SAVE_DIR RESULTS_DIR NUM_ROLLOUT +export DATA_DIR MODEL_PATH SAVE_DIR RESULTS_DIR NUM_ROLLOUT RUN_NAME if [[ "${SKIP_PREPARE:-0}" != "1" ]]; then bash "${SCRIPT_DIR}/prepare-data.sh" fi bash "${SCRIPT_DIR}/run-qwen3-4B-train.sh" +# The rollout logger writes every raw reward mean to the Ray job text log. +# Materialize the complete per-step series and first/last-window statistics +# before conversion so a successful pipeline always leaves effect evidence. +python3 "${SCRIPT_DIR}/summarize_reward.py" \ + --log-file "${SCRIPT_DIR}/../../logs/${RUN_NAME}.log" \ + --expected-steps "${NUM_ROLLOUT}" \ + --output "${RESULTS_DIR}/training-reward.summary.json" + export CHECKPOINT_DIR="${SAVE_DIR}" # ReLax numbers checkpoints from zero, so a two-step pipeline produces # iter_0000001 while the frozen 100-step recipe produces iter_0000099. diff --git a/examples/mem_agent/summarize_reward.py b/examples/mem_agent/summarize_reward.py new file mode 100644 index 000000000..b25aca009 --- /dev/null +++ b/examples/mem_agent/summarize_reward.py @@ -0,0 +1,134 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Turn ReLax rollout reward logs into an auditable training summary.""" + +from __future__ import annotations + +import argparse +import ast +import json +import math +import re +from collections.abc import Iterable +from pathlib import Path +from statistics import fmean +from typing import Any + + +DEFAULT_METRIC = "rollout/mem_agent_raw_reward/mean" +SCHEMA_VERSION = "mem-agent-reward-summary-v1" +_PERF_LINE = re.compile(r"\bperf\s+(\d+):\s+(\{.*\})") + + +def extract_reward_points(lines: Iterable[str], metric: str = DEFAULT_METRIC) -> list[tuple[int, float]]: + """Extract one trajectory-level reward mean for every logged rollout. + + ReLax logs the complete rollout metric dictionary as ``perf N: {...}`` + before sending the same values to TensorBoard. Parsing that durable text + log keeps the acceptance artifact independent of a tracking backend. + Identical duplicated lines are tolerated because Ray may replay driver + output, while conflicting values for one rollout fail closed. + """ + by_rollout: dict[int, float] = {} + for line_number, line in enumerate(lines, start=1): + match = _PERF_LINE.search(line) + if match is None: + continue + try: + payload = ast.literal_eval(match.group(2)) + except (SyntaxError, ValueError) as exc: + raise ValueError(f"Malformed ReLax perf payload at line {line_number}.") from exc + if not isinstance(payload, dict) or metric not in payload: + continue + + rollout_id = int(match.group(1)) + raw_value = payload[metric] + if isinstance(raw_value, bool) or not isinstance(raw_value, (int, float)): + raise ValueError(f"Reward metric {metric!r} at rollout {rollout_id} is not numeric: {raw_value!r}.") + value = float(raw_value) + if not math.isfinite(value) or not 0.0 <= value <= 1.0: + raise ValueError(f"Reward metric {metric!r} at rollout {rollout_id} is outside [0, 1]: {value!r}.") + if rollout_id in by_rollout and by_rollout[rollout_id] != value: + raise ValueError( + f"Conflicting reward values for rollout {rollout_id}: {by_rollout[rollout_id]} and {value}." + ) + by_rollout[rollout_id] = value + + return sorted(by_rollout.items()) + + +def summarize_reward_points( + points: list[tuple[int, float]], + *, + expected_steps: int | None = None, + window_size: int = 10, + metric: str = DEFAULT_METRIC, +) -> dict[str, Any]: + """Summarize first/last windows without inventing an effect threshold.""" + if not points: + raise ValueError(f"No {metric!r} points were found in the training log.") + if window_size <= 0: + raise ValueError("window_size must be positive.") + + rollout_ids = [rollout_id for rollout_id, _ in points] + if len(set(rollout_ids)) != len(rollout_ids): + raise ValueError("Reward points contain duplicate rollout ids.") + if expected_steps is not None: + if expected_steps <= 0: + raise ValueError("expected_steps must be positive.") + expected_ids = list(range(expected_steps)) + if rollout_ids != expected_ids: + raise ValueError(f"Reward rollout ids are incomplete: expected {expected_ids}, got {rollout_ids}.") + + effective_window = min(window_size, len(points)) + first_values = [value for _, value in points[:effective_window]] + last_values = [value for _, value in points[-effective_window:]] + first_mean = fmean(first_values) + last_mean = fmean(last_values) + delta = last_mean - first_mean + peak_rollout_id, peak_reward = max(points, key=lambda item: (item[1], -item[0])) + return { + "schema_version": SCHEMA_VERSION, + "metric": metric, + "num_steps": len(points), + "window_size_requested": window_size, + "window_size_used": effective_window, + "first_window_mean": first_mean, + "last_window_mean": last_mean, + "last_minus_first": delta, + # "Clearly improved" has no frozen numeric margin in Task 36. Report + # the strict direction here and leave any later threshold explicit in + # the experiment record rather than silently choosing one in code. + "strictly_improved": delta > 0.0, + "peak_reward": peak_reward, + "peak_rollout_id": peak_rollout_id, + "points": [{"rollout_id": rollout_id, "reward": value} for rollout_id, value in points], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--log-file", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--metric", default=DEFAULT_METRIC) + parser.add_argument("--expected-steps", type=int) + parser.add_argument("--window-size", type=int, default=10) + args = parser.parse_args() + + with args.log_file.open(encoding="utf-8", errors="replace") as source: + points = extract_reward_points(source, metric=args.metric) + summary = summarize_reward_points( + points, + expected_steps=args.expected_steps, + window_size=args.window_size, + metric=args.metric, + ) + summary["log_file"] = str(args.log_file.resolve()) + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as destination: + json.dump(summary, destination, ensure_ascii=False, indent=2) + destination.write("\n") + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index 4357716f5..1fb3fbd66 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -48,8 +48,16 @@ def test_pipeline_and_acceptance_scripts_cover_required_stages_and_metrics(): paired = (EXAMPLE / "run-paired-eval.sh").read_text(encoding="utf-8") evaluator = (EXAMPLE / "run-eval.sh").read_text(encoding="utf-8") - for stage in ("prepare-data.sh", "run-qwen3-4B-train.sh", "convert-to-hf.sh", "run-eval.sh"): + for stage in ( + "prepare-data.sh", + "run-qwen3-4B-train.sh", + "summarize_reward.py", + "convert-to-hf.sh", + "run-eval.sh", + ): assert stage in pipeline + assert '--expected-steps "${NUM_ROLLOUT}"' in pipeline + assert "training-reward.summary.json" in pipeline assert 'BASE_MODEL_PATH="${BASE_MODEL_PATH:?' in paired assert paired.count("MODE=recurrent") == 3 assert '--pair "ruler-hqa-${length}"' in paired diff --git a/tests/examples/mem_agent/test_summarize_reward.py b/tests/examples/mem_agent/test_summarize_reward.py new file mode 100644 index 000000000..1820ec5c8 --- /dev/null +++ b/tests/examples/mem_agent/test_summarize_reward.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Tests for the durable MemAgent training-reward acceptance artifact.""" + +from __future__ import annotations + +import pytest + +from examples.mem_agent.summarize_reward import extract_reward_points, summarize_reward_points + + +def test_extract_and_summarize_complete_reward_series(): + lines = [ + "unrelated startup output\n", + "2026-08-04 | INFO | worker perf 0: {'rollout/response_len/mean': 4.0, " + "'rollout/mem_agent_raw_reward/mean': 0.25}\n", + "perf 1: {'rollout/mem_agent_raw_reward/mean': 0.5}\n", + "perf 2: {'rollout/mem_agent_raw_reward/mean': 0.75}\n", + ] + + points = extract_reward_points(lines) + summary = summarize_reward_points(points, expected_steps=3, window_size=1) + + assert points == [(0, 0.25), (1, 0.5), (2, 0.75)] + assert summary["first_window_mean"] == 0.25 + assert summary["last_window_mean"] == 0.75 + assert summary["last_minus_first"] == 0.5 + assert summary["strictly_improved"] is True + assert summary["peak_rollout_id"] == 2 + assert summary["points"][-1] == {"rollout_id": 2, "reward": 0.75} + + +def test_identical_replayed_line_is_deduplicated_but_conflict_fails(): + line = "perf 0: {'rollout/mem_agent_raw_reward/mean': 0.5}\n" + assert extract_reward_points([line, line]) == [(0, 0.5)] + + conflict = "perf 0: {'rollout/mem_agent_raw_reward/mean': 0.75}\n" + with pytest.raises(ValueError, match="Conflicting reward values"): + extract_reward_points([line, conflict]) + + +def test_summary_rejects_missing_rollout_and_invalid_reward(): + with pytest.raises(ValueError, match="incomplete"): + summarize_reward_points([(0, 0.0), (2, 1.0)], expected_steps=3) + + with pytest.raises(ValueError, match=r"outside \[0, 1\]"): + extract_reward_points(["perf 0: {'rollout/mem_agent_raw_reward/mean': 1.1}\n"]) From f2f53b9aa5f745d617f8281d4b83e3d4111664b3 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 18:21:48 +0800 Subject: [PATCH 06/16] feat: add Qwen3 0.6B MemAgent pilot --- examples/mem_agent/README.md | 29 ++ examples/mem_agent/compare_results.py | 6 +- .../mem_agent/config-pilot-qwen3-0.6b.yaml | 16 + examples/mem_agent/eval_ruler_hqa.py | 89 ++++- .../mem_agent/prepare-pilot-candidates.sh | 36 ++ examples/mem_agent/prepare_data.py | 52 ++- examples/mem_agent/prepare_pilot_data.py | 377 ++++++++++++++++++ examples/mem_agent/prompts.py | 20 +- examples/mem_agent/rollout.py | 8 +- examples/mem_agent/run-qwen3-0.6B-baseline.sh | 109 +++++ examples/mem_agent/run-qwen3-0.6B-eval.sh | 77 ++++ examples/mem_agent/run-qwen3-0.6B-train.sh | 149 +++++++ examples/mem_agent/summarize_reward.py | 63 +++ .../mem_agent/test_compare_results.py | 2 + .../mem_agent/test_data_and_metrics.py | 19 + tests/examples/mem_agent/test_eval.py | 60 ++- .../mem_agent/test_mock_integration.py | 3 +- tests/examples/mem_agent/test_pilot_data.py | 98 +++++ .../mem_agent/test_recipe_contract.py | 31 ++ tests/examples/mem_agent/test_rollout.py | 33 +- .../mem_agent/test_summarize_reward.py | 22 +- 21 files changed, 1271 insertions(+), 28 deletions(-) create mode 100644 examples/mem_agent/config-pilot-qwen3-0.6b.yaml create mode 100755 examples/mem_agent/prepare-pilot-candidates.sh create mode 100644 examples/mem_agent/prepare_pilot_data.py create mode 100755 examples/mem_agent/run-qwen3-0.6B-baseline.sh create mode 100755 examples/mem_agent/run-qwen3-0.6B-eval.sh create mode 100755 examples/mem_agent/run-qwen3-0.6B-train.sh create mode 100644 tests/examples/mem_agent/test_pilot_data.py diff --git a/examples/mem_agent/README.md b/examples/mem_agent/README.md index a86032d75..96cf7af56 100644 --- a/examples/mem_agent/README.md +++ b/examples/mem_agent/README.md @@ -80,3 +80,32 @@ bash examples/mem_agent/run-pipeline.sh ``` GPU execution is intentionally not started by the CPU test suite. The caller remains responsible for starting the ReLax/Ray environment described by the repository deployment guide. + +## Qwen3-0.6B single-4090 pilot + +The 0.6B recipe is a low-cost pipeline and learnability diagnostic; it does not replace the frozen Qwen3-4B VIME/ReLax acceptance run. It fixes `Qwen/Qwen3-0.6B@c1899de289a04d12100db370d81485cdf75e47ca`, disables Qwen3 thinking, uses 512-token chunks, a 128-token memory, a 64-token final answer, and at most four chunks. + +Prepare 24 immutable 2--4-chunk candidates before allocating a GPU: + +```bash +TOKENIZER_PATH=/data/models/Qwen3-0.6B \ +DATA_DIR=/data/task36-pilot \ +bash examples/mem_agent/prepare-pilot-candidates.sh +``` + +The baseline samples every candidate eight times. Selection fails unless at least 12 prompts have both a success and a failure; 2--6 successes out of 8 are preferred. Eight prompts become the training split and four disjoint prompts become the held-out pilot split. This Pass@N-screened set deliberately supplies GRPO reward variance and must not be reported as an unbiased HotpotQA metric. + +```bash +MODEL_PATH=/data/models/Qwen3-0.6B \ +DATA_DIR=/data/task36-pilot \ +RESULTS_DIR=/data/task36-runs/baseline \ +bash examples/mem_agent/run-qwen3-0.6B-baseline.sh + +MODEL_PATH=/data/models/Qwen3-0.6B \ +DATA_DIR=/data/task36-pilot \ +RUN_ROOT=/data/task36-runs/train \ +NUM_ROLLOUT=2 \ +bash examples/mem_agent/run-qwen3-0.6B-train.sh +``` + +The train script is TP=1 and produces the complete Ray job log, TensorBoard events, `training-reward.summary.json`, exact `training-reward.csv` points, and `training-reward.svg`. Run the two-step smoke first. Only after its generated/transferred/consumed row counts agree should a longer run start. Converted checkpoints are evaluated against the frozen pilot split with the same seed, sampling parameters, prompt path, and tokenizer via `run-qwen3-0.6B-eval.sh`. diff --git a/examples/mem_agent/compare_results.py b/examples/mem_agent/compare_results.py index 67c5a004d..02a34561f 100644 --- a/examples/mem_agent/compare_results.py +++ b/examples/mem_agent/compare_results.py @@ -18,6 +18,8 @@ "temperature", "top_p", "sampling_count", + "seed", + "enable_thinking", "chunk_tokens", "max_memory_tokens", "max_final_tokens", @@ -112,7 +114,7 @@ def main() -> None: nargs=3, action="append", metavar=("LABEL", "VIME_SUMMARY", "RELAX_SUMMARY"), - required=True, + default=[], help="Repeat for every RULER-HQA length selected for acceptance.", ) parser.add_argument("--metric", default="sub_em_pct") @@ -127,6 +129,8 @@ def main() -> None: ) parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() + if not args.pair and not args.baseline_pair: + parser.error("At least one --pair or --baseline-pair is required.") comparisons = [] for label, vime_path, relax_path in args.pair: diff --git a/examples/mem_agent/config-pilot-qwen3-0.6b.yaml b/examples/mem_agent/config-pilot-qwen3-0.6b.yaml new file mode 100644 index 000000000..23bc9a5df --- /dev/null +++ b/examples/mem_agent/config-pilot-qwen3-0.6b.yaml @@ -0,0 +1,16 @@ +# Qwen3-0.6B/RTX 4090 diagnostic recipe. This does not replace config.yaml's +# frozen Qwen3-4B/VIME acceptance contract. +mem_agent_chunk_tokens: 512 +mem_agent_max_memory_tokens: 128 +mem_agent_max_final_tokens: 64 +mem_agent_max_chunks: 4 +mem_agent_enable_thinking: false +mem_agent_credit_assignment: split +mem_agent_strict_alignment: true +# Four memory turns plus one final turn is the maximum pilot expansion. +custom_train_sample_expansion_factor: 5 +custom_train_data_group_size: 1 +custom_train_expanded_batch: true +mem_agent_train_rows_multiple: 1 +model_id: Qwen/Qwen3-0.6B +model_revision: c1899de289a04d12100db370d81485cdf75e47ca diff --git a/examples/mem_agent/eval_ruler_hqa.py b/examples/mem_agent/eval_ruler_hqa.py index 2561f52cc..bd4812b6e 100644 --- a/examples/mem_agent/eval_ruler_hqa.py +++ b/examples/mem_agent/eval_ruler_hqa.py @@ -8,6 +8,7 @@ import asyncio import hashlib import json +from collections import defaultdict from pathlib import Path from typing import Any @@ -54,7 +55,7 @@ def load_data(path: Path) -> list[dict[str, Any]]: else: metadata = item.get("metadata") or {} normalized = { - "_id": index, + "_id": item.get("_id", index), "input": item.get("prompt", metadata.get("question", "")), "answers": metadata.get("ground_truth", [item.get("label", "")]), "context": metadata.get("context", ""), @@ -74,6 +75,8 @@ async def _chat_once( temperature: float, top_p: float, max_tokens: int, + seed: int | None = None, + enable_thinking: bool | None = None, ) -> str: payload = { "model": model, @@ -82,6 +85,13 @@ async def _chat_once( "top_p": top_p, "max_tokens": max_tokens, } + if seed is not None: + payload["seed"] = seed + if enable_thinking is not None: + # vLLM and SGLang expose Qwen's chat-template controls through this + # OpenAI-compatible extension. Formal VIME evaluation leaves it unset; + # the short-response 0.6B pilot disables thinking explicitly. + payload["chat_template_kwargs"] = {"enable_thinking": enable_thinking} async with session.post( f"{base_url.rstrip('/')}/chat/completions", headers={"Authorization": f"Bearer {api_key}"}, @@ -99,6 +109,7 @@ async def recurrent_infer( args: argparse.Namespace, tokenizer: Any, session: aiohttp.ClientSession, + sample_index: int = 0, ) -> tuple[str, dict[str, Any]]: question = str(item["input"]).strip() context_ids = tokenizer.encode(str(item["context"]).strip(), add_special_tokens=False) @@ -108,7 +119,7 @@ async def recurrent_infer( chunks = all_chunks[: args.max_chunks] memory = NO_MEMORY memory_lengths = [] - for chunk_ids in chunks: + for chunk_index, chunk_ids in enumerate(chunks): chunk = tokenizer.decode(chunk_ids, skip_special_tokens=True) generated_memory = strip_stop_tokens( await _chat_once( @@ -120,6 +131,8 @@ async def recurrent_infer( args.temperature, args.top_p, args.max_memory_tokens, + _request_seed(args, item["_id"], sample_index, "memory", chunk_index), + args.enable_thinking, ) ) # Match training exactly: the next turn only sees the re-tokenized, @@ -135,6 +148,8 @@ async def recurrent_infer( args.temperature, args.top_p, args.max_final_tokens, + _request_seed(args, item["_id"], sample_index, "final", len(chunks)), + args.enable_thinking, ) return answer, { "num_chunks": len(chunks), @@ -148,6 +163,7 @@ async def base_infer( args: argparse.Namespace, tokenizer: Any, session: aiohttp.ClientSession, + sample_index: int = 0, ) -> tuple[str, dict[str, Any]]: suffix = f"\n\nQuestion: {item['input']}\nPlease answer the question and put the answer in \\boxed{{}}." suffix_ids = tokenizer.encode(suffix, add_special_tokens=False) @@ -172,10 +188,50 @@ async def base_infer( args.temperature, args.top_p, args.max_final_tokens, + _request_seed(args, item["_id"], sample_index, "base", 0), + args.enable_thinking, ) return answer, {"num_chunks": 1, "context_truncated": context_truncated, "memory_token_lengths": []} +def _request_seed( + args: argparse.Namespace, + item_id: Any, + sample_index: int, + stage: str, + turn_index: int, +) -> int | None: + """Derive a stable per-request seed without coupling concurrent tasks.""" + base_seed = getattr(args, "seed", None) + if base_seed is None: + return None + material = f"{base_seed}|{item_id}|{sample_index}|{stage}|{turn_index}".encode() + return int.from_bytes(hashlib.sha256(material).digest()[:4], "big") & 0x7FFFFFFF + + +def summarize_pass_at_n(records: list[dict[str, Any]], samples_per_item: int) -> dict[str, float | int]: + """Summarize boxed-reward Pass@N and GRPO-useful group variance.""" + grouped: dict[str, list[dict[str, Any]]] = defaultdict(list) + for record in records: + grouped[str(record["_id"])].append(record) + + complete_groups = [group for group in grouped.values() if len(group) == samples_per_item] + success_counts = [ + sum(float(record.get("judge_boxed_em", 0.0)) == 1.0 for record in group) for group in complete_groups + ] + pass_count = sum(count > 0 for count in success_counts) + variance_count = sum(0 < count < samples_per_item for count in success_counts) + denominator = len(complete_groups) + return { + "pass_at_n": pass_count / denominator if denominator else 0.0, + "pass_at_n_pct": 100 * pass_count / denominator if denominator else 0.0, + "complete_prompt_groups": denominator, + "reward_variance_groups": variance_count, + "reward_variance_group_pct": 100 * variance_count / denominator if denominator else 0.0, + "mean_successes_per_prompt": sum(success_counts) / denominator if denominator else 0.0, + } + + async def run_evaluation( data: list[dict[str, Any]], args: argparse.Namespace, tokenizer: Any ) -> tuple[list[dict[str, Any]], dict[str, Any]]: @@ -188,22 +244,27 @@ async def run_evaluation( async with aiohttp.ClientSession(timeout=timeout) as session: - async def evaluate_one(item: dict[str, Any]) -> dict[str, Any]: + async def evaluate_one(item: dict[str, Any], sample_index: int) -> dict[str, Any]: answers = item["answers"] if isinstance(item["answers"], list) else [item["answers"]] answers = [str(answer) for answer in answers] ground_truth = answers[0] try: async with semaphore: if args.mode == "recurrent": - response, diagnostics = await recurrent_infer(item, args, tokenizer, session) + response, diagnostics = await recurrent_infer( + item, args, tokenizer, session, sample_index=sample_index + ) else: - response, diagnostics = await base_infer(item, args, tokenizer, session) + response, diagnostics = await base_infer( + item, args, tokenizer, session, sample_index=sample_index + ) # RULER-HQA's VIME-compatible metrics score the first reference. # boxed_em additionally mirrors the HotpotQA training reward and # accepts any annotated answer. prediction = extract_last_boxed(response[-300:]) return { "_id": item["_id"], + "sample_index": sample_index, "answer": ground_truth, "answers": answers, "pred": prediction, @@ -220,6 +281,7 @@ async def evaluate_one(item: dict[str, Any]) -> dict[str, Any]: # serving request itself failed. return { "_id": item["_id"], + "sample_index": sample_index, "answer": ground_truth, "answers": answers, "pred": "", @@ -231,10 +293,13 @@ async def evaluate_one(item: dict[str, Any]) -> dict[str, Any]: "error": f"{type(exc).__name__}: {exc}", } - records = await asyncio.gather(*(evaluate_one(item) for item in data)) + records = await asyncio.gather( + *(evaluate_one(item, sample_index) for item in data for sample_index in range(args.samples_per_item)) + ) summary = { **aggregate(records), + **summarize_pass_at_n(records, args.samples_per_item), "mode": args.mode, "model": args.model, "tokenizer": args.tokenizer, @@ -246,7 +311,9 @@ async def evaluate_one(item: dict[str, Any]) -> dict[str, Any]: "evaluator_schema_version": EVALUATOR_SCHEMA_VERSION, "temperature": args.temperature, "top_p": args.top_p, - "sampling_count": 1, + "sampling_count": args.samples_per_item, + "seed": args.seed, + "enable_thinking": args.enable_thinking, "chunk_tokens": args.chunk_tokens, "max_memory_tokens": args.max_memory_tokens, "max_final_tokens": args.max_final_tokens, @@ -267,6 +334,12 @@ def main() -> None: parser.add_argument("--mode", choices=("recurrent", "base"), default="recurrent") parser.add_argument("--base-url", default="http://127.0.0.1:8000/v1") parser.add_argument("--api-key", default="EMPTY") + parser.add_argument("--samples-per-item", type=int, default=1) + parser.add_argument("--seed", type=int, default=None) + thinking = parser.add_mutually_exclusive_group() + thinking.add_argument("--enable-thinking", dest="enable_thinking", action="store_true") + thinking.add_argument("--disable-thinking", dest="enable_thinking", action="store_false") + parser.set_defaults(enable_thinking=None) parser.add_argument("--temperature", type=float, default=0.7) parser.add_argument("--top-p", type=float, default=0.95) parser.add_argument("--chunk-tokens", type=int, default=2048) @@ -278,6 +351,8 @@ def main() -> None: parser.add_argument("--concurrency", type=int, default=16) parser.add_argument("--timeout", type=int, default=86400) args = parser.parse_args() + if args.samples_per_item <= 0: + parser.error("--samples-per-item must be positive") from transformers import AutoTokenizer diff --git a/examples/mem_agent/prepare-pilot-candidates.sh b/examples/mem_agent/prepare-pilot-candidates.sh new file mode 100755 index 000000000..8ece5ab12 --- /dev/null +++ b/examples/mem_agent/prepare-pilot-candidates.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +DATA_DIR="${DATA_DIR:?Set DATA_DIR to the Task 36 pilot data directory.}" +TOKENIZER_PATH="${TOKENIZER_PATH:?Set TOKENIZER_PATH to the frozen Qwen3-0.6B tokenizer.}" +SOURCE_DATA="${SOURCE_DATA:-${DATA_DIR}/train.jsonl}" +CANDIDATE_COUNT="${CANDIDATE_COUNT:-24}" + +mkdir -p "${DATA_DIR}" +if [[ ! -f "${SOURCE_DATA}" ]]; then + # Download only the training parquet before GPU allocation. The long RULER + # files are not needed for the short 0.6B pilot candidate screen. + python3 "${SCRIPT_DIR}/prepare_data.py" \ + --hf-file hotpotqa_train_32k.parquet \ + --output "${SOURCE_DATA}" \ + --repo-id BytedTsinghua-SIA/hotpotqa \ + --revision 27275ff4fee67ac0acb6478e405e7ac07efbdc1a \ + --cache-dir "${DATA_DIR}/hf-cache" \ + --manifest "${DATA_DIR}/source-manifest.json" +fi + +python3 "${SCRIPT_DIR}/prepare_pilot_data.py" candidates \ + --input "${SOURCE_DATA}" \ + --tokenizer "${TOKENIZER_PATH}" \ + --output "${DATA_DIR}/pilot-candidates.jsonl" \ + --manifest "${DATA_DIR}/pilot-candidates.manifest.json" \ + --chunk-tokens 512 \ + --min-chunks 2 \ + --max-chunks 4 \ + --candidate-count "${CANDIDATE_COUNT}" \ + --seed 42 + +echo "Prepared ${DATA_DIR}/pilot-candidates.jsonl before GPU allocation." diff --git a/examples/mem_agent/prepare_data.py b/examples/mem_agent/prepare_data.py index 200df86d9..ba4d8e38e 100644 --- a/examples/mem_agent/prepare_data.py +++ b/examples/mem_agent/prepare_data.py @@ -36,13 +36,44 @@ def _first_user_content(prompt: Any) -> str: return "" +def _context_text(value: Any) -> str: + """Flatten both MemAgent strings and official HotpotQA paragraph + structs.""" + value = _json_safe(value) + if isinstance(value, str): + return value.strip() + if isinstance(value, dict): + titles = value.get("title", []) + sentences = value.get("sentences", []) + if isinstance(titles, str): + titles = [titles] + if sentences and isinstance(sentences[0], str): + sentences = [sentences] + paragraphs = [] + for index, paragraph_sentences in enumerate(sentences): + title = str(titles[index]).strip() if index < len(titles) else "" + body = " ".join(str(sentence).strip() for sentence in paragraph_sentences if str(sentence).strip()) + paragraphs.append("\n".join(part for part in (title, body) if part)) + return "\n\n".join(part for part in paragraphs if part).strip() + if isinstance(value, list): + paragraphs = [] + for part in value: + if isinstance(part, dict): + paragraphs.append(_context_text(part)) + elif isinstance(part, (list, tuple)) and len(part) == 2: + title, sentences = part + body = " ".join(str(sentence).strip() for sentence in sentences if str(sentence).strip()) + paragraphs.append("\n".join(item for item in (str(title).strip(), body) if item)) + else: + paragraphs.append(str(part).strip()) + return "\n\n".join(part for part in paragraphs if part).strip() + return str(value).strip() + + def convert_row(row: dict[str, Any]) -> dict[str, Any] | None: """Normalize either training parquet rows or RULER-HQA JSON rows.""" row = _json_safe(row) - context = row.get("context", "") - if isinstance(context, list): - context = "\n\n".join(str(part) for part in context) - context = str(context).strip() + context = _context_text(row.get("context", "")) if not context: return None @@ -50,6 +81,13 @@ def convert_row(row: dict[str, Any]) -> dict[str, Any] | None: question = str(row["input"]).strip() ground_truth = row.get("answers", []) extra_info = row + elif row.get("question") and row.get("answer") is not None: + # The resource-constrained pilot uses the official HotpotQA distractor + # rows directly. Formal VIME reproduction still uses the frozen 32K + # MemAgent parquet path above; both normalize to the same public schema. + question = str(row["question"]).strip() + ground_truth = [row["answer"]] + extra_info = row else: extra_info = row.get("extra_info") or {} question = _first_user_content(row.get("prompt")) or str(extra_info.get("question", "")) @@ -65,7 +103,7 @@ def convert_row(row: dict[str, Any]) -> dict[str, Any] | None: if not question or not ground_truth: return None - return { + converted = { "prompt": question, "label": ground_truth[0], "metadata": { @@ -76,6 +114,10 @@ def convert_row(row: dict[str, Any]) -> dict[str, Any] | None: "data_source": str(row.get("data_source", "hotpotqa")), }, } + source_id = row.get("id", row.get("_id")) + if source_id is not None: + converted["_id"] = str(source_id) + return converted def read_rows(path: Path) -> Iterable[dict[str, Any]]: diff --git a/examples/mem_agent/prepare_pilot_data.py b/examples/mem_agent/prepare_pilot_data.py new file mode 100644 index 000000000..3253a765b --- /dev/null +++ b/examples/mem_agent/prepare_pilot_data.py @@ -0,0 +1,377 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Build a short-context, non-degenerate Qwen3-0.6B MemAgent pilot set. + +The workflow is deliberately two-stage. ``candidates`` filters only immutable +input properties such as token length. After an untrained recurrent Pass@N +run, ``select`` keeps prompts with both successes and failures. This gives GRPO +useful within-group reward variance without pretending the screened pilot is a +formal, unbiased HotpotQA benchmark. +""" + +from __future__ import annotations + +import argparse +import copy +import hashlib +import json +import math +import random +from collections import Counter, defaultdict +from pathlib import Path +from typing import Any + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def read_jsonl(path: Path) -> list[dict[str, Any]]: + with path.open(encoding="utf-8") as source: + return [json.loads(line) for line in source if line.strip()] + + +def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as destination: + for row in rows: + destination.write(json.dumps(row, ensure_ascii=False) + "\n") + + +def _metadata(row: dict[str, Any]) -> dict[str, Any]: + metadata = row.get("metadata") + return metadata if isinstance(metadata, dict) else {} + + +def _answers(row: dict[str, Any]) -> list[str]: + metadata = _metadata(row) + answers = metadata.get("ground_truth", [row.get("label", "")]) + if isinstance(answers, str): + answers = [answers] + return [str(answer).strip() for answer in answers if str(answer).strip()] + + +def _contains_subsequence(sequence: list[int], subsequence: list[int]) -> bool: + if not subsequence or len(subsequence) > len(sequence): + return False + width = len(subsequence) + return any(sequence[offset : offset + width] == subsequence for offset in range(len(sequence) - width + 1)) + + +def build_candidates( + rows: list[dict[str, Any]], + tokenizer: Any, + *, + chunk_tokens: int, + min_chunks: int, + max_chunks: int, + candidate_count: int, + seed: int, + require_answer_in_context: bool = True, +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Filter by immutable length/content properties and sample deterministically.""" + if not 0 < min_chunks <= max_chunks: + raise ValueError("Expected 0 < min_chunks <= max_chunks.") + if chunk_tokens <= 0 or candidate_count <= 0: + raise ValueError("chunk_tokens and candidate_count must be positive.") + + eligible = [] + rejected = Counter() + chunk_histogram = Counter() + for source_index, row in enumerate(rows): + metadata = _metadata(row) + context = str(metadata.get("context", "")).strip() + answers = _answers(row) + if not context or not answers: + rejected["missing_context_or_answer"] += 1 + continue + context_ids = tokenizer.encode(context, add_special_tokens=False) + num_chunks = math.ceil(len(context_ids) / chunk_tokens) + if num_chunks < min_chunks: + rejected["too_short"] += 1 + continue + if num_chunks > max_chunks: + rejected["too_long"] += 1 + continue + answer_visible = any( + _contains_subsequence(context_ids, tokenizer.encode(answer, add_special_tokens=False)) + for answer in answers + ) + if require_answer_in_context and not answer_visible: + rejected["answer_not_in_context"] += 1 + continue + + candidate = copy.deepcopy(row) + candidate_id = str(candidate.get("_id", f"train-{source_index:06d}")) + candidate["_id"] = candidate_id + candidate_metadata = candidate.setdefault("metadata", {}) + candidate_metadata["pilot"] = { + "source_index": source_index, + "context_tokens": len(context_ids), + "num_chunks": num_chunks, + "answer_in_context": answer_visible, + } + eligible.append(candidate) + chunk_histogram[num_chunks] += 1 + + random.Random(seed).shuffle(eligible) + selected = eligible[:candidate_count] + if len(selected) < candidate_count: + raise ValueError(f"Only {len(selected)} rows satisfy pilot filters; requested {candidate_count}.") + return selected, { + "input_rows": len(rows), + "eligible_rows": len(eligible), + "candidate_rows": len(selected), + "rejected": dict(sorted(rejected.items())), + "eligible_chunk_histogram": {str(key): value for key, value in sorted(chunk_histogram.items())}, + "chunk_tokens": chunk_tokens, + "min_chunks": min_chunks, + "max_chunks": max_chunks, + "candidate_count": candidate_count, + "seed": seed, + "require_answer_in_context": require_answer_in_context, + } + + +def _screening_groups( + candidates: list[dict[str, Any]], + records: list[dict[str, Any]], + *, + samples_per_item: int, + min_successes: int, + max_successes: int, + preferred_min_successes: int, + preferred_max_successes: int, +) -> list[dict[str, Any]]: + grouped: dict[str, list[dict[str, Any]]] = defaultdict(list) + for record in records: + grouped[str(record["_id"])].append(record) + + screening = [] + for candidate in candidates: + candidate_id = str(candidate["_id"]) + group = grouped.get(candidate_id, []) + errors = sum(bool(record.get("error")) for record in group) + sample_indices = {int(record.get("sample_index", -1)) for record in group} + successes = sum(float(record.get("judge_boxed_em", 0.0)) == 1.0 for record in group) + complete = len(group) == samples_per_item and sample_indices == set(range(samples_per_item)) + if not complete: + status = "incomplete" + elif errors: + status = "request_error" + elif preferred_min_successes <= successes <= preferred_max_successes: + status = "preferred" + elif min_successes <= successes <= max_successes: + status = "eligible_boundary" + else: + status = "no_reward_variance" + screening.append( + { + "_id": candidate_id, + "samples": len(group), + "successes": successes, + "failures": len(group) - successes, + "errors": errors, + "status": status, + "selected_split": None, + } + ) + return screening + + +def select_pilot_sets( + candidates: list[dict[str, Any]], + records: list[dict[str, Any]], + *, + samples_per_item: int, + train_count: int, + eval_count: int, + seed: int, + min_successes: int = 1, + max_successes: int | None = None, + preferred_min_successes: int = 2, + preferred_max_successes: int | None = None, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]: + """Select disjoint train/eval diagnostics with Pass@N and reward variance.""" + if samples_per_item <= 1: + raise ValueError("Pass@N screening requires samples_per_item > 1.") + max_successes = samples_per_item - 1 if max_successes is None else max_successes + preferred_max_successes = samples_per_item - 2 if preferred_max_successes is None else preferred_max_successes + if not 1 <= min_successes <= max_successes < samples_per_item: + raise ValueError("Success bounds must guarantee at least one success and one failure.") + + screening = _screening_groups( + candidates, + records, + samples_per_item=samples_per_item, + min_successes=min_successes, + max_successes=max_successes, + preferred_min_successes=preferred_min_successes, + preferred_max_successes=preferred_max_successes, + ) + by_id = {str(row["_id"]): row for row in candidates} + eligible = [entry for entry in screening if entry["status"] in ("preferred", "eligible_boundary")] + tie_break = {entry["_id"]: random.Random(f"{seed}:{entry['_id']}").random() for entry in eligible} + eligible.sort( + key=lambda entry: ( + entry["status"] != "preferred", + abs(entry["successes"] - samples_per_item / 2), + tie_break[entry["_id"]], + ) + ) + required = train_count + eval_count + if len(eligible) < required: + raise ValueError( + f"Only {len(eligible)} complete prompts have non-degenerate Pass@{samples_per_item}; " + f"need {required}. Run another candidate shard before training." + ) + + chosen = eligible[:required] + random.Random(seed).shuffle(chosen) + eval_entries = chosen[:eval_count] + train_entries = chosen[eval_count:] + selected_lookup = {entry["_id"]: entry for entry in chosen} + for entry in screening: + if entry["_id"] in {item["_id"] for item in train_entries}: + entry["selected_split"] = "train" + elif entry["_id"] in {item["_id"] for item in eval_entries}: + entry["selected_split"] = "eval" + + def materialize(entries: list[dict[str, Any]], split: str) -> list[dict[str, Any]]: + output = [] + for entry in entries: + row = copy.deepcopy(by_id[entry["_id"]]) + pilot = row.setdefault("metadata", {}).setdefault("pilot", {}) + pilot.update( + { + "split": split, + "baseline_samples": samples_per_item, + "baseline_successes": entry["successes"], + "baseline_failures": entry["failures"], + "baseline_pass_at_n": True, + "baseline_reward_variance": True, + } + ) + output.append(row) + return output + + train_rows = materialize(train_entries, "train") + eval_rows = materialize(eval_entries, "eval") + assert not ({row["_id"] for row in train_rows} & {row["_id"] for row in eval_rows}) + return ( + train_rows, + eval_rows, + { + "samples_per_item": samples_per_item, + "train_count": train_count, + "eval_count": eval_count, + "seed": seed, + "min_successes": min_successes, + "max_successes": max_successes, + "preferred_min_successes": preferred_min_successes, + "preferred_max_successes": preferred_max_successes, + "screening": screening, + "selected_ids": { + "train": [row["_id"] for row in train_rows], + "eval": [row["_id"] for row in eval_rows], + }, + "selected_successes": { + entry_id: selected_lookup[entry_id]["successes"] for entry_id in sorted(selected_lookup) + }, + }, + ) + + +def _write_manifest(path: Path, manifest: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as destination: + json.dump(manifest, destination, ensure_ascii=False, indent=2) + destination.write("\n") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + + candidates = subparsers.add_parser("candidates") + candidates.add_argument("--input", type=Path, required=True) + candidates.add_argument("--tokenizer", required=True) + candidates.add_argument("--output", type=Path, required=True) + candidates.add_argument("--manifest", type=Path, required=True) + candidates.add_argument("--chunk-tokens", type=int, default=512) + candidates.add_argument("--min-chunks", type=int, default=2) + candidates.add_argument("--max-chunks", type=int, default=4) + candidates.add_argument("--candidate-count", type=int, default=24) + candidates.add_argument("--seed", type=int, default=42) + candidates.add_argument("--allow-answer-not-in-context", action="store_true") + + select = subparsers.add_parser("select") + select.add_argument("--candidates", type=Path, required=True) + select.add_argument("--baseline-records", type=Path, required=True) + select.add_argument("--train-output", type=Path, required=True) + select.add_argument("--eval-output", type=Path, required=True) + select.add_argument("--manifest", type=Path, required=True) + select.add_argument("--samples-per-item", type=int, default=8) + select.add_argument("--train-count", type=int, default=8) + select.add_argument("--eval-count", type=int, default=4) + select.add_argument("--seed", type=int, default=42) + + args = parser.parse_args() + if args.command == "candidates": + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + output_rows, manifest = build_candidates( + read_jsonl(args.input), + tokenizer, + chunk_tokens=args.chunk_tokens, + min_chunks=args.min_chunks, + max_chunks=args.max_chunks, + candidate_count=args.candidate_count, + seed=args.seed, + require_answer_in_context=not args.allow_answer_not_in_context, + ) + write_jsonl(args.output, output_rows) + manifest.update( + { + "source_file": str(args.input), + "source_sha256": sha256_file(args.input), + "output_file": str(args.output), + "output_sha256": sha256_file(args.output), + "tokenizer": args.tokenizer, + } + ) + _write_manifest(args.manifest, manifest) + return + + candidate_rows = read_jsonl(args.candidates) + train_rows, eval_rows, manifest = select_pilot_sets( + candidate_rows, + read_jsonl(args.baseline_records), + samples_per_item=args.samples_per_item, + train_count=args.train_count, + eval_count=args.eval_count, + seed=args.seed, + ) + write_jsonl(args.train_output, train_rows) + write_jsonl(args.eval_output, eval_rows) + manifest.update( + { + "candidate_file": str(args.candidates), + "candidate_sha256": sha256_file(args.candidates), + "baseline_records_file": str(args.baseline_records), + "baseline_records_sha256": sha256_file(args.baseline_records), + "train_file": str(args.train_output), + "train_sha256": sha256_file(args.train_output), + "eval_file": str(args.eval_output), + "eval_sha256": sha256_file(args.eval_output), + } + ) + _write_manifest(args.manifest, manifest) + + +if __name__ == "__main__": + main() diff --git a/examples/mem_agent/prompts.py b/examples/mem_agent/prompts.py index b4ba11dcb..8d9487b8d 100644 --- a/examples/mem_agent/prompts.py +++ b/examples/mem_agent/prompts.py @@ -90,7 +90,21 @@ def final_instruction(question: str, memory: str, *, evaluation: bool = False) - ) -def render_chat_prompt(tokenizer: Any, instruction: str) -> str: - """Render one independent user turn with the model's chat template.""" +def render_chat_prompt(tokenizer: Any, instruction: str, *, enable_thinking: bool | None = None) -> str: + """Render one independent user turn with the model's chat template. + + The formal Qwen3-4B reproduction leaves ``enable_thinking`` unspecified to + match fixed VIME. The resource-constrained Qwen3-0.6B pilot can disable it + explicitly so a short response budget contains memory/answer text instead + of only a truncated reasoning trace. + """ messages = [{"role": "user", "content": instruction}] - return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + template_kwargs = {} + if enable_thinking is not None: + template_kwargs["enable_thinking"] = enable_thinking + return tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + **template_kwargs, + ) diff --git a/examples/mem_agent/rollout.py b/examples/mem_agent/rollout.py index 4419eead3..aaa005b75 100644 --- a/examples/mem_agent/rollout.py +++ b/examples/mem_agent/rollout.py @@ -125,6 +125,7 @@ async def generate_trajectory( max_memory_tokens = int(getattr(args, "mem_agent_max_memory_tokens", 1024)) max_final_tokens = int(getattr(args, "mem_agent_max_final_tokens", 256)) max_chunks = int(getattr(args, "mem_agent_max_chunks", 64)) + enable_thinking = getattr(args, "mem_agent_enable_thinking", None) if not getattr(args, "mem_agent_strict_alignment", True): raise ValueError("MemAgent training requires mem_agent_strict_alignment=true.") if max_memory_tokens <= 0 or max_final_tokens <= 0: @@ -150,6 +151,7 @@ async def generate_trajectory( prompt = render_chat_prompt( tokenizer, memory_instruction(question, memory, chunk, evaluation=evaluation), + enable_thinking=enable_thinking, ) turn_sample = await _run_turn( args, @@ -179,7 +181,11 @@ async def generate_trajectory( # The final request deliberately excludes context/chunks. This enforces # the question + latest-memory information boundary from the task spec. - final_prompt = render_chat_prompt(tokenizer, final_instruction(question, memory, evaluation=evaluation)) + final_prompt = render_chat_prompt( + tokenizer, + final_instruction(question, memory, evaluation=evaluation), + enable_thinking=enable_thinking, + ) final_sample = await _run_turn( args, sample, diff --git a/examples/mem_agent/run-qwen3-0.6B-baseline.sh b/examples/mem_agent/run-qwen3-0.6B-baseline.sh new file mode 100755 index 000000000..07464d99c --- /dev/null +++ b/examples/mem_agent/run-qwen3-0.6B-baseline.sh @@ -0,0 +1,109 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to the frozen Qwen3-0.6B BF16 checkpoint.}" +TOKENIZER_PATH="${TOKENIZER_PATH:-${MODEL_PATH}}" +DATA_DIR="${DATA_DIR:?Set DATA_DIR to the Task 36 pilot data directory.}" +RESULTS_DIR="${RESULTS_DIR:?Set RESULTS_DIR to the Task 36 baseline output directory.}" +GPU_ID="${GPU_ID:-0}" +SERVE_PORT="${SERVE_PORT:-30000}" +SAMPLES_PER_ITEM="${SAMPLES_PER_ITEM:-8}" +CONCURRENCY="${CONCURRENCY:-2}" + +[[ -f "${DATA_DIR}/pilot-candidates.jsonl" ]] || { + echo "Missing pilot candidates; run prepare-pilot-candidates.sh before starting a GPU." >&2 + exit 1 +} +mkdir -p "${RESULTS_DIR}" +export CUDA_VISIBLE_DEVICES="${GPU_ID}" +export NO_PROXY="127.0.0.1,localhost,::1" +export no_proxy="${NO_PROXY}" + +SERVER_LOG="${RESULTS_DIR}/qwen3-0.6b-baseline-server.log" +python3 -m sglang.launch_server \ + --model-path "${MODEL_PATH}" \ + --host 127.0.0.1 \ + --port "${SERVE_PORT}" \ + --api-key EMPTY \ + --tp-size 1 \ + --context-length 1536 \ + --mem-fraction-static 0.55 \ + --trust-remote-code \ + >"${SERVER_LOG}" 2>&1 & +SERVER_PID=$! +trap 'kill -TERM "${SERVER_PID}" 2>/dev/null || true; wait "${SERVER_PID}" 2>/dev/null || true' EXIT INT TERM + +for _ in $(seq 1 120); do + if ! kill -0 "${SERVER_PID}" 2>/dev/null; then + echo "SGLang exited early; see ${SERVER_LOG}" >&2 + exit 1 + fi + if curl --noproxy '*' -fsS "http://127.0.0.1:${SERVE_PORT}/health" >/dev/null; then + break + fi + sleep 5 +done +curl --noproxy '*' -fsS "http://127.0.0.1:${SERVE_PORT}/health" >/dev/null + +python3 "${SCRIPT_DIR}/eval_ruler_hqa.py" \ + --data-file "${DATA_DIR}/pilot-candidates.jsonl" \ + --model "${MODEL_PATH}" \ + --tokenizer "${TOKENIZER_PATH}" \ + --output-dir "${RESULTS_DIR}" \ + --run-name qwen3-0.6b-untrained-pass8 \ + --mode recurrent \ + --base-url "http://127.0.0.1:${SERVE_PORT}/v1" \ + --api-key EMPTY \ + --samples-per-item "${SAMPLES_PER_ITEM}" \ + --seed 42 \ + --disable-thinking \ + --temperature 1.0 \ + --top-p 1.0 \ + --chunk-tokens 512 \ + --max-memory-tokens 128 \ + --max-final-tokens 64 \ + --max-chunks 4 \ + --max-input-tokens 1472 \ + --server-max-model-len 1536 \ + --concurrency "${CONCURRENCY}" + +python3 "${SCRIPT_DIR}/prepare_pilot_data.py" select \ + --candidates "${DATA_DIR}/pilot-candidates.jsonl" \ + --baseline-records "${RESULTS_DIR}/qwen3-0.6b-untrained-pass8.jsonl" \ + --train-output "${DATA_DIR}/pilot-train.jsonl" \ + --eval-output "${DATA_DIR}/pilot-eval.jsonl" \ + --manifest "${DATA_DIR}/pilot-selection.manifest.json" \ + --samples-per-item "${SAMPLES_PER_ITEM}" \ + --train-count 8 \ + --eval-count 4 \ + --seed 42 + +# Re-evaluate the disjoint diagnostic split while the same untrained server is +# still resident. This is the baseline compared with trained checkpoints; the +# larger candidate run exists only for Pass@N screening. +python3 "${SCRIPT_DIR}/eval_ruler_hqa.py" \ + --data-file "${DATA_DIR}/pilot-eval.jsonl" \ + --model "${MODEL_PATH}" \ + --tokenizer "${TOKENIZER_PATH}" \ + --output-dir "${RESULTS_DIR}" \ + --run-name qwen3-0.6b-untrained-heldout-pass8 \ + --mode recurrent \ + --base-url "http://127.0.0.1:${SERVE_PORT}/v1" \ + --api-key EMPTY \ + --samples-per-item "${SAMPLES_PER_ITEM}" \ + --seed 4242 \ + --disable-thinking \ + --temperature 1.0 \ + --top-p 1.0 \ + --chunk-tokens 512 \ + --max-memory-tokens 128 \ + --max-final-tokens 64 \ + --max-chunks 4 \ + --max-input-tokens 1472 \ + --server-max-model-len 1536 \ + --concurrency "${CONCURRENCY}" + +echo "Baseline and non-degenerate pilot split completed under ${RESULTS_DIR} and ${DATA_DIR}." diff --git a/examples/mem_agent/run-qwen3-0.6B-eval.sh b/examples/mem_agent/run-qwen3-0.6B-eval.sh new file mode 100755 index 000000000..ad1011253 --- /dev/null +++ b/examples/mem_agent/run-qwen3-0.6B-eval.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to a converted Qwen3-0.6B checkpoint.}" +TOKENIZER_PATH="${TOKENIZER_PATH:?Set TOKENIZER_PATH to the frozen base tokenizer.}" +EVAL_DATA="${EVAL_DATA:?Set EVAL_DATA to the frozen pilot-eval.jsonl.}" +RESULTS_DIR="${RESULTS_DIR:?Set RESULTS_DIR to the checkpoint evaluation output directory.}" +RUN_NAME="${RUN_NAME:?Set RUN_NAME to an immutable checkpoint/run label.}" +GPU_ID="${GPU_ID:-0}" +SERVE_PORT="${SERVE_PORT:-30000}" +SAMPLES_PER_ITEM="${SAMPLES_PER_ITEM:-8}" +CONCURRENCY="${CONCURRENCY:-2}" +BASELINE_SUMMARY="${BASELINE_SUMMARY:-}" + +mkdir -p "${RESULTS_DIR}" +export CUDA_VISIBLE_DEVICES="${GPU_ID}" +export NO_PROXY="127.0.0.1,localhost,::1" +export no_proxy="${NO_PROXY}" + +SERVER_LOG="${RESULTS_DIR}/${RUN_NAME}.server.log" +python3 -m sglang.launch_server \ + --model-path "${MODEL_PATH}" \ + --host 127.0.0.1 \ + --port "${SERVE_PORT}" \ + --api-key EMPTY \ + --tp-size 1 \ + --context-length 1536 \ + --mem-fraction-static 0.55 \ + --trust-remote-code \ + >"${SERVER_LOG}" 2>&1 & +SERVER_PID=$! +trap 'kill -TERM "${SERVER_PID}" 2>/dev/null || true; wait "${SERVER_PID}" 2>/dev/null || true' EXIT INT TERM + +for _ in $(seq 1 120); do + if ! kill -0 "${SERVER_PID}" 2>/dev/null; then + echo "SGLang exited early; see ${SERVER_LOG}" >&2 + exit 1 + fi + if curl --noproxy '*' -fsS "http://127.0.0.1:${SERVE_PORT}/health" >/dev/null; then + break + fi + sleep 5 +done +curl --noproxy '*' -fsS "http://127.0.0.1:${SERVE_PORT}/health" >/dev/null + +python3 "${SCRIPT_DIR}/eval_ruler_hqa.py" \ + --data-file "${EVAL_DATA}" \ + --model "${MODEL_PATH}" \ + --tokenizer "${TOKENIZER_PATH}" \ + --output-dir "${RESULTS_DIR}" \ + --run-name "${RUN_NAME}" \ + --mode recurrent \ + --base-url "http://127.0.0.1:${SERVE_PORT}/v1" \ + --api-key EMPTY \ + --samples-per-item "${SAMPLES_PER_ITEM}" \ + --seed 4242 \ + --disable-thinking \ + --temperature 1.0 \ + --top-p 1.0 \ + --chunk-tokens 512 \ + --max-memory-tokens 128 \ + --max-final-tokens 64 \ + --max-chunks 4 \ + --max-input-tokens 1472 \ + --server-max-model-len 1536 \ + --concurrency "${CONCURRENCY}" + +if [[ -n "${BASELINE_SUMMARY}" ]]; then + [[ -f "${BASELINE_SUMMARY}" ]] || { echo "Missing baseline summary: ${BASELINE_SUMMARY}" >&2; exit 1; } + python3 "${SCRIPT_DIR}/compare_results.py" \ + --baseline-pair pilot-boxed-em boxed_em_pct "${BASELINE_SUMMARY}" "${RESULTS_DIR}/${RUN_NAME}.summary.json" \ + --baseline-pair pilot-pass-at-n pass_at_n_pct "${BASELINE_SUMMARY}" "${RESULTS_DIR}/${RUN_NAME}.summary.json" \ + --output "${RESULTS_DIR}/${RUN_NAME}.vs-baseline.json" +fi diff --git a/examples/mem_agent/run-qwen3-0.6B-train.sh b/examples/mem_agent/run-qwen3-0.6B-train.sh new file mode 100755 index 000000000..d7fe286eb --- /dev/null +++ b/examples/mem_agent/run-qwen3-0.6B-train.sh @@ -0,0 +1,149 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail +set -x + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +RELAX_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" +MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to the frozen Qwen3-0.6B BF16 checkpoint.}" +DATA_DIR="${DATA_DIR:?Set DATA_DIR to the screened Task 36 pilot data directory.}" +RUN_ROOT="${RUN_ROOT:?Set RUN_ROOT to the Task 36 experiment output directory.}" +SAVE_DIR="${SAVE_DIR:-${RUN_ROOT}/checkpoints}" +TRAIN_DATA="${TRAIN_DATA:-${DATA_DIR}/pilot-train.jsonl}" +NUM_ROLLOUT="${NUM_ROLLOUT:-20}" +SAVE_INTERVAL="${SAVE_INTERVAL:-5}" +ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-1}" +N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" +GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}" +GPU_ID="${GPU_ID:-0}" +RAY_NUM_CPUS="${RAY_NUM_CPUS:-12}" +RUN_NAME="${RUN_NAME:-mem-agent-qwen3-0.6b-pilot}" + +[[ -f "${TRAIN_DATA}" ]] || { + echo "Missing Pass@N-screened pilot data: ${TRAIN_DATA}" >&2 + exit 1 +} +[[ -f "${DATA_DIR}/pilot-selection.manifest.json" ]] || { + echo "Missing pilot selection manifest; baseline screening must run before training." >&2 + exit 1 +} +[[ -f "${MODEL_PATH}/config.json" ]] || { echo "Missing model config: ${MODEL_PATH}/config.json" >&2; exit 1; } +mkdir -p "${RUN_ROOT}/logs" "${RUN_ROOT}/tensorboard" "${SAVE_DIR}" + +export CUDA_VISIBLE_DEVICES="${GPU_ID}" +export NUM_GPUS=1 +export NO_PROXY="127.0.0.1,localhost,::1" +export no_proxy="${NO_PROXY}" +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export NUMEXPR_NUM_THREADS=1 +export TOKENIZERS_PARALLELISM=false + +# PPIO exposes many logical CPUs under a much smaller cgroup pids.max. Limit +# only `ray start`; all other Ray CLI calls retain their original arguments. +ray() { + if [[ "${1:-}" == "start" ]]; then + command ray "$@" --num-cpus "${RAY_NUM_CPUS}" + else + command ray "$@" + fi +} +if [[ -z "${RELAX_ENTRYPOINT_MODE:-}" ]]; then + source "${RELAX_ROOT}/scripts/entrypoint/local.sh" +fi +unset -f ray +source "${MODEL_CONFIG_DIR}/qwen3-0.6B.sh" + +NOW="$(date '+%Y%m%dT%H%M%S%z')" +LOG_FILE="${RUN_ROOT}/logs/${RUN_NAME}-${NOW}.log" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"actor": [1, 1], "rollout": [1, 1]}' \ + "${MODEL_ARGS[@]}" \ + --hf-checkpoint "${MODEL_PATH}" \ + --ref-load "${MODEL_PATH}" \ + --megatron-to-hf-mode bridge \ + --warm-hf-checkpoint-page-cache \ + --save "${SAVE_DIR}" \ + --save-interval "${SAVE_INTERVAL}" \ + --max-actor-ckpt-to-keep 4 \ + --prompt-data "${TRAIN_DATA}" \ + --input-key prompt \ + --label-key label \ + --metadata-key metadata \ + --custom-generate-function-path examples.mem_agent.rollout.generate \ + --custom-rm-path examples.mem_agent.reward.reward_func \ + --custom-convert-samples-to-train-data-path examples.mem_agent.convert.convert_samples \ + --custom-config-path "${SCRIPT_DIR}/config-pilot-qwen3-0.6b.yaml" \ + --reward-key score \ + --reward-num-workers 2 \ + --num-rollout "${NUM_ROLLOUT}" \ + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" \ + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" \ + --rollout-max-response-len 128 \ + --rollout-max-context-len 1536 \ + --rollout-temperature 1.0 \ + --rollout-top-p 1.0 \ + --rollout-seed 42 \ + --rollout-shuffle \ + --global-batch-size "${GLOBAL_BATCH_SIZE}" \ + --balance-data \ + --log-passrate \ + --use-rollout-logprobs \ + --advantage-estimator grpo \ + --use-kl-loss \ + --kl-loss-coef 0.001 \ + --kl-loss-type low_var_kl \ + --entropy-coef 0.0 \ + --eps-clip 0.2 \ + --eps-clip-high 0.3 \ + --optimizer adam \ + --lr 1e-6 \ + --lr-decay-style constant \ + --weight-decay 0.1 \ + --adam-beta1 0.9 \ + --adam-beta2 0.98 \ + --tensor-model-parallel-size 1 \ + --pipeline-model-parallel-size 1 \ + --context-parallel-size 1 \ + --expert-model-parallel-size 1 \ + --expert-tensor-parallel-size 1 \ + --recompute-granularity full \ + --recompute-method uniform \ + --recompute-num-layers 1 \ + --use-dynamic-batch-size \ + --max-tokens-per-gpu 1536 \ + --log-probs-max-tokens-per-gpu 1536 \ + --rollout-num-gpus-per-engine 1 \ + --sglang-mem-fraction-static 0.35 \ + --seed 1234 \ + --attention-dropout 0.0 \ + --hidden-dropout 0.0 \ + --accumulate-allreduce-grads-in-fp32 \ + --attention-softmax-in-fp32 \ + --attention-backend flash \ + --skip-eval-before-train \ + --max-staleness 0 \ + --num-data-storage-units 1 \ + --colocate \ + --use-health-check \ + --use-metrics-service \ + --tb-project-name Relax/task36-mem-agent-0.6b \ + --tb-experiment-name "${RUN_NAME}-${NOW}" \ + 2>&1 | tee "${LOG_FILE}" + +ray job list >"${RUN_ROOT}/logs/${RUN_NAME}-${NOW}-ray-job-list.txt" 2>&1 + +# Keep the durable text log, exact CSV points and a dependency-free SVG curve +# together. Missing/duplicated/conflicting rollout ids make this step fail. +python3 "${SCRIPT_DIR}/summarize_reward.py" \ + --log-file "${LOG_FILE}" \ + --output "${RUN_ROOT}/training-reward.summary.json" \ + --csv-output "${RUN_ROOT}/training-reward.csv" \ + --svg-output "${RUN_ROOT}/training-reward.svg" \ + --expected-steps "${NUM_ROLLOUT}" \ + --window-size 5 diff --git a/examples/mem_agent/summarize_reward.py b/examples/mem_agent/summarize_reward.py index b25aca009..c996e0d19 100644 --- a/examples/mem_agent/summarize_reward.py +++ b/examples/mem_agent/summarize_reward.py @@ -5,6 +5,7 @@ import argparse import ast +import csv import json import math import re @@ -105,6 +106,62 @@ def summarize_reward_points( } +def write_reward_csv(path: Path, points: list[tuple[int, float]]) -> None: + """Write the exact plotted points in a spreadsheet-friendly form.""" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8", newline="") as destination: + writer = csv.writer(destination) + writer.writerow(["rollout_id", "reward"]) + writer.writerows(points) + + +def write_reward_svg(path: Path, points: list[tuple[int, float]]) -> None: + """Render a dependency-free reward curve whose source remains the CSV.""" + if not points: + raise ValueError("Cannot render an empty reward curve.") + width, height = 800, 420 + left, right, top, bottom = 70, 30, 35, 60 + plot_width = width - left - right + plot_height = height - top - bottom + min_step, max_step = points[0][0], points[-1][0] + + def x_position(step: int) -> float: + if min_step == max_step: + return left + plot_width / 2 + return left + (step - min_step) * plot_width / (max_step - min_step) + + def y_position(reward: float) -> float: + return top + (1.0 - reward) * plot_height + + polyline = " ".join(f"{x_position(step):.2f},{y_position(value):.2f}" for step, value in points) + circles = "\n".join( + f'' for step, value in points + ) + grid = "\n".join( + ( + f'' + f'{level:.2f}' + ) + for level in (0.0, 0.25, 0.5, 0.75, 1.0) + ) + svg = f""" + + +MemAgent rollout reward +{grid} + + + +{circles} +rollout_id ({min_step}..{max_step}) +reward + +""" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(svg, encoding="utf-8") + + def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--log-file", type=Path, required=True) @@ -112,6 +169,8 @@ def main() -> None: parser.add_argument("--metric", default=DEFAULT_METRIC) parser.add_argument("--expected-steps", type=int) parser.add_argument("--window-size", type=int, default=10) + parser.add_argument("--csv-output", type=Path) + parser.add_argument("--svg-output", type=Path) args = parser.parse_args() with args.log_file.open(encoding="utf-8", errors="replace") as source: @@ -127,6 +186,10 @@ def main() -> None: with args.output.open("w", encoding="utf-8") as destination: json.dump(summary, destination, ensure_ascii=False, indent=2) destination.write("\n") + if args.csv_output is not None: + write_reward_csv(args.csv_output, points) + if args.svg_output is not None: + write_reward_svg(args.svg_output, points) print(json.dumps(summary, ensure_ascii=False, indent=2)) diff --git a/tests/examples/mem_agent/test_compare_results.py b/tests/examples/mem_agent/test_compare_results.py index 44fa9b1ca..d4b0c81cc 100644 --- a/tests/examples/mem_agent/test_compare_results.py +++ b/tests/examples/mem_agent/test_compare_results.py @@ -33,6 +33,8 @@ def test_comparator_rejects_control_variable_mismatch(): "temperature": 0.7, "top_p": 0.95, "sampling_count": 1, + "seed": None, + "enable_thinking": None, "chunk_tokens": 2048, "max_memory_tokens": 1024, "max_final_tokens": 256, diff --git a/tests/examples/mem_agent/test_data_and_metrics.py b/tests/examples/mem_agent/test_data_and_metrics.py index f6b4eb34f..0f17b768d 100644 --- a/tests/examples/mem_agent/test_data_and_metrics.py +++ b/tests/examples/mem_agent/test_data_and_metrics.py @@ -36,6 +36,25 @@ def test_convert_row_supports_training_and_ruler_formats(): assert ruler["metadata"]["question"] == "Where?" +def test_convert_row_supports_official_hotpotqa_distractor_struct(): + row = convert_row( + { + "id": "hp-1", + "question": "Who?", + "answer": "Alice", + "level": "medium", + "context": { + "title": ["First", "Second"], + "sentences": [["Alice appears here."], ["Another paragraph."]], + }, + } + ) + + assert row["_id"] == "hp-1" + assert row["label"] == "Alice" + assert row["metadata"]["context"] == "First\nAlice appears here.\n\nSecond\nAnother paragraph." + + def test_convert_file_and_manifest_are_deterministic(tmp_path): source = tmp_path / "eval.json" source.write_text( diff --git a/tests/examples/mem_agent/test_eval.py b/tests/examples/mem_agent/test_eval.py index 0056e50aa..99e75eca2 100644 --- a/tests/examples/mem_agent/test_eval.py +++ b/tests/examples/mem_agent/test_eval.py @@ -6,7 +6,7 @@ import pytest -from examples.mem_agent.eval_ruler_hqa import base_infer, run_evaluation +from examples.mem_agent.eval_ruler_hqa import _request_seed, base_infer, run_evaluation, summarize_pass_at_n class CharacterTokenizer: @@ -23,8 +23,19 @@ def decode(self, token_ids, skip_special_tokens=True): async def test_base_infer_truncates_context_without_dropping_question(monkeypatch): captured = {} - async def fake_chat_once(session, base_url, api_key, model, instruction, temperature, top_p, max_tokens): - del session, base_url, api_key, model, temperature, top_p, max_tokens + async def fake_chat_once( + session, + base_url, + api_key, + model, + instruction, + temperature, + top_p, + max_tokens, + seed, + enable_thinking, + ): + del session, base_url, api_key, model, temperature, top_p, max_tokens, seed, enable_thinking captured["instruction"] = instruction return r"\boxed{x}" @@ -37,9 +48,14 @@ async def fake_chat_once(session, base_url, api_key, model, instruction, tempera temperature=0.0, top_p=1.0, max_final_tokens=16, + seed=42, + enable_thinking=False, ) _, diagnostics = await base_infer( - {"context": "c" * 500, "input": "Which answer?"}, args, CharacterTokenizer(), object() + {"_id": "q1", "context": "c" * 500, "input": "Which answer?"}, + args, + CharacterTokenizer(), + object(), ) assert diagnostics["context_truncated"] is True @@ -49,8 +65,8 @@ async def fake_chat_once(session, base_url, api_key, model, instruction, tempera @pytest.mark.asyncio async def test_evaluation_error_keeps_ground_truth_and_counts_as_zero(monkeypatch, tmp_path): - async def failed_infer(item, args, tokenizer, session): - del item, args, tokenizer, session + async def failed_infer(item, args, tokenizer, session, sample_index=0): + del item, args, tokenizer, session, sample_index raise RuntimeError("server unavailable") monkeypatch.setattr("examples.mem_agent.eval_ruler_hqa.recurrent_infer", failed_infer) @@ -71,6 +87,9 @@ async def failed_infer(item, args, tokenizer, session): max_chunks=64, max_input_tokens=7936, server_max_model_len=8192, + samples_per_item=2, + seed=42, + enable_thinking=False, ) records, summary = await run_evaluation( [{"_id": "q1", "input": "Question", "context": "Context", "answers": ["A", "Alias"]}], @@ -82,8 +101,33 @@ async def failed_infer(item, args, tokenizer, session): assert records[0]["pred"] == "" assert records[0]["judge_boxed_em"] == 0.0 assert "server unavailable" in records[0]["error"] - assert summary["total"] == 1 - assert summary["errors"] == 1 + assert summary["total"] == 2 + assert summary["errors"] == 2 assert summary["sub_em_pct"] == 0.0 assert len(summary["data_sha256"]) == 64 assert summary["evaluator_schema_version"] == "mem-agent-vime-eval-v1" + assert summary["sampling_count"] == 2 + assert summary["pass_at_n"] == 0.0 + + +def test_pass_at_n_requires_nonzero_reward_and_reports_group_variance(): + records = [{"_id": "q1", "sample_index": index, "judge_boxed_em": float(index == 0)} for index in range(4)] + records += [{"_id": "q2", "sample_index": index, "judge_boxed_em": 0.0} for index in range(4)] + + summary = summarize_pass_at_n(records, samples_per_item=4) + + assert summary["pass_at_n"] == 0.5 + assert summary["pass_at_n_pct"] == 50.0 + assert summary["reward_variance_groups"] == 1 + assert summary["mean_successes_per_prompt"] == 0.5 + + +def test_request_seed_is_stable_and_distinguishes_samples_and_turns(): + args = SimpleNamespace(seed=42) + + first = _request_seed(args, "q1", 0, "memory", 0) + + assert first == _request_seed(args, "q1", 0, "memory", 0) + assert first != _request_seed(args, "q1", 1, "memory", 0) + assert first != _request_seed(args, "q1", 0, "final", 1) + assert _request_seed(SimpleNamespace(seed=None), "q1", 0, "final", 0) is None diff --git a/tests/examples/mem_agent/test_mock_integration.py b/tests/examples/mem_agent/test_mock_integration.py index eaf67ab51..75efdf868 100644 --- a/tests/examples/mem_agent/test_mock_integration.py +++ b/tests/examples/mem_agent/test_mock_integration.py @@ -23,7 +23,8 @@ def decode(self, token_ids, skip_special_tokens=True): del skip_special_tokens return "".join(chr(token_id) for token_id in token_ids) - def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True): + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True, **kwargs): + del kwargs assert not tokenize and add_generation_prompt return f"{messages[0]['content']}" diff --git a/tests/examples/mem_agent/test_pilot_data.py b/tests/examples/mem_agent/test_pilot_data.py new file mode 100644 index 000000000..47396d3d2 --- /dev/null +++ b/tests/examples/mem_agent/test_pilot_data.py @@ -0,0 +1,98 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import pytest + +from examples.mem_agent.prepare_pilot_data import build_candidates, select_pilot_sets + + +class CharacterTokenizer: + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [ord(character) for character in text] + + +def _row(index: int, context: str, answer: str = "x"): + return { + "prompt": f"Question {index}?", + "label": answer, + "metadata": {"question": f"Question {index}?", "context": context, "ground_truth": [answer]}, + } + + +def test_candidate_filter_requires_short_multichunk_answer_visible_rows(): + rows = [ + _row(0, "abcx"), # one chunk: too short + _row(1, "abcdxefg"), # two chunks: eligible + _row(2, "abcdefghxijk"), # three chunks: eligible + _row(3, "abcdefghijklmnopx"), # five chunks: too long + _row(4, "abcdefgh", answer="z"), # answer absent + ] + + candidates, manifest = build_candidates( + rows, + CharacterTokenizer(), + chunk_tokens=4, + min_chunks=2, + max_chunks=3, + candidate_count=2, + seed=42, + ) + + assert {row["metadata"]["pilot"]["source_index"] for row in candidates} == {1, 2} + assert all(2 <= row["metadata"]["pilot"]["num_chunks"] <= 3 for row in candidates) + assert manifest["rejected"] == {"answer_not_in_context": 1, "too_long": 1, "too_short": 1} + + +def test_pass_at_n_selection_guarantees_success_failure_and_disjoint_split(): + candidates = [_row(index, f"abcdx{index}ef") | {"_id": f"q{index}"} for index in range(5)] + records = [] + success_counts = {"q0": 0, "q1": 1, "q2": 2, "q3": 3, "q4": 4} + for candidate in candidates: + for sample_index in range(4): + records.append( + { + "_id": candidate["_id"], + "sample_index": sample_index, + "judge_boxed_em": float(sample_index < success_counts[candidate["_id"]]), + } + ) + + train_rows, eval_rows, manifest = select_pilot_sets( + candidates, + records, + samples_per_item=4, + train_count=2, + eval_count=1, + seed=7, + preferred_min_successes=2, + preferred_max_successes=2, + ) + + selected = train_rows + eval_rows + assert len(train_rows) == 2 + assert len(eval_rows) == 1 + assert not ({row["_id"] for row in train_rows} & {row["_id"] for row in eval_rows}) + assert all(0 < row["metadata"]["pilot"]["baseline_successes"] < 4 for row in selected) + assert all(row["metadata"]["pilot"]["baseline_pass_at_n"] is True for row in selected) + assert {entry["status"] for entry in manifest["screening"]} >= {"preferred", "no_reward_variance"} + + +def test_pass_at_n_selection_fails_before_gpu_training_when_pool_is_too_hard(): + candidates = [_row(index, f"abcdx{index}ef") | {"_id": f"q{index}"} for index in range(2)] + records = [ + {"_id": candidate["_id"], "sample_index": sample_index, "judge_boxed_em": 0.0} + for candidate in candidates + for sample_index in range(4) + ] + + with pytest.raises(ValueError, match="non-degenerate Pass@4"): + select_pilot_sets( + candidates, + records, + samples_per_item=4, + train_count=1, + eval_count=1, + seed=42, + ) diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index 1fb3fbd66..b77e5279e 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -38,6 +38,37 @@ def test_train_script_keeps_trajectory_loss_and_real_turn_context_envelope(): assert "--calculate-per-token-loss" not in script +def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): + config = yaml.safe_load((EXAMPLE / "config-pilot-qwen3-0.6b.yaml").read_text(encoding="utf-8")) + train = (EXAMPLE / "run-qwen3-0.6B-train.sh").read_text(encoding="utf-8") + baseline = (EXAMPLE / "run-qwen3-0.6B-baseline.sh").read_text(encoding="utf-8") + evaluation = (EXAMPLE / "run-qwen3-0.6B-eval.sh").read_text(encoding="utf-8") + + assert config["model_id"] == "Qwen/Qwen3-0.6B" + assert config["model_revision"] == "c1899de289a04d12100db370d81485cdf75e47ca" + assert config["mem_agent_chunk_tokens"] == 512 + assert config["mem_agent_max_memory_tokens"] == 128 + assert config["mem_agent_max_final_tokens"] == 64 + assert config["mem_agent_max_chunks"] == 4 + assert config["mem_agent_enable_thinking"] is False + assert config["custom_train_sample_expansion_factor"] == 5 + assert config["mem_agent_train_rows_multiple"] == 1 + assert "qwen3-0.6B.sh" in train + assert '--resource \'{"actor": [1, 1], "rollout": [1, 1]}\'' in train + assert "--tensor-model-parallel-size 1" in train + assert "--rollout-max-context-len 1536" in train + assert "--max-tokens-per-gpu 1536" in train + assert "--log-passrate" in train + assert "pilot-selection.manifest.json" in train + assert "training-reward.svg" in train + assert '--expected-steps "${NUM_ROLLOUT}"' in train + assert '--samples-per-item "${SAMPLES_PER_ITEM}"' in baseline + assert 'prepare_pilot_data.py" select' in baseline + assert "--disable-thinking" in baseline + assert "pilot-boxed-em boxed_em_pct" in evaluation + assert "pilot-pass-at-n pass_at_n_pct" in evaluation + + def test_reward_exposes_a_step_level_raw_reward_metric(): reward_source = (EXAMPLE / "reward.py").read_text(encoding="utf-8") assert '"mem_agent_raw_reward": score' in reward_source diff --git a/tests/examples/mem_agent/test_rollout.py b/tests/examples/mem_agent/test_rollout.py index d849a9fef..52136df6e 100644 --- a/tests/examples/mem_agent/test_rollout.py +++ b/tests/examples/mem_agent/test_rollout.py @@ -13,6 +13,9 @@ class FakeTokenizer: + def __init__(self): + self.template_kwargs = [] + def encode(self, text, add_special_tokens=False): del add_special_tokens return [ord(character) for character in text] @@ -21,9 +24,10 @@ def decode(self, token_ids, skip_special_tokens=True): del skip_special_tokens return "".join(chr(token_id) for token_id in token_ids) - def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True): + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True, **kwargs): assert tokenize is False assert add_generation_prompt is True + self.template_kwargs.append(kwargs) return f"{messages[0]['content']}" @@ -136,6 +140,33 @@ async def fake_generate(args, turn, sampling_params, evaluation): assert all(" \nQ\n" not in prompt for prompt in prompts) +@pytest.mark.asyncio +async def test_generate_trajectory_can_disable_qwen_thinking_for_small_pilot(): + tokenizer = FakeTokenizer() + responses = iter(["MEM", r"\boxed{x}"]) + + async def fake_generate(args, turn, sampling_params, evaluation): + del args, sampling_params, evaluation + response = next(responses) + response_ids = tokenizer.encode(response) + turn.response = response + turn.tokens = tokenizer.encode(turn.prompt) + response_ids + turn.rollout_tokens = list(turn.tokens) + turn.response_length = len(response_ids) + turn.loss_mask = [1] * len(response_ids) + turn.rollout_log_probs = [-0.1] * len(response_ids) + turn.status = Sample.Status.COMPLETED + return turn + + args = _args() + args.mem_agent_enable_thinking = False + sample = Sample(index=0, group_index=0, prompt="Q", metadata={"context": "abc"}) + result = await generate_trajectory(args, sample, {}, tokenizer, generator=fake_generate) + + assert result.status == Sample.Status.COMPLETED + assert tokenizer.template_kwargs == [{"enable_thinking": False}, {"enable_thinking": False}] + + def test_truncate_text_to_tokens_retokenizes_to_the_hard_limit(): text, length = truncate_text_to_tokens(FakeTokenizer(), "MEMORY", 3) assert text == "MEM" diff --git a/tests/examples/mem_agent/test_summarize_reward.py b/tests/examples/mem_agent/test_summarize_reward.py index 1820ec5c8..d3107f088 100644 --- a/tests/examples/mem_agent/test_summarize_reward.py +++ b/tests/examples/mem_agent/test_summarize_reward.py @@ -5,7 +5,12 @@ import pytest -from examples.mem_agent.summarize_reward import extract_reward_points, summarize_reward_points +from examples.mem_agent.summarize_reward import ( + extract_reward_points, + summarize_reward_points, + write_reward_csv, + write_reward_svg, +) def test_extract_and_summarize_complete_reward_series(): @@ -44,3 +49,18 @@ def test_summary_rejects_missing_rollout_and_invalid_reward(): with pytest.raises(ValueError, match=r"outside \[0, 1\]"): extract_reward_points(["perf 0: {'rollout/mem_agent_raw_reward/mean': 1.1}\n"]) + + +def test_reward_csv_and_svg_keep_auditable_points(tmp_path): + points = [(0, 0.25), (1, 0.75)] + csv_path = tmp_path / "reward.csv" + svg_path = tmp_path / "reward.svg" + + write_reward_csv(csv_path, points) + write_reward_svg(svg_path, points) + + assert csv_path.read_text(encoding="utf-8").splitlines() == ["rollout_id,reward", "0,0.25", "1,0.75"] + svg = svg_path.read_text(encoding="utf-8") + assert "MemAgent rollout reward" in svg + assert ' Date: Tue, 4 Aug 2026 18:49:39 +0800 Subject: [PATCH 07/16] chore: capture Qwen3 pilot GPU telemetry --- examples/mem_agent/prepare_pilot_data.py | 10 ++++++---- examples/mem_agent/run-qwen3-0.6B-baseline.sh | 14 +++++++++++++- examples/mem_agent/run-qwen3-0.6B-eval.sh | 12 +++++++++++- examples/mem_agent/run-qwen3-0.6B-train.sh | 7 +++++++ tests/examples/mem_agent/test_recipe_contract.py | 3 +++ 5 files changed, 40 insertions(+), 6 deletions(-) diff --git a/examples/mem_agent/prepare_pilot_data.py b/examples/mem_agent/prepare_pilot_data.py index 3253a765b..02a2d90ec 100644 --- a/examples/mem_agent/prepare_pilot_data.py +++ b/examples/mem_agent/prepare_pilot_data.py @@ -2,8 +2,8 @@ """Build a short-context, non-degenerate Qwen3-0.6B MemAgent pilot set. The workflow is deliberately two-stage. ``candidates`` filters only immutable -input properties such as token length. After an untrained recurrent Pass@N -run, ``select`` keeps prompts with both successes and failures. This gives GRPO +input properties such as token length. After an untrained recurrent Pass@N run, +``select`` keeps prompts with both successes and failures. This gives GRPO useful within-group reward variance without pretending the screened pilot is a formal, unbiased HotpotQA benchmark. """ @@ -72,7 +72,8 @@ def build_candidates( seed: int, require_answer_in_context: bool = True, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: - """Filter by immutable length/content properties and sample deterministically.""" + """Filter by immutable length/content properties and sample + deterministically.""" if not 0 < min_chunks <= max_chunks: raise ValueError("Expected 0 < min_chunks <= max_chunks.") if chunk_tokens <= 0 or candidate_count <= 0: @@ -195,7 +196,8 @@ def select_pilot_sets( preferred_min_successes: int = 2, preferred_max_successes: int | None = None, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]: - """Select disjoint train/eval diagnostics with Pass@N and reward variance.""" + """Select disjoint train/eval diagnostics with Pass@N and reward + variance.""" if samples_per_item <= 1: raise ValueError("Pass@N screening requires samples_per_item > 1.") max_successes = samples_per_item - 1 if max_successes is None else max_successes diff --git a/examples/mem_agent/run-qwen3-0.6B-baseline.sh b/examples/mem_agent/run-qwen3-0.6B-baseline.sh index 07464d99c..c4fa4889b 100755 --- a/examples/mem_agent/run-qwen3-0.6B-baseline.sh +++ b/examples/mem_agent/run-qwen3-0.6B-baseline.sh @@ -23,6 +23,14 @@ export NO_PROXY="127.0.0.1,localhost,::1" export no_proxy="${NO_PROXY}" SERVER_LOG="${RESULTS_DIR}/qwen3-0.6b-baseline-server.log" +GPU_LOG="${RESULTS_DIR}/qwen3-0.6b-baseline-gpu.csv" +# Keep low-frequency GPU telemetry beside the raw model outputs so startup, +# OOM and idle/hang diagnoses do not depend on an external dashboard. +nvidia-smi \ + --query-gpu=timestamp,index,name,memory.used,memory.total,utilization.gpu,power.draw \ + --format=csv,nounits \ + --loop=5 >"${GPU_LOG}" 2>&1 & +GPU_MONITOR_PID=$! python3 -m sglang.launch_server \ --model-path "${MODEL_PATH}" \ --host 127.0.0.1 \ @@ -34,7 +42,11 @@ python3 -m sglang.launch_server \ --trust-remote-code \ >"${SERVER_LOG}" 2>&1 & SERVER_PID=$! -trap 'kill -TERM "${SERVER_PID}" 2>/dev/null || true; wait "${SERVER_PID}" 2>/dev/null || true' EXIT INT TERM +cleanup() { + kill -TERM "${SERVER_PID}" "${GPU_MONITOR_PID}" 2>/dev/null || true + wait "${SERVER_PID}" "${GPU_MONITOR_PID}" 2>/dev/null || true +} +trap cleanup EXIT INT TERM for _ in $(seq 1 120); do if ! kill -0 "${SERVER_PID}" 2>/dev/null; then diff --git a/examples/mem_agent/run-qwen3-0.6B-eval.sh b/examples/mem_agent/run-qwen3-0.6B-eval.sh index ad1011253..45c722823 100755 --- a/examples/mem_agent/run-qwen3-0.6B-eval.sh +++ b/examples/mem_agent/run-qwen3-0.6B-eval.sh @@ -21,6 +21,12 @@ export NO_PROXY="127.0.0.1,localhost,::1" export no_proxy="${NO_PROXY}" SERVER_LOG="${RESULTS_DIR}/${RUN_NAME}.server.log" +GPU_LOG="${RESULTS_DIR}/${RUN_NAME}.gpu.csv" +nvidia-smi \ + --query-gpu=timestamp,index,name,memory.used,memory.total,utilization.gpu,power.draw \ + --format=csv,nounits \ + --loop=5 >"${GPU_LOG}" 2>&1 & +GPU_MONITOR_PID=$! python3 -m sglang.launch_server \ --model-path "${MODEL_PATH}" \ --host 127.0.0.1 \ @@ -32,7 +38,11 @@ python3 -m sglang.launch_server \ --trust-remote-code \ >"${SERVER_LOG}" 2>&1 & SERVER_PID=$! -trap 'kill -TERM "${SERVER_PID}" 2>/dev/null || true; wait "${SERVER_PID}" 2>/dev/null || true' EXIT INT TERM +cleanup() { + kill -TERM "${SERVER_PID}" "${GPU_MONITOR_PID}" 2>/dev/null || true + wait "${SERVER_PID}" "${GPU_MONITOR_PID}" 2>/dev/null || true +} +trap cleanup EXIT INT TERM for _ in $(seq 1 120); do if ! kill -0 "${SERVER_PID}" 2>/dev/null; then diff --git a/examples/mem_agent/run-qwen3-0.6B-train.sh b/examples/mem_agent/run-qwen3-0.6B-train.sh index d7fe286eb..bb5921277 100755 --- a/examples/mem_agent/run-qwen3-0.6B-train.sh +++ b/examples/mem_agent/run-qwen3-0.6B-train.sh @@ -58,6 +58,13 @@ source "${MODEL_CONFIG_DIR}/qwen3-0.6B.sh" NOW="$(date '+%Y%m%dT%H%M%S%z')" LOG_FILE="${RUN_ROOT}/logs/${RUN_NAME}-${NOW}.log" +GPU_LOG="${RUN_ROOT}/logs/${RUN_NAME}-${NOW}-gpu.csv" +nvidia-smi \ + --query-gpu=timestamp,index,name,memory.used,memory.total,utilization.gpu,power.draw \ + --format=csv,nounits \ + --loop=5 >"${GPU_LOG}" 2>&1 & +GPU_MONITOR_PID=$! +trap 'kill -TERM "${GPU_MONITOR_PID}" 2>/dev/null || true; wait "${GPU_MONITOR_PID}" 2>/dev/null || true' EXIT INT TERM ray job submit --address="http://127.0.0.1:8265" \ --runtime-env-json="${RUNTIME_ENV_JSON}" \ diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index b77e5279e..2254316e2 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -67,6 +67,9 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): assert "--disable-thinking" in baseline assert "pilot-boxed-em boxed_em_pct" in evaluation assert "pilot-pass-at-n pass_at_n_pct" in evaluation + assert "qwen3-0.6b-baseline-gpu.csv" in baseline + assert 'GPU_LOG="${RUN_ROOT}/logs/${RUN_NAME}-${NOW}-gpu.csv"' in train + assert 'GPU_LOG="${RESULTS_DIR}/${RUN_NAME}.gpu.csv"' in evaluation def test_reward_exposes_a_step_level_raw_reward_metric(): From 16bf15336c24acac33be165a328cd9937f5335cb Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 19:23:06 +0800 Subject: [PATCH 08/16] fix: expose MemAgent package to pilot scripts --- examples/mem_agent/run-qwen3-0.6B-baseline.sh | 4 ++++ examples/mem_agent/run-qwen3-0.6B-eval.sh | 4 ++++ tests/examples/mem_agent/test_recipe_contract.py | 2 ++ 3 files changed, 10 insertions(+) diff --git a/examples/mem_agent/run-qwen3-0.6B-baseline.sh b/examples/mem_agent/run-qwen3-0.6B-baseline.sh index c4fa4889b..0b064847f 100755 --- a/examples/mem_agent/run-qwen3-0.6B-baseline.sh +++ b/examples/mem_agent/run-qwen3-0.6B-baseline.sh @@ -4,6 +4,7 @@ set -euo pipefail SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +RELAX_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to the frozen Qwen3-0.6B BF16 checkpoint.}" TOKENIZER_PATH="${TOKENIZER_PATH:-${MODEL_PATH}}" DATA_DIR="${DATA_DIR:?Set DATA_DIR to the Task 36 pilot data directory.}" @@ -21,6 +22,9 @@ mkdir -p "${RESULTS_DIR}" export CUDA_VISIBLE_DEVICES="${GPU_ID}" export NO_PROXY="127.0.0.1,localhost,::1" export no_proxy="${NO_PROXY}" +# The evaluator is executed by file path, so Python otherwise places only +# examples/mem_agent on sys.path and cannot import examples.mem_agent.*. +export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" SERVER_LOG="${RESULTS_DIR}/qwen3-0.6b-baseline-server.log" GPU_LOG="${RESULTS_DIR}/qwen3-0.6b-baseline-gpu.csv" diff --git a/examples/mem_agent/run-qwen3-0.6B-eval.sh b/examples/mem_agent/run-qwen3-0.6B-eval.sh index 45c722823..9f815a195 100755 --- a/examples/mem_agent/run-qwen3-0.6B-eval.sh +++ b/examples/mem_agent/run-qwen3-0.6B-eval.sh @@ -4,6 +4,7 @@ set -euo pipefail SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +RELAX_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to a converted Qwen3-0.6B checkpoint.}" TOKENIZER_PATH="${TOKENIZER_PATH:?Set TOKENIZER_PATH to the frozen base tokenizer.}" EVAL_DATA="${EVAL_DATA:?Set EVAL_DATA to the frozen pilot-eval.jsonl.}" @@ -19,6 +20,9 @@ mkdir -p "${RESULTS_DIR}" export CUDA_VISIBLE_DEVICES="${GPU_ID}" export NO_PROXY="127.0.0.1,localhost,::1" export no_proxy="${NO_PROXY}" +# Keep standalone file execution import-compatible with the package modules +# used by eval_ruler_hqa.py. +export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" SERVER_LOG="${RESULTS_DIR}/${RUN_NAME}.server.log" GPU_LOG="${RESULTS_DIR}/${RUN_NAME}.gpu.csv" diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index 2254316e2..58c9a8799 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -70,6 +70,8 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): assert "qwen3-0.6b-baseline-gpu.csv" in baseline assert 'GPU_LOG="${RUN_ROOT}/logs/${RUN_NAME}-${NOW}-gpu.csv"' in train assert 'GPU_LOG="${RESULTS_DIR}/${RUN_NAME}.gpu.csv"' in evaluation + assert 'export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"' in baseline + assert 'export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"' in evaluation def test_reward_exposes_a_step_level_raw_reward_metric(): From 83db8610a0db5c6a3fbf4bfd59534ac80ddcea06 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 19:29:13 +0800 Subject: [PATCH 09/16] feat: prepare thousand-prompt MemAgent pilot --- .../prepare-qwen3-0.6B-formal-data.sh | 33 ++++ examples/mem_agent/prepare_pilot_data.py | 150 ++++++++++++++++++ .../mem_agent/run-qwen3-0.6B-formal-screen.sh | 103 ++++++++++++ examples/mem_agent/run-qwen3-0.6B-train.sh | 5 +- tests/examples/mem_agent/test_pilot_data.py | 30 +++- .../mem_agent/test_recipe_contract.py | 12 +- 6 files changed, 329 insertions(+), 4 deletions(-) create mode 100755 examples/mem_agent/prepare-qwen3-0.6B-formal-data.sh create mode 100755 examples/mem_agent/run-qwen3-0.6B-formal-screen.sh diff --git a/examples/mem_agent/prepare-qwen3-0.6B-formal-data.sh b/examples/mem_agent/prepare-qwen3-0.6B-formal-data.sh new file mode 100755 index 000000000..c5e8a3de7 --- /dev/null +++ b/examples/mem_agent/prepare-qwen3-0.6B-formal-data.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +NORMALIZED_DATA="${NORMALIZED_DATA:?Set NORMALIZED_DATA to the frozen normalized HotpotQA JSONL.}" +TOKENIZER_PATH="${TOKENIZER_PATH:?Set TOKENIZER_PATH to the frozen Qwen3-0.6B tokenizer.}" +DATA_DIR="${DATA_DIR:?Set DATA_DIR to the formal Task 36 pilot directory.}" +TRAIN_CANDIDATE_COUNT="${TRAIN_CANDIDATE_COUNT:-4000}" +SMOKE_COUNT="${SMOKE_COUNT:-48}" +DIAGNOSTIC_COUNT="${DIAGNOSTIC_COUNT:-128}" +HELDOUT_COUNT="${HELDOUT_COUNT:-500}" + +mkdir -p "${DATA_DIR}" +python3 "${SCRIPT_DIR}/prepare_pilot_data.py" freeze \ + --input "${NORMALIZED_DATA}" \ + --tokenizer "${TOKENIZER_PATH}" \ + --train-candidates-output "${DATA_DIR}/formal-train-candidates.jsonl" \ + --smoke-output "${DATA_DIR}/formal-smoke-candidates.jsonl" \ + --diagnostic-output "${DATA_DIR}/formal-diagnostic.jsonl" \ + --heldout-output "${DATA_DIR}/formal-heldout.jsonl" \ + --manifest "${DATA_DIR}/formal-static-splits.manifest.json" \ + --chunk-tokens 512 \ + --min-chunks 2 \ + --max-chunks 4 \ + --train-candidate-count "${TRAIN_CANDIDATE_COUNT}" \ + --smoke-count "${SMOKE_COUNT}" \ + --diagnostic-count "${DIAGNOSTIC_COUNT}" \ + --heldout-count "${HELDOUT_COUNT}" \ + --seed 42 + +echo "Frozen mutually disjoint Task 36 smoke, train candidates, diagnostic and held-out IDs under ${DATA_DIR}." diff --git a/examples/mem_agent/prepare_pilot_data.py b/examples/mem_agent/prepare_pilot_data.py index 02a2d90ec..b32c7d5ed 100644 --- a/examples/mem_agent/prepare_pilot_data.py +++ b/examples/mem_agent/prepare_pilot_data.py @@ -137,6 +137,98 @@ def build_candidates( } +def freeze_static_splits( + rows: list[dict[str, Any]], + tokenizer: Any, + *, + chunk_tokens: int, + min_chunks: int, + max_chunks: int, + train_candidate_count: int, + smoke_count: int, + diagnostic_count: int, + heldout_count: int, + seed: int, +) -> tuple[ + list[dict[str, Any]], + list[dict[str, Any]], + list[dict[str, Any]], + list[dict[str, Any]], + dict[str, Any], +]: + """Freeze disjoint IDs before any outcome-based Pass@N screening. + + Held-out and diagnostic rows are placed first in the deterministic shuffle. + Increasing only the training candidate count therefore cannot silently + change either evaluation split. + """ + counts = { + "train_candidate": train_candidate_count, + "smoke": smoke_count, + "diagnostic": diagnostic_count, + "heldout": heldout_count, + } + if any(count <= 0 for count in counts.values()): + raise ValueError("All static split counts must be positive.") + + selected, filter_manifest = build_candidates( + rows, + tokenizer, + chunk_tokens=chunk_tokens, + min_chunks=min_chunks, + max_chunks=max_chunks, + candidate_count=sum(counts.values()), + seed=seed, + ) + ids = [str(row["_id"]) for row in selected] + if len(ids) != len(set(ids)): + raise ValueError("Static pilot candidates contain duplicate IDs.") + + heldout_end = heldout_count + diagnostic_end = heldout_end + diagnostic_count + smoke_end = diagnostic_end + smoke_count + heldout_rows = selected[:heldout_end] + diagnostic_rows = selected[heldout_end:diagnostic_end] + smoke_rows = selected[diagnostic_end:smoke_end] + train_candidates = selected[smoke_end:] + split_rows = { + "train_candidate": train_candidates, + "smoke": smoke_rows, + "diagnostic": diagnostic_rows, + "heldout": heldout_rows, + } + for split, split_items in split_rows.items(): + for row in split_items: + row.setdefault("metadata", {}).setdefault("pilot", {})["static_split"] = split + + split_ids = {split: [str(row["_id"]) for row in split_items] for split, split_items in split_rows.items()} + assert not (set(split_ids["train_candidate"]) & set(split_ids["diagnostic"])) + assert not (set(split_ids["train_candidate"]) & set(split_ids["heldout"])) + assert not (set(split_ids["train_candidate"]) & set(split_ids["smoke"])) + assert not (set(split_ids["smoke"]) & set(split_ids["diagnostic"])) + assert not (set(split_ids["smoke"]) & set(split_ids["heldout"])) + assert not (set(split_ids["diagnostic"]) & set(split_ids["heldout"])) + return ( + train_candidates, + smoke_rows, + diagnostic_rows, + heldout_rows, + { + "schema_version": "mem-agent-pilot-static-splits-v1", + "seed": seed, + "counts": counts, + "selected_ids": split_ids, + "pass_at_n_screening": { + "train_candidate": "pending", + "smoke": "pending", + "diagnostic": "forbidden", + "heldout": "forbidden", + }, + "filter": filter_manifest, + }, + ) + + def _screening_groups( candidates: list[dict[str, Any]], records: list[dict[str, Any]], @@ -310,6 +402,23 @@ def main() -> None: candidates.add_argument("--seed", type=int, default=42) candidates.add_argument("--allow-answer-not-in-context", action="store_true") + freeze = subparsers.add_parser("freeze") + freeze.add_argument("--input", type=Path, required=True) + freeze.add_argument("--tokenizer", required=True) + freeze.add_argument("--train-candidates-output", type=Path, required=True) + freeze.add_argument("--smoke-output", type=Path, required=True) + freeze.add_argument("--diagnostic-output", type=Path, required=True) + freeze.add_argument("--heldout-output", type=Path, required=True) + freeze.add_argument("--manifest", type=Path, required=True) + freeze.add_argument("--chunk-tokens", type=int, default=512) + freeze.add_argument("--min-chunks", type=int, default=2) + freeze.add_argument("--max-chunks", type=int, default=4) + freeze.add_argument("--train-candidate-count", type=int, default=4000) + freeze.add_argument("--smoke-count", type=int, default=48) + freeze.add_argument("--diagnostic-count", type=int, default=128) + freeze.add_argument("--heldout-count", type=int, default=500) + freeze.add_argument("--seed", type=int, default=42) + select = subparsers.add_parser("select") select.add_argument("--candidates", type=Path, required=True) select.add_argument("--baseline-records", type=Path, required=True) @@ -349,6 +458,47 @@ def main() -> None: _write_manifest(args.manifest, manifest) return + if args.command == "freeze": + from transformers import AutoTokenizer + + train_candidates, smoke_rows, diagnostic_rows, heldout_rows, manifest = freeze_static_splits( + read_jsonl(args.input), + AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True), + chunk_tokens=args.chunk_tokens, + min_chunks=args.min_chunks, + max_chunks=args.max_chunks, + train_candidate_count=args.train_candidate_count, + smoke_count=args.smoke_count, + diagnostic_count=args.diagnostic_count, + heldout_count=args.heldout_count, + seed=args.seed, + ) + outputs = { + "train_candidate": (args.train_candidates_output, train_candidates), + "smoke": (args.smoke_output, smoke_rows), + "diagnostic": (args.diagnostic_output, diagnostic_rows), + "heldout": (args.heldout_output, heldout_rows), + } + for _, (path, output_rows) in outputs.items(): + write_jsonl(path, output_rows) + manifest.update( + { + "source_file": str(args.input), + "source_sha256": sha256_file(args.input), + "outputs": { + split: { + "file": str(path), + "rows": len(output_rows), + "sha256": sha256_file(path), + } + for split, (path, output_rows) in outputs.items() + }, + "tokenizer": args.tokenizer, + } + ) + _write_manifest(args.manifest, manifest) + return + candidate_rows = read_jsonl(args.candidates) train_rows, eval_rows, manifest = select_pilot_sets( candidate_rows, diff --git a/examples/mem_agent/run-qwen3-0.6B-formal-screen.sh b/examples/mem_agent/run-qwen3-0.6B-formal-screen.sh new file mode 100755 index 000000000..49eb47f1d --- /dev/null +++ b/examples/mem_agent/run-qwen3-0.6B-formal-screen.sh @@ -0,0 +1,103 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +# Batch-screen only the pre-frozen training candidate IDs. Diagnostic and +# held-out files never enter this command, preventing baseline-outcome leakage. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +RELAX_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" +MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to the frozen Qwen3-0.6B checkpoint.}" +TOKENIZER_PATH="${TOKENIZER_PATH:-${MODEL_PATH}}" +DATA_DIR="${DATA_DIR:?Set DATA_DIR to the frozen formal pilot directory.}" +RESULTS_DIR="${RESULTS_DIR:?Set RESULTS_DIR to the formal screening output directory.}" +GPU_ID="${GPU_ID:-0}" +SERVE_PORT="${SERVE_PORT:-30000}" +SAMPLES_PER_ITEM="${SAMPLES_PER_ITEM:-8}" +TRAIN_COUNT="${TRAIN_COUNT:-1000}" +CONCURRENCY="${CONCURRENCY:-32}" +TRAIN_CANDIDATES="${DATA_DIR}/formal-train-candidates.jsonl" +STATIC_MANIFEST="${DATA_DIR}/formal-static-splits.manifest.json" + +[[ -f "${TRAIN_CANDIDATES}" ]] || { echo "Missing ${TRAIN_CANDIDATES}" >&2; exit 1; } +[[ -f "${STATIC_MANIFEST}" ]] || { echo "Missing ${STATIC_MANIFEST}" >&2; exit 1; } +mkdir -p "${RESULTS_DIR}" +export CUDA_VISIBLE_DEVICES="${GPU_ID}" +export NO_PROXY="127.0.0.1,localhost,::1" +export no_proxy="${NO_PROXY}" +export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +SERVER_LOG="${RESULTS_DIR}/qwen3-0.6b-formal-screen-server.log" +GPU_LOG="${RESULTS_DIR}/qwen3-0.6b-formal-screen-gpu.csv" +nvidia-smi \ + --query-gpu=timestamp,index,name,memory.used,memory.total,utilization.gpu,power.draw \ + --format=csv,nounits \ + --loop=5 >"${GPU_LOG}" 2>&1 & +GPU_MONITOR_PID=$! +python3 -m sglang.launch_server \ + --model-path "${MODEL_PATH}" \ + --host 127.0.0.1 \ + --port "${SERVE_PORT}" \ + --api-key EMPTY \ + --tp-size 1 \ + --context-length 1536 \ + --mem-fraction-static 0.55 \ + --trust-remote-code \ + >"${SERVER_LOG}" 2>&1 & +SERVER_PID=$! +cleanup() { + kill -TERM "${SERVER_PID}" "${GPU_MONITOR_PID}" 2>/dev/null || true + wait "${SERVER_PID}" "${GPU_MONITOR_PID}" 2>/dev/null || true +} +trap cleanup EXIT INT TERM + +for _ in $(seq 1 120); do + if ! kill -0 "${SERVER_PID}" 2>/dev/null; then + echo "SGLang exited early; see ${SERVER_LOG}" >&2 + exit 1 + fi + if curl --noproxy '*' -fsS "http://127.0.0.1:${SERVE_PORT}/health" >/dev/null; then + break + fi + sleep 5 +done +curl --noproxy '*' -fsS "http://127.0.0.1:${SERVE_PORT}/health" >/dev/null + +python3 "${SCRIPT_DIR}/eval_ruler_hqa.py" \ + --data-file "${TRAIN_CANDIDATES}" \ + --model "${MODEL_PATH}" \ + --tokenizer "${TOKENIZER_PATH}" \ + --output-dir "${RESULTS_DIR}" \ + --run-name qwen3-0.6b-formal-train-candidates-pass8 \ + --mode recurrent \ + --base-url "http://127.0.0.1:${SERVE_PORT}/v1" \ + --api-key EMPTY \ + --samples-per-item "${SAMPLES_PER_ITEM}" \ + --seed 42 \ + --disable-thinking \ + --temperature 1.0 \ + --top-p 1.0 \ + --chunk-tokens 512 \ + --max-memory-tokens 128 \ + --max-final-tokens 64 \ + --max-chunks 4 \ + --max-input-tokens 1472 \ + --server-max-model-len 1536 \ + --concurrency "${CONCURRENCY}" + +python3 "${SCRIPT_DIR}/prepare_pilot_data.py" select \ + --candidates "${TRAIN_CANDIDATES}" \ + --baseline-records "${RESULTS_DIR}/qwen3-0.6b-formal-train-candidates-pass8.jsonl" \ + --train-output "${DATA_DIR}/formal-train.jsonl" \ + --eval-output "${DATA_DIR}/formal-screen-unused-eval.jsonl" \ + --manifest "${DATA_DIR}/formal-selection.manifest.json" \ + --samples-per-item "${SAMPLES_PER_ITEM}" \ + --train-count "${TRAIN_COUNT}" \ + --eval-count 0 \ + --seed 42 + +[[ "$(wc -l < "${DATA_DIR}/formal-train.jsonl")" -eq "${TRAIN_COUNT}" ]] || { + echo "Formal training row count did not match ${TRAIN_COUNT}." >&2 + exit 1 +} +echo "Selected ${TRAIN_COUNT} independent training questions; diagnostic/held-out IDs remained outcome-independent." diff --git a/examples/mem_agent/run-qwen3-0.6B-train.sh b/examples/mem_agent/run-qwen3-0.6B-train.sh index bb5921277..7a64df8ad 100755 --- a/examples/mem_agent/run-qwen3-0.6B-train.sh +++ b/examples/mem_agent/run-qwen3-0.6B-train.sh @@ -11,6 +11,7 @@ DATA_DIR="${DATA_DIR:?Set DATA_DIR to the screened Task 36 pilot data directory. RUN_ROOT="${RUN_ROOT:?Set RUN_ROOT to the Task 36 experiment output directory.}" SAVE_DIR="${SAVE_DIR:-${RUN_ROOT}/checkpoints}" TRAIN_DATA="${TRAIN_DATA:-${DATA_DIR}/pilot-train.jsonl}" +SELECTION_MANIFEST="${SELECTION_MANIFEST:-${DATA_DIR}/pilot-selection.manifest.json}" NUM_ROLLOUT="${NUM_ROLLOUT:-20}" SAVE_INTERVAL="${SAVE_INTERVAL:-5}" ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-1}" @@ -24,8 +25,8 @@ RUN_NAME="${RUN_NAME:-mem-agent-qwen3-0.6b-pilot}" echo "Missing Pass@N-screened pilot data: ${TRAIN_DATA}" >&2 exit 1 } -[[ -f "${DATA_DIR}/pilot-selection.manifest.json" ]] || { - echo "Missing pilot selection manifest; baseline screening must run before training." >&2 +[[ -f "${SELECTION_MANIFEST}" ]] || { + echo "Missing selection manifest: ${SELECTION_MANIFEST}; baseline screening must run before training." >&2 exit 1 } [[ -f "${MODEL_PATH}/config.json" ]] || { echo "Missing model config: ${MODEL_PATH}/config.json" >&2; exit 1; } diff --git a/tests/examples/mem_agent/test_pilot_data.py b/tests/examples/mem_agent/test_pilot_data.py index 47396d3d2..fe070b6ed 100644 --- a/tests/examples/mem_agent/test_pilot_data.py +++ b/tests/examples/mem_agent/test_pilot_data.py @@ -4,7 +4,7 @@ import pytest -from examples.mem_agent.prepare_pilot_data import build_candidates, select_pilot_sets +from examples.mem_agent.prepare_pilot_data import build_candidates, freeze_static_splits, select_pilot_sets class CharacterTokenizer: @@ -96,3 +96,31 @@ def test_pass_at_n_selection_fails_before_gpu_training_when_pool_is_too_hard(): eval_count=1, seed=42, ) + + +def test_static_splits_are_disjoint_and_freeze_heldout_before_screening(): + rows = [_row(index, f"abcdx{index}ef") | {"_id": f"q{index}"} for index in range(10)] + + train, smoke, diagnostic, heldout, manifest = freeze_static_splits( + rows, + CharacterTokenizer(), + chunk_tokens=4, + min_chunks=2, + max_chunks=3, + train_candidate_count=4, + smoke_count=1, + diagnostic_count=2, + heldout_count=2, + seed=42, + ) + + split_ids = [{row["_id"] for row in split} for split in (train, smoke, diagnostic, heldout)] + assert [len(split) for split in (train, smoke, diagnostic, heldout)] == [4, 1, 2, 2] + assert sum(len(left & right) for index, left in enumerate(split_ids) for right in split_ids[index + 1 :]) == 0 + assert all(row["metadata"]["pilot"]["static_split"] == "heldout" for row in heldout) + assert manifest["pass_at_n_screening"] == { + "train_candidate": "pending", + "smoke": "pending", + "diagnostic": "forbidden", + "heldout": "forbidden", + } diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index 58c9a8799..3fff1f30f 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -43,6 +43,8 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): train = (EXAMPLE / "run-qwen3-0.6B-train.sh").read_text(encoding="utf-8") baseline = (EXAMPLE / "run-qwen3-0.6B-baseline.sh").read_text(encoding="utf-8") evaluation = (EXAMPLE / "run-qwen3-0.6B-eval.sh").read_text(encoding="utf-8") + formal_prepare = (EXAMPLE / "prepare-qwen3-0.6B-formal-data.sh").read_text(encoding="utf-8") + formal_screen = (EXAMPLE / "run-qwen3-0.6B-formal-screen.sh").read_text(encoding="utf-8") assert config["model_id"] == "Qwen/Qwen3-0.6B" assert config["model_revision"] == "c1899de289a04d12100db370d81485cdf75e47ca" @@ -59,7 +61,7 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): assert "--rollout-max-context-len 1536" in train assert "--max-tokens-per-gpu 1536" in train assert "--log-passrate" in train - assert "pilot-selection.manifest.json" in train + assert 'SELECTION_MANIFEST="${SELECTION_MANIFEST:-${DATA_DIR}/pilot-selection.manifest.json}"' in train assert "training-reward.svg" in train assert '--expected-steps "${NUM_ROLLOUT}"' in train assert '--samples-per-item "${SAMPLES_PER_ITEM}"' in baseline @@ -72,6 +74,14 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): assert 'GPU_LOG="${RESULTS_DIR}/${RUN_NAME}.gpu.csv"' in evaluation assert 'export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"' in baseline assert 'export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"' in evaluation + assert 'TRAIN_CANDIDATE_COUNT="${TRAIN_CANDIDATE_COUNT:-4000}"' in formal_prepare + assert 'SMOKE_COUNT="${SMOKE_COUNT:-48}"' in formal_prepare + assert 'DIAGNOSTIC_COUNT="${DIAGNOSTIC_COUNT:-128}"' in formal_prepare + assert 'HELDOUT_COUNT="${HELDOUT_COUNT:-500}"' in formal_prepare + assert 'TRAIN_COUNT="${TRAIN_COUNT:-1000}"' in formal_screen + assert 'CONCURRENCY="${CONCURRENCY:-32}"' in formal_screen + assert "--eval-count 0" in formal_screen + assert "formal-heldout.jsonl" not in formal_screen def test_reward_exposes_a_step_level_raw_reward_metric(): From b06356433b275312861ef52a31fa91695581d5d0 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 20:12:08 +0800 Subject: [PATCH 10/16] feat: resume MemAgent pilot training --- examples/mem_agent/run-qwen3-0.6B-train.sh | 25 ++++++++++++++++++- examples/mem_agent/summarize_reward.py | 14 ++++++++++- .../mem_agent/test_recipe_contract.py | 6 ++++- .../mem_agent/test_summarize_reward.py | 14 +++++++++++ 4 files changed, 56 insertions(+), 3 deletions(-) diff --git a/examples/mem_agent/run-qwen3-0.6B-train.sh b/examples/mem_agent/run-qwen3-0.6B-train.sh index 7a64df8ad..1ebfdd934 100755 --- a/examples/mem_agent/run-qwen3-0.6B-train.sh +++ b/examples/mem_agent/run-qwen3-0.6B-train.sh @@ -20,6 +20,8 @@ GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}" GPU_ID="${GPU_ID:-0}" RAY_NUM_CPUS="${RAY_NUM_CPUS:-12}" RUN_NAME="${RUN_NAME:-mem-agent-qwen3-0.6b-pilot}" +LOAD_PATH="${LOAD_PATH:-}" +START_ROLLOUT_ID="${START_ROLLOUT_ID:-0}" [[ -f "${TRAIN_DATA}" ]] || { echo "Missing Pass@N-screened pilot data: ${TRAIN_DATA}" >&2 @@ -30,6 +32,18 @@ RUN_NAME="${RUN_NAME:-mem-agent-qwen3-0.6b-pilot}" exit 1 } [[ -f "${MODEL_PATH}/config.json" ]] || { echo "Missing model config: ${MODEL_PATH}/config.json" >&2; exit 1; } +((START_ROLLOUT_ID >= 0 && START_ROLLOUT_ID < NUM_ROLLOUT)) || { + echo "START_ROLLOUT_ID must be in [0, NUM_ROLLOUT)." >&2 + exit 1 +} +if ((START_ROLLOUT_ID > 0)) && [[ -z "${LOAD_PATH}" ]]; then + echo "A positive START_ROLLOUT_ID requires LOAD_PATH." >&2 + exit 1 +fi +if [[ -n "${LOAD_PATH}" && ! -d "${LOAD_PATH}" ]]; then + echo "Missing resume checkpoint directory: ${LOAD_PATH}" >&2 + exit 1 +fi mkdir -p "${RUN_ROOT}/logs" "${RUN_ROOT}/tensorboard" "${SAVE_DIR}" export CUDA_VISIBLE_DEVICES="${GPU_ID}" @@ -57,6 +71,13 @@ fi unset -f ray source "${MODEL_CONFIG_DIR}/qwen3-0.6B.sh" +RESUME_ARGS=() +if [[ -n "${LOAD_PATH}" ]]; then + # --num-rollout remains the total target. ReLax restores optimizer/model + # state from LOAD_PATH and continues at this explicit rollout id. + RESUME_ARGS+=(--load "${LOAD_PATH}" --start-rollout-id "${START_ROLLOUT_ID}") +fi + NOW="$(date '+%Y%m%dT%H%M%S%z')" LOG_FILE="${RUN_ROOT}/logs/${RUN_NAME}-${NOW}.log" GPU_LOG="${RUN_ROOT}/logs/${RUN_NAME}-${NOW}-gpu.csv" @@ -73,6 +94,7 @@ ray job submit --address="http://127.0.0.1:8265" \ --resource '{"actor": [1, 1], "rollout": [1, 1]}' \ "${MODEL_ARGS[@]}" \ --hf-checkpoint "${MODEL_PATH}" \ + "${RESUME_ARGS[@]}" \ --ref-load "${MODEL_PATH}" \ --megatron-to-hf-mode bridge \ --warm-hf-checkpoint-page-cache \ @@ -153,5 +175,6 @@ python3 "${SCRIPT_DIR}/summarize_reward.py" \ --output "${RUN_ROOT}/training-reward.summary.json" \ --csv-output "${RUN_ROOT}/training-reward.csv" \ --svg-output "${RUN_ROOT}/training-reward.svg" \ - --expected-steps "${NUM_ROLLOUT}" \ + --expected-steps "$((NUM_ROLLOUT - START_ROLLOUT_ID))" \ + --expected-start "${START_ROLLOUT_ID}" \ --window-size 5 diff --git a/examples/mem_agent/summarize_reward.py b/examples/mem_agent/summarize_reward.py index c996e0d19..807e30b25 100644 --- a/examples/mem_agent/summarize_reward.py +++ b/examples/mem_agent/summarize_reward.py @@ -61,6 +61,7 @@ def summarize_reward_points( points: list[tuple[int, float]], *, expected_steps: int | None = None, + expected_start: int = 0, window_size: int = 10, metric: str = DEFAULT_METRIC, ) -> dict[str, Any]: @@ -76,7 +77,9 @@ def summarize_reward_points( if expected_steps is not None: if expected_steps <= 0: raise ValueError("expected_steps must be positive.") - expected_ids = list(range(expected_steps)) + if expected_start < 0: + raise ValueError("expected_start must be non-negative.") + expected_ids = list(range(expected_start, expected_start + expected_steps)) if rollout_ids != expected_ids: raise ValueError(f"Reward rollout ids are incomplete: expected {expected_ids}, got {rollout_ids}.") @@ -91,6 +94,8 @@ def summarize_reward_points( "schema_version": SCHEMA_VERSION, "metric": metric, "num_steps": len(points), + "first_rollout_id": rollout_ids[0], + "last_rollout_id": rollout_ids[-1], "window_size_requested": window_size, "window_size_used": effective_window, "first_window_mean": first_mean, @@ -168,6 +173,12 @@ def main() -> None: parser.add_argument("--output", type=Path, required=True) parser.add_argument("--metric", default=DEFAULT_METRIC) parser.add_argument("--expected-steps", type=int) + parser.add_argument( + "--expected-start", + type=int, + default=0, + help="First expected rollout id for a resumed phase (default: 0).", + ) parser.add_argument("--window-size", type=int, default=10) parser.add_argument("--csv-output", type=Path) parser.add_argument("--svg-output", type=Path) @@ -178,6 +189,7 @@ def main() -> None: summary = summarize_reward_points( points, expected_steps=args.expected_steps, + expected_start=args.expected_start, window_size=args.window_size, metric=args.metric, ) diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index 3fff1f30f..1c056a486 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -63,7 +63,11 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): assert "--log-passrate" in train assert 'SELECTION_MANIFEST="${SELECTION_MANIFEST:-${DATA_DIR}/pilot-selection.manifest.json}"' in train assert "training-reward.svg" in train - assert '--expected-steps "${NUM_ROLLOUT}"' in train + assert "--start-rollout-id" in train + assert "--load" in train + assert "NUM_ROLLOUT - START_ROLLOUT_ID" in train + assert '--expected-steps "$((NUM_ROLLOUT - START_ROLLOUT_ID))"' in train + assert '--expected-start "${START_ROLLOUT_ID}"' in train assert '--samples-per-item "${SAMPLES_PER_ITEM}"' in baseline assert 'prepare_pilot_data.py" select' in baseline assert "--disable-thinking" in baseline diff --git a/tests/examples/mem_agent/test_summarize_reward.py b/tests/examples/mem_agent/test_summarize_reward.py index d3107f088..9e825c59e 100644 --- a/tests/examples/mem_agent/test_summarize_reward.py +++ b/tests/examples/mem_agent/test_summarize_reward.py @@ -51,6 +51,20 @@ def test_summary_rejects_missing_rollout_and_invalid_reward(): extract_reward_points(["perf 0: {'rollout/mem_agent_raw_reward/mean': 1.1}\n"]) +def test_summary_accepts_complete_resumed_rollout_range(): + summary = summarize_reward_points( + [(20, 0.25), (21, 0.5), (22, 0.75)], + expected_steps=3, + expected_start=20, + window_size=1, + ) + + assert summary["first_rollout_id"] == 20 + assert summary["last_rollout_id"] == 22 + with pytest.raises(ValueError, match="incomplete"): + summarize_reward_points([(20, 0.25), (22, 0.75)], expected_steps=3, expected_start=20) + + def test_reward_csv_and_svg_keep_auditable_points(tmp_path): points = [(0, 0.25), (1, 0.75)] csv_path = tmp_path / "reward.csv" From a1d4e1a79282bce9eb0451a27af5f24e697b6550 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 20:14:53 +0800 Subject: [PATCH 11/16] feat: pair formal pilot heldout baseline --- .../mem_agent/run-qwen3-0.6B-formal-screen.sh | 29 +++++++++++++++++++ .../mem_agent/test_recipe_contract.py | 4 ++- 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/examples/mem_agent/run-qwen3-0.6B-formal-screen.sh b/examples/mem_agent/run-qwen3-0.6B-formal-screen.sh index 49eb47f1d..4cc6db446 100755 --- a/examples/mem_agent/run-qwen3-0.6B-formal-screen.sh +++ b/examples/mem_agent/run-qwen3-0.6B-formal-screen.sh @@ -16,10 +16,13 @@ SERVE_PORT="${SERVE_PORT:-30000}" SAMPLES_PER_ITEM="${SAMPLES_PER_ITEM:-8}" TRAIN_COUNT="${TRAIN_COUNT:-1000}" CONCURRENCY="${CONCURRENCY:-32}" +HELDOUT_SAMPLES_PER_ITEM="${HELDOUT_SAMPLES_PER_ITEM:-1}" TRAIN_CANDIDATES="${DATA_DIR}/formal-train-candidates.jsonl" +HELDOUT_DATA="${DATA_DIR}/formal-heldout.jsonl" STATIC_MANIFEST="${DATA_DIR}/formal-static-splits.manifest.json" [[ -f "${TRAIN_CANDIDATES}" ]] || { echo "Missing ${TRAIN_CANDIDATES}" >&2; exit 1; } +[[ -f "${HELDOUT_DATA}" ]] || { echo "Missing ${HELDOUT_DATA}" >&2; exit 1; } [[ -f "${STATIC_MANIFEST}" ]] || { echo "Missing ${STATIC_MANIFEST}" >&2; exit 1; } mkdir -p "${RESULTS_DIR}" export CUDA_VISIBLE_DEVICES="${GPU_ID}" @@ -100,4 +103,30 @@ python3 "${SCRIPT_DIR}/prepare_pilot_data.py" select \ echo "Formal training row count did not match ${TRAIN_COUNT}." >&2 exit 1 } + +# Measure the already-frozen held-out split once while the untrained server is +# resident. These outcomes never feed selection; the trained checkpoint later +# uses this same file, seed and sampling count for a paired pilot comparison. +python3 "${SCRIPT_DIR}/eval_ruler_hqa.py" \ + --data-file "${HELDOUT_DATA}" \ + --model "${MODEL_PATH}" \ + --tokenizer "${TOKENIZER_PATH}" \ + --output-dir "${RESULTS_DIR}" \ + --run-name qwen3-0.6b-formal-heldout-baseline \ + --mode recurrent \ + --base-url "http://127.0.0.1:${SERVE_PORT}/v1" \ + --api-key EMPTY \ + --samples-per-item "${HELDOUT_SAMPLES_PER_ITEM}" \ + --seed 4242 \ + --disable-thinking \ + --temperature 1.0 \ + --top-p 1.0 \ + --chunk-tokens 512 \ + --max-memory-tokens 128 \ + --max-final-tokens 64 \ + --max-chunks 4 \ + --max-input-tokens 1472 \ + --server-max-model-len 1536 \ + --concurrency "${CONCURRENCY}" + echo "Selected ${TRAIN_COUNT} independent training questions; diagnostic/held-out IDs remained outcome-independent." diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index 1c056a486..af3e5134f 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -85,7 +85,9 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): assert 'TRAIN_COUNT="${TRAIN_COUNT:-1000}"' in formal_screen assert 'CONCURRENCY="${CONCURRENCY:-32}"' in formal_screen assert "--eval-count 0" in formal_screen - assert "formal-heldout.jsonl" not in formal_screen + assert 'HELDOUT_DATA="${DATA_DIR}/formal-heldout.jsonl"' in formal_screen + assert "qwen3-0.6b-formal-heldout-baseline" in formal_screen + assert 'HELDOUT_SAMPLES_PER_ITEM="${HELDOUT_SAMPLES_PER_ITEM:-1}"' in formal_screen def test_reward_exposes_a_step_level_raw_reward_metric(): From da7675b0b2af955143b67252904af3da91ceb8f9 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 20:34:04 +0800 Subject: [PATCH 12/16] fix: initialize single-node NCCL override --- examples/mem_agent/run-qwen3-0.6B-train.sh | 4 ++++ tests/examples/mem_agent/test_recipe_contract.py | 1 + 2 files changed, 5 insertions(+) diff --git a/examples/mem_agent/run-qwen3-0.6B-train.sh b/examples/mem_agent/run-qwen3-0.6B-train.sh index 1ebfdd934..a9d53a379 100755 --- a/examples/mem_agent/run-qwen3-0.6B-train.sh +++ b/examples/mem_agent/run-qwen3-0.6B-train.sh @@ -55,6 +55,10 @@ export MKL_NUM_THREADS=1 export OPENBLAS_NUM_THREADS=1 export NUMEXPR_NUM_THREADS=1 export TOKENIZERS_PARALLELISM=false +# local.sh currently reads this variable without a nounset-safe default before +# replacing it with the detected HAS_NVLINK value. Preserve the same empty +# semantic explicitly so this `set -u` example can source the common entrypoint. +export NCCL_NVLS_ENABLE="${NCCL_NVLS_ENABLE:-}" # PPIO exposes many logical CPUs under a much smaller cgroup pids.max. Limit # only `ray start`; all other Ray CLI calls retain their original arguments. diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index af3e5134f..f1750124f 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -75,6 +75,7 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): assert "pilot-pass-at-n pass_at_n_pct" in evaluation assert "qwen3-0.6b-baseline-gpu.csv" in baseline assert 'GPU_LOG="${RUN_ROOT}/logs/${RUN_NAME}-${NOW}-gpu.csv"' in train + assert 'export NCCL_NVLS_ENABLE="${NCCL_NVLS_ENABLE:-}"' in train assert 'GPU_LOG="${RESULTS_DIR}/${RUN_NAME}.gpu.csv"' in evaluation assert 'export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"' in baseline assert 'export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"' in evaluation From 767fdbe59821c58bbd5a16cb5cf395b4d9398544 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 20:38:47 +0800 Subject: [PATCH 13/16] fix: source local entrypoint without nounset --- examples/mem_agent/run-qwen3-0.6B-train.sh | 10 ++++++---- tests/examples/mem_agent/test_recipe_contract.py | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/examples/mem_agent/run-qwen3-0.6B-train.sh b/examples/mem_agent/run-qwen3-0.6B-train.sh index a9d53a379..42996d1a1 100755 --- a/examples/mem_agent/run-qwen3-0.6B-train.sh +++ b/examples/mem_agent/run-qwen3-0.6B-train.sh @@ -55,10 +55,6 @@ export MKL_NUM_THREADS=1 export OPENBLAS_NUM_THREADS=1 export NUMEXPR_NUM_THREADS=1 export TOKENIZERS_PARALLELISM=false -# local.sh currently reads this variable without a nounset-safe default before -# replacing it with the detected HAS_NVLINK value. Preserve the same empty -# semantic explicitly so this `set -u` example can source the common entrypoint. -export NCCL_NVLS_ENABLE="${NCCL_NVLS_ENABLE:-}" # PPIO exposes many logical CPUs under a much smaller cgroup pids.max. Limit # only `ray start`; all other Ray CLI calls retain their original arguments. @@ -70,7 +66,13 @@ ray() { fi } if [[ -z "${RELAX_ENTRYPOINT_MODE:-}" ]]; then + # The shared entrypoint treats several unset networking overrides as empty, + # but reads them without nounset-safe expansion. Suspend only `-u` while it + # is sourced; `-e`/pipefail remain active, and strict mode is restored before + # any Task36 argument construction or job submission. + set +u source "${RELAX_ROOT}/scripts/entrypoint/local.sh" + set -u fi unset -f ray source "${MODEL_CONFIG_DIR}/qwen3-0.6B.sh" diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index f1750124f..10e82f8eb 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -75,7 +75,7 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): assert "pilot-pass-at-n pass_at_n_pct" in evaluation assert "qwen3-0.6b-baseline-gpu.csv" in baseline assert 'GPU_LOG="${RUN_ROOT}/logs/${RUN_NAME}-${NOW}-gpu.csv"' in train - assert 'export NCCL_NVLS_ENABLE="${NCCL_NVLS_ENABLE:-}"' in train + assert 'set +u\n source "${RELAX_ROOT}/scripts/entrypoint/local.sh"\n set -u' in train assert 'GPU_LOG="${RESULTS_DIR}/${RUN_NAME}.gpu.csv"' in evaluation assert 'export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"' in baseline assert 'export PYTHONPATH="${RELAX_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"' in evaluation From cdf04756f4ebd9dc1fc2e421c0ae6feb8aa01732 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 21:06:10 +0800 Subject: [PATCH 14/16] fix: derive reward curve from rollout artifacts --- examples/mem_agent/README.md | 2 +- examples/mem_agent/run-qwen3-0.6B-train.sh | 1 + examples/mem_agent/summarize_reward.py | 76 ++++++++++++++++++- .../mem_agent/test_summarize_reward.py | 28 +++++++ 4 files changed, 103 insertions(+), 4 deletions(-) diff --git a/examples/mem_agent/README.md b/examples/mem_agent/README.md index 96cf7af56..e0ef1e71c 100644 --- a/examples/mem_agent/README.md +++ b/examples/mem_agent/README.md @@ -2,7 +2,7 @@ This example trains Qwen3-4B to update a bounded textual memory while reading a long document chunk by chunk. Every memory-update turn and the final-answer turn is saved as an independent training row. Only the final boxed answer receives a rule-based reward; GRPO normalization happens before the trajectory is expanded. -The training log reports the trajectory-level 0/1 outcome as `rollout/mem_agent_raw_reward/mean` on every rollout step. This diagnostic mirrors the primary `score` exactly but is not consumed by GRPO. `run-pipeline.sh` then runs `summarize_reward.py` and writes `training-reward.summary.json`, containing every raw point, the first/last-window means, their delta, and the peak. It rejects a run whose rollout ids are incomplete instead of producing a partial trend. +The reward summary reads the trajectory-level 0/1 outcome from `rollout_result/train/.jsonl`, before memory turns are expanded. This avoids weighting long trajectories more heavily and is robust to Ray log de-duplication. `run-pipeline.sh` then runs `summarize_reward.py` and writes `training-reward.summary.json`, containing every raw point, the first/last-window means, their delta, and the peak. It rejects a run whose rollout ids are incomplete instead of producing a partial trend. The reproducibility contract is frozen to: diff --git a/examples/mem_agent/run-qwen3-0.6B-train.sh b/examples/mem_agent/run-qwen3-0.6B-train.sh index 42996d1a1..d3b16a734 100755 --- a/examples/mem_agent/run-qwen3-0.6B-train.sh +++ b/examples/mem_agent/run-qwen3-0.6B-train.sh @@ -178,6 +178,7 @@ ray job list >"${RUN_ROOT}/logs/${RUN_NAME}-${NOW}-ray-job-list.txt" 2>&1 # together. Missing/duplicated/conflicting rollout ids make this step fail. python3 "${SCRIPT_DIR}/summarize_reward.py" \ --log-file "${LOG_FILE}" \ + --rollout-result-dir "${SAVE_DIR}/rollout_result/train" \ --output "${RUN_ROOT}/training-reward.summary.json" \ --csv-output "${RUN_ROOT}/training-reward.csv" \ --svg-output "${RUN_ROOT}/training-reward.svg" \ diff --git a/examples/mem_agent/summarize_reward.py b/examples/mem_agent/summarize_reward.py index 807e30b25..e55afb010 100644 --- a/examples/mem_agent/summarize_reward.py +++ b/examples/mem_agent/summarize_reward.py @@ -16,6 +16,7 @@ DEFAULT_METRIC = "rollout/mem_agent_raw_reward/mean" +DEFAULT_ROLLOUT_RESULT_REWARD_KEY = "mem_agent_raw_reward" SCHEMA_VERSION = "mem-agent-reward-summary-v1" _PERF_LINE = re.compile(r"\bperf\s+(\d+):\s+(\{.*\})") @@ -57,6 +58,55 @@ def extract_reward_points(lines: Iterable[str], metric: str = DEFAULT_METRIC) -> return sorted(by_rollout.items()) +def extract_reward_points_from_rollout_results( + directory: Path, + *, + reward_key: str = DEFAULT_ROLLOUT_RESULT_REWARD_KEY, + expected_rollout_ids: Iterable[int] | None = None, +) -> list[tuple[int, float]]: + """Read trajectory-level reward means from ReLax rollout JSONL dumps. + + The actor log's ``rollout/raw_reward`` is computed after MemAgent expands + each trajectory into turns, so documents with more chunks can receive more + weight. Rollout-result JSONL keeps exactly one record per trajectory and is + also immune to Ray log de-duplication, making it the authoritative source + for the reward curve when available. + """ + if expected_rollout_ids is None: + rollout_ids = sorted(int(path.stem) for path in directory.glob("*.jsonl") if path.stem.isdigit()) + else: + rollout_ids = list(expected_rollout_ids) + points: list[tuple[int, float]] = [] + for rollout_id in rollout_ids: + path = directory / f"{rollout_id}.jsonl" + if not path.is_file(): + raise ValueError(f"Missing rollout result for rollout {rollout_id}: {path}") + values: list[float] = [] + with path.open(encoding="utf-8") as source: + for line_number, line in enumerate(source, start=1): + if not line.strip(): + continue + record = json.loads(line) + reward = record.get("reward") + raw_value = reward.get(reward_key) if isinstance(reward, dict) else reward + if isinstance(raw_value, bool) or not isinstance(raw_value, (int, float)): + raise ValueError( + f"Reward {reward_key!r} at rollout {rollout_id}, line {line_number} " + f"is not numeric: {raw_value!r}." + ) + value = float(raw_value) + if not math.isfinite(value) or not 0.0 <= value <= 1.0: + raise ValueError( + f"Reward {reward_key!r} at rollout {rollout_id}, line {line_number} " + f"is outside [0, 1]: {value!r}." + ) + values.append(value) + if not values: + raise ValueError(f"Rollout result {path} contains no reward records.") + points.append((rollout_id, fmean(values))) + return points + + def summarize_reward_points( points: list[tuple[int, float]], *, @@ -172,6 +222,12 @@ def main() -> None: parser.add_argument("--log-file", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--metric", default=DEFAULT_METRIC) + parser.add_argument( + "--rollout-result-dir", + type=Path, + help="Prefer one-record-per-trajectory rollout JSONL files over potentially de-duplicated log metrics.", + ) + parser.add_argument("--rollout-result-reward-key", default=DEFAULT_ROLLOUT_RESULT_REWARD_KEY) parser.add_argument("--expected-steps", type=int) parser.add_argument( "--expected-start", @@ -184,16 +240,30 @@ def main() -> None: parser.add_argument("--svg-output", type=Path) args = parser.parse_args() - with args.log_file.open(encoding="utf-8", errors="replace") as source: - points = extract_reward_points(source, metric=args.metric) + expected_rollout_ids = None + if args.expected_steps is not None: + expected_rollout_ids = range(args.expected_start, args.expected_start + args.expected_steps) + summary_metric = args.metric + if args.rollout_result_dir is not None: + points = extract_reward_points_from_rollout_results( + args.rollout_result_dir, + reward_key=args.rollout_result_reward_key, + expected_rollout_ids=expected_rollout_ids, + ) + summary_metric = f"rollout_result/reward/{args.rollout_result_reward_key}/mean" + else: + with args.log_file.open(encoding="utf-8", errors="replace") as source: + points = extract_reward_points(source, metric=args.metric) summary = summarize_reward_points( points, expected_steps=args.expected_steps, expected_start=args.expected_start, window_size=args.window_size, - metric=args.metric, + metric=summary_metric, ) summary["log_file"] = str(args.log_file.resolve()) + if args.rollout_result_dir is not None: + summary["rollout_result_dir"] = str(args.rollout_result_dir.resolve()) args.output.parent.mkdir(parents=True, exist_ok=True) with args.output.open("w", encoding="utf-8") as destination: json.dump(summary, destination, ensure_ascii=False, indent=2) diff --git a/tests/examples/mem_agent/test_summarize_reward.py b/tests/examples/mem_agent/test_summarize_reward.py index 9e825c59e..3ba5dcb23 100644 --- a/tests/examples/mem_agent/test_summarize_reward.py +++ b/tests/examples/mem_agent/test_summarize_reward.py @@ -7,6 +7,7 @@ from examples.mem_agent.summarize_reward import ( extract_reward_points, + extract_reward_points_from_rollout_results, summarize_reward_points, write_reward_csv, write_reward_svg, @@ -65,6 +66,33 @@ def test_summary_accepts_complete_resumed_rollout_range(): summarize_reward_points([(20, 0.25), (22, 0.75)], expected_steps=3, expected_start=20) +def test_rollout_result_rewards_are_trajectory_means_and_support_resume(tmp_path): + rollout_dir = tmp_path / "rollout_result" / "train" + rollout_dir.mkdir(parents=True) + (rollout_dir / "20.jsonl").write_text( + "\n".join( + [ + '{"reward": {"score": 1.0, "mem_agent_raw_reward": 1.0}}', + '{"reward": {"score": 0.0, "mem_agent_raw_reward": 0.0}}', + ] + ) + + "\n", + encoding="utf-8", + ) + (rollout_dir / "21.jsonl").write_text( + '{"reward": {"score": 1.0, "mem_agent_raw_reward": 1.0}}\n', encoding="utf-8" + ) + + points = extract_reward_points_from_rollout_results( + rollout_dir, + expected_rollout_ids=range(20, 22), + ) + + assert points == [(20, 0.5), (21, 1.0)] + with pytest.raises(ValueError, match="Missing rollout result"): + extract_reward_points_from_rollout_results(rollout_dir, expected_rollout_ids=range(20, 23)) + + def test_reward_csv_and_svg_keep_auditable_points(tmp_path): points = [(0, 0.25), (1, 0.75)] csv_path = tmp_path / "reward.csv" From 22a1b478344a3f90dba78e2f18583931c4d6fe85 Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Tue, 4 Aug 2026 23:21:21 +0800 Subject: [PATCH 15/16] fix: allow MemAgent pilot scheduler resume --- examples/mem_agent/run-qwen3-0.6B-train.sh | 10 ++++++++-- tests/examples/mem_agent/test_recipe_contract.py | 1 + 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/examples/mem_agent/run-qwen3-0.6B-train.sh b/examples/mem_agent/run-qwen3-0.6B-train.sh index d3b16a734..fcf27f6fb 100755 --- a/examples/mem_agent/run-qwen3-0.6B-train.sh +++ b/examples/mem_agent/run-qwen3-0.6B-train.sh @@ -80,8 +80,14 @@ source "${MODEL_CONFIG_DIR}/qwen3-0.6B.sh" RESUME_ARGS=() if [[ -n "${LOAD_PATH}" ]]; then # --num-rollout remains the total target. ReLax restores optimizer/model - # state from LOAD_PATH and continues at this explicit rollout id. - RESUME_ARGS+=(--load "${LOAD_PATH}" --start-rollout-id "${START_ROLLOUT_ID}") + # state from LOAD_PATH and continues at this explicit rollout id. The trend + # gate and full pilot intentionally use different total rollout targets, so + # keep the new constant-LR scheduler horizon while loading the saved state. + RESUME_ARGS+=( + --load "${LOAD_PATH}" + --start-rollout-id "${START_ROLLOUT_ID}" + --override-opt-param-scheduler + ) fi NOW="$(date '+%Y%m%dT%H%M%S%z')" diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index 10e82f8eb..793043a5a 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -65,6 +65,7 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): assert "training-reward.svg" in train assert "--start-rollout-id" in train assert "--load" in train + assert "--override-opt-param-scheduler" in train assert "NUM_ROLLOUT - START_ROLLOUT_ID" in train assert '--expected-steps "$((NUM_ROLLOUT - START_ROLLOUT_ID))"' in train assert '--expected-start "${START_ROLLOUT_ID}"' in train From 9a0f3a3c74de11697eb3d9e14607795ef7c515fa Mon Sep 17 00:00:00 2001 From: manager_of_pre_star_li Date: Thu, 13 Aug 2026 18:31:48 +0800 Subject: [PATCH 16/16] fix: separate rollout row count from metrics --- docs/en/guide/customize-training.md | 6 ++ docs/zh/guide/customize-training.md | 5 ++ .../mem_agent/config-pilot-qwen3-0.6b.yaml | 1 - examples/mem_agent/config.yaml | 1 - examples/mem_agent/contracts.py | 18 ++++++ examples/mem_agent/convert.py | 8 +-- examples/mem_agent/rollout.py | 4 +- relax/core/controller.py | 5 +- relax/distributed/ray/rollout.py | 3 +- relax/engine/rollout/base_types.py | 14 +++++ relax/engine/rollout/sglang_rollout.py | 9 +-- tests/engine/rollout/test_base_types.py | 61 +++++++++++++++++++ tests/examples/mem_agent/test_convert.py | 8 +++ .../mem_agent/test_recipe_contract.py | 2 + tests/examples/mem_agent/test_rollout.py | 35 +++++++++++ 15 files changed, 164 insertions(+), 16 deletions(-) create mode 100644 examples/mem_agent/contracts.py diff --git a/docs/en/guide/customize-training.md b/docs/en/guide/customize-training.md index da0627e50..b4fb96be1 100644 --- a/docs/en/guide/customize-training.md +++ b/docs/en/guide/customize-training.md @@ -211,6 +211,12 @@ generate.manages_inference_permit = True Specify via launch script (`--custom-generate-function-path examples.deepeyes.rollout.generate`), or per eval dataset via `custom_generate_function_path` in eval config. +If you replace the higher-level function selected by `--rollout-function-path`, return +`RolloutFnTrainOutput`. Its `metrics` member is for observability. A rollout that enables +`custom_train_expanded_batch` and transfers a data-dependent number of rows must also set +`train_row_count` to the exact post-conversion row count placed in the current training partition. +Ordinary 1:1 rollout functions should leave `train_row_count` as `None`. + ### Per-request concurrency scheduling for multi-turn rollout By default `generate_and_rm` holds one session-level concurrency permit (`GenerateState.semaphore`) for the entire custom `generate` call — a multi-turn rollout keeps the slot even while running env/tool steps, which hurts engine utilization. diff --git a/docs/zh/guide/customize-training.md b/docs/zh/guide/customize-training.md index bfa86f3da..ee4aa57f8 100644 --- a/docs/zh/guide/customize-training.md +++ b/docs/zh/guide/customize-training.md @@ -207,6 +207,11 @@ generate.manages_inference_permit = True 通过启动脚本指定（`--custom-generate-function-path examples.deepeyes.rollout.generate`），或在评估数据集配置中通过 `custom_generate_function_path` 按数据集设置。 +如果替换 `--rollout-function-path` 指向的更高层 rollout 函数，应返回 +`RolloutFnTrainOutput`。其中 `metrics` 只用于可观测性；若 rollout 开启 +`custom_train_expanded_batch`，且实际传输行数由 converter 动态展开决定，还必须把当前训练分区中 +实际写入的转换后行数填入 `train_row_count`。普通 1:1 rollout 应保持 `train_row_count=None`。 + ### 多轮 Rollout 的请求级并发调度 默认情况下，`generate_and_rm` 会为整个自定义 `generate` 调用持有一把会话级并发锁（`GenerateState.semaphore`）——多轮 rollout 在环境/工具执行期间也一直占用名额，降低推理引擎利用率。 diff --git a/examples/mem_agent/config-pilot-qwen3-0.6b.yaml b/examples/mem_agent/config-pilot-qwen3-0.6b.yaml index 23bc9a5df..6217714fc 100644 --- a/examples/mem_agent/config-pilot-qwen3-0.6b.yaml +++ b/examples/mem_agent/config-pilot-qwen3-0.6b.yaml @@ -6,7 +6,6 @@ mem_agent_max_final_tokens: 64 mem_agent_max_chunks: 4 mem_agent_enable_thinking: false mem_agent_credit_assignment: split -mem_agent_strict_alignment: true # Four memory turns plus one final turn is the maximum pilot expansion. custom_train_sample_expansion_factor: 5 custom_train_data_group_size: 1 diff --git a/examples/mem_agent/config.yaml b/examples/mem_agent/config.yaml index 8ead9f1b5..5a3b5afd6 100644 --- a/examples/mem_agent/config.yaml +++ b/examples/mem_agent/config.yaml @@ -3,7 +3,6 @@ mem_agent_max_memory_tokens: 1024 mem_agent_max_final_tokens: 256 mem_agent_max_chunks: 64 mem_agent_credit_assignment: split -mem_agent_strict_alignment: true # One trajectory has at most 64 memory turns plus one final-answer turn. # These fields reserve enough queue capacity, disable post-expansion GRPO # regrouping, and tell the actor to consume every converted row in one step. diff --git a/examples/mem_agent/contracts.py b/examples/mem_agent/contracts.py new file mode 100644 index 000000000..aaf97876b --- /dev/null +++ b/examples/mem_agent/contracts.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Compatibility guards for the MemAgent trajectory contract.""" + +from __future__ import annotations + +from typing import Any + + +def require_strict_alignment(args: Any) -> None: + """Preserve rejection of the retired false-valued compatibility option. + + Alignment validation is always strict and unconditional. The bundled + recipes therefore no longer advertise a boolean switch, while older + external configs that explicitly requested an unsupported relaxed mode + continue to fail instead of silently changing behavior. + """ + if not getattr(args, "mem_agent_strict_alignment", True): + raise ValueError("MemAgent training requires mem_agent_strict_alignment=true.") diff --git a/examples/mem_agent/convert.py b/examples/mem_agent/convert.py index 542d69121..a8c09b9b4 100644 --- a/examples/mem_agent/convert.py +++ b/examples/mem_agent/convert.py @@ -5,6 +5,7 @@ from typing import Any +from examples.mem_agent.contracts import require_strict_alignment from relax.utils.logging_utils import get_logger from relax.utils.types import Sample from relax.utils.utils import dict_to_tensordict, post_process_rewards @@ -65,8 +66,7 @@ def convert_samples(args: Any, samples: list[Sample]): """Normalize trajectory rewards first, then expand every saved turn.""" if not samples: raise ValueError("MemAgent converter received an empty sample list.") - if not getattr(args, "mem_agent_strict_alignment", True): - raise ValueError("MemAgent training requires mem_agent_strict_alignment=true.") + require_strict_alignment(args) for sample in samples: if sample.status in (Sample.Status.ABORTED, Sample.Status.FAILED): raise ValueError(f"Cannot train from sample index={sample.index} with status={sample.status.value}.") @@ -116,8 +116,8 @@ def convert_samples(args: Any, samples: list[Sample]): zip(samples, validated_turns, raw_rewards, advantages, strict=True) ): turn_credit = float(advantage) / len(turns) if credit_assignment == "split" else float(advantage) - # This expansion is deliberately lossless. Unlike the fixed VIME - # helper, no tail rows are trimmed to a global-batch multiple. + # Expansion is lossless: every turn remains an independent row and no + # tail rows are trimmed merely to manufacture divisibility. for fallback_turn_index, turn in enumerate(turns): tokens = list(turn["tokens"]) response_length = int(turn["response_length"]) diff --git a/examples/mem_agent/rollout.py b/examples/mem_agent/rollout.py index aaa005b75..58e63f643 100644 --- a/examples/mem_agent/rollout.py +++ b/examples/mem_agent/rollout.py @@ -7,6 +7,7 @@ from collections.abc import Awaitable, Callable from typing import Any +from examples.mem_agent.contracts import require_strict_alignment from examples.mem_agent.prompts import ( NO_MEMORY, final_instruction, @@ -126,8 +127,7 @@ async def generate_trajectory( max_final_tokens = int(getattr(args, "mem_agent_max_final_tokens", 256)) max_chunks = int(getattr(args, "mem_agent_max_chunks", 64)) enable_thinking = getattr(args, "mem_agent_enable_thinking", None) - if not getattr(args, "mem_agent_strict_alignment", True): - raise ValueError("MemAgent training requires mem_agent_strict_alignment=true.") + require_strict_alignment(args) if max_memory_tokens <= 0 or max_final_tokens <= 0: raise ValueError("MemAgent memory and final response limits must be positive.") diff --git a/relax/core/controller.py b/relax/core/controller.py index 089711cc1..c93e968f5 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -258,8 +258,9 @@ def _initialize_data_system(self): * self.config.n_samples_per_prompt * get_train_sample_expansion_factor(self.config) ) - # MemAgent completes GRPO normalization before turn expansion, so its - # converted rows opt into sampler group size 1. + # The grouping unit is a physical TransferQueue sampling contract. + # Custom converters that emit independently consumable rows may opt + # into size 1 without teaching the controller their reward semantics. train_data_group_size = get_train_data_group_size(self.config) if getattr(self.config, "fully_async", False) and getattr(self.config, "use_dynamic_batch_size", False): # Fully-async + dynamic-batch path streams data per DP via token diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index cea8163bf..df67937a9 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -1045,8 +1045,7 @@ async def generate(self, rollout_id): # The transfer helper reports the exact rows it actually put. # Re-converting output.samples here is unsafe because custom # converters may filter rows or be stateful. - metrics = output.metrics or {} - row_count = metrics.get("rollout/train_batch_row_count") + row_count = output.train_row_count if row_count is None: raise RuntimeError("Expanded rollout did not report its transferred train row count.") # This count controls TQ consumption only. global_batch_size diff --git a/relax/engine/rollout/base_types.py b/relax/engine/rollout/base_types.py index 404bc4561..1d7adee07 100644 --- a/relax/engine/rollout/base_types.py +++ b/relax/engine/rollout/base_types.py @@ -9,6 +9,10 @@ class RolloutFnTrainOutput: samples: list[list[Sample]] metrics: dict[str, Any] = None + # Exact post-conversion row count transferred to the training partition. + # This is transport metadata for data-dependent 1:N converters; keeping it + # separate prevents orchestration from depending on monitoring key names. + train_row_count: int | None = None @dataclass @@ -24,6 +28,16 @@ def call_rollout_fn(fn, *args, evaluation: bool, **kwargs): if not isinstance(output, (RolloutFnTrainOutput, RolloutFnEvalOutput)): output = RolloutFnEvalOutput(data=output) if evaluation else RolloutFnTrainOutput(samples=output) + if not evaluation and isinstance(output, RolloutFnTrainOutput): + # One-way adapter for custom rollout functions written against the + # initial expanded-batch contract. Keep the legacy metric observable, + # but normalize its control value into the typed transport field here. + legacy_row_count = (output.metrics or {}).get("rollout/train_batch_row_count") + if output.train_row_count is None: + output.train_row_count = legacy_row_count + elif legacy_row_count is not None and output.train_row_count != legacy_row_count: + raise ValueError("Rollout train_row_count conflicts with the legacy rollout/train_batch_row_count metric.") + # Apply --rollout-sample-filter-path (train only). The filter sets # sample.remove_sample=True in-place; downstream (relax/utils/utils.py:126) # zeros loss_mask for those samples so they don't contribute gradient, while diff --git a/relax/engine/rollout/sglang_rollout.py b/relax/engine/rollout/sglang_rollout.py index 36760fd8b..23c43e041 100644 --- a/relax/engine/rollout/sglang_rollout.py +++ b/relax/engine/rollout/sglang_rollout.py @@ -1151,12 +1151,13 @@ def target_reached() -> bool: state.reset() metrics = metric_gatherer.collect() + train_row_count = None if getattr(args, "custom_train_expanded_batch", False): - # This private coordination metric is consumed by RolloutManager. It is - # emitted only for expanded converters so ordinary rollout metrics stay - # byte-for-byte compatible. + train_row_count = transferred_train_rows + # Retain the published metric key for external observers and return-shape + # compatibility. The structured field carries the control signal. metrics["rollout/train_batch_row_count"] = transferred_train_rows - return RolloutFnTrainOutput(samples=data, metrics=metrics), aborted_samples + return RolloutFnTrainOutput(samples=data, metrics=metrics, train_row_count=train_row_count), aborted_samples EVAL_PROMPT_DATASET = {} diff --git a/tests/engine/rollout/test_base_types.py b/tests/engine/rollout/test_base_types.py index 35f0184c5..8d54f756a 100644 --- a/tests/engine/rollout/test_base_types.py +++ b/tests/engine/rollout/test_base_types.py @@ -12,6 +12,8 @@ from argparse import Namespace +import pytest + from relax.engine.rollout.base_types import ( RolloutFnEvalOutput, RolloutFnTrainOutput, @@ -46,6 +48,65 @@ def _build_groups(n: int) -> list[list[Sample]]: return [[Sample(index=i, response_length=1)] for i in range(n)] +def test_train_output_keeps_metrics_and_transport_metadata_independent(): + groups = _build_groups(2) + output = RolloutFnTrainOutput( + samples=groups, + metrics={"rollout/train_batch_row_count": 6}, + train_row_count=6, + ) + + assert output.samples is groups + assert output.metrics == {"rollout/train_batch_row_count": 6} + assert output.train_row_count == 6 + + +def test_train_output_defaults_preserve_existing_callers(): + output = RolloutFnTrainOutput(samples=_build_groups(1)) + + assert output.metrics is None + assert output.train_row_count is None + + +def test_call_rollout_fn_adapts_legacy_row_count_metric(): + groups = _build_groups(2) + + def rollout_fn(args, evaluation): + return RolloutFnTrainOutput(samples=groups, metrics={"rollout/train_batch_row_count": 6}) + + output = call_rollout_fn(rollout_fn, Namespace(), evaluation=False) + + assert output.metrics == {"rollout/train_batch_row_count": 6} + assert output.train_row_count == 6 + + +def test_call_rollout_fn_prefers_structured_row_count(): + groups = _build_groups(2) + + def rollout_fn(args, evaluation): + return RolloutFnTrainOutput( + samples=groups, + metrics={"rollout/train_batch_row_count": 6}, + train_row_count=6, + ) + + output = call_rollout_fn(rollout_fn, Namespace(), evaluation=False) + + assert output.train_row_count == 6 + + +def test_call_rollout_fn_rejects_conflicting_row_count_contracts(): + def rollout_fn(args, evaluation): + return RolloutFnTrainOutput( + samples=_build_groups(2), + metrics={"rollout/train_batch_row_count": 5}, + train_row_count=6, + ) + + with pytest.raises(ValueError, match="train_row_count conflicts"): + call_rollout_fn(rollout_fn, Namespace(), evaluation=False) + + class TestCallRolloutFnFilter: def test_filter_applied_on_train_output(self): groups = _build_groups(4) diff --git a/tests/examples/mem_agent/test_convert.py b/tests/examples/mem_agent/test_convert.py index 681bb7766..38504c159 100644 --- a/tests/examples/mem_agent/test_convert.py +++ b/tests/examples/mem_agent/test_convert.py @@ -115,3 +115,11 @@ def test_converter_rejects_structurally_invalid_or_overlong_turns(): args.mem_agent_max_memory_tokens = 1 with pytest.raises(ValueError, match="response_length=2 exceeds 1"): convert_samples(args, [overlong, _sample(6, 1.0, 2)]) + + +def test_converter_preserves_explicit_false_strict_alignment_compatibility(): + args = _args() + args.mem_agent_strict_alignment = False + + with pytest.raises(ValueError, match="mem_agent_strict_alignment=true"): + convert_samples(args, [_sample(3, 0.0, 1), _sample(4, 1.0, 1)]) diff --git a/tests/examples/mem_agent/test_recipe_contract.py b/tests/examples/mem_agent/test_recipe_contract.py index 793043a5a..53d4f55dc 100644 --- a/tests/examples/mem_agent/test_recipe_contract.py +++ b/tests/examples/mem_agent/test_recipe_contract.py @@ -22,6 +22,7 @@ def test_custom_config_freezes_model_memory_and_expanded_batch_contract(): assert config["mem_agent_max_final_tokens"] == 256 assert config["mem_agent_max_chunks"] == 64 assert config["mem_agent_credit_assignment"] == "split" + assert "mem_agent_strict_alignment" not in config assert config["custom_train_sample_expansion_factor"] == 65 assert config["custom_train_data_group_size"] == 1 assert config["custom_train_expanded_batch"] is True @@ -53,6 +54,7 @@ def test_qwen06_pilot_is_short_context_single_gpu_and_pass_at_n_gated(): assert config["mem_agent_max_final_tokens"] == 64 assert config["mem_agent_max_chunks"] == 4 assert config["mem_agent_enable_thinking"] is False + assert "mem_agent_strict_alignment" not in config assert config["custom_train_sample_expansion_factor"] == 5 assert config["mem_agent_train_rows_multiple"] == 1 assert "qwen3-0.6B.sh" in train diff --git a/tests/examples/mem_agent/test_rollout.py b/tests/examples/mem_agent/test_rollout.py index 52136df6e..050bcca7d 100644 --- a/tests/examples/mem_agent/test_rollout.py +++ b/tests/examples/mem_agent/test_rollout.py @@ -167,6 +167,16 @@ async def fake_generate(args, turn, sampling_params, evaluation): assert tokenizer.template_kwargs == [{"enable_thinking": False}, {"enable_thinking": False}] +@pytest.mark.asyncio +async def test_generate_trajectory_preserves_explicit_false_strict_alignment_rejection(): + args = _args() + args.mem_agent_strict_alignment = False + sample = Sample(index=0, group_index=0, prompt="Q", metadata={"context": "abc"}) + + with pytest.raises(ValueError, match="mem_agent_strict_alignment=true"): + await generate_trajectory(args, sample, {}, FakeTokenizer(), generator=lambda *args: None) + + def test_truncate_text_to_tokens_retokenizes_to_the_hard_limit(): text, length = truncate_text_to_tokens(FakeTokenizer(), "MEMORY", 3) assert text == "MEM" @@ -284,6 +294,31 @@ async def fake_sglang_generate(args, turn, sampling_params, evaluation): assert [turn["kind"] for turn in result.train_metadata["mem_agent_turns"]] == ["memory", "final"] +@pytest.mark.asyncio +async def test_public_generate_preserves_false_strict_alignment_as_aborted(monkeypatch): + class FakeGenerateState: + def __init__(self, args): + del args + self.tokenizer = FakeTokenizer() + + async def unused_generate(*args, **kwargs): + raise AssertionError("strict contract must reject before inference") + + fake_module = ModuleType("relax.engine.rollout.sglang_rollout") + fake_module.GenerateState = FakeGenerateState + fake_module.generate = unused_generate + monkeypatch.setitem(sys.modules, fake_module.__name__, fake_module) + + args = _args() + args.mem_agent_strict_alignment = False + sample = Sample(index=3, group_index=0, prompt="Q", metadata={"context": "abc"}) + result = await generate(args, sample, {"temperature": 1.0}) + + assert result.status == Sample.Status.ABORTED + assert result.train_metadata is None + assert result.rollout_log_probs == [] + + @pytest.mark.asyncio async def test_generate_trajectory_rejects_response_longer_than_turn_limit(): tokenizer = FakeTokenizer()