132 lines
8.0 KiB
Python
132 lines
8.0 KiB
Python
"""Strict real MTP oracle: compare the same snapshot stage as the native bank.
|
|
|
|
#114: snapshot_cache on a lazy trunk adds zero, creating a false signed-zero
|
|
mismatch. Preserve lazy trunk bits and deep-clone only the committed MTP head.
|
|
"""
|
|
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
|
|
os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(root),generation_mode='mtp'),None))
|
|
os.environ['MTPLX_SUSTAINED_PREFILL']='1'
|
|
os.environ['MTPLX_PREFILL_CHUNK_SIZE']='2048'
|
|
os.environ['MTPLX_CONTEXT_COPY']='0'
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
from mlx_lm.utils import load_model
|
|
from mtplx.models import qwen4_exp as qwen
|
|
from mtplx.runtime import _load_tokenizer_resilient,MTPLXRuntime
|
|
from mtplx.mtp_patch import MTPContract
|
|
from mtplx.generation import generate_mtpk
|
|
from mtplx.sampling import SamplerConfig
|
|
from mtplx.session_bank import SessionBank,_lazy_snapshot_enabled
|
|
from mtplx.cache_state import snapshot_cache,snapshot_cache_lazy_hybrid
|
|
from mtplx.attention_context import current_attention_phase
|
|
out.mkdir(parents=True,exist_ok=True)
|
|
print(json.dumps(dict(case='environment',overrides={k:v for k,v in os.environ.items() if k.startswith('MTPLX_')})),flush=True)
|
|
tokenizer=_load_tokenizer_resilient(root,json.loads((root/'config.json').read_text()))
|
|
model,_=load_model(root,lazy=False,strict=True,get_model_classes=lambda **_:(qwen.Model,qwen.ModelArgs))
|
|
model.post_weight_load(root)
|
|
assert model.attach_mtp(root)
|
|
rt=MTPLXRuntime(model,tokenizer,root,True,MTPContract())
|
|
phase='cold'; index=0; nested=0
|
|
|
|
def wrap(fn):
|
|
def call(input_ids,*args,**kwargs):
|
|
global index,nested
|
|
outer=nested==0
|
|
nested+=1
|
|
try: result=fn(input_ids,*args,**kwargs)
|
|
finally: nested-=1
|
|
if outer and current_attention_phase() == 'prefill' and isinstance(result,tuple) and result[0] is not None and result[1] is not None:
|
|
cache=kwargs.get('cache',args[0] if args else None)
|
|
if cache is not None:
|
|
key=np.asarray((snapshot_cache_lazy_hybrid(cache) if _lazy_snapshot_enabled() else snapshot_cache(cache)).states[35][0].astype(mx.float32))
|
|
(out/f'{phase}-prefill-key35.f32').write_bytes(key.tobytes())
|
|
for part,a in zip(['logits','hidden'],result[:2]):
|
|
# The suffix owner retains only the final hidden row (generation.py:4689).
|
|
if part=='hidden':a=a[:,-1:,:]
|
|
data=np.asarray(a.astype(mx.float32)).tobytes()
|
|
(out/f'{phase}-prefill-{part}.f32').write_bytes(data)
|
|
path=native/f'{phase}-prefill-{part}.f32'
|
|
print(json.dumps(dict(case='prefill_compare',phase=phase,part=part,shape=list(a.shape),exact=path.exists() and path.read_bytes()==data)),flush=True)
|
|
if outer and current_attention_phase()=='decode_verify':
|
|
expected=json.loads((native/phase/f'{index}.json').read_text())
|
|
ids=np.asarray(input_ids).reshape(-1).tolist()
|
|
assert expected['verify'] and ids==expected['ids'], (phase,index,ids,expected['ids'])
|
|
receipt=dict(case='verify_forward',phase=phase,index=index,ids=ids)
|
|
for part,a in zip(['logits','hidden'],result[:2]):
|
|
data=np.asarray(a.astype(mx.float32)).tobytes()
|
|
(out/f'{phase}-{index}-{part}.f32').write_bytes(data)
|
|
ref=(native/phase/f'{index}-{part}.f32').read_bytes()
|
|
receipt[part+'_exact']=data==ref
|
|
receipt[part+'_sha256']=hashlib.sha256(data).hexdigest()
|
|
if data!=ref:
|
|
x=np.frombuffer(data,np.float32); y=np.frombuffer(ref,np.float32)
|
|
receipt[part+'_different']=int(np.count_nonzero(x!=y))
|
|
receipt[part+'_max_abs']=float(np.max(np.abs(x-y)))
|
|
print(json.dumps(receipt),flush=True)
|
|
assert receipt['logits_exact'] and receipt['hidden_exact'],receipt
|
|
index+=1
|
|
return result
|
|
return call
|
|
rt.forward_ar=wrap(rt.forward_ar)
|
|
rt.forward_ar_capture=wrap(rt.forward_ar_capture)
|
|
bank=SessionBank()
|
|
for phase in ['cold','warm']:
|
|
index=0
|
|
expected=json.loads((native/f'{phase}.json').read_text())
|
|
ids=expected['prompt']
|
|
result=generate_mtpk(rt,ids,max_tokens=12,sampler=SamplerConfig(temperature=0.,top_p=.95,top_k=20),speculative_depth=3,
|
|
seed=1,stop_token_ids=set(tokenizer.eos_token_ids),mtp_history_policy='committed',verify_strategy='batched',capture_final_state=True,
|
|
session_bank=bank,session_id='issue109',store_prefix_snapshot=False,
|
|
token_callback=lambda ids:print(json.dumps(dict(case='tokens',phase=phase,ids=ids)),flush=True))
|
|
assert result.tokens==expected['tokens'],(result.tokens,expected['tokens'])
|
|
assert result.stats.cached_tokens==expected['cached'],(result.stats.cached_tokens,expected['cached'])
|
|
assert index==len(list((native/phase).glob('*.json')))
|
|
final=result.final_state
|
|
assert final.safe_to_commit
|
|
for part,a in [('logits',final.final_logits),('hidden',final.final_hidden)]:
|
|
data=np.asarray(a.astype(mx.float32)).tobytes()
|
|
assert data==(native/f'{phase}-{part}.f32').read_bytes(),(phase,part)
|
|
def states(cache,lazy=False):
|
|
def array(a):
|
|
if a is None: return None
|
|
values=np.asarray(a.astype(mx.float32)).copy()
|
|
bits=hashlib.sha256(values.tobytes()).hexdigest()
|
|
return dict(shape=list(a.shape),sha256=bits)
|
|
return [[array(a) for a in state] if state is not None else None for state in (snapshot_cache_lazy_hybrid(cache) if lazy else snapshot_cache(cache)).states]
|
|
arrays=dict(trunk=states(final.final_trunk_cache,_lazy_snapshot_enabled()),mtp=states(final.final_committed_mtp_cache))
|
|
(out/f'{phase}-arrays.json').write_text(json.dumps(arrays))
|
|
native_arrays=json.loads((native/f'{phase}-arrays.json').read_text())
|
|
for family in arrays:
|
|
assert len(arrays[family])==len(native_arrays[family])
|
|
for layer,(actual,want) in enumerate(zip(arrays[family],native_arrays[family])):
|
|
if actual is None or want is None:
|
|
assert actual==want
|
|
continue
|
|
assert len(actual)==len(want)
|
|
for slot,(a,b) in enumerate(zip(actual,want)):
|
|
if a is None or b is None:
|
|
assert a==b
|
|
continue
|
|
assert a['shape']==b['shape'] and a['sha256']==b['sha256'],(phase,family,layer,slot,a,b)
|
|
# Compare the same lazy trunk snapshot used by the server/native bank.
|
|
# An extra deep clone would add zero and change negative-zero bits.
|
|
key=np.asarray((snapshot_cache_lazy_hybrid(final.final_trunk_cache) if _lazy_snapshot_enabled() else snapshot_cache(final.final_trunk_cache)).states[35][0].astype(mx.float32))
|
|
(out/f'{phase}-key35.f32').write_bytes(key.tobytes())
|
|
want=np.fromfile(native/f'{phase}-key35.f32',np.float32).reshape(key.shape)
|
|
different=key.view(np.uint32)!=want.view(np.uint32)
|
|
assert not np.any(different)
|
|
for part in ["key35","logits","hidden"]:
|
|
assert (out/f"{phase}-prefill-{part}.f32").read_bytes()==(native/f"{phase}-prefill-{part}.f32").read_bytes(),(phase,part,"prefill differs")
|
|
print(json.dumps(dict(case='cache_arrays',phase=phase,raw_bits_exact=True,key35_bit_coordinates=np.argwhere(different).tolist())),flush=True)
|
|
head=snapshot_cache(final.final_committed_mtp_cache)
|
|
key=ids+result.tokens
|
|
assert bank.put(runtime=rt,token_ids=key,cache=final.final_trunk_cache,logits=final.final_logits,hidden=final.final_hidden,
|
|
hidden_variant='post_norm',session_id='issue109',mtp_history_policy='committed',mtp_history_snapshot=head,
|
|
snapshot_epoch=len(key),mtp_snapshot_epoch=len(key)) is not None
|
|
print(json.dumps(dict(case='turn',phase=phase,tokens=result.tokens,cached=result.stats.cached_tokens,verify_calls=index,final_exact=True)),flush=True)
|