(ICML26) TileQ: Efficient Low-Rank Quantization of Mixture-of-Experts with 2D Tiling

TileQ is a low-rank quantization method of Mixture-of-Experts (MoE) models by 2D-Tiling.

Installation

Prerequisites

• Python 3.10 or higher
• CUDA 12.4 compatible environment
• Sufficient GPU memory for model processing

Setup

cd TileQ
pip install -r requirements.txt

Tips: The requirement list may not be complete. You may first install the latest version of gptqmodel, which includes most of our required libraries.

Quick Start

1. Prepare your model: Ensure your HuggingFace-format model is available locally or specify the path.

2. Run quantization:

   cd example
   # Edit tileq.sh to set your model paths and parameters
   bash tileq.sh

   This will quantize the specified models, perform low-rank decomposition, and run evaluations on perplexity text of wikitext2 and 6 down-stream tasks.

Usage

Command Line Arguments

The main quantization script tileq/run_quantize.py accepts the following main arguments:

• --model_path: Path to the input model directory
• --output_path: Directory to save the quantized model
• --qbit: Quantization bit width (2, 3, 4, etc.)
• --groupsize: Group size for quantization (default: 128)
• --fix_rank: Fixed rank for low-rank approximation
• --lora_bit: Bit width for low-rank components (8 or 16)
• --method: Quantization method ('gptq' or 'gptvq')
• --tile_row: Row size for 2D tiling

Example Usage

python tileq/run_quantize.py \
    --model_path /path/to/model \
    --output_path /path/to/output \
    --qbit 2 \
    --groupsize 128 \
    --fix_rank 32 \
    --method gptq \
    --tile_row 8

Inference and Benchmarking

Use the scripts in the Inference/ directory for evaluation on inference latency of MoE-MLP block:

# Run inference benchmark
python Inference/inference_bench.py

Algorithm in Codes

As detailed in Section 3, the TileQ algorithm comprises two main components: the Quantization phase and the Inference phase.

Quantization

Get Low-Rank in 2D-Tiling

• Algo 1, lines 1–3:
  run_quantize.py, lines 175-182
• Algo 1, lines 4–7:
  get_scale_quant.py, lines 175-213
• Algo 1, lines 8–13:
  get_scale_quant.py, lines 43-85
• Algo 1, lines 14–16:
  get_scale_quant.py, lines 91-108
• Algo 1, lines 17-18:
  get_scale_quant.py, lines 120-163

Quantize the residual matrix (Take scalar quantization as an example)

• Algo 1, lines 19:
  get_scale_quant.py, lines 226
• Algo 1, lines 22-24:
  scalar_quant_utils_lora.py, lines 83-88
• Algo 1, lines 22-25:
  scalar_quant_utils_lora.py, lines 83-89
• Algo 1, lines 26:
  scalar_quant_utils_lora.py, lines 310

Inference

Inference in TileQ are concise and highly efficient. We also provide a vanilla implementation without applying the tricks from Section 3."

• The optimized implementation: tileq_mlp.py, lines 71-100
• The original implementation: tileq_mlp.py, lines 24-68

Directory Structure

TileQ/
├── tileq/                    # Core quantization modules
│   ├── quantizer/           # Quantization utilities
│   │   ├── quantizer.py
│   │   ├── quantizer_v.py
│   │   └── vector_quant_utils_lora.py
│   ├── sketch/              # Low-rank decomposition
│   │   └── r1_sketch.py
│   └── utils/               # General utilities
│       ├── eval_util.py
│       ├── get_calib_data.py
│       └── ...
├── eval/                    # Evaluation scripts
│   ├── evaluation.py
│   └── get_tokens.py
├── Inference/               # Inference and benchmarking
│   ├── inference_bench.py
│   └── tileq_mlp.py
├── example/                 # Example scripts
│   └── tileq.sh
└── README.md               # This file

License

This project is licensed under the MIT License - see the LICENSE file for details.