vllm-omni/probe/issue-5/quantize_dit_int8wo.py

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