Improvements to Gemma4 quant functions (#20768)

Differential Revision: D110919317

Pull Request resolved: #20768

Author: metascroy

backends/cuda/TARGETS

@@ -2,6 +2,20 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 oncall("executorch")
 
+runtime.python_library(
+    name = "int4_dispatch",
+    srcs = [
+        "int4_dispatch.py",
+    ],
+    visibility = [
+        "//executorch/...",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//pytorch/ao:torchao",
+    ],
+)
+
 runtime.python_library(
     name = "cuda_passes",
     srcs = [

examples/models/gemma4_31b/quant/quantize.py

@@ -7,7 +7,8 @@
 """Quantize weights to torchao tensor subclasses.
 
 ``quantize_weight`` quantizes a single tensor given a ``QuantConfig``,
-returning an ``Int4Tensor`` (4-bit) or ``IntxUnpackedToInt8Tensor`` (8-bit).
+returning an ``Int4Tensor`` (4-bit) or ``IntxUnpackedToInt8Tensor`` (6- or
+8-bit).
 
 ``quantize_model`` walks a model's parameters, applies a ``QuantRecipe``,
 and returns a single state dict containing both quantized subclass tensors
@@ -152,11 +153,14 @@ def _to_intx_tensor(
     weight: torch.Tensor,
     config: QuantConfig,
 ) -> torch.Tensor:
-    """Quantize 8-bit and wrap in IntxUnpackedToInt8Tensor.
+    """Quantize to 6- or 8-bit and wrap in IntxUnpackedToInt8Tensor.
 
     Quantizes in float32 for numerical precision, then constructs the
     subclass directly. We avoid ``from_hp`` because it quantizes in the
     input dtype (bf16), which loses precision for small-magnitude weights.
+
+    Sub-byte data (e.g. 6-bit) is still stored unpacked in an int8 container;
+    ``target_dtype`` records the true bit width for the export/runtime path.
     """
     from torchao.quantization import IntxUnpackedToInt8Tensor
     from torchao.quantization.quant_primitives import (
@@ -165,19 +169,23 @@ def _to_intx_tensor(
         quantize_affine,
     )
 
+    qmin = -(1 << (config.bits - 1))
+    qmax = (1 << (config.bits - 1)) - 1
+    target_dtype = getattr(torch, f"int{config.bits}")
+
     if config.method == "hqq":
         if not config.symmetric:
             raise ValueError(
-                "8-bit HQQ only supports symmetric quantization "
-                "(HQQ_SCALE_ONLY). Use method='min_max' for asymmetric 8-bit."
+                "intx HQQ only supports symmetric quantization "
+                "(HQQ_SCALE_ONLY). Use method='min_max' for asymmetric intx."
             )
         from torchao.quantization.quant_primitives import (
             _choose_qparams_and_quantize_scale_only_hqq,
         )
 
         w2d = weight.float().reshape(-1, weight.shape[-1])
         qdata, scale = _choose_qparams_and_quantize_scale_only_hqq(
-            w2d, [1, config.group_size], -128, 127
+            w2d, [1, config.group_size], qmin, qmax
         )
         qdata = qdata.to(torch.int8).reshape(weight.shape)
         scale = scale.to(torch.bfloat16).reshape(weight.shape[0], -1)
@@ -190,8 +198,8 @@ def _to_intx_tensor(
             mapping,
             block_size,
             target_dtype=torch.int8,
-            quant_min=-128,
-            quant_max=127,
+            quant_min=qmin,
+            quant_max=qmax,
             scale_dtype=torch.bfloat16,
             zero_point_dtype=torch.int8,
         )
@@ -201,8 +209,8 @@ def _to_intx_tensor(
             scale,
             zero_point,
             output_dtype=torch.int8,
-            quant_min=-128,
-            quant_max=127,
+            quant_min=qmin,
+            quant_max=qmax,
         )
         N, n_groups = weight.shape[0], weight.shape[-1] // config.group_size
         scale = scale.reshape(N, n_groups)
@@ -212,7 +220,7 @@ def _to_intx_tensor(
         qdata=qdata,
         scale=scale,
         zero_point=zero_point,
-        target_dtype=torch.int8,
+        target_dtype=target_dtype,
         block_size=(1, config.group_size),
         dtype=torch.bfloat16,
         activation_quantization=None,
@@ -222,9 +230,10 @@ def _to_intx_tensor(
 def quantize_weight(weight: torch.Tensor, config: QuantConfig) -> torch.Tensor:
     """Quantize ``weight`` to a torchao tensor subclass.
 
-    Returns ``Int4Tensor`` for 4-bit or ``IntxUnpackedToInt8Tensor`` for 8-bit.
+    Returns ``Int4Tensor`` for 4-bit or ``IntxUnpackedToInt8Tensor`` for 6-
+    and 8-bit.
     """
-    if config.bits == 8:
+    if config.bits in (6, 8):
         return _to_intx_tensor(weight, config)
 
     if config.bits != 4:

examples/models/gemma4_31b/quant/recipe.py

@@ -24,7 +24,7 @@ class QuantConfig:
     carry their own metadata. This is purely for driving ``quantize_weight``.
     """
 
-    bits: int  # 4 or 8
+    bits: int  # 4, 6, or 8
     group_size: int
     symmetric: bool  # True = no zero point
     method: str  # "min_max" | "hqq"