| |
| """Losslessly repack official weights using the existing fused tensor names. |
| |
| INT4 codes and BF16 group scales remain unchanged. Execution precision is chosen |
| by the experimental runtime, not baked into this source-weight package. |
| """ |
| import argparse |
| import hashlib |
| import json |
| import mmap |
| from pathlib import Path |
| import struct |
| import zlib |
|
|
| import torch |
| from safetensors import safe_open |
| from pack_model import HEADER, ENTRY, align, pad_to |
| from quantize_model import source_linear_names |
| from official_int4_reference import unpack4, REVISION |
|
|
| p=argparse.ArgumentParser(description=__doc__) |
| p.add_argument('--template',type=Path,required=True) |
| p.add_argument('--official',type=Path,required=True) |
| p.add_argument('--output',type=Path,required=True) |
| a=p.parse_args() |
| if a.output.exists(): raise FileExistsError(a.output) |
| torch.set_num_threads(4) |
| index=json.loads((a.official/'model.safetensors.index.json').read_text())['weight_map'] |
| handles={s:safe_open(a.official/s,framework='pt') for s in set(index.values())} |
| def get(name): return handles[index[name]].get_tensor(name) |
| with a.template.open('rb') as src, mmap.mmap(src.fileno(),0,access=mmap.ACCESS_READ) as data: |
| header=list(HEADER.unpack_from(data)); count=header[4]; table=header[7]; strings=header[8]; payload=header[10] |
| entries=[] |
| for i in range(count): |
| e=list(ENTRY.unpack_from(data,table+i*ENTRY.size)) |
| name=data[strings+e[0]:strings+e[0]+e[1]].decode() |
| entries.append((name,e)) |
| a.output.parent.mkdir(parents=True,exist_ok=True) |
| with a.output.open('w+b') as out: |
| out.write(data[:payload]); written=[]; linear_count=0 |
| for name,e in entries: |
| raw=None |
| if e[4]==3: |
| base=name.removesuffix('.weight'); names=source_linear_names(base) |
| k,n=e[8:10] |
| if '.mlp.experts.' in base: |
| codes=[] |
| for key in names: |
| shape=tuple(get(key+'_shape').tolist()) |
| codes.append(unpack4(get(key+'_packed'),shape)) |
| q=torch.cat(codes).t().contiguous().reshape(-1).numpy().view('uint8') & 15 |
| raw=(q[::2] | (q[1::2]<<4)).tobytes() |
| e[2]=5;e[5]=4;e[6]=2;e[7]=32 |
| assert len(raw)==k*n//2 |
| else: |
| weight=torch.cat([get(key) for key in names]) |
| assert weight.dtype==torch.bfloat16 and weight.shape[1]==k and weight.shape[0]<=n |
| if weight.shape[0]<n:weight=torch.nn.functional.pad(weight,(0,0,0,n-weight.shape[0])) |
| raw=weight.contiguous().view(torch.uint16).numpy().tobytes() |
| e[2]=1;e[5]=0;e[6]=0;e[7]=0 |
| linear_count+=1 |
| if linear_count%256==0:print(f'packed {linear_count} linears',flush=True) |
| elif name.endswith('.scales') and '.mlp.experts.' in name: |
| base=name.removesuffix('.scales') |
| scales=torch.cat([get(key+'_scale') for key in source_linear_names(base)]) |
| raw=scales.contiguous().view(torch.uint16).numpy().tobytes() |
| e[2]=1;e[3]=2;e[5]=4;e[6]=0;e[7]=32;e[8:12]=[scales.shape[0],scales.shape[1],0,0] |
| if raw is None:raw=data[e[12]:e[12]+e[13]] |
| pad_to(out,align(out.tell()));e[12]=out.tell();e[13]=len(raw);out.write(raw) |
| e[14]=e[15]=0;e[22]=hashlib.sha256(raw).digest();written.append(ENTRY.pack(*e)) |
| pad_to(out,align(out.tell()));header[11]=out.tell()-payload;header[12]=out.tell();header[13]=bytes.fromhex(REVISION) |
| header[6]|=0x100;header[14]=0;header[14]=zlib.crc32(HEADER.pack(*header)) & 0xffffffff |
| out.seek(0);out.write(HEADER.pack(*header));out.seek(table) |
| for entry in written:out.write(entry) |
| print(json.dumps({'output':str(a.output),'bytes':a.output.stat().st_size,'linears':linear_count,'source_revision':REVISION,'quantization':'unchanged official INT4 group32 and high precision weights'}),flush=True) |
|
|
