Working E2E pybind runner

Author: jackzhxng

optimum/executorch/modeling.py

@@ -35,14 +35,15 @@
     add_start_docstrings,
 )
 from transformers.configuration_utils import PretrainedConfig
+from transformers.tokenization_utils import PreTrainedTokenizer
 from transformers.utils import is_offline_mode
 
 from executorch.extension.pybindings.portable_lib import ExecuTorchModule, _load_for_executorch
 from executorch.kernels import quantized  # noqa
 
 from ..exporters import TasksManager
 from ..exporters.executorch import main_export
-from ..exporters.executorch.utils import verify_eos_tokens_in_tokenizer
+from ..exporters.executorch.utils import verify_eos_tokens_in_pretrained_tokenizer
 from ..modeling_base import FROM_PRETRAINED_START_DOCSTRING, OptimizedModel
 from ..utils.file_utils import find_files_matching_pattern
 from .stats import Stats
@@ -525,7 +526,7 @@ def generate(
 
     def text_generation(
         self,
-        tokenizer: "PreTrainedTokenizer",
+        tokenizer: PreTrainedTokenizer,
         prompt: str,
         echo: bool = True,
         max_seq_len: Optional[int] = None,
@@ -745,7 +746,7 @@ def generate(
 
     def text_generation(
         self,
-        tokenizer: "PreTrainedTokenizer",
+        tokenizer: PreTrainedTokenizer,
         prompt: str,
         echo: bool = True,
         max_seq_len: Optional[int] = None,
@@ -772,7 +773,7 @@ def text_generation(
             raise ValueError(
                 f"The tokenizer's bos_token_id={self.tokenizer.bos_token_id} must be the same as the model's bos_token_id={self.bos_token_id}."
             )
-        if not verify_eos_tokens_in_tokenizer(self.eos_token_ids, self.tokenizer):
+        if not verify_eos_tokens_in_pretrained_tokenizer(self.eos_token_ids, self.tokenizer):
             raise ValueError(
                 f"The tokenizer's eos_token_id does not match with the model's eos_token_ids={self.eos_token_ids}."
             )
@@ -1066,7 +1067,7 @@ def generate(
 
     def transcribe(
         self,
-        tokenizer: "PreTrainedTokenizer",
+        tokenizer: PreTrainedTokenizer,
         input_features: torch.Tensor,
         echo: bool = True,
         max_seq_len: Optional[int] = None,
@@ -1197,7 +1198,7 @@ def forward(
 
     def generate(
         self,
-        tokenizer: "PretrainedTokenizer",
+        tokenizer: PreTrainedTokenizer,
         input_ids: torch.LongTensor,
         pixel_values: Optional[torch.FloatTensor] = None,
         max_new_tokens: int = 100,
@@ -1237,31 +1238,37 @@ class ExecuTorchModelForMultiModalToText(ExecuTorchModelBase):
 
     def __init__(self, models: Dict[str, "ExecuTorchModule"], config: "PretrainedConfig"):
         super().__init__(models=models, config=config)
-        if not hasattr(self, "decoder"):
-            raise AttributeError("Expected attribute 'decoder' not found in the instance.")
-        if not hasattr(self, "token_embeddings"):
-            raise AttributeError("Expected attribute 'token_embeddings' not found in the instance.")
-        if not hasattr(self, "audio_encoder"):
-            raise AttributeError("Expected attribute 'audio_encoder' not found in the instance.")
-        metadata = self.decoder.method_names()
+        # if not hasattr(self, "decoder"):
+        #     raise AttributeError("Expected attribute 'decoder' not found in the instance.")
+        # if not hasattr(self, "token_embeddings"):
+        #     raise AttributeError("Expected attribute 'token_embeddings' not found in the instance.")
+        # if not hasattr(self, "audio_encoder"):
+        #     raise AttributeError("Expected attribute 'audio_encoder' not found in the instance.")
+
+        # required_methods = ["decoder", "token_embeddings", "audio_encoder"]
+        # for required_method in required_methods:
+        #     if required_method not in self.model.method_names():
+        #         raise ValueError("Exported .pte file needs to containt 'decoder', 'token_embeddings', and 'audio_encoder' methods.")
+        
+        metadata = self.model.method_names()
         if "use_kv_cache" in metadata:
-            self.use_kv_cache = self.decoder.run_method("use_kv_cache")[0]
+            self.use_kv_cache = self.model.run_method("use_kv_cache")[0]
         if "get_max_seq_len" in metadata:
-            self.max_cache_size = self.decoder.run_method("get_max_seq_len")[0]
+            self.max_cache_size = self.model.run_method("get_max_seq_len")[0]
         if "get_max_batch_size" in metadata:
-            self.max_batch_size = self.decoder.run_method("get_max_batch_size")[0]
+            self.max_batch_size = self.model.run_method("get_max_batch_size")[0]
         if "get_dtype" in metadata:
-            self.dtype = self.decoder.run_method("get_dtype")[0]
+            self.dtype = self.model.run_method("get_dtype")[0]
         if "get_bos_id" in metadata:
-            self.bos_token_id = self.decoder.run_method("get_bos_id")[0]
+            self.bos_token_id = self.model.run_method("get_bos_id")[0]
         if "get_eos_id" in metadata:
-            self.eos_token_id = self.decoder.run_method("get_eos_id")[0]
+            self.eos_token_id = self.model.run_method("get_eos_id")[0]
         if "get_vocab_size" in metadata:
-            self.vocab_size = self.decoder.run_method("get_vocab_size")[0]
+            self.vocab_size = self.model.run_method("get_vocab_size")[0]
         if "max_hidden_seq_length" in metadata:
-            self.max_hidden_seq_length = self.decoder.run_method("max_hidden_seq_length")[0]
+            self.max_hidden_seq_length = self.model.run_method("max_hidden_seq_length")[0]
         if "decoder_start_token_id" in metadata:
-            self.decoder_start_token_id = self.decoder.run_method("decoder_start_token_id")[0]
+            self.decoder_start_token_id = self.model.run_method("decoder_start_token_id")[0]
 
     def forward(
         self,
@@ -1300,25 +1307,28 @@ def generate(
             )
             max_seq_len = self.max_cache_size
 
+        # Prefill.
         self.stats.on_sampling_begin()
         logits = self.forward(
-            input_ids=torch.tensor(prompt_tokens, dtype=torch.long, device=self.device).unsqueeze(0),
-            cache_position=torch.arange(len(prompt_tokens), dtype=torch.long, device=self.device),
+            input_ids=torch.tensor(prompt_tokens, dtype=torch.long, device=self.device),
+            cache_position=torch.arange(len(prompt_tokens[0]), dtype=torch.long, device=self.device),
             input_features=input_features,
         )
         self.stats.on_sampling_end()
-        next_token = torch.argmax(logits, dim=-1)[0, -1].item()
         self.stats.on_prompt_eval_end()
-        first_token_generated = False
 
-        generated_tokens = prompt_tokens + [next_token]
+        next_token = torch.argmax(logits[:, -1, :], dim=-1).item()
+        generated_tokens = [next_token]
+        print(self.tokenizer.decode([next_token]), end="")
 
-        while len(generated_tokens) < max_seq_len:
+        # Token-by-token generation.
+        first_token_generated = False
+        while len(generated_tokens) + len(prompt_tokens) < max_seq_len:
             self.stats.on_sampling_begin()
             logits = self.forward(
                 input_ids=torch.tensor([next_token], dtype=torch.long, device=self.device).unsqueeze(0),
                 cache_position=torch.tensor(
-                    [pos_base + len(generated_tokens) - 1],
+                    [pos_base + len(generated_tokens) + len(prompt_tokens) - 1],
                     dtype=torch.long,
                     device=self.device,
                 ),
@@ -1328,20 +1338,20 @@ def generate(
                 self.stats.on_first_token()
                 first_token_generated = True
 
-            next_token = torch.argmax(logits, dim=-1).item()
+            next_token = torch.argmax(logits[:, -1, :], dim=-1).item()
             generated_tokens.append(next_token)
+            print(self.tokenizer.decode([next_token]), end="")
 
-            if next_token in self.eos_token_ids:
+            if next_token == self.eos_token_id:
                 break
 
         self.stats.set_num_generated_tokens(len(generated_tokens) - len(prompt_tokens))
-
         return generated_tokens if echo else generated_tokens[len(prompt_tokens) :]
 
     def text_generation(
         self,
         processor: "ProcessorMixin",
-        tokenizer: "PreTrainedTokenizer",
+        tokenizer: PreTrainedTokenizer,
         input_conversation: List[Dict],
         echo: bool = True,
         max_seq_len: Optional[int] = None,
@@ -1368,22 +1378,21 @@ def text_generation(
             raise ValueError(
                 f"The tokenizer's bos_token_id={self.tokenizer.bos_token_id} must be the same as the model's bos_token_id={self.bos_token_id}."
             )
-        if not verify_eos_tokens_in_tokenizer(self.eos_token_ids, self.tokenizer):
+        if isinstance(self.tokenizer, PreTrainedTokenizer) and verify_eos_tokens_in_pretrained_tokenizer(self.eos_token_id, self.tokenizer):
             raise ValueError(
-                f"The tokenizer's eos_token_id does not match with the model's eos_token_ids={self.eos_token_ids}."
+                f"The tokenizer's eos_token_id does not match with the model's eos_token_id={self.eos_token_id}."
             )
 
         # Reset stats for a new generation
         self.stats.reset()
         self.stats.on_inference_start()
 
         inputs = processor.apply_chat_template(input_conversation)
-        prompt_tokens = self.tokenizer.encode(inputs["input_ids"])
         self.stats.on_token_encode_end()
-        self.stats.set_num_prompt_tokens(len(prompt_tokens))
+        self.stats.set_num_prompt_tokens(len(inputs["input_ids"][0]))
 
         generated_tokens = self.generate(
-            prompt_tokens=prompt_tokens,
+            prompt_tokens=inputs["input_ids"],
             input_features=inputs["input_features"],
             echo=echo,
             max_seq_len=max_seq_len,

optimum/exporters/executorch/integrations.py

@@ -671,17 +671,19 @@ def export(
             exported_programs["token_embeddings"] = token_embeddings_exported_program
 
             # 3. Export encoder.
+            input_ids = torch.zeros_like(inputs_embeds[:, :, 0], dtype=torch.long)
+            input_ids[0, 1] = self.config.audio_token_id  # Make sure we don't have an all-false mask for the imput_embeds.
             if isinstance(self.model, VoxtralForConditionalGeneration):
                 # TODO(JZ): specific to Voxtral, should generalize.
                 chunk_length = self.model.audio_tower.config.max_source_positions * self.model.audio_tower.conv1.stride[0] * self.model.audio_tower.conv2.stride[0]
                 encoder_input_kwargs = {
                     "input_features": torch.rand(3, 128, chunk_length),  # (bsz, features, seq_len)
                     "inputs_embeds": inputs_embeds,
-                    "input_ids": inputs_embeds[:, :, 0],
+                    "input_ids": input_ids,
                 }
 
                 max_audio_len = 150  # In s, should be a multiple of 30. TODO(JZ): make this configurable top-level.
-                max_seq_len = self.metadata.get("get_max_seq_len") - 1  # TODO(JZ): why - 1? Copied from Gemma3 draft PR.
+                max_seq_len = self.metadata.get("get_max_seq_len")
                 dynamic_shapes = {
                     "input_features": {
                         0: torch.export.Dim("enc_batch_size_dim", min=1, max=max_audio_len//30),

optimum/exporters/executorch/recipes/xnnpack.py

@@ -70,7 +70,7 @@ def _lower_to_executorch(
     ) -> Dict[str, ExecutorchProgram]:
         backend_config_dict = {
             "extract_delegate_segments": True,
-            # "memory_planning_pass": MemoryPlanningPass(alloc_graph_input=False),
+            "memory_planning_pass": MemoryPlanningPass(alloc_graph_input=False),
         }
         if parse(executorch_version.__version__).base_version > "0.6.0":
             backend_config_dict["do_quant_fusion_and_const_prop"] = True
@@ -89,14 +89,16 @@ def _lower_to_executorch(
         et_prog = et_prog.to_executorch(
             config=ExecutorchBackendConfig(**backend_config_dict),
         )
-        logging.debug(
-            f"\nExecuTorch program for {pte_name}.pte: {et_prog.exported_program().graph_module}"
-        )
-        delegation_info = get_delegation_info(et_prog.exported_program().graph_module)
-        logging.debug(f"\nDelegation info Summary for {pte_name}.pte: {delegation_info.get_summary()}")
-        logging.debug(
-            f"\nDelegation info for {pte_name}.pte: {tabulate(delegation_info.get_operator_delegation_dataframe(), headers='keys', tablefmt='fancy_grid')}"
-        )
+        for method in et_prog.methods:
+            logging.debug(f"---------------------- Method: {method} ----------------------")
+            logging.debug(
+                f"\nExecuTorch program for {pte_name}.pte: {et_prog.exported_program(method).graph_module}"
+            )
+            delegation_info = get_delegation_info(et_prog.exported_program(method).graph_module)
+            logging.debug(f"\nDelegation info Summary for {pte_name}.pte: {delegation_info.get_summary()}")
+            logging.debug(
+                f"\nDelegation info for {pte_name}.pte: {tabulate(delegation_info.get_operator_delegation_dataframe(), headers='keys', tablefmt='fancy_grid')}"
+            )
         return {pte_name: et_prog}
 
     exported_progs = model.export()

optimum/exporters/executorch/tasks/multimodal_text_to_text.py

@@ -61,6 +61,8 @@ def load_multimodal_text_to_text_model(model_name_or_path: str, **kwargs):
     attn_implementation = kwargs.get("attn_implementation", "custom_sdpa" if use_custom_sdpa else "sdpa")
     cache_implementation = kwargs.get("cache_implementation", "static")
     use_custom_sdpa = use_custom_sdpa or attn_implementation == "custom_sdpa"
+    qlinear_config = kwargs.get("qlinear", None)
+    qembedding_config = kwargs.get("qembedding", None)
     max_length = kwargs.get("max_length", 2048)
     config = kwargs.get("config") or AutoConfig.from_pretrained(model_name_or_path)
 
@@ -111,8 +113,6 @@ def load_multimodal_text_to_text_model(model_name_or_path: str, **kwargs):
 
     # TODO: Move quantization recipe out for better composability.
     # TODO: Should switch to `TorchAoConfig` once the quant issue on final lm_head layer is fixed.
-    qlinear_config = kwargs.get("qlinear", None)
-    qembedding_config = kwargs.get("qembedding", None)
     if qlinear_config or qembedding_config:
         # TODO: Update torchao to use 0.11.0 once released
         if parse(torchao.__version__) < parse("0.11.0.dev0"):

optimum/exporters/executorch/utils.py

@@ -16,6 +16,7 @@
 
 import torch
 from transformers import GenerationConfig, PretrainedConfig
+from transformers.tokenization_utils import PreTrainedTokenizer
 
 
 def save_config_to_constant_methods(
@@ -65,7 +66,7 @@ def save_config_to_constant_methods(
     return {k: v for k, v in {**metadata, **kwargs}.items() if v is not None}
 
 
-def verify_eos_tokens_in_tokenizer(model_eos_ids: List[int], tokenizer) -> bool:
+def verify_eos_tokens_in_pretrained_tokenizer(model_eos_ids: List[int], tokenizer: PreTrainedTokenizer) -> bool:
     """
     Verifies that the model's EOS token IDs are present in the tokenizer's
     set of potential end-of-sequence tokens.

tests/models/test_modeling_voxtral.py

@@ -27,7 +27,7 @@
 import transformers
 from executorch.extension.pybindings.portable_lib import ExecuTorchModule
 from packaging.version import parse
-from transformers import AutoTokenizer, AutoProcessor
+from transformers import AutoConfig, AutoTokenizer, AutoProcessor
 from transformers.testing_utils import slow
 
 from optimum.utils.import_utils import is_transformers_version
@@ -42,7 +42,7 @@
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.DEBUG)
 
 
 @pytest.mark.skipif(
@@ -71,15 +71,16 @@ def __init__(self, *args, **kwargs):
     #     reason="Only available on transformers >= 4.53.0.dev0 and torchao >= 0.11.0",
     # )
     # @pytest.mark.skipif(is_linux_ci, reason="OOM on linux runner")
-    @pytest.mark.skip()
+    # @pytest.mark.skip()
     def test_voxtral_audio_text_to_text_generation_with_custom_sdpa_kv_cache_8da4w_8we_exported_program(self):
         model_id = "mistralai/Voxtral-Mini-3B-2507"
+        config = AutoConfig.from_pretrained(model_id)
         module = load_multimodal_text_to_text_model(
             model_id,
             use_custom_sdpa=True,
             use_custom_kv_cache=True,
             qlinear=True,
-            qembedding_config=True,
+            qembedding=True,
         )
 
         res = module.export()
@@ -166,11 +167,12 @@ def test_voxtral_audio_text_to_text_generation_with_custom_sdpa_kv_cache_8da4w_8
         ]
 
         model = ExecuTorchModelForMultiModalToText.from_pretrained(
-            model_id,
+            # model_id,
+            "/Users/jackzhxng/Documents/voxtral",  # Load already exported model in local file path.
             recipe="xnnpack",
             attn_implementation="custom_sdpa",
             use_custom_kv_cache=True,
-            **{"qlinear": True, "qembeeding": True, "task": "multimodal-text-to-text"},
+            **{"qlinear": True, "qembedding": True, "task": "multimodal-text-to-text"},
         )
         self.assertIsInstance(model, ExecuTorchModelForMultiModalToText)
         self.assertIsInstance(model.model, ExecuTorchModule)