Fix quantized embedding export logic (#3095)

Summary:
Add patches to make 4bit quantized embedding work for export. Fixed:
* Schema mismatch between functional embedding_4bit and out variant
* Set `packed=True` for 4bit quantization

Pull Request resolved: #3095

Reviewed By: mikekgfb

Differential Revision: D56340670

Pulled By: larryliu0820

fbshipit-source-id: c98623a9b7633fc5a6c390be1557213c719fa95a

Author: larryliu0820

examples/models/llama2/export_llama_lib.py

@@ -614,7 +614,10 @@ def _prepare_for_llama_export(modelname: str, args) -> LlamaEdgeManager:
         bitwidth = int(bitwidth)
         transforms.append(
             lambda model: EmbeddingQuantHandler(
-                model, bitwidth=bitwidth, group_size=group_size
+                model,
+                bitwidth=bitwidth,
+                group_size=group_size,
+                packed=(bitwidth == 4),
             ).quantized_model()
         )
 

exir/passes/_quant_patterns_and_replacements.py

@@ -189,7 +189,7 @@ def embedding_byte_dtype_out_meta(
 
 quantized_decomposed_lib.define(
     "embedding_4bit.dtype_out(Tensor weight, Tensor weight_scales, Tensor? weight_zero_points, "
-    "int weight_quant_min, int weight_quant_max, Tensor indices, ScalarType? dtype=None, *, Tensor(a!) out) -> Tensor(a!)",
+    "int weight_quant_min, int weight_quant_max, Tensor indices, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!)",
 )