From f19d3d3f95473d8f1f038b28c57b4513f639febb Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Mon, 24 Aug 2026 10:37:20 -0700 Subject: [PATCH 1/2] Report skipped LLM quantization layers --- examples/models/llama/BUCK | 12 +++++++ .../models/llama/config/test_llm_config.py | 5 +++ examples/models/llama/export_llama_lib.py | 7 ++-- .../llama/source_transformation/quantize.py | 28 ++++++++++++++++ .../source_transformation/test_quantize.py | 33 +++++++++++++++++++ extension/llm/export/config/llm_config.py | 1 + 6 files changed, 84 insertions(+), 2 deletions(-) create mode 100644 examples/models/llama/source_transformation/test_quantize.py diff --git a/examples/models/llama/BUCK b/examples/models/llama/BUCK index a2dabb954b7..ee961344bbc 100644 --- a/examples/models/llama/BUCK +++ b/examples/models/llama/BUCK @@ -257,6 +257,18 @@ fbcode_target(_kind = runtime.python_test, ], ) +fbcode_target(_kind = runtime.python_test, + name = "quantize_source_transform_test", + srcs = [ + "source_transformation/test_quantize.py", + ], + deps = [ + ":source_transformation", + "//caffe2:torch", + "//pytorch/ao:torchao", + ], +) + fbcode_target(_kind = runtime.python_test, name = "quantized_sdpa_with_kv_cache_test", srcs = [ diff --git a/examples/models/llama/config/test_llm_config.py b/examples/models/llama/config/test_llm_config.py index 4bbd56cdb66..70b0f9db0db 100644 --- a/examples/models/llama/config/test_llm_config.py +++ b/examples/models/llama/config/test_llm_config.py @@ -110,6 +110,11 @@ def test_shared_embedding_without_lowbit(self): class TestValidConstruction(unittest.TestCase): + def test_8da8w_qmode(self): + qcfg = QuantizationConfig(qmode="8da8w") + + self.assertEqual(qcfg.qmode, "8da8w") + def test_valid_llm_config(self): LlmConfig( base=BaseConfig( diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index 60a2c2e87c9..096af3e4a49 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -32,7 +32,10 @@ from executorch.exir.backend.partitioner import Partitioner from executorch.exir.passes.init_mutable_pass import InitializedMutableBufferPass from executorch.extension.llm.export.builder import DType, LLMEdgeManager -from executorch.extension.llm.export.config.llm_config import LlmConfig +from executorch.extension.llm.export.config.llm_config import ( + LlmConfig, + QuantizationConfig, +) from executorch.extension.llm.export.partitioner_lib import ( get_coreml_partitioner, get_mps_partitioner, @@ -1052,7 +1055,7 @@ def get_quantizer_and_quant_params(llm_config): def _qmode_type(value): - choices = ["int8", "8da4w", "8da4w-gptq", "4w"] + choices = QuantizationConfig.QMODE_OPTIONS patterns = [r"torchao:8da(\d+)w", r"torchao:fpa(\d+)w"] if value in choices: diff --git a/examples/models/llama/source_transformation/quantize.py b/examples/models/llama/source_transformation/quantize.py index 6bcf35b2a69..15af48fcc4a 100644 --- a/examples/models/llama/source_transformation/quantize.py +++ b/examples/models/llama/source_transformation/quantize.py @@ -153,6 +153,18 @@ def filter_fn(m, fqn): return True return m.weight.shape[1] % group_size == 0 + linear_quantization_decisions = [ + filter_fn(module, fqn) + for fqn, module in model.named_modules() + if isinstance(module, nn.Linear) + and "lora_a" not in fqn.split(".") + and "lora_b" not in fqn.split(".") + ] + quantized_linear_count = sum(linear_quantization_decisions) + skipped_linear_count = ( + len(linear_quantization_decisions) - quantized_linear_count + ) + weight_dtype = torch.int4 if qmode == "8da4w" else torch.int8 quantize_( model, @@ -169,6 +181,22 @@ def filter_fn(m, fqn): ), filter_fn=filter_fn, ) + if skipped_linear_count: + logging.warning( + "%s quantization: quantized %d linear layer(s), skipped %d linear " + "layer(s) because in_features is not divisible by group_size=%d.", + qmode, + quantized_linear_count, + skipped_linear_count, + group_size, + ) + else: + logging.info( + "%s quantization: quantized %d linear layer(s), skipped 0 linear " + "layer(s).", + qmode, + quantized_linear_count, + ) # TODO: deal with checkpoint / computation dtype decoupling. if verbose: diff --git a/examples/models/llama/source_transformation/test_quantize.py b/examples/models/llama/source_transformation/test_quantize.py new file mode 100644 index 00000000000..050a790ca0e --- /dev/null +++ b/examples/models/llama/source_transformation/test_quantize.py @@ -0,0 +1,33 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import unittest +from unittest.mock import patch + +import torch.nn as nn + +from executorch.examples.models.llama.source_transformation.quantize import quantize + + +class TestQuantize(unittest.TestCase): + @patch("torchao.quantization.quantize_") + def test_reports_linears_skipped_by_group_size(self, mock_quantize): + model = nn.Sequential(nn.Linear(128, 16), nn.Linear(96, 16)) + + with self.assertLogs(level=logging.WARNING) as logs: + quantize(model, qmode="8da4w", group_size=128) + + self.assertIn( + "8da4w quantization: quantized 1 linear layer(s), skipped 1 linear " + "layer(s) because in_features is not divisible by group_size=128.", + logs.output[0], + ) + mock_quantize.assert_called_once() + + +if __name__ == "__main__": + unittest.main() diff --git a/extension/llm/export/config/llm_config.py b/extension/llm/export/config/llm_config.py index 8e04f359935..bd7f87f70ec 100644 --- a/extension/llm/export/config/llm_config.py +++ b/extension/llm/export/config/llm_config.py @@ -465,6 +465,7 @@ class QuantizationConfig: QMODE_OPTIONS: ClassVar[List[str]] = [ "int8", "8da4w", + "8da8w", "8da4w-gptq", "4w", ] From fecee84c1c37f04b8a5fb0fee7f930bc50967dff Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Mon, 24 Aug 2026 10:41:09 -0700 Subject: [PATCH 2/2] Test 8da8w quantization export --- .../llama/source_transformation/test_quantize.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/examples/models/llama/source_transformation/test_quantize.py b/examples/models/llama/source_transformation/test_quantize.py index 050a790ca0e..d05941c7c5a 100644 --- a/examples/models/llama/source_transformation/test_quantize.py +++ b/examples/models/llama/source_transformation/test_quantize.py @@ -8,12 +8,24 @@ import unittest from unittest.mock import patch +import torch import torch.nn as nn from executorch.examples.models.llama.source_transformation.quantize import quantize class TestQuantize(unittest.TestCase): + def test_8da8w_quantizes_and_exports(self): + model = nn.Sequential(nn.Linear(128, 32), nn.ReLU(), nn.Linear(32, 16)) + inputs = (torch.randn(2, 128),) + + quantize(model, qmode="8da8w", group_size=0) + exported = torch.export.export(model, inputs) + output = exported.module()(*inputs) + + self.assertEqual(output.shape, (2, 16)) + self.assertTrue(torch.isfinite(output).all()) + @patch("torchao.quantization.quantize_") def test_reports_linears_skipped_by_group_size(self, mock_quantize): model = nn.Sequential(nn.Linear(128, 16), nn.Linear(96, 16))