| 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)) |
|
|
