oblib=vious

Author: xadupre

_unittests/ut_torch_models/test_tiny_llms_bypassed.py

@@ -62,6 +62,7 @@ def test_export_phi2_2_bypassed(self):
             {"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
         )
         model(**torch_deepcopy(inputs))
+        ds = use_dyn_not_str(ds)
         with torch_export_patches(patch_transformers=True, stop_if_static=1) as modificator:
             inputs = modificator(inputs)
             ep = torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds, strict=False)

onnx_diagnostic/torch_models/untrained/llm_phi2.py

@@ -9,6 +9,7 @@ def get_phi2(
     sequence_length: int = 30,
     sequence_length2: int = 3,
     dynamic_rope: bool = False,
+    use_dim_not_dynamic: bool = False,
     **kwargs,
 ) -> Dict[str, Any]:
     """
@@ -18,6 +19,8 @@ def get_phi2(
     :param sequence_length: sequence length
     :param sequence_length2: new sequence length
     :param dynamic_rope: use dynamic rope (see :class:`transformers.LlamaConfig`)
+    :param use_dim_not_dynamic: uses ``torch.export.Dim`` and not a string for the batch size,
+        the sequence length and the cache length
     :param kwargs: to overwrite the configuration, example ``num_hidden_layers=1``
     :return: dictionary
 
@@ -62,7 +65,7 @@ def get_phi2(
     n_layers = config["num_hidden_layers"]
     num_key_value_heads = config["num_key_value_heads"]
 
-    if batch_size == 1:
+    if use_dim_not_dynamic:
         batch = torch.export.Dim("batch", min=1, max=1024)
         seq_length = torch.export.Dim("seq_length", min=1, max=4096)
         cache_length = torch.export.Dim("cache_length", min=1, max=4096)