Patch

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: larryliu0820

examples/models/llama2/export_llama_lib.py

@@ -29,7 +29,7 @@
     Transformer,
 )
 from executorch.exir.backend.backend_details import CompileSpec
-
+from executorch.exir.passes import *
 from executorch.sdk.etrecord import generate_etrecord
 from executorch.util.activation_memory_profiler import generate_memory_trace
 from sentencepiece import SentencePieceProcessor
@@ -541,7 +541,7 @@ def _prepare_for_llama_export(modelname: str, args) -> LlamaEdgeManager:
         bitwidth = int(bitwidth)
         transforms.append(
             lambda model: EmbeddingOnlyInt8QuantHandler(
-                model, bitwidth=bitwidth, group_size=group_size
+                model, bitwidth=bitwidth, group_size=group_size, packed=(bitwidth==4),
             ).quantized_model()
         )
 

examples/models/llama2/quantize.py

@@ -438,6 +438,6 @@ def forward(self, indices: torch.Tensor) -> torch.Tensor:
                 self.weight, self.scales, None, 0, 0, indices, dtype=self.dtype
             )
         else:  # 4bit packed
-            return torch.ops.llama_quantized.embedding_4bit.dtype(
+            return torch.ops.quantized_decomposed.embedding_4bit.dtype(
                 self.weight, self.scales, None, 0, 0, indices, dtype=self.dtype
             )