ensuring small memory footprint for parallel test execution on A30

Signed-off-by: Lucas Liebenwein <[email protected]>

Author: lucaslie

tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py

@@ -435,15 +435,33 @@ def apply_rotary_pos_emb_ds(q, k, cos, sin, position_ids, unsqueeze_dim=1):
         "llm_models_subdir": "Qwen2.5-3B-Instruct",
         "model_kwargs": {
             "num_hidden_layers": 2,
+            "hidden_size": 64,
+            "intermediate_size": 128,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
         },
     },
     "mistralai/Mistral-Small-3.1-24B-Instruct-2503": {
         "llm_models_subdir": "Mistral-Small-3.1-24B-Instruct-2503",
         "model_factory": "AutoModelForImageTextToText",
         "compile_backend": "torch-simple",
         "model_kwargs": {
-            "text_config": {"num_hidden_layers": 2},
-            "vision_config": {"num_hidden_layers": 2},
+            "text_config": {
+                "num_hidden_layers": 2,
+                "head_dim": 64,
+                "hidden_size": 64,
+                "intermediate_size": 128,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 2,
+            },
+            "vision_config": {
+                "num_hidden_layers": 1,
+                "hidden_size": 64,
+                "head_dim": 32,
+                "image_size": 128,
+                "intermediate_size": 128,
+                "num_attention_heads": 2,
+            },
         },
     },
     "ibm-ai-platform/Bamba-9B-v2": {