Verify native rollout and held-out outputs after recovery

benchmarks/audit_native_outputs.py

@@ -0,0 +1,228 @@
+"""Compare complete native training and held-out outputs across fresh recovery."""
+
+import argparse
+import hashlib
+import json
+from pathlib import Path
+
+import torch
+
+from benchmarks.audit_native_learning_signal import audit as audit_signal
+from benchmarks.audit_native_runtime import audit as audit_runtime
+from vime.rollout.native_eval import evaluation_seed
+from vime.rollout.rm_hub.math_utils import grade_answer_verl
+
+SCORE_ATOL, SCORE_RTOL = 1e-5, 3e-5
+DATASET = "qualification-development"
+
+
+def compare_samples(left: dict, right: dict) -> float:
+    for key in (
+        "index",
+        "group_index",
+        "prompt",
+        "tokens",
+        "response_length",
+        "response",
+        "label",
+        "reward",
+        "status",
+        "metadata",
+        "weight_versions",
+    ):
+        assert left[key] == right[key], f"Changed sample {left['index']}: {key}"
+    traces = [
+        {key: value for key, value in sample["recurrent_trace"].items() if key not in {"runtime_epoch", "request_id"}}
+        for sample in (left, right)
+    ]
+    assert traces[0] == traces[1], "Changed semantic trace"
+    return compare_scores(left["rollout_log_probs"], right["rollout_log_probs"])
+
+
+def compare_scores(left, right) -> float:
+    a, b = torch.as_tensor(left), torch.as_tensor(right)
+    assert a.shape == b.shape and a.numel() > 0 and bool(torch.isfinite(a).all()) and bool(torch.isfinite(b).all())
+    torch.testing.assert_close(a, b, atol=SCORE_ATOL, rtol=SCORE_RTOL)
+    return float((a - b).abs().max())
+
+
+def validate_response(sample: dict, limit: int, tokenizer) -> None:
+    length, trace = sample["response_length"], sample["recurrent_trace"]
+    assert 0 < length <= limit and len(sample["tokens"]) > length
+    assert trace["finish_reason"] in ("stop", "length")
+    assert sample["status"] == ("completed" if trace["finish_reason"] == "stop" else "truncated")
+    if trace["finish_reason"] == "length":
+        assert length == limit
+    else:
+        assert sample["tokens"][-1] == tokenizer.eos_token_id
+    assert sample["response"] == tokenizer.decode(sample["tokens"][-length:], skip_special_tokens=True), (
+        "Saved text does not decode from its tokens"
+    )
+    assert (
+        sample["reward"] in (0, 1) and int(grade_answer_verl(sample["response"], sample["label"])) == sample["reward"]
+    )
+    scores = torch.as_tensor(sample["rollout_log_probs"])
+    assert scores.shape == (length,) and bool(torch.isfinite(scores).all())
+
+
+def audit(continuous: Path, resumed: Path, packet: Path, input_audit: Path, algorithm: str, tokenizer) -> dict:
+    runtime = audit_runtime(continuous, resumed, packet)
+    signals = [audit_signal(run, packet, input_audit, algorithm) for run in (continuous, resumed)]
+    consumed = dict(runtime["consumed_sha256"])
+
+    def consume(path):
+        consumed[str(path.resolve())] = hashlib.sha256(path.read_bytes()).hexdigest()
+        return path
+
+    def load(path):
+        return torch.load(consume(path), map_location="cpu", weights_only=False)
+
+    profile = json.loads(consume(packet / "qualification.json").read_text())
+    assert profile["eval_completions_per_prompt"] == 1
+    assert profile["prompts_per_rollout"] == 4 and profile["completions_per_prompt"] == 8
+    train_rows = [json.loads(line) for line in consume(packet / "train.jsonl").read_text().splitlines()]
+    train_ids = [row["metadata"]["id"] for row in train_rows]
+    assert len(train_ids) == len(set(train_ids)) == 12 and train_ids == profile["selected_ids"]["train"]
+    inputs = json.loads(consume(input_audit).read_text())
+    development = packet / "development.jsonl"
+    assert (
+        hashlib.sha256(consume(development).read_bytes()).hexdigest() == profile["files"][development.name]["sha256"]
+    )
+    rows = [json.loads(line) for line in development.read_text().splitlines()]
+    ids = [row["metadata"]["id"] for row in rows]
+    assert len(ids) == len(set(ids)) == profile["files"][development.name]["rows"] == 8
+    assert ids == profile["selected_ids"]["development"] and not set(ids) & set(profile["selected_ids"]["train"])
+    assert len(inputs["rows"]) == 20 and {row["id"] for row in inputs["rows"]} == set(train_ids + ids)
+    token_rows = [row for row in inputs["rows"] if row["split"] == "development"]
+    assert len(token_rows) == 8 and [row["id"] for row in token_rows] == ids
+    token_hashes = {row["id"]: row["token_ids_sha256"] for row in token_rows}
+    training, evaluation, summaries = [], [], []
+    for run_index, run in enumerate((continuous, resumed)):
+        train_rounds, eval_rounds, curve = [], [], []
+        for step in range(3):
+            data = load(run / f"details/rollout_data/{step}.pt")
+            trained = load(run / f"details/train_data/{step}.pt")
+            for sample in data["samples"]:
+                validate_response(sample, profile["response_limit"], tokenizer)
+            train_rounds.append((data["samples"], trained["samples"]))
+            consume(run / f"grad-actor-{step}-0.pt")
+        for step in range(4):
+            data = load(run / f"details/rollout_data/eval_{step}.pt")
+            samples = data["samples"]
+            assert data["rollout_id"] == step and len(samples) == len(rows)
+            phase = "continuous" if run_index == 0 else ("split" if step <= 2 else "resume")
+            epoch = runtime["phases"][phase]["runtime_epoch"]
+            digests, requests = set(), set()
+            for index, (sample, row) in enumerate(zip(samples, rows, strict=True)):
+                assert sample["index"] == sample["group_index"] == index
+                assert sample["metadata"]["id"] == row["metadata"]["id"] and sample["label"] == row["label"]
+                assert sample["metadata"]["native_eval"] == {
+                    "dataset": DATASET,
+                    "prompt_index": index,
+                    "completion": 0,
+                    "seed_profile": "native-heldout-v1",
+                    "top_p": 1,
+                    "top_k": -1,
+                    "max_tokens": profile["response_limit"],
+                    "stop_token_ids": [],
+                }
+                length, trace = sample["response_length"], sample["recurrent_trace"]
+                validate_response(sample, profile["response_limit"], tokenizer)
+                prefix = hashlib.sha256(json.dumps(sample["tokens"][:-length]).encode()).hexdigest()
+                assert prefix == token_hashes[row["metadata"]["id"]]
+                assert trace["runtime_epoch"] == epoch and trace["policy_version"] == step + 1
+                assert sample["weight_versions"] == [str(step + 1)]
+                assert trace["model_family"] == "ouro" and trace["schema_version"] == 1
+                assert (
+                    trace["model_revision"] == profile["model_revision"]
+                    and trace["engine_revision"] == profile["engine_revision"]
+                )
+                assert trace["prefill_depth"] == 4 and trace["decode_depths"] == [4] * length
+                assert trace["temperature"] == 0 and trace["seed"] == evaluation_seed(
+                    profile["seed"], DATASET, index, 0
+                )
+                assert trace["latent_seed"] is None and trace["latent_profile"] is None
+                assert (
+                    isinstance(trace["request_id"], str)
+                    and trace["request_id"]
+                    and trace["request_id"] not in requests
+                )
+                requests.add(trace["request_id"])
+                assert len(trace["publication_digest"]) == 64
+                digests.add(trace["publication_digest"])
+            assert len(digests) == 1, "Mixed held-out policy publications"
+            digest = digests.pop()
+            if step < 3:
+                assert digest == signals[run_index]["rounds"][step]["publication_digest"], (
+                    "Held-out and training policy differ"
+                )
+            curve.append(
+                dict(
+                    completed_updates=step,
+                    reward_mean=sum(sample["reward"] for sample in samples) / len(samples),
+                    response_tokens=sum(sample["response_length"] for sample in samples),
+                    truncated=sum(sample["status"] == "truncated" for sample in samples),
+                    publication_digest=digest,
+                )
+            )
+            eval_rounds.append(samples)
+        assert len({row["publication_digest"] for row in curve}) == 4, "An update reused an old publication"
+        training.append(train_rounds)
+        evaluation.append(eval_rounds)
+        summaries.append(curve)
+    errors = {"training_scores": 0.0, "heldout_scores": 0.0, "training_tensors": 0.0, "gradient_norm": 0.0}
+    for left, right in zip(signals[0]["rounds"], signals[1]["rounds"], strict=True):
+        errors["gradient_norm"] = max(
+            errors["gradient_norm"], compare_scores(left["gradient_norm"], right["gradient_norm"])
+        )
+    for left, right in zip(*training, strict=True):
+        for a, b in zip(left[0], right[0], strict=True):
+            errors["training_scores"] = max(errors["training_scores"], compare_samples(a, b))
+        for a, b in zip(left[1], right[1], strict=True):
+            assert torch.equal(torch.as_tensor(a["loss_masks"]), torch.as_tensor(b["loss_masks"]))
+            for key in ("advantages", "returns", "kl", "ref_log_probs" if algorithm == "rltt" else "log_probs"):
+                errors["training_tensors"] = max(errors["training_tensors"], compare_scores(a[key], b[key]))
+    for left, right in zip(*evaluation, strict=True):
+        for a, b in zip(left, right, strict=True):
+            errors["heldout_scores"] = max(errors["heldout_scores"], compare_samples(a, b))
+    useful = sorted(set(signals[0]["useful_updates"]) & set(signals[1]["useful_updates"]))
+    return {
+        "output_equivalence_pass": True,
+        "joint_signal_updates": useful,
+        "output_and_signal_pass": bool(useful),
+        "training_samples_per_run": 96,
+        "heldout_samples_per_run": 32,
+        "numeric_tolerance": {"atol": SCORE_ATOL, "rtol": SCORE_RTOL},
+        "max_absolute_errors": errors,
+        "continuous_curve": summaries[0],
+        "resumed_curve": summaries[1],
+        "runtime": runtime,
+        "signals": signals,
+        "consumed_sha256": consumed,
+        "scope": "complete three-update native output/held-out recovery qualification, not a convergence experiment",
+        "weights_optimizer_frozen_sources": "REQUIRES_SEPARATE_AUDITS",
+        "reward_convergence": "NOT_ESTABLISHED",
+    }
+
+
+def main():
+    from transformers import AutoTokenizer
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    for name in ("continuous", "resumed", "packet", "input-audit", "tokenizer", "output"):
+        parser.add_argument("--" + name, type=Path, required=True)
+    parser.add_argument("--algorithm", choices=("grpo", "rltt"), required=True)
+    args = parser.parse_args()
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, local_files_only=True, trust_remote_code=False)
+    report = audit(args.continuous, args.resumed, args.packet, args.input_audit, args.algorithm, tokenizer)
+    report["tokenizer_sha256"] = {
+        name: hashlib.sha256((args.tokenizer / name).read_bytes()).hexdigest()
+        for name in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json", "merges.txt")
+        if (args.tokenizer / name).is_file()
+    }
+    args.output.write_text(json.dumps(report, indent=2) + "\n")
+    raise SystemExit(0 if report["output_and_signal_pass"] else 1)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/results/native-output-recovery/audit_native_response_tokenizer.py

@@ -0,0 +1,55 @@
+"""Validate answer-replay controls with the retained official tokenizer, not model output."""
+
+import hashlib
+import json
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer
+
+from benchmarks.audit_native_outputs import validate_response
+from vime.rollout.rm_hub.math_utils import grade_answer_verl
+from vime.utils.data import Dataset
+
+task = Path(__file__).resolve().parents[1]
+packet = task / "data/native-qualification-v1"
+model = task / "data/math-v1/model-config/ouro-thinking"
+profile = json.loads((packet / "qualification.json").read_text())
+inputs = json.loads((task / "evidence/native-qualification-input-audit.json").read_text())
+assert inputs["qualification_sha256"] == hashlib.sha256((packet / "qualification.json").read_bytes()).hexdigest()
+token_hashes = {row["id"]: row["token_ids_sha256"] for row in inputs["rows"]}
+tokenizer = AutoTokenizer.from_pretrained(model, local_files_only=True, trust_remote_code=False)
+assert tokenizer.eos_token_id == json.loads((model / "config.json").read_text())["eos_token_id"] == 2
+rows = []
+for split in ("train", "development"):
+    source = packet / f"{split}.jsonl"
+    assert hashlib.sha256(source.read_bytes()).hexdigest() == profile["files"][source.name]["sha256"]
+    dataset = Dataset(str(source), tokenizer, None, None, prompt_key="prompt", label_key="label",
+                      metadata_key="metadata", apply_chat_template=True, apply_chat_template_kwargs={})
+    for sample in dataset.samples:
+        prefix = tokenizer.encode(sample.prompt, add_special_tokens=False)
+        assert hashlib.sha256(json.dumps(prefix).encode()).hexdigest() == token_hashes[sample.metadata["id"]]
+        wrong = r"\boxed{0}" if not grade_answer_verl(r"\boxed{0}", sample.label) else r"\boxed{1}"
+        for expected, response in ((1, r"\boxed{" + sample.label + "}"), (0, wrong)):
+            output = tokenizer.encode(response, add_special_tokens=False) + [tokenizer.eos_token_id]
+            decoded = tokenizer.decode(output, skip_special_tokens=True)
+            assert int(grade_answer_verl(decoded, sample.label)) == expected
+            record = dict(tokens=prefix + output, response_length=len(output), response=decoded, label=sample.label,
+                          reward=expected, status="completed", recurrent_trace={"finish_reason": "stop"},
+                          rollout_log_probs=[-1.0] * len(output))
+            validate_response(record, profile["response_limit"], tokenizer)
+            rows.append(dict(split=split, id=sample.metadata["id"], expected_reward=expected,
+                             response_tokens=len(output), token_sha256=hashlib.sha256(json.dumps(output).encode()).hexdigest()))
+assert len(rows) == 40 and not torch.cuda.is_initialized()
+result = {
+    "scope": "answer-replay controls from known labels with invented scores; no model generation or training",
+    "qualification_sha256": inputs["qualification_sha256"], "eos_token_id": tokenizer.eos_token_id,
+    "tokenizer_class": type(tokenizer).__name__, "trust_remote_code": False, "local_files_only": True,
+    "rows": rows, "cuda_initialized": False, "model_generation": "NOT_RUN",
+    "source_sha256": hashlib.sha256((task / "vime/benchmarks/audit_native_outputs.py").read_bytes()).hexdigest(),
+    "tokenizer_sha256": {name: hashlib.sha256((model / name).read_bytes()).hexdigest() for name in
+                         ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json", "merges.txt")},
+}
+(task / "evidence/native-output-tokenizer.json").write_text(json.dumps(result, indent=2) + "\n")
+print(json.dumps({"answer_replay_controls": len(rows), "eos_token_id": tokenizer.eos_token_id,
+                  "model_generation": "NOT_RUN", "cuda_initialized": False}))

benchmarks/results/native-output-recovery/manifest.json

@@ -0,0 +1,33 @@
+{
+  "source_sha256": {
+    "benchmarks/audit_native_outputs.py": "e55d2fa5f89a45d7afdb8fa22724f33f4fb70396830e4584957bcffbcb6fb598",
+    "benchmarks/audit_native_runtime.py": "b676295b99a7489d7d80810186e5b414bbeb36a3258d073fd696cc92dad40e50",
+    "benchmarks/audit_native_learning_signal.py": "48dda7c53d488d826e6ee19dfa8dff4342c6f68c7b104109ba02f9ea33930fbd",
+    "tests/test_native_output_audit.py": "07eed856325720541eb67842fea529c6decf7fb9f86932e95d85035e4b24515a",
+    "tests/plugins/test_native_heldout_training.py": "d18dca50a47aaa4a7bd49d644e8b74a4651225a79e23a3935bd833f4ff54f738",
+    "tests/plugins/test_rltt_online_cycle.py": "d2e36ca16f5fed6d229612aa42da0c97cfc53147f9936b89e5c614814190d893",
+    "tests/plugins/test_looped_grpo.py": "51af0f220b21331b745a0a3fdd21e403dec31cc8053c5f1ac41e9839b4479610",
+    "tests/test_native_learning_signal_audit.py": "60674e58a5a6ab1e42c8e080d02ae2fc255baa73e5d5ec6137b7d70b0eba505c",
+    "tests/test_native_runtime_audit.py": "583de0496f5b3424c617dd7238f7c1500da68a4864fa400d9548ac207a0da15f",
+    "vime/rollout/native_eval.py": "06a7aab463bd888f4d7946544f8f390d1d462504fe88120305e0fa5f4ca4a363",
+    "vime/rollout/vllm_rlt_rollout.py": "aa18f8806f6bfa0a894b3b0f67ecad0e669f7695351309e6f2f4e5b55ddbbeee",
+    "docs/native_output_recovery.md": "d3e1df1441a0e87ffdd5bfd128b61d142fb67f06989f4522cc538d42e26e9429"
+  },
+  "artifact_sha256": {
+    "audit_native_response_tokenizer.py": "0216710e96363e0c9d32f83bf53701e4acae71f37906b2c12cd9de529f79d321",
+    "native-output-final.log": "2fd155bf8c51f66251af04e218270ec383a6db962e63500495ad0411bb25f273",
+    "native-output-final.xml": "f5c4a9f88986fe8a674c9a346fbc7d5e32cd8e94c15cb1eb92e4b25a2654b7a0",
+    "native-output-initial.log": "d4867a66fdc061af434e11cb92ec551223f1190a2ac581eced1b876361eb022e",
+    "native-output-initial.xml": "66b8a1636dd56ea3df2e3f935f147d5371fdc91ddd1912a00b301caeaaf9af5f",
+    "native-output-tiny.log": "3352127eebc7b499ccfeb03af616890924ee0f9bb8a97cd9e92117067a0bd46f",
+    "native-output-tiny.xml": "f74f3d826fa947e9c34f7815a9e81aac8dfceae8a4954a74c5bedf68bacbda5e",
+    "native-output-tokenizer-source.json": "bdd16bb7077b9548dc9f0a690b41ef6e439072331980413bcb289942f52c1542",
+    "native-output-tokenizer-v1.json": "812e0cfa18af358ff9ca1e490bbb1c038bca1fdf09b3d3f9454069b5ceedcb90",
+    "native-output-tokenizer-v1.log": "c17bfecb66747c0e4a67ee13d6490400fa03dcf682d7545c80d1297d5cb38514",
+    "native-output-tokenizer.json": "6ce20866b2fc7dd5f0b0f3f3d4f0cb485c9a73ff0d0651d62fc39c598fe591bd",
+    "native-output-tokenizer.log": "c17bfecb66747c0e4a67ee13d6490400fa03dcf682d7545c80d1297d5cb38514",
+    "ouro-thinking-model-manifest.json": "22b0ab9886d283181acb6e51a4dbd19420be8c2da4474b653a042d2bb2e36081",
+    "resident_pytest.py": "2af5ce445da14f656c71e8d4817dcf85cd57402a77f50f612898c47f52f4d4f4",
+    "validation.json": "ef6cbaeb1f4b2f2ebb090cec90d689c25b7f4ca26dc5129cc6305dd1e12c9583"
+  }
+}