* Equal Contribution ✉ Corresponding author
Mixture-of-Experts (MoE) large language models (LLMs) suffer significant inference and storage overheads, which makes extremely low-bit quantization (sub 2-bit) highly desirable. However, existing methods for dense LLMs typically quantize each expert independently and require it to approximate the original mapping over the entire input space. We argue that this formulation is overly conservative and structurally mismatched for MoE: 1) experts in the same MoE layer are not fully independent but operate on a shared hidden representation space. Under low-bit budgets, independently binarizing each expert will distort the activation-sensitive input geometry, which leads to severe performance degradation, and 2) due to routing-induced specialization, each expert only serves a concentrated subset of tokens, whose activations typically occupy a much narrower subspace than that of dense layers. More importantly, when the weights approach 1-bit, the dominant deployment bottleneck is no longer the weights themselves, but the metadata (e.g., scaling factors, bitmap, grouping index) which can raise the real effective inference precision to around 3-bit. This issue, however, has largely been overlooked in prior works. Building on these perspectives, we introduce a novel extremely low-bit quantization framework for MoE-LLMs, called TriS-MoE, which considers Structure, Subspace and System. Specifically, we first extract a shared high-precision input-side backbone to preserve activation-sensitive input geometry, while binarizing the remaining components to save costs. Secondly, we propose to redirect expert quantization errors into the null space of routed activations to minimize output perturbation on the subspace actually used by each expert. Finally, we co-design a bitmap compression method based on Golomb-Rice coding which approaches the Shannon entropy lower bound and a specialized streaming loading mechanism to reduce metadata overhead during inference. Extensive experiments on MoE-LLMs demonstrate our TriS-MoE outperforms the strongest baseline by 8.31% average accuracies, and also achieves an average 6.32x reduction in inference memory on real systems, enabling Qwen3.5-35B-A3B and Mixtral-8x7B to run on a single consumer-GPU.
- Environment setup
- Kernel compilation
- Quantization
- Checkpoint export: block-GR and streaming
- Benchmark Evaluation
- Acknowledgements
- License
Use the default environment for the OLMoE example below. For Qwen3.5, use the
separate TriS_Qwen35 environment instead.
Use Python 3.12 with the TriSMoE Conda environment. If the environment already
exists, activate it instead of creating it again.
conda create -n TriSMoE python=3.12 pip setuptools wheel -y
conda activate TriSMoEInstall PyTorch from the official CUDA 11.8 wheel index, then install the remaining core dependencies:
python -m pip install "torch==2.6.0+cu118" \
--index-url https://download.pytorch.org/whl/cu118
python -m pip install \
"transformers==4.55.0" \
"accelerate==1.13.0" \
"datasets==4.8.4" \
"numpy==2.4.3" \
"scipy==1.17.1" \
"safetensors==0.7.0" \
"huggingface-hub==0.36.2" \
"packaging==26.0" \
"tqdm==4.67.3" \
"pyparsing==3.3.2" \
"sentencepiece==0.2.1" \
"protobuf==7.34.1"The following core dependency versions match the TriS_Qwen35 environment. It
uses Python 3.12, transformers==5.7.0, and huggingface-hub==1.12.0 for Qwen3.5
support. If this environment already exists, activate it instead of creating it
again.
conda create -n TriS_Qwen35 python=3.12 pip setuptools wheel -y
conda activate TriS_Qwen35
python -m pip install "torch==2.6.0+cu118" \
--index-url https://download.pytorch.org/whl/cu118
python -m pip install \
"transformers==5.7.0" \
"accelerate==1.13.0" \
"datasets==4.8.4" \
"numpy==2.4.3" \
"scipy==1.17.1" \
"safetensors==0.7.0" \
"huggingface-hub==1.12.0" \
"packaging==26.0" \
"tqdm==4.67.3" \
"pyparsing==3.3.2" \
"sentencepiece==0.2.1" \
"protobuf==7.34.1"Both environments use PyTorch 2.6.0 with CUDA 11.8 and share the kernel compilation instructions below.
Use the bundled, modified AutoGPTQ source rather than a separate PyPI installation:
export REPO_ROOT="$PWD"
export PYTHONPATH="$REPO_ROOT/AutoGPTQ:$REPO_ROOT${PYTHONPATH:+:$PYTHONPATH}"Repeat these exports from the repository root when starting a new shell.
Install a complete CUDA 11.8 toolkit, including nvcc, headers, and libraries,
and a compatible host compiler, such as GCC/G++ 11.
Point CUDA_HOME to your toolkit installation. The path below is an example;
TORCH_CUDA_ARCH_LIST=8.0 targets an NVIDIA A100 and should be adjusted for your GPU.
The build uses g++ from PATH.
export CUDA_HOME=/path/to/cuda-11.8
export PATH="$CUDA_HOME/bin:$PATH"
export TORCH_CUDA_ARCH_LIST="8.0"
nvcc --version
python -c 'import torch; print("PyTorch:", torch.__version__, "CUDA:", torch.version.cuda)'
(
cd "$REPO_ROOT/AutoGPTQ"
BUILD_TRIS_ONLY=1 python setup.py build_ext --inplace
)Before building, check that both nvcc and torch.version.cuda report CUDA 11.8.
BUILD_TRIS_ONLY=1 builds only autogptq_TriS_cuda, skipping unrelated AutoGPTQ
extensions. Rebuild after changing the CUDA sources.
Check that the compiled extension can be imported; the current kernel ABI is 2:
python - <<'PY'
import torch
import autogptq_TriS_cuda as kernel
print("Extension:", kernel.__file__)
print("Kernel ABI:", kernel.TRIS_COMPACT3_KERNEL_ABI_VERSION)
PYSet MODEL_PATH to a local model directory or a Hugging Face model ID. This OLMoE
example uses WikiText-2 calibration, fit_nohead rank selection, and vector
correction. It keeps router weights in their original precision, evaluates
perplexity, and saves a real-quantized checkpoint for export.
MODEL_PATH="allenai/OLMoE-1B-7B-0125"
QUANT_DIR="checkpoints/olmoe-quant"
CUDA_VISIBLE_DEVICES=0 TRIS_COMPACT3_BLOCK_GR=0 \
python tris_main.py "$MODEL_PATH" wikitext2 \
--device cuda:0 \
--nsamples 128 \
--seqlen 4096 \
--seed 0 \
--blocksize 128 \
--percdamp 0.01 \
--idx_rule fit_nohead \
--ns_rule vector \
--quant_router false \
--real_quant true \
--save "$QUANT_DIR" \
--eval--real_quant truepacks weights into the deployment representation.--savewrites the checkpoint as safetensors; the default shard limit is5GB, configurable with--save_max_shard_size.--evalevaluates WikiText-2 perplexity after quantization. Results and run configuration are written to the console andruntime_log/.TRIS_COMPACT3_BLOCK_GR=0leaves block-GR conversion to the export step, so the checkpoint can serve as the source for both output formats below.- For fake-quantization and PPL evaluation without saving, change
--real_quant trueto--real_quant falseand omit--save "$QUANT_DIR". Fake-quantized floating-point weights are not an input to the export workflow.
Calibration data is cached under cache/ by default; set TRIS_EVAL_CACHE_DIR
to use another location. The first run needs access to the model and dataset,
either online or through an existing local cache.
First build a grouped sparse checkpoint from the real-quantized checkpoint. Then export two separate variants from that same sparse source:
Real-quantized checkpoint
└── Grouped sparse checkpoint
├── block-GR checkpoint
└── Streaming sparse checkpoint
Use a different, empty output directory for each step. The initial grouped export
accepts sharded checkpoints; each export below writes a single
grouped_runtime_payload.safetensors file and copies the model configuration.
Keep the complete output directory when moving a checkpoint.
This step loads the quantized model onto the selected GPU and builds grouped expert payloads. Keep block-GR and CPU staging disabled while creating the sparse source.
GROUPED_DIR="checkpoints/olmoe-grouped-sparse"
BLOCK_GR_DIR="checkpoints/olmoe-block-gr"
STREAMING_DIR="checkpoints/olmoe-streaming"
CUDA_VISIBLE_DEVICES=0 \
TRIS_COMPACT3_BLOCK_GR=0 TRIS_GROUPED_QGMAP_CPU_STAGING=0 \
python tools/export_grouped_sparse_checkpoint.py \
--source-dir "$QUANT_DIR" \
--output-dir "$GROUPED_DIR" \
--device cuda:0python tools/export_grouped_blockgr_checkpoint.py \
--source-dir "$GROUPED_DIR" \
--output-dir "$BLOCK_GR_DIR" \
--blockgr-mode compressedThis produces block-wise Golomb-Rice maps for grouped expert projections and
non-grouped compact3 modules. By default, the exported checkpoint removes the
original sparse maps; add --keep-sparse only if you also need those tensors in
the block-GR output. The grouped sparse source is unchanged.
python tools/export_grouped_streaming_sparse_checkpoint.py \
--source-dir "$GROUPED_DIR" \
--output-dir "$STREAMING_DIR"Use GROUPED_DIR, not BLOCK_GR_DIR, as the source. The streaming exporter
requires sparse maps and does not convert block-GR maps back to sparse form.
Both final export steps operate on CPU; allow enough host memory and disk space
to read and write the complete checkpoint.
Before loading either exported checkpoint with the bundled AutoGPTQ runtime, set:
export TRIS_GROUPED_EXPERTS=1
export TRIS_DEPLOYMENT_ONLY=1Select the format-specific environment values:
| Environment variable | block-GR | Streaming sparse |
|---|---|---|
TRIS_COMPACT3_BLOCK_GR |
1 |
0 |
TRIS_GROUPED_QGMAP_CPU_STAGING |
0 |
1 |
Streaming keeps sparse map bitmaps and prefix data in pinned CPU memory and stages them to the GPU for execution. It does not offload all model weights or the KV cache. Exporting the streaming checkpoint alone does not enable streaming; set the runtime variables before model loading.
We use lm-evaluation-harness to evaluate performance on QA datasets.
Use ./tools/benchmark_e2e_grid.py as an example for model loading and inferencing, we have only supported OLMoE for fast decoding yet.
Parts of this codebase are adapted from AutoGPTQ (https://github.com/AutoGPTQ/AutoGPTQ) and ARB-LLM (https://github.com/ZHITENGLI/ARB-LLM). We sincerely thank them for their outstanding contributions.
TriS-MoE is licensed under the Apache License 2.0. Third-party code retains its original licenses; see AutoGPTQ/LICENSE.