397 lines
19 KiB
Python
397 lines
19 KiB
Python
"""Bounded CPU-only screening, exact verification and conservative sharing groups."""
|
|
import argparse
|
|
import collections
|
|
import csv
|
|
import hashlib
|
|
import itertools
|
|
import json
|
|
import math
|
|
import pathlib
|
|
import platform
|
|
import resource
|
|
import time
|
|
|
|
import numpy as np
|
|
|
|
from .tensors import decode, finish, load_tensors, metrics
|
|
|
|
THRESHOLDS = [.01, .02, .05, .10, .20]
|
|
|
|
|
|
def write_json(path, obj):
|
|
path.write_text(json.dumps(obj, indent=2, allow_nan=False))
|
|
|
|
|
|
def write_csv(path, rows):
|
|
if not rows:
|
|
path.write_text('')
|
|
return
|
|
with path.open('w', newline='') as f:
|
|
writer = csv.DictWriter(f, fieldnames=list(rows[0]))
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
|
|
|
|
def stable_seed(key, seed):
|
|
return int.from_bytes(hashlib.sha256((str(seed) + str(key)).encode()).digest()[:8], 'little')
|
|
|
|
|
|
def read_object(obj, tensors):
|
|
t = tensors[obj['tensor']]
|
|
if obj['kind'] == 'tensor':
|
|
return t.read()
|
|
slices = tuple(slice(s, s + b) for s, b in zip(obj['start'], obj['shape']))
|
|
return decode(t.raw.reshape(t.shape)[slices], t.dtype).ravel()
|
|
|
|
|
|
def sample_object(obj, tensors, sample_size=2048):
|
|
t = tensors[obj['tensor']]
|
|
n = math.prod(obj['shape'])
|
|
indices = np.linspace(0, n - 1, min(n, sample_size), dtype=np.int64)
|
|
if obj['kind'] == 'tensor':
|
|
return t.read(indices)
|
|
coords = np.unravel_index(indices, obj['shape'])
|
|
global_coords = tuple(c + s for c, s in zip(coords, obj['start']))
|
|
flat = np.ravel_multi_index(global_coords, t.shape)
|
|
return t.read(flat)
|
|
|
|
|
|
def exact_pair(a, b, tensors, deadline):
|
|
if a['kind'] != 'tensor':
|
|
result = metrics(read_object(a, tensors), read_object(b, tensors))
|
|
tensors[a['tensor']].evict()
|
|
tensors[b['tensor']].evict()
|
|
return result
|
|
ta, tb = tensors[a['tensor']], tensors[b['tensor']]
|
|
xx = yy = xy = delta = sx = sy = 0.
|
|
for start in range(0, ta.numel, 262144):
|
|
if time.monotonic() > deadline:
|
|
raise TimeoutError('Analysis deadline reached')
|
|
sl = slice(start, min(start + 262144, ta.numel))
|
|
x, y = ta.read(sl).astype(np.float64), tb.read(sl).astype(np.float64)
|
|
xx += float(x @ x)
|
|
yy += float(y @ y)
|
|
xy += float(x @ y)
|
|
d = x - y
|
|
delta += float(d @ d)
|
|
sx += float(x.sum())
|
|
sy += float(y.sum())
|
|
ta.evict()
|
|
tb.evict()
|
|
ta.evict()
|
|
tb.evict()
|
|
return finish(xx, yy, xy, delta, sx, sy, ta.numel)
|
|
|
|
|
|
def sharing_groups(objects, pairs, threshold):
|
|
"""Star groups; every member must have an exact verified edge to its representative."""
|
|
adjacent = collections.defaultdict(dict)
|
|
for pair in pairs:
|
|
a, b = pair['a'], pair['b']
|
|
if pair['error_member_to_rep'] <= threshold:
|
|
adjacent[b][a] = pair['error_member_to_rep']
|
|
if pair['error_rep_to_member'] <= threshold:
|
|
adjacent[a][b] = pair['error_rep_to_member']
|
|
unassigned = set(range(len(objects)))
|
|
order = sorted(unassigned, key=lambda i: (-len(adjacent[i]), objects[i]['id']))
|
|
groups = []
|
|
for rep in order:
|
|
if rep not in unassigned:
|
|
continue
|
|
members = sorted(set(adjacent[rep]) & unassigned)
|
|
members.append(rep)
|
|
unassigned.difference_update(members)
|
|
groups.append({'representative': rep, 'members': members,
|
|
'max_error': max([adjacent[rep][i] for i in members if i != rep] or [0.])})
|
|
return groups
|
|
|
|
|
|
def nearest_candidates(objects, features, cap, neighbors=10):
|
|
buckets = collections.defaultdict(list)
|
|
for i, obj in enumerate(objects):
|
|
buckets[(obj['role'], tuple(obj['shape']), obj['scale'], obj['dtype'])].append(i)
|
|
candidates = []
|
|
eligible_pairs = 0
|
|
for key, ids in sorted(buckets.items(), key=lambda item: str(item[0])):
|
|
eligible_pairs += len(ids) * (len(ids) - 1) // 2
|
|
if len(ids) < 2:
|
|
continue
|
|
x = np.stack([features[i] for i in ids]).astype(np.float64)
|
|
norm = np.linalg.norm(x, axis=1)
|
|
u = x / np.maximum(norm[:, None], 1e-30)
|
|
# Cap each bucket's reservoir to bound the screening matrix.
|
|
cos = u @ u.T
|
|
sq = np.maximum(0., norm[:, None] ** 2 + norm[None, :] ** 2 - 2 * (x @ x.T))
|
|
err = np.sqrt(sq / np.maximum(norm[:, None] ** 2, 1e-30))
|
|
np.fill_diagonal(cos, -np.inf)
|
|
np.fill_diagonal(err, np.inf)
|
|
seen = set()
|
|
for row, i in enumerate(ids):
|
|
near = set(np.argsort(err[row])[:neighbors]) | set(np.argsort(-cos[row])[:neighbors])
|
|
for col in near:
|
|
j = ids[col]
|
|
edge = (min(i, j), max(i, j))
|
|
if i == j or edge in seen:
|
|
continue
|
|
seen.add(edge)
|
|
candidates.append({'a': edge[0], 'b': edge[1], 'screen_error': float(min(err[row, col], err[col, row])),
|
|
'screen_cosine': float(cos[row, col])})
|
|
# Low-error candidates first, then guarantee role/scale diversity round-robin.
|
|
queues = collections.defaultdict(list)
|
|
for c in sorted(candidates, key=lambda c: (c['screen_error'], c['a'], c['b'])):
|
|
o = objects[c['a']]
|
|
queues[(o['role'], tuple(o['shape']), o['scale'])].append(c)
|
|
queues = [collections.deque(q) for _, q in sorted(queues.items(), key=lambda x: str(x[0]))]
|
|
chosen = []
|
|
while queues and len(chosen) < cap:
|
|
for q in queues:
|
|
if q and len(chosen) < cap:
|
|
chosen.append(q.popleft())
|
|
queues = [q for q in queues if q]
|
|
return chosen, {'eligible_pairs_in_screened_objects': eligible_pairs,
|
|
'candidate_pairs_before_cap': len(candidates), 'selected_pairs': len(chosen)}
|
|
|
|
|
|
def object_record(t, kind='tensor', shape=None, start=None, scale='whole'):
|
|
shape = tuple(shape or t.shape)
|
|
return {'id': t.name if kind == 'tensor' else f'{t.name}|{scale}|{start}',
|
|
'tensor': t.name, 'role': t.role, 'layer': t.layer, 'kind': kind, 'dtype': t.dtype,
|
|
'shape': shape, 'start': start, 'scale': scale,
|
|
'bytes': math.prod(shape) * (t.bytes // t.numel)}
|
|
|
|
|
|
def build_blocks(tensors, seed, deadline, reservoir_limit=256):
|
|
reservoirs, counts, rngs = {}, collections.Counter(), {}
|
|
totals = collections.defaultdict(lambda: {'blocks': 0, 'bytes': 0, 'edge_bytes': 0})
|
|
for t in tensors.values():
|
|
if len(t.shape) == 2:
|
|
specs = [(f'2d-{s}', (s, s)) for s in (64, 128) if min(t.shape) >= s]
|
|
elif len(t.shape) == 3:
|
|
specs = [('3d-axis0-slice', (1, min(64, t.shape[1]), min(64, t.shape[2]))),
|
|
('3d-volume', (min(4, t.shape[0]), min(32, t.shape[1]), min(32, t.shape[2])))]
|
|
else:
|
|
continue
|
|
for scale, shape in specs:
|
|
nblocks = math.prod(d // b for d, b in zip(t.shape, shape))
|
|
totals[scale]['blocks'] += nblocks
|
|
block_bytes = math.prod(shape) * (t.bytes // t.numel)
|
|
totals[scale]['bytes'] += nblocks * block_bytes
|
|
totals[scale]['edge_bytes'] += t.bytes - nblocks * block_bytes
|
|
for grid in itertools.product(*(range(d // b) for d, b in zip(t.shape, shape))):
|
|
if time.monotonic() > deadline:
|
|
raise TimeoutError('Deadline during block enumeration')
|
|
start = tuple(i * b for i, b in zip(grid, shape))
|
|
obj = object_record(t, 'block', shape, start, scale)
|
|
key = (t.role, shape, scale)
|
|
if key not in reservoirs:
|
|
reservoirs[key], rngs[key] = [], np.random.default_rng(stable_seed(key, seed))
|
|
counts[key] += 1
|
|
reservoir = reservoirs[key]
|
|
if len(reservoir) < reservoir_limit:
|
|
reservoir.append(obj)
|
|
else:
|
|
j = int(rngs[key].integers(counts[key]))
|
|
if j < reservoir_limit:
|
|
reservoir[j] = obj
|
|
objects = [obj for key in sorted(reservoirs, key=str) for obj in reservoirs[key]]
|
|
coverage = [{'role': key[0], 'shape': str(key[1]), 'scale': key[2],
|
|
'total_blocks': counts[key], 'sampled_blocks': len(reservoirs[key])}
|
|
for key in sorted(reservoirs, key=str)]
|
|
return objects, coverage, dict(totals)
|
|
|
|
|
|
def structural_probe(a, b, tensors):
|
|
"""Bounded mode unfoldings of aligned subvolumes; exploratory, never a grouping criterion."""
|
|
def small(obj):
|
|
t = tensors[obj['tensor']]
|
|
shape = obj['shape']
|
|
starts = obj['start'] or [0] * len(shape)
|
|
slices = tuple(slice(s, s + min(d, 32)) for s, d in zip(starts, shape))
|
|
return decode(t.raw.reshape(t.shape)[slices], t.dtype)
|
|
x, y = small(a), small(b)
|
|
result = []
|
|
if x.ndim < 2:
|
|
return result
|
|
for axis in range(x.ndim):
|
|
mx = np.moveaxis(x, axis, 0).reshape(x.shape[axis], -1)
|
|
my = np.moveaxis(y, axis, 0).reshape(y.shape[axis], -1)
|
|
if min(mx.shape) < 2:
|
|
result.append({'axis': axis, 'probe_shape': list(x.shape), 'trivial_mode_skipped': True})
|
|
continue
|
|
ux, sx, _ = np.linalg.svd(mx, full_matrices=False)
|
|
uy, sy, _ = np.linalg.svd(my, full_matrices=False)
|
|
k = min(4, max(1, min(len(sx), len(sy)) // 2))
|
|
result.append({'axis': axis, 'probe_shape': list(x.shape),
|
|
'k': k,
|
|
'normalized_spectrum_distance': float(np.linalg.norm(sx / max(np.linalg.norm(sx), 1e-30) - sy / max(np.linalg.norm(sy), 1e-30))),
|
|
'left_subspace_overlap': float(np.linalg.norm(ux[:, :k].T @ uy[:, :k]) ** 2 / k)})
|
|
return result
|
|
|
|
|
|
def analyze_mode(mode, objects, tensors, root, deadline, cap, seed, baseline_cap=100):
|
|
print(f'{mode}: sampling {len(objects)} objects', flush=True)
|
|
features = []
|
|
for obj in objects:
|
|
features.append(sample_object(obj, tensors))
|
|
tensors[obj['tensor']].evict()
|
|
candidates, coverage = nearest_candidates(objects, features, cap)
|
|
write_json(root / f'{mode}_objects.json', objects)
|
|
pairs, nulls, structures = [], [], []
|
|
rng = np.random.default_rng(seed)
|
|
random_pairs = []
|
|
comparable = collections.defaultdict(list)
|
|
for j, obj in enumerate(objects):
|
|
comparable[(obj['role'], tuple(obj['shape']), obj['scale'], obj['dtype'])].append(j)
|
|
for key, ids in sorted(comparable.items(), key=lambda item: str(item[0])):
|
|
if len(ids) < 2:
|
|
continue
|
|
for _ in range(3):
|
|
a, b = rng.choice(ids, size=2, replace=False).tolist()
|
|
random_pairs.append({'a': a, 'b': b, 'role': key[0], 'sampled': True,
|
|
**metrics(features[a], features[b])})
|
|
write_json(root / f'{mode}_random_pairs.json', random_pairs)
|
|
timeout = False
|
|
structural_counts = collections.Counter()
|
|
for i, candidate in enumerate(candidates):
|
|
if time.monotonic() > deadline:
|
|
timeout = True
|
|
break
|
|
a, b = candidate['a'], candidate['b']
|
|
try:
|
|
measured = exact_pair(objects[a], objects[b], tensors, deadline)
|
|
except TimeoutError:
|
|
timeout = True
|
|
break
|
|
row = {**candidate, **measured,
|
|
'a_id': objects[a]['id'], 'b_id': objects[b]['id'],
|
|
'role': objects[a]['role'], 'scale_name': objects[a]['scale'],
|
|
'layer_distance': abs(objects[a]['layer'] - objects[b]['layer'])}
|
|
row['screen_vs_exact_min_error_absolute_difference'] = abs(candidate['screen_error'] - min(measured['error_member_to_rep'], measured['error_rep_to_member']))
|
|
# Tail diagnostics are sampled, explicitly labelled, not claimed exact.
|
|
fx, fy = features[a].astype(np.float64), features[b].astype(np.float64)
|
|
row['sampled_error_q99_over_member_rms'] = float(np.quantile(np.abs(fx - fy), .99) / max(np.sqrt(np.mean(fx * fx)), 1e-30))
|
|
pairs.append(row)
|
|
if i < baseline_cap:
|
|
for method, y in [('shuffled', rng.permutation(fy)), ('gaussian_norm_matched', rng.normal(size=fy.size))]:
|
|
if method == 'gaussian_norm_matched':
|
|
y *= np.linalg.norm(fy) / max(np.linalg.norm(y), 1e-30)
|
|
nulls.append({'a': a, 'b': b, 'method': method, 'sampled': True, **metrics(fx, y)})
|
|
ndim = len(objects[a]['shape'])
|
|
if ndim >= 2 and structural_counts[ndim] < 20:
|
|
structures.append({'a': a, 'b': b, 'modes': structural_probe(objects[a], objects[b], tensors)})
|
|
tensors[objects[a]['tensor']].evict()
|
|
tensors[objects[b]['tensor']].evict()
|
|
structural_counts[ndim] += 1
|
|
if (i + 1) % 100 == 0:
|
|
write_csv(root / f'{mode}_pairs.csv', pairs)
|
|
print(f'{mode}: verified {i + 1}/{len(candidates)}', flush=True)
|
|
write_csv(root / f'{mode}_pairs.csv', pairs)
|
|
write_json(root / f'{mode}_nulls.json', nulls)
|
|
write_json(root / f'{mode}_structure.json', structures)
|
|
coverage.update({'objects_screened': len(objects), 'pairs_verified': len(pairs), 'timed_out': timeout,
|
|
'screen_sample_elements': 2048, 'candidate_neighbors_per_metric': 10})
|
|
write_json(root / f'{mode}_coverage.json', coverage)
|
|
return pairs
|
|
|
|
|
|
def summarize(mode, objects, pairs, model_bytes, root, partition_totals=None):
|
|
rows = []
|
|
scales = sorted({o['scale'] for o in objects})
|
|
for scale in scales:
|
|
ids = [i for i, o in enumerate(objects) if o['scale'] == scale]
|
|
selected = set(ids)
|
|
for threshold in THRESHOLDS:
|
|
groups = sharing_groups(objects, [p for p in pairs if p['a'] in selected], threshold)
|
|
groups = [g for g in groups if g['representative'] in selected]
|
|
original = sum(objects[i]['bytes'] for i in ids)
|
|
rep_bytes = sum(objects[g['representative']]['bytes'] for g in groups)
|
|
# Concrete theoretical layout: 16 bytes per object mapping, plus 64 bytes per shared group.
|
|
# Identity case falls back to original checkpoint at zero additional cost.
|
|
mapping_bytes = 16 * len(ids) + 64 * sum(len(g['members']) > 1 for g in groups)
|
|
proposed = rep_bytes + mapping_bytes
|
|
saved = max(0, original - proposed)
|
|
partition_count = partition_totals[scale]['blocks'] if partition_totals else len(ids)
|
|
rows.append({'mode': mode, 'scale': scale, 'threshold': threshold,
|
|
'objects': len(ids), 'groups': len(groups),
|
|
'total_partition_objects': partition_count,
|
|
'unexamined_partition_objects': partition_count - len(ids),
|
|
'conservative_partition_groups': partition_count - len(ids) + len(groups),
|
|
'shared_groups': sum(len(g['members']) > 1 for g in groups),
|
|
'max_group_size': max([len(g['members']) for g in groups] or [0]),
|
|
'verified_shared_bytes': sum(sum(objects[i]['bytes'] for i in g['members']) for g in groups if len(g['members']) > 1),
|
|
'examined_object_bytes': original, 'representative_bytes': rep_bytes,
|
|
'mapping_bytes_if_encoded': mapping_bytes, 'net_saved_bytes': saved,
|
|
'net_saved_model_fraction': saved / model_bytes,
|
|
'max_replacement_error': max([g['max_error'] for g in groups] or [0.])})
|
|
write_json(root / f'groups_{mode}_{scale}_{threshold:.2f}.json', groups)
|
|
return rows
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument('--root', type=pathlib.Path, required=True)
|
|
p.add_argument('--seconds', type=int, default=7200)
|
|
p.add_argument('--seed', type=int, default=20261109)
|
|
args = p.parse_args()
|
|
started = time.monotonic()
|
|
deadline = started + args.seconds
|
|
root = args.root
|
|
output = root / 'analysis'
|
|
output.mkdir(exist_ok=True)
|
|
tensors = load_tensors(root / 'weights')
|
|
inventory = []
|
|
for t in tensors.values():
|
|
if time.monotonic() > deadline:
|
|
raise TimeoutError('Deadline during inventory')
|
|
n = 0
|
|
total = squared = 0.
|
|
max_abs = 0.
|
|
for x in t.chunked():
|
|
if not np.isfinite(x).all():
|
|
raise ValueError(f'Nonfinite weights: {t.name}')
|
|
z = x.astype(np.float64)
|
|
n += z.size
|
|
total += float(z.sum())
|
|
squared += float(z @ z)
|
|
max_abs = max(max_abs, float(np.abs(z).max()))
|
|
t.evict()
|
|
inventory.append({'name': t.name, 'shape': str(t.shape), 'ndim': len(t.shape), 'dtype': t.dtype,
|
|
'numel': t.numel, 'bytes': t.bytes, 'role': t.role, 'layer': t.layer,
|
|
'mean': total / n, 'rms': math.sqrt(squared / n), 'max_abs': max_abs})
|
|
t.evict()
|
|
write_csv(output / 'inventory.csv', inventory)
|
|
print(f'Inventory: {len(tensors)} tensors, {sum(t.bytes for t in tensors.values())} bytes', flush=True)
|
|
whole = [object_record(t) for t in tensors.values()]
|
|
# Stage budgets prevent the whole-tensor stage from consuming block-analysis time.
|
|
whole_deadline = min(deadline, started + args.seconds * .50)
|
|
whole_pairs = analyze_mode('tensor', whole, tensors, output, whole_deadline, 2000, args.seed)
|
|
blocks, block_coverage, block_totals = build_blocks(tensors, args.seed, deadline)
|
|
write_csv(output / 'block_sampling.csv', block_coverage)
|
|
write_json(output / 'block_partitions.json', block_totals)
|
|
block_pairs = analyze_mode('block', blocks, tensors, output, deadline, 20000, args.seed)
|
|
model_bytes = sum(t.bytes for t in tensors.values())
|
|
rows = summarize('tensor', whole, whole_pairs, model_bytes, output)
|
|
rows += summarize('block', blocks, block_pairs, model_bytes, output, block_totals)
|
|
write_csv(output / 'compression_curve.csv', rows)
|
|
summary = {'model': json.loads((root / 'manifest.json').read_text())['repo'],
|
|
'revision': json.loads((root / 'manifest.json').read_text())['revision'],
|
|
'tensor_count': len(tensors), 'parameter_count': sum(t.numel for t in tensors.values()),
|
|
'model_weight_bytes': model_bytes, 'ndim_counts': dict(collections.Counter(len(t.shape) for t in tensors.values())),
|
|
'elapsed_seconds': time.monotonic() - started, 'analysis_budget_seconds': args.seconds,
|
|
'peak_rss_kib': resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
|
|
'python': platform.python_version(), 'numpy': np.__version__, 'seed': args.seed,
|
|
'limitations': ['No inference or functional equivalence measured', 'No exhaustive pair search',
|
|
'Block search uses per-role/shape/scale reservoirs of 256',
|
|
'Different block scales overlap; savings must not be added',
|
|
'Block savings count only verified disjoint sampled blocks, no extrapolation',
|
|
'Spectral probes use aligned bounded subvolumes, not whole tensors',
|
|
'Shared encoding is theoretical, not a runnable compressed checkpoint',
|
|
'No cross-role grouping, neuron permutation alignment, or training'],
|
|
'completed': True}
|
|
write_json(output / 'summary.json', summary)
|
|
print(json.dumps(summary), flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|