File size: 3,657 Bytes
795f737 | 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 | """Isolated-process CPU policy benchmark and optional dynamic INT8 artifact."""
import argparse
import hashlib
import json
from pathlib import Path
import platform
import statistics
from time import perf_counter, process_time
import warnings
import psutil
import torch
from .features import encode
from .policy import LearnedPolicy
from .synthetic import load
from .train import logits, metrics
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint',required=True)
parser.add_argument('--data',default='datasets/synthetic-v1')
parser.add_argument('--quantized',action='store_true')
parser.add_argument('--output',required=True)
args = parser.parse_args()
torch.set_num_threads(2)
torch.set_num_interop_threads(1)
started = perf_counter()
with warnings.catch_warnings():
warnings.simplefilter('ignore',DeprecationWarning)
policy = LearnedPolicy(args.checkpoint,quantized=args.quantized)
load_ms = (perf_counter()-started)*1000
root = Path(args.checkpoint)
artifact = root/'model.safetensors'
if args.quantized:
artifact = root/'model-linear-int8.pt'
torch.save(policy.model.state_dict(),artifact)
# Verify serialization with restricted loading; no arbitrary pickle globals.
state = torch.load(artifact,map_location='cpu',weights_only=True)
policy.model.load_state_dict(state)
evaluation = {}
rows = None
for split in ['validation','test','novel_wording']:
rows = load(Path(args.data)/f'{split}.jsonl')
inputs,a,t,_ = encode(rows,policy.vocab)
la,lt = logits(policy.model,inputs)
evaluation[split] = metrics(la,lt,a,t,policy.temperatures)
wall, model_ms, cpu_ms, encode_ms = [], [], [], []
process = psutil.Process()
observed_rss = process.memory_info().rss
with torch.inference_mode():
for index,row in enumerate(rows[:240]):
start,cpu = perf_counter(),process_time()
inputs,*_ = encode([row],policy.vocab)
encoded = perf_counter()
policy.model(*inputs)
finished = perf_counter()
if index>=20:
wall.append((finished-start)*1000)
model_ms.append((finished-encoded)*1000)
encode_ms.append((encoded-start)*1000)
cpu_ms.append((process_time()-cpu)*1000)
observed_rss = max(observed_rss,process.memory_info().rss)
def stats(values):
return dict(median=statistics.median(values),p95=sorted(values)[int(.95*len(values))])
report = dict(checkpoint=args.checkpoint,quantization='dynamic INT8 Linear only; FP32 embeddings/encoder' if args.quantized else 'FP32',
platform=platform.platform(),torch_version=torch.__version__,threads=2,
load_ms=load_ms,disk_bytes=artifact.stat().st_size,
artifact_sha256=hashlib.sha256(artifact.read_bytes()).hexdigest(),
observed_process_rss_bytes=observed_rss,
memory_scope='Observed Python RSS including training-library imports and evaluation tensors, not browser or exact peak',
end_to_end_policy_ms=stats(wall),neural_forward_ms=stats(model_ms),
feature_encoding_ms=stats(encode_ms),python_cpu_ms=stats(cpu_ms),
evaluation=evaluation,target_vps_validated=False,
scope='Synthetic single-step benchmark. Two torch threads, not two-vCPU CPU affinity or target EPYC.')
Path(args.output).parent.mkdir(parents=True,exist_ok=True)
Path(args.output).write_text(json.dumps(report,indent=2),encoding='utf-8')
print(json.dumps(report,indent=2))
if __name__ == '__main__':
main()
|