MOSS-Image fused family, fused single projections, optional GPTQ rounding (0.14.0)

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).
 

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"

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()
 

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",

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)

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]

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

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: