paroquant 0.1.0__tar.gz

paroquant-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Z Lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

paroquant-0.1.0/PKG-INFO

@@ -0,0 +1,151 @@
+Metadata-Version: 2.4
+Name: paroquant
+Version: 0.1.0
+Summary: ParoQuant — Pairwise Rotation Quantization for LLMs
+Author: Z Lab
+License-Expression: MIT
+Project-URL: Homepage, https://paroquant.z-lab.ai
+Project-URL: Paper, https://arxiv.org/abs/2511.10645
+Project-URL: Models, https://huggingface.co/collections/z-lab/paroquant
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: rich
+Provides-Extra: transformers
+Requires-Dist: torch>=2.8; extra == "transformers"
+Requires-Dist: transformers>=4.55; extra == "transformers"
+Requires-Dist: autoawq; extra == "transformers"
+Provides-Extra: vllm
+Requires-Dist: vllm>=0.15; extra == "vllm"
+Provides-Extra: mlx
+Requires-Dist: mlx; extra == "mlx"
+Requires-Dist: mlx-lm; extra == "mlx"
+Requires-Dist: mlx-vlm; extra == "mlx"
+Provides-Extra: optim
+Requires-Dist: paroquant[transformers]; extra == "optim"
+Requires-Dist: datasets; extra == "optim"
+Requires-Dist: simple_parsing; extra == "optim"
+Requires-Dist: tqdm; extra == "optim"
+Provides-Extra: eval
+Requires-Dist: lm_eval; extra == "eval"
+Requires-Dist: zstandard; extra == "eval"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pre-commit; extra == "dev"
+Dynamic: license-file
+
+# ParoQuant
+
+**Pairwise Rotation Quantization for Efficient Reasoning LLM Inference**
+
+<p align="center">
+  <a href="https://arxiv.org/abs/2511.10645"><img src="https://img.shields.io/badge/arXiv-2511.10645-b31b1b.svg" alt="Paper"></a>
+  <a href="https://paroquant.z-lab.ai"><img src="https://img.shields.io/badge/Blog-ParoQuant-blue" alt="Blog"></a>
+  <a href="https://huggingface.co/collections/z-lab/paroquant"><img src="https://img.shields.io/badge/%F0%9F%A4%97-Models-yellow" alt="Models"></a>
+</p>
+
+State-of-the-art INT4 quantization for LLMs. ParoQuant uses learned pairwise rotations to suppress weight outliers, closing the accuracy gap with FP16 while running at near-AWQ speed. Supports NVIDIA GPUs (vLLM, Transformers) and Apple Silicon (MLX).
+
+<p align="center">
+  <a href="https://youtu.be/fISG4CkizLM">
+    <img src="https://img.youtube.com/vi/fISG4CkizLM/maxresdefault.jpg" width="80%">
+  </a>
+</p>
+
+## Quick Start
+
+**NVIDIA GPU:**
+
+```bash
+pip install paroquant[vllm]
+python -m paroquant.cli.chat --model z-lab/Qwen3-8B-PARO
+
+# or with Docker
+docker run --pull=always --rm -it --gpus all --ipc=host \
+  ghcr.io/z-lab/paroquant:chat --model z-lab/Qwen3-8B-PARO
+```
+
+**Apple Silicon:**
+
+```bash
+pip install paroquant[mlx]
+python -m paroquant.cli.chat --model z-lab/Qwen3-8B-PARO
+```
+
+## Models
+
+All models are available on [Hugging Face](https://huggingface.co/collections/z-lab/paroquant). Swap the model name in the commands above to try any of them.
+
+**Qwen3**
+
+| Model | Checkpoint |
+|---|---|
+| Qwen3-0.6B | [`z-lab/Qwen3-0.6B-PARO`](https://huggingface.co/z-lab/Qwen3-0.6B-PARO) |
+| Qwen3-1.7B | [`z-lab/Qwen3-1.7B-PARO`](https://huggingface.co/z-lab/Qwen3-1.7B-PARO) |
+| Qwen3-4B | [`z-lab/Qwen3-4B-PARO`](https://huggingface.co/z-lab/Qwen3-4B-PARO) |
+| Qwen3-8B | [`z-lab/Qwen3-8B-PARO`](https://huggingface.co/z-lab/Qwen3-8B-PARO) |
+| Qwen3-14B | [`z-lab/Qwen3-14B-PARO`](https://huggingface.co/z-lab/Qwen3-14B-PARO) |
+| Qwen3-4B-Thinking-2507 | [`z-lab/Qwen3-4B-Thinking-2507-PARO`](https://huggingface.co/z-lab/Qwen3-4B-Thinking-2507-PARO) |
+
+**Llama**
+
+| Model | Checkpoint |
+|---|---|
+| Llama-2-7B | [`z-lab/Llama-2-7b-hf-PARO`](https://huggingface.co/z-lab/Llama-2-7b-hf-PARO) |
+| Llama-3-8B | [`z-lab/Meta-Llama-3-8B-PARO`](https://huggingface.co/z-lab/Meta-Llama-3-8B-PARO) |
+| Llama-3-70B | [`z-lab/Meta-Llama-3-70B-PARO`](https://huggingface.co/z-lab/Meta-Llama-3-70B-PARO) |
+| Llama-3.1-8B-Instruct | [`z-lab/Llama-3.1-8B-Instruct-PARO`](https://huggingface.co/z-lab/Llama-3.1-8B-Instruct-PARO) |
+
+Want a model that's not listed? [Open an issue](https://github.com/z-lab/paroquant/issues/new) and let us know.
+
+## Installation
+
+```bash
+git clone https://github.com/z-lab/paroquant && cd paroquant
+
+pip install -e ".[vllm]"            # vLLM backend (GPU, recommended)
+pip install -e ".[transformers]"    # Transformers backend (GPU)
+pip install -e ".[mlx]"             # MLX backend (Apple Silicon)
+pip install -e ".[optim,eval]"      # Optimization & evaluation
+```
+
+Or use Docker: `docker run -it --gpus all --ipc=host ghcr.io/z-lab/paroquant:latest`
+
+## Quantize Your Own Model
+
+```bash
+# 1. Optimize rotation parameters
+experiments/optimize/4bit.sh Qwen/Qwen3-8B
+
+# 2. Export to HF checkpoint (--mode real for INT4, --mode pseudo for FP16)
+python -m paroquant.cli.convert \
+  --model Qwen/Qwen3-8B \
+  --result-dir output/Qwen3-8B \
+  --output-path models/Qwen3-8B-PARO
+```
+
+## Reproduction
+
+See [`experiments/README.md`](./experiments/README.md) for scripts to reproduce all results in the paper.
+
+## Docker Images
+
+| Image | Purpose |
+|---|---|
+| `ghcr.io/z-lab/paroquant:latest` | Optimization & evaluation |
+| `ghcr.io/z-lab/paroquant:chat` | Interactive chat |
+| `ghcr.io/z-lab/paroquant:chat-cu130` | Interactive chat (CUDA 13.0 / ARM64) |
+| `ghcr.io/z-lab/paroquant:eval-reasoning` | Reasoning task evaluation |
+
+## Citation
+
+```bibtex
+@inproceedings{liang2026paroquant,
+  title     = {{ParoQuant: Pairwise Rotation Quantization for Efficient Reasoning LLM Inference}},
+  author    = {Liang, Yesheng and Chen, Haisheng and Zhang, Zihan and Han, Song and Liu, Zhijian},
+  booktitle = {International Conference on Learning Representations (ICLR)},
+  year      = {2026}
+}
+```

paroquant-0.1.0/README.md

@@ -0,0 +1,113 @@
+# ParoQuant
+
+**Pairwise Rotation Quantization for Efficient Reasoning LLM Inference**
+
+<p align="center">
+  <a href="https://arxiv.org/abs/2511.10645"><img src="https://img.shields.io/badge/arXiv-2511.10645-b31b1b.svg" alt="Paper"></a>
+  <a href="https://paroquant.z-lab.ai"><img src="https://img.shields.io/badge/Blog-ParoQuant-blue" alt="Blog"></a>
+  <a href="https://huggingface.co/collections/z-lab/paroquant"><img src="https://img.shields.io/badge/%F0%9F%A4%97-Models-yellow" alt="Models"></a>
+</p>
+
+State-of-the-art INT4 quantization for LLMs. ParoQuant uses learned pairwise rotations to suppress weight outliers, closing the accuracy gap with FP16 while running at near-AWQ speed. Supports NVIDIA GPUs (vLLM, Transformers) and Apple Silicon (MLX).
+
+<p align="center">
+  <a href="https://youtu.be/fISG4CkizLM">
+    <img src="https://img.youtube.com/vi/fISG4CkizLM/maxresdefault.jpg" width="80%">
+  </a>
+</p>
+
+## Quick Start
+
+**NVIDIA GPU:**
+
+```bash
+pip install paroquant[vllm]
+python -m paroquant.cli.chat --model z-lab/Qwen3-8B-PARO
+
+# or with Docker
+docker run --pull=always --rm -it --gpus all --ipc=host \
+  ghcr.io/z-lab/paroquant:chat --model z-lab/Qwen3-8B-PARO
+```
+
+**Apple Silicon:**
+
+```bash
+pip install paroquant[mlx]
+python -m paroquant.cli.chat --model z-lab/Qwen3-8B-PARO
+```
+
+## Models
+
+All models are available on [Hugging Face](https://huggingface.co/collections/z-lab/paroquant). Swap the model name in the commands above to try any of them.
+
+**Qwen3**
+
+| Model | Checkpoint |
+|---|---|
+| Qwen3-0.6B | [`z-lab/Qwen3-0.6B-PARO`](https://huggingface.co/z-lab/Qwen3-0.6B-PARO) |
+| Qwen3-1.7B | [`z-lab/Qwen3-1.7B-PARO`](https://huggingface.co/z-lab/Qwen3-1.7B-PARO) |
+| Qwen3-4B | [`z-lab/Qwen3-4B-PARO`](https://huggingface.co/z-lab/Qwen3-4B-PARO) |
+| Qwen3-8B | [`z-lab/Qwen3-8B-PARO`](https://huggingface.co/z-lab/Qwen3-8B-PARO) |
+| Qwen3-14B | [`z-lab/Qwen3-14B-PARO`](https://huggingface.co/z-lab/Qwen3-14B-PARO) |
+| Qwen3-4B-Thinking-2507 | [`z-lab/Qwen3-4B-Thinking-2507-PARO`](https://huggingface.co/z-lab/Qwen3-4B-Thinking-2507-PARO) |
+
+**Llama**
+
+| Model | Checkpoint |
+|---|---|
+| Llama-2-7B | [`z-lab/Llama-2-7b-hf-PARO`](https://huggingface.co/z-lab/Llama-2-7b-hf-PARO) |
+| Llama-3-8B | [`z-lab/Meta-Llama-3-8B-PARO`](https://huggingface.co/z-lab/Meta-Llama-3-8B-PARO) |
+| Llama-3-70B | [`z-lab/Meta-Llama-3-70B-PARO`](https://huggingface.co/z-lab/Meta-Llama-3-70B-PARO) |
+| Llama-3.1-8B-Instruct | [`z-lab/Llama-3.1-8B-Instruct-PARO`](https://huggingface.co/z-lab/Llama-3.1-8B-Instruct-PARO) |
+
+Want a model that's not listed? [Open an issue](https://github.com/z-lab/paroquant/issues/new) and let us know.
+
+## Installation
+
+```bash
+git clone https://github.com/z-lab/paroquant && cd paroquant
+
+pip install -e ".[vllm]"            # vLLM backend (GPU, recommended)
+pip install -e ".[transformers]"    # Transformers backend (GPU)
+pip install -e ".[mlx]"             # MLX backend (Apple Silicon)
+pip install -e ".[optim,eval]"      # Optimization & evaluation
+```
+
+Or use Docker: `docker run -it --gpus all --ipc=host ghcr.io/z-lab/paroquant:latest`
+
+## Quantize Your Own Model
+
+```bash
+# 1. Optimize rotation parameters
+experiments/optimize/4bit.sh Qwen/Qwen3-8B
+
+# 2. Export to HF checkpoint (--mode real for INT4, --mode pseudo for FP16)
+python -m paroquant.cli.convert \
+  --model Qwen/Qwen3-8B \
+  --result-dir output/Qwen3-8B \
+  --output-path models/Qwen3-8B-PARO
+```
+
+## Reproduction
+
+See [`experiments/README.md`](./experiments/README.md) for scripts to reproduce all results in the paper.
+
+## Docker Images
+
+| Image | Purpose |
+|---|---|
+| `ghcr.io/z-lab/paroquant:latest` | Optimization & evaluation |
+| `ghcr.io/z-lab/paroquant:chat` | Interactive chat |
+| `ghcr.io/z-lab/paroquant:chat-cu130` | Interactive chat (CUDA 13.0 / ARM64) |
+| `ghcr.io/z-lab/paroquant:eval-reasoning` | Reasoning task evaluation |
+
+## Citation
+
+```bibtex
+@inproceedings{liang2026paroquant,
+  title     = {{ParoQuant: Pairwise Rotation Quantization for Efficient Reasoning LLM Inference}},
+  author    = {Liang, Yesheng and Chen, Haisheng and Zhang, Zihan and Han, Song and Liu, Zhijian},
+  booktitle = {International Conference on Learning Representations (ICLR)},
+  year      = {2026}
+}
+```

paroquant-0.1.0/paroquant/cli/benchmark.py

@@ -0,0 +1,53 @@
+import argparse
+import asyncio
+import time
+
+from paroquant.inference import GenerationParams, build_prompt, create_generator
+
+
+async def bench(gen, prompt: str, max_tokens: int, warmup: int, runs: int):
+    params = GenerationParams(max_tokens=max_tokens, temperature=0.0)
+    formatted = build_prompt(gen.tokenizer, [{"role": "user", "content": prompt}])
+
+    for _ in range(warmup):
+        await gen.generate(formatted, params)
+
+    times, tokens = [], []
+    for _ in range(runs):
+        start = time.perf_counter()
+        result = await gen.generate(formatted, params)
+        elapsed = time.perf_counter() - start
+        times.append(elapsed)
+        tokens.append(result.stats.num_tokens)
+
+    avg_tokens = sum(tokens) / len(tokens)
+    avg_time = sum(times) / len(times)
+    avg_tps = avg_tokens / avg_time if avg_time > 0 else 0
+    return avg_tokens, avg_time, avg_tps
+
+
+async def main():
+    parser = argparse.ArgumentParser(description="Benchmark generation throughput")
+    parser.add_argument("--model", type=str, required=True)
+    parser.add_argument("--backend", type=str, default="auto", choices=["auto", "mlx", "vllm", "transformers"])
+    parser.add_argument("--max-tokens", type=int, default=128)
+    parser.add_argument("--warmup", type=int, default=2)
+    parser.add_argument("--runs", type=int, default=5)
+    parser.add_argument("--prompt", type=str, default="Write a short essay about the history of computing.")
+    args = parser.parse_args()
+
+    print(f"Loading {args.model} with {args.backend} backend...")
+    start = time.perf_counter()
+    gen = create_generator(args.backend, args.model)
+    load_time = time.perf_counter() - start
+    print(f"Loaded in {load_time:.1f}s")
+
+    print(f"Benchmarking: {args.warmup} warmup + {args.runs} runs, {args.max_tokens} tokens each")
+    avg_tokens, avg_time, avg_tps = await bench(gen, args.prompt, args.max_tokens, args.warmup, args.runs)
+    print(f"  Avg tokens: {avg_tokens:.0f}")
+    print(f"  Avg time:   {avg_time:.2f}s")
+    print(f"  Throughput:  {avg_tps:.1f} tok/s")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

paroquant-0.1.0/paroquant/cli/chat.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+import contextlib
+import importlib
+import io
+import os
+import re
+import time
+import warnings
+
+from rich.console import Console
+from rich.live import Live
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich.theme import Theme
+
+from paroquant.inference import (
+    GenerationParams,
+    build_prompt,
+    create_generator,
+    detect_backend,
+)
+
+
+@contextlib.contextmanager
+def _silence_stderr():
+    with contextlib.redirect_stderr(io.StringIO()):
+        yield
+
+
+_THINKING_LINES = 4
+_SPECIAL_RE = re.compile(r"<\|[^|]+\|>")
+
+
+class _ThinkingTracker:
+    """Parses raw streamed text, shows thinking in a small live box, then folds it."""
+
+    _REFRESH_INTERVAL = 1.0 / 24
+
+    def __init__(self, console: Console, enable_thinking: bool):
+        self.console = console
+        self.enable_thinking = enable_thinking
+        self.raw = ""
+        self.in_thinking = enable_thinking
+        self._live: Live | None = None
+        self._think_start = time.perf_counter()
+        self._last_update = 0.0
+
+    def on_token(self, text: str):
+        self.raw += text
+
+        if self.in_thinking:
+            if "</think>" in self.raw:
+                self.in_thinking = False
+                self.stop()
+                self._start_response_live()
+            else:
+                self._update_thinking_box()
+        else:
+            if self._live is None:
+                self._start_response_live()
+            else:
+                self._update_response()
+
+    def _update_thinking_box(self):
+        now = time.perf_counter()
+        if self._live is not None and now - self._last_update < self._REFRESH_INTERVAL:
+            return
+        self._last_update = now
+
+        lines = self.raw.splitlines()
+        tail = "\n".join(lines[-_THINKING_LINES:])
+        panel = Panel(
+            tail or "...",
+            title=f"thinking ({now - self._think_start:.1f}s)",
+            border_style="dim",
+            width=min(self.console.width, 80),
+            height=_THINKING_LINES + 2,
+        )
+        if self._live is None:
+            self._live = Live(panel, console=self.console, transient=True)
+            self._live.start()
+        else:
+            self._live.update(panel)
+
+    def _start_response_live(self):
+        self._live = Live(Markdown(""), console=self.console, vertical_overflow="visible")
+        self._live.start()
+        self._update_response()
+
+    def _update_response(self):
+        now = time.perf_counter()
+        if now - self._last_update < self._REFRESH_INTERVAL:
+            return
+        self._last_update = now
+
+        response = self._get_response()
+        if response and self._live is not None:
+            self._live.update(Markdown(response))
+
+    def _get_response(self) -> str:
+        if self.enable_thinking and "</think>" in self.raw:
+            text = self.raw.split("</think>", 1)[1].lstrip("\n")
+        elif self.in_thinking:
+            text = ""
+        else:
+            text = self.raw
+        return _SPECIAL_RE.sub("", text)
+
+    def stop(self):
+        if self._live is not None:
+            response = self._get_response()
+            if response:
+                self._live.update(Markdown(response))
+            self._live.stop()
+            self._live = None
+
+    @property
+    def output_text(self) -> str:
+        return self._get_response().strip()
+
+
+def _suppress_library_noise():
+    warnings.filterwarnings("ignore")
+    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+    os.environ.setdefault("VLLM_LOGGING_LEVEL", "ERROR")
+    for mod, fn in [
+        ("huggingface_hub", "disable_progress_bars"),
+        ("transformers.utils.logging", "set_verbosity_error"),
+    ]:
+        try:
+            getattr(importlib.import_module(mod), fn)()
+        except Exception:
+            pass
+
+
+async def run_chat_app(model: str, backend: str, params: GenerationParams):
+    _suppress_library_noise()
+
+    console = Console(theme=Theme({"hint": "dim"}))
+
+    backend = detect_backend() if backend == "auto" else backend
+
+    console.print(f"[hint]Loading model ({backend})...[/hint]")
+    generator = create_generator(backend, model)
+    console.clear()
+
+    enable_thinking = False
+
+    banner = (
+        f"[bold]ParoQuant Chat[/bold]\n"
+        f"Backend: [bold]{backend}[/bold]\n"
+        f"Model: [bold]{model}[/bold]\n\n"
+        f"[bold]/think[/bold] · [bold]/clear[/bold] · [bold]/quit[/bold]"
+    )
+    console.print(Panel.fit(banner, border_style="bright_blue"))
+
+    history: list[dict[str, str]] = []
+
+    try:
+        while True:
+            try:
+                user_prompt = console.input(">>> ").strip()
+            except (KeyboardInterrupt, EOFError):
+                console.print("\n[hint]Session closed.[/hint]")
+                break
+
+            if not user_prompt:
+                continue
+            cmd = user_prompt.lower()
+            if cmd in {"/quit", "quit", "/exit", "exit"}:
+                break
+            if cmd == "/clear":
+                history.clear()
+                console.clear()
+                console.print("[hint]Conversation history cleared.[/hint]\n")
+                continue
+            if cmd == "/think":
+                enable_thinking = not enable_thinking
+                console.print(f"[hint]Thinking {'on' if enable_thinking else 'off'}.[/hint]\n")
+                continue
+
+            history.append({"role": "user", "content": user_prompt})
+
+            prompt = build_prompt(generator.tokenizer, history, enable_thinking)
+
+            tracker = _ThinkingTracker(console, enable_thinking)
+            with _silence_stderr():
+                result = await generator.generate(prompt, params, on_text=tracker.on_token)
+            tracker.stop()
+
+            history.append({"role": "assistant", "content": tracker.output_text})
+
+            s = result.stats
+            parts = []
+            if s.ttft is not None:
+                parts.append(f"ttft {s.ttft:.2f}s")
+            parts += [f"{s.num_tokens} tokens", f"{s.tps:.1f} tok/s", f"{s.latency:.1f}s total"]
+            console.print(f"  {' · '.join(parts)}", style="hint", highlight=False)
+            console.print()
+    finally:
+        await generator.close()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, required=True)
+    parser.add_argument("--backend", type=str, default="auto", choices=["auto", "vllm", "transformers", "mlx"])
+    parser.add_argument("--max-tokens", type=int, default=8192)
+    parser.add_argument("--temperature", type=float, default=0.6)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument("--top-k", type=int, default=0)
+    parser.add_argument("--repetition-penalty", type=float, default=1.0)
+    args = parser.parse_args()
+
+    asyncio.run(
+        run_chat_app(
+            model=args.model,
+            backend=args.backend,
+            params=GenerationParams(
+                max_tokens=args.max_tokens,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                top_k=args.top_k,
+                repetition_penalty=args.repetition_penalty,
+            ),
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()