253 lines
10 KiB
Python
253 lines
10 KiB
Python
#!/usr/bin/env python
|
|
"""Issue 5: quantize the Qwen-Image-2.1 DiT to int8 weight-only (W8A16) and
|
|
save it as a compressed-tensors pack-quantized checkpoint.
|
|
|
|
Quantizes exactly the linears the vllm-omni transformer constructs with
|
|
``quant_config`` — per-block attention (to_q/to_k/to_v/to_out.0) and
|
|
img_mlp (proj/gate_layer/out) — and copies every boundary layer (img_in,
|
|
txt_in, time_text_embed, modulation, norm_out, proj_out) plus the RMSNorm
|
|
q/k scales through in BF16. Per-output-channel symmetric int8, scales in
|
|
fp32, weights packed 4x uint8 (biased by 128) per int32 via
|
|
compressed-tensors' own ``pack_to_int32`` so the artifact matches the
|
|
llm-compressor storage convention bit for bit.
|
|
|
|
The artifact is a full diffusers-layout model directory: the quantized
|
|
``transformer/`` is real, the remaining components are symlinks into the
|
|
sealed snapshot. ``transformer/config.json`` carries the quantization
|
|
metadata wrapped per-component (``{"transformer": {...}}``) so vLLM-Omni's
|
|
auto-detection quantizes the DiT only — the text encoder and VAE stay BF16.
|
|
|
|
A MANIFEST.json (inputs, versions, per-file sha256, quantization stats) is
|
|
written beside the weights. Nothing here enters Git; the artifact lives
|
|
under the helper-system state root.
|
|
|
|
Usage:
|
|
quantize_dit_int8wo.py --snapshot SNAPSHOT_DIR --out ARTIFACT_DIR
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import datetime as dt
|
|
import hashlib
|
|
import json
|
|
import platform
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from compressed_tensors.compressors.pack_quantized.helpers import pack_to_int32
|
|
from compressed_tensors.config import CompressionFormat
|
|
from compressed_tensors.quantization import (
|
|
QuantizationArgs,
|
|
QuantizationConfig,
|
|
QuantizationScheme,
|
|
QuantizationStatus,
|
|
QuantizationStrategy,
|
|
QuantizationType,
|
|
)
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
SNAPSHOT_COMMIT = "790c92633540aa0cb11d9abf19eb46d861714758"
|
|
# Boundary layers the vllm-omni model pins to quant_config=None (BF16), plus
|
|
# the non-linear tensors; everything else in a block is quantized.
|
|
IGNORED_TARGETS = ["img_in", "txt_in", "time_text_embed", "modulation", "norm_out", "proj_out"]
|
|
QUANTIZED_KEY = re.compile(
|
|
r"^transformer_blocks\.\d+\."
|
|
r"(attn\.(to_q|to_k|to_v|to_out\.0)|img_mlp\.(proj|gate_layer|out))\.weight$"
|
|
)
|
|
SYMLINKED_COMPONENTS = ("scheduler", "text_encoder", "vae", "processor")
|
|
SYMLINKED_FILES = ("model_index.json",)
|
|
|
|
|
|
def build_quantization_dict(scheme: str) -> dict:
|
|
"""Serialized compressed-tensors metadata (llm-compressor schema).
|
|
|
|
scheme="int8wo": int8 weight-only, pack-quantized storage.
|
|
scheme="w8a8dyn": int8 weights (static per-channel) + dynamic per-token
|
|
symmetric int8 activations, int-quantized storage — the non-Marlin
|
|
cutlass path on sm_86.
|
|
"""
|
|
weights = QuantizationArgs(
|
|
num_bits=8,
|
|
type=QuantizationType.INT,
|
|
symmetric=True,
|
|
strategy=QuantizationStrategy.CHANNEL,
|
|
dynamic=False,
|
|
observer="minmax",
|
|
)
|
|
input_activations = (
|
|
QuantizationArgs(
|
|
num_bits=8,
|
|
type=QuantizationType.INT,
|
|
symmetric=True,
|
|
strategy=QuantizationStrategy.TOKEN,
|
|
dynamic=True,
|
|
)
|
|
if scheme == "w8a8dyn"
|
|
else None
|
|
)
|
|
scheme_obj = QuantizationScheme(targets=["Linear"], weights=weights, input_activations=input_activations)
|
|
config = QuantizationConfig(
|
|
config_groups={"group_0": scheme_obj},
|
|
format=CompressionFormat.int_quantized if scheme == "w8a8dyn" else CompressionFormat.pack_quantized,
|
|
ignore=list(IGNORED_TARGETS),
|
|
quant_method="compressed-tensors",
|
|
quantization_status=QuantizationStatus.FROZEN,
|
|
)
|
|
return config.to_dict()
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1 << 20), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def quantize_weight(weight: torch.Tensor, scheme: str) -> tuple[dict[str, torch.Tensor], dict]:
|
|
"""Per-output-channel symmetric int8.
|
|
|
|
scheme="int8wo": q = round(W/s), s = rowmax(|W|)/127, packed via
|
|
compressed-tensors' packer (uint8 biased by 128, 4 per int32).
|
|
scheme="w8a8dyn": same quantization, stored unpacked as plain int8
|
|
(int-quantized format) with fp32 scales.
|
|
"""
|
|
w = weight.detach().to(torch.float32)
|
|
scale = w.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
|
|
q = torch.round(w / scale).clamp(-127, 127).to(torch.int8)
|
|
dequant = q.to(torch.float32) * scale
|
|
err = (w - dequant).norm() / w.norm().clamp_min(1e-12)
|
|
stats = {"rel_l2": float(err), "out": int(w.shape[0]), "in": int(w.shape[1])}
|
|
if scheme == "w8a8dyn":
|
|
return {"weight": q.contiguous(), "weight_scale": scale.to(torch.float32).contiguous()}, stats
|
|
return {
|
|
"weight_packed": pack_to_int32(q.contiguous(), 8).contiguous(),
|
|
"weight_scale": scale.to(torch.float32).contiguous(),
|
|
"weight_shape": torch.tensor(weight.shape, dtype=torch.int64),
|
|
}, stats
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--snapshot", type=Path, required=True, help="sealed snapshot dir")
|
|
parser.add_argument("--out", type=Path, required=True, help="artifact dir to create")
|
|
parser.add_argument(
|
|
"--scheme",
|
|
choices=("int8wo", "w8a8dyn"),
|
|
default="int8wo",
|
|
help="int8 weight-only (pack-quantized, Marlin path) or W8A8 int8-dynamic (int-quantized, cutlass path)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
snapshot = args.snapshot.resolve()
|
|
if not (snapshot / "transformer" / "config.json").exists():
|
|
sys.exit(f"snapshot transformer config not found under {snapshot}")
|
|
if args.out.exists():
|
|
sys.exit(f"refusing to overwrite existing {args.out}")
|
|
artifact = args.out
|
|
(artifact / "transformer").mkdir(parents=True)
|
|
|
|
index = json.loads(
|
|
(snapshot / "transformer" / "diffusion_pytorch_model.safetensors.index.json").read_text()
|
|
)
|
|
shards = sorted(set(index["weight_map"].values()))
|
|
state: dict[str, torch.Tensor] = {}
|
|
quant_stats: dict[str, dict] = {}
|
|
quant_params = bf16_params = 0
|
|
for shard in shards:
|
|
tensors = load_file(str(snapshot / "transformer" / shard))
|
|
for name, tensor in tensors.items():
|
|
if QUANTIZED_KEY.match(name):
|
|
emitted, stats = quantize_weight(tensor, args.scheme)
|
|
if args.scheme == "w8a8dyn":
|
|
# int-quantized storage: plain int8 weight + scale, mirroring
|
|
# the vLLM layer parameter names (weight / weight_scale).
|
|
state[name] = emitted["weight"]
|
|
state[f"{name}_scale"] = emitted["weight_scale"]
|
|
else:
|
|
state[f"{name}_packed"] = emitted["weight_packed"]
|
|
state[f"{name}_scale"] = emitted["weight_scale"]
|
|
state[f"{name}_shape"] = emitted["weight_shape"]
|
|
quant_stats[name] = stats
|
|
quant_params += tensor.numel()
|
|
print(f"q {name} {tuple(tensor.shape)} rel_l2={stats['rel_l2']:.5f}", flush=True)
|
|
else:
|
|
state[name] = tensor
|
|
bf16_params += tensor.numel()
|
|
del tensors
|
|
|
|
weights_path = artifact / "transformer" / "diffusion_pytorch_model.safetensors"
|
|
save_file(state, str(weights_path), metadata={"format": "pt"})
|
|
total_bytes = weights_path.stat().st_size
|
|
|
|
config = json.loads((snapshot / "transformer" / "config.json").read_text())
|
|
config["quantization_config"] = {"transformer": build_quantization_dict(args.scheme)}
|
|
(artifact / "transformer" / "config.json").write_text(json.dumps(config, indent=2) + "\n")
|
|
|
|
for name in SYMLINKED_FILES:
|
|
target = snapshot / name
|
|
(artifact / name).symlink_to(target)
|
|
for name in SYMLINKED_COMPONENTS:
|
|
target = snapshot / name
|
|
(artifact / name).symlink_to(target, target_is_directory=True)
|
|
|
|
files = [
|
|
{"path": "transformer/config.json", "bytes": (artifact / "transformer" / "config.json").stat().st_size},
|
|
{"path": "transformer/diffusion_pytorch_model.safetensors", "bytes": total_bytes},
|
|
]
|
|
for entry in files:
|
|
entry["sha256"] = sha256_file(artifact / entry["path"])
|
|
tree_json = snapshot.parent.parent / "trees" / f"{SNAPSHOT_COMMIT}.json"
|
|
manifest = {
|
|
"schema": f"issue5-{args.scheme}-v1",
|
|
"created_utc": dt.datetime.now(dt.UTC).replace(microsecond=0).isoformat(),
|
|
"description": (
|
|
f"Qwen-Image-2.1 DiT {args.scheme} "
|
|
+ (
|
|
"W8A8 int8-dynamic, int-quantized"
|
|
if args.scheme == "w8a8dyn"
|
|
else "int8 weight-only W8A16, pack-quantized"
|
|
)
|
|
+ ", compressed-tensors"
|
|
),
|
|
"inputs": {
|
|
"snapshot": str(snapshot),
|
|
"revision": SNAPSHOT_COMMIT,
|
|
"tree_manifest": str(tree_json),
|
|
"tree_sha256": sha256_file(tree_json) if tree_json.exists() else None,
|
|
},
|
|
"versions": {
|
|
"python": platform.python_version(),
|
|
"torch": torch.__version__,
|
|
"compressed_tensors": __import__("compressed_tensors").__version__,
|
|
},
|
|
"quantization": {
|
|
"method": "compressed-tensors",
|
|
"format": "pack-quantized",
|
|
"weights": "int8 symmetric per-output-channel, static, minmax",
|
|
"config": build_quantization_dict(args.scheme),
|
|
"ignored_layers": list(IGNORED_TARGETS),
|
|
"note": "ignore list mirrors the vllm-omni model's quant_config=None boundary layers",
|
|
},
|
|
"stats": {
|
|
"quantized_tensors": len(quant_stats),
|
|
"quantized_params": quant_params,
|
|
"bf16_params": bf16_params,
|
|
"mean_rel_l2": sum(s["rel_l2"] for s in quant_stats.values()) / max(len(quant_stats), 1),
|
|
"max_rel_l2": max((s["rel_l2"] for s in quant_stats.values()), default=0.0),
|
|
"weights_file_bytes": total_bytes,
|
|
},
|
|
"files": files,
|
|
}
|
|
(artifact / "MANIFEST.json").write_text(json.dumps(manifest, indent=1) + "\n")
|
|
print(json.dumps(manifest["stats"], indent=1))
|
|
print(f"artifact: {artifact}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|