onnx-diagnostic 0.5.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

onnx_diagnostic/tasks/__init__.py

@@ -39,9 +39,30 @@ def supported_tasks() -> List[str]:
 
 def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
     """Reduces a model size."""
+    head_size0 = (
+        config.head_dim
+        if hasattr(config, "head_dim") and config.head_dim
+        else (
+            config.hidden_size // config.num_attention_heads
+            if hasattr(config, "hidden_size") and hasattr(config, "num_attention_heads")
+            else None
+        )
+    )
     tasks = {mod.__TASK__: mod.reduce_model_config for mod in __TASKS__}
     assert task in tasks, f"Task {task!r} not found in {sorted(tasks)}"
-    return tasks[task](config)
+    res = tasks[task](config)
+    if head_size0 and "head_dim" in res:
+        head_size = (
+            config.head_dim
+            if hasattr(config, "head_dim") and config.head_dim
+            else config.hidden_size // config.num_attention_heads
+        )
+        assert head_size0 == head_size or head_size % 16 == 0, (
+            f"head_size should be a multiple of 16 "
+            f"(head_size0={head_size0}), res={res}, "
+            f"config=\n{config}"
+        )
+    return res
 
 
 def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:

onnx_diagnostic/tasks/image_classification.py

@@ -58,8 +58,8 @@ def get_inputs(
     shapes = {
         "pixel_values": {
             0: torch.export.Dim("batch", min=1, max=1024),
-            2: torch.export.Dim("width", min=1, max=4096),
-            3: torch.export.Dim("height", min=1, max=4096),
+            2: "width",
+            3: "height",
         },
     }
     inputs = dict(

onnx_diagnostic/tasks/text_generation.py

@@ -27,7 +27,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
         kwargs = dict(
             num_hidden_layers=min(config.num_hidden_layers, 2),
             intermediate_size=256 if config is None else min(512, config.intermediate_size),
-            hidden_size=256 if config is None else min(256, config.hidden_size),
+            hidden_size=512 if config is None else min(512, config.hidden_size),
             cls_cache="MambaCache",
             state_size=8 if config is None else getattr(config, "state_size", None),
             conv_kernel=4 if config is None else getattr(config, "conv_kernel", None),
@@ -44,8 +44,8 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
                 else config.num_attention_heads
             ),
             hidden_size=(
-                min(config.hidden_size, 3072 // 4)
-                if config.hidden_size % 4 == 0
+                min(config.hidden_size, 4096 // 4)
+                if config.hidden_size % 64 == 0
                 else config.hidden_size
             ),
         )