File size: 10,597 Bytes
ffc54ec | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | """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()
|