| 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.' |
|
|
