Works

Author: jackzhxng

optimum/commands/export/executorch.py

@@ -101,7 +101,7 @@ def parse_args_executorch(parser):
             "Options:\n"
             "  8da4w - 8-bit dynamic activation, 4-bit weight\n"
             "  8da8w - 8-bit dynamic activation, 8-bit weight\n"
-            "  8da4w,8da8w - 8-bit dynamic activation, 4-bit weight and 8-bit weight\n"
+            "  8da4w,8da8w - 8-bit dynamic activation, 4-bit weight; fallback on 8-bit dynamic activation, 8-bit weight per-channel where group size doesn't divide block size cleanly \n"
             "  4w    - 4-bit weight only\n"
             "  8w    - 8-bit weight only"
         ),

optimum/exporters/executorch/quantization.py

@@ -40,7 +40,7 @@ def quantize_model_(
         if qlinear_config == "8w":
             assert (
                 qembedding_group_size == 0
-            ), "8-bit embedding quantization only supports per-channel at the moment, please use qembedding_group_size = 0."
+            ), "8-bit embedding quantization only supports per-token at the moment, please use qembedding_group_size = 0."
         if qembedding_group_size == 0:
             embedding_weight_granularity = PerAxis(0)
         else:
@@ -67,6 +67,7 @@ def quantize_model_(
         )
 
     if qlinear_config:
+
         def build_linear_config(config_key: str, granularity):
             if config_key == "8da4w":
                 return Int8DynamicActivationIntxWeightConfig(
@@ -94,9 +95,7 @@ def build_linear_config(config_key: str, granularity):
         if any(cfg == "" for cfg in qlinear_configs):
             raise ValueError("Linear quantization config entries must be non-empty.")
         if len(qlinear_configs) > 2:
-            raise ValueError(
-                "Expected at most one fallback linear quantization config, got more than one comma."
-            )
+            raise ValueError("Expected at most one fallback linear quantization config, got more than one comma.")
 
         primary_linear_config_key = qlinear_configs[0]
         fallback_linear_config_key = qlinear_configs[1] if len(qlinear_configs) == 2 else None
@@ -109,16 +108,16 @@ def build_linear_config(config_key: str, granularity):
                 )
                 fallback_linear_config_key = None
         else:
-            assert qlinear_group_size % 2 == 0, f"Linear quantization group size must be a multiple of 2, got {qlinear_group_size}."
+            assert (
+                qlinear_group_size % 2 == 0
+            ), f"Linear quantization group size must be a multiple of 2, got {qlinear_group_size}."
             linear_weight_granularity = PerGroup(qlinear_group_size)
 
         logging.info("Quantizing linear layers.")
-        primary_linear_config = build_linear_config(
-            primary_linear_config_key, linear_weight_granularity
-        )
+        primary_linear_config = build_linear_config(primary_linear_config_key, linear_weight_granularity)
 
         # First, quantize layers that are compatible with group quantization
-        def quant_filter(module, fqn):
+        def per_group_filter(module, fqn):
             if isinstance(module, torch.nn.Linear):
                 # Check if hidden dimension is divisible by group size
                 # For Linear layers, weight shape is [out_features, in_features]
@@ -129,20 +128,16 @@ def quant_filter(module, fqn):
         quantize_(
             eager_model,
             primary_linear_config,
-            filter_fn=quant_filter,
+            filter_fn=per_group_filter,
         )
 
         # Then, quantize incompatible layers using the fallback per-axis config
         if fallback_linear_config_key is not None:
-            fallback_linear_config = build_linear_config(
-                fallback_linear_config_key, PerAxis(0)
-            )
-            
-            def per_channel_filter(module, fqn):
+            fallback_linear_config = build_linear_config(fallback_linear_config_key, PerAxis(0))
+
+            def per_token_filter(module, fqn):
                 if isinstance(module, torch.nn.Linear):
-                    # Only quantize layers that are NOT compatible with group quantization
-                    # and haven't been quantized yet
-                    return not quant_filter(module, fqn)
+                    return module.weight.shape[1] % qlinear_group_size != 0
                 return False
 
             logging.info(
@@ -152,7 +147,7 @@ def per_channel_filter(module, fqn):
             quantize_(
                 eager_model,
                 fallback_linear_config,
-                filter_fn=per_channel_filter,
+                filter_fn=per_token_filter,
             )
 
     unwrap_tensor_subclass(eager_model)

optimum/exporters/executorch/tasks/multimodal_text_to_text.py

@@ -14,6 +14,7 @@
 
 
 import json
+import logging
 import os.path
 
 import torchao
@@ -201,15 +202,24 @@ def load_multimodal_text_to_text_model(model_name_or_path: str, **kwargs):
     qembedding_group_size = kwargs.get("qembedding_group_size", None)
 
     # Quantize decoder linear weights.
+    if qlinear_config:
+        logging.info("Quantizing decoder linears...")
     quantize_decoder_kwargs = {
         "eager_model": getattr(eager_model, decoder_name),
         "qlinear_config": qlinear_config,
     }
+    quantize_lm_head_kwargs = {
+        "eager_model": eager_model.lm_head,
+        "qlinear_config": qlinear_config,
+    }
     if qlinear_group_size is not None:
         quantize_decoder_kwargs["qlinear_group_size"] = qlinear_group_size
     quantize_model_(**quantize_decoder_kwargs)
+    quantize_model_(**quantize_lm_head_kwargs)
 
     # Quantize encoder linear weights.
+    if qlinear_encoder_config:
+        logging.info("Quantizing encoder linears...")
     quantize_encoder_kwargs = {
         "eager_model": getattr(eager_model, encoder_name),
         "qlinear_config": qlinear_encoder_config,
@@ -219,6 +229,8 @@ def load_multimodal_text_to_text_model(model_name_or_path: str, **kwargs):
     quantize_model_(**quantize_encoder_kwargs)
 
     # Quantize decoder embeddings.
+    if qembedding_config:
+        logging.info("Quantizing decoder embeddings...")
     quantize_decoder_embedding_kwargs = {
         "eager_model": eager_model,
         "qembedding_config": qembedding_config,
@@ -227,14 +239,6 @@ def load_multimodal_text_to_text_model(model_name_or_path: str, **kwargs):
         quantize_decoder_embedding_kwargs["qembedding_group_size"] = qembedding_group_size
     quantize_model_(**quantize_decoder_embedding_kwargs)
 
-    # Quantize lm_head
-    if hasattr(eager_model, "lm_head") and qlinear_config is not None:
-        quantize_model_(
-            eager_model=eager_model.lm_head,
-            qlinear_config=qlinear_config,
-            qlinear_group_size=qlinear_group_size if qlinear_group_size is not None else 0,
-        )
-    print(eager_model)
     return MultiModalTextToTextExportableModule(
         model=eager_model,
         modality="audio" if audio_encoder_name else "vision",

tests/models/test_modeling_gemma3.py

@@ -309,9 +309,9 @@ def test_gemma3_image_vision_with_custom_sdpa_kv_cache_8da4w_8we(self):
             use_custom_kv_cache=True,
             qlinear="8da4w",
             qlinear_group_size=32,
-            # Can't quantize the encoder a the moment, hidden dim of 4304 doesn't fit ExecuTorch's
-            # XNNPack 32-group size quantized kernels. See https://github.com/pytorch/executorch/issues/14221.
-            qembedding_config="8w",
+            qlinear_encoder="8da4w,8da8w",
+            qlinear_encoder_group_size=32,
+            qembedding="8w",
         )
 
         # Generate