File size: 2,376 Bytes
3fd1a35 | 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 | #!/usr/bin/env python3
"""Re-export only calibrated linear assets; reuse baseline raw tensors/islands."""
import argparse
import copy
import json
from pathlib import Path
import time
import torch
from quantize_model import AssetWriter, TensorSource, source_linear_names
p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--source',type=Path,required=True)
p.add_argument('--baseline-manifest',type=Path,required=True)
p.add_argument('--scales',type=Path,required=True)
p.add_argument('--output',type=Path,required=True)
a=p.parse_args()
torch.set_num_threads(4)
baseline=a.baseline_manifest.resolve()
manifest=json.loads(baseline.read_text())
assert manifest.get('flags')==1, 'baseline must be a complete model'
writer=AssetWriter(a.output.resolve(),a.scales)
source=TensorSource(a.source.resolve())
start=time.monotonic()
count=0
for item in manifest['tensors']:
if item['role']!='linear_weight':
continue
name=item['name'].removesuffix('.weight')
tensors=[source.tensor(n) for n in source_linear_names(name)]
n=sum(t.shape[0] for t in tensors)
writer.linear(name,tensors,layer=item.get('layer',0xFFFFFFFF),
expert=item.get('expert',0xFFFFFFFF),pad_n=n!=item['dims'][1],
required_k_splits=item['flags'])
count+=1
if count%200==0:
print(json.dumps({'exported_linears':count,'seconds':time.monotonic()-start}),flush=True)
replacement={t['name']:t for t in writer.entries}
result=copy.deepcopy(manifest)
result['tensors']=[]
for item in manifest['tensors']:
if item['name'] in replacement:
result['tensors'].append(replacement.pop(item['name']))
else:
original=copy.deepcopy(item)
for key in ('path','aux_path'):
if original.get(key):
original[key]=str((baseline.parent/original[key]).resolve())
assert Path(original[key]).is_file()
result['tensors'].append(original)
assert not replacement, 'candidate unexpectedly added tensors'
result['calibration']={'method':'activation_diagonal_weighted_scale_search_v1',
'scales_sha256':writer.calibration_sha256,
'preserve_baseline_k_partitions':True}
(a.output/'package-manifest.json').write_text(json.dumps(result,indent=2)+'\n')
print('PASS: complete candidate manifest; baseline assets unchanged',flush=True)
|