Upload 17 files

README.md

@@ -1,3 +1,131 @@
----
-license: mit
----
+# 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).

RESULTS.md

@@ -0,0 +1,40 @@
+# Results
+
+CPU (2 cores), WikiText-2 test, 32 windows of 1024 tokens, one run per setting. Final package defaults:
+shaped coding with λ = 1 and κ = 1 (nothing tuned) and dynamic group grids. RTN uses the same GPTQ grid.
+
+| model | group | dense | RTN INT4 | shaped INT4 | RTN mixed | shaped mixed | mixed bits/w |
+|---|---|---|---|---|---|---|---|
+| GPT-2 124M | 128 | 26.84 | 29.91 | **29.14** | 29.62 | **28.86** | 4.75 |
+| OPT-125m | 128 | 29.27 | 31.98 | **31.39** | 30.71 | **30.44** | 4.75 |
+| SmolLM2-135M (Llama) | 64 | 15.46 | **19.28** | 19.29 | **17.82** | 17.83 | 5.0 |
+
+Uniform INT4 is 4 + 32/group bits per quantized weight: 4.25 for groups of 128, 4.5 for groups of 64.
+"Mixed" is INT4/INT8/FP16 per module chosen by the type rule at the stated budget.
+
+Change in perplexity from shaped coding vs RTN:
+* GPT-2 124M: INT4 -2.6%, mixed -2.6%
+* OPT-125m: INT4 -1.9%, mixed -0.9%
+* SmolLM2-135M (Llama): INT4 +0.1%, mixed +0.1%
+
+Static group grids (grids fixed from the original weights), for comparison:
+* GPT-2 124M: shaped INT4 29.09, shaped mixed 28.82
+* OPT-125m: shaped INT4 32.03, shaped mixed 31.02
+* SmolLM2-135M (Llama): shaped INT4 19.19, shaped mixed 17.74
+
+Exact per-layer output error vs RTN under the true input covariance (OPT-125m, layers 0/4/8, evaluation only):
+geometric-mean ratio 0.868 (data-free H), 0.282 (true covariance). The layer-level gain only partly reaches perplexity.
+
+## 70B scaling
+
+Estimates from the real configs (`shapequant memory`):
+
+| model | params | device peak | host buffer | INT4 checkpoint | largest H |
+|---|---|---|---|---|---|
+| Llama-3.1-70B | 70.6B | 21.2 GB | 5.6 GB | 37.8 GB | 28672 |
+| Qwen2.5-72B | 72.7B | 22.1 GB | 5.6 GB | 39.4 GB | 29568 |
+| Llama-2-70B | 69.0B | 21.2 GB | 5.6 GB | 34.8 GB | 28672 |
+
+Per-layer core at Llama-3.1-70B shapes, executed on cpu (2 cores): stacked q/k/v (10240×8192) 26 s, o_proj (8192×8192) 19 s, attention-output covariance 1.7 s.
+The MLP-sized steps (28672-wide covariance, Monte Carlo) need ~7 GB for two 28672² fp32 matrices and were not run
+on this 7.6 GB machine. 70B needs a GPU (≥ 24 GB, 40–80 GB comfortable). Work per layer is the same order as GPTQ's.

pyproject.toml

@@ -0,0 +1,23 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "shapequant"
+version = "0.1.0"
+description = "Data-free, training-free noise-shaped quantization of Hugging Face causal LMs (GPTQ-format output)"
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "Apache-2.0" }
+authors = [{ name = "Angshul Majumdar", email = "angshul@iiitd.ac.in" }]
+dependencies = ["torch>=2.1", "safetensors>=0.4", "transformers>=4.40", "accelerate>=0.26", "huggingface_hub>=0.20"]
+
+[project.optional-dependencies]
+eval = ["datasets"]
+test = ["pytest"]
+
+[project.scripts]
+shapequant = "shapequant.cli:main"
+
+[tool.setuptools]
+packages = ["shapequant"]

shapequant/__init__.py

@@ -0,0 +1,9 @@
+"""shapequant: data-free, training-free noise-shaped quantization of Hugging Face causal LMs."""
+from .quantize import quantize_model, allocate
+from .qlinear import load_quantized, QuantLinear
+from .evaluate import perplexity, wikitext2_windows, windows_from_text
+from .memory import estimate as estimate_memory
+
+__version__ = "0.1.0"
+__all__ = ["quantize_model", "allocate", "load_quantized", "QuantLinear", "perplexity", "wikitext2_windows",
+           "windows_from_text", "estimate_memory"]

shapequant/arch.py

@@ -0,0 +1,98 @@
+"""Architecture adapters: tensor names and structure for the supported Hugging Face causal LMs.
+
+Supported model_type: llama, mistral, qwen2 (incl. Llama-2/3/3.1/3.3-70B, Qwen2.5-72B), opt (pre-LN), gpt2.
+A linear module is described by (module_name, input_key, role):
+  input_key in {"ln1", "attn_out", "ln2", "act"}  - which activation it reads
+  role      in {"q","k","v","qkv","o","gate","up","down"}
+Weights are returned in torch Linear layout (out, in); GPT-2 Conv1D weights are transposed on read.
+"""
+from dataclasses import dataclass, field
+from typing import List, Optional, Tuple
+
+
+@dataclass
+class Linear:
+    name: str              # module path, e.g. "model.layers.3.self_attn.q_proj"
+    input_key: str
+    role: str
+    conv1d: bool = False   # stored (in, out)
+
+
+@dataclass
+class Arch:
+    model_type: str
+    n_layer: int
+    d: int
+    n_head: int
+    n_kv: int
+    head_dim: int
+    inter: int
+    norm: str              # "rms" | "ln"
+    norm_eps: float
+    act: str               # "swiglu" | "relu" | "gelu_new" | "gelu"
+    embed: str
+    pos_embed: Optional[str]
+    pos_offset: int
+    final_norm: str
+    lm_head: Optional[str]
+    vocab: int
+    max_pos: int
+    layers: List[List[Linear]] = field(default_factory=list)
+    norms: List[Tuple[str, str]] = field(default_factory=list)       # per layer: (ln1 prefix, ln2 prefix)
+
+    def norm_names(self, L, which):
+        p = self.norms[L][0 if which == "ln1" else 1]
+        return p + ".weight", (p + ".bias" if self.norm == "ln" else None)
+
+
+def from_config(cfg: dict) -> Arch:
+    mt = cfg.get("model_type")
+    if mt in ("llama", "mistral", "qwen2"):
+        d, H = cfg["hidden_size"], cfg["num_attention_heads"]
+        act = cfg.get("hidden_act", "silu")
+        if act != "silu":
+            raise NotImplementedError(f"{mt} with hidden_act={act}")
+        a = Arch(mt, cfg["num_hidden_layers"], d, H, cfg.get("num_key_value_heads", H), cfg.get("head_dim") or d // H,
+                 cfg["intermediate_size"], "rms", cfg.get("rms_norm_eps", 1e-6), "swiglu",
+                 "model.embed_tokens.weight", None, 0, "model.norm.weight",
+                 None if cfg.get("tie_word_embeddings", False) else "lm_head.weight",
+                 cfg["vocab_size"], cfg.get("max_position_embeddings", 4096))
+        for L in range(a.n_layer):
+            p = f"model.layers.{L}."
+            a.layers.append([Linear(p + "self_attn.q_proj", "ln1", "q"), Linear(p + "self_attn.k_proj", "ln1", "k"),
+                             Linear(p + "self_attn.v_proj", "ln1", "v"), Linear(p + "self_attn.o_proj", "attn_out", "o"),
+                             Linear(p + "mlp.gate_proj", "ln2", "gate"), Linear(p + "mlp.up_proj", "ln2", "up"),
+                             Linear(p + "mlp.down_proj", "act", "down")])
+            a.norms.append((p + "input_layernorm", p + "post_attention_layernorm"))
+        return a
+    if mt == "opt":
+        if not cfg.get("do_layer_norm_before", True):
+            raise NotImplementedError("post-LN OPT (opt-350m) is not supported")
+        if cfg.get("word_embed_proj_dim", cfg["hidden_size"]) != cfg["hidden_size"]:
+            raise NotImplementedError("OPT with project_in/out is not supported")
+        d, H = cfg["hidden_size"], cfg["num_attention_heads"]
+        a = Arch(mt, cfg["num_hidden_layers"], d, H, H, d // H, cfg["ffn_dim"], "ln", 1e-5,
+                 {"relu": "relu"}.get(cfg.get("activation_function", "relu"), cfg.get("activation_function")),
+                 "model.decoder.embed_tokens.weight", "model.decoder.embed_positions.weight", 2,
+                 "model.decoder.final_layer_norm.weight", None if cfg.get("tie_word_embeddings", True) else "lm_head.weight",
+                 cfg["vocab_size"], cfg.get("max_position_embeddings", 2048))
+        for L in range(a.n_layer):
+            p = f"model.decoder.layers.{L}."
+            a.layers.append([Linear(p + "self_attn.q_proj", "ln1", "q"), Linear(p + "self_attn.k_proj", "ln1", "k"),
+                             Linear(p + "self_attn.v_proj", "ln1", "v"), Linear(p + "self_attn.out_proj", "attn_out", "o"),
+                             Linear(p + "fc1", "ln2", "up"), Linear(p + "fc2", "act", "down")])
+            a.norms.append((p + "self_attn_layer_norm", p + "final_layer_norm"))
+        return a
+    if mt == "gpt2":
+        d, H = cfg["n_embd"], cfg["n_head"]
+        a = Arch(mt, cfg["n_layer"], d, H, H, d // H, cfg.get("n_inner") or 4 * d, "ln", cfg.get("layer_norm_epsilon", 1e-5),
+                 {"gelu_new": "gelu_new"}.get(cfg.get("activation_function", "gelu_new"), cfg.get("activation_function")),
+                 "transformer.wte.weight", "transformer.wpe.weight", 0, "transformer.ln_f.weight", None,
+                 cfg["vocab_size"], cfg.get("n_positions", 1024))
+        for L in range(a.n_layer):
+            p = f"transformer.h.{L}."
+            a.layers.append([Linear(p + "attn.c_attn", "ln1", "qkv", True), Linear(p + "attn.c_proj", "attn_out", "o", True),
+                             Linear(p + "mlp.c_fc", "ln2", "up", True), Linear(p + "mlp.c_proj", "act", "down", True)])
+            a.norms.append((p + "ln_1", p + "ln_2"))
+        return a
+    raise NotImplementedError(f"model_type={mt!r} is not supported (supported: llama, mistral, qwen2, opt, gpt2)")

shapequant/cli.py

@@ -0,0 +1,48 @@
+import argparse, json, sys
+
+
+def main(argv=None):
+    p = argparse.ArgumentParser("shapequant", description="Data-free, training-free quantization of Hugging Face causal LMs")
+    sub = p.add_subparsers(dest="cmd", required=True)
+    q = sub.add_parser("quantize", help="quantize a model (hub id or local dir) to a GPTQ-format checkpoint")
+    q.add_argument("model"); q.add_argument("out")
+    q.add_argument("--bits", type=int, default=4, choices=[4, 8])
+    q.add_argument("--group-size", type=int, default=128)
+    q.add_argument("--method", default="shaped", choices=["shaped", "rtn"])
+    q.add_argument("--budget", type=float, default=None, help="average bits/weight for mixed INT4/INT8/FP16 (e.g. 4.75)")
+    q.add_argument("--lam", type=float, default=1.0); q.add_argument("--kappa", type=float, default=1.0)
+    q.add_argument("--samples", type=int, default=16384); q.add_argument("--device", default=None)
+    q.add_argument("--revision", default=None); q.add_argument("--shard-gb", type=float, default=4.0)
+    q.add_argument("--static-groups", action="store_true", help="fit group grids to the original weights instead of the error-fed ones")
+    e = sub.add_parser("eval", help="perplexity of a quantized checkpoint on WikiText-2 (or a text file)")
+    e.add_argument("ckpt"); e.add_argument("--text", default=None); e.add_argument("--seq-len", type=int, default=1024)
+    e.add_argument("--windows", type=int, default=32); e.add_argument("--device", default="cpu"); e.add_argument("--bos", action="store_true")
+    m = sub.add_parser("memory", help="peak-memory estimate from a config (hub id, dir or config.json)")
+    m.add_argument("model"); m.add_argument("--method", default="shaped", choices=["shaped", "rtn"])
+    a = p.parse_args(argv)
+    if a.cmd == "quantize":
+        from .quantize import quantize_model
+        s = quantize_model(a.model, a.out, bits=a.bits, group_size=a.group_size, method=a.method, budget=a.budget, lam=a.lam,
+                           kappa=a.kappa, n_samples=a.samples, device=a.device, revision=a.revision, max_shard_size_gb=a.shard_gb,
+                           dynamic_groups=not a.static_groups)
+        print(json.dumps({k: v for k, v in s.items() if k != "bits_per_module"}, indent=2))
+    elif a.cmd == "eval":
+        from transformers import AutoTokenizer
+        from .qlinear import load_quantized
+        from .evaluate import wikitext2_windows, perplexity
+        tok = AutoTokenizer.from_pretrained(a.ckpt)
+        model = load_quantized(a.ckpt, device=a.device)
+        W = wikitext2_windows(tok, a.seq_len, a.windows, a.text)
+        print(f"perplexity {perplexity(model, W, tok.bos_token_id if a.bos else None):.3f} on {W.numel()} tokens")
+    elif a.cmd == "memory":
+        import os
+        from .memory import estimate
+        src = a.model
+        if not os.path.exists(src):
+            from huggingface_hub import hf_hub_download
+            src = hf_hub_download(a.model, "config.json")
+        print(json.dumps(estimate(src, method=a.method), indent=2))
+
+
+if __name__ == "__main__":
+    sys.exit(main())

shapequant/core.py

@@ -0,0 +1,116 @@
+"""Quantizer core (torch, CPU or GPU).
+
+Grid: standard asymmetric GPTQ grid per (output row, input group): range includes 0,
+scale = (max - min) / (2^b - 1), integer zero-point z = round(-min / scale), q = clamp(round(w/scale) + z, 0, 2^b - 1),
+w_hat = (q - z) * scale. Grids are fixed from the original weights ("static groups").
+Zero-points are kept >= 1 so the checkpoint can use the GPTQ v1 storage convention (stores z - 1).
+
+Rounding: sequential error feedback that minimises tr(R H R^T) for a design matrix H (in x in), i.e. the
+OBQ/GPTQ recursion with lazy block updates. H = None gives round-to-nearest. With H = I it is identical to RTN.
+"""
+import torch
+
+
+def group_params(W, bits, group):
+    """W: (out, in). Returns scale, zero of shape (out, in // group)."""
+    out, inn = W.shape
+    assert inn % group == 0, f"input dim {inn} not divisible by group size {group}"
+    maxq = 2 ** bits - 1
+    Wg = W.reshape(out, inn // group, group)
+    xmin = torch.clamp(Wg.amin(-1), max=0.0)
+    xmax = torch.clamp(Wg.amax(-1), min=0.0)
+    both0 = (xmin == 0) & (xmax == 0)
+    xmin, xmax = torch.where(both0, -torch.ones_like(xmin), xmin), torch.where(both0, torch.ones_like(xmax), xmax)
+    scale = (xmax - xmin) / maxq
+    zero = torch.round(-xmin / scale)
+    # keep z >= 1 (GPTQ v1 storage): if the group has no negative part, use z = 1 and rescale so max still fits
+    low = zero < 1
+    scale = torch.where(low, xmax / (maxq - 1), scale)
+    zero = torch.where(low, torch.ones_like(zero), zero)
+    zero = torch.clamp(zero, 1, maxq)
+    return scale, zero
+
+
+def _q(w, scale, zero, maxq):
+    return torch.clamp(torch.round(w / scale) + zero, 0, maxq)
+
+
+def quantize(W, bits, group, H=None, block=128, dynamic_groups=True):
+    """W: (out, in) float tensor (torch Linear layout). Returns (q uint8 (out, in), scale (out, G), zero (out, G)).
+    H: (in, in) positive-definite design matrix, or None for round-to-nearest.
+    dynamic_groups: each group's grid is fitted to the error-fed (updated) weights when the group starts,
+    as in standard GPTQ, so feedback never pushes weights outside their grid. False = grids from the original W."""
+    W = W.to(torch.float32)
+    out, inn = W.shape
+    maxq = 2 ** bits - 1
+    scale, zero = group_params(W, bits, group)
+    gid = torch.arange(inn, device=W.device) // group
+    if H is None:
+        Q = _q(W, scale[:, gid], zero[:, gid], maxq)
+        return Q.to(torch.uint8), scale, zero
+    W = W.clone()
+    H = H.to(device=W.device, dtype=torch.float32)
+    L = torch.linalg.cholesky(H)
+    Hinv = torch.cholesky_inverse(L); del L
+    U = torch.linalg.cholesky(Hinv, upper=True); del Hinv
+    Q = torch.zeros((out, inn), dtype=torch.uint8, device=W.device)
+    for i1 in range(0, inn, block):
+        i2 = min(i1 + block, inn); n = i2 - i1
+        W1 = W[:, i1:i2].clone(); E1 = torch.zeros_like(W1); U1 = U[i1:i2, i1:i2]
+        for i in range(n):
+            c = i1 + i; g = c // group
+            if dynamic_groups and c % group == 0:
+                seg = W1[:, i:i + group] if i + group <= n else torch.cat([W1[:, i:], W[:, i2:c + group]], 1)
+                sg, zg = group_params(seg, bits, group)
+                scale[:, g], zero[:, g] = sg[:, 0], zg[:, 0]
+            w = W1[:, i]
+            q = _q(w, scale[:, g], zero[:, g], maxq)
+            Q[:, c] = q.to(torch.uint8)
+            err = (w - (q - zero[:, g]) * scale[:, g]) / U1[i, i]
+            W1[:, i:] -= err.unsqueeze(1) * U1[i, i:].unsqueeze(0)
+            E1[:, i] = err
+        if i2 < inn:
+            W[:, i2:] -= E1 @ U[i1:i2, i2:]
+    return Q, scale, zero
+
+
+def dequantize(Q, scale, zero, group):
+    gid = torch.arange(Q.shape[1], device=Q.device) // group
+    return (Q.to(torch.float32) - zero[:, gid]) * scale[:, gid]
+
+
+# --------------------------------------------------------------- GPTQ v1 packing (AutoGPTQ / GPTQModel "gptq" format)
+def pack(Q, scale, zero, bits, group):
+    """Returns dict(qweight int32 (in*bits/32, out), qzeros int32 (G, out*bits/32), scales fp16 (G, out), g_idx int32 (in,))."""
+    assert bits in (2, 4, 8)
+    per = 32 // bits
+    out, inn = Q.shape
+    q = Q.t().contiguous().to(torch.int64)                      # (in, out)
+    assert inn % per == 0 and out % per == 0
+    qw = torch.zeros((inn // per, out), dtype=torch.int64, device=Q.device)
+    for k in range(per):
+        qw |= q[k::per, :] << (bits * k)
+    z = (zero.t().contiguous().to(torch.int64) - 1)             # v1 convention: store z - 1
+    qz = torch.zeros((z.shape[0], out // per), dtype=torch.int64, device=Q.device)
+    for k in range(per):
+        qz |= z[:, k::per] << (bits * k)
+    to32 = lambda t: (t & 0xFFFFFFFF).to(torch.int64).where(t < 2 ** 31, t - 2 ** 32).to(torch.int32)
+    return {"qweight": to32(qw), "qzeros": to32(qz), "scales": scale.t().contiguous().to(torch.float16),
+            "g_idx": (torch.arange(inn, device=Q.device) // group).to(torch.int32)}
+
+
+def unpack(qweight, qzeros, scales, g_idx, bits):
+    """Inverse of pack: dequantized weight in torch Linear layout (out, in), float32."""
+    per = 32 // bits; mask = 2 ** bits - 1
+    qw = qweight.to(torch.int64) & 0xFFFFFFFF
+    inn = qw.shape[0] * per
+    q = torch.empty((inn, qw.shape[1]), dtype=torch.int64, device=qw.device)
+    for k in range(per):
+        q[k::per] = (qw >> (bits * k)) & mask
+    qz = qzeros.to(torch.int64) & 0xFFFFFFFF
+    z = torch.empty((qz.shape[0], qz.shape[1] * per), dtype=torch.int64, device=qz.device)
+    for k in range(per):
+        z[:, k::per] = ((qz >> (bits * k)) & mask) + 1
+    g = g_idx.to(torch.int64)
+    W = (q.to(torch.float32) - z[g].to(torch.float32)) * scales[g].to(torch.float32)   # (in, out)
+    return W.t().contiguous()

shapequant/evaluate.py

@@ -0,0 +1,33 @@
+"""Perplexity on fixed windows (same protocol as the paper experiments)."""
+import math
+import torch
+
+
+def windows_from_text(tokenizer, text, seq_len=1024, n=None, offset=0):
+    ids = tokenizer(text, add_special_tokens=False, return_tensors="pt").input_ids[0]
+    k = len(ids) // seq_len
+    W = ids[: k * seq_len].view(k, seq_len)[offset:]
+    return W if n is None else W[:n]
+
+
+def wikitext2_windows(tokenizer, seq_len=1024, n=None, path=None):
+    if path is not None:
+        text = open(path).read()
+    else:
+        from datasets import load_dataset
+        text = "\n\n".join(load_dataset("wikitext", "wikitext-2-raw-v1", split="test")["text"])
+    return windows_from_text(tokenizer, text, seq_len, n)
+
+
+@torch.no_grad()
+def perplexity(model, windows, bos_token_id=None, device=None):
+    """windows: LongTensor (n, L). If bos_token_id is set it is prepended to every window and all L tokens are scored."""
+    device = device or next(model.parameters()).device
+    nll, cnt = 0.0, 0
+    for w in windows:
+        ids = w if bos_token_id is None else torch.cat([torch.tensor([bos_token_id]), w])
+        ids = ids.to(device)[None]
+        logits = model(ids).logits[0, :-1].float()
+        tgt = ids[0, 1:]
+        nll += float(torch.nn.functional.cross_entropy(logits, tgt, reduction="sum")); cnt += len(tgt)
+    return math.exp(nll / cnt)

shapequant/io.py

@@ -0,0 +1,113 @@
+"""Lazy weight reading (one tensor at a time) and sharded safetensors writing. Nothing loads the whole model."""
+import os, json, glob, shutil
+import torch
+
+
+def resolve(path_or_repo, revision=None):
+    """Local directory, or a Hugging Face repo id (downloads config, tokenizer and weight files)."""
+    if os.path.isdir(path_or_repo):
+        return path_or_repo
+    from huggingface_hub import snapshot_download
+    return snapshot_download(path_or_repo, revision=revision,
+                             allow_patterns=["*.json", "*.safetensors", "*.model", "*.txt", "*.tiktoken", "*.bin"],
+                             ignore_patterns=["*onnx*", "*.msgpack", "*.h5", "*consolidated*", "original/*"])
+
+
+class Reader:
+    """Reads individual tensors from safetensors shards (or a small pytorch_model.bin) by name.
+    Accepts names with or without a leading 'model.' / 'transformer.' prefix."""
+
+    def __init__(self, path):
+        self.path = path
+        st = sorted(glob.glob(os.path.join(path, "*.safetensors")))
+        self.where, self._bin = {}, None
+        if st:
+            from safetensors import safe_open
+            self._open = {}
+            for f in st:
+                with safe_open(f, framework="pt") as h:
+                    for k in h.keys():
+                        self.where[k] = f
+        else:
+            bins = sorted(glob.glob(os.path.join(path, "pytorch_model*.bin")))
+            if not bins:
+                raise FileNotFoundError(f"no *.safetensors or pytorch_model*.bin in {path}")
+            self._bin = {}
+            for f in bins:
+                sd = torch.load(f, map_location="cpu", mmap=True, weights_only=True)
+                for k, v in sd.items():
+                    self._bin[k] = v; self.where[k] = f
+
+    def _key(self, name):
+        if name in self.where:
+            return name
+        for cand in (name.split(".", 1)[1] if "." in name else name, "model." + name, "transformer." + name):
+            if cand in self.where:
+                return cand
+        tail = name.split(".", 1)[1] if name.startswith(("model.", "transformer.")) else name
+        for k in self.where:
+            if k.endswith(tail) and (len(k) == len(tail) or k[-len(tail) - 1] == "."):
+                return k
+        return None
+
+    def has(self, name):
+        return self._key(name) is not None
+
+    def get(self, name, device="cpu", dtype=torch.float32):
+        k = self._key(name)
+        if k is None:
+            raise KeyError(name)
+        if self._bin is not None:
+            t = self._bin[k]
+        else:
+            from safetensors import safe_open
+            with safe_open(self.where[k], framework="pt", device="cpu") as h:
+                t = h.get_tensor(k)
+        return t.to(device=device, dtype=dtype)
+
+    def all_keys(self):
+        return list(self.where)
+
+
+class ShardWriter:
+    """Accumulates tensors and writes model-0000x-of-0000y.safetensors shards plus the index json."""
+
+    def __init__(self, out_dir, max_shard_bytes=4 * 2 ** 30):
+        os.makedirs(out_dir, exist_ok=True)
+        self.out, self.max, self.buf, self.size, self.shards, self.index = out_dir, max_shard_bytes, {}, 0, [], {}
+
+    def add(self, name, t):
+        t = t.detach().to("cpu").contiguous().clone()
+        nb = t.numel() * t.element_size()
+        if self.buf and self.size + nb > self.max:
+            self.flush()
+        self.buf[name] = t; self.size += nb
+
+    def flush(self):
+        if not self.buf:
+            return
+        from safetensors.torch import save_file
+        fn = f"tmp-{len(self.shards):05d}.safetensors"
+        save_file(self.buf, os.path.join(self.out, fn), metadata={"format": "pt"})
+        self.shards.append((fn, list(self.buf)))
+        self.buf, self.size = {}, 0
+
+    def close(self):
+        self.flush()
+        n = len(self.shards); total = 0; wm = {}
+        for i, (fn, names) in enumerate(self.shards):
+            final = "model.safetensors" if n == 1 else f"model-{i + 1:05d}-of-{n:05d}.safetensors"
+            os.replace(os.path.join(self.out, fn), os.path.join(self.out, final))
+            total += os.path.getsize(os.path.join(self.out, final))
+            for k in names:
+                wm[k] = final
+        if n > 1:
+            json.dump({"metadata": {"total_size": total}, "weight_map": wm},
+                      open(os.path.join(self.out, "model.safetensors.index.json"), "w"), indent=2)
+        return total
+
+
+def copy_side_files(src, dst):
+    for f in os.listdir(src):
+        if f.startswith(("tokenizer", "special_tokens_map", "generation_config", "vocab", "merges")) or f.endswith((".model", ".tiktoken")):
+            shutil.copy(os.path.join(src, f), os.path.join(dst, f))

shapequant/memory.py

@@ -0,0 +1,45 @@
+"""Peak-memory estimate for quantize_model from a config (no weights needed)."""
+import json, os
+from .arch import from_config
+
+
+def estimate(config, n_samples=16384, batch=2048, method="shaped"):
+    """config: dict, path to config.json, or a directory. Returns bytes for the accelerator and for host RAM."""
+    if isinstance(config, str):
+        p = os.path.join(config, "config.json") if os.path.isdir(config) else config
+        config = json.load(open(p))
+    a = from_config(config)
+    f32, f64 = 4, 8
+    lin = sum(_numel(a, l) for l in a.layers[0])                       # one decoder layer, fp32 on device
+    ins = {"ln1": a.d, "attn_out": a.n_head * a.head_dim, "ln2": a.d, "act": a.inter}
+    # quantizing the largest input group: stacked W (fp32) + H + Cholesky chain (2 alive) + error buffer
+    worst = 0
+    for key, n in ins.items():
+        rows = sum(_out(a, l) for l in a.layers[0] if l.input_key == key)
+        worst = max(worst, rows * n * f32 * 2 + 2 * n * n * f32 + rows * 128 * f32 * 2)
+    shaped = 0
+    if method == "shaped":
+        shaped = (3 * a.d * a.d * f64                                   # residual Sig, moments
+                  + 2 * a.inter * a.inter * f32                         # S_act accumulator + H
+                  + batch * (a.d + 2 * a.inter) * f32                   # Monte Carlo batch
+                  + 4 * (a.n_head * a.head_dim) ** 2 * f64)             # attention-output S
+    embed = a.vocab * a.d * f32
+    device_peak = lin * f32 + max(worst + shaped, embed + 2 * a.d * a.d * f64)
+    n_params = a.n_layer * lin + a.vocab * a.d * (1 if a.lm_head is None else 2)
+    out_bytes = a.n_layer * lin * (4.25 / 8) + a.vocab * a.d * 2 * (1 if a.lm_head is None else 2)
+    return {"model_type": a.model_type, "params_B": n_params / 1e9, "device_peak_GB": device_peak / 2 ** 30,
+            "host_buffer_GB": 4 + lin * 2 / 2 ** 30, "int4_checkpoint_GB": out_bytes / 2 ** 30,
+            "largest_H_dim": max(ins.values())}
+
+
+def _out(a, l):
+    return {"q": a.n_head * a.head_dim, "k": a.n_kv * a.head_dim, "v": a.n_kv * a.head_dim, "qkv": 3 * a.d,
+            "o": a.d, "gate": a.inter, "up": a.inter, "down": a.d}[l.role]
+
+
+def _in(a, l):
+    return {"ln1": a.d, "attn_out": a.n_head * a.head_dim, "ln2": a.d, "act": a.inter}[l.input_key]
+
+
+def _numel(a, l):
+    return _out(a, l) * _in(a, l)

shapequant/qlinear.py

@@ -0,0 +1,114 @@
+"""Loading shapequant / GPTQ-v1 checkpoints with plain transformers + torch (no CUDA kernels needed).
+
+The checkpoint is standard GPTQ format, so GPTQ-capable stacks (GPTQModel, vLLM, TGI) can load it with their
+fast kernels. This loader is the dependency-light path: each quantized layer keeps packed int32 weights and
+dequantizes on the fly (optionally caching the fp16 weight).
+"""
+import os, re, json, glob
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from . import core
+
+
+class QuantLinear(nn.Module):
+    def __init__(self, in_features, out_features, bits, group_size, bias, dtype=torch.float16, cache=False):
+        super().__init__()
+        self.in_features, self.out_features, self.bits, self.group_size, self.cache = in_features, out_features, bits, group_size, cache
+        per = 32 // bits; G = in_features // group_size
+        self.register_buffer("qweight", torch.zeros((in_features // per, out_features), dtype=torch.int32))
+        self.register_buffer("qzeros", torch.zeros((G, out_features // per), dtype=torch.int32))
+        self.register_buffer("scales", torch.zeros((G, out_features), dtype=torch.float16))
+        self.register_buffer("g_idx", torch.zeros(in_features, dtype=torch.int32))
+        self.register_buffer("bias", torch.zeros(out_features, dtype=dtype) if bias else None)
+        self._w = None
+
+    def weight_fp(self, dtype):
+        if self._w is not None and self._w.dtype == dtype:
+            return self._w
+        w = core.unpack(self.qweight, self.qzeros, self.scales, self.g_idx, self.bits).to(dtype)
+        if self.cache:
+            self._w = w
+        return w
+
+    def forward(self, x):
+        b = None if self.bias is None else self.bias.to(x.dtype)
+        return F.linear(x, self.weight_fp(x.dtype), b)
+
+    def extra_repr(self):
+        return f"in={self.in_features}, out={self.out_features}, bits={self.bits}, group={self.group_size}"
+
+
+def _bits_for(name, qc):
+    bits = qc.get("bits", 4)
+    for pat, ov in (qc.get("dynamic") or {}).items():
+        sign, rx = (pat[0], pat[2:]) if pat[:2] in ("+:", "-:") else ("+", pat)
+        if re.match(rx, name):
+            if sign == "-":
+                return None
+            bits = ov.get("bits", bits)
+    return bits
+
+
+def _set_module(model, name, mod):
+    parent, _, child = name.rpartition(".")
+    setattr(model.get_submodule(parent) if parent else model, child, mod)
+
+
+def load_quantized(path, device="cpu", dtype=torch.float16, cache_weights=False, device_map=None):
+    """Load a checkpoint written by quantize_model. device_map="auto" dispatches layers over GPUs/CPU (accelerate)."""
+    from transformers import AutoConfig, AutoModelForCausalLM
+    from accelerate import init_empty_weights
+    from safetensors.torch import load_file
+    cfg = AutoConfig.from_pretrained(path)
+    qc = cfg.quantization_config if isinstance(cfg.quantization_config, dict) else cfg.quantization_config.to_dict()
+    del cfg.quantization_config
+    with init_empty_weights(include_buffers=False):
+        try:
+            model = AutoModelForCausalLM.from_config(cfg, dtype=dtype)          # transformers >= 4.56
+        except TypeError:
+            model = AutoModelForCausalLM.from_config(cfg, torch_dtype=dtype)
+    files = sorted(glob.glob(os.path.join(path, "*.safetensors")))
+    sd = {}
+    for f in files:
+        sd.update(load_file(f))
+    qnames = sorted({k[: -len(".qweight")] for k in sd if k.endswith(".qweight")})
+    msd = model.state_dict()
+    def find(name):                      # map checkpoint module name to model module name (prefix differences)
+        for cand in (name, name.split(".", 1)[1], "model." + name):
+            try:
+                model.get_submodule(cand); return cand
+            except AttributeError:
+                continue
+        raise KeyError(name)
+    ren = {}
+    for qn in qnames:
+        mn = find(qn); old = model.get_submodule(mn)
+        if hasattr(old, "in_features"):
+            inn, out = old.in_features, old.out_features
+        else:                            # GPT-2 Conv1D: weight (in, out)
+            inn, out = old.weight.shape
+        b = _bits_for(qn, qc)
+        _set_module(model, mn, QuantLinear(inn, out, b, qc["group_size"], (qn + ".bias") in sd, dtype, cache_weights))
+        ren[qn] = mn
+    fixed = {}
+    for k, v in sd.items():
+        mod = k.rsplit(".", 1)[0]
+        if mod in ren:
+            k = ren[mod] + "." + k.rsplit(".", 1)[1]
+        elif k not in msd:
+            for cand in (k.split(".", 1)[1], "model." + k):
+                if cand in msd:
+                    k = cand; break
+        fixed[k] = v.to(dtype) if v.is_floating_point() and not k.endswith(".scales") else v
+    missing, unexpected = model.load_state_dict(fixed, strict=False, assign=True)
+    model.tie_weights()
+    still = [n for n, p in model.named_parameters() if p.device.type == "meta"]
+    if still:
+        raise RuntimeError(f"weights missing from checkpoint: {still[:5]}")
+    model.eval()
+    if device_map is not None:
+        from accelerate import dispatch_model, infer_auto_device_map
+        return dispatch_model(model, infer_auto_device_map(model, no_split_module_classes=model._no_split_modules))
+    return model.to(device)

shapequant/quantize.py

@@ -0,0 +1,214 @@
+"""Data-free, training-free quantization of Hugging Face causal LMs, streamed one decoder layer at a time.
+
+    quantize_model("meta-llama/Llama-3.3-70B-Instruct", "out/llama70b-sq-int4", bits=4, group_size=128)
+
+method="shaped": noise-shaped rounding with the model-implied input covariance H = I + lam*S (default, lam = 1).
+method="rtn":    plain round-to-nearest on the same grid (baseline).
+budget=<bits/weight>: mixed INT4 / INT8 / FP16 per module by the type rule (writers into the residual stream,
+                  i.e. V, attention-output and MLP-output projections, count 4x). None = uniform `bits`.
+Output: GPTQ-format checkpoint (quant_method "gptq", checkpoint_format "gptq", sym=False, desc_act=False),
+per-module bit overrides in quantization_config["dynamic"], plus everything else unquantized in fp16.
+"""
+import os, json, time, math
+import torch
+
+from . import core, sigma
+from .arch import from_config
+from .io import resolve, Reader, ShardWriter, copy_side_files
+
+__all__ = ["quantize_model", "allocate"]
+
+TYPE_WEIGHT = {"q": 1.0, "k": 1.0, "qkv": 1.0, "gate": 1.0, "up": 1.0, "v": 4.0, "o": 4.0, "down": 4.0}
+
+
+def _get_linear(reader, lin, device):
+    W = reader.get(lin.name + ".weight", device=device)
+    if lin.conv1d:
+        W = W.T.contiguous()
+    b = reader.get(lin.name + ".bias", device=device) if reader.has(lin.name + ".bias") else None
+    return W, b
+
+
+def _default_device(device):
+    if device is not None:
+        return torch.device(device)
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def allocate(arch, reader, budget, group_size, device, log=print):
+    """Per-module bits in {4, 8, 16} minimising sum_l s_l D_l(b_l) under an average bits/weight budget.
+    s_l = type weight / ||W_l||_F^2, D_l(b) = ||W_l - RTN_b(W_l)||_F^2 (exact). One streaming pass over the weights."""
+    over = 32.0 / group_size
+    cost = torch.tensor([4 + over, 8 + over, 16.0], dtype=torch.float64)
+    names, D, n, s = [], [], [], []
+    for L in range(arch.n_layer):
+        for lin in arch.layers[L]:
+            W, _ = _get_linear(reader, lin, device)
+            row = []
+            for b in (4, 8):
+                q, sc, z = core.quantize(W, b, group_size)
+                row.append(float(((W - core.dequantize(q, sc, z, group_size)) ** 2).sum()))
+            row.append(float(((W - W.half().float()) ** 2).sum()))
+            names.append(lin.name); D.append(row); n.append(W.numel()); s.append(TYPE_WEIGHT[lin.role] / float((W ** 2).sum()))
+        log(f"[allocate] scanned layer {L + 1}/{arch.n_layer}")
+    D, n, s = torch.tensor(D, dtype=torch.float64), torch.tensor(n, dtype=torch.float64), torch.tensor(s, dtype=torch.float64)
+    N = n.sum()
+    used = lambda k: float((n * cost[k]).sum() / N)
+    pick = lambda mu: (s[:, None] * D + mu * n[:, None] * cost[None, :]).argmin(1)
+    if used(torch.zeros(len(n), dtype=torch.long)) > budget + 1e-9:
+        raise ValueError(f"budget {budget} is below the INT4 floor {cost[0]:.3f} bits/weight")
+    lo, hi = 0.0, 1.0
+    while used(pick(hi)) > budget:
+        hi *= 2
+    for _ in range(200):
+        mid = (lo + hi) / 2
+        lo, hi = (mid, hi) if used(pick(mid)) > budget else (lo, mid)
+    k = pick(hi)
+    while True:                                        # spend leftover budget on the best marginal upgrades
+        best, gain = None, 0.0
+        for l in range(len(k)):
+            if k[l] < 2:
+                extra = float(n[l] * (cost[k[l] + 1] - cost[k[l]]) / N)
+                if used(k) + extra <= budget + 1e-12:
+                    r = float(s[l] * (D[l, k[l]] - D[l, k[l] + 1]) / (n[l] * (cost[k[l] + 1] - cost[k[l]])))
+                    if r > gain:
+                        best, gain = l, r
+        if best is None:
+            break
+        k[best] += 1
+    bits = {nm: (4, 8, 16)[int(x)] for nm, x in zip(names, k)}
+    return bits, used(k)
+
+
+def quantize_model(model, out_dir, bits=4, group_size=128, method="shaped", budget=None, lam=1.0, kappa=1.0,
+                   n_samples=16384, device=None, max_shard_size_gb=4.0, revision=None, seed=0, dynamic_groups=True, log=print):
+    """Quantize a Hugging Face causal LM (local dir or hub id) and write a GPTQ-format checkpoint to out_dir.
+    Peak memory: one decoder layer of weights plus a few (in x in) matrices; see shapequant.memory.estimate."""
+    t0 = time.time()
+    assert method in ("shaped", "rtn")
+    assert bits in (4, 8) or budget is not None
+    dev = _default_device(device)
+    src = resolve(model, revision)
+    cfg = json.load(open(os.path.join(src, "config.json")))
+    arch = from_config(cfg)
+    reader = Reader(src)
+    if budget is not None:
+        bits_of, bpw = allocate(arch, reader, budget, group_size, dev, log)
+        log(f"[allocate] {sum(v == 4 for v in bits_of.values())} INT4, {sum(v == 8 for v in bits_of.values())} INT8, "
+            f"{sum(v == 16 for v in bits_of.values())} FP16 modules; {bpw:.3f} bits/weight")
+    else:
+        bits_of = None
+
+    writer = ShardWriter(out_dir, int(max_shard_size_gb * 2 ** 30))
+    quantized = set()
+    prefixes = tuple(_layer_prefixes(arch))
+    for key in reader.all_keys():                          # everything outside the decoder layers: fp16 copy
+        ck = _canon(key, arch)
+        if ck.startswith("lm_head") and arch.lm_head is None:     # tied to the input embedding
+            continue
+        if not ck.startswith(prefixes):
+            writer.add(ck, reader.get(key, dtype=torch.float16))
+
+    res = sigma.ResidualModel(arch, reader, dev) if method == "shaped" else None
+    for L in range(arch.n_layer):
+        tl = time.time()
+        lins = {lin.name: lin for lin in arch.layers[L]}
+        W = {nm: _get_linear(reader, lin, dev) for nm, lin in lins.items()}
+        S = {}
+        if method == "shaped":
+            m1 = res.norm_moments(L, "ln1")
+            S["ln1"] = sigma.S_norm(m1)
+            v = [lin for lin in lins.values() if lin.role in ("v", "qkv")][0]
+            Wv, bv = W[v.name]
+            if v.role == "qkv":
+                Wv, bv = Wv[2 * arch.d:3 * arch.d], (None if bv is None else bv[2 * arch.d:3 * arch.d])
+            S["attn_out"] = sigma.S_attn(arch, m1, Wv, bv, kappa)
+            o = [lin for lin in lins.values() if lin.role == "o"][0]
+            res.add_writer(W[o.name][0])
+            m2 = res.norm_moments(L, "ln2")
+            S["ln2"] = sigma.S_norm(m2)
+            up = [lin for lin in lins.values() if lin.role == "up"][0]
+            gate = [lin for lin in lins.values() if lin.role == "gate"]
+            S["act"] = sigma.S_act(arch, m2, W[up.name][0], W[up.name][1], W[gate[0].name][0] if gate else None,
+                                   n_samples=n_samples, seed=seed + L, device=dev)
+            dn = [lin for lin in lins.values() if lin.role == "down"][0]
+            res.add_writer(W[dn.name][0])
+        # group modules that read the same input and use the same bit width: one shaped pass on the stacked rows
+        groups = {}
+        for nm, lin in lins.items():
+            b = bits if bits_of is None else bits_of[nm]
+            groups.setdefault((lin.input_key, b), []).append(nm)
+        uses = {}
+        for (ik, b) in groups:
+            if b != 16:
+                uses[ik] = uses.get(ik, 0) + 1
+        for (ik, b), nms in groups.items():
+            if b == 16:
+                for nm in nms:
+                    writer.add(nm + ".weight", (W[nm][0].T if lins[nm].conv1d else W[nm][0]).half())
+                    if W[nm][1] is not None:
+                        writer.add(nm + ".bias", W[nm][1].half())
+                continue
+            Wcat = torch.cat([W[nm][0] for nm in nms], 0)
+            H = None
+            if method == "shaped":
+                uses[ik] -= 1                                     # last use: reuse S's memory for H
+                H = S.pop(ik) if uses[ik] == 0 else S[ik].clone()
+                H = H.to(device=dev, dtype=torch.float32); H.mul_(lam); H.diagonal().add_(1.0)
+            q, sc, z = core.quantize(Wcat, b, group_size, H=H, dynamic_groups=dynamic_groups)
+            del H
+            r0 = 0
+            for nm in nms:
+                r1 = r0 + W[nm][0].shape[0]
+                pk = core.pack(q[r0:r1], sc[r0:r1], z[r0:r1], b, group_size)
+                for k_, t in pk.items():
+                    writer.add(nm + "." + k_, t)
+                if W[nm][1] is not None:
+                    writer.add(nm + ".bias", W[nm][1].half())
+                quantized.add(nm); r0 = r1
+        for which in ("ln1", "ln2"):
+            gname, bname = arch.norm_names(L, which)
+            writer.add(gname, reader.get(gname, dtype=torch.float16))
+            if bname and reader.has(bname):
+                writer.add(bname, reader.get(bname, dtype=torch.float16))
+        del W, S
+        if dev.type == "cuda":
+            torch.cuda.empty_cache()
+        log(f"[quantize] layer {L + 1}/{arch.n_layer} done in {time.time() - tl:.1f}s")
+    total = writer.close()
+
+    qc = {"quant_method": "gptq", "bits": bits if bits_of is None else 4, "group_size": group_size, "desc_act": False,
+          "sym": False, "static_groups": not dynamic_groups, "true_sequential": False, "damp_percent": 0.0, "checkpoint_format": "gptq",
+          "pack_dtype": "int32", "lm_head": False,
+          "quantizer": ["shapequant:0.1.0"],
+          "shapequant": {"method": method, "lam": lam, "kappa": kappa, "dynamic_groups": dynamic_groups, "n_samples": n_samples, "budget": budget,
+                         "data_free": True, "training_free": True}}
+    if bits_of is not None:
+        qc["dynamic"] = {}
+        for nm, b in bits_of.items():
+            if b == 8:
+                qc["dynamic"]["+:^" + nm.replace(".", r"\.") + "$"] = {"bits": 8}
+            elif b == 16:
+                qc["dynamic"]["-:^" + nm.replace(".", r"\.") + "$"] = {}
+    cfg_out = dict(cfg); cfg_out["quantization_config"] = qc; cfg_out["torch_dtype"] = "float16"
+    json.dump(cfg_out, open(os.path.join(out_dir, "config.json"), "w"), indent=2)
+    copy_side_files(src, out_dir)
+    summary = {"model": model, "out_dir": out_dir, "method": method, "group_size": group_size,
+               "bits": bits if bits_of is None else "mixed", "checkpoint_bytes": total, "minutes": (time.time() - t0) / 60}
+    if bits_of is not None:
+        summary["bits_per_module"] = bits_of
+    json.dump(summary, open(os.path.join(out_dir, "shapequant_summary.json"), "w"), indent=2)
+    log(f"[done] {out_dir}: {total / 2 ** 30:.2f} GiB in {summary['minutes']:.1f} min")
+    return summary
+
+
+def _layer_prefixes(arch):
+    return [arch.norms[L][0].rsplit(".", 1)[0] + "." for L in range(arch.n_layer)]
+
+
+def _canon(key, arch):
+    """Key with the canonical root used by the architecture ('model.' for llama/opt, 'transformer.' for gpt2)."""
+    root = "transformer." if arch.model_type == "gpt2" else "model."
+    if key.startswith((root, "lm_head")):
+        return key
+    return root + key

shapequant/sigma.py

@@ -0,0 +1,112 @@
+"""Model-implied input second moments, built from the weights only and streamed layer by layer.
+
+Residual stream:  Sig_0 = E_{t~pi}[e_t e_t^T] (+ positional embeddings), pi = Zipf prior on token id
+                  (BPE ids are assigned in merge order, i.e. by frequency); then Sig += W^T W for every writer
+                  (attention output and MLP output) in order.
+Norm outputs:     LayerNorm: C = P Sig P / (tr/n), cov = g C g, mean = beta, M2 = cov + beta beta^T.
+                  RMSNorm:   C = Sig / (tr/n),     M2 = g C g, mean = g m / s, cov = M2 - mean mean^T.
+MLP 2nd input:    Monte Carlo of the block's own nonlinearity on x ~ N(mean, cov).
+Attention output: E[o o^T] = mu_o mu_o^T + kappa * blockdiag_heads(C_v of the head's KV group).
+Every S is normalised to trace n. Design H = I + lam * S.
+"""
+import math
+import torch
+
+F64 = torch.float64
+
+
+class ResidualModel:
+    def __init__(self, arch, reader, device, ctx=1024, zipf_shift=10.0, chunk=8192):
+        self.a, self.r, self.dev = arch, reader, device
+        E = reader.get(arch.embed, device=device, dtype=torch.float32)
+        V = E.shape[0]
+        pi = 1.0 / (torch.arange(V, device=device, dtype=F64) + zipf_shift); pi /= pi.sum()
+        Sig = torch.zeros((arch.d, arch.d), dtype=F64, device=device); m = torch.zeros(arch.d, dtype=F64, device=device)
+        for s in range(0, V, chunk):
+            e = E[s:s + chunk].to(F64); w = pi[s:s + chunk]
+            Sig += (e * w[:, None]).T @ e; m += w @ e
+        del E
+        if arch.pos_embed is not None and reader.has(arch.pos_embed):
+            P = reader.get(arch.pos_embed, device=device, dtype=F64)
+            T = min(ctx, P.shape[0] - arch.pos_offset)
+            P = P[arch.pos_offset: arch.pos_offset + T]
+            pm = P.mean(0)
+            Sig += P.T @ P / T + torch.outer(m, pm) + torch.outer(pm, m); m = m + pm
+        self.Sig, self.m = Sig, m
+
+    def add_writer(self, W_out_in):
+        """W in torch layout (out=d, in); the writer adds W h with isotropic h."""
+        W = W_out_in.to(device=self.dev, dtype=F64)
+        self.Sig += W @ W.T
+
+    def norm_moments(self, L, which):
+        a = self.a; n = a.d
+        gname, bname = a.norm_names(L, which)
+        g = self.r.get(gname, device=self.dev, dtype=F64)
+        if a.norm == "ln":
+            Sg = self.Sig - self.Sig.mean(0, keepdim=True)
+            Sg = Sg - Sg.mean(1, keepdim=True)                    # P Sig P
+            C = Sg / (torch.trace(Sg) / n)
+            cov = g[:, None] * C * g[None, :]
+            b = self.r.get(bname, device=self.dev, dtype=F64) if bname and self.r.has(bname) else torch.zeros(n, dtype=F64, device=self.dev)
+            return b, cov, cov + torch.outer(b, b)
+        s2 = torch.trace(self.Sig) / n
+        M2 = g[:, None] * (self.Sig / s2) * g[None, :]
+        mu = g * self.m / torch.sqrt(s2)
+        w, Q = torch.linalg.eigh(M2 - torch.outer(mu, mu))
+        cov = (Q * w.clamp(min=0)) @ Q.T
+        return mu, cov, M2
+
+
+def normalise(M):
+    return M.mul_(M.shape[0] / torch.trace(M))
+
+
+def S_norm(moments):
+    return normalise(moments[2].clone())
+
+
+def S_act(arch, moments, W_up, b_up, W_gate, n_samples=16384, batch=2048, seed=0, device="cpu"):
+    """W_up/W_gate: torch layout (inter, d). Returns (inter, inter) float32."""
+    mu, cov, _ = moments
+    w, Q = torch.linalg.eigh(cov)
+    R = (Q * w.clamp(min=0).sqrt()).to(torch.float32)
+    mu = mu.to(torch.float32)
+    Wu = W_up.to(device=device, dtype=torch.float32).T
+    Wg = None if W_gate is None else W_gate.to(device=device, dtype=torch.float32).T
+    bu = None if b_up is None else b_up.to(device=device, dtype=torch.float32)
+    gen = torch.Generator(device=device).manual_seed(seed)
+    M = torch.zeros((Wu.shape[1], Wu.shape[1]), dtype=torch.float32, device=device); done = 0
+    while done < n_samples:
+        b = min(batch, n_samples - done)
+        x = mu + torch.randn((b, R.shape[0]), generator=gen, device=device) @ R.T
+        u = x @ Wu
+        if bu is not None:
+            u = u + bu
+        if arch.act == "swiglu":
+            z = x @ Wg; h = z * torch.sigmoid(z) * u
+        elif arch.act == "relu":
+            h = torch.relu(u)
+        elif arch.act == "gelu_new":
+            h = 0.5 * u * (1 + torch.tanh(math.sqrt(2 / math.pi) * (u + 0.044715 * u ** 3)))
+        elif arch.act == "gelu":
+            h = torch.nn.functional.gelu(u)
+        else:
+            raise NotImplementedError(arch.act)
+        M += h.T @ h; done += b
+    return normalise(M.div_(done))
+
+
+def S_attn(arch, moments, W_v, b_v, kappa=1.0):
+    """W_v: torch layout (n_kv*head_dim, d). Returns (n_head*head_dim)^2 float64."""
+    mu, cov, _ = moments
+    Wv = W_v.to(device=cov.device, dtype=F64)
+    mv = Wv @ mu + (0 if b_v is None else b_v.to(device=cov.device, dtype=F64))
+    Cv = Wv @ cov @ Wv.T
+    H, dh = arch.n_head, arch.head_dim; per = H // arch.n_kv
+    n = H * dh
+    mo = torch.zeros(n, dtype=F64, device=cov.device); M = torch.zeros((n, n), dtype=F64, device=cov.device)
+    for h in range(H):
+        gk = h // per; s, t = slice(h * dh, (h + 1) * dh), slice(gk * dh, (gk + 1) * dh)
+        mo[s] = mv[t]; M[s, s] = kappa * Cv[t, t]
+    return normalise(M + torch.outer(mo, mo))

tests/real_models.py

@@ -0,0 +1,22 @@
+"""Real-model check: quantize -> load -> WikiText-2 perplexity, for RTN and shaped, uniform INT4 and mixed.
+usage: python tests/real_models.py <hf_model_dir_or_id> <out_dir> [--group 128] [--budget 4.75] [--windows 32] [--bos ID] [--text wt2_test.txt]"""
+import argparse, json, os, time, torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from shapequant import quantize_model, load_quantized, perplexity, wikitext2_windows
+
+p = argparse.ArgumentParser()
+p.add_argument("model"); p.add_argument("out"); p.add_argument("--group", type=int, default=128)
+p.add_argument("--budget", type=float, default=4.75); p.add_argument("--windows", type=int, default=32)
+p.add_argument("--bos", type=int, default=None); p.add_argument("--text", default=None); p.add_argument("--device", default="cpu")
+a = p.parse_args()
+tok = AutoTokenizer.from_pretrained(a.model)
+W = wikitext2_windows(tok, 1024, a.windows, a.text)
+R = {"dense": perplexity(AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.float32).eval().to(a.device), W, a.bos)}
+for tag, kw in (("rtn_int4", dict(bits=4, method="rtn")), ("shaped_int4", dict(bits=4)),
+                ("rtn_mixed", dict(budget=a.budget, method="rtn")), ("shaped_mixed", dict(budget=a.budget))):
+    t = time.time(); d = os.path.join(a.out, tag)
+    s = quantize_model(a.model, d, group_size=a.group, device=a.device, log=lambda *x: None, **kw)
+    R[tag] = {"ppl": perplexity(load_quantized(d, dtype=torch.float32, device=a.device), W, a.bos),
+              "minutes": (time.time() - t) / 60, "checkpoint_MB": s["checkpoint_bytes"] / 2 ** 20}
+    print(tag, R[tag], flush=True)
+json.dump(R, open(os.path.join(a.out, "results.json"), "w"), indent=1)

tests/scale_70b.py

@@ -0,0 +1,35 @@
+"""70B scaling check without the weights: memory estimates from real configs, and timing of the per-layer core at
+Llama-3.1-70B shapes on whatever device is available (random weights; the arithmetic is shape-dependent only)."""
+import json, sys, time, torch
+from shapequant import estimate_memory, core, sigma
+from shapequant.arch import from_config
+
+CFGS = {
+ "Llama-3.1-70B": {"model_type": "llama", "hidden_size": 8192, "intermediate_size": 28672, "num_hidden_layers": 80,
+                   "num_attention_heads": 64, "num_key_value_heads": 8, "vocab_size": 128256, "hidden_act": "silu",
+                   "rms_norm_eps": 1e-5, "max_position_embeddings": 131072, "tie_word_embeddings": False},
+ "Qwen2.5-72B":   {"model_type": "qwen2", "hidden_size": 8192, "intermediate_size": 29568, "num_hidden_layers": 80,
+                   "num_attention_heads": 64, "num_key_value_heads": 8, "vocab_size": 152064, "hidden_act": "silu",
+                   "rms_norm_eps": 1e-6, "max_position_embeddings": 32768, "tie_word_embeddings": False},
+ "Llama-2-70B":   {"model_type": "llama", "hidden_size": 8192, "intermediate_size": 28672, "num_hidden_layers": 80,
+                   "num_attention_heads": 64, "num_key_value_heads": 8, "vocab_size": 32000, "hidden_act": "silu",
+                   "rms_norm_eps": 1e-5, "max_position_embeddings": 4096, "tie_word_embeddings": False},
+}
+out = {k: estimate_memory(v) for k, v in CFGS.items()}
+print(json.dumps(out, indent=1))
+dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+a = from_config(CFGS["Llama-3.1-70B"]); d = a.d; t = {}
+torch.manual_seed(0)
+# attention-input group (q,k,v stacked: 10240 x 8192) and attention-output (8192 x 8192) at full 70B size
+for name, rows in (("qkv_stacked", 64 * 128 + 2 * 8 * 128), ("o_proj", d)):
+    W = torch.randn(rows, d, device=dev) * 0.02
+    A = torch.randn(d, 256, device=dev) / 16; H = A @ A.T; H.diagonal().add_(1.0)
+    t0 = time.time(); q, s, z = core.quantize(W, 4, 128, H=H); t[name] = time.time() - t0
+    pk = core.pack(q, s, z, 4, 128); assert pk["qweight"].shape == (d // 8, rows)
+    del W, H, A, q, s, z, pk
+# attention-output covariance construction at 70B size (GQA 64 query / 8 KV heads)
+mu = torch.zeros(d, dtype=torch.float64, device=dev); cov = torch.eye(d, dtype=torch.float64, device=dev)
+t0 = time.time(); S = sigma.S_attn(a, (mu, cov, cov), torch.randn(8 * 128, d, device=dev) * 0.02, None); t["S_attn"] = time.time() - t0
+assert S.shape == (d, d); del S, cov
+print(json.dumps({"device": str(dev), "seconds": t}, indent=1))
+json.dump({"estimates": out, "timing": {"device": str(dev), "seconds": t}}, open(sys.argv[1] if len(sys.argv) > 1 else "scale_70b.json", "w"), indent=1)

tests/test_core.py

@@ -0,0 +1,64 @@
+import torch
+from shapequant import core
+
+
+def _ref_sequential(W, bits, group, H, dynamic=False):
+    """Unblocked reference of the error-feedback recursion."""
+    W = W.clone().float(); out, inn = W.shape; maxq = 2 ** bits - 1
+    sc, z = core.group_params(W, bits, group)
+    U = torch.linalg.cholesky(torch.cholesky_inverse(torch.linalg.cholesky(H)), upper=True)
+    Q = torch.zeros_like(W)
+    for c in range(inn):
+        g = c // group
+        if dynamic and c % group == 0:
+            s_, z_ = core.group_params(W[:, c:c + group], bits, group); sc[:, g], z[:, g] = s_[:, 0], z_[:, 0]
+        q = torch.clamp(torch.round(W[:, c] / sc[:, g]) + z[:, g], 0, maxq); Q[:, c] = q
+        err = (W[:, c] - (q - z[:, g]) * sc[:, g]) / U[c, c]
+        W[:, c:] -= err[:, None] * U[c, c:][None, :]
+    return Q
+
+
+def test_identity_design_is_rtn():
+    torch.manual_seed(0)
+    W = torch.randn(64, 256)
+    q0, s0, z0 = core.quantize(W, 4, 64)
+    for dyn in (False, True):
+        q1, s1, z1 = core.quantize(W, 4, 64, H=torch.eye(256), dynamic_groups=dyn)
+        assert torch.equal(q0, q1) and torch.equal(s0, s1)
+
+
+def test_blocked_equals_sequential():
+    torch.manual_seed(1)
+    W = torch.randn(48, 384); A = torch.randn(384, 384); H = A @ A.T / 384 + torch.eye(384)
+    for dyn in (False, True):
+        for blk in (32, 64, 128, 256):
+            q, _, _ = core.quantize(W, 4, 128, H=H, block=blk, dynamic_groups=dyn)
+            assert (q.float() - _ref_sequential(W, 4, 128, H, dyn)).abs().max() == 0, (dyn, blk)
+
+
+def test_pack_roundtrip():
+    torch.manual_seed(2)
+    for bits in (4, 8):
+        W = torch.randn(96, 256) * 0.05
+        q, s, z = core.quantize(W, bits, 128)
+        pk = core.pack(q, s, z, bits, 128)
+        Wd = core.unpack(pk["qweight"], pk["qzeros"], pk["scales"], pk["g_idx"], bits)
+        ref = core.dequantize(q, s.half().float(), z, 128)
+        assert torch.allclose(Wd, ref, atol=0, rtol=0)
+        assert pk["qweight"].dtype == torch.int32 and pk["qweight"].shape == (256 * bits // 32, 96)
+
+
+def test_zero_points_v1_safe():
+    W = torch.rand(32, 128)                       # all positive groups: zero-point must still be >= 1
+    q, s, z = core.quantize(W, 4, 64)
+    assert z.min() >= 1
+    err = (W - core.dequantize(q, s, z, 64)).abs().max() / W.abs().max()
+    assert err < 0.05
+
+
+def test_shaping_reduces_weighted_error():
+    torch.manual_seed(3)
+    W = torch.randn(128, 256); A = torch.randn(256, 32); H = torch.eye(256) + A @ A.T
+    def cost(q, s, z):
+        R = W - core.dequantize(q, s, z, 64); return float(torch.trace(R @ H @ R.T))
+    assert cost(*core.quantize(W, 4, 64, H=H)) < cost(*core.quantize(W, 4, 64))

tests/test_end_to_end.py

@@ -0,0 +1,55 @@
+"""Tiny random models of every supported architecture: quantize -> save -> load -> logits must equal a model whose
+weights are the dequantized values (exactness of format, loader and module mapping)."""
+import json, os, tempfile
+import pytest
+import torch
+from transformers import AutoModelForCausalLM, LlamaConfig, OPTConfig, GPT2Config, Qwen2Config
+
+from shapequant import quantize_model, load_quantized, core
+
+CFGS = {
+    "llama": LlamaConfig(hidden_size=128, intermediate_size=256, num_hidden_layers=2, num_attention_heads=4,
+                         num_key_value_heads=2, vocab_size=512, max_position_embeddings=256, tie_word_embeddings=False),
+    "qwen2": Qwen2Config(hidden_size=128, intermediate_size=256, num_hidden_layers=2, num_attention_heads=4,
+                         num_key_value_heads=2, vocab_size=512, max_position_embeddings=256, tie_word_embeddings=True),
+    "opt": OPTConfig(hidden_size=128, ffn_dim=256, num_hidden_layers=2, num_attention_heads=4, vocab_size=512,
+                     max_position_embeddings=256, word_embed_proj_dim=128),
+    "gpt2": GPT2Config(n_embd=128, n_layer=2, n_head=4, vocab_size=512, n_positions=256),
+}
+
+
+@pytest.mark.parametrize("kind", list(CFGS))
+@pytest.mark.parametrize("budget", [None, 6.5])
+def test_roundtrip(kind, budget):
+    torch.manual_seed(0)
+    src = AutoModelForCausalLM.from_config(CFGS[kind]).eval()
+    with tempfile.TemporaryDirectory() as d:
+        src.save_pretrained(os.path.join(d, "src"), safe_serialization=True)
+        s = quantize_model(os.path.join(d, "src"), os.path.join(d, "q"), bits=4, group_size=64, budget=budget,
+                           n_samples=2048, device="cpu", log=lambda *a: None)
+        qm = load_quantized(os.path.join(d, "q"), dtype=torch.float32)
+        # reference: same architecture, weights replaced by the dequantized values
+        ref = AutoModelForCausalLM.from_config(CFGS[kind]).eval()
+        ref.load_state_dict(src.state_dict())
+        sd = ref.state_dict()
+        from safetensors.torch import load_file
+        ck = load_file(os.path.join(d, "q", "model.safetensors"))
+        qc = json.load(open(os.path.join(d, "q", "config.json")))["quantization_config"]
+        from shapequant.qlinear import _bits_for
+        n_q = 0
+        for k in ck:
+            if k.endswith(".qweight"):
+                mod = k[:-8]; bits = _bits_for(mod, qc)
+                Wd = core.unpack(ck[mod + ".qweight"], ck[mod + ".qzeros"], ck[mod + ".scales"], ck[mod + ".g_idx"], bits)
+                key = next(c for c in (mod + ".weight", mod.split(".", 1)[1] + ".weight") if c in sd)
+                sd[key] = Wd.T if kind == "gpt2" else Wd
+                n_q += 1
+        ref.load_state_dict(sd)
+        ids = torch.randint(0, 512, (1, 40))
+        with torch.no_grad():
+            a, b, c = qm(ids).logits, ref(ids).logits, src(ids).logits
+        assert n_q > 0
+        assert torch.allclose(a, b, atol=2e-3), (a - b).abs().max()      # loader == dequantized reference
+        assert (a - c).abs().max() < 5 * (c.std() + 1e-3)                  # close to the original model
+        if budget is not None:
+            assert "dynamic" in qc and s["bits"] == "mixed"