Small fixes for ModelBuilder

CHANGELOGS.rst

@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.6.1
++++++
+
+* :pr:`113`: fixes a couple of issues with ModelBuilder
+
 0.6.0
 +++++
 

_doc/index.rst

@@ -206,6 +206,7 @@ The function replaces dynamic dimensions defined as strings by
 Older versions
 ++++++++++++++
 
+* `0.6.1 <../v0.6.1/index.html>`_
 * `0.6.0 <../v0.6.0/index.html>`_
 * `0.5.0 <../v0.5.0/index.html>`_
 * `0.4.4 <../v0.4.4/index.html>`_

_unittests/ut_tasks/test_tasks.py

@@ -206,7 +206,7 @@ def test_falcon_mamba_dev(self):
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
         model(**inputs)
         model(**data["inputs2"])
-        self.assertIn((data["size"], data["n_weights"]), [(138640384, 34660096)])
+        self.assertIn((data["size"], data["n_weights"]), [(274958336, 68739584)])
         if not has_transformers("4.55"):
             raise unittest.SkipTest("The model has control flow.")
         with torch_export_patches(patch_transformers=True, verbose=10, stop_if_static=1):

_unittests/ut_torch_models/test_hghub_model.py

@@ -83,7 +83,7 @@ def test_get_untrained_model_with_inputs_codellama(self):
         model, inputs = data["model"], data["inputs"]
         model(**inputs)
         # different expected value for different version of transformers
-        self.assertIn((data["size"], data["n_weights"]), [(410532864, 102633216)])
+        self.assertIn((data["size"], data["n_weights"]), [(547377152, 136844288)])
 
     @hide_stdout()
     @ignore_errors(OSError)

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.6.0"
+__version__ = "0.6.1"
 __author__ = "Xavier Dupré"

onnx_diagnostic/helpers/model_builder_helper.py

@@ -237,18 +237,24 @@ def create_model_builder(
         "OlmoForCausalLM": builder.OLMoModel,
         "PhiForCausalLM": builder.PhiModel,
         "Phi3ForCausalLM": (
-            lambda config, *_: (
-                builder.Phi3MiniModel
-                if config.max_position_embeddings == config.original_max_position_embeddings
-                else builder.Phi3MiniLongRoPEModel
+            lambda config, *args: (
+                (
+                    builder.Phi3MiniModel
+                    if config.max_position_embeddings
+                    == config.original_max_position_embeddings
+                    else builder.Phi3MiniLongRoPEModel
+                )(config, *args)
             )
         ),
         "PhiMoEForCausalLM": builder.Phi3MoELongRoPEModel,
         "Phi3SmallForCausalLM": (
-            lambda config, *_: (
-                builder.Phi3SmallModel
-                if config.max_position_embeddings == config.original_max_position_embeddings
-                else builder.Phi3SmallLongRoPEModel
+            lambda config, *args: (
+                (
+                    builder.Phi3SmallModel
+                    if config.max_position_embeddings
+                    == config.original_max_position_embeddings
+                    else builder.Phi3SmallLongRoPEModel
+                )(config, *args)
             )
         ),
         "Phi3VForCausalLM": builder.Phi3VModel,
@@ -317,7 +323,17 @@ def _post(onnx_model):
         )
 
     cls = arch_map[config.architectures[0]]
+
+    # ModelBuilder does not like None values for some parameters.
+    remove = set()
+    for c in ["head_dim"]:
+        if hasattr(config, c) and getattr(config, c) is None:
+            remove.add(c)
+    for c in remove:
+        delattr(config, c)
+
     onnx_model = cls(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
+
     if post:
         post(onnx_model)
     _make_model(onnx_model, model, verbose=verbose)

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/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
             ),
         )