67 lines
3.1 KiB
Python
67 lines
3.1 KiB
Python
"""Pinned MTPLX full-model classic AR receipts. Reference only; no downloads.
|
|
|
|
Run serially under test-supervisor with MLX 0.32.2 and pinned MTPLX on PYTHONPATH.
|
|
"""
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
print("#104 importing pinned MTPLX", file=sys.stderr, flush=True)
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
from mlx_lm.utils import load_model
|
|
from mtplx.models import qwen4_exp
|
|
from mtplx.runtime import MTPLXRuntime, _load_tokenizer_resilient
|
|
from mtplx.sampling import SamplerConfig
|
|
from mtplx.generation import generate_ar
|
|
from mtplx.cache_state import snapshot_cache
|
|
from mtplx.session_bank import SessionBank
|
|
|
|
assert mx.__version__ == "0.32.2"
|
|
root = Path(sys.argv[1])
|
|
print("#104 loading unchanged Bare-Speed model", file=sys.stderr, flush=True)
|
|
model, config = load_model(root, lazy=False, strict=True,
|
|
get_model_classes=lambda **_: (qwen4_exp.Model, qwen4_exp.ModelArgs))
|
|
model.post_weight_load(root)
|
|
tokenizer = _load_tokenizer_resilient(root, config)
|
|
rt = MTPLXRuntime(model, tokenizer, root, False, None)
|
|
|
|
def fingerprint(array):
|
|
values = np.asarray(array.astype(mx.float32)).astype('<f4')
|
|
assert np.isfinite(values).all()
|
|
return dict(shape=list(array.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
|
|
|
|
def final_state(out, prompt):
|
|
assert out.final_state is not None and out.final_state.safe_to_commit
|
|
final = out.final_state
|
|
snapshot = snapshot_cache(final.final_trunk_cache)
|
|
states = [[None if a is None else fingerprint(a) for a in arrays]
|
|
if arrays is not None else None for arrays in snapshot.states]
|
|
return dict(ids=prompt + out.tokens, logits=fingerprint(final.final_logits), states=states)
|
|
|
|
records, warm_records = [], []
|
|
for temperature in (.6, 1.):
|
|
prompt = [9707, 11, 1075]
|
|
sampler = SamplerConfig(temperature=temperature, top_p=.95, top_k=20)
|
|
bank = SessionBank()
|
|
for warm in (False, True):
|
|
out = generate_ar(rt, prompt, max_tokens=8,
|
|
sampler=sampler, seed=0, stop_token_ids=set(), capture_final_state=True,
|
|
session_bank=bank, session_id="issue104", store_prefix_snapshot=False,
|
|
token_callback=lambda ids: print(f"#104 t={temperature} warm={warm}: {ids}", file=sys.stderr, flush=True))
|
|
value = final_state(out, prompt)
|
|
(warm_records if warm else records).append(dict(temperature=temperature, state=value))
|
|
print(json.dumps(dict(event="issue104_oracle", temperature=temperature, warm=warm,
|
|
tokens=out.tokens, state=value, cached_tokens=out.stats.cached_tokens,
|
|
seconds=out.stats.elapsed_s)), flush=True)
|
|
if not warm:
|
|
final = out.final_state
|
|
assert bank.put(runtime=rt, token_ids=value["ids"], cache=final.final_trunk_cache,
|
|
logits=final.final_logits, hidden=None, hidden_variant="post_norm",
|
|
session_id="issue104", mtp_history_policy="cycle", snapshot_epoch=len(value["ids"])) is not None
|
|
prompt = value["ids"] + [1075]
|
|
else:
|
|
assert out.stats.cached_tokens == len(prompt) - 1
|
|
print(json.dumps(dict(event="issue104_classic_baseline", records=records, warm_records=warm_records)), flush=True)
|