Added experimental orbit quant

toolkit/util/ostris_quant.py

@@ -0,0 +1,133 @@
+"""
+Quantization-agnostic custom quantized linear.
+
+OstrisLinear is a drop-in nn.Linear replacement whose weight is held by a pluggable
+quantizer backend (OstrisQuantizer). Backends own the quantized representation
+(buffers + per-module attributes) and how the forward pass computes W x from it; the
+module and the rest of the toolkit stay backend agnostic. The first backend is
+OrbitQuant (toolkit/util/orbit_quant.py) via the orbit2/orbit3/orbit4 qtypes; add new
+backends by implementing OstrisQuantizer and resolving them in get_ostris_quantizer.
+
+Modules are converted in place by convert_linear_to_ostris via class swap, so the
+original module object (and any references to it, e.g. LoRA org_module or parent
+containers) stays valid.
+"""
+
+from typing import Dict, Optional
+
+import torch
+import torch.nn.functional as F
+
+
+class OstrisQuantizer:
+    """Base class for weight quantization backends used by OstrisLinear.
+
+    Backends are stateless with respect to tensors: everything tensor-shaped must be
+    registered as a buffer on the module inside quantize_ (so device moves and dtype
+    casts through nn.Module._apply keep working), and read back off the module in the
+    other methods. One backend instance may be shared by many modules.
+    """
+
+    def can_quantize(self, module: torch.nn.Linear) -> bool:
+        """Whether this backend can quantize the given linear (e.g. shape constraints)."""
+        return True
+
+    def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None:
+        """Build the quantized representation of weight_fp32 and attach it to the
+        module (register_buffer for tensors, plain attributes for scalars). Called
+        while the module is still an nn.Linear, before the weight param is removed."""
+        raise NotImplementedError
+
+    def dequantize(self, module: "OstrisLinear") -> torch.Tensor:
+        """Reconstruct the full weight in the original basis, in float32."""
+        raise NotImplementedError
+
+    def requantize_(self, module: "OstrisLinear", fp_weight: torch.Tensor) -> None:
+        """Re-quantize in place from a full precision weight in the original basis
+        (used by the continuous merge/reset method)."""
+        raise NotImplementedError
+
+    def forward(self, module: "OstrisLinear", x: torch.Tensor) -> torch.Tensor:
+        # default: dequantize per forward and run a plain linear. backends can
+        # override with a cheaper formulation. the weight is frozen, so build it
+        # outside autograd; gradients still flow to x through the matmul
+        with torch.no_grad():
+            w = self.dequantize(module).to(x.dtype)
+        return F.linear(x, w, module.bias)
+
+
+class OstrisLinear(torch.nn.Linear):
+    """A linear layer whose weight is quantized by an OstrisQuantizer backend.
+
+    Never instantiate directly: created in place by convert_linear_to_ostris. The
+    weight parameter is removed; the quantized representation lives in backend-owned
+    buffers, plus:
+      ostris_quantizer   the backend instance
+      ostris_orig_dtype  dtype of the original weight (used for dequantized views)
+    """
+
+    is_ostris_quantized = True
+
+    @torch.no_grad()
+    def dequantize_weight(self) -> torch.Tensor:
+        """Reconstruct the weight in the original basis and dtype."""
+        return self.ostris_quantizer.dequantize(self).to(self.ostris_orig_dtype)
+
+    @property
+    def weight(self):
+        # materializes the full dequantized weight. kept for code that inspects the
+        # weight (shape/dtype/device) and for the network merge paths, which detect
+        # the marker via toolkit.util.quantize.is_quantized_tensor
+        w = self.dequantize_weight()
+        w._is_ostris_weight = True
+        return w
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.ostris_quantizer.forward(self, x)
+
+    @torch.no_grad()
+    def requantize_(self, fp_weight: torch.Tensor) -> None:
+        self.ostris_quantizer.requantize_(self, fp_weight)
+
+    def _save_to_state_dict(self, destination, prefix, keep_vars):
+        # emit a plain full precision weight so full-model saves need no special casing
+        destination[prefix + "weight"] = self.dequantize_weight()
+        if self.bias is not None:
+            destination[prefix + "bias"] = self.bias if keep_vars else self.bias.detach()
+
+
+def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]:
+    """Resolve a qtype string to a quantizer backend instance, or None if the qtype
+    does not belong to a custom backend. Add new backends here."""
+    from toolkit.util.orbit_quant import ORBIT_QTYPES, OrbitQuantizer
+    from toolkit.util.orbit_vq_quant import ORBIT_VQ_QTYPES, OrbitVQQuantizer
+
+    if qtype in ORBIT_QTYPES:
+        return OrbitQuantizer(ORBIT_QTYPES[qtype])
+    if qtype in ORBIT_VQ_QTYPES:
+        return OrbitVQQuantizer(**ORBIT_VQ_QTYPES[qtype])
+    return None
+
+
+@torch.no_grad()
+def convert_linear_to_ostris(module: torch.nn.Linear, quantizer: OstrisQuantizer) -> bool:
+    """Quantize an nn.Linear in place (class swap). Returns True if the module was
+    converted (or already was), False if it is not a candidate."""
+    if isinstance(module, OstrisLinear):
+        return True
+    weight = getattr(module, "weight", None)
+    if not isinstance(weight, torch.nn.Parameter) or not weight.dtype.is_floating_point:
+        return False
+    if type(weight.data) is not torch.Tensor:
+        # already holds a quantized tensor subclass (e.g. torchao)
+        return False
+    if not quantizer.can_quantize(module):
+        return False
+    quantizer.quantize_(module, weight.data.to(torch.float32))
+    module.ostris_quantizer = quantizer
+    module.ostris_orig_dtype = weight.dtype
+    del module._parameters["weight"]
+    if module.bias is not None:
+        module.bias.requires_grad_(False)
+    module.__class__ = OstrisLinear
+    return True