From c17d9d0d2189afaa62138d37a61206f3878dda5f Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Fri, 30 Oct 2026 16:28:01 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/native/test/samples/BUCK | 3 + backends/native/test/samples/linear.py | 19 ++-- backends/native/test/samples/record.py | 17 ++-- backends/native/test/samples/schemes.py | 114 +++++++++++++++++++++--- 4 files changed, 123 insertions(+), 30 deletions(-) diff --git a/backends/native/test/samples/BUCK b/backends/native/test/samples/BUCK index 326d33d26de..d1d4ab455ec 100644 --- a/backends/native/test/samples/BUCK +++ b/backends/native/test/samples/BUCK @@ -17,8 +17,11 @@ fbcode_target( "//caffe2:torch", "//executorch/backends/test/suite/operators:test_operators", "//executorch/exir/native:native", + "//executorch/extension/llm/export:gguf", "//pytorch/ao:torchao", "//pytorch/facto:facto", + "fbsource//third-party/pypi/gguf:gguf", + "fbsource//third-party/pypi/numpy:numpy", ], ) diff --git a/backends/native/test/samples/linear.py b/backends/native/test/samples/linear.py index c1f1dc60e48..9ee902e681c 100644 --- a/backends/native/test/samples/linear.py +++ b/backends/native/test/samples/linear.py @@ -15,6 +15,7 @@ suite_cases, ) from executorch.backends.test.suite.operators import test_linear +from executorch.extension.llm.export.gguf import QK_K from facto.inputgen.argument.type import ArgType from facto.inputgen.specs.model import ConstraintProducer as cp, InPosArg, OutArg, Spec @@ -29,8 +30,9 @@ def _unit(*constraints) -> list: ] -# in_features is a multiple of every scheme's group size, and out_features is -# large enough to fit an error's abs_bias and rel_bias from one sample. +# in_features is a multiple of every torchao scheme's group size and, in some +# tuples, of GGUF's block; out_features is large enough to fit an error's +# abs_bias and rel_bias from one sample. _FACTO_SPEC = Spec( op="linear.default", inspec=[ @@ -98,8 +100,13 @@ def cases() -> list[Case]: ), Case("outlier_channel", torch.nn.Linear(256, 128), outlier), Case("zero_row", torch.nn.Linear(256, 64), zero_row), - # FACTO varies one argument at a time, so only every fifth tuple is kept. - # Its values lie on a 1/8 grid, which makes float32 math exact: these - # cases cover shapes more than rounding. - *facto_cases(_FACTO_SPEC, _facto_case, lambda i, _: i % 5 == 0), + # FACTO varies one argument at a time, so only every fifth tuple is kept, + # plus those whose in_features fit GGUF blocks. Its values lie on a 1/8 + # grid, which makes float32 math exact: these cases cover shapes more + # than rounding. + *facto_cases( + _FACTO_SPEC, + _facto_case, + lambda i, args: i % 5 == 0 or args[1].size(1) % QK_K == 0, + ), ] diff --git a/backends/native/test/samples/record.py b/backends/native/test/samples/record.py index d3565a01667..ef439ce541f 100644 --- a/backends/native/test/samples/record.py +++ b/backends/native/test/samples/record.py @@ -22,7 +22,7 @@ import torch from executorch.backends.native.test.samples.cases import Case, Inputs, outputs -from executorch.backends.native.test.samples.schemes import Scheme +from executorch.backends.native.test.samples.schemes import int8_activations, Scheme from executorch.exir.native import Call, Sample, SampleSet, to_native, write_samples from torch.func import functional_call @@ -51,7 +51,6 @@ def record( dtypes: Iterable[str] = tuple(DTYPES), ) -> None: """Writes no package, removing any at `path`, when the scheme fits no case.""" - act_scheme = scheme.act_quant() programs, golden, precise, golden_act = {}, [], [], [] for case in cases: if not scheme.fits(case.model): @@ -65,9 +64,6 @@ def record( model = copy.deepcopy(case.model).to(dtype).eval() scheme.quantize(model) method = _Method({name: model}) - if act_scheme is not None: - act_model = copy.deepcopy(case.model).to(dtype).eval() - act_scheme.quantize(act_model) state = { k: scheme.decode(v).double() for k, v in method.state_dict().items() } @@ -81,10 +77,11 @@ def record( t.float() if t.is_floating_point() else t for t in outputs(ref) ] precise.append(Sample([Call(name, list(x), ref)])) - if act_scheme is not None: - golden_act.append( - Sample([Call(name, list(x), outputs(act_model(*x)))]) - ) + if scheme.weight_only: + with int8_activations(method): + golden_act.append( + Sample([Call(name, list(x), outputs(method(*x)))]) + ) if not programs: logger.info(f"{path.name}: no case fits the scheme, not written") path.unlink(missing_ok=True) @@ -99,7 +96,7 @@ def record( precise, "float64, weights dequantized, activations unquantized", source ), ) - if act_scheme is not None: + if scheme.weight_only: write_samples( str(path), "golden-int8_act", diff --git a/backends/native/test/samples/schemes.py b/backends/native/test/samples/schemes.py index 82d9c2d8cec..af4183da9b1 100644 --- a/backends/native/test/samples/schemes.py +++ b/backends/native/test/samples/schemes.py @@ -7,9 +7,15 @@ """Quantization schemes. Every op is recorded under every scheme, keeping the cases the scheme fits.""" +import contextlib +import tempfile +from collections.abc import Iterator from dataclasses import dataclass +import numpy as np import torch +from executorch.extension.llm.export.gguf import iter_gguf, QK_K +from gguf import GGML_QUANT_SIZES, GGMLQuantizationType, GGUFWriter from torchao.core.config import AOBaseConfig from torchao.quantization import ( Int8DynamicActivationIntxWeightConfig, @@ -17,12 +23,24 @@ quantize_, ) from torchao.quantization.granularity import PerGroup + +# Private, but it is the input quantize-dequantize that +# Int8DynamicActivationIntxWeightConfig linears run. +from torchao.quantization.quantize_.workflows.intx.intx_unpacked_to_int8_tensor import ( + _apply_int8_act_asym_per_token_quant_dequant, +) from torchao.utils import TorchAOBaseTensor class Scheme: """No quantization.""" + @property + def weight_only(self) -> bool: + """Whether the scheme quantizes only weights; `record` then adds + `golden-int8_act`.""" + return False + def fits(self, model: torch.nn.Module) -> bool: return True @@ -33,11 +51,6 @@ def decode(self, value: torch.Tensor) -> torch.Tensor: """A state tensor of a quantized model as a plain tensor.""" return value - def act_quant(self) -> "Scheme | None": - """This weight-only scheme's weights with int8 dynamic activations, or - None when the scheme is not weight-only.""" - return None - def _fits(model: torch.nn.Module, group: int | None) -> bool: """Whether `model` has an nn.Linear and every one's in_features split into @@ -57,6 +70,10 @@ class Torchao(Scheme): config: AOBaseConfig + @property + def weight_only(self) -> bool: + return isinstance(self.config, IntxWeightOnlyConfig) + def fits(self, model: torch.nn.Module) -> bool: granularity = ( self.config.granularity @@ -73,17 +90,84 @@ def quantize(self, model: torch.nn.Module) -> None: def decode(self, value: torch.Tensor) -> torch.Tensor: return value.dequantize() if isinstance(value, TorchAOBaseTensor) else value - def act_quant(self) -> Scheme | None: - if not isinstance(self.config, IntxWeightOnlyConfig): - return None - return Torchao( - Int8DynamicActivationIntxWeightConfig( - weight_dtype=self.config.weight_dtype, - weight_granularity=self.config.granularity, - weight_mapping_type=self.config.mapping_type, - weight_scale_dtype=self.config.scale_dtype, + +# fp16 super-block scales by byte offset in each block, kept in fp16's normal +# range; q4_k's dmin of 8 d centers random weights near zero. +_GGUF_SUPER_SCALES = {"q4_k": {0: 2**-12, 2: 2**-9}, "q6_k": {208: 2**-12}} + + +def _random_blocks( + ggml_type: str, rows: int, cols: int, rng: np.random.Generator +) -> np.ndarray: + """Valid `ggml_type` blocks for a [rows, cols] weight from random bytes.""" + _, block_bytes = GGML_QUANT_SIZES[GGMLQuantizationType[ggml_type.upper()]] + raw = rng.integers(0, 256, (rows, cols // QK_K, block_bytes), dtype=np.uint8) + for offset, scale in _GGUF_SUPER_SCALES[ggml_type].items(): + raw[..., offset : offset + 2] = np.array([scale], np.float16).view(np.uint8) + return raw.reshape(rows, -1) + + +def _load_gguf(model: torch.nn.Module, ggml_type: str) -> None: + """Writes random `ggml_type` weights for every nn.Linear to a .gguf file and + loads them back as a GGUF checkpoint would be.""" + linears = [m for m in model.modules() if isinstance(m, torch.nn.Linear)] + rng = np.random.default_rng(0) + with tempfile.TemporaryDirectory() as tmp: + path = f"{tmp}/weights.gguf" + writer = GGUFWriter(path, arch="linear") + for i, m in enumerate(linears): + writer.add_tensor( + f"linear.{i}.weight", + _random_blocks(ggml_type, *m.weight.shape, rng), + raw_dtype=GGMLQuantizationType[ggml_type.upper()], + ) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + for m, (_, weight) in zip(linears, iter_gguf(path), strict=True): + m.weight = torch.nn.Parameter( + weight.to(m.weight.dtype), requires_grad=False ) + + +@dataclass(frozen=True) +class Gguf(Scheme): + """Weight-only GGUF k-quant blocks (`q4_k` or `q6_k`) for every nn.Linear, + loaded from a .gguf file.""" + + ggml_type: str + + @property + def weight_only(self) -> bool: + return True + + def fits(self, model: torch.nn.Module) -> bool: + return _fits(model, QK_K) + + def quantize(self, model: torch.nn.Module) -> None: + _load_gguf(model, self.ggml_type) + + def decode(self, value: torch.Tensor) -> torch.Tensor: + return value.dequantize() if isinstance(value, TorchAOBaseTensor) else value + + +@contextlib.contextmanager +def int8_activations(model: torch.nn.Module) -> Iterator[None]: + """Quantizes every nn.Linear input as torchao's int8 dynamic activation + schemes do: int8, asymmetric, per token.""" + hooks = [ + m.register_forward_pre_hook( + lambda _, args: (_apply_int8_act_asym_per_token_quant_dequant(args[0]),) ) + for m in model.modules() + if isinstance(m, torch.nn.Linear) + ] + try: + yield + finally: + for hook in hooks: + hook.remove() def schemes() -> dict[str, Scheme]: @@ -105,4 +189,6 @@ def schemes() -> dict[str, Scheme]: weight_dtype=torch.int4, weight_granularity=PerGroup(128) ) ), + "gguf_q4_k": Gguf("q4_k"), + "gguf_q6_k": Gguf("q6_k"), }