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