#!/usr/bin/env python3 """Measure warmed free-T4 throughput using repeated TRAINING prompts only. Writes an incremental Drive report for the pinned base or a completed adapter. These are capacity checks, not evidence about the frozen blind benchmark. """ from __future__ import annotations import argparse from datetime import datetime, timezone import gc import hashlib import importlib.metadata import json import os from pathlib import Path import sys import time ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT / "src")) from vibegame import BASE_MODEL from generate import UnslothGenerator BASE_REVISION = "4858886f896bd05db823557fc5cbc4ac28342af7" BATCH_SIZES = [1, 4, 8, 16] MEMORY_ESTIMATE_POLICY = "kv_full_context_v1" def require_gpu(torch, pro_a100=False): expected = "A100" if pro_a100 else "T4" if (not torch.cuda.is_available() or torch.cuda.device_count() != 1 or expected not in torch.cuda.get_device_name()): raise RuntimeError(f"Requires one already allocated {expected} CUDA GPU; no GPU is provisioned") def capacity_prompts(adapter=None, training_jsonl=None): """The opt-in A100 path reads only the exact admitted training split.""" path = training_jsonl or ROOT / "data/teacher_prompts.jsonl" if training_jsonl is not None: signature = json.loads((adapter / "run_manifest.json").read_text())["signature"] if hashlib.sha256(path.read_bytes()).hexdigest() != signature["train_sha256"]: raise ValueError("Capacity training file differs from the completed adapter's training split") rows = [] with path.open() as handle: for line in handle: if line.strip(): row = json.loads(line) rows.append({"id": row.get("id", f"training-row-{len(rows)}"), "instruction": row["instruction"]}) if len(rows) == 4: break if len(rows) != 4 or any(not isinstance(row["instruction"], str) or not row["instruction"] for row in rows): raise ValueError("Four training prompts are required") return rows def save_report(path: Path, report: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") with temporary.open("w", encoding="utf-8") as handle: json.dump(report, handle, indent=2) handle.write("\n") handle.flush() os.fsync(handle.fileno()) temporary.replace(path) def select_candidate(measurements: list[dict], memory_limit_bytes: int) -> dict | None: candidates = [row for row in measurements if row.get("status") == "success" and row["peak_allocated_bytes"] < memory_limit_bytes and row["peak_reserved_bytes"] < memory_limit_bytes and row.get("full_context_memory_estimate", {}).get("within_limit") is True] return max(candidates, key=lambda row: row["sum_tokens_per_second"]) if candidates else None def estimate_full_context_memory(row: dict, config: dict, gpu_memory_bytes: int, memory_limit_bytes: int, cache_element_bytes: int = 2, growth_tokens: int = 512) -> dict: """Conservative KV-growth extrapolation, not a proof against CUDA OOM.""" batch = row["batch_size"] inputs, outputs = row["per_prompt_input_tokens"], row["per_prompt_tokens"] if len(inputs) != batch or len(outputs) != batch or any(value <= 0 for value in inputs + outputs): raise ValueError("Memory estimate requires actual input/output lengths for every batch row") layers = int(config["num_hidden_layers"]) heads = int(config["num_attention_heads"]) kv_heads = int(config["num_key_value_heads"]) head_dim = int(config.get("head_dim") or (int(config["hidden_size"]) // heads)) if min(layers, heads, kv_heads, head_dim, cache_element_bytes) <= 0: raise ValueError("Invalid loaded model cache dimensions") growth_tokens = max(512, growth_tokens) bytes_per_sequence_token = 2 * layers * kv_heads * head_dim * cache_element_bytes # Left padding and completed rows remain in the same decoding batch until # the longest output finishes. The last emitted token need not be cached. observed_cache_tokens = max(inputs) + max(outputs) - 1 target_tokens = 4096 + growth_tokens extra_cache_tokens = max(0, target_tokens - observed_cache_tokens) delta_kv_bytes = batch * bytes_per_sequence_token * extra_cache_tokens observed_peak = max(row["peak_allocated_bytes"], row["peak_reserved_bytes"]) reserve_bytes = 2_000_000_000 effective_limit = min(memory_limit_bytes, gpu_memory_bytes - reserve_bytes) estimated_peak = observed_peak + delta_kv_bytes return {"policy": MEMORY_ESTIMATE_POLICY, "assessed_batch_size": batch, "target_context_tokens": 4096, "cache_growth_allowance_tokens": growth_tokens, "cache_element_bytes": cache_element_bytes, "bytes_per_sequence_token": bytes_per_sequence_token, "model_dimensions": {"layers": layers, "attention_heads": heads, "kv_heads": kv_heads, "head_dim": head_dim}, "observed_padded_prompt_tokens": max(inputs), "observed_max_generated_tokens": max(outputs), "observed_cache_tokens": observed_cache_tokens, "observed_peak_bytes": observed_peak, "additional_cache_tokens": extra_cache_tokens, "additional_kv_bytes": delta_kv_bytes, "estimated_peak_bytes": estimated_peak, "effective_limit_bytes": effective_limit, "reserve_bytes": reserve_bytes, "within_limit": estimated_peak < effective_limit, "limitations": "Estimate only, not a measured full-context peak or an OOM guarantee. Adds KV growth to max observed allocated/reserved memory, with an extra cache growth block. Other workspace growth, prefill shapes and allocator fragmentation are not modeled; at least 2 GB of physical GPU memory remains reserved."} def full_check_reason(best: dict | None, speedup: float | None, batch_sizes: list[int], short_only: bool) -> str | None: """A single requested size checks capacity without claiming comparative speed.""" if short_only: return "short-only requested" if best is None: return "No successful short benchmark has the required memory headroom" if len(batch_sizes) == 1: return None if speedup is None or speedup < 1.05: return "No eligible batch improves aggregate throughput by at least 5% over batch one" return None def measure(backend, prompts: list[str], token_limit: int, torch) -> dict: # Warm this batch shape separately; shape-specific JIT is excluded from timing. warmup_started = time.monotonic() backend.generate_many(prompts, 4) torch.cuda.synchronize() warmup_seconds = time.monotonic() - warmup_started gc.collect() torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() outer_started = time.monotonic() results = backend.generate_many(prompts, token_limit) outer_seconds = time.monotonic() - outer_started elapsed = max(row["generation_seconds"] for row in results) call_elapsed = max(row.get("generation_call_seconds", row["generation_seconds"]) for row in results) total_tokens = sum(row["tokens"] for row in results) if elapsed <= 0 or len(results) != len(prompts): raise RuntimeError("Invalid benchmark output count or elapsed time") return {"status": "success", "batch_size": len(prompts), "max_new_tokens": token_limit, "warmup_seconds": warmup_seconds, "generation_seconds": elapsed, "generation_call_seconds": call_elapsed, "end_to_end_call_seconds": outer_seconds, "sum_generated_tokens": total_tokens, "sum_tokens_per_second": total_tokens / elapsed, "sum_tokens_per_call_second": total_tokens / outer_seconds, "per_prompt_tokens": [row["tokens"] for row in results], "per_prompt_input_tokens": [row["prompt_tokens"] for row in results], "per_prompt_done_reason": [row["done_reason"] for row in results], "peak_allocated_bytes": torch.cuda.max_memory_allocated(), "peak_reserved_bytes": torch.cuda.max_memory_reserved(), "latency_note": "Every prompt waits the whole batch call, including tokenization and decoding. sum_tokens_per_second uses model.generate; sum_tokens_per_call_second uses the full outer call. Both are aggregate throughput, not request latency."} def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("output", type=Path, help="JSON path under /content/drive/MyDrive") parser.add_argument("--adapter", type=Path, help="Completed local adapter; load its exact saved base revision and tokenizer") parser.add_argument("--batch-sizes", type=int, nargs="+", choices=[1, 2, 4, 8, 16, 32], default=BATCH_SIZES, help="Ascending batch sizes to measure; pass one size for a direct capacity check") parser.add_argument("--short-tokens", type=int, choices=[128, 256], default=256) parser.add_argument("--memory-limit-gb", type=float, default=12.5, help="Conservative decimal-GB threshold before the full check") parser.add_argument("--short-only", action="store_true", help="Skip automatic full-length check") parser.add_argument("--pro-a100", action="store_true", help="Explicit separate Pro A100 evaluation capacity check; never changes T4 defaults") parser.add_argument("--training-jsonl", type=Path, help="Exact admitted training split, required only for --pro-a100") args = parser.parse_args() output = args.output.expanduser().resolve() drive = Path("/content/drive/MyDrive") if not drive.is_dir() or not output.is_relative_to(drive.resolve()): parser.error("Benchmark evidence must persist under mounted /content/drive/MyDrive") if output.exists(): parser.error("Report already exists; use a new path to preserve measurements") if not 1 <= args.memory_limit_gb <= 12.5: parser.error("memory-limit-gb must be within 1..12.5") if args.batch_sizes != sorted(set(args.batch_sizes)): parser.error("batch-sizes must be unique and ascending") if args.pro_a100: if not args.adapter or args.batch_sizes != [8] or args.short_only or not args.training_jsonl: parser.error("--pro-a100 requires a completed adapter, --batch-sizes 8, --training-jsonl, and the full check") elif args.training_jsonl is not None: parser.error("--training-jsonl is reserved for the opt-in --pro-a100 path") pinned_base, pinned_revision = BASE_MODEL, BASE_REVISION if args.adapter: args.adapter = args.adapter.expanduser().resolve(strict=True) if not (args.adapter / "TRAINING_COMPLETE.json").is_file(): parser.error("Adapter capacity check requires TRAINING_COMPLETE.json from the finished training run") signature = json.loads((args.adapter / "run_manifest.json").read_text())["signature"] pinned_base, pinned_revision = signature["base_model"], signature["base_revision"] from unsloth import FastLanguageModel import torch from bitsandbytes.nn import Linear4bit require_gpu(torch, args.pro_a100) # Read exactly four teacher TRAINING prompts. Never open a blind-eval file. prompt_rows = capacity_prompts(args.adapter, args.training_jsonl) prompts = [row["instruction"] for row in prompt_rows] memory_limit = int(args.memory_limit_gb * 1_000_000_000) report = {"kind": ("trained_adapter" if args.adapter else "untrained_base") + "_warmed_batch_throughput_on_training_prompts_not_model_evaluation", "started_utc": datetime.now(timezone.utc).isoformat(), "status": "loading_adapter" if args.adapter else "loading_base", "base_model": pinned_base, "base_revision": pinned_revision, "adapter": str(args.adapter) if args.adapter else None, "gpu": torch.cuda.get_device_name(), "gpu_memory_bytes": torch.cuda.get_device_properties(0).total_memory, "batch_sizes": args.batch_sizes, "short_max_new_tokens": args.short_tokens, "full_max_new_tokens": 3072, "memory_limit_bytes": memory_limit, "memory_estimate_policy": MEMORY_ESTIMATE_POLICY, "training_prompt_ids": [row["id"] for row in prompt_rows], "training_prompts_sha256": hashlib.sha256(json.dumps(prompts, ensure_ascii=False).encode()).hexdigest(), "blind_prompts_read": False, "versions": {name: importlib.metadata.version(name) for name in ["torch", "unsloth", "transformers", "bitsandbytes"]}, "measurements": [], "full_check": None, "limitations": ("Completed adapter capacity check. " if args.adapter else "Original base only; adapter still needs capacity check. ") + "A 3072-token cap can end at EOS earlier and is not a forced worst-case allocation test. No quality or gameplay conclusion."} if args.pro_a100: report["evaluation_profile"] = "pro-a100-opt-in" report["training_jsonl_sha256"] = hashlib.sha256(args.training_jsonl.read_bytes()).hexdigest() save_report(output, report) if args.adapter: backend = UnslothGenerator(str(args.adapter)) report["generation"] = backend.metadata else: model, tokenizer = FastLanguageModel.from_pretrained(model_name=BASE_MODEL, revision=BASE_REVISION, max_seq_length=4096, dtype=torch.float16, load_in_4bit=True, trust_remote_code=False, use_exact_model_name=True) for module in model.modules(): if isinstance(module, Linear4bit): module.compute_dtype = torch.float16 module.compute_type_is_set = True FastLanguageModel.for_inference(model) backend = UnslothGenerator.__new__(UnslothGenerator) backend.model, backend.tokenizer, backend.torch = model, tokenizer, torch config = backend.model.config.to_dict() cache_element_bytes = backend.model.get_input_embeddings().weight.element_size() from unsloth.models.llama import KV_CACHE_INCREMENT cache_growth_tokens = max(512, int(KV_CACHE_INCREMENT)) report["memory_estimate_cache_dtype_basis"] = str(backend.model.get_input_embeddings().weight.dtype) def add_memory_estimate(row): row["full_context_memory_estimate"] = estimate_full_context_memory( row, config, report["gpu_memory_bytes"], memory_limit, cache_element_bytes=cache_element_bytes, growth_tokens=cache_growth_tokens) report["status"] = "short_benchmarks" save_report(output, report) for size in args.batch_sizes: repeated = [prompts[index % 4] for index in range(size)] try: row = measure(backend, repeated, args.short_tokens, torch) add_memory_estimate(row) except Exception as exc: row = {"batch_size": size, "status": "failed", "error": type(exc).__name__ + ": " + str(exc)} gc.collect() torch.cuda.empty_cache() report["measurements"].append(row) save_report(output, report) print(json.dumps(row), flush=True) if row["status"] == "failed": # Avoid making an OOM worse by trying a still-larger batch. break best = select_candidate(report["measurements"], memory_limit) report["best_short_batch_size"] = best["batch_size"] if best else None baseline = next((row for row in report["measurements"] if row.get("status") == "success" and row["batch_size"] == 1), None) speedup = best["sum_tokens_per_second"] / baseline["sum_tokens_per_second"] if best and baseline else None report["short_speedup_vs_batch1"] = speedup skip_reason = full_check_reason(best, speedup, args.batch_sizes, args.short_only) if skip_reason is None: report["status"] = "full_length_check" save_report(output, report) size = best["batch_size"] try: full = measure(backend, [prompts[index % 4] for index in range(size)], 3072, torch) add_memory_estimate(full) full["within_memory_limit"] = full["peak_allocated_bytes"] < memory_limit and full["peak_reserved_bytes"] < memory_limit report["full_check"] = full report["recommended_batch_size"] = size if full["within_memory_limit"] and full["full_context_memory_estimate"]["within_limit"] else None except Exception as exc: report["full_check"] = {"batch_size": size, "status": "failed", "error": type(exc).__name__ + ": " + str(exc)} report["recommended_batch_size"] = None print(json.dumps({"full_check": report["full_check"]}), flush=True) else: report["full_check_skip_reason"] = skip_reason report["recommended_batch_size"] = None report["status"] = "complete" save_report(output, report) print(json.dumps({"report": str(output), "best_short_batch_size": report["best_short_batch_size"], "recommended_batch_size": report["recommended_batch_size"]}), flush=True) if __name__ == "__main__": main()