Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
42 lines
2.1 KiB
Python
42 lines
2.1 KiB
Python
"""Check all updated packed LM matrices against the full upstream weights."""
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
import numpy as np
|
|
import torch
|
|
from safetensors import safe_open
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / 'tools/community_models'))
|
|
import kitten_tts2_weights as converter
|
|
from test_updated_weights import dequantize
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('--model', type=Path, required=True)
|
|
parser.add_argument('--output', type=Path, required=True)
|
|
args = parser.parse_args()
|
|
torch.set_num_threads(8)
|
|
report = {}
|
|
with safe_open(args.model / 'lm/model.safetensors', framework='pt') as full:
|
|
with safe_open(args.model / 'lm/model-ternary.safetensors', framework='pt') as source:
|
|
names = json.loads(source.metadata()['packed'])
|
|
for name, info in names.items():
|
|
reference = full.get_tensor(name)
|
|
actual = converter.unpack_ternary(source.get_tensor(name+'::codes'), source.get_tensor(name+'::scale'), info['shape'])
|
|
if not torch.equal(actual, reference): raise RuntimeError(f'unpacking changed {name}')
|
|
for bits in (4, 8):
|
|
encoded = converter.exact_ternary_blocks(reference, bits)
|
|
if not np.array_equal(dequantize(encoded, bits).reshape(reference.shape), reference.float().numpy()):
|
|
raise RuntimeError(f'Q{bits}_0 packing changed {name}')
|
|
if len(names) != 196: raise RuntimeError('unexpected ternary body matrix count')
|
|
if set(names) != {name for name in full.keys() if converter.is_body_projection(name)}:
|
|
raise RuntimeError('full and packed body projection sets differ')
|
|
report = {'exact_body_matrices': len(names), 'q8_repacking_exact': True, 'q4_repacking_exact': True}
|
|
print('196 matrices match full BF16 exactly before and after Q8_0/Q4_0 packing', flush=True)
|
|
args.output.write_text(json.dumps(report,indent=2),encoding='utf-8')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|