File size: 4,376 Bytes
f642d1f | 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 | """Calibrated W8A8 E4M3 inference using native CUTLASS SM120 scaled GEMMs.
Built with Qwen. Quantization modifications by ProCreations, 2026.
The upstream model is subject to the included Qwen Research License.
"""
import json
from pathlib import Path
import torch
from torch import nn
import triton
import triton.language as tl
from safetensors.torch import load_file
@triton.jit
def _quantize_rows(X, SMOOTH, Q, SCALE, K: tl.constexpr, BLOCK: tl.constexpr):
row = tl.program_id(0)
cols = tl.arange(0, BLOCK)
x = tl.load(X + row * K + cols, cols < K, 0).to(tl.float32)
s = tl.load(SMOOTH + cols, cols < K, 1).to(tl.float32)
z = x / s
scale = tl.maximum(tl.max(tl.abs(z), 0) / 448.0, 1.e-12)
tl.store(Q + row * K + cols, z / scale, cols < K)
tl.store(SCALE + row, scale)
def quantize_activation(x, smooth):
x = x.reshape(-1, x.shape[-1]).contiguous()
q = torch.empty_like(x, dtype=torch.float8_e4m3fn)
scale = torch.empty((x.shape[0], 1), device=x.device, dtype=torch.float32)
_quantize_rows[(x.shape[0],)](x, smooth, q, scale, x.shape[1],
triton.next_power_of_2(x.shape[1]), num_warps=8)
return q, scale
class CalibratedFP8Linear(nn.Module):
def __init__(self, weight, scale, smooth):
super().__init__()
self.register_buffer("weight", weight)
self.register_buffer("weight_scale", scale.float())
self.register_buffer("smooth", smooth.float())
self.in_features = weight.shape[1]
self.out_features = weight.shape[0]
self.bias = None
def forward(self, x):
shape = x.shape[:-1]
q, scale = quantize_activation(x, self.smooth)
y = torch._scaled_mm(q, self.weight.t(), scale_a=scale,
scale_b=self.weight_scale.t(),
out_dtype=torch.bfloat16, use_fast_accum=False)
return y.reshape(*shape, self.out_features)
def eligible_modules(model):
return {n:m for n,m in model.named_modules()
if isinstance(m, nn.Linear) and n.startswith("transformer_blocks.")
and m.bias is None}
def replace_module(root, name, replacement):
parent_name, leaf = name.rsplit(".", 1)
setattr(root.get_submodule(parent_name), leaf, replacement)
def load_fp8_transformer(directory, device="cuda"):
from diffusers import QwenImage21Transformer2DModel
directory = Path(directory)
config = json.loads((directory / "config.json").read_text())
meta = json.loads((directory / "quantization_config.json").read_text())
# Construct without allocating an intermediate BF16 model.
with torch.device("meta"):
model = QwenImage21Transformer2DModel.from_config(config)
for n, spec in meta["quantized_modules"].items():
out_f, in_f = spec["shape"]
replace_module(model, n, CalibratedFP8Linear(
torch.empty((out_f,in_f),dtype=torch.float8_e4m3fn),
torch.empty((out_f,1),dtype=torch.float32),
torch.empty(in_f,dtype=torch.float32)))
state = {}
for path in sorted(directory.glob("model-*.safetensors")):
state.update(load_file(str(path), device="cpu"))
model.load_state_dict(state, strict=True, assign=True)
# The original rotary module has nonpersistent buffers / Python tensor lists.
# Reconstruct these by creating a lightweight ordinary position module.
from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21Rope
model.pos_embed = QwenImage21Rope(theta=10000, axes_dim=list(model.config.axes_dims_rope))
from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21TemporalTimesteps
model.time_text_embed.time_proj = QwenImage21TemporalTimesteps(timestep_dim=256)
return model.eval().to(device=device)
def load_pipeline(base, quant=None):
from diffusers import QwenImage21Pipeline
args = {}
if not Path(base).exists():
args["revision"] = "b3179ad355be050328e483a9dfdd9e60cd62adfa"
if quant:
args["transformer"] = load_fp8_transformer(quant, device="cuda")
pipe = QwenImage21Pipeline.from_pretrained(base, torch_dtype=torch.bfloat16, **args)
# Do not call pipe.to(dtype=...) after FP8 loading; scales are FP32.
pipe.to(device="cuda")
pipe.set_progress_bar_config(disable=True)
return pipe
|