source/quantize_nvfp4.py

"""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()