onecomp 1.0.0__tar.gz

onecomp-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright 2025-2026 Fujitsu Ltd.
+
+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.

onecomp-1.0.0/PKG-INFO

@@ -0,0 +1,303 @@
+Metadata-Version: 2.4
+Name: onecomp
+Version: 1.0.0
+Summary: Python package for LLM compression
+Author: Keiji Kimura
+License: MIT License
+        
+        Copyright 2025-2026 Fujitsu Ltd.
+        
+        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.
+        
+Project-URL: Homepage, https://github.com/FujitsuResearch/OneCompression
+Project-URL: Documentation, https://FujitsuResearch.github.io/OneCompression/
+Project-URL: Repository, https://github.com/FujitsuResearch/OneCompression
+Project-URL: Bug Tracker, https://github.com/FujitsuResearch/OneCompression/issues
+Project-URL: Changelog, https://github.com/FujitsuResearch/OneCompression/blob/main/CHANGELOG.md
+Keywords: llm,quantization,compression,post-training-quantization,gptq
+Classifier: Development Status :: 4 - Beta
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Intended Audience :: Science/Research
+Requires-Python: <3.14,>=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: transformers>=5.3.0
+Requires-Dist: accelerate
+Requires-Dist: datasets
+Requires-Dist: lm-eval
+Requires-Dist: primefac
+Requires-Dist: scipy
+Requires-Dist: gemlite>=0.5.1
+Requires-Dist: hqq>=0.2.8
+Requires-Dist: safetensors
+Requires-Dist: ortools>=9.15.6755
+Provides-Extra: cpu
+Requires-Dist: torch; extra == "cpu"
+Requires-Dist: torchvision; extra == "cpu"
+Provides-Extra: cu118
+Requires-Dist: torch; extra == "cu118"
+Requires-Dist: torchvision; extra == "cu118"
+Provides-Extra: cu121
+Requires-Dist: torch; extra == "cu121"
+Requires-Dist: torchvision; extra == "cu121"
+Provides-Extra: cu124
+Requires-Dist: torch; extra == "cu124"
+Requires-Dist: torchvision; extra == "cu124"
+Provides-Extra: cu126
+Requires-Dist: torch; extra == "cu126"
+Requires-Dist: torchvision; extra == "cu126"
+Provides-Extra: cu128
+Requires-Dist: torch; extra == "cu128"
+Requires-Dist: torchvision; extra == "cu128"
+Provides-Extra: dev
+Requires-Dist: black; extra == "dev"
+Requires-Dist: matplotlib>=3.10.8; extra == "dev"
+Requires-Dist: pylint; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Provides-Extra: vllm
+Requires-Dist: vllm; extra == "vllm"
+Provides-Extra: docs
+Requires-Dist: mkdocs-material; extra == "docs"
+Requires-Dist: mkdocstrings[python]; extra == "docs"
+Requires-Dist: mkdocs-gen-files; extra == "docs"
+Requires-Dist: mkdocs-literate-nav; extra == "docs"
+Dynamic: license-file
+
+# Fujitsu One Compression
+
+Fujitsu One Compression (OneComp) is a Python package for LLM compression.
+
+## 📖 Documentation
+
+Full documentation is available at **[https://FujitsuResearch.github.io/OneCompression/](https://FujitsuResearch.github.io/OneCompression/)**.
+
+## 📦 Features
+
+- **Quantization Error Propagation (QEP)**: A post-training quantization method that corrects quantization errors by propagating them to subsequent layers, improving the accuracy of quantized LLMs. See [Arai & Ichikawa, NeurIPS 2025](https://openreview.net/forum?id=a3l3K9khbL) for details. The original reference implementation is available at [FujitsuResearch/qep](https://github.com/FujitsuResearch/qep).
+- **vLLM Plugin Integration**: Serve OneComp-quantized models with [vLLM](https://docs.vllm.ai/) via built-in plugins for DBF and Mixed-GPTQ quantization methods.
+- **AutoBit**: Mixed-precision quantization with ILP-based bitwidth assignment. Automatically estimates the target bitwidth from available VRAM and assigns per-layer bitwidths to minimize quantization error under the memory budget.
+- **JointQ**: Joint quantization method that optimizes weight assignments and scale parameters simultaneously for improved quantization accuracy. Supports group-wise quantization (e.g., 4-bit, groupsize=128).
+- **LoRA SFT Post-Process**: Fine-tune quantized models with LoRA adapters for accuracy recovery or domain-specific knowledge injection. Supports SFT loss, teacher distillation, and intermediate block alignment.
+- **Rotation Preprocessing**: SpinQuant/OstQuant-based rotation preprocessing that reduces quantization error by learning optimal rotation matrices before quantization. Rotation/scaling matrices are absorbed into model weights, with online Hadamard hooks automatically registered at load time. Supports Llama and Qwen3 architectures.
+- (TBD)
+
+## 🤖 Supported Models
+
+OneComp has been verified with the following model architectures.
+Other Hugging Face-compatible models may work but are currently untested.
+
+| # | Architecture | Verified Models | Status |
+|---|-------------|-----------------|--------|
+| 1 | Llama | TinyLlama, Llama-2, Llama-3 | ✅ Verified |
+| 2 | Qwen3 | Qwen3-0.6B ~ 32B | ✅ Verified |
+
+> **Note:** Support for additional architectures is planned. Contributions and test reports are welcome.
+
+## 🔧 Installation
+
+### for users (pip)
+
+#### 1. Install PyTorch
+
+Please install the appropriate version of PyTorch.
+
+#### ✅ CPU-only
+```bash
+pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+```
+
+#### ✅ CUDA-enabled
+
+Choose the appropriate CUDA version for your system:
+
+| CUDA Version | Installation Command |
+|--------------|------------------------|
+| CUDA 11.8    | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118` |
+| CUDA 12.1    | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121` |
+| CUDA 12.4    | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124` |
+| CUDA 12.6    | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126` |
+| CUDA 12.8    | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128` |
+
+Check your CUDA version:
+```bash
+nvcc --version
+```
+
+or
+```bash
+nvidia-smi
+```
+
+Verify PyTorch GPU support:
+```python
+import torch
+print(torch.cuda.is_available())
+```
+
+#### 2. Install `onecomp`
+
+Once PyTorch is installed, you can install `onecomp`:
+
+```bash
+pip install onecomp
+```
+
+### for developers (uv : recommended)
+
+#### Install `uv`
+
+[`uv`](https://docs.astral.sh/uv/getting-started/installation/) is a fast Python package and project manager written in Rust.
+It offers a drop-in replacement for pip and pip-tools while also managing virtual environments and Python installations.
+With its Rust-based dependency resolver and the `uv.lock` lockfile, uv provides deterministic and reproducible environments across development machines and CI pipelines.
+
+```bash
+# install uv (for macOS or Linux)
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+git clone <git repository URL>
+cd OneCompression
+uv sync --extra cu128 --extra dev
+```
+
+The `uv sync` command creates a Python virtual environment and installs all dependent libraries.
+
+The `--extra cu128` option installs the CUDA-enabled version of PyTorch (along with `torchvision` from the same CUDA index).
+Replace `cu128` with the appropriate variant for your environment: `cpu`, `cu118`, `cu121`, `cu124`, `cu126`, or `cu128`.
+PyTorch will be automatically downloaded by `uv`, so you do not need to install it beforehand.
+
+Adding `--extra dev` installs additional packages for development.
+
+To use vLLM for serving quantized models, add `--extra vllm`:
+
+```bash
+uv sync --extra cu128 --extra dev --extra vllm
+```
+
+> **Note:** `--extra vllm` may take a long time on the first run if a pre-built `xformers` wheel is not available for your Python/CUDA combination (e.g. Python 3.13). Using Python 3.12 typically avoids this.
+
+#### Running commands (uv environment)
+
+In the environment created by `uv sync`, you can run commands in two ways:
+
+##### Option 1: Use `uv run` (no activation needed)
+
+```bash
+uv run pytest tests/ -v
+uv run python example/example1.py
+uv run black --check onecomp/
+```
+
+##### Option 2: Activate the virtual environment (traditional approach)
+
+```bash
+source .venv/bin/activate
+pytest tests/ -v
+python example/example1.py
+black --check onecomp/
+```
+
+### for developers (pip)
+
+```bash
+git clone <git repository URL>
+cd OneCompression
+
+# First, install PyTorch with CUDA support for your environment
+pip install torch --index-url https://download.pytorch.org/whl/cu128
+# Then install onecomp with development dependencies
+pip install -e ".[dev]"
+```
+
+Replace `cu128` with the appropriate variant for your environment: `cpu`, `cu118`, `cu121`, `cu124`, `cu126`, or `cu128`.
+
+
+### Building Documentation Locally
+
+```bash
+uv sync --extra cu128 --extra dev --extra docs
+uv run mkdocs serve
+```
+
+Then open [http://127.0.0.1:8000](http://127.0.0.1:8000) in your browser.
+
+## 🚀 Examples
+
+| Category | Script | Description |
+|----------|--------|-------------|
+| Quantization | [example_gptq.py](./example/example_gptq.py) | GPTQ quantization |
+| | [example_qep_gptq.py](./example/example_qep_gptq.py) | GPTQ + QEP (error propagation) |
+| | [example_jointq.py](./example/example_jointq.py) | JointQ quantization |
+| | [example_autobit.py](./example/example_autobit.py) | AutoBit mixed-precision quantization |
+| | [example_auto_run.py](./example/example_auto_run.py) | AutoBit with automatic VRAM estimation |
+| Save / Load | [example_save_load.py](./example/example_save_load.py) | Save and load quantized models |
+| Rotation Preprocessing | [example_llama_preprocess_rtn.py](./example/pre_process/example_llama_preprocess_rtn.py) | Rotation preprocessing + RTN (TinyLlama) |
+| | [example_preprocess_save_load.py](./example/pre_process/example_preprocess_save_load.py) | Save and load rotation-preprocessed quantized models |
+| Post-Process | [example_lora_sft.py](./example/post_process/example_lora_sft.py) | LoRA SFT post-quantization fine-tuning |
+| | [example_lora_sft_knowledge.py](./example/post_process/example_lora_sft_knowledge.py) | LoRA SFT knowledge injection |
+| vLLM | [example_gptq_vllm_inference.py](./example/vllm_inference/example_gptq_vllm_inference.py) | GPTQ + QEP quantization and vLLM inference |
+| | [example_autobit_vllm_inference.py](./example/vllm_inference/example_autobit_vllm_inference.py) | AutoBit quantization and vLLM inference |
+
+## 🔌 vLLM Inference
+
+OneComp-quantized models can be served with [vLLM](https://docs.vllm.ai/) via built-in plugins (DBF, Mixed-GPTQ).
+
+```bash
+# uv users
+uv sync --extra cu128 --extra vllm
+
+# pip users
+pip install vllm
+```
+
+See the [vLLM Inference guide](https://FujitsuResearch.github.io/OneCompression/user-guide/vllm-inference/) for details.
+
+
+## 📄 License
+
+See [LICENSE](./LICENSE) for more details.
+
+## Citation
+
+OneComp technical report (coming soon on ArXiv):
+
+```
+@misc{onecomp2026,
+  title={TBD},
+  author={TBD},
+  year={2026},
+  note={arXiv preprint coming soon}
+}
+```
+
+QEP (Quantization Error Propagation):
+
+```
+@inproceedings{
+arai2025quantization,
+title={Quantization Error Propagation: Revisiting Layer-Wise Post-Training Quantization},
+author={Yamato Arai and Yuma Ichikawa},
+booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
+year={2025},
+url={https://openreview.net/forum?id=a3l3K9khbL}
+}
+```

onecomp-1.0.0/README.md

@@ -0,0 +1,218 @@
+# Fujitsu One Compression
+
+Fujitsu One Compression (OneComp) is a Python package for LLM compression.
+
+## 📖 Documentation
+
+Full documentation is available at **[https://FujitsuResearch.github.io/OneCompression/](https://FujitsuResearch.github.io/OneCompression/)**.
+
+## 📦 Features
+
+- **Quantization Error Propagation (QEP)**: A post-training quantization method that corrects quantization errors by propagating them to subsequent layers, improving the accuracy of quantized LLMs. See [Arai & Ichikawa, NeurIPS 2025](https://openreview.net/forum?id=a3l3K9khbL) for details. The original reference implementation is available at [FujitsuResearch/qep](https://github.com/FujitsuResearch/qep).
+- **vLLM Plugin Integration**: Serve OneComp-quantized models with [vLLM](https://docs.vllm.ai/) via built-in plugins for DBF and Mixed-GPTQ quantization methods.
+- **AutoBit**: Mixed-precision quantization with ILP-based bitwidth assignment. Automatically estimates the target bitwidth from available VRAM and assigns per-layer bitwidths to minimize quantization error under the memory budget.
+- **JointQ**: Joint quantization method that optimizes weight assignments and scale parameters simultaneously for improved quantization accuracy. Supports group-wise quantization (e.g., 4-bit, groupsize=128).
+- **LoRA SFT Post-Process**: Fine-tune quantized models with LoRA adapters for accuracy recovery or domain-specific knowledge injection. Supports SFT loss, teacher distillation, and intermediate block alignment.
+- **Rotation Preprocessing**: SpinQuant/OstQuant-based rotation preprocessing that reduces quantization error by learning optimal rotation matrices before quantization. Rotation/scaling matrices are absorbed into model weights, with online Hadamard hooks automatically registered at load time. Supports Llama and Qwen3 architectures.
+- (TBD)
+
+## 🤖 Supported Models
+
+OneComp has been verified with the following model architectures.
+Other Hugging Face-compatible models may work but are currently untested.
+
+| # | Architecture | Verified Models | Status |
+|---|-------------|-----------------|--------|
+| 1 | Llama | TinyLlama, Llama-2, Llama-3 | ✅ Verified |
+| 2 | Qwen3 | Qwen3-0.6B ~ 32B | ✅ Verified |
+
+> **Note:** Support for additional architectures is planned. Contributions and test reports are welcome.
+
+## 🔧 Installation
+
+### for users (pip)
+
+#### 1. Install PyTorch
+
+Please install the appropriate version of PyTorch.
+
+#### ✅ CPU-only
+```bash
+pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+```
+
+#### ✅ CUDA-enabled
+
+Choose the appropriate CUDA version for your system:
+
+| CUDA Version | Installation Command |
+|--------------|------------------------|
+| CUDA 11.8    | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118` |
+| CUDA 12.1    | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121` |
+| CUDA 12.4    | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124` |
+| CUDA 12.6    | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126` |
+| CUDA 12.8    | `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128` |
+
+Check your CUDA version:
+```bash
+nvcc --version
+```
+
+or
+```bash
+nvidia-smi
+```
+
+Verify PyTorch GPU support:
+```python
+import torch
+print(torch.cuda.is_available())
+```
+
+#### 2. Install `onecomp`
+
+Once PyTorch is installed, you can install `onecomp`:
+
+```bash
+pip install onecomp
+```
+
+### for developers (uv : recommended)
+
+#### Install `uv`
+
+[`uv`](https://docs.astral.sh/uv/getting-started/installation/) is a fast Python package and project manager written in Rust.
+It offers a drop-in replacement for pip and pip-tools while also managing virtual environments and Python installations.
+With its Rust-based dependency resolver and the `uv.lock` lockfile, uv provides deterministic and reproducible environments across development machines and CI pipelines.
+
+```bash
+# install uv (for macOS or Linux)
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+git clone <git repository URL>
+cd OneCompression
+uv sync --extra cu128 --extra dev
+```
+
+The `uv sync` command creates a Python virtual environment and installs all dependent libraries.
+
+The `--extra cu128` option installs the CUDA-enabled version of PyTorch (along with `torchvision` from the same CUDA index).
+Replace `cu128` with the appropriate variant for your environment: `cpu`, `cu118`, `cu121`, `cu124`, `cu126`, or `cu128`.
+PyTorch will be automatically downloaded by `uv`, so you do not need to install it beforehand.
+
+Adding `--extra dev` installs additional packages for development.
+
+To use vLLM for serving quantized models, add `--extra vllm`:
+
+```bash
+uv sync --extra cu128 --extra dev --extra vllm
+```
+
+> **Note:** `--extra vllm` may take a long time on the first run if a pre-built `xformers` wheel is not available for your Python/CUDA combination (e.g. Python 3.13). Using Python 3.12 typically avoids this.
+
+#### Running commands (uv environment)
+
+In the environment created by `uv sync`, you can run commands in two ways:
+
+##### Option 1: Use `uv run` (no activation needed)
+
+```bash
+uv run pytest tests/ -v
+uv run python example/example1.py
+uv run black --check onecomp/
+```
+
+##### Option 2: Activate the virtual environment (traditional approach)
+
+```bash
+source .venv/bin/activate
+pytest tests/ -v
+python example/example1.py
+black --check onecomp/
+```
+
+### for developers (pip)
+
+```bash
+git clone <git repository URL>
+cd OneCompression
+
+# First, install PyTorch with CUDA support for your environment
+pip install torch --index-url https://download.pytorch.org/whl/cu128
+# Then install onecomp with development dependencies
+pip install -e ".[dev]"
+```
+
+Replace `cu128` with the appropriate variant for your environment: `cpu`, `cu118`, `cu121`, `cu124`, `cu126`, or `cu128`.
+
+
+### Building Documentation Locally
+
+```bash
+uv sync --extra cu128 --extra dev --extra docs
+uv run mkdocs serve
+```
+
+Then open [http://127.0.0.1:8000](http://127.0.0.1:8000) in your browser.
+
+## 🚀 Examples
+
+| Category | Script | Description |
+|----------|--------|-------------|
+| Quantization | [example_gptq.py](./example/example_gptq.py) | GPTQ quantization |
+| | [example_qep_gptq.py](./example/example_qep_gptq.py) | GPTQ + QEP (error propagation) |
+| | [example_jointq.py](./example/example_jointq.py) | JointQ quantization |
+| | [example_autobit.py](./example/example_autobit.py) | AutoBit mixed-precision quantization |
+| | [example_auto_run.py](./example/example_auto_run.py) | AutoBit with automatic VRAM estimation |
+| Save / Load | [example_save_load.py](./example/example_save_load.py) | Save and load quantized models |
+| Rotation Preprocessing | [example_llama_preprocess_rtn.py](./example/pre_process/example_llama_preprocess_rtn.py) | Rotation preprocessing + RTN (TinyLlama) |
+| | [example_preprocess_save_load.py](./example/pre_process/example_preprocess_save_load.py) | Save and load rotation-preprocessed quantized models |
+| Post-Process | [example_lora_sft.py](./example/post_process/example_lora_sft.py) | LoRA SFT post-quantization fine-tuning |
+| | [example_lora_sft_knowledge.py](./example/post_process/example_lora_sft_knowledge.py) | LoRA SFT knowledge injection |
+| vLLM | [example_gptq_vllm_inference.py](./example/vllm_inference/example_gptq_vllm_inference.py) | GPTQ + QEP quantization and vLLM inference |
+| | [example_autobit_vllm_inference.py](./example/vllm_inference/example_autobit_vllm_inference.py) | AutoBit quantization and vLLM inference |
+
+## 🔌 vLLM Inference
+
+OneComp-quantized models can be served with [vLLM](https://docs.vllm.ai/) via built-in plugins (DBF, Mixed-GPTQ).
+
+```bash
+# uv users
+uv sync --extra cu128 --extra vllm
+
+# pip users
+pip install vllm
+```
+
+See the [vLLM Inference guide](https://FujitsuResearch.github.io/OneCompression/user-guide/vllm-inference/) for details.
+
+
+## 📄 License
+
+See [LICENSE](./LICENSE) for more details.
+
+## Citation
+
+OneComp technical report (coming soon on ArXiv):
+
+```
+@misc{onecomp2026,
+  title={TBD},
+  author={TBD},
+  year={2026},
+  note={arXiv preprint coming soon}
+}
+```
+
+QEP (Quantization Error Propagation):
+
+```
+@inproceedings{
+arai2025quantization,
+title={Quantization Error Propagation: Revisiting Layer-Wise Post-Training Quantization},
+author={Yamato Arai and Yuma Ichikawa},
+booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
+year={2025},
+url={https://openreview.net/forum?id=a3l3K9khbL}
+}
+```

onecomp-1.0.0/benchmark/llama3-8b-gptq/quant_benchmark.py

@@ -0,0 +1,95 @@
+"""GPTQ Benchmark
+
+Run GPTQ for all combinations of bits × group_size in a single pass.
+Shares calibration data accumulation across quantizers for efficiency.
+Results are saved under output_dir.
+
+Copyright 2025-2026 Fujitsu Ltd.
+
+Usage:
+    python quant_benchmark.py
+"""
+
+import itertools
+
+import hydra
+from omegaconf import DictConfig, OmegaConf
+
+from onecomp import GPTQ, ModelConfig, Runner
+
+
+def create_quantizers(cfg: DictConfig):
+    """Create a list of GPTQ quantizers for all combinations of bits × group_size."""
+    quantizers = []
+    sym = cfg.gptq.symmetric
+    sym_label = "sym" if sym else "asym"
+
+    for bits, gs in itertools.product(cfg.gptq.bits, cfg.gptq.group_size):
+        # Label strings
+        gs_label = "pc" if gs is None else f"gs{gs}"
+
+        # GPTQ: groupsize=-1 means per-channel
+        gptq_groupsize = -1 if gs is None else gs
+        quantizers.append(
+            GPTQ(
+                num_layers=cfg.gptq.num_layers,
+                wbits=bits,
+                sym=sym,
+                groupsize=gptq_groupsize,
+                blocksize=cfg.gptq.blocksize,
+                percdamp=cfg.gptq.percdamp,
+                actorder=cfg.gptq.actorder,
+                mse=cfg.gptq.mse,
+                q_grid=cfg.gptq.q_grid,
+                q_norm=cfg.gptq.q_norm,
+                calc_quant_error=True,
+                name=f"GPTQ_{bits}bit_{gs_label}_{sym_label}",
+            )
+        )
+
+    return quantizers
+
+
+@hydra.main(version_base=None, config_path="conf", config_name="benchmark_llama3-8b")
+def main(cfg: DictConfig):
+    print(OmegaConf.to_yaml(cfg))
+
+    model_config = ModelConfig(path=cfg.model_path, device=cfg.model_device)
+
+    quantizers = create_quantizers(cfg)
+
+    print(f"Number of quantizers: {len(quantizers)}")
+    for q in quantizers:
+        print(f"  - {q.name}")
+
+    # Build Runner
+    runner = Runner(
+        model_config=model_config,
+        quantizers=quantizers,
+        max_length=cfg.max_length,
+        num_calibration_samples=cfg.num_calibration_samples,
+        calibration_strategy=cfg.calibration_strategy,
+        calibration_seed=cfg.calibration_seed,
+        calibration_batch_size=cfg.calibration_batch_size,
+    )
+
+    # Run quantization
+    runner.run()
+
+    # Save results
+    for q in quantizers:
+        runner.save_quantization_statistics(
+            f"quantization_statistics_{q.name}.json", quantizer=q
+        )
+
+    # Perplexity evaluation
+    if cfg.calc_ppl:
+        runner.benchmark_perplexity(original_model=cfg.calc_original_ppl)
+
+    # Accuracy evaluation
+    if cfg.calc_acc:
+        runner.benchmark_accuracy(original_model=cfg.calc_original_acc)
+
+
+if __name__ == "__main__":
+    main()