| |
| """Activation-diagonal weighted W4 scale search; unchanged per-row INT4 layout. |
| |
| Calibration and held-out fixtures must be different. This is not GPTQ/AWQ: |
| no full Hessian compensation, channel rescaling, or optimizer training is used. |
| """ |
| import argparse |
| import hashlib |
| import json |
| import time |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import save_file |
| from transformers import AutoTokenizer |
| from quant_loss_reference import load_model, evaluate |
|
|
|
|
| def safe_rows(codes, splits): |
| """Exact worst-case signed INT4 products, matching SplitAligned(K, splits, 32).""" |
| blocks = codes.shape[1] // 32 |
| if codes.shape[1] % 32 or not 1 <= splits <= blocks: |
| raise ValueError('invalid aligned K partition') |
| size, extra = divmod(blocks, splits) |
| begin = 0 |
| safe = torch.ones(codes.shape[0], dtype=torch.bool, device=codes.device) |
| for part in range(splits): |
| end = begin + (size + (part < extra)) * 32 |
| q = codes[:,begin:end].float() |
| positive = q.clamp_min(0).sum(1) |
| negative = (-q).clamp_min(0).sum(1) |
| safe &= (7*positive+8*negative <= 32767) & (8*positive+7*negative <= 32768) |
| begin = end |
| return safe |
|
|
|
|
| def choose_scales(weight, diagonal, rows=256, k_splits=None): |
| """Search clipping and least-squares scales; retain baseline if no improvement.""" |
| selected = [] |
| before, after = 0.0, 0.0 |
| for offset in range(0, weight.shape[0], rows): |
| w = weight[offset:offset+rows].float() |
| base = torch.maximum(w.amax(1).clamp_min(0)/7, (-w.amin(1)).clamp_min(0)/8) |
| base = torch.where(base > 0, base, torch.ones_like(base)) |
| def error(scale): |
| q = torch.round(w/scale[:, None]).clamp(-8, 7) |
| return ((w-q*scale[:, None]).square()*diagonal).sum(1), q |
| best_error, baseline_codes = error(base) |
| if k_splits is not None and not safe_rows(baseline_codes,k_splits).all(): |
| raise ValueError('baseline scale is unsafe for required K partitions') |
| before += best_error.double().sum().item() |
| best = base.clone() |
| |
| for fraction in (1.0, .95, .9, .85, .8, .75, .7, .65, .6, .55, .5): |
| scale = base * fraction |
| for iteration in range(3): |
| loss, q = error(scale) |
| improve = loss < best_error |
| if k_splits is not None: |
| improve &= safe_rows(q,k_splits) |
| best = torch.where(improve, scale, best) |
| best_error = torch.where(improve, loss, best_error) |
| if iteration < 2: |
| scale = ((w*q*diagonal).sum(1) / |
| (q.square()*diagonal).sum(1).clamp_min(1e-30)).clamp_min(1e-20) |
| after += best_error.double().sum().item() |
| selected.append(best) |
| return torch.cat(selected), before, after |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument('--source', type=Path, required=True) |
| p.add_argument('--corpus', type=Path, required=True) |
| p.add_argument('--suite', type=Path, required=True) |
| p.add_argument('--output', type=Path, required=True) |
| p.add_argument('--max-tokens', type=int, default=512) |
| p.add_argument('--reuse-activations', action='store_true') |
| p.add_argument('--baseline-manifest', type=Path, help='Preserve baseline INT16-safe K partitions') |
| a = p.parse_args() |
| torch.set_num_threads(4) |
| torch.backends.cuda.matmul.allow_tf32 = False |
| a.output.mkdir(parents=True, exist_ok=True) |
| corpus = json.loads(a.corpus.read_text()) |
| suite = json.loads(a.suite.read_text()) |
| |
| for c in corpus: |
| for e in suite['cases']: |
| assert c['prompt'] not in e['prompt_text'] and c['answer'] not in e['target_text'] |
| model = load_model(a.source) |
| modules = {name: m for name, m in model.named_modules() if isinstance(m, torch.nn.Linear)} |
| partitions = {} |
| if a.baseline_manifest: |
| from quantize_model import source_linear_names |
| for item in json.loads(a.baseline_manifest.read_text())['tensors']: |
| if item['role'] == 'linear_weight': |
| for n in source_linear_names(item['name'].removesuffix('.weight')): |
| partitions[n.removesuffix('.weight')] = item['flags'] |
| assert set(partitions) == set(modules) |
| stats_path = a.output/'activations.pt' |
| if a.reuse_activations: |
| stats = torch.load(stats_path, map_location='cuda', weights_only=True) |
| assert stats['corpus_sha256'] == hashlib.sha256(a.corpus.read_bytes()).hexdigest() |
| else: |
| sums, counts = {}, {} |
| def hook(name): |
| def accumulate(module, inputs): |
| x = inputs[0].detach().reshape(-1, inputs[0].shape[-1]).float() |
| if not x.shape[0]: |
| return |
| value = x.square().sum(0) |
| if name in sums: |
| sums[name].add_(value) |
| else: |
| sums[name] = value |
| counts[name] = counts.get(name, 0) + x.shape[0] |
| return accumulate |
| hooks = [m.register_forward_pre_hook(hook(n)) for n,m in modules.items()] |
| tokenizer = AutoTokenizer.from_pretrained(a.source, trust_remote_code=True, local_files_only=True) |
| lengths = [] |
| with torch.inference_mode(): |
| for i, c in enumerate(corpus): |
| for thinking in (False, True): |
| prompt = tokenizer.apply_chat_template([{'role':'user','content':c['prompt']}], |
| tokenize=False, add_generation_prompt=True, enable_thinking=thinking) |
| target = (c['reasoning']+'</think>' if thinking else '') + c['answer'] + tokenizer.eos_token |
| ids = tokenizer.encode(prompt+target, add_special_tokens=False)[:a.max_tokens] |
| output = model(input_ids=torch.tensor([ids],device='cuda'), use_cache=False) |
| del output |
| lengths.append(len(ids)) |
| print(json.dumps({'calibration_cases':i+1, 'total_tokens':sum(lengths), |
| 'matrices_observed':len(sums)}), flush=True) |
| for h in hooks: |
| h.remove() |
| stats = {'sum_squares':{k:v.cpu() for k,v in sums.items()},'counts':counts, |
| 'lengths':lengths, 'corpus_sha256':hashlib.sha256(a.corpus.read_bytes()).hexdigest()} |
| torch.save(stats, stats_path) |
| del sums |
| scales, records = {}, [] |
| start = time.monotonic() |
| with torch.inference_mode(): |
| for i,(name,m) in enumerate(modules.items()): |
| count = stats['counts'].get(name,0) |
| diag = stats['sum_squares'].get(name) |
| if diag is None or count == 0: |
| |
| w = m.weight.float() |
| scale = torch.maximum(w.amax(1).clamp_min(0)/7,(-w.amin(1)).clamp_min(0)/8) |
| scale = torch.where(scale>0,scale,torch.ones_like(scale)) |
| old = new = None |
| else: |
| diag = diag.to('cuda') / count |
| mean = diag.mean().clamp_min(1e-20) |
| |
| strength = .05 + .95 * max(0., 1.-count/64.) |
| diag = (1-strength)*diag/mean + strength |
| scale, old, new = choose_scales(m.weight,diag,k_splits=partitions.get(name)) |
| scales[name+'.weight'] = scale.cpu().contiguous() |
| records.append({'name':name,'observed_tokens':count,'weighted_error_before':old, |
| 'weighted_error_after':new}) |
| for row in range(0,m.weight.shape[0],1024): |
| w = m.weight[row:row+1024].float() |
| s = scale[row:row+1024,None] |
| m.weight[row:row+1024].copy_((torch.round(w/s).clamp(-8,7)*s).to(m.weight.dtype)) |
| if (i+1)%200 == 0: |
| print(json.dumps({'calibrated_matrices':i+1,'seconds':time.monotonic()-start}),flush=True) |
| save_file(scales,str(a.output/'scales.safetensors')) |
| report = {'method':'activation_diagonal_weighted_scale_search_v1','corpus_sha256':stats['corpus_sha256'], |
| 'calibration_sequences':len(stats['lengths']),'calibration_tokens':sum(stats['lengths']), |
| 'matrices':records,'seconds':time.monotonic()-start,'max_tokens':a.max_tokens, |
| 'baseline_manifest_sha256':hashlib.sha256(a.baseline_manifest.read_bytes()).hexdigest() if a.baseline_manifest else None, |
| 'preserve_baseline_k_partitions':bool(a.baseline_manifest)} |
| (a.output/'calibration.json').write_text(json.dumps(report,indent=2)+'\n') |
| evaluate(model,suite['cases'],a.output,'calibrated_bf16') |
| print('PASS: calibrated scales and held-out logits exported',flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
