Upload 17 files

c675b8d verified 2 days ago

6.82 kB

# shapequant

Data-free, training-free quantization of Hugging Face causal language models.

No calibration text and no fine-tuning are used. The method builds each layer's input covariance from the weights
themselves and uses it to place rounding error where the inputs don't reach (noise-shaped rounding). The output is a
standard **GPTQ-format** checkpoint (INT4/INT8, asymmetric, group-wise), so GPTQ-capable stacks can serve it.
A dependency-light loader is included.

Supported `model_type`s:
* `llama`, `mistral`, `qwen2`: includes Llama-2/3/3.1/3.3-70B and Qwen2.5-72B.
* `opt` (pre-LN), `gpt2`.

## Install

```bash
pip install -e .            # torch, transformers, accelerate, safetensors, huggingface_hub
```

## Python API

```python
from shapequant import quantize_model, load_quantized, perplexity, wikitext2_windows, estimate_memory
# uniform INT4, groups of 128 (default method: data-free shaped rounding)
quantize_model("meta-llama/Llama-3.3-70B-Instruct", "llama70b-sq-int4", bits=4, group_size=128, device="cuda")
# mixed INT4/INT8/FP16 at 4.75 bits/weight (type rule: V, attention-output and MLP-output projections weighted 4x)
quantize_model("HuggingFaceTB/SmolLM2-135M", "smol-sq-mixed", budget=4.75, group_size=64)
# baseline on the same grid
quantize_model("facebook/opt-125m", "opt-rtn-int4", bits=4, method="rtn")
model = load_quantized("smol-sq-mixed", device="cpu")                 # or device_map="auto"
```

Key arguments of `quantize_model(model, out_dir, ...)`:

| argument | default | meaning |
|---|---|---|
| `bits` | 4 | 4 or 8, uniform |
| `group_size` | 128 | must divide every input dimension (use 64 for hidden size 576, etc.) |
| `method` | `"shaped"` | `"shaped"` (data-free covariance) or `"rtn"` (round to nearest, same grid) |
| `budget` | None | average bits/weight for mixed INT4/INT8/FP16 (must be ≥ INT4 cost, 4 + 32/group) |
| `lam`, `kappa` | 1.0, 1.0 | design H = I + lam·S; kappa = attention concentration (1 = theory value) |
| `n_samples` | 16384 | Monte Carlo samples for the MLP-input covariance |
| `dynamic_groups` | True | fit each group's grid to the error-fed weights (standard GPTQ); False = grids from original weights |
| `device` | cuda if available | where the per-layer work runs |

## CLI

```bash
shapequant memory meta-llama/Llama-3.1-70B          # peak-memory estimate from the config only
shapequant quantize meta-llama/Llama-3.1-70B out/ --bits 4 --group-size 128 --device cuda
shapequant quantize HuggingFaceTB/SmolLM2-135M out/ --budget 5.0 --group-size 64
shapequant eval out/ --windows 32                   # WikiText-2 perplexity (needs `datasets`) or --text file.txt
```

## How it scales to 70B

* **Streaming.** Weights are read tensor by tensor from the safetensors shards. Only one decoder layer is ever
  resident, and it is written out before the next one is read. Host RAM is not proportional to model size.
* **Per-layer device work.**
  * The residual-stream second moment (d×d) and the norm-output moments are updated layer by layer.
  * Monte Carlo of the layer's own MLP nonlinearity gives the (inter×inter) MLP-input covariance.
  * Error-feedback rounding uses lazy block updates (block 128). Modules that read the same input
    (q/k/v, gate/up) are stacked and rounded in one pass.
* **Peak device memory** is about one layer in fp32 plus two (in×in) fp32 matrices for the largest input.
  For Llama-3.1-70B the largest input is the 28672-wide MLP input. `shapequant memory <model>` prints the estimate.
* **Output.** Sharded `model-0000x-of-0000y.safetensors` (4 GB shards by default) + index + `config.json`
  with `quantization_config` (`quant_method: gptq`, `checkpoint_format: gptq`, per-module overrides under `dynamic`).
## Checkpoint format

The checkpoint is GPTQ v1, as written by AutoGPTQ/GPTQModel:
* `qweight` int32 (in·b/32, out); `qzeros` int32 (G, out·b/32), stored as z − 1; `scales` fp16 (G, out);
  `g_idx` = arange(in)//group; `sym=False`, `desc_act=False`.
* Zero-points are kept ≥ 1, so the v1 "z − 1" convention never underflows.
* Mixed-precision modules are listed in `quantization_config["dynamic"]`:
  * `"+:<regex>": {"bits": 8}` for INT8 modules;
  * `"-:<regex>": {}` for modules left in FP16 (GPTQModel convention).
* Everything else (embeddings, norms, untied lm_head) is stored in fp16.
## Results
`RESULTS.md` has the full numbers. Summary (WikiText-2, 32 × 1024 tokens; perplexity, lower is better):
| model | dense | RTN INT4 | shaped INT4 | RTN mixed | shaped mixed |
|---|---|---|---|---|---|
| GPT-2 124M (g128, mixed 4.75 b/w) | 26.84 | 29.91 | **29.14** | 29.62 | **28.86** |
| OPT-125m (g128, mixed 4.75 b/w) | 29.27 | 31.98 | **31.39** | 30.71 | **30.44** |
| SmolLM2-135M (g64, mixed 5.0 b/w) | 15.46 | **19.28** | 19.29 | **17.82** | 17.83 |
At standard INT4 the gain over RTN is modest (−2.6%, −1.9%, +0.1%). It was much larger at 3 bits in the research
experiments, but 3-bit is not a standard packed format and is not exported here.
Reproduce with `python tests/real_models.py <model> <out> --group 128 --budget 4.75`.

## Verified here, and not

* **Verified:**
  * unit tests and tiny-model round trips for llama / qwen2 / opt / gpt2;
  * real GPT-2, OPT-125m and SmolLM2-135M runs (quantize → save → load → perplexity) on CPU;
  * memory estimates from the real 70B configs;
  * the attention-sized per-layer steps at full Llama-3.1-70B shapes on CPU.
* **Not verified:**
  * an end-to-end 70B run (needs a GPU);
  * the MLP-sized 70B steps (need ~7 GB for two 28672² matrices);
  * loading the checkpoints with GPTQModel / vLLM kernels. The format follows the GPTQ v1 layout and the bundled
    loader reads it, but no GPU kernel was run here.
## Method

For a linear block y = xW with quantization error R, the output error is tr(RᵀΣR) with Σ = E[xxᵀ].
No data is available, so Σ is built from the model:
* **LayerNorm/RMSNorm inputs:** token embeddings under a Zipf prior on token id (plus positions), then every
  earlier attention-output and MLP-output projection added to the residual stream.
* **MLP second input:** the layer's own nonlinearity applied by Monte Carlo to that Gaussian.
* **Attention output:** a mixture of value vectors, E[ooᵀ] = μμᵀ + κ·blockdiag(C_v).
Rounding is the sequential error-feedback recursion that minimises tr(R H Rᵀ) with H = I + λS. With H = I it is
exactly round-to-nearest.
## Tests
```bash
pytest -q tests/test_core.py tests/test_end_to_end.py
```
* `test_core.py`: H = I reproduces RTN; blocked == sequential recursion; GPTQ pack/unpack is exact;
  zero-point safety; shaping lowers the H-weighted error.
* `test_end_to_end.py`: tiny Llama/Qwen2/OPT/GPT-2 models; quantize → save → load gives logits identical to the
  dequantized reference (uniform and mixed bits).

	# shapequant

	Data-free, training-free quantization of Hugging Face causal language models.

	No calibration text and no fine-tuning are used. The method builds each layer's input covariance from the weights
	themselves and uses it to place rounding error where the inputs don't reach (noise-shaped rounding). The output is a
	standard GPTQ-format checkpoint (INT4/INT8, asymmetric, group-wise), so GPTQ-capable stacks can serve it.
	A dependency-light loader is included.

	Supported `model_type`s:
	* `llama`, `mistral`, `qwen2`: includes Llama-2/3/3.1/3.3-70B and Qwen2.5-72B.
	* `opt` (pre-LN), `gpt2`.

	## Install

	```bash
	pip install -e . # torch, transformers, accelerate, safetensors, huggingface_hub
	```

	## Python API

	```python
	from shapequant import quantize_model, load_quantized, perplexity, wikitext2_windows, estimate_memory

	# uniform INT4, groups of 128 (default method: data-free shaped rounding)
	quantize_model("meta-llama/Llama-3.3-70B-Instruct", "llama70b-sq-int4", bits=4, group_size=128, device="cuda")

	# mixed INT4/INT8/FP16 at 4.75 bits/weight (type rule: V, attention-output and MLP-output projections weighted 4x)
	quantize_model("HuggingFaceTB/SmolLM2-135M", "smol-sq-mixed", budget=4.75, group_size=64)

	# baseline on the same grid
	quantize_model("facebook/opt-125m", "opt-rtn-int4", bits=4, method="rtn")

	model = load_quantized("smol-sq-mixed", device="cpu") # or device_map="auto"
	```

	Key arguments of `quantize_model(model, out_dir, ...)`:

	\| argument \| default \| meaning \|
	\|---\|---\|---\|
	\| `bits` \| 4 \| 4 or 8, uniform \|
	\| `group_size` \| 128 \| must divide every input dimension (use 64 for hidden size 576, etc.) \|
	\| `method` \| `"shaped"` \| `"shaped"` (data-free covariance) or `"rtn"` (round to nearest, same grid) \|
	\| `budget` \| None \| average bits/weight for mixed INT4/INT8/FP16 (must be ≥ INT4 cost, 4 + 32/group) \|
	\| `lam`, `kappa` \| 1.0, 1.0 \| design H = I + lam·S; kappa = attention concentration (1 = theory value) \|
	\| `n_samples` \| 16384 \| Monte Carlo samples for the MLP-input covariance \|
	\| `dynamic_groups` \| True \| fit each group's grid to the error-fed weights (standard GPTQ); False = grids from original weights \|
	\| `device` \| cuda if available \| where the per-layer work runs \|

	## CLI

	```bash
	shapequant memory meta-llama/Llama-3.1-70B # peak-memory estimate from the config only
	shapequant quantize meta-llama/Llama-3.1-70B out/ --bits 4 --group-size 128 --device cuda
	shapequant quantize HuggingFaceTB/SmolLM2-135M out/ --budget 5.0 --group-size 64
	shapequant eval out/ --windows 32 # WikiText-2 perplexity (needs `datasets`) or --text file.txt
	```

	## How it scales to 70B

	* Streaming. Weights are read tensor by tensor from the safetensors shards. Only one decoder layer is ever
	resident, and it is written out before the next one is read. Host RAM is not proportional to model size.
	* Per-layer device work.
	* The residual-stream second moment (d×d) and the norm-output moments are updated layer by layer.
	* Monte Carlo of the layer's own MLP nonlinearity gives the (inter×inter) MLP-input covariance.
	* Error-feedback rounding uses lazy block updates (block 128). Modules that read the same input
	(q/k/v, gate/up) are stacked and rounded in one pass.
	* Peak device memory is about one layer in fp32 plus two (in×in) fp32 matrices for the largest input.
	For Llama-3.1-70B the largest input is the 28672-wide MLP input. `shapequant memory <model>` prints the estimate.
	* Output. Sharded `model-0000x-of-0000y.safetensors` (4 GB shards by default) + index + `config.json`
	with `quantization_config` (`quant_method: gptq`, `checkpoint_format: gptq`, per-module overrides under `dynamic`).

	## Checkpoint format

	The checkpoint is GPTQ v1, as written by AutoGPTQ/GPTQModel:
	* `qweight` int32 (in·b/32, out); `qzeros` int32 (G, out·b/32), stored as z − 1; `scales` fp16 (G, out);
	`g_idx` = arange(in)//group; `sym=False`, `desc_act=False`.
	* Zero-points are kept ≥ 1, so the v1 "z − 1" convention never underflows.
	* Mixed-precision modules are listed in `quantization_config["dynamic"]`:
	* `"+:<regex>": {"bits": 8}` for INT8 modules;
	* `"-:<regex>": {}` for modules left in FP16 (GPTQModel convention).
	* Everything else (embeddings, norms, untied lm_head) is stored in fp16.

	## Results

	`RESULTS.md` has the full numbers. Summary (WikiText-2, 32 × 1024 tokens; perplexity, lower is better):

	\| model \| dense \| RTN INT4 \| shaped INT4 \| RTN mixed \| shaped mixed \|
	\|---\|---\|---\|---\|---\|---\|
	\| GPT-2 124M (g128, mixed 4.75 b/w) \| 26.84 \| 29.91 \| 29.14 \| 29.62 \| 28.86 \|
	\| OPT-125m (g128, mixed 4.75 b/w) \| 29.27 \| 31.98 \| 31.39 \| 30.71 \| 30.44 \|
	\| SmolLM2-135M (g64, mixed 5.0 b/w) \| 15.46 \| 19.28 \| 19.29 \| 17.82 \| 17.83 \|

	At standard INT4 the gain over RTN is modest (−2.6%, −1.9%, +0.1%). It was much larger at 3 bits in the research
	experiments, but 3-bit is not a standard packed format and is not exported here.
	Reproduce with `python tests/real_models.py <model> <out> --group 128 --budget 4.75`.

	## Verified here, and not

	* Verified:
	* unit tests and tiny-model round trips for llama / qwen2 / opt / gpt2;
	* real GPT-2, OPT-125m and SmolLM2-135M runs (quantize → save → load → perplexity) on CPU;
	* memory estimates from the real 70B configs;
	* the attention-sized per-layer steps at full Llama-3.1-70B shapes on CPU.
	* Not verified:
	* an end-to-end 70B run (needs a GPU);
	* the MLP-sized 70B steps (need ~7 GB for two 28672² matrices);
	* loading the checkpoints with GPTQModel / vLLM kernels. The format follows the GPTQ v1 layout and the bundled
	loader reads it, but no GPU kernel was run here.

	## Method

	For a linear block y = xW with quantization error R, the output error is tr(RᵀΣR) with Σ = E[xxᵀ].
	No data is available, so Σ is built from the model:
	* LayerNorm/RMSNorm inputs: token embeddings under a Zipf prior on token id (plus positions), then every
	earlier attention-output and MLP-output projection added to the residual stream.
	* MLP second input: the layer's own nonlinearity applied by Monte Carlo to that Gaussian.
	* Attention output: a mixture of value vectors, E[ooᵀ] = μμᵀ + κ·blockdiag(C_v).

	Rounding is the sequential error-feedback recursion that minimises tr(R H Rᵀ) with H = I + λS. With H = I it is
	exactly round-to-nearest.

	## Tests

	```bash
	pytest -q tests/test_core.py tests/test_end_to_end.py
	```
	* `test_core.py`: H = I reproduces RTN; blocked == sequential recursion; GPTQ pack/unpack is exact;
	zero-point safety; shaping lowers the H-weighted error.
	* `test_end_to_end.py`: tiny Llama/Qwen2/OPT/GPT-2 models; quantize → save → load gives logits identical to the
	dequantized reference (uniform and mixed bits).