update to 4 layers (#204)

titaiwangms · web-flow · commit 945e778ca9d3 · 2025-08-27T10:23:44.000+02:00
diff --git a/onnx_diagnostic/tasks/automatic_speech_recognition.py b/onnx_diagnostic/tasks/automatic_speech_recognition.py
@@ -15,7 +15,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if hasattr(config, "decoder_layers"):
         config.decoder_layers = min(config.decoder_layers, 2)
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     update_config(config, kwargs)
     return kwargs
 
diff --git a/onnx_diagnostic/tasks/feature_extraction.py b/onnx_diagnostic/tasks/feature_extraction.py
@@ -9,7 +9,7 @@
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_hidden_layers")
-    kwargs = dict(num_hidden_layers=min(config.num_hidden_layers, 2))
+    kwargs = dict(num_hidden_layers=min(config.num_hidden_layers, 4))
     update_config(config, kwargs)
     return kwargs
 
diff --git a/onnx_diagnostic/tasks/fill_mask.py b/onnx_diagnostic/tasks/fill_mask.py
@@ -9,7 +9,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
-        num_hidden_layers=min(config.num_hidden_layers, 2),
+        num_hidden_layers=min(config.num_hidden_layers, 4),
         num_attention_heads=min(config.num_attention_heads, 4),
     )
     update_config(config, kwargs)
diff --git a/onnx_diagnostic/tasks/image_classification.py b/onnx_diagnostic/tasks/image_classification.py
@@ -17,7 +17,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
-            min(config.num_hidden_layers, 2)
+            min(config.num_hidden_layers, 4)
             if hasattr(config, "num_hidden_layers")
             else len(config.hidden_sizes)
         )
diff --git a/onnx_diagnostic/tasks/image_text_to_text.py b/onnx_diagnostic/tasks/image_text_to_text.py
@@ -10,7 +10,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     if hasattr(config, "mm_tokens_per_image"):
         config.mm_tokens_per_image = min(config.mm_tokens_per_image, 2)
     if hasattr(config, "vision_config"):
diff --git a/onnx_diagnostic/tasks/mask_generation.py b/onnx_diagnostic/tasks/mask_generation.py
@@ -9,7 +9,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     if hasattr(config, "vision_config") and hasattr(config.vision_config, "num_hidden_layers"):
         config.vision_config.num_hidden_layers = min(config.vision_config.num_hidden_layers, 2)
     update_config(config, kwargs)
diff --git a/onnx_diagnostic/tasks/mixture_of_expert.py b/onnx_diagnostic/tasks/mixture_of_expert.py
@@ -11,7 +11,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     if hasattr(config, "vision_config") and hasattr(config.vision_config, "num_hidden_layers"):
         config.vision_config.num_hidden_layers = min(config.vision_config.num_hidden_layers, 2)
     if hasattr(config, "audio_processor") and hasattr(
diff --git a/onnx_diagnostic/tasks/object_detection.py b/onnx_diagnostic/tasks/object_detection.py
@@ -10,7 +10,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
-            min(config.num_hidden_layers, 2)
+            min(config.num_hidden_layers, 4)
             if hasattr(config, "num_hidden_layers")
             else len(config.hidden_sizes)
         )
diff --git a/onnx_diagnostic/tasks/sentence_similarity.py b/onnx_diagnostic/tasks/sentence_similarity.py
@@ -9,7 +9,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
-        num_hidden_layers=min(config.num_hidden_layers, 2),
+        num_hidden_layers=min(config.num_hidden_layers, 4),
         num_attention_heads=min(config.num_attention_heads, 4),
     )
     update_config(config, kwargs)
diff --git a/onnx_diagnostic/tasks/summarization.py b/onnx_diagnostic/tasks/summarization.py
@@ -12,7 +12,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if hasattr(config, "num_decoder_layers"):
         config.num_decoder_layers = min(config.num_decoder_layers, 2)
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     update_config(config, kwargs)
     return kwargs
 
diff --git a/onnx_diagnostic/tasks/text2text_generation.py b/onnx_diagnostic/tasks/text2text_generation.py
@@ -12,7 +12,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if hasattr(config, "num_decoder_layers"):
         config.num_decoder_layers = min(config.num_decoder_layers, 2)
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     update_config(config, kwargs)
     return kwargs
 
diff --git a/onnx_diagnostic/tasks/text_classification.py b/onnx_diagnostic/tasks/text_classification.py
@@ -9,7 +9,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
-        num_hidden_layers=min(config.num_hidden_layers, 2),
+        num_hidden_layers=min(config.num_hidden_layers, 4),
         num_attention_heads=min(config.num_attention_heads, 4),
     )
     update_config(config, kwargs)
diff --git a/onnx_diagnostic/tasks/text_generation.py b/onnx_diagnostic/tasks/text_generation.py
@@ -25,7 +25,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if config.__class__.__name__ == "FalconMambaConfig":
         check_hasattr(config, "conv_kernel", "state_size", "intermediate_size")  # 4 and 8
         kwargs = dict(
-            num_hidden_layers=min(config.num_hidden_layers, 2),
+            num_hidden_layers=min(config.num_hidden_layers, 4),
             intermediate_size=256 if config is None else min(512, config.intermediate_size),
             hidden_size=512 if config is None else min(512, config.hidden_size),
             cls_cache="MambaCache",
@@ -37,7 +37,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
             head_dim=getattr(
                 config, "head_dim", config.hidden_size // config.num_attention_heads
             ),
-            num_hidden_layers=min(config.num_hidden_layers, 2),
+            num_hidden_layers=min(config.num_hidden_layers, 4),
             num_key_value_heads=(
                 config.num_key_value_heads
                 if hasattr(config, "num_key_value_heads")

Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:`
`9`	`9`	`"""Reduces a model size."""`
`10`	`10`	`check_hasattr(config, "num_attention_heads", "num_hidden_layers")`
`11`	`11`	`kwargs = dict(`
`12`		`- num_hidden_layers=min(config.num_hidden_layers, 2),`
	`12`	`+ num_hidden_layers=min(config.num_hidden_layers, 4),`
`13`	`13`	`num_attention_heads=min(config.num_attention_heads, 4),`
`14`	`14`	`)`
`15`	`15`	`update_config(config, kwargs)`
Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:`
`17`	`17`	`check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))`
`18`	`18`	`kwargs = dict(`
`19`	`19`	`num_hidden_layers=(`
`20`		`- min(config.num_hidden_layers, 2)`
	`20`	`+ min(config.num_hidden_layers, 4)`
`21`	`21`	`if hasattr(config, "num_hidden_layers")`
`22`	`22`	`else len(config.hidden_sizes)`
`23`	`23`	`)`
Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:`
`10`	`10`	`check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))`
`11`	`11`	`kwargs = dict(`
`12`	`12`	`num_hidden_layers=(`
`13`		`- min(config.num_hidden_layers, 2)`
	`13`	`+ min(config.num_hidden_layers, 4)`
`14`	`14`	`if hasattr(config, "num_hidden_layers")`
`15`	`15`	`else len(config.hidden_sizes)`
`16`	`16`	`)`