From c4228f6f2d48892f5941a743e8b017feb23b264c Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Fri, 8 Aug 2025 21:58:39 +0000 Subject: [PATCH 1/2] fix --- onnx_diagnostic/tasks/text_generation.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/onnx_diagnostic/tasks/text_generation.py b/onnx_diagnostic/tasks/text_generation.py index 1a617bdd..b7717dfe 100644 --- a/onnx_diagnostic/tasks/text_generation.py +++ b/onnx_diagnostic/tasks/text_generation.py @@ -34,27 +34,13 @@ def reduce_model_config(config: Any) -> Dict[str, Any]: ) else: kwargs = dict( - head_dim=getattr( - config, "head_dim", config.hidden_size // config.num_attention_heads - ), num_hidden_layers=min(config.num_hidden_layers, 2), num_key_value_heads=( config.num_key_value_heads if hasattr(config, "num_key_value_heads") else config.num_attention_heads ), - hidden_size=( - min(config.hidden_size, 4096 // 4) - if config.hidden_size % 64 == 0 - else config.hidden_size - ), ) - if config is None or hasattr(config, "intermediate_size"): - kwargs["intermediate_size"] = ( - min(config.intermediate_size, 24576 // 4) - if config.intermediate_size % 4 == 0 - else config.intermediate_size - ) update_config(config, kwargs) return kwargs From 7d8c2615ee5a459be83983c15289471f5611a2da Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Fri, 8 Aug 2025 22:05:42 +0000 Subject: [PATCH 2/2] keep head_dim --- onnx_diagnostic/tasks/text_generation.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/onnx_diagnostic/tasks/text_generation.py b/onnx_diagnostic/tasks/text_generation.py index b7717dfe..601f729b 100644 --- a/onnx_diagnostic/tasks/text_generation.py +++ b/onnx_diagnostic/tasks/text_generation.py @@ -34,6 +34,9 @@ def reduce_model_config(config: Any) -> Dict[str, Any]: ) else: kwargs = dict( + head_dim=getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ), num_hidden_layers=min(config.num_hidden_layers, 2), num_key_value_heads=( config.num_key_value_heads