grpo_colocated_benchmark.py

"""Compare complete local generation/transfer/fixed-update cycles.

Run each backend in a separate process from the same local base model. Generated
GSM8K responses are timed, but training uses a common synthetic B8/G8 workload so
zero-reward skips cannot masquerade as speedups. This is not an accuracy run or a
measurement of the cost of training on all 256 generated responses.
"""

from __future__ import annotations

import argparse
from collections.abc import Callable
import hashlib
import json
import math
import os
from pathlib import Path
import random
import sys
import time
from typing import Any

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--backend", choices=("transformers", "vllm-colocated"), required=True)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--vllm-python", type=Path)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--prompts", type=int, default=32)
    parser.add_argument("--validation-prompts", type=int, default=64)
    parser.add_argument("--cycles", type=int, choices=(2, 3), default=3)
    parser.add_argument("--max-tokens", type=int, default=512)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--max-num-seqs", type=int, default=32)
    parser.add_argument("--max-model-len", type=int, default=1024)
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.65)
    parser.add_argument("--train-data", type=Path, default=REPO_ROOT / "data/gsm8k/train.jsonl")
    parser.add_argument("--val-data", type=Path, default=REPO_ROOT / "data/gsm8k/test.jsonl")
    parser.add_argument("--prompt-path", type=Path, default=REPO_ROOT / "cs336_alignment/prompts/r1_zero.prompt")
    args = parser.parse_args(argv)
    if os.path.lexists(args.output_dir):
        parser.error(f"Refusing existing output directory: {args.output_dir}")
    if not args.model.is_dir():
        parser.error("--model must be an existing local directory")
    if args.backend == "vllm-colocated" and (args.vllm_python is None or not args.vllm_python.is_file()):
        parser.error("vllm-colocated requires an existing --vllm-python executable")
    for name in ("prompts", "validation_prompts", "max_tokens", "max_num_seqs", "max_model_len"):
        if getattr(args, name) <= 0:
            parser.error(f"--{name.replace('_', '-')} must be positive")
    if not 0 < args.gpu_memory_utilization < 1:
        parser.error("--gpu-memory-utilization must be between zero and one")
    for name in ("model", "output_dir", "train_data", "val_data", "prompt_path"):
        setattr(args, name, getattr(args, name).resolve())
    if args.vllm_python is not None:
        # Python uses the venv executable location to find pyvenv.cfg; resolving
        # this symlink to /usr/bin would silently select the system environment.
        args.vllm_python = args.vllm_python.absolute()
    return args


def sample_prompt_sets(
    train_path: Path, validation_path: Path, template_path: Path, *,
    train_count: int, validation_count: int, seed: int,
) -> dict[str, Any]:
    template = template_path.read_text(encoding="utf-8")

    def sample(path: Path, count: int, sample_seed: int, limit: int) -> tuple[list[int], list[str]]:
        rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()][:limit]
        if count <= 0 or count > len(rows):
            raise ValueError(f"Requested {count} prompts from {len(rows)} available examples in {path}")
        indices = random.Random(sample_seed).sample(range(len(rows)), count)
        return indices, [template.format(question=rows[index]["question"]) for index in indices]

    training_indices, training_prompts = sample(train_path, train_count, seed, 6400)
    validation_indices, validation_prompts = sample(validation_path, validation_count, seed + 100_000, 1024)
    return {
        "training_indices": training_indices, "training_prompts": training_prompts,
        "validation_indices": validation_indices, "validation_prompts": validation_prompts,
        "input_sha256": {str(path): hashlib.sha256(path.read_bytes()).hexdigest() for path in (train_path, validation_path, template_path)},
    }


def low_memory_config(args: argparse.Namespace) -> argparse.Namespace:
    from scripts.grpo_local_memory_probe import training_args

    return training_args(argparse.Namespace(
        model=str(args.model), device=args.device, algorithm="offpolicy_gspo",
        batch_size=8, group_size=8, response_tokens=args.max_tokens,
        learning_rate=1e-5, attention_implementation="sdpa", logits_chunk_size=32,
        seed=args.seed,
    ))


def _require_finite(value: Any) -> None:
    if isinstance(value, float) and not math.isfinite(value):
        raise RuntimeError("Benchmark produced a non-finite diagnostic")
    if isinstance(value, dict):
        for child in value.values():
            _require_finite(child)
    elif isinstance(value, (list, tuple)):
        for child in value:
            _require_finite(child)


def _validate_update(record: dict[str, Any], expected_step: int) -> None:
    _require_finite(record)
    if record.get("parameter_delta_max", 0.0) <= 0:
        raise RuntimeError("A benchmark cycle must perform a nonzero real AdamW update")
    state = record.get("optimizer_state", {})
    if state.get("initialized_parameter_tensors", 0) <= 0 or state.get("step_min") != expected_step or state.get("step_max") != expected_step:
        raise RuntimeError("AdamW state must remain initialized and advance exactly once per cycle")


class FixedSyntheticUpdate:
    """One real BF16 AdamW update over the same synthetic B8/G8 token fixture."""

    def __init__(self, *, model: Any, tokenizer: Any, optimizer: Any, config: argparse.Namespace, prompt: str, response_tokens: int, seed: int):
        import torch
        from scripts import grpo
        from scripts import grpo_local_memory_probe as probe

        self.model, self.tokenizer, self.optimizer, self.config = model, tokenizer, optimizer, config
        self.seed, self.steps = seed, 0
        self.parameters = list(model.parameters())
        if not isinstance(optimizer, torch.optim.AdamW) or optimizer.defaults["foreach"] is not False:
            raise ValueError("Benchmark requires AdamW with foreach=False")
        if not self.parameters or any(not p.requires_grad or p.dtype != torch.bfloat16 for p in self.parameters):
            raise ValueError("Benchmark requires the full policy in native BF16")
        fixture = argparse.Namespace(prompt=prompt, response_tokens=response_tokens, batch_size=8, group_size=8)
        self.prompts, self.responses, self.truths = probe._synthetic_batch(tokenizer, fixture)
        tokenized = grpo.tokenize_prompt_and_output(self.prompts, self.responses, tokenizer)
        lengths = tokenized["response_mask"].sum(dim=1).tolist()
        if lengths != [response_tokens] * 8:
            raise ValueError(f"Synthetic response token lengths {lengths} do not match requested {response_tokens}")
        self.batch_metadata = {
            "batch_size": 8, "group_size": 8, "microbatch_size": 1,
            "response_token_lengths": lengths, "actual_input_width": tokenized["input_ids"].shape[1],
            "rewards": [1.0, 0.0] * 4,
        }

    def __call__(self) -> dict[str, Any]:
        import torch
        from scripts import grpo
        from scripts import grpo_local_memory_probe as probe
        from cs336_alignment.grpo import grpo_train_step

        # Generation implementations consume RNG differently; keep the common
        # update independent of that difference and of the sampled responses.
        grpo.seed_everything(self.seed + 200_000 + self.steps)
        self.model.train()
        old = grpo.score_fixed_old_log_probs(self.model, self.tokenizer, self.prompts, self.responses, batch_size=1)
        if old.dtype != torch.float32 or not probe._finite_tensor(old):
            raise RuntimeError("Old policy log probabilities must be finite FP32")
        before = probe._parameter_samples(self.parameters)
        gradient_checks: dict[str, Any] = {}

        def check_gradients(optimizer: Any, positional: tuple, keyword: dict) -> None:
            nonzero = 0
            for parameter in self.parameters:
                if parameter.grad is None or not probe._finite_tensor(parameter.grad):
                    raise RuntimeError("Missing or non-finite full-policy gradient")
                nonzero += int(torch.count_nonzero(parameter.grad).item() > 0)
            if nonzero == 0:
                raise RuntimeError("Synthetic update has zero gradient")
            gradient_checks.update(all_finite=True, nonzero_tensors=nonzero)

        handle = self.optimizer.register_step_pre_hook(check_gradients)
        try:
            loss, metrics = grpo_train_step(
                model=self.model, tokenizer=self.tokenizer, optimizer=self.optimizer,
                gradient_accumulation_steps=8, max_grad_norm=self.config.max_grad_norm,
                reward_fn=probe._reward, repeated_prompts=self.prompts,
                rollout_responses=self.responses, repeated_ground_truths=self.truths,
                group_size=8, baseline="mean", advantage_eps=self.config.advantage_eps,
                advantage_normalizer="std", importance_reweighting_method="gspo",
                old_log_probs=old, cliprange=self.config.cliprange, loss_normalization="sequence",
            )
        finally:
            handle.remove()
        numeric = {key: float(value.detach().item()) if isinstance(value, torch.Tensor) else float(value) for key, value in metrics.items()}
        numeric["loss"] = float(loss.item())
        _require_finite(numeric)
        if numeric["optimizer_step"] != 1.0 or numeric["gradient_norm"] <= 0:
            raise RuntimeError("Synthetic update was skipped or has zero gradient norm")
        after = probe._parameter_samples(self.parameters)
        deltas = [float((new - previous).abs().max().item()) for previous, new in zip(before, after, strict=True)]
        _require_finite(deltas)
        record = {
            "metrics": numeric, "gradient_checks": gradient_checks, "batch": self.batch_metadata,
            "old_log_probs_dtype": str(old.dtype), "parameter_delta_max": max(deltas),
            "optimizer_state": probe._optimizer_state(self.optimizer, self.parameters),
        }
        _validate_update(record, self.steps + 1)
        self.steps += 1
        return record


def run_benchmark_loop(
    *, backend: str, generator_factory: Callable[[], Any], update: Callable[[], dict[str, Any]],
    train_prompts: list[str], validation_prompts: list[str], cycles: int,
    group_size: int, seed: int, synchronize: Callable[[], None],
    count_tokens: Callable[[str], int], clock: Callable[[], float] = time.perf_counter,
    policy_load_seconds: float = 0.0, memory: Callable[[], dict[str, Any]] | None = None,
    report: dict[str, Any] | None = None,
) -> dict[str, Any]:
    if cycles < 2 or not train_prompts or not validation_prompts:
        raise ValueError("Need at least two cycles and nonempty train/validation prompts")
    result = report if report is not None else {}
    result.update(status="running", diagnostic="generation_plus_fixed_synthetic_update_not_accuracy", backend=backend, cycles=[])
    result["scope"] = "Generation is GSM8K; each update is a fixed synthetic B8/G8 workload. These timings do not measure full-rollout GRPO training or accuracy."
    result["policy_load_seconds"] = policy_load_seconds
    result["token_count_note"] = "Common retokenization uses the same policy tokenizer for both backends; backend token counts are also retained. EOS/stop can shorten generation."
    previous_fingerprint: str | None = None
    generator = None
    benchmark_started = clock()

    def generate(prompts: list[str], n: int, generation_seed: int) -> dict[str, Any]:
        nonlocal previous_fingerprint
        started = clock()
        generator.sync()
        texts, lengths = generator.generate(prompts, n=n, seed=generation_seed)
        synchronize()
        expected = len(prompts) * n
        if len(texts) != expected or len(lengths) != expected:
            raise RuntimeError(f"Expected {expected} responses, got {len(texts)} texts and {len(lengths)} lengths")
        metadata = getattr(generator, "last_generation_metadata", {})
        _require_finite(metadata)
        metadata = json.loads(json.dumps(metadata, allow_nan=False))
        if backend == "vllm-colocated":
            fingerprint = metadata.get("fingerprint")
            if not fingerprint or fingerprint == previous_fingerprint:
                raise RuntimeError("vLLM must verify a changed policy fingerprint after each real update")
            previous_fingerprint = fingerprint
        common_tokens = sum(count_tokens(text) for text in texts)
        _require_finite(lengths)
        if any(length < 0 for length in lengths):
            raise RuntimeError("Backend returned negative token lengths")
        generated = {
            "elapsed_seconds": clock() - started, "prompt_count": len(prompts), "n": n,
            "response_count": len(texts), "backend_reported_tokens": sum(lengths),
            "common_retokenized_tokens": common_tokens, "backend_metadata": metadata,
        }
        if memory is not None:
            generated["memory_after_restore"] = memory()
        return generated

    try:
        constructor_started = clock()
        generator = generator_factory()
        synchronize()
        result["generator_constructor_seconds"] = clock() - constructor_started
        for cycle_index in range(1, cycles + 1):
            row: dict[str, Any] = {"cycle": cycle_index, "cold": cycle_index == 1, "status": "running"}
            result["cycles"].append(row)
            started = clock()
            print(json.dumps({"phase": "generation_start", "cycle": cycle_index, "backend": backend}), flush=True)
            row["generation"] = generate(train_prompts, group_size, seed + cycle_index)
            print(json.dumps({"phase": "generation_complete", "cycle": cycle_index, **row["generation"]}), flush=True)
            update_started = clock()
            update_record = update()
            synchronize()
            _validate_update(update_record, cycle_index)
            row["fixed_update"] = {**update_record, "elapsed_seconds": clock() - update_started}
            row.update(status="passed", elapsed_seconds=clock() - started)
            if memory is not None:
                row["memory_after_update"] = memory()
            print(json.dumps({"phase": "cycle_complete", "cycle": cycle_index, "elapsed_seconds": row["elapsed_seconds"], "fixed_update_seconds": row["fixed_update"]["elapsed_seconds"], "parameter_delta_max": update_record["parameter_delta_max"]}), flush=True)
        result["validation"] = generate(validation_prompts, 1, seed + 100_000)
    except BaseException as error:
        result.update(status="failed", error_type=type(error).__name__, error=str(error))
        raise
    finally:
        close_started = clock()
        try:
            if generator is not None:
                generator.close()
                synchronize()
        finally:
            result["teardown_seconds"] = clock() - close_started
            result["wall_seconds"] = clock() - benchmark_started
    steady = result["cycles"][1:]
    result["steady"] = {
        "cycle_count": len(steady),
        "total_cycle_seconds": sum(row["elapsed_seconds"] for row in steady),
        "generation_seconds": sum(row["generation"]["elapsed_seconds"] for row in steady),
        "fixed_update_seconds": sum(row["fixed_update"]["elapsed_seconds"] for row in steady),
        "common_retokenized_tokens": sum(row["generation"]["common_retokenized_tokens"] for row in steady),
    }
    result["steady"]["mean_cycle_seconds"] = result["steady"]["total_cycle_seconds"] / len(steady)
    result["first_server_start_seconds"] = result["cycles"][0]["generation"]["backend_metadata"].get("server_start_seconds", 0.0)
    result["accounted_seconds"] = result["generator_constructor_seconds"] + sum(row["elapsed_seconds"] for row in result["cycles"]) + result["validation"]["elapsed_seconds"] + result["teardown_seconds"]
    result["unattributed_overhead_seconds"] = max(0.0, result["wall_seconds"] - result["accounted_seconds"])
    result["end_to_end_seconds_including_policy_load"] = policy_load_seconds + result["wall_seconds"]
    _require_finite(result)
    result["status"] = "passed"
    return result


def main() -> None:
    args = parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=False)
    os.environ["HF_HUB_OFFLINE"] = "1"
    os.environ["TRANSFORMERS_OFFLINE"] = "1"
    import torch
    from scripts import grpo
    from scripts import grpo_local_memory_probe as probe

    report: dict[str, Any] = {"status": "loading", "config": {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()}}
    try:
        selected = sample_prompt_sets(args.train_data, args.val_data, args.prompt_path, train_count=args.prompts, validation_count=args.validation_prompts, seed=args.seed)
        report["sampling"] = selected
        config = low_memory_config(args)
        grpo.seed_everything(args.seed)
        started = time.perf_counter()
        model, tokenizer = grpo.load_policy_and_tokenizer(config)
        device = torch.device(args.device)
        probe._synchronize(device)
        load_seconds = time.perf_counter() - started
        optimizer = grpo.make_optimizer(config, model)
        report["training_config"] = vars(config)
        report["parameter_count"] = sum(parameter.numel() for parameter in model.parameters())
        report["after_load_memory"] = probe.memory_snapshot(device)
        report["memory_note"] = "Allocated is contained in reserved. Allocator metrics cover the training process, not the separate vLLM worker; driver-free snapshots are not a cross-process peak."
        prompt_lengths = [len(tokenizer.encode(text, add_special_tokens=False)) for text in selected["training_prompts"] + selected["validation_prompts"]]
        report["maximum_prompt_tokens"] = max(prompt_lengths)
        if max(prompt_lengths) + args.max_tokens > args.max_model_len:
            raise ValueError("Selected prompt plus generation budget exceeds --max-model-len; use the same larger limit for both backends")
        update = FixedSyntheticUpdate(model=model, tokenizer=tokenizer, optimizer=optimizer, config=config, prompt=selected["training_prompts"][0], response_tokens=args.max_tokens, seed=args.seed)

        def factory() -> Any:
            if args.backend == "transformers":
                return grpo.TransformersGenerator(model=model, tokenizer=tokenizer, device=device, batch_size=1, temperature=1.0, top_p=1.0, max_tokens=args.max_tokens, stop="</answer>")
            from cs336_alignment.vllm_colocated import SingleGPUVLLMGenerator

            return SingleGPUVLLMGenerator(
                policy=model, optimizer=optimizer, model_id=str(args.model),
                vllm_python=str(args.vllm_python), work_dir=args.output_dir / "vllm_worker",
                device=args.device, max_tokens=args.max_tokens, temperature=1.0,
                top_p=1.0, stop="</answer>", seed=args.seed,
                max_num_seqs=args.max_num_seqs, max_model_len=args.max_model_len,
                gpu_memory_utilization=args.gpu_memory_utilization,
            )

        run_benchmark_loop(
            backend=args.backend, generator_factory=factory, update=update,
            train_prompts=selected["training_prompts"], validation_prompts=selected["validation_prompts"],
            cycles=args.cycles, group_size=8, seed=args.seed,
            synchronize=lambda: probe._synchronize(device),
            count_tokens=lambda text: len(tokenizer.encode(text, add_special_tokens=False)),
            policy_load_seconds=load_seconds, memory=lambda: probe.memory_snapshot(device), report=report,
        )
        report["final_optimizer_state"] = probe._optimizer_state(optimizer, update.parameters)
    except BaseException as error:
        report.update(status="failed", error_type=type(error).__name__, error=str(error))
        with (args.output_dir / "benchmark.json").open("x", encoding="utf-8") as output:
            json.dump(report, output, indent=2, sort_keys=True, allow_nan=False)
            output.write("\n")
        raise
    with (args.output_dir / "benchmark.json").open("x", encoding="utf-8") as output:
        json.dump(report, output, indent=2, sort_keys=True, allow_nan=False)
        output.write("\n")
    print(json.dumps({"status": report["status"], "output_file": str(args.output_dir / "benchmark.json"), "steady": report["steady"]}), flush=True)


if __name__ == "__main__":
    main()