diff --git a/README.md b/README.md index 2bbbf25..e556570 100644 --- a/README.md +++ b/README.md @@ -18,12 +18,14 @@ The implementation is clean-room and Apache-2.0 licensed. - RPBH rotation, exact unit-sphere Lloyd-Max codebooks, packed 2/3/4/6/8-bit weights, and online activation quantization without calibration data. - Model-specific policies for paper-sensitive AdaLN and output-layer handling. +- Optional GPTQ rounding (`orbitquant.gptq`): calibrated codes in the same codebook and + artifact format, for models that lose too much with nearest rounding. - Compact `safetensors` artifacts and Hugging Face `save_pretrained()` / `from_pretrained()` integration. - Packed-weight CUDA, Triton, and Metal inference paths that avoid a full dequantized weight matrix. - Fused transformer checkpoints for Krea 2, FLUX.2, Ideogram 4, Qwen-Image 2.1, - MiniMax-H3, Boogu-Image and FLUX.3 Klein (`orbitquant.fused`): each block runs + MiniMax-H3, Boogu-Image, FLUX.3 Klein and MOSS-Image (`orbitquant.fused`): each block runs grouped INT8 GEMMs with norm/modulation prologues and SwiGLU/gated-residual epilogues plus an INT8 attention kernel, and the checkpoint stores only the fused groups. @@ -57,6 +59,43 @@ larger batches or when the native package is unavailable. All paths share the ep changes logits at the INT8 rounding level (~0.4% relative); measure sampled outputs before adopting it for a model. +## GPTQ rounding (opt-in calibration) + +Plain OrbitQuant rounds every rotated weight to its nearest Lloyd-Max centroid and uses no +data. At 4 bits some models lose too much this way. `orbitquant.gptq` keeps the codebook, +rotation, bit width, BF16 row-norm scale and packed layout, and chooses the codes against the +Hessian `X^T X` of each projection's rotated calibration inputs: columns are rounded one at a +time (largest Hessian diagonal first) and each rounding error is compensated in the columns not +yet rounded. The result is an ordinary OrbitQuant artifact that loads and runs through every +runtime path unchanged. + +```python +from orbitquant import quantize_model +from orbitquant.gptq import HessianCollector, load_rounded_codes, round_groups + +groups = [["blocks.0.attn.q", "blocks.0.attn.k", "blocks.0.attn.v"], ["blocks.0.attn.o"]] +with HessianCollector(model, groups, config) as hessians: # BF16 model on the GPU + for batch in calibration: + model(**batch) +codes = round_groups(model, groups, hessians, config) +quantize_model(model, config) +load_rounded_codes(model, codes) +``` + +- Projections in one group read the same input and share one Hessian. +- `update_state_dict(state, codes)` rewrites the tensors of an existing artifact's state dict. +- For a large model, collect and round a few blocks at a time. Groups of the same shape round + together on the GPU (`batch`, `memory_budget`). `sequential=True` writes the rounded weights + back, so later blocks are calibrated on rounded inputs. +- Hessians and updates run in full FP32 (TF32 is disabled inside). A Hessian that does not + factorize gets 10x more damping, up to six times; `RoundedCodes.damping` reports the value + used, and `rtn_error` / `gptq_error` report the Hessian-weighted output error of nearest and + GPTQ rounding. + +On the MOSS-Image 4-bit transformer, calibrated on 10 prompts at 12 sampling steps, the mean +LPIPS of 8 generated images against BF16 fell from 0.371 (nearest rounding) to 0.209; 6-bit +nearest rounding gives 0.166. + ## Install ```bash @@ -514,6 +553,7 @@ the same for an `orbitquant-v1` component artifact, and | MiniMax-H3 | `MiniMaxH3Transformer3DModel` | per-row (timestep, modality) AdaLN tables, partial rotate-half RoPE | | Boogu-Image | `BooguImageTransformer2DModel` | refiner, single- and double-stream blocks; 8-bit activations | | FLUX.3 Klein | `flux3.dit.DiT` (Black Forest Labs' `flux3` package) | per-mode and single-stream blocks: Q\|K\|V, SwiGLU, gated residuals; one sample at a time | +| MOSS-Image | `MOSSImageTransformer` (the model's `pipeline_moss_image.py`) | joint image/text blocks: RMSNorm and shared AdaLN prologues, Q\|K\|V, SiTU gate\|up, gated residuals; one sample at a time | Groups take the activation width of their projections unless the layout overrides it: `fuse(..., activation_bits={"qkv": 8})` gives a group 8-bit @@ -523,6 +563,21 @@ block; the GEMMs stay INT8, so the wider codes cost almost no speed. Low-bit artifacts with mixed weight widths (boundary and interior protection) fuse into row segments of one GEMM. +Components without a block family can still run their OrbitQuant projections +on the fused runtime: `fused.fuse_linears(model)` replaces every loaded +`OrbitQuantLinear` with a `FusedLinear` (one INT8 GEMM over the projection's +packed rows) and quantizes its input per token, by default with 8-bit absmax +activations. Text encoders that run once per prompt use it: + +```python +import orbitquant +from orbitquant import fused +from transformers import AutoModel + +encoder = AutoModel.from_pretrained(text_encoder_dir, device_map={"": "cuda"}) +fused.fuse_linears(encoder) # activation_bits=4 keeps the Lloyd-Max codes +``` + The fused runtime needs CUDA and Triton. Run inference under `torch.no_grad()` or `torch.inference_mode()`. Fused weights are frozen parameters, so model CPU offload and Diffusers block-level group offloading (including streamed @@ -571,6 +626,9 @@ The implementation follows the paper's shared data-agnostic basis: no prompt, timestep, or calibration statistics. - INT4 group-64 RTN for model policies that identify dynamic AdaLN projections. +GPTQ rounding (`orbitquant.gptq`) is not part of the paper's method: it uses calibration data +to choose codes within the same codebook and format. The default quantizer never calls it. + The detailed requirement matrix is in [`docs/paper-methodology-audit.md`](docs/paper-methodology-audit.md). diff --git a/pyproject.toml b/pyproject.toml index 27647bb..d2c384f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "orbitquant" -version = "0.13.0" +version = "0.14.0" description = "Calibration-free OrbitQuant for transformer linear projections" readme = "README.md" requires-python = ">=3.11" diff --git a/src/orbitquant/__init__.py b/src/orbitquant/__init__.py index f5f489f..b2d2d44 100644 --- a/src/orbitquant/__init__.py +++ b/src/orbitquant/__init__.py @@ -24,7 +24,7 @@ from orbitquant.quantizer import register_hf_quantizers from orbitquant.recipes import recipe -__version__ = "0.13.0" +__version__ = "0.14.0" register_hf_quantizers() diff --git a/src/orbitquant/fused/__init__.py b/src/orbitquant/fused/__init__.py index 58d4de6..c11760b 100644 --- a/src/orbitquant/fused/__init__.py +++ b/src/orbitquant/fused/__init__.py @@ -11,10 +11,12 @@ from importlib import import_module _EXPORTS = { + "FusedLinear": "orbitquant.fused.linear", "FusedRuntime": "orbitquant.fused.runtime", "finalize": "orbitquant.fused.install", "fuse": "orbitquant.fused.install", "fuse_component_artifact": "orbitquant.fused.convert", + "fuse_linears": "orbitquant.fused.linear", "prepare_skeleton": "orbitquant.fused.install", "release": "orbitquant.fused.install", "save_pretrained": "orbitquant.fused.install", diff --git a/src/orbitquant/fused/families/__init__.py b/src/orbitquant/fused/families/__init__.py index c89057b..758b1ca 100644 --- a/src/orbitquant/fused/families/__init__.py +++ b/src/orbitquant/fused/families/__init__.py @@ -18,7 +18,16 @@ from orbitquant.fused.groups import RowSource -_MODULES = ("krea2", "flux2", "ideogram4", "minimax_h3", "qwenimage21", "boogu", "flux3") +_MODULES = ( + "krea2", + "flux2", + "ideogram4", + "minimax_h3", + "qwenimage21", + "boogu", + "flux3", + "moss_image", +) @dataclass(frozen=True) diff --git a/src/orbitquant/fused/families/moss_image.py b/src/orbitquant/fused/families/moss_image.py new file mode 100644 index 0000000..48be669 --- /dev/null +++ b/src/orbitquant/fused/families/moss_image.py @@ -0,0 +1,137 @@ +"""MOSS-Image (OpenMOSS ``MOSSImageTransformer``) joint image and text blocks. + +Every block is a sequential attention + SiTU feed-forward block over the concatenated image and +text tokens. Both RMSNorms and the shared AdaLN modulation (one row per sample, the same for all +42 blocks) run in the activation-quantization prologue; Q|K|V is one GEMM, Q/K RMSNorm and the +rotate-half rotary embedding are one kernel, the gate|up rows of ``ffn_up_gate`` are one GEMM +with the SiTU epilogue, and both output projections add their gated residual in the GEMM +epilogue. The transformer evaluates its samples one at a time, so the blocks see one sample. +""" + +from __future__ import annotations + +import torch + +from orbitquant.fused.families import int8_attention_allowed, specs +from orbitquant.fused.groups import RowSource +from orbitquant.kernels import triton_activation as act +from orbitquant.kernels import triton_int8_gemm as gemm + +FAMILY = "moss_image" +MODEL_CLASSES = ("MOSSImageTransformer",) +EPS = 1e-6 + + +def blocks(model): + return [(f"layers.{i}", "block") for i in range(len(model.layers))] + + +def groups(kind, block): + # ffn_up_gate rows are [hidden | gate]; the gated-unit epilogue reads (gate, up) pairs. + ffn = block.proj["ffn_down"].in_features + return specs( + qkv=["proj.q", "proj.k", "proj.v"], + mlp=( + [RowSource("proj.ffn_up_gate", ffn, 2 * ffn), RowSource("proj.ffn_up_gate", 0, ffn)], + {"swiglu": True}, + ), + attn_out=["proj.o"], + mlp_out=["proj.ffn_down"], + ) + + +def prepare(block, kind): + block.oq_int8_attention = int8_attention_allowed(block.q_norm.weight, block.k_norm.weight) + + +def forward(kind): + return _block + + +def _rope(runtime, rope): + """The ``[1, L, 128]`` bf16 ``sin | cos`` table of 64 frequencies as per-channel ``[L, 128]`` + cos and sin.""" + cached = runtime.__dict__.get("moss_image_rope") + if cached is not None and cached[0] is rope: + return cached[1], cached[2] + sin, cos = rope[0].float().chunk(2, dim=-1) + cos = torch.cat((cos, cos), dim=-1).contiguous() + sin = torch.cat((sin, sin), dim=-1).contiguous() + runtime.__dict__["moss_image_rope"] = (rope, cos, sin) + return cos, sin + + +def _block(self, x, rope, modulation): + if x.shape[0] != 1: + raise NotImplementedError("fused MOSS-Image blocks run one sample at a time") + runtime, fused = self.oq_runtime, self.oq_fused + shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = ( + m.reshape(-1).contiguous() for m in modulation.to(torch.bfloat16).unbind(1) + ) + h = x[0].to(torch.bfloat16).clone() + seq, dim = h.shape + source = runtime.input( + h, + dim, + norm=act.NORM_RMS, + norm_weight=self.norm1.weight, + mod=act.MOD_SHIFT_SCALE, + mod_scale=scale_a, + mod_shift=shift_a, + rms_eps=EPS, + ) + inner = fused.qkv.out_features // 3 + heads = self.heads + head_dim = inner // heads + qkv = runtime.linear(fused.qkv, source) + cos, sin = _rope(runtime, rope) + query = torch.empty((1, seq, heads, head_dim), device=h.device, dtype=torch.bfloat16) + key = torch.empty_like(query) + for src, weight, out in ( + (qkv[:, :inner], self.q_norm.weight, query), + (qkv[:, inner : 2 * inner], self.k_norm.weight, key), + ): + act.qk_norm_rope_with_triton( + src, + heads, + head_dim, + weight, + cos, + sin, + EPS, + weight_mode=act.QK_WEIGHT_DIFFUSERS, + style=act.ROPE_HALF, + rope_bf16=True, + out=out, + ) + value = qkv[:, 2 * inner :].unflatten(-1, (heads, head_dim))[None] + attention = runtime.attention(query, key, value, int8_allowed=self.oq_int8_attention) + attention = attention.reshape(seq, inner) + runtime.linear( + fused.attn_out, + runtime.input(attention, inner), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=gate_a, + out=h, + ) + source = runtime.input( + h, + dim, + norm=act.NORM_RMS, + norm_weight=self.norm2.weight, + mod=act.MOD_SHIFT_SCALE, + mod_scale=scale_f, + mod_shift=shift_f, + rms_eps=EPS, + ) + mid = runtime.linear(fused.mlp, source, epilogue=gemm.EPILOGUE_SITU) + runtime.linear( + fused.mlp_out, + runtime.input(mid, mid.shape[1]), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=gate_f, + out=h, + ) + return h[None] diff --git a/src/orbitquant/fused/linear.py b/src/orbitquant/fused/linear.py new file mode 100644 index 0000000..ba3b607 --- /dev/null +++ b/src/orbitquant/fused/linear.py @@ -0,0 +1,67 @@ +"""Single OrbitQuant projections as fused INT8 GEMMs, for models without a block family. + +``fuse_linears(model)`` swaps every loaded ``OrbitQuantLinear`` for a ``FusedLinear``: the +projection's packed rows become one ``PackedGroup`` that decodes into the shared INT8 workspace +right before its GEMM, and its input is quantized per token by the fused activation kernels. +With ``activation_bits=8`` (the default) the rotated input becomes INT8 with one absmax scale +per token instead of the projection's Lloyd-Max codes; the GEMM is INT8 either way, so the +wider activations cost no GEMM time. Use it for components that run a few times per +generation, such as text encoders; block families also fuse norms, modulation and epilogues. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from orbitquant.fused.groups import PackedGroup, RowSource +from orbitquant.fused.runtime import FusedRuntime +from orbitquant.layers import OrbitQuantLinear + + +class FusedLinear(nn.Module): + """One packed projection on the fused runtime; bf16 output of the input's shape.""" + + def __init__(self, group: PackedGroup, runtime: FusedRuntime): + super().__init__() + self.group = group + self.in_features = group.in_features + self.out_features = group.out_features + self.__dict__["oq_runtime"] = runtime + + def forward(self, x: torch.Tensor) -> torch.Tensor: + rows = x.reshape(-1, self.in_features) + if rows.dtype != torch.bfloat16 or not rows.is_contiguous(): + rows = rows.to(torch.bfloat16).contiguous() + runtime = self.oq_runtime + out = runtime.linear(self.group, runtime.input(rows, self.in_features)) + return out.reshape(*x.shape[:-1], self.out_features) + + +def fuse_linears(model: nn.Module, *, activation_bits: int = 8, config=None) -> int: + """Replace the loaded ``OrbitQuantLinear`` modules of ``model`` in place; returns how many. + + ``config`` is the OrbitQuant config when the model does not carry one. Call + ``orbitquant.fused.release(model)`` before moving the weights off the GPU.""" + from orbitquant.fused.install import _quant_config + + runtime = FusedRuntime.from_config(_quant_config(model, config), attention="sdpa") + count = 0 + for name, module in list(model.named_modules()): + if not isinstance(module, OrbitQuantLinear) or module.packed_weight_indices is None: + continue + holder = nn.Module() + holder.proj = module + group = PackedGroup.from_sources( + holder, [RowSource("proj")], activation_bits=activation_bits + ) + parent_path, _, child = name.rpartition(".") + parent = model.get_submodule(parent_path) if parent_path else model + fused = FusedLinear(group, runtime) + if child.isdigit(): + parent[int(child)] = fused + else: + setattr(parent, child, fused) + count += 1 + model.__dict__["oq_runtime"] = runtime + return count diff --git a/src/orbitquant/fused/runtime.py b/src/orbitquant/fused/runtime.py index 3bb603b..e63292d 100644 --- a/src/orbitquant/fused/runtime.py +++ b/src/orbitquant/fused/runtime.py @@ -198,11 +198,14 @@ def _dense_epilogue( ): """Eager equivalents of the GEMM epilogues for BF16 groups (bias is already applied).""" del bias - if epilogue == gemm.EPILOGUE_SWIGLU: + if epilogue in (gemm.EPILOGUE_SWIGLU, gemm.EPILOGUE_SITU): rows = y.shape[0] chunks = y.reshape(rows, -1, 2, gemm.SWIGLU_INTERLEAVE) g, u = chunks[:, :, 0].reshape(rows, -1), chunks[:, :, 1].reshape(rows, -1) - y = F.silu(g) * u + if epilogue == gemm.EPILOGUE_SWIGLU: + y = F.silu(g) * u + else: + y = (4.0 * torch.tanh(g / 4.0) * torch.sigmoid(g)) * (25.0 * torch.tanh(u / 25.0)) elif epilogue == gemm.EPILOGUE_GELU_TANH: y = F.gelu(y, approximate="tanh") elif epilogue == gemm.EPILOGUE_SIGMOID_TAIL: diff --git a/src/orbitquant/gptq.py b/src/orbitquant/gptq.py new file mode 100644 index 0000000..21f72d3 --- /dev/null +++ b/src/orbitquant/gptq.py @@ -0,0 +1,359 @@ +"""Optional calibration: GPTQ rounding of OrbitQuant weight codes. + +Plain OrbitQuant rounds every rotated, row-normalized weight to its nearest Lloyd-Max centroid +and needs no data. GPTQ keeps the codebook, rotation, bit width, row-norm scale and packed layout, +and chooses the codes column by column against the Hessian ``X^T X`` of the projection's rotated +calibration inputs: the rounding error of each column is compensated in the columns not yet +rounded. The packed result loads and runs through every existing runtime path unchanged. + +Flow, with the BF16 model on the GPU:: + + groups = [["blocks.0.attn.q", "blocks.0.attn.k", "blocks.0.attn.v"], ["blocks.0.attn.o"]] + with HessianCollector(model, groups, config) as hessians: + for batch in calibration: + model(**batch) + codes = round_groups(model, groups, hessians, config) + quantize_model(model, config) + load_rounded_codes(model, codes) + +Projections in one group read the same input and share one Hessian. For a model too large for all +Hessians at once, collect and round a few blocks at a time; ``sequential=True`` writes the rounded +weights back into the BF16 model so that later blocks are calibrated on rounded inputs. +""" + +from __future__ import annotations + +import contextlib +from dataclasses import dataclass + +import torch +from torch import nn + +from orbitquant.codebooks.lloyd_max import get_codebook +from orbitquant.config import OrbitQuantConfig +from orbitquant.packing import pack_lowbit +from orbitquant.rotations.rpbh import get_rpbh_rotation + +__all__ = [ + "HessianCollector", + "RoundedCodes", + "gptq_round", + "load_rounded_codes", + "round_groups", + "update_state_dict", +] + + +@dataclass +class RoundedCodes: + """Packed codes and BF16 row norms of one projection, with the Hessian-weighted relative + output error of nearest rounding (``rtn_error``) and of GPTQ rounding (``gptq_error``) and + the relative damping its Hessian needed.""" + + packed: torch.Tensor + row_norms: torch.Tensor + bits: int + rtn_error: float + gptq_error: float + damping: float = 0.0 + + +@contextlib.contextmanager +def _full_fp32(): + """Hessians and GPTQ updates need full float32 matmuls: with TF32 a few large-norm tokens + leave rounding noise above the damping in the weak directions, and Cholesky fails.""" + previous = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = previous + + +def _inverse_factor(H: torch.Tensor, damp: float, attempts: int = 6): + """Upper Cholesky factor of the damped inverse Hessian per group, and the damping it took: + a group that is not positive definite at ``damp`` retries with ten times the damping.""" + eye = torch.eye(H.shape[-1], device=H.device) + mean = torch.diagonal(H, dim1=1, dim2=2).mean(dim=1) + level = torch.full_like(mean, damp) + for _ in range(attempts): + L, info = torch.linalg.cholesky_ex(H + (level * mean)[:, None, None] * eye) + if not info.any(): + U, info = torch.linalg.cholesky_ex(torch.cholesky_inverse(L), upper=True) + if not info.any(): + return U, level + level = torch.where(info != 0, level * 10, level) + raise torch.linalg.LinAlgError(f"Hessian not positive definite at damping {level.tolist()}") + + +def _weighted_error(weight: torch.Tensor, approx: torch.Tensor, hessian: torch.Tensor) -> float: + diff = weight - approx + with _full_fp32(): + return (torch.sum((diff @ hessian) * diff) / torch.sum((weight @ hessian) * weight)).item() + + +def gptq_round( + weight: torch.Tensor, + hessian: torch.Tensor, + *, + bits: int, + codebook_version: int = 2, + damp: float = 0.01, + block: int = 128, + act_order: bool = True, + eps: float = 1e-10, + return_damping: bool = False, +): + """Codes ``[rows, cols]`` (uint8) and row scales ``[rows]`` (BF16-representable float32) of a + rotated weight against the Hessian of its rotated input; with ``return_damping`` also the + relative damping each Hessian needed (``damp`` or a tenfold multiple of it). + + ``weight`` is the rotated weight in float32. The row scale is the BF16-rounded row norm, as in + nearest rounding, so only the codes differ. A leading batch dimension (``[G, rows, cols]`` + weights with ``[G, cols, cols]`` Hessians) rounds G same-shape projections in one pass, which + keeps the GPU busy: the column loop launches the same number of kernels for any G.""" + with _full_fp32(): + return _gptq_round( + weight, hessian, bits, codebook_version, damp, block, act_order, eps, return_damping + ) + + +def _gptq_round( + weight, hessian, bits, codebook_version, damp, block, act_order, eps, return_damping +): + batched = weight.ndim == 3 + W = (weight if batched else weight[None]).clone() + H = (hessian if batched else hessian[None]).to(torch.float32).clone() + groups, rows, cols = W.shape + if H.shape != (groups, cols, cols): + raise ValueError("weight must be [(G,) rows, cols] and hessian [(G,) cols, cols]") + codebook = get_codebook(cols, bits, codebook_version) + centroids = codebook.centroids.to(W.device) + boundaries = codebook.boundaries.to(W.device) + eye = torch.eye(cols, device=H.device) + diag = torch.diagonal(H, dim1=1, dim2=2) + dead = diag == 0 + H = H + eye * dead[:, None, :] + W = W.masked_fill(dead[:, None, :], 0) + scale = W.new_empty(groups, rows) + scale[:] = ( + (weight if batched else weight[None]).norm(dim=2).to(torch.bfloat16).to(torch.float32) + ) + scale = scale.clamp_min(eps) + if act_order: + perm = torch.argsort(torch.diagonal(H, dim1=1, dim2=2), dim=1, descending=True) + else: + perm = torch.arange(cols, device=W.device).expand(groups, cols) + W = torch.gather(W, 2, perm[:, None, :].expand(groups, rows, cols)) + H = torch.gather( + torch.gather(H, 1, perm[:, :, None].expand(groups, cols, cols)), + 2, + perm[:, None, :].expand(groups, cols, cols), + ) + Hinv, used_damp = _inverse_factor(H, damp) + codes = torch.zeros(groups, rows, cols, dtype=torch.uint8, device=W.device) + for start in range(0, cols, block): + stop = min(start + block, cols) + W1 = W[:, :, start:stop].clone() + E1 = torch.zeros_like(W1) + Hinv1 = Hinv[:, start:stop, start:stop] + for i in range(stop - start): + column = W1[:, :, i] + index = torch.bucketize(column / scale, boundaries) + codes[:, :, start + i] = index.to(torch.uint8) + err = (column - centroids[index] * scale) / Hinv1[:, i, i][:, None] + W1[:, :, i:] -= err[:, :, None] * Hinv1[:, i, i:][:, None, :] + E1[:, :, i] = err + W[:, :, stop:] -= E1 @ Hinv[:, start:stop, stop:] + out = torch.empty_like(codes) + out.scatter_(2, perm[:, None, :].expand(groups, rows, cols), codes) + result = (out, scale, used_damp) if batched else (out[0], scale[0], used_damp[0]) + return result if return_damping else result[:2] + + +class HessianCollector: + """Accumulates ``X^T X`` of the rotated input of the first projection of every group while the + context is open. Index it by any projection name of a group.""" + + def __init__(self, model: nn.Module, groups: list[list[str]], config: OrbitQuantConfig): + self.model = model + self.groups = [list(group) for group in groups] + self.config = config + self._hessians: dict[str, torch.Tensor] = {} + self._owner: dict[str, str] = {} + self._handles = [] + self.tokens: dict[str, int] = {} + + def __enter__(self) -> HessianCollector: + for group in self.groups: + head = group[0] + module = self.model.get_submodule(head) + dim = module.in_features + rotation = get_rpbh_rotation( + dim=dim, seed=self.config.rotation_seed, block_size=self.config.block_size + ) + device = module.weight.device + self._hessians[head] = torch.zeros(dim, dim, dtype=torch.float32, device=device) + self.tokens[head] = 0 + for name in group: + self._owner[name] = head + + def hook(_module, args, head=head, dim=dim, rotation=rotation): + x = rotation.apply_to_activations(args[0].reshape(-1, dim).to(torch.float32)) + with _full_fp32(): + self._hessians[head].addmm_(x.T, x) + self.tokens[head] += x.shape[0] + + self._handles.append(module.register_forward_pre_hook(hook)) + return self + + def __exit__(self, *exc) -> None: + for handle in self._handles: + handle.remove() + self._handles.clear() + + def __getitem__(self, name: str) -> torch.Tensor: + return self._hessians[self._owner[name]] + + def release(self, name: str) -> None: + """Free the Hessian of ``name``'s group.""" + self._hessians.pop(self._owner[name], None) + + +def round_groups( + model: nn.Module, + groups: list[list[str]], + hessians: HessianCollector, + config: OrbitQuantConfig, + *, + module_bits: dict[str, int] | None = None, + damp: float = 0.01, + block: int = 128, + act_order: bool = True, + sequential: bool = False, + release: bool = True, + batch: int = 8, + memory_budget: int = 6 << 30, +) -> dict[str, RoundedCodes]: + """GPTQ codes of every projection in ``groups`` from the BF16 weights of ``model``. + + ``module_bits`` overrides ``config.weight_bits`` per projection name (one width per group). + Up to ``batch`` groups of the same shape and width round together in one pass, fewer when + their float32 work tensors would exceed ``memory_budget`` bytes. + With ``sequential`` the BF16 weights are replaced by their rounded values.""" + module_bits = module_bits or {} + prepared = [] + for group in groups: + modules = [model.get_submodule(name) for name in group] + bits = {module_bits.get(name, config.weight_bits) for name in group} + if len(bits) != 1: + raise ValueError(f"group {group} mixes weight widths {sorted(bits)}") + sizes = [m.out_features for m in modules] + prepared.append((group, modules, sizes, bits.pop(), sum(sizes), modules[0].in_features)) + buckets: dict[tuple[int, int, int], list] = {} + for item in prepared: + buckets.setdefault((item[3], item[4], item[5]), []).append(item) + result: dict[str, RoundedCodes] = {} + for (bits, rows, dim), items in buckets.items(): + rotation = get_rpbh_rotation( + dim=dim, seed=config.rotation_seed, block_size=config.block_size + ) + codebook = get_codebook(dim, bits, config.codebook_version) + # Per group: about five [cols, cols] (Hessian copies, Cholesky factors) and eight + # [rows, cols] float32 work tensors; the batch also stays within half the free GPU memory. + per_group = 4 * dim * (5 * dim + 8 * rows) + budget = memory_budget + if torch.cuda.is_available() and items[0][1][0].weight.is_cuda: + budget = min(budget, torch.cuda.mem_get_info(items[0][1][0].weight.device)[0] // 2) + size = max(1, min(batch, budget // per_group)) + for first in range(0, len(items), size): + chunk = items[first : first + size] + H = torch.stack([hessians[item[0][0]] for item in chunk]) + weight = torch.stack( + [ + rotation.apply_to_weight( + torch.cat( + [ + m.weight.detach().to(device=H.device, dtype=torch.float32) + for m in item[1] + ] + ) + ) + for item in chunk + ] + ) + codes, scale, used = gptq_round( + weight, + H, + bits=bits, + codebook_version=config.codebook_version, + damp=damp, + block=block, + act_order=act_order, + eps=config.activation_eps, + return_damping=True, + ) + centroids = codebook.centroids.to(weight.device) + boundaries = codebook.boundaries.to(weight.device) + for g, (group, modules, sizes, _, _, _) in enumerate(chunk): + rounded = centroids[codes[g].long()] * scale[g, :, None] + nearest = centroids[torch.bucketize(weight[g] / scale[g, :, None], boundaries)] + rtn_error = _weighted_error(weight[g], nearest * scale[g, :, None], H[g]) + gptq_error = _weighted_error(weight[g], rounded, H[g]) + del nearest + for name, part, part_scale in zip( + group, codes[g].split(sizes), scale[g].split(sizes), strict=True + ): + result[name] = RoundedCodes( + packed=pack_lowbit(part.contiguous(), bits=bits), + row_norms=part_scale.to(torch.bfloat16), + bits=bits, + rtn_error=rtn_error, + gptq_error=gptq_error, + damping=float(used[g]), + ) + if sequential: + original = rotation.apply_inverse_to_weight(rounded) + for module, part in zip(modules, original.split(sizes), strict=True): + module.weight.data = part.to( + device=module.weight.device, dtype=module.weight.dtype + ) + if release: + hessians.release(group[0]) + del rounded + del H, weight, codes, scale + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return result + + +def load_rounded_codes(model: nn.Module, codes: dict[str, RoundedCodes]) -> None: + """Copy rounded codes into the ``OrbitQuantLinear`` modules of a quantized ``model``.""" + from orbitquant.layers import OrbitQuantLinear + + for name, rounded in codes.items(): + module = model.get_submodule(name) + if not isinstance(module, OrbitQuantLinear) or module.packed_weight_indices is None: + raise TypeError(f"{name} is not a packed OrbitQuantLinear") + if module.weight_bits != rounded.bits: + raise ValueError(f"{name}: {module.weight_bits}-bit module, {rounded.bits}-bit codes") + target = module.packed_weight_indices + if target.numel() != rounded.packed.numel(): + raise ValueError(f"{name}: packed size {rounded.packed.numel()} != {target.numel()}") + target.data.copy_(rounded.packed.reshape(target.shape).to(target.device)) + module.row_norms.data.copy_( + rounded.row_norms.to(module.row_norms.device, module.row_norms.dtype) + ) + module.clear_dequantized_cache() + + +def update_state_dict(state: dict[str, torch.Tensor], codes: dict[str, RoundedCodes]) -> None: + """Write rounded codes into a quantized state dict (``.packed_weight_indices`` and + ``.row_norms``), keeping each tensor's shape and dtype.""" + for name, rounded in codes.items(): + packed_key, norms_key = f"{name}.packed_weight_indices", f"{name}.row_norms" + packed = state[packed_key] + if packed.numel() != rounded.packed.numel(): + raise ValueError(f"{name}: packed size {rounded.packed.numel()} != {packed.numel()}") + state[packed_key] = rounded.packed.reshape(packed.shape).to("cpu") + state[norms_key] = rounded.row_norms.to("cpu", state[norms_key].dtype) diff --git a/src/orbitquant/kernels/triton_int8_gemm.py b/src/orbitquant/kernels/triton_int8_gemm.py index 2a46fa1..48c4418 100644 --- a/src/orbitquant/kernels/triton_int8_gemm.py +++ b/src/orbitquant/kernels/triton_int8_gemm.py @@ -18,6 +18,8 @@ gate may be selected per row from a table (``gate_index``). * ``EPILOGUE_SIGMOID_TAIL``: columns ``>= sig_from`` store ``sigmoid(bf16(y))``. * ``EPILOGUE_GELU_TANH``: ``gelu_tanh(bf16(y))``. +* ``EPILOGUE_SITU``: the SwiGLU layout with MOSS-Image's soft-capped unit, + ``bf16(4 tanh(g / 4) * sigmoid(g)) * bf16(25 tanh(u / 25))``. """ from __future__ import annotations @@ -37,6 +39,8 @@ EPILOGUE_RESIDUAL_GATE = 2 EPILOGUE_SIGMOID_TAIL = 3 EPILOGUE_GELU_TANH = 4 +EPILOGUE_SITU = 5 +_PAIRED = (EPILOGUE_SWIGLU, EPILOGUE_SITU) # The SwiGLU epilogue pairs gate/up halves inside one N tile; weights packed for it interleave # rows in chunks of this size. @@ -151,10 +155,17 @@ def _int8_scaled_gemm_body( if HAS_BIAS: y += tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)[None, :] rows = offs_m[:, None].to(tl.int64) * stride_om - if EPILOGUE == 1: + if EPILOGUE == 1 or EPILOGUE == 5: pairs = tl.permute(tl.reshape(_bf16(y), (BLOCK_M, 2, BLOCK_N // 2)), (0, 2, 1)) g, u = tl.split(pairs) - h = _bf16(libdevice.div_rn(g, 1.0 + libdevice.exp(-g))) * u + if EPILOGUE == 1: + h = _bf16(libdevice.div_rn(g, 1.0 + libdevice.exp(-g))) * u + else: + # Each eager bf16 op of the reference rounds once; so does each step here. + soft = _bf16(4.0 * _bf16(libdevice.tanh(g * 0.25))) + sig = _bf16(libdevice.div_rn(1.0, 1.0 + libdevice.exp(-g))) + cap = _bf16(25.0 * _bf16(libdevice.tanh(_bf16(libdevice.div_rn(u, 25.0))))) + h = _bf16(soft * sig) * cap offs_h = col_offset + pid_n * (BLOCK_N // 2) + tl.arange(0, BLOCK_N // 2) tl.store( out_ptr + rows + offs_h[None, :], @@ -240,11 +251,11 @@ def matmul_int8_scaled_with_triton( n = b.shape[0] if b.shape[1] != k: raise ValueError(f"inner dimensions differ: {k} vs {b.shape[1]}") - if epilogue == EPILOGUE_SWIGLU and n % (2 * SWIGLU_INTERLEAVE): - raise ValueError(f"SwiGLU epilogue needs N divisible by {2 * SWIGLU_INTERLEAVE}") + if epilogue in _PAIRED and n % (2 * SWIGLU_INTERLEAVE): + raise ValueError(f"gated-unit epilogues need N divisible by {2 * SWIGLU_INTERLEAVE}") if epilogue == EPILOGUE_RESIDUAL_GATE and (residual is None or gate is None): raise ValueError("residual-gate epilogue needs residual and gate") - out_cols = n // 2 if epilogue == EPILOGUE_SWIGLU else n + out_cols = n // 2 if epilogue in _PAIRED else n if out is None: out = torch.empty((rows, out_cols), device=a.device, dtype=output_dtype) if rows == 0: @@ -256,7 +267,7 @@ def matmul_int8_scaled_with_triton( raise ValueError("residual must share the output layout") a_scale = a_scale.to(torch.float32).contiguous() b_scale = b_scale.to(torch.float32).contiguous() - kernel = _int8_scaled_gemm_swiglu if epilogue == EPILOGUE_SWIGLU else _int8_scaled_gemm + kernel = _int8_scaled_gemm_swiglu if epilogue in _PAIRED else _int8_scaled_gemm def grid(meta): return (triton.cdiv(rows, meta["BLOCK_M"]) * triton.cdiv(n, meta["BLOCK_N"]),) diff --git a/tests/test_fused_moss_image.py b/tests/test_fused_moss_image.py new file mode 100644 index 0000000..79cd242 --- /dev/null +++ b/tests/test_fused_moss_image.py @@ -0,0 +1,230 @@ +import types + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +pytest.importorskip("triton") +from orbitquant.fused.families import ( # noqa: E402 + adapter_for_class, + adapter_for_family, + moss_image, +) +from orbitquant.kernels import available_backends # noqa: E402 + + +def _require_cuda(): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + +# A copy of MOSS-Image's block (pipeline_moss_image.py), the eager reference of the fused forward. +class _RMSNorm(nn.Module): + def __init__(self, dim): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6) + return y.to(x.dtype) * self.weight + + +def _situ(hidden, gate): + return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * (25.0 * torch.tanh(hidden / 25.0)) + + +def _rotate(x, rope): + sin, cos = rope.chunk(2, -1) + sin = torch.cat((sin, sin), -1).unsqueeze(-2) + cos = torch.cat((cos, cos), -1).unsqueeze(-2) + a, b = x.chunk(2, -1) + return x * cos + torch.cat((-b, a), -1) * sin + + +def _attention(q, k, v, rope, heads): + shape = (*q.shape[:2], heads, q.shape[-1] // heads) + q = _rotate(q.reshape(shape), rope).transpose(1, 2) + k = _rotate(k.reshape(shape), rope).transpose(1, 2) + v = v.reshape(shape).transpose(1, 2) + return F.scaled_dot_product_attention(q, k, v).transpose(1, 2).flatten(2) + + +class _Block(nn.Module): + def __init__(self, dim, ffn_dim, heads): + super().__init__() + self.heads = heads + self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim) + self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads) + self.proj = nn.ModuleDict( + {name: nn.Linear(dim, dim, bias=False) for name in ("q", "k", "v", "o")} + ) + self.proj["ffn_up_gate"] = nn.Linear(dim, 2 * ffn_dim, bias=False) + self.proj["ffn_down"] = nn.Linear(ffn_dim, dim, bias=False) + + def forward(self, x, rope, modulation): + shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = modulation.unbind(1) + y = self.norm1(x) * (1 + scale_a[:, None]) + shift_a[:, None] + q = self.q_norm(self.proj["q"](y).unflatten(-1, (self.heads, -1))).flatten(2) + k = self.k_norm(self.proj["k"](y).unflatten(-1, (self.heads, -1))).flatten(2) + x = ( + x + + self.proj["o"](_attention(q, k, self.proj["v"](y), rope, self.heads)) + * gate_a[:, None] + ) + y = self.norm2(x) * (1 + scale_f[:, None]) + shift_f[:, None] + hidden, gate = self.proj["ffn_up_gate"](y).chunk(2, -1) + return x + self.proj["ffn_down"](_situ(hidden, gate)) * gate_f[:, None] + + +def _rope_table(seq, bands=16, device="cpu"): + """The model's ``[1, L, 128]`` sin | cos table (4 axes of 16 bands), cast to bf16 like the + model.""" + positions = torch.randn(seq, 4, device=device) * 20 + freqs = 1 / (10000 ** (torch.arange(bands, device=device, dtype=torch.float32) / bands)) + phase = positions.unsqueeze(-1) * freqs + return torch.cat((phase.sin(), phase.cos()), -2).flatten(-2).unsqueeze(0).to(torch.bfloat16) + + +def test_family_is_registered_for_the_transformer_class(): + assert adapter_for_family("moss_image") is moss_image + assert adapter_for_class("MOSSImageTransformer") is moss_image + + +def test_blocks_and_groups(): + model = nn.Module() + model.layers = nn.ModuleList([_Block(256, 512, 2) for _ in range(3)]) + assert [p for p, _ in moss_image.blocks(model)] == ["layers.0", "layers.1", "layers.2"] + specs = moss_image.groups("block", model.layers[0]) + assert list(specs) == ["qkv", "mlp", "attn_out", "mlp_out"] + gate, up = specs["mlp"].sources + # ffn_up_gate rows are [hidden | gate]: the epilogue's gate input is the second half. + assert (gate.path, gate.start, gate.stop) == ("proj.ffn_up_gate", 512, 1024) + assert (up.path, up.start, up.stop) == ("proj.ffn_up_gate", 0, 512) + assert specs["mlp"].swiglu and not specs["qkv"].swiglu + + +def test_rope_table_becomes_per_channel_cos_and_sin(): + rope = _rope_table(5) + runtime = types.SimpleNamespace(__dict__={}) + cos, sin = moss_image._rope(runtime, rope) + assert cos.shape == sin.shape == (5, 128) + assert torch.equal(cos[:, :64], cos[:, 64:]) and torch.equal(sin[:, :64], sin[:, 64:]) + assert torch.equal(sin[:, :64], rope[0, :, :64].float()) + assert torch.equal(cos[:, :64], rope[0, :, 64:].float()) + assert moss_image._rope(runtime, rope)[0] is cos + + +def test_multiple_samples_are_rejected(): + with pytest.raises(NotImplementedError, match="one sample"): + moss_image._block(nn.Module(), torch.zeros(2, 4, 8), None, None) + + +def test_situ_epilogue_matches_eager_rounding(): + _require_cuda() + from orbitquant.kernels import triton_int8_gemm as gemm + + torch.manual_seed(3) + rows, k, n = 300, 512, 1024 + a = torch.randint(-127, 128, (rows, k), device="cuda", dtype=torch.int8) + b = torch.randint(-127, 128, (n, k), device="cuda", dtype=torch.int8) + a_s = torch.rand(rows, device="cuda") * 1e-3 + 1e-4 + b_s = torch.rand(n, device="cuda") + 0.5 + alpha = 0.37 + y = ((((a.float() @ b.float().T) * a_s[:, None]) * b_s[None, :]) * alpha).to(torch.bfloat16) + half = gemm.SWIGLU_INTERLEAVE + chunks = y.reshape(rows, -1, 2, half) + gate, up = chunks[:, :, 0].reshape(rows, -1), chunks[:, :, 1].reshape(rows, -1) + expected = _situ(up, gate) + actual = gemm.matmul_int8_scaled_with_triton( + a, b, a_s, b_s, alpha=alpha, epilogue=gemm.EPILOGUE_SITU + ) + assert actual.shape == expected.shape + # libdevice and torch tanh agree to an ulp of fp32, so a rare bf16 tie can round apart. + mismatch = (actual != expected).float().mean().item() + assert mismatch < 1e-3 + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=1e-2) + + +def _quantized_block(dim=256, ffn=512, heads=2): + from orbitquant import OrbitQuantConfig + from orbitquant.layers import OrbitQuantLinear + + torch.manual_seed(0) + block = _Block(dim, ffn, heads).to(torch.bfloat16) + with torch.no_grad(): + for norm in (block.norm1, block.norm2, block.q_norm, block.k_norm): + norm.weight.copy_(torch.rand_like(norm.weight) * 0.5 + 0.75) + reference = _Block(dim, ffn, heads).to(torch.bfloat16) + reference.load_state_dict(block.state_dict()) + config = OrbitQuantConfig() + for name, linear in list(block.proj.items()): + block.proj[name] = OrbitQuantLinear.from_linear( + linear, config=config, module_name=f"proj.{name}" + ) + model = nn.Module() + model.layers = nn.ModuleList([block]) + model.quantization_config = config + type(model).__name__ # noqa: B018 + return model, reference + + +def test_fused_block_tracks_the_eager_block(): + _require_cuda() + from orbitquant.fused.groups import PackedGroup + from orbitquant.fused.install import _attach + from orbitquant.fused.runtime import FusedRuntime + + model, reference = _quantized_block() + model, reference = model.cuda(), reference.cuda() + block = model.layers[0] + runtime = FusedRuntime.from_config(model.quantization_config, attention="sdpa") + modules = nn.ModuleDict( + { + name: PackedGroup.from_sources( + block, list(spec.sources), swiglu=spec.swiglu, activation_bits=8 + ) + for name, spec in moss_image.groups("block", block).items() + } + ) + moss_image.prepare(block, "block") + _attach(model, block, moss_image, "block", modules, runtime) + seq, dim = 300, 256 + x = torch.randn(1, seq, dim, device="cuda", dtype=torch.bfloat16) + rope = _rope_table(seq, device="cuda") + modulation = (torch.randn(1, 6, dim, device="cuda") * 0.3).to(torch.bfloat16) + with torch.no_grad(): + expected = reference(x, rope, modulation) + actual = block(x, rope, modulation) + assert actual.shape == expected.shape and actual.dtype == torch.bfloat16 + update, ref_update = (actual - x).float().flatten(), (expected - x).float().flatten() + assert F.cosine_similarity(update, ref_update, dim=0) > 0.98 + + +def test_fuse_linears_replaces_projections_with_int8_gemms(): + _require_cuda() + from orbitquant import OrbitQuantConfig + from orbitquant.fused import FusedLinear, fuse_linears, release + from orbitquant.layers import OrbitQuantLinear + + torch.manual_seed(1) + config = OrbitQuantConfig() + dense = nn.Sequential( + nn.Linear(4096, 1024, bias=False), nn.SiLU(), nn.Linear(1024, 4096, bias=True) + ).to(torch.bfloat16) + model = nn.Sequential(*[m for m in dense]) + for index in (0, 2): + model[index] = OrbitQuantLinear.from_linear( + dense[index], config=config, module_name=str(index) + ) + model = model.cuda() + dense = dense.cuda() + model.quantization_config = config + assert fuse_linears(model) == 2 + assert isinstance(model[0], FusedLinear) and isinstance(model[2], FusedLinear) + x = torch.randn(2, 37, 4096, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + actual, expected = model(x), dense(x) + assert actual.shape == expected.shape == (2, 37, 4096) + assert F.cosine_similarity(actual.float().flatten(), expected.float().flatten(), dim=0) > 0.98 + release(model) diff --git a/tests/test_gptq.py b/tests/test_gptq.py new file mode 100644 index 0000000..e4f7a97 --- /dev/null +++ b/tests/test_gptq.py @@ -0,0 +1,185 @@ +import pytest +import torch + +from orbitquant.codebooks.lloyd_max import get_codebook +from orbitquant.config import OrbitQuantConfig +from orbitquant.gptq import ( + HessianCollector, + gptq_round, + load_rounded_codes, + round_groups, + update_state_dict, +) +from orbitquant.layers import OrbitQuantLinear +from orbitquant.modeling import quantize_model +from orbitquant.packing import unpack_lowbit +from orbitquant.rotations.rpbh import get_rpbh_rotation + +DIM = 64 + + +def _config(): + return OrbitQuantConfig( + weight_bits=4, + activation_bits=4, + block_size=DIM, + target_policy="universal", + runtime_mode="dequant_bf16", + ) + + +class _Block(torch.nn.Module): + def __init__(self): + super().__init__() + self.q = torch.nn.Linear(DIM, DIM, bias=False) + self.k = torch.nn.Linear(DIM, 2 * DIM, bias=False) + self.o = torch.nn.Linear(DIM, DIM, bias=False) + + def forward(self, x): + return self.o(torch.tanh(self.q(x)) + self.k(x)[..., :DIM]) + + +def _calibration(n=512): + # Correlated inputs: a few strong directions plus noise, as in real activations. + torch.manual_seed(1) + basis = torch.randn(4, DIM) + return torch.randn(n, 4) @ basis * 3 + 0.1 * torch.randn(n, DIM) + + +def test_gptq_beats_nearest_rounding_on_correlated_inputs(): + torch.manual_seed(0) + weight = torch.randn(32, DIM) + rotation = get_rpbh_rotation(dim=DIM, seed=0, block_size=DIM) + x = rotation.apply_to_activations(_calibration()) + hessian = x.T @ x + rotated = rotation.apply_to_weight(weight) + codes, scale = gptq_round(rotated, hessian, bits=4) + codebook = get_codebook(DIM, 4, 2) + gptq = codebook.centroids[codes.long()] * scale[:, None] + nearest = ( + codebook.centroids[codebook.quantize_indices(rotated / scale[:, None]).long()] + * scale[:, None] + ) + + def error(approx): + diff = rotated - approx + return torch.sum((diff @ hessian) * diff) / torch.sum((rotated @ hessian) * rotated) + + assert codes.dtype == torch.uint8 and codes.shape == rotated.shape + assert torch.equal(scale, rotated.norm(dim=1).to(torch.bfloat16).float()) + assert error(gptq) < 0.5 * error(nearest) + + +def test_hessian_collector_shares_one_hessian_per_group(): + block = _Block() + config = _config() + x = _calibration(64) + with HessianCollector(block, [["q", "k"], ["o"]], config) as hessians: + block(x) + rotation = get_rpbh_rotation(dim=DIM, seed=config.rotation_seed, block_size=config.block_size) + xr = rotation.apply_to_activations(x) + torch.testing.assert_close(hessians["q"], xr.T @ xr) + assert hessians["k"] is hessians["q"] + assert hessians.tokens["q"] == 64 + assert not block.q._forward_pre_hooks + + +def test_rounded_codes_load_into_the_quantized_model(): + torch.manual_seed(0) + block = _Block() + config = _config() + groups = [["q", "k"], ["o"]] + with HessianCollector(block, groups, config) as hessians: + block(_calibration()) + codes = round_groups(block, groups, hessians, config) + assert codes["q"].gptq_error < codes["q"].rtn_error + quantize_model(block, config, quantization_device=None) + load_rounded_codes(block, codes) + module = block.k + assert isinstance(module, OrbitQuantLinear) + unpacked = unpack_lowbit(module.packed_weight_indices, bits=4, length=2 * DIM * DIM).reshape( + 2 * DIM, DIM + ) + expected = unpack_lowbit(codes["k"].packed, bits=4, length=2 * DIM * DIM).reshape(2 * DIM, DIM) + assert torch.equal(unpacked, expected) + centroids = get_codebook(DIM, 4, 2).centroids + dequantized = module._dequantize_weight( + device=torch.device("cpu"), dtype=torch.float32, remember=False + ) + torch.testing.assert_close( + dequantized, centroids[expected.long()] * codes["k"].row_norms.float()[:, None] + ) + + +def test_sequential_rounding_writes_rounded_weights_back(): + torch.manual_seed(0) + block = _Block() + config = _config() + with HessianCollector(block, [["o"]], config) as hessians: + block(_calibration()) + codes = round_groups(block, [["o"]], hessians, config, sequential=True) + rotation = get_rpbh_rotation(dim=DIM, seed=config.rotation_seed, block_size=config.block_size) + indices = unpack_lowbit(codes["o"].packed, bits=4, length=DIM * DIM).reshape(DIM, DIM).long() + rounded = get_codebook(DIM, 4, 2).centroids[indices] * codes["o"].row_norms.float()[:, None] + torch.testing.assert_close(block.o.weight.data, rotation.apply_inverse_to_weight(rounded)) + + +def test_update_state_dict_keeps_shapes_and_dtypes(): + torch.manual_seed(0) + block = _Block() + config = _config() + with HessianCollector(block, [["o"]], config) as hessians: + block(_calibration()) + codes = round_groups(block, [["o"]], hessians, config) + state = { + "o.packed_weight_indices": torch.zeros(DIM * DIM // 2, dtype=torch.uint8), + "o.row_norms": torch.zeros(DIM, dtype=torch.bfloat16), + } + update_state_dict(state, codes) + assert state["o.packed_weight_indices"].shape == (DIM * DIM // 2,) + assert state["o.row_norms"].dtype == torch.bfloat16 + assert torch.equal(state["o.packed_weight_indices"], codes["o"].packed) + + +def test_batched_rounding_matches_one_group_at_a_time(): + torch.manual_seed(0) + rotation = get_rpbh_rotation(dim=DIM, seed=0, block_size=DIM) + weights, hessians = [], [] + for seed in range(3): + torch.manual_seed(seed) + weights.append(rotation.apply_to_weight(torch.randn(16, DIM))) + x = rotation.apply_to_activations(_calibration(256) + seed) + hessians.append(x.T @ x) + codes, scale = gptq_round(torch.stack(weights), torch.stack(hessians), bits=4) + for g in range(3): + one_codes, one_scale = gptq_round(weights[g], hessians[g], bits=4) + assert torch.equal(codes[g], one_codes) + assert torch.equal(scale[g], one_scale) + + +def test_round_groups_batches_same_shape_groups(): + torch.manual_seed(0) + blocks = torch.nn.ModuleList([_Block() for _ in range(3)]) + config = _config() + groups = [[f"{i}.q", f"{i}.k"] for i in range(3)] + [[f"{i}.o"] for i in range(3)] + x = _calibration() + with HessianCollector(blocks, groups, config) as hessians: + for block in blocks: + block(x) + batched = round_groups(blocks, groups, hessians, config, release=False, batch=8) + single = round_groups(blocks, groups, hessians, config, batch=1) + for name, codes in batched.items(): + assert torch.equal(codes.packed, single[name].packed) + assert torch.equal(codes.row_norms, single[name].row_norms) + + +def test_damping_grows_until_the_hessian_factorizes(): + torch.manual_seed(0) + rotation = get_rpbh_rotation(dim=DIM, seed=0, block_size=DIM) + weight = rotation.apply_to_weight(torch.randn(8, DIM)) + hessian = torch.eye(DIM) + # Four negative eigenvalues of -0.5: positive definite only from about 0.55 relative damping. + hessian[range(DIM - 4, DIM), range(DIM - 4, DIM)] = -0.5 + codes, scale, damping = gptq_round(weight, hessian, bits=4, return_damping=True) + assert codes.shape == weight.shape + assert damping.item() == pytest.approx(1.0) diff --git a/uv.lock b/uv.lock index 08d3609..be721b1 100644 --- a/uv.lock +++ b/uv.lock @@ -844,7 +844,7 @@ wheels = [ [[package]] name = "orbitquant" -version = "0.13.0" +version = "0.14.0" source = { editable = "." } dependencies = [ { name = "huggingface-hub" },