Files

125 lines
7.1 KiB
Python

"""CPU-only, read-only parameter cloud census with paired uniform references."""
import argparse, hashlib, json, math, os, time
from pathlib import Path
from types import SimpleNamespace
import numpy as np
from pilot.tensors import load_tensors
from finetune_compare.lora import install
from model_shell.groups import enumerate_groups
MAX_GRID=np.array([.5,.75,.9,.95,.99,1.,1.05,1.1,1.25])
RMS_GRID=np.linspace(.5,1.5,41)
def seed(key):return int.from_bytes(hashlib.sha256(key.encode()).digest()[:8],'little')
def guard(start):
mem=dict((s.split(':')[0],int(s.split()[1])) for s in Path('/proc/meminfo').read_text().splitlines() if s.startswith('MemAvailable:'))
rss=int(Path('/proc/self/statm').read_text().split()[1])*os.sysconf('SC_PAGE_SIZE')
if rss>512*1024**2 or mem['MemAvailable']<128*1024 or time.monotonic()-start>7200:
raise RuntimeError(f'Resource guard: RSS={rss}, available_kB={mem}, elapsed={time.monotonic()-start:.0f}')
def sample(t,g):
rng=np.random.default_rng(seed(g['id']+'/indices'))
n=min(512,g['population']);d=min(8192,g['dimension'])
points=np.sort(rng.choice(g['population'],n,replace=False))
features=np.sort(rng.choice(g['dimension'],d,replace=False)) if d<g['dimension'] else np.arange(d)
width=math.prod(t.shape[1:]);out=np.empty((n,d),np.float32)
if g['orientation']=='rows':
for j,i in enumerate(points):
row=i+g['start'];out[j]=t.read(slice(row*width,(row+1)*width))[features]
if j%32==0:t.evict()
else:
for j,i in enumerate(features):
out[:,j]=t.read(slice(i*width,(i+1)*width))[points]
if j%32==0:t.evict()
t.evict();out=out[rng.permutation(n)]
digest=hashlib.sha256(points.tobytes()+features.tobytes()).hexdigest()
return out,digest
def participation(y):
gram=y@y.T;energy=float(np.trace(gram));den=float(np.square(gram).sum())
return energy**2/den if den else 0.
def geometry(x,ncal,center,raw=False):
# Float64 arithmetic; bounded sample matrix, no model-wide conversion.
z=x.astype(np.float64);mu=z[:ncal].mean(0) if center=='mean' else np.zeros(z.shape[1])
y=z-mu;norm=np.linalg.norm(y,axis=1);cal=norm[:ncal];r=norm[ncal:]
maximum=float(cal.max());rms=float(np.sqrt(np.mean(cal**2)))
valid=maximum>0 and rms>0
a=r/maximum if maximum else np.zeros_like(r);b=r/rms if rms else np.zeros_like(r)
coordvar=z[:ncal].var(0,ddof=1);cv=float(coordvar.sum()**2/np.square(coordvar).sum()) if coordvar.any() else 0.
e=np.square(y[ncal:]);den=np.square(e).sum(1)
pe=np.divide(e.sum(1)**2,den,out=np.zeros(len(e)),where=den>0)
out=dict(valid=bool(valid),ncal=ncal,neval=len(r),d=z.shape[1],rmax=maximum,rms=rms,
median_max=float(np.median(a)),median_rms=float(np.median(b)),
p10_rms=float(np.quantile(b,.1)),p90_rms=float(np.quantile(b,.9)),
outside=float(np.mean(a>1)),cdf_max=[float(np.mean(a<=v)) for v in MAX_GRID],
cdf_rms=[float(np.mean(b<=v)) for v in RMS_GRID],
covariance_participation=participation(y[:ncal]),rank_ceiling=ncal-(center=='mean'),
coordinate_variance_participation=cv,point_energy_participation_median=float(np.median(pe)),
zero_fraction=float(np.mean(z==0)))
if raw:out.update(cal_norm=cal.tolist(),eval_norm=r.tolist())
return out
BALL_CACHE={}
def references(x,ncal,key):
d=x.shape[1];cachekey=(d,ncal)
if cachekey not in BALL_CACHE:
rng=np.random.default_rng(seed(f'ball/{d}/{ncal}'))
y=rng.standard_normal((ncal+512,d),dtype=np.float32)
y*=((rng.random(len(y))**(1/d))/np.linalg.norm(y,axis=1))[:,None]
BALL_CACHE[cachekey]={c:geometry(y,ncal,c) for c in ('origin','mean')}
cal=x[:ncal].astype(np.float64);lo=cal.min(0);hi=cal.max(0);mu=cal.mean(0);sd=cal.std(0,ddof=1)
rng=np.random.default_rng(seed(key+'/uniform'))
u=rng.random((ncal+512,d),dtype=np.float32)
refs={'ball':BALL_CACHE[cachekey]}
for name,left,right in [('range_box',lo,hi),('moment_box',mu-np.sqrt(3)*sd,mu+np.sqrt(3)*sd)]:
y=(left+(right-left)*u).astype(np.float32)
refs[name]={c:geometry(y,ncal,c) for c in ('origin','mean')}
return refs,dict(heldout_outside_coordinate_range=float(np.mean((x[ncal:]<lo)|(x[ncal:]>hi))),
range_box_expected_energy=float(np.sum((lo**2+lo*hi+hi**2)/3)),
observed_cal_energy=float(np.mean(np.sum(cal**2,axis=1))))
def analyze(x,key):
ncal=len(x)//2;refs,extra=references(x,ncal,key);out={}
for center in ('origin','mean'):
p=geometry(x,ncal,center,True);p['references']=refs; p.update(extra)
p['comparisons']={}
for name,bycenter in refs.items():
ref=bycenter[center];diff=np.array(p['cdf_rms'])-ref['cdf_rms']
p['comparisons'][name]=dict(interior90_gap=p['cdf_max'][2]-ref['cdf_max'][2],
interior95_gap=p['cdf_max'][3]-ref['cdf_max'][3],
rms_cdf_positive_gap=float(max(0,diff.max())),rms_cdf_negative_gap=float(max(0,-diff.min())),
median_rms_delta=p['median_rms']-ref['median_rms'],
participation_ratio=p['covariance_participation']/ref['covariance_participation'] if ref['covariance_participation'] else None)
p['ball_origin_log_expected_cdf90']=float(x.shape[1]*np.log(.9)+np.log(ncal/(ncal+1)))
out[center]=p
return out
def main():
ap=argparse.ArgumentParser();ap.add_argument('--root',default='runs');ap.add_argument('--out',default='runs/model-shell');args=ap.parse_args()
root=Path(args.root);out=Path(args.out);out.mkdir(exist_ok=True,parents=True);(out/'checkpoints').mkdir(exist_ok=True)
cfg=json.loads((root/'cpu-qwen35-2b/model_config.json').read_text())['text_config']
base=load_tensors(root/'cpu-qwen35-2b/weights');ft=SimpleNamespace(tensors=dict(base),config=cfg)
meta=json.loads((root/'finetune-head-comparison/adapter_meta.json').read_text())
pairs,scale=install(ft,root/'finetune-head-comparison',meta)
groups,skips=enumerate_groups(base,cfg)
manifest=dict(groups=groups,skips=skips,total_tensors=len(base),parameters=sum(t.numel for t in base.values()),changed_tensors=sorted(pairs),adapter_meta=meta,
max_grid=MAX_GRID.tolist(),rms_grid=RMS_GRID.tolist(),null_eval=512,point_cap=512,feature_cap=8192)
(out/'inventory.json').write_text(json.dumps(manifest,indent=2));start=time.monotonic()
for i,g in enumerate(groups):
target=out/'checkpoints'/f'{i:04d}.json'
if target.exists():continue
guard(start);x,digest=sample(base[g['tensor']],g);p=analyze(x,g['id']);del x
changed=g['tensor'] in pairs
if changed:
guard(start);x,digest2=sample(ft.tensors[g['tensor']],g);assert digest==digest2
q=analyze(x,g['id']);del x
else:q=None
record=dict(group=g,sampling_sha256=digest,changed=changed,original=p,finetuned=q,finetuned_reuses_original=not changed)
temp=target.with_suffix('.tmp');temp.write_text(json.dumps(record,allow_nan=False));temp.replace(target)
if i%10==0:print(f'{i+1}/{len(groups)} elapsed={time.monotonic()-start:.1f}s',flush=True)
print('COMPLETE',len(groups),flush=True)
if __name__=='__main__':main()