File size: 3,543 Bytes
7bf0e3b | 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 | """Measure standalone packed loading and a short System One request in a fresh process.
Use an empty HF_HOME and HF_HUB_OFFLINE=1. The base loaders are explicitly
blocked. RSS alone misses Metal buffers, so also sample physical footprint.
"""
import argparse
import json
import platform
import sys
import time
from pathlib import Path
import mlx.core as mx
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from kev.api import SystemOneRequest, to_record
from kev.checkpoint import Checkpoint, LoadOptions
from kev.model import SERVE_MAX_BRANCH, SERVE_MAX_STATE
from scripts.mlx_long_states import RSSPeak, footprint
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--run", required=True)
parser.add_argument("--out", required=True)
parser.add_argument("--mlx-limit-gb", type=float, default=3, help="MLX memory guideline, not a hard RAM cap")
parser.add_argument("--filler-repeats", type=int, default=0, help="append neutral filler to measure longer states; this is not an accuracy benchmark")
args = parser.parse_args()
import kev.mlx_model
import huggingface_hub
def forbidden(*unused, **kwargs):
raise AssertionError("standalone profiling must never call a base loader or download")
kev.mlx_model.load_base = kev.mlx_model.load_full = kev.mlx_model.merge_lora = forbidden
huggingface_hub.snapshot_download = forbidden
mx.set_memory_limit(int(args.mlx_limit_gb * 1e9))
mx.reset_peak_memory()
start = time.perf_counter()
with RSSPeak() as load_peak:
ck = Checkpoint(args.run)
tok, model = ck.load("mps", LoadOptions())
result = {"run": str(Path(args.run).resolve()), "backend": model.backend, "dtype": model.dtype,
"platform": platform.platform(), "original_loaders_blocked": True,
"mlx_limit_bytes": int(args.mlx_limit_gb * 1e9), "mlx_limit_is_hard_cap": False, "load_seconds": time.perf_counter() - start,
"mlx_load_peak_bytes": mx.get_peak_memory(), "mlx_after_load_bytes": mx.get_active_memory(),
"load_footprint_peak_bytes": max(load_peak.footprint, footprint()), "load_rss_peak_bytes": load_peak.peak}
request = SystemOneRequest(state="I was charged twice for order 8812. Please refund the duplicate charge." + " This is a routine account note." * args.filler_repeats,
questions={"team": {"type": "choice", "instructions": "Which team should handle this ticket?",
"criteria": {"billing": "Charges and payments", "shipping": "Delivery and delays"}}})
record, _ = to_record(request)
enc = model.encode(tok, record, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH, strict=True)
start = time.perf_counter()
with RSSPeak() as infer_peak:
probs = model.probs(enc)
result.update(state_tokens=enc["seg"].count(0), input_tokens=len(enc["ids"]), inference_seconds=time.perf_counter() - start,
probabilities={"billing": float(probs[0][0]), "shipping": float(probs[0][1])},
mlx_inference_peak_bytes=mx.get_peak_memory(), inference_footprint_peak_bytes=max(infer_peak.footprint, footprint()),
inference_rss_peak_bytes=infer_peak.peak)
output = Path(args.out)
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
print(json.dumps(result, indent=2))
if __name__ == "__main__":
main()
|