Export with no kv cache + non-strict load checkpoint

Author: jackzhxng

examples/models/llama3_2_vision/model.py

@@ -109,7 +109,7 @@ def __init__(self, **kwargs):
         # Load checkpoint.
         missing, unexpected = self.model_.load_state_dict(
             checkpoint,
-            strict=True,
+            strict=False,
             assign=True,
         )
         if kwargs.get("verbose", False):
@@ -139,6 +139,7 @@ def __init__(self, **kwargs):
             self.model_.setup_caches(
                 batch_size=1,
                 dtype=self.dtype,
+                decoder_max_seq_len=self.max_seq_len,
             )
 
     def get_eager_model(self) -> torch.nn.Module:
@@ -153,21 +154,29 @@ def get_example_inputs(self):
     def get_example_kwarg_inputs(self):
         # For export we must use the prefill versions of the
         # causal mask and input_pos.
-        return {
-            "mask": self.causal_mask[None, :32],
-            # "encoder_input": None,
-            # "encoder_mask": None,
-            "input_pos": self.input_pos[None, :32]
-        }
+        if self.use_kv_cache:
+            return {
+                "input_pos": self.input_pos[None, :32],
+                "mask": self.causal_mask[None, :32],
+                # "encoder_input": None,
+                # "encoder_mask": None,
+            }
+        else:
+            return None
 
     def get_dynamic_shapes(self):
         batch_size = 1
         dim_seq_len = torch.export.Dim("token_dim", min=1, max=self.max_seq_len)
-        dynamic_shapes = {
-            "tokens": {0: batch_size, 1: dim_seq_len},
-            # "encoder_input": {0: 1, 1: dim_enc, 2: 4096},
-            # "encoder_mask": {0: 1, 1: dim, 2: dim_enc},
-            "mask": {0: batch_size, 1: dim_seq_len, 2: dim_seq_len},
-            "input_pos" : {0: batch_size, 1: dim_seq_len},
-        }
+        if self.use_kv_cache:
+            dynamic_shapes = {
+                "tokens": {0: batch_size, 1: dim_seq_len},
+                # "encoder_input": {0: 1, 1: dim_enc, 2: 4096},
+                # "encoder_mask": {0: 1, 1: dim, 2: dim_enc},
+                "mask": {0: batch_size, 1: dim_seq_len, 2: None},
+                "input_pos" : {0: batch_size, 1: dim_seq_len},
+            }
+        else:
+            dynamic_shapes = {
+                "tokens": {0: batch_size, 1: dim_seq_len},
+            }
         return dynamic_shapes