| |
| """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 |
| |
| |
| 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: |
| |
| 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) |
| |
| 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": |
| |
| 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() |
|
|
