"""Calibrated NVFP4 W4A4 + BF16 rank-128 correction for Image 2.1 Turbo; no evaluation prompts used. Per layer: five SmoothQuant exponents x three residual-weight rounding methods (plain FlashInfer quantizer; activation-weighted 15-point block-scale search; GPTQ error feedback with the same block-scale search, using full-token second moments). Low-rank correction: rank-128 randomized SVD of the smoothed weight (seed 1234, q160, niter 3). Every candidate is scored with the actual native SM120 W4A4 kernel and dynamic tensor-wide activation range on 512 fit rows; the selected candidate is reported on 512 disjoint diagnostic rows. """ import json, math, os, shutil, statistics from pathlib import Path import torch from safetensors import safe_open from safetensors.torch import load_file, save_file from flashinfer import nvfp4_quantize from flashinfer.gemm import nvfp4_quantize_smooth, mm_nvfp4_svdquant from common import ROOT, BASE, BASE_REVISION, CALIBRATION from dynamic_scale import scale from weight_quant import quantize_weight from hessian_store import hessian from gptq import gptq_nvfp4, pack_nvfp4 TRANSFORMER = BASE / 'transformer' INDEX = json.loads((TRANSFORMER / 'diffusion_pytorch_model.safetensors.index.json').read_text())['weight_map'] OUT_ROOT = ROOT / 'nvfp4' / 'release' OUT = OUT_ROOT / 'transformer' CACHE = ROOT / 'nvfp4' / 'calibrated-layers' RANK = 128 NOTICE = 'Built with Qwen. ProCreations modified Qwen-Image-2.1-Turbo transformer weights with calibrated NVFP4 quantization, 2026-10-09.' def weight(name): with safe_open(TRANSFORMER / INDEX[name + '.weight'], framework='pt', device='cuda') as f: return f.get_tensor(name + '.weight') def split(x): g = torch.arange(len(x) // 4, device=x.device) return x[g * 4 + g % 4].contiguous(), x[g * 4 + (g + 1) % 4].contiguous() def error(y, ref): return ((y.float() - ref.float()).square().mean() / ref.float().square().mean()).sqrt().item() def prepare(w, x, H, smoothing, method): amax = x.float().abs().amax(0).clamp_min(1e-5); wmax = w.float().abs().amax(0).clamp_min(1e-5) smooth = torch.ones_like(amax) if smoothing == 0 else (amax ** smoothing / wmax ** (1 - smoothing)).clamp(.01, 100) smooth /= smooth.log().mean().exp() pre = (1 / smooth).to(torch.bfloat16) p32 = pre.float() ws = w.float() / p32[None, :] torch.manual_seed(1234) u, s, v = torch.svd_lowrank(ws, q=RANK + 32, niter=3) up, down = u[:, :RANK] * s[:RANK], v[:, :RANK] residual = ws - up @ down.T gx = (2688 / ((x.float() * p32).abs().max() * 1.2)).reshape(1) gw = (2688 / residual.abs().max()).reshape(1) hs = H.diagonal() * p32 * p32 h = hs + .01 * hs.mean() if method == 'plain': wq, wsc = nvfp4_quantize(residual.to(torch.bfloat16), gw, backend='cute-dsl') wq, wsc = wq.view(torch.uint8), wsc.view(torch.uint8).reshape(-1) elif method == 'block-scale-search': wq, wsc = quantize_weight(residual, gw, h) else: Hs = H * p32[:, None] * p32[None, :] codes, scales = gptq_nvfp4(residual, Hs, gw.item(), h) wq, wsc = pack_nvfp4(codes, scales) alpha = (1 / (gx * gw)).reshape(1) return {'weight': wq.view(torch.uint8), 'sf': wsc.view(torch.uint8).reshape(-1), 'pre': pre, 'gx': gx, 'alpha': alpha, 'down': (p32[:, None] * down).to(torch.bfloat16).contiguous(), 'up': (up / alpha).to(torch.bfloat16).contiguous()} def run(x, p): gx, alpha, up = scale(x, p['pre'], p['gx'], p['alpha'], p['up']) q, sf = nvfp4_quantize_smooth(x, p['pre'], gx, backend='cute-dsl') return mm_nvfp4_svdquant(q, p['weight'], sf, p['sf'], alpha, x @ p['down'], up, backend='cute-dsl') @torch.inference_mode() def calibrate_layer(name): w = weight(name) x = load_file(str(CALIBRATION / f'{name}.safetensors'))['activations'].cuda() fit, check = split(x) H, tokens = hessian(name) rf, rc = fit @ w.T, check @ w.T best, candidates = None, [] for smoothing in [0., .25, .5, .75, 1.]: for method in ['plain', 'block-scale-search', 'gptq']: p = prepare(w, fit, H, smoothing, method) err = error(run(fit, p), rf) item = {'smoothing_alpha': smoothing, 'weight_rounding': method, 'fit_nrmse': err} candidates.append(item) if best is None or err < best[0]: best = (err, {k: v.clone() for k, v in p.items()}, item) err, p, item = best check_error = error(run(check, p), rc) same = {c['weight_rounding']: c['fit_nrmse'] for c in candidates if c['smoothing_alpha'] == item['smoothing_alpha']} result = {**item, 'diagnostic_nrmse': check_error, 'rank': RANK, 'shape': list(w.shape), 'sample_rows': len(x), 'fit_rows': len(fit), 'diagnostic_rows': len(check), 'hessian_tokens': tokens, 'same_alpha_fit_nrmse_by_rounding': same, 'candidates': candidates} assert check_error < .2 and all(torch.isfinite(v.float()).all() for v in p.values()), (name, check_error) return result, p @torch.inference_mode() def main(): CACHE.mkdir(parents=True, exist_ok=True); OUT.mkdir(parents=True, exist_ok=True) names = json.loads((CALIBRATION / 'manifest.json').read_text())['modules'] import sys shard = tuple(map(int, sys.argv[2].split('/'))) if len(sys.argv) == 3 and sys.argv[1] == '--shard' else None if shard: for i, name in enumerate(names): record, saved = CACHE / (name + '.json'), CACHE / (name + '.safetensors') if i % shard[1] != shard[0] or (record.exists() and saved.exists()): continue result, p = calibrate_layer(name) save_file({k: v.cpu().contiguous() for k, v in p.items()}, str(saved)) record.write_text(json.dumps(result, indent=2)) print(json.dumps({'event': 'layer', 'shard': shard[0], 'done': i + 1, 'name': name, 'nrmse': result['diagnostic_nrmse'], 'rounding': result['weight_rounding'], 'alpha': result['smoothing_alpha']}), flush=True) return results, specs = {}, {} for i, name in enumerate(names): record, saved = CACHE / (name + '.json'), CACHE / (name + '.safetensors') if record.exists() and saved.exists(): result = json.loads(record.read_text()) else: result, p = calibrate_layer(name) save_file({k: v.cpu().contiguous() for k, v in p.items()}, str(saved)) record.write_text(json.dumps(result, indent=2)) results[name] = result specs[name] = {'shape': result['shape'], 'rank': RANK, 'smoothing_alpha': result['smoothing_alpha'], 'weight_rounding': result['weight_rounding']} print(json.dumps({'event': 'quantize', 'done': i + 1, 'total': len(names), 'name': name, 'nrmse': result['diagnostic_nrmse'], 'fit': result['fit_nrmse'], 'by_rounding': result['same_alpha_fit_nrmse_by_rounding'], 'alpha': result['smoothing_alpha'], 'rounding': result['weight_rounding']}), flush=True) state = {} for shard in sorted(set(INDEX.values())): for name, v in load_file(str(TRANSFORMER / shard)).items(): if name.removesuffix('.weight') not in specs: state[name] = v for name in names: for k, v in load_file(str(CACHE / (name + '.safetensors'))).items(): state[name + '.' + k] = v shutil.copy(TRANSFORMER / 'config.json', OUT / 'config.json') cfg = {'format': 'procreations-calibrated-nvfp4-svd-v1', 'base_model': 'Qwen/Qwen-Image-2.1-Turbo', 'base_revision': BASE_REVISION, 'quantized_modules': specs, 'weight_dtype': 'NVFP4 E2M1 packed uint8', 'activation_dtype': 'NVFP4 E2M1', 'scale_format': 'E4M3 per 16 elements, 128x4 swizzle, FP32 global multiplier', 'low_rank_correction': 'BF16 rank128, native fused FP32 accumulator epilogue', 'activation_scale': 'dynamic tensor-wide activation amax at every call; dynamic per16 block scales; no static activation clipping', 'weight_rounding': 'per layer: plain, activation-weighted block-scale search, or GPTQ error feedback with full-token second moments', 'flashinfer': 'flashinfer-python 0.7.1', 'kernel': 'SM120 CuTe DSL native block-scaled FP4 GEMM with fused BF16 rank correction', 'modified_notice': NOTICE} (OUT / 'quantization_config.json').write_text(json.dumps(cfg, indent=2)) for old in OUT.glob('model-*.safetensors'): old.unlink() shards, s, size = [], {}, 0 for n, v in state.items(): nb = v.numel() * v.element_size() if s and size + nb > 3_000_000_000: shards.append(s); s, size = {}, 0 s[n] = v.contiguous(); size += nb if s: shards.append(s) for i, s in enumerate(shards): save_file(s, str(OUT / f'model-{i + 1:05d}-of-{len(shards):05d}.safetensors'), metadata={'format': 'pt', 'notice': NOTICE}) total = sum(v.numel() * v.element_size() for v in state.values()) vals = list(results.values()) summary = {'modules': len(names), 'transformer_bytes': total, 'mean_diagnostic_nrmse': statistics.mean(r['diagnostic_nrmse'] for r in vals), 'max_diagnostic_nrmse': max(r['diagnostic_nrmse'] for r in vals), 'selected_rounding': {m: sum(r['weight_rounding'] == m for r in vals) for m in ['plain', 'block-scale-search', 'gptq']}, 'mean_fit_nrmse_by_rounding_at_selected_alpha': {m: statistics.mean(r['same_alpha_fit_nrmse_by_rounding'][m] for r in vals) for m in ['plain', 'block-scale-search', 'gptq']}, 'protocol': '64 BF16 Turbo trajectories (8 saved-schedule steps); 512 fit and 512 disjoint diagnostic rows per layer. ' 'Full-token second moments from the same trajectories for GPTQ/scale weighting. Native W4A4 kernels with ' 'dynamic activation amax used for all selection/diagnostic errors. No held-out evaluation prompts used.'} report = OUT_ROOT / 'reports'; report.mkdir(parents=True, exist_ok=True) (report / 'calibration_search.json').write_text(json.dumps(results, indent=2)) (report / 'quantization_summary.json').write_text(json.dumps(summary, indent=2)) shutil.copy(CALIBRATION / 'manifest.json', report / 'calibration_manifest.json') print('QUANTIZATION_COMPLETE', json.dumps(summary), flush=True) if __name__ == '__main__': main()