File size: 4,582 Bytes
fc32096 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | import argparse
import hashlib
import json
import numpy as np
import torch
from common import ROOT, inside, permit_compute, save_json
from weights import load, save_edit, export, project_writers
permit_compute()
p = argparse.ArgumentParser()
p.add_argument('--analysis', required=True)
p.add_argument('--layers', required=True)
p.add_argument('--alpha', type=float, required=True)
p.add_argument('--topk', type=int, default=64)
p.add_argument('--name', required=True)
p.add_argument('--control', choices=['targeted','orthogonal','random_columns'], default='targeted')
p.add_argument('--seed', type=int, default=1701)
p.add_argument('--hypothesis', required=True)
p.add_argument('--protect-coefficients', action='store_true')
p.add_argument('--match-protected-norm', action='store_true')
a = p.parse_args()
if not a.name or any(c not in 'abcdefghijklmnopqrstuvwxyz0123456789_-' for c in a.name):
raise ValueError('Используй простое уникальное имя кандидата')
directory, model = ROOT/'edits'/a.name, ROOT/'models'/a.name
if directory.exists() or model.exists():
raise ValueError('Имя уже занято')
z = np.load(inside(a.analysis)/'directions_and_writers.npz')
records = []
for layer in map(int,a.layers.split(',')):
if not 0 <= layer < 24:
raise ValueError('Недопустимый слой')
name = f'model.language_model.layers.{layer}.mlp.down_proj.weight'
weight = load(ROOT/'models/original', name).float()
u = torch.from_numpy(z[f'u_{layer}'])
scores = z[f'scores_{layer}']
k = a.topk or weight.shape[1]
if not 1 <= k <= weight.shape[1]:
raise ValueError('Недопустимое число столбцов')
columns = np.argsort(scores)[-k:].copy()
targeted = project_writers(weight, u, columns, a.alpha)
candidate = targeted
rng = np.random.default_rng(a.seed+layer)
if a.control == 'orthogonal':
v = torch.tensor(rng.normal(size=u.numel()),dtype=torch.float32)
v = v-u*torch.dot(u,v)
v = v/v.norm()
delta = weight-project_writers(weight,v,columns,1)
candidate = weight-delta*((weight-targeted).norm()/delta.norm())
elif a.control == 'random_columns':
selected = rng.choice(weight.shape[1],size=k,replace=False)
delta = weight-project_writers(weight,u,selected,1)
candidate = weight-delta*((weight-targeted).norm()/delta.norm())
columns = selected
protection = None
if a.protect_coefficients:
if a.control != 'targeted':
raise ValueError('Coefficient protection currently requires targeted control')
productive = z[f'productive_mlp_{layer}'][:,columns].astype(np.float64)
_, singular, basis = np.linalg.svd(productive,full_matrices=False)
basis = basis[singular>singular[0]*1e-6]
coefficients = (u@weight).numpy()[columns].astype(np.float64)
projected = coefficients-basis.T@(basis@coefficients)
retained = float(np.linalg.norm(projected)/np.linalg.norm(coefficients))
if a.match_protected_norm:
if retained < 1e-8:
raise ValueError('Protected coefficient vector vanishes')
projected = projected/retained
vector = torch.zeros(weight.shape[1],dtype=torch.float32)
vector[columns] = torch.from_numpy(projected.astype(np.float32))
candidate = weight-a.alpha*torch.outer(u,vector)
protection = {'constraint':'coefficient vector orthogonal to original productive MLP means on selected columns',
'retained_coefficient_norm_fraction_before_matching':retained,
'norm_matched':a.match_protected_norm,
'max_constraint_residual':float(np.max(np.abs(productive@projected)))}
record = save_edit(ROOT/'models/original',directory,name,candidate,a.hypothesis)
record.update({'layer':layer,'alpha':a.alpha,'control':a.control,'columns':columns.tolist(),'coefficient_protection':protection})
records.append(record)
save_json(directory/'candidate.json',{'origin':'independent_sol','analysis':a.analysis,'hypothesis':a.hypothesis,
'analysis_sha256':hashlib.sha256((inside(a.analysis)/'directions_and_writers.npz').read_bytes()).hexdigest(),
'control':a.control,'records':records,'status':'unconfirmed'})
export(ROOT/'models/original',directory,model)
print(json.dumps({'candidate':a.name,'model':str(model),'changed_entries':sum(r['changed_entries'] for r in records)},ensure_ascii=False))
|