54 lines
3.4 KiB
Python
54 lines
3.4 KiB
Python
"""Full installed-model cold/warm chunk-guard reference, with exact state checks."""
|
|
import argparse,hashlib,json,os,subprocess,sys
|
|
from pathlib import Path
|
|
checkout,root,native,out=map(Path,sys.argv[1:5])
|
|
assert subprocess.check_output(['git','-C',str(checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
|
|
sys.path.insert(0,str(checkout))
|
|
from mtplx.server.openai import _server_runtime_env_overrides
|
|
if os.environ.get('DS4_ISSUE110_STOCK_CONTROL')!='1':
|
|
os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(root),generation_mode='ar'),None))
|
|
os.environ.update(MTPLX_COMPILED_GDN='1',MTPLX_QWEN4EXP_COMPILE='1',MTPLX_SUSTAINED_PREFILL='1')
|
|
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
|
|
from mtplx import moe_sorted_gather as msg
|
|
from mtplx.kernels import moe_sorted_gather_nax as kernel
|
|
assert mx.__version__=='0.32.2'
|
|
assert os.environ.get('MTPLX_GDN_BOUNDARY_TAIL_INTERVAL')=='0'
|
|
out.mkdir(parents=True,exist_ok=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(a):
|
|
values=np.asarray(a.astype(mx.float32)).astype('<f4')
|
|
assert np.isfinite(values).all()
|
|
return dict(shape=list(a.shape),sha256=hashlib.sha256(values.tobytes()).hexdigest())
|
|
for chunk in [2048,4096]:
|
|
os.environ['MTPLX_PREFILL_CHUNK_SIZE']=str(chunk)
|
|
bank=SessionBank()
|
|
for phase in ['cold','warm']:
|
|
expected=json.loads((native/f'{chunk}-{phase}.json').read_text())
|
|
prompt=expected['state']['ids'][:expected['prompt_len']]
|
|
guard_before=msg.stats()
|
|
result=generate_ar(rt,prompt,max_tokens=3,sampler=SamplerConfig(temperature=0.,top_p=.95,top_k=20),seed=1,
|
|
stop_token_ids=set(),capture_final_state=True,session_bank=bank,session_id='issue110',store_prefix_snapshot=False,
|
|
token_callback=lambda ids:print(json.dumps(dict(case='token',chunk=chunk,phase=phase,ids=ids)),flush=True))
|
|
final=result.final_state
|
|
assert final.safe_to_commit
|
|
states=[[None if a is None else fingerprint(a) for a in state] if state is not None else None for state in snapshot_cache(final.final_trunk_cache).states]
|
|
state=dict(ids=prompt+result.tokens,logits=fingerprint(final.final_logits),states=states)
|
|
(out/f'{chunk}-{phase}.json').write_text(json.dumps(state))
|
|
assert result.stats.cached_tokens==expected['cached']
|
|
assert state==expected['state'],(chunk,phase,'state differs')
|
|
print(json.dumps(dict(case='turn',chunk=chunk,phase=phase,exact=True,cached=result.stats.cached_tokens,tokens=result.tokens,guard_before=guard_before,guard_after=msg.stats(),kernel=kernel.stats())),flush=True)
|
|
if phase=='cold':
|
|
assert bank.put(runtime=rt,token_ids=state['ids'],cache=final.final_trunk_cache,logits=final.final_logits,hidden=None,
|
|
hidden_variant='post_norm',session_id='issue110',mtp_history_policy='cycle',snapshot_epoch=len(state['ids'])) is not None
|