fix

Author: xadupre

_unittests/ut_torch_models/test_llm_phi2.py

@@ -22,7 +22,7 @@ def test_get_phi2(self):
     @ignore_warnings(UserWarning)
     @requires_transformers("4.54")
     @requires_torch("2.9.99")
-    def test_export_phi2_1(self):
+    def test_export_phi2_1_batch_size_1(self):
         # exporting vmap does not work
         data = get_phi2(num_hidden_layers=2)
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
@@ -35,6 +35,22 @@ def test_export_phi2_1(self):
             )
         assert ep
 
+    @ignore_warnings(UserWarning)
+    @requires_transformers("4.54")
+    @requires_torch("2.9.99")
+    def test_export_phi2_1_batch_size_2(self):
+        # exporting vmap does not work
+        data = get_phi2(num_hidden_layers=2, batch=2)
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        self.assertEqual(
+            {"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
+        )
+        with torch_export_patches(patch_transformers=True):
+            ep = torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
+            )
+        assert ep
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/torch_models/untrained/llm_phi2.py

@@ -62,9 +62,14 @@ def get_phi2(
     n_layers = config["num_hidden_layers"]
     num_key_value_heads = config["num_key_value_heads"]
 
-    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)
+    if batch_size == 1:
+        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)
+    else:
+        batch = "batch"
+        seq_length = "seq_length"
+        cache_length = "cache_length"
 
     shapes = {
         "input_ids": {0: batch, 1: seq_length},