#!/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() -> dict: """Serialized compressed-tensors int8 weight-only metadata (llm-compressor schema).""" weights = QuantizationArgs( num_bits=8, type=QuantizationType.INT, symmetric=True, strategy=QuantizationStrategy.CHANNEL, dynamic=False, observer="minmax", ) scheme = QuantizationScheme(targets=["Linear"], weights=weights) config = QuantizationConfig( config_groups={"group_0": scheme}, format=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) -> tuple[torch.Tensor, torch.Tensor, dict]: """Per-output-channel symmetric int8: q = round(W/s), s = rowmax(|W|)/127. Returns (weight_packed int32 [out, in/4], weight_scale fp32 [out, 1], stats) using compressed-tensors' packer (uint8 biased by 128, 4 per int32). """ 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) packed = pack_to_int32(q.contiguous(), 8).contiguous() 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])} return packed, scale.to(torch.float32).contiguous(), 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") 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): packed, scale, stats = quantize_weight(tensor) state[f"{name}_packed"] = packed state[f"{name}_scale"] = scale state[f"{name}_shape"] = torch.tensor(tensor.shape, dtype=torch.int64) 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()} (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": "issue5-int8wo-v1", "created_utc": dt.datetime.now(dt.UTC).replace(microsecond=0).isoformat(), "description": "Qwen-Image-2.1 DiT int8 weight-only (W8A16), compressed-tensors pack-quantized", "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(), "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())