#!/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())