File size: 9,345 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 | """Activation-aware calibrated FP8 (W8A8 E4M3) for Image 2.1 Turbo; no evaluation prompts used.
Per layer: five SmoothQuant exponents x {three weight-clipping factors with round-to-nearest,
GPTQ error-feedback rounding with full-token second moments}. Every candidate is scored with the
actual runtime path (Triton per-token activation quantization + native SM120 FP8 GEMM, FP32
accumulation) on 512 fit rows; the selected candidate is reported on 512 disjoint diagnostic rows.
Layers above 5% diagnostic output NRMSE would stay BF16.
"""
import json, math, shutil, statistics, sys
from pathlib import Path
import torch
from safetensors.torch import load_file, save_file
from diffusers import QwenImage21Transformer2DModel
from common import ROOT, BASE, BASE_REVISION, CALIBRATION, eligible_modules
from fp8_runtime import CalibratedFP8Linear, replace_module, quantize_activation
from hessian_store import hessian
from gptq import gptq_fp8
OUT_ROOT = ROOT / 'fp8' / 'release'
OUT = OUT_ROOT / 'transformer'
CACHE = ROOT / 'fp8' / 'calibrated-layers'
NOTICE = 'Built with Qwen. ProCreations modified Qwen-Image-2.1-Turbo transformer weights through calibrated FP8 quantization, 2026-10-09.'
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 fp8_forward(x, q, scales, smooth):
a, asc = quantize_activation(x, smooth)
return torch._scaled_mm(a, q.T, scale_a=asc, scale_b=scales.T, out_dtype=torch.bfloat16, use_fast_accum=False).float()
def nmse(y, ref):
return ((y - ref).square().mean() / ref.square().mean().clamp_min(1e-12)).item()
def bf16_weight(name):
from quantize_nvfp4 import weight
return weight(name)
@torch.inference_mode()
def calibrate_layer(name, layer_weight):
x = load_file(str(CALIBRATION / (name + '.safetensors')))['activations'].to('cuda')
fit, check = split(x)
H, tokens = hessian(name)
w = layer_weight.float()
amax = fit.float().abs().amax(0).clamp_min(1e-5)
wmax = w.abs().amax(0).clamp_min(1e-5)
ref = torch.nn.functional.linear(fit, layer_weight).float()
candidates, best = [], None
for alpha in [0., .25, .5, .75, 1.]:
smooth = torch.ones_like(amax) if alpha == 0 else (amax.pow(alpha) / wmax.pow(1 - alpha)).clamp(.01, 100.)
smooth = smooth / smooth.log().mean().exp()
ws = w * smooth[None, :]
options = []
for clip in [1., .995, .99]:
scales = (ws.abs().amax(1, keepdim=True) * clip / 448).clamp_min(1e-12)
options.append(({'rounding': 'nearest', 'weight_amax_clip_factor': clip}, scales,
(ws / scales).clamp(-448, 448).to(torch.float8_e4m3fn)))
scales = (ws.abs().amax(1, keepdim=True) / 448).clamp_min(1e-12)
Hz = H / (smooth[:, None] * smooth[None, :])
options.append(({'rounding': 'gptq', 'weight_amax_clip_factor': 1.}, scales, gptq_fp8(ws, Hz, scales)))
for item, scales, q in options:
err = nmse(fp8_forward(fit, q, scales, smooth), ref)
item = {'alpha': alpha, **item, 'fit_nmse': err}
candidates.append(item)
if best is None or err < best[0]:
best = (err, q.clone(), scales.clone(), smooth.clone(), item)
_, q, scales, smooth, item = best
ref_check = torch.nn.functional.linear(check, layer_weight).float()
diag = nmse(fp8_forward(check, q, scales, smooth), ref_check)
# Round-to-nearest reference at the same smoothing, for reporting the GPTQ contribution.
rtn = min((c for c in candidates if c['alpha'] == item['alpha'] and c['rounding'] == 'nearest'), key=lambda c: c['fit_nmse'])
result = {**item, 'diagnostic_nrmse': math.sqrt(diag), 'fit_nrmse': math.sqrt(item['fit_nmse']),
'nearest_fit_nrmse_same_alpha': math.sqrt(rtn['fit_nmse']), 'shape': list(q.shape),
'sample_rows': len(x), 'fit_rows': len(fit), 'diagnostic_rows': len(check),
'hessian_tokens': tokens, 'candidates': candidates}
return result, {'weight': q, 'weight_scale': scales.float(), 'smooth': smooth.float()}
def compute_shard(index, count):
names = json.loads((CALIBRATION / 'manifest.json').read_text())['modules']
for i, name in enumerate(names):
record, saved = CACHE / (name + '.json'), CACHE / (name + '.safetensors')
if i % count != index or (record.exists() and saved.exists()):
continue
result, state = calibrate_layer(name, bf16_weight(name))
save_file({k: v.cpu().contiguous() for k, v in state.items()}, str(saved))
record.write_text(json.dumps(result, indent=2))
print(json.dumps({'event': 'layer', 'shard': index, 'done': i + 1, 'layer': name, 'nrmse': result['diagnostic_nrmse'],
'alpha': result['alpha'], 'rounding': result['rounding']}), flush=True)
@torch.inference_mode()
def main():
CACHE.mkdir(parents=True, exist_ok=True); OUT.mkdir(parents=True, exist_ok=True)
if len(sys.argv) == 3 and sys.argv[1] == '--shard':
return compute_shard(*map(int, sys.argv[2].split('/')))
model = QwenImage21Transformer2DModel.from_pretrained(BASE, subfolder='transformer', torch_dtype=torch.bfloat16).to('cuda').eval()
modules = eligible_modules(model)
results, quantized, held = {}, {}, {}
for i, (name, layer) in enumerate(modules.items()):
record, saved = CACHE / (name + '.json'), CACHE / (name + '.safetensors')
if record.exists() and saved.exists():
result = json.loads(record.read_text()); state = {k: v.cuda() for k, v in load_file(str(saved)).items()}
else:
result, state = calibrate_layer(name, layer.weight)
save_file({k: v.cpu().contiguous() for k, v in state.items()}, str(saved))
record.write_text(json.dumps(result, indent=2))
preserve = not math.isfinite(result['diagnostic_nrmse']) or result['diagnostic_nrmse'] > .05
result['kept_bf16'] = preserve
results[name] = result
if preserve:
held[name] = result
else:
quantized[name] = {'shape': result['shape'], 'alpha': result['alpha'], 'rounding': result['rounding'],
'weight_amax_clip_factor': result['weight_amax_clip_factor']}
replace_module(model, name, CalibratedFP8Linear(state['weight'], state['weight_scale'], state['smooth']))
print(json.dumps({'event': 'quantize', 'done': i + 1, 'total': len(modules), 'layer': name,
'nrmse': result['diagnostic_nrmse'], 'alpha': result['alpha'], 'rounding': result['rounding'],
'nearest_fit': result['nearest_fit_nrmse_same_alpha'], 'fit': result['fit_nrmse']}), flush=True)
model.save_config(OUT)
config = {'format': 'procreations-calibrated-fp8-v1', 'base_model': 'Qwen/Qwen-Image-2.1-Turbo',
'base_revision': BASE_REVISION, 'weight_dtype': 'float8_e4m3fn', 'activation_dtype': 'float8_e4m3fn',
'weight_scaling': 'per-output-channel FP32', 'activation_scaling': 'dynamic per-token FP32',
'weight_rounding': 'per layer: round-to-nearest with clipping search, or GPTQ error feedback (full-token second moments)',
'accumulation': 'FP32, use_fast_accum=False', 'output_dtype': 'bfloat16',
'kernel': 'torch._scaled_mm native CUTLASS SM120, Triton fused smoothing and activation quantization',
'quantized_modules': quantized, 'kept_bf16_modules': list(held), 'modified_notice': NOTICE}
(OUT / 'quantization_config.json').write_text(json.dumps(config, indent=2))
report = OUT_ROOT / 'reports'; report.mkdir(parents=True, exist_ok=True)
(report / 'calibration_search.json').write_text(json.dumps(results, indent=2))
shutil.copy(CALIBRATION / 'manifest.json', report / 'calibration_manifest.json')
state = model.state_dict(); shards, shard, size = [], {}, 0
for n, v in state.items():
nb = v.numel() * v.element_size()
if shard and size + nb > 3_500_000_000:
shards.append(shard); shard, size = {}, 0
shard[n] = v.detach().cpu().contiguous(); size += nb
if shard:
shards.append(shard)
for old in OUT.glob('model-*.safetensors'):
old.unlink()
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 = {'quantized_modules': len(quantized), 'bf16_outliers': list(held), 'transformer_bytes': total,
'mean_diagnostic_nrmse': statistics.mean(v['diagnostic_nrmse'] for v in vals),
'max_diagnostic_nrmse': max(v['diagnostic_nrmse'] for v in vals),
'gptq_selected': sum(v['rounding'] == 'gptq' for v in vals),
'mean_fit_nrmse': statistics.mean(v['fit_nrmse'] for v in vals),
'mean_nearest_fit_nrmse_same_alpha': statistics.mean(v['nearest_fit_nrmse_same_alpha'] for v in vals)}
(report / 'quantization_summary.json').write_text(json.dumps(summary, indent=2))
print('QUANTIZATION_COMPLETE', json.dumps(summary), flush=True)
if __name__ == '__main__':
main()
|