30 lines
2.2 KiB
Python
30 lines
2.2 KiB
Python
"""Deterministic finite family of matrix/vector clouds; no data-dependent choice."""
|
|
import math,re
|
|
|
|
def enumerate_groups(tensors,cfg):
|
|
groups=[];skips=[]
|
|
for name,t in sorted(tensors.items()):
|
|
if len(t.shape)<2:skips.append({'name':name,'reason':'one vector, no vector cloud','shape':t.shape});continue
|
|
rows=t.shape[0];cols=math.prod(t.shape[1:]);count=0
|
|
for axis,n,d in [('rows',rows,cols),('columns',cols,rows)]:
|
|
if n>=16 and d>=32:
|
|
groups.append({'id':name+'::'+axis,'tensor':name,'kind':'matrix','orientation':axis,'start':0,'stop':rows,'population':n,'dimension':d,'stratum':'primary' if n>=64 else 'small_sample'});count+=1
|
|
if not count:skips.append({'name':name,'reason':'insufficient points or dimension <32','shape':t.shape})
|
|
if not (name.startswith('model.language_model.layers.') or name.startswith('mtp.layers.')):continue
|
|
segments=[]
|
|
if name.endswith('.linear_attn.in_proj_qkv.weight'):
|
|
kh,kd,vh,vd=[cfg[k] for k in ('linear_num_key_heads','linear_key_head_dim','linear_num_value_heads','linear_value_head_dim')]
|
|
for role,num,dim,offset in [('Q',kh,kd,0),('K',kh,kd,kh*kd),('V',vh,vd,2*kh*kd)]:
|
|
for h in range(num):segments.append((role,h,offset+h*dim,offset+(h+1)*dim))
|
|
elif name.endswith('.self_attn.q_proj.weight'):
|
|
dim=cfg['head_dim'];assert rows==cfg['num_attention_heads']*2*dim
|
|
for h in range(cfg['num_attention_heads']):
|
|
for role,offset in [('Q',0),('gate',dim)]:segments.append((role,h,h*2*dim+offset,h*2*dim+offset+dim))
|
|
elif name.endswith(('.self_attn.k_proj.weight','.self_attn.v_proj.weight')):
|
|
role='K' if '.k_proj.' in name else 'V';dim=cfg['head_dim'];assert rows==cfg['num_key_value_heads']*dim
|
|
for h in range(cfg['num_key_value_heads']):segments.append((role,h,h*dim,(h+1)*dim))
|
|
for role,h,start,stop in segments:
|
|
assert stop<=rows
|
|
groups.append({'id':name+f'::{role}_head{h}','tensor':name,'kind':'head_rows','orientation':'rows','start':start,'stop':stop,'population':stop-start,'dimension':cols,'stratum':'primary' if stop-start>=64 else 'small_sample'})
|
|
return groups,skips
|