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