Remove vocab_size from dynamic axes

Signed-off-by: quic-sanising <[email protected]>

Author: quic-sanising

QEfficient/transformers/models/modeling_auto.py

@@ -255,7 +255,6 @@ def export(self, export_dir: Optional[str] = None) -> str:
                 fbs if self.continuous_batching else bs, self.model.config.vocab_size, dtype=torch.int32)
             dynamic_axes["repetition_penalty_retain_state"] = {
                 0: "full_batch_size" if self.continuous_batching else "batch_size",
-                1: "vocab_size",
             }
             output_names.append("repetition_penalty_retain_state_RetainedState")
 
@@ -266,7 +265,6 @@ def export(self, export_dir: Optional[str] = None) -> str:
                 fbs if self.continuous_batching else bs, self.model.config.vocab_size, dtype=torch.int32)
             dynamic_axes["presence_penalty_retain_state"] = {
                 0: "full_batch_size" if self.continuous_batching else "batch_size",
-                1: "vocab_size",
             }
             output_names.append("presence_penalty_retain_state_RetainedState")
 
@@ -374,7 +372,6 @@ def compile(
         }
         if self.include_sampler:
              prefill_specialization.update({
-                 "vocab_size": self.model.config.vocab_size,
                  "max_top_k_ids": constants.Constants.MAX_TOP_K_IDS,
              })
         prefill_specialization.update({"num_logits_to_keep": 1})
@@ -396,7 +393,6 @@ def compile(
             }
             if self.include_sampler:
                 decode_specialization.update({
-                    "vocab_size": self.model.config.vocab_size,
                     "max_top_k_ids": constants.Constants.MAX_TOP_K_IDS,
                 })
             if self.continuous_batching:

QEfficient/utils/constants.py

@@ -44,7 +44,7 @@ def get_models_dir():
 
 QEFF_MODELS_DIR = get_models_dir()
 
-ONNX_EXPORT_EXAMPLE_BATCH_SIZE = 1
+ONNX_EXPORT_EXAMPLE_BATCH_SIZE = 2
 ONNX_EXPORT_EXAMPLE_SEQ_LEN = 32
 ONNX_EXPORT_EXAMPLE_FBS = 4
 ONNX_EXPORT_EXAMPLE_NLK = 2  # Number of Logits to Keep