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