Compare Qwen cache snapshots at matching stages
This commit is contained in:
Georg Bauer committed
1 parent
d9026aa7bf
commit
5a8782b9b6
21 files changed
+3397
-149
No files matched your search
@@ -1,132 +1,5 @@
|
||||
"""Strict default/off head-chunk oracle for real verify rows and matching bank snapshots."""
|
||||
import argparse,hashlib,json,os,subprocess,sys
|
||||
"""Use the shared strict real-request oracle for the head-chunk on/off controls."""
|
||||
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
|
||||
import runpy
|
||||
|
||||
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='issue113',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()
|
||||
values[values==0]=0
|
||||
return dict(shape=list(a.shape),sha256=bits,value_sha256=hashlib.sha256(values.tobytes()).hexdigest())
|
||||
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())
|
||||
zero_sign_differences=[]
|
||||
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['value_sha256']==b['value_sha256'],(phase,family,layer,slot,a,b)
|
||||
if a['sha256']!=b['sha256']: zero_sign_differences.append([family,layer,slot])
|
||||
assert zero_sign_differences==[], zero_sign_differences
|
||||
# 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 np.all(key[different]==0) and np.all(want[different]==0)
|
||||
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,values_exact=True,zero_sign_differences=zero_sign_differences,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='issue113',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)
|
||||
runpy.run_path(str(Path(__file__).with_name("issue109-model-reference.py")), run_name="__main__")
|
||||
Reference in new issue
Block a user