fix: set quantization recipe for the softmax with embedding's value

PiperOrigin-RevId: 756931600

Author: zichuan-wei

ai_edge_torch/generative/quantize/quant_recipe.py

@@ -16,9 +16,12 @@
 from dataclasses import dataclass
 from typing import Optional, Union
 
+from ai_edge_torch.generative.layers import model_config
 from ai_edge_torch.generative.quantize import quant_attrs
 from ai_edge_torch.generative.quantize import supported_schemes
 
+ModelConfig = model_config.ModelConfig
+
 
 @dataclass
 class LayerQuantRecipe:
@@ -52,7 +55,7 @@ def __str__(self):
         f'w:{self.weight_dtype.name}, '
         f'{self.mode.name}, '
         f'{self.algorithm.name}, '
-        f'{self.granularity.name}'
+        f'{self.granularity.name}, '
         f'{self.block_size}'
     )
     return f'{base_str})'
@@ -133,6 +136,7 @@ class GenerativeQuantRecipe:
   feedforward: Union[
       Optional[LayerQuantRecipe], Optional[dict[int, LayerQuantRecipe]]
   ] = None
+  _model_config: Optional[ModelConfig] = None
 
   def __str__(self):
     return f"""GenerativeQuantRecipe(

ai_edge_torch/generative/quantize/quant_recipes.py

@@ -63,6 +63,7 @@ def all_supported_int4_dynamic_block_recipe(
       generative_recipe=quant_recipe.GenerativeQuantRecipe(
           default=quant_recipe_utils.create_layer_quant_int4_dynamic_block(
               block_size
-          )
+          ),
+          embedding=quant_recipe_utils.create_layer_quant_int8_dynamic(),
       )
   )

ai_edge_torch/generative/test/test_quantize.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 
 import ai_edge_torch
-from ai_edge_torch import config
 from ai_edge_torch.generative.examples.test_models import toy_model  # NOQA
 from ai_edge_torch.generative.quantize import quant_recipe
 from ai_edge_torch.generative.quantize import quant_recipe_utils

ai_edge_torch/generative/utilities/converter.py

@@ -270,6 +270,7 @@ def _export_helper(
   )
 
   quant_config = quant_recipes.full_int8_dynamic_recipe() if quantize else None
+  quant_config._model_config = config
 
   # For export, we create a module that captures any non-exportable,
   # arugments, e.g. the generation config object.
@@ -334,5 +335,7 @@ def _export_helper(
         sample_kwargs=sample_kwargs,
     )
 
-  edge_model = converter.convert(quant_config=quant_config)
+  edge_model = converter.convert(
+      quant_config=quant_config,
+  )
   edge_model.export(output_file)

ai_edge_torch/lowertools/_shim.py

@@ -50,7 +50,7 @@ def exported_programs_to_tflite(
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     _tfl_converter_flags: Optional[dict[str, Any]] = None,
-    _saved_model_dir: Optional[str] = None
+    _saved_model_dir: Optional[str] = None,
 ):
   """Converts a list of ExportedProgram to a TFLite model.
 

ai_edge_torch/lowertools/translate_recipe.py

@@ -29,6 +29,8 @@
 _ATTENTION_REGEX_STR = 'ai_edge_torch.generative.layers.attention'
 _FEEDFORWARD_REGEX_STR = 'ai_edge_torch.generative.layers.feed_forward'
 _EMBEDDING_REGEX_STR = 'Embedding_tok_embedding'
+# TODO: b/415833584 - Improve the regex for pre-softmax layer.
+_DECODE_LOGITS_REGEX_STR = 'StatefulPartitionedCall'
 _ANY_TWO_DIGITS_REGEX_STR = '\d{1,2}'
 
 
@@ -95,10 +97,11 @@ def _set_quant_config(
     rm: quantizer.recipe_manager.RecipeManager,
     layer_recipe: quant_recipe.LayerQuantRecipe,
     regex: str,
+    operation_name: _OpName = _OpName.ALL_SUPPORTED,
 ):
   rm.add_quantization_config(
       regex=regex,
-      operation_name=_OpName.ALL_SUPPORTED,
+      operation_name=operation_name,
       op_config=_OpQuantConfig(
           weight_tensor_config=_TensorQuantConfig(
               num_bits=_get_nbits_from_dtype(layer_recipe.weight_dtype),
@@ -126,6 +129,16 @@ def translate_to_ai_edge_recipe(
 
   if recipe.embedding is not None:
     _set_quant_config(rm, recipe.embedding, _EMBEDDING_REGEX_STR)
+    if (
+        recipe._model_config is not None
+        and recipe._model_config.lm_head_share_weight_with_embedding
+    ):
+      _set_quant_config(
+          rm,
+          recipe.embedding,
+          _DECODE_LOGITS_REGEX_STR,
+          _OpName.FULLY_CONNECTED,
+      )
 
   if recipe.attention is not None:
     if isinstance(recipe.attention, dict):