File size: 4,113 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """Solo-only synthetic top32 qualification and graph timing; no serve requests."""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import resource
from collections.abc import Callable
from pathlib import Path
import torch
from ..sampled import MASS, Top32
UUID = "GPU-bc62df2a-59cf-9828-7d43-558d1355bab7"
def timing(fn: Callable[[], object], repetitions: int = 100) -> dict[str, float]:
for _ in range(3):
fn()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
for _ in range(10):
fn()
samples: list[float] = []
start, stop = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
for _ in range(repetitions):
start.record()
graph.replay()
stop.record()
stop.synchronize()
samples.append(start.elapsed_time(stop) / 10)
samples.sort()
return {"median_ms": samples[len(samples)//2], "p95_ms": samples[int(.95*len(samples))]}
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--out", type=Path, required=True)
a = p.parse_args()
if os.environ.get("CUDA_VISIBLE_DEVICES") != UUID:
raise RuntimeError("PP3 solo UUID required")
import fcntl
fd = int(os.environ["PP3_LOCK_FD"])
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
if str(torch.cuda.get_device_properties(0).uuid).removeprefix("GPU-") != UUID.removeprefix("GPU-"):
raise RuntimeError("wrong device")
torch.set_num_threads(8)
torch.set_num_interop_threads(1)
torch.manual_seed(4103)
dev = torch.device("cuda:0")
out: dict[str, object] = {"schema": "pp3.producer.v1", "classification": "MEASURED solo synthetic logits",
"complete": False, "torch": torch.__version__, "uuid": UUID,
"source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest()}
out["producer_sha256"] = hashlib.sha256((Path(__file__).parents[1]/'sampled.py').read_bytes()).hexdigest()
rows: list[dict[str, object]] = []
out["rows"] = rows
def save() -> None:
out["maxrss_kib"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
out["gpu_peak_bytes"] = torch.cuda.max_memory_allocated()
a.out.write_text(json.dumps(out, indent=2)+"\n")
for batch in range(1, 5):
x = torch.randn(batch, 129280, device=dev).to(torch.bfloat16)
keys = torch.tensor([[4103+j, 1, 32768] for j in range(batch)], device=dev, dtype=torch.int64)
tau = torch.full((batch,), .7, device=dev)
sampler = Top32(batch, 129280, dev)
tok, ids, q = sampler(x, keys, tau)
torch.cuda.synchronize()
if sampler.error.any().item():
raise AssertionError("valid proposal refused")
expected = torch.argsort(x.float(), descending=True, stable=True)[:, :32].sort(-1).values
torch.testing.assert_close(ids.long(), expected, rtol=0, atol=0)
if not torch.equal((q.double()*MASS).sum(-1), torch.full((batch,), MASS, device=dev, dtype=torch.float64)):
raise AssertionError("nonconserving mass")
if not ((ids == tok[:, None]) & (q > 0)).any(-1).all().item():
raise AssertionError("proposal outside support")
for depth in (4, 5):
def sampled(depth: int = depth, sampler: Top32 = sampler, x: torch.Tensor = x,
keys: torch.Tensor = keys, tau: torch.Tensor = tau) -> None:
for _ in range(depth):
sampler(x, keys, tau)
def greedy(depth: int = depth, x: torch.Tensor = x) -> None:
for _ in range(depth):
x.argmax(-1)
row: dict[str, object] = {"batch": batch, "depth": depth, "sampled": timing(sampled),
"argmax": timing(greedy), "scope": "logit selection only; excludes Markov, TP exchange, wire"}
rows.append(row)
save()
print(json.dumps(row), flush=True)
out["complete"] = True
save()
if __name__ == "__main__":
main()
|