Small fixes for ModelBuilder

Author: xadupre

CHANGELOGS.rst

@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.6.1
++++++
+
+* :pr:`112`: 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>`_

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

@@ -41,7 +41,17 @@ def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
     """Reduces a model size."""
     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_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_size % 16 == 0
+        ), f"head_size should be a multiple of 16, res={res}, 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
             ),
         )