Files
exp-model-parameter-similarity/functional/storage.py
T

63 lines
3.6 KiB
Python

"""Measure an explicit two-component alias storage layout, not inference speed.
Payloads remain on the server under weights/ and must not be committed or cleaned
without the user's explicit deletion request. Manifest/byte/hash evidence may be archived.
"""
import argparse
import hashlib
import json
from pathlib import Path
import numpy as np
from pilot.tensors import load_tensors,decode
from .run import write_json
def payload(t,path):
digest=hashlib.sha256()
with path.open('wb') as f:
for start in range(0,t.raw.size,262144):
b=t.raw[start:start+262144].tobytes();f.write(b);digest.update(b)
t.evict()
return {'file':path.name,'dtype':t.dtype,'shape':list(t.shape),'bytes':path.stat().st_size,'sha256':digest.hexdigest()}
def main():
p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);a=p.parse_args();root=a.root
selected=json.loads((root/'selection.json').read_text())['selected']
if not selected:
write_json(root/'storage-evidence.json',{'measured':False,'reason':'No calibration candidate selected'});return
c=selected[0]
if c['kind']!='qk_norm':
write_json(root/'storage-evidence.json',{'measured':False,'reason':'Only two-gain alias layout implemented in this bounded pilot; larger-module savings remain theoretical'});return
tensors=load_tensors(root/'weights')
prefix=f'model.language_model.layers.{c["i"]}.self_attn.'
names=[prefix+'q_norm.weight',prefix+'k_norm.weight']
# Original pair plus two alternative layouts; each is counted independently.
path=root/'component-layout'/'weights';path.mkdir(parents=True,exist_ok=True)
original={'format':'two_component_raw_v1','tensors':{}}
for short,name in zip(('q','k'),names):original['tensors'][name]=payload(tensors[name],path/f'{short}.bin')
original_bytes=json.dumps(original,separators=(',',':'),sort_keys=True).encode()
(path/'original-layout.json').write_bytes(original_bytes)
evidence=[]
for target,donor in ((names[0],names[1]),(names[1],names[0])):
entry=original['tensors'][donor]
layout={'format':'two_component_raw_v1','tensors':{donor:entry},'aliases':{target:donor}}
encoded=json.dumps(layout,separators=(',',':'),sort_keys=True).encode()
file=path/('share-'+('k' if donor==names[1] else 'q')+'.json');file.write_bytes(encoded)
t=tensors[donor]
raw=np.fromfile(path/entry['file'],dtype=t.raw.dtype)
reconstructed=decode(raw,t.dtype).reshape(t.shape)
np.testing.assert_array_equal(reconstructed,t.read().reshape(t.shape));t.evict()
before=sum(e['bytes'] for e in original['tensors'].values())+len(original_bytes)
after=entry['bytes']+len(encoded)
evidence.append({'target':target,'donor':donor,'original_payload_bytes':sum(e['bytes'] for e in original['tensors'].values()),
'shared_payload_bytes':entry['bytes'],'original_descriptor_bytes':len(original_bytes),
'shared_descriptor_bytes':len(encoded),'original_total_bytes':before,'shared_total_bytes':after,
'net_saved_bytes':before-after,'shared_descriptor_sha256':hashlib.sha256(encoded).hexdigest(),
'decoded_donor_exactly_matches_intervention_gain':True})
write_json(root/'storage-evidence.json',{'measured':True,'candidate':c['id'],'layouts':evidence,
'claim_limit':'Measured standalone component layout including descriptors, not an integrated compressed checkpoint, latency gain or full-model RSS reduction. Two sharing directions are alternatives and must not be added.'})
if __name__=='__main__':main()