| import argparse |
| import hashlib |
| import json |
| import math |
| from pathlib import Path |
| import shutil |
| import struct |
| import time |
| import numpy as np |
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
| from common import ROOT, inside, save_json, append, permit_compute |
|
|
| torch.set_grad_enabled(False) |
| FROZEN = ('embed_tokens', 'lm_head', 'model.language_model.norm.weight', 'model.norm.weight') |
|
|
|
|
| def headers(model): |
| result = {} |
| for file in sorted(inside(model).glob('*.safetensors')): |
| with file.open('rb') as f: |
| length = struct.unpack('<Q', f.read(8))[0] |
| head = json.loads(f.read(length)) |
| for name, info in head.items(): |
| if name != '__metadata__': |
| if name in result: |
| raise ValueError(f'Duplicate tensor: {name}') |
| result[name] = dict(info, file=str(file)) |
| return result |
|
|
|
|
| def load(model, name): |
| info = headers(model)[name] |
| with safe_open(info['file'], framework='pt', device='cpu') as f: |
| return f.get_tensor(name) |
|
|
|
|
| def validate(name, original, value): |
| if not name.startswith('model.language_model.layers.') or any(x in name for x in FROZEN): |
| raise ValueError('Разрешены только внутренние слои языковой модели') |
| if value.shape != original.shape or not torch.isfinite(value).all(): |
| raise ValueError('Неверная форма или нечисловые значения') |
| w, v = original.float(), value.float() |
| if torch.equal(w, v): |
| raise ValueError('Правка не меняет веса') |
| ratio = (w * v).sum() / w.square().sum().clamp_min(1e-30) |
| if (v - ratio * w).norm() <= 1e-6 * v.norm().clamp_min(1e-30): |
| raise ValueError('Масштабирование всего тензора запрещено') |
| cast = value.to(original.dtype).contiguous() |
| if not torch.isfinite(cast).all() or torch.equal(cast, original): |
| raise ValueError('Правка исчезла или стала нечисловой после приведения dtype') |
| return cast |
|
|
|
|
| def digest(tensor): |
| return hashlib.sha256(tensor.contiguous().view(torch.uint8).numpy().tobytes()).hexdigest() |
|
|
|
|
| def save_edit(model, directory, name, value, hypothesis): |
| permit_compute() |
| directory = inside(directory) |
| if not directory.is_relative_to(ROOT / 'edits'): |
| raise ValueError('Правки сохраняются только в edits/') |
| if not hypothesis.strip(): |
| raise ValueError('Нужна техническая гипотеза') |
| original = load(model, name) |
| cast = validate(name, original, value) |
| directory.mkdir(parents=True, exist_ok=True) |
| file = directory / (name + '.safetensors') |
| save_file({name: cast}, str(file)) |
| record = {'time': time.time(), 'source_model': str(inside(model)), 'tensor': name, 'hypothesis': hypothesis, 'status': 'unconfirmed', |
| 'changed_entries': int((cast != original).sum()), 'tensor_entries': original.numel(), |
| 'base_sha256': digest(original), 'edited_sha256': digest(cast), |
| 'relative_frobenius': float((cast.float()-original.float()).norm()/original.float().norm())} |
| save_json(directory / (name + '.json'), record) |
| append(ROOT / 'logs/experiments.jsonl', record) |
| return record |
|
|
|
|
| def project_writers(w, direction, columns, alpha): |
| w = w.float() |
| u = torch.as_tensor(direction, dtype=torch.float32) |
| if not torch.isfinite(u).all() or u.norm() == 0: |
| raise ValueError('Некорректное направление') |
| u = u / u.norm() |
| coeff = u @ w |
| mask = torch.zeros_like(coeff) |
| mask[columns] = coeff[columns] |
| return w - float(alpha) * torch.outer(u, mask) |
|
|
|
|
| def lineage(model): |
| current, ancestors, inherited = inside(model), [], set() |
| while True: |
| if str(current) in ancestors: |
| raise ValueError('Цикл в происхождении модели') |
| ancestors.append(str(current)) |
| manifest = current / 'wsp_manifest.json' |
| if not manifest.exists(): |
| return current, ancestors, inherited |
| record = json.loads(manifest.read_text()) |
| inherited.update(record['tensors']) |
| current = inside(record['source']) |
|
|
|
|
| def weight_fingerprint(model): |
| result = [] |
| for path in sorted(inside(model).glob('*.safetensors')): |
| with path.open('rb') as stream: |
| result.append([path.name, hashlib.file_digest(stream, 'sha256').hexdigest()]) |
| return result |
|
|
|
|
| def export(model, edits, out): |
| permit_compute() |
| model, edits, out = inside(model), inside(edits), inside(out) |
| if out.exists(): |
| raise ValueError('Каталог назначения должен быть новым') |
| if not out.is_relative_to(ROOT / 'models') or out == ROOT / 'models/original': |
| raise ValueError('Недопустимый каталог модели') |
| metadata = headers(model) |
| changes = {} |
| for file in edits.glob('*.safetensors'): |
| with safe_open(str(file), framework='pt', device='cpu') as f: |
| for name in f.keys(): |
| if name in changes: |
| raise ValueError('Повторная правка тензора') |
| base = load(model, name) |
| changes[name] = validate(name, base, f.get_tensor(name)) |
| manifest = json.loads((edits / (name + '.json')).read_text()) |
| if manifest['base_sha256'] != digest(base) or manifest['edited_sha256'] != digest(changes[name]): |
| raise ValueError('Хеш правки не совпал') |
| if not changes: |
| raise ValueError('Нет правок для экспорта') |
| original, ancestors, inherited = lineage(model) |
| original_metadata = headers(original) |
| if metadata.keys() != original_metadata.keys() or any( |
| metadata[name]['shape'] != original_metadata[name]['shape'] or |
| metadata[name]['dtype'] != original_metadata[name]['dtype'] for name in metadata): |
| raise ValueError('Архитектура родителя отличается от исходной модели') |
| cumulative = {} |
| for name in sorted(inherited | changes.keys()): |
| value = changes[name] if name in changes else load(model, name) |
| cumulative[name] = int((value != load(original, name)).sum()) |
| parent_fingerprint = weight_fingerprint(model) |
| original_fingerprint = parent_fingerprint if original == model else weight_fingerprint(original) |
| staging = out.with_name(out.name + '.building') |
| if staging.exists(): |
| raise ValueError('Обнаружен незавершённый экспорт') |
| staging.mkdir(parents=True) |
| for file in model.iterdir(): |
| if file.is_file() and file.suffix != '.safetensors': |
| shutil.copyfile(file, staging / file.name) |
| for shard in sorted({x['file'] for x in metadata.values()}): |
| with safe_open(shard, framework='pt', device='cpu') as f: |
| tensors = {name: changes[name] if name in changes else f.get_tensor(name) for name in f.keys()} |
| save_file(tensors, str(staging / Path(shard).name), metadata={'format': 'pt'}) |
| del tensors |
| total = sum(math.prod(x['shape']) for x in metadata.values() if x['dtype'].startswith(('BF', 'F'))) |
| language_total = sum(math.prod(x['shape']) for name, x in metadata.items() |
| if name.startswith('model.language_model.') and x['dtype'].startswith(('BF', 'F'))) |
| changed = sum(int((value != load(model, name)).sum()) for name, value in changes.items()) |
| cumulative_changed = sum(cumulative.values()) |
| save_json(staging / 'wsp_manifest.json', {'schema_version': 2, |
| 'source': str(model), 'edits': str(edits), 'tensors': list(changes), |
| 'changed_entries': changed, 'stored_parameter_entries': total, |
| 'changed_percent': 100 * changed / total, |
| 'language_parameter_entries': language_total, |
| 'changed_percent_language': 100 * changed / language_total, |
| 'changed_entries_reference': 'immediate_parent', |
| 'parent_weight_fingerprint': parent_fingerprint, |
| 'original_source': str(original), 'ancestor_chain_parent_first': ancestors, |
| 'original_weight_fingerprint': original_fingerprint, |
| 'changed_entries_vs_original': cumulative_changed, |
| 'changed_percent_vs_original': 100 * cumulative_changed / total, |
| 'changed_percent_language_vs_original': 100 * cumulative_changed / language_total, |
| 'changed_entries_vs_original_by_tensor': cumulative, |
| 'status': 'unconfirmed'}) |
| staging.rename(out) |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('command', choices=['inventory', 'export']) |
| parser.add_argument('--model', default=str(ROOT / 'models/original')) |
| parser.add_argument('--edits') |
| parser.add_argument('--out') |
| args = parser.parse_args() |
| if args.command == 'inventory': |
| save_json(ROOT / 'reports/tensor_inventory.json', headers(args.model)) |
| else: |
| export(args.model, args.edits, args.out) |
|
|
