fix quantization_config error when running convert with new quantization flag

PiperOrigin-RevId: 757886360

Author: zichuan-wei

ai_edge_torch/generative/quantize/quant_recipes.py

@@ -27,43 +27,56 @@
   )
 """
 
+from typing import Optional
+from ai_edge_torch.generative.layers import model_config
 from ai_edge_torch.generative.quantize import quant_recipe
 from ai_edge_torch.generative.quantize import quant_recipe_utils
 from ai_edge_torch.quantize import quant_config
 
 
-def full_int8_dynamic_recipe() -> quant_config.QuantConfig:
+def full_int8_dynamic_recipe(
+    mcfg: Optional[model_config.ModelConfig] = None,
+) -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
       generative_recipe=quant_recipe.GenerativeQuantRecipe(
           default=quant_recipe_utils.create_layer_quant_int8_dynamic(),
+          _model_config=mcfg,
       )
   )
 
 
-def full_int8_weight_only_recipe() -> quant_config.QuantConfig:
+def full_int8_weight_only_recipe(
+    mcfg: Optional[model_config.ModelConfig] = None,
+) -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
       generative_recipe=quant_recipe.GenerativeQuantRecipe(
           default=quant_recipe_utils.create_layer_quant_int8_weight_only(),
+          _model_config=mcfg,
       )
   )
 
 
-def full_fp16_recipe() -> quant_config.QuantConfig:
+def full_fp16_recipe(
+    mcfg: Optional[model_config.ModelConfig] = None,
+) -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
       generative_recipe=quant_recipe.GenerativeQuantRecipe(
-          default=quant_recipe_utils.create_layer_quant_fp16()
+          default=quant_recipe_utils.create_layer_quant_fp16(),
+          _model_config=mcfg,
       )
   )
 
 
 def all_supported_int4_dynamic_block_recipe(
     block_size: int,
+    mcfg: Optional[model_config.ModelConfig] = None,
 ) -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
       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(),
+          _model_config=mcfg,
       )
   )

ai_edge_torch/generative/utilities/converter.py

@@ -26,6 +26,7 @@
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.quantize import quant_recipes
 from ai_edge_torch.generative.utilities import export_config
+from ai_edge_torch.quantize import quant_config as qcfg
 import torch
 
 ExportConfig = export_config.ExportConfig
@@ -123,7 +124,8 @@ def define_conversion_flags(
 
 def get_quant_recipe_from_flag(
     quantize: str,
-) -> Optional[quant_recipes.QuantizationRecipe]:
+    model_config: cfg.ModelConfig,
+) -> Optional[qcfg.QuantConfig]:
   """Processes the quantization flag and returns the corresponding recipe.
 
   Args:
@@ -139,15 +141,19 @@ def get_quant_recipe_from_flag(
     case QuantizationName.NONE:
       return None
     case QuantizationName.DYNAMIC_INT8:
-      return quant_recipes.full_int8_dynamic_recipe()
+      return quant_recipes.full_int8_dynamic_recipe(mcfg=model_config)
     case QuantizationName.WEIGHT_ONLY_INT8:
-      return quant_recipes.full_int8_weight_only_recipe()
+      return quant_recipes.full_int8_weight_only_recipe(mcfg=model_config)
     case QuantizationName.FP16:
       return quant_recipes.full_fp16_recipe()
     case QuantizationName.DYNAMIC_INT4_BLOCK32:
-      return quant_recipes.full_int4_dynamic_block_recipe(32)
+      return quant_recipes.all_supported_int4_dynamic_block_recipe(
+          32, mcfg=model_config
+      )
     case QuantizationName.DYNAMIC_INT4_BLOCK128:
-      return quant_recipes.full_int4_dynamic_block_recipe(128)
+      return quant_recipes.all_supported_int4_dynamic_block_recipe(
+          128, mcfg=model_config
+      )
     case _:
       raise ValueError(f'Unsupported quantization flag: {quantize}')
 
@@ -351,8 +357,7 @@ def _export_helper(
       kv_layout=export_config.kvcache_layout,
   )
 
-  quant_config = get_quant_recipe_from_flag(quantize)
-  quant_config._model_config = config
+  quant_config = get_quant_recipe_from_flag(quantize, config)
 
   # For export, we create a module that captures any non-exportable,
   # arugments, e.g. the generation config object.