[LLM] Report skipped quantization layers and enable 8da8w

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 = [

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(

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:

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:

examples/models/llama/source_transformation/test_quantize.py

@@ -0,0 +1,45 @@
+# 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
+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))
+
+        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()

extension/llm/export/config/llm_config.py

@@ -465,6 +465,7 @@ class QuantizationConfig:
     QMODE_OPTIONS: ClassVar[List[str]] = [
         "int8",
         "8da4w",
+        "8da8w",
         "8da4w-gptq",
         "4w",
     ]