File size: 5,249 Bytes
64f6172 | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import argparse
import json
import re
import time
from pathlib import Path
import torch
from safetensors import safe_open
from transformers import AutoTokenizer
from packed_int4 import PackedLinear, load_int4
parser = argparse.ArgumentParser()
parser.add_argument('--folder', required=True)
parser.add_argument('--source', required=True)
parser.add_argument('--reload-only', action='store_true')
args = parser.parse_args()
folder, source = Path(args.folder), Path(args.source)
torch.set_num_threads(4)
torch.manual_seed(20261020)
model = load_int4(folder)
tokenizer = AutoTokenizer.from_pretrained(folder, trust_remote_code=True, local_files_only=True)
probe = tokenizer('The sum of 17 and 23 is', return_tensors='pt').to('cuda')
with torch.inference_mode():
actual = model(**probe, use_cache=False, logits_to_keep=1).logits.float().cpu()
assert torch.isfinite(actual).all()
if args.reload_only:
expected = torch.load(folder/'reference_probe.pt', weights_only=True)
error = (actual-expected).abs().max().item()
assert torch.equal(actual, expected), error
validation = json.loads((folder/'validation.json').read_text())
validation['fresh_process_reload_max_logit_error'] = error
validation['complete'] = True
(folder/'validation.json').write_text(json.dumps(validation, indent=2, ensure_ascii=False)+'\n')
print(json.dumps({'stage': 'reload_verified', 'max_logit_error': error}), flush=True)
raise SystemExit(0)
torch.save(actual, folder/'reference_probe.pt')
manifest = json.loads((folder/'edgerock_quantization.json').read_text())
decoded = len(manifest['modules'])
assert sum(isinstance(module, PackedLinear) for module in model.modules()) == decoded
print(json.dumps({'stage': 'packed_model_loaded', 'quantized_modules': decoded}), flush=True)
index = json.loads((source/'model.safetensors.index.json').read_text())['weight_map']
protected = 0
for shard in sorted(folder.glob('model*.safetensors')):
with safe_open(shard, framework='pt', device='cpu') as checkpoint:
ordinary = [name for name in checkpoint.keys() if name in index and
name.removesuffix('.weight') not in manifest['modules']]
for original_name in sorted({index[name] for name in ordinary}):
with safe_open(source/original_name, framework='pt', device='cpu') as original:
for name in ordinary:
if index[name] == original_name:
candidate = checkpoint.get_tensor(name)
reference = original.get_tensor(name)
assert candidate.dtype == reference.dtype and torch.equal(candidate, reference), name
protected += 1
assert protected > 0
cases = [
('multiply', 'Calculate 17 * 23. Give only the answer.', r'\b391\b'),
('equation', 'Solve x: 3*x + 7 = 31. Give only the answer.', r'\b8\b'),
('memory_update', 'The active project code is COBALT-7391.\n'
+ 'The archive contains ordinary observations about the weather and landscape. '*120
+ '\nUpdate: the active project code is now AMBER-5286. The earlier code is obsolete.\n'
+ 'Additional notes record the temperature, wind, and clouds each afternoon. '*160
+ '\nWhat is the active project code now? Give only that code.', r'AMBER-5286'),
]
rows = []
for name, prompt, expected in cases:
text = tokenizer.apply_chat_template([{'role': 'user', 'content': prompt}],
tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors='pt', add_special_tokens=False).to('cuda')
started = time.perf_counter()
torch.cuda.reset_peak_memory_stats()
with torch.inference_mode():
output = model.generate(**inputs, max_new_tokens=2048, do_sample=False, use_cache=True)
torch.cuda.synchronize()
answer = tokenizer.decode(output[0, inputs.input_ids.shape[1]:], skip_special_tokens=True)
final = answer.split('</think>')[-1].strip()
row = {
'name': name, 'prompt': prompt, 'prompt_tokens': inputs.input_ids.shape[1],
'output_tokens': output.shape[1]-inputs.input_ids.shape[1],
'output': answer, 'final_answer': final,
'matched_expected': bool(re.search(expected, final)),
'finished_by_eos': int(output[0, -1]) in (248046, 248044),
'seconds': time.perf_counter()-started, 'peak_allocated_gib': torch.cuda.max_memory_allocated()/2**30,
}
rows.append(row)
print(json.dumps({'stage': 'generation', **row}, ensure_ascii=False), flush=True)
validation = {
'complete': False, 'format': manifest['method'].upper()+' W4A16', 'unquantized_tensors_exact': protected,
'packed_linear_modules': decoded,
'generations': rows, 'all_generation_checks_passed': all(row['matched_expected'] and row['finished_by_eos'] for row in rows),
'runtime': 'GPU packed INT4 linear weights, fused Triton BF16-activation matmul',
'hardware': torch.cuda.get_device_name(0), 'max_new_tokens': 2048, 'do_sample': False,
}
(folder/'validation.json').write_text(json.dumps(validation, indent=2, ensure_ascii=False)+'\n')
assert validation['all_generation_checks_passed'], 'Generation checks did not pass; publication has not been started.'
|