File size: 4,135 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 | import argparse
import base64
import hashlib
import json
from pathlib import Path
import shutil
import struct
import numpy as np
def file_hash(path):
with path.open('rb') as stream:
return hashlib.file_digest(stream, 'sha256').hexdigest()
def apply_patch(original, manifest_path, patch_path, output):
original, manifest_path, patch_path, output = map(Path, (original, manifest_path, patch_path, output))
if output.exists() or output.resolve() == original.resolve():
raise ValueError('Output must be a new file; original weights are never overwritten')
manifest = json.loads(manifest_path.read_text())
if file_hash(original) != manifest['original_sha256'] or file_hash(patch_path) != manifest['patch_sha256']:
raise ValueError('Original checkpoint or patch checksum differs')
with original.open('rb') as stream:
original_length = struct.unpack('<Q', stream.read(8))[0]
original_header = json.loads(stream.read(original_length))
target_header_bytes = base64.b64decode(manifest['target_header_base64'], validate=True)
target_header = json.loads(target_header_bytes)
if {k: v for k, v in original_header.items() if k != '__metadata__'} != {
k: v for k, v in target_header.items() if k != '__metadata__'}:
raise ValueError('Tensor layouts differ')
output.parent.mkdir(parents=True, exist_ok=True)
temporary = output.with_name(output.name + '.partial')
if temporary.exists():
raise ValueError('Partial output already exists; inspect it before retrying')
with original.open('rb') as source, temporary.open('xb') as target:
target.write(struct.pack('<Q', len(target_header_bytes)))
target.write(target_header_bytes)
source.seek(8 + original_length)
shutil.copyfileobj(source, target, length=8 * 1024 * 1024)
changed = 0
with np.load(patch_path, allow_pickle=False) as arrays:
for entry in manifest['tensors']:
name = entry['name']
if name not in {f'model.language_model.layers.{layer}.mlp.down_proj.weight' for layer in (15, 19)}:
raise ValueError('Unexpected release tensor')
info = target_header[name]
if info['dtype'] != 'BF16' or info['shape'] != entry['shape']:
raise ValueError('Tensor shape/dtype differs')
indices, bits = arrays[entry['indices_key']], arrays[entry['bits_key']]
size = int(np.prod(info['shape']))
if indices.ndim != 1 or indices.dtype != np.dtype('uint64') or bits.dtype != np.dtype('uint16'):
raise ValueError('Invalid sparse array format')
if bits.shape != indices.shape or len(indices) != entry['changed_bitpatterns']:
raise ValueError('Sparse array lengths differ')
if len(indices) and (indices[-1] >= size or np.any(indices[1:] <= indices[:-1])):
raise ValueError('Sparse indices must be sorted, unique, and in range')
values = np.memmap(temporary, mode='r+', dtype='<u2',
offset=8 + len(target_header_bytes) + info['data_offsets'][0], shape=(size,))
values[indices] = bits
values.flush()
del values
changed += len(indices)
if changed != manifest['changed_bitpatterns'] or file_hash(temporary) != manifest['target_sha256']:
raise ValueError('Reconstructed weight file differs from the evaluated candidate')
temporary.rename(output)
return {'passed': True, 'changed_bitpatterns': changed, 'target_sha256': manifest['target_sha256'],
'byte_identical_to_evaluated_weights': True, 'model_forward_executed': False}
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--original', required=True)
parser.add_argument('--manifest', required=True)
parser.add_argument('--patch', required=True)
parser.add_argument('--output', required=True)
args = parser.parse_args()
print(json.dumps(apply_patch(args.original, args.manifest, args.patch, args.output)))
if __name__ == '__main__':
main()
|