change

Author: xadupre

_unittests/ut_tasks/try_tasks.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 import torch
 from onnx_diagnostic.ext_test_case import ExtTestCase, never_test
@@ -799,17 +800,19 @@ def test_imagetext2text_generation_gemma3_4b_it(self):
         from transformers import AutoProcessor, Gemma3ForConditionalGeneration
 
         model_id = "google/gemma-3-4b-it"
-        # model_id = "google/gemma-3n-e4b-it"
-        # model_id = "qnaug/gemma-3-4b-med"
-        # model_id = "hf-internal-testing/tiny-random-Gemma3ForCausalLM"
-        # data = get_untrained_model_with_inputs(
-        #     model_id, verbose=1, add_second_input=True,
-        #     same_as_pretrained=True, use_pretrained=True
-        # )
-        # model = data["model"]
-        model = Gemma3ForConditionalGeneration.from_pretrained(
-            model_id, device_map="cpu"
-        ).eval()
+        if os.environ.get("PRETRAINED", ""):
+            model = Gemma3ForConditionalGeneration.from_pretrained(
+                model_id, device_map="cpu"
+            ).eval()
+        else:
+            data = get_untrained_model_with_inputs(
+                model_id,
+                verbose=1,
+                add_second_input=True,
+                # same_as_pretrained=True, #use_pretrained=True
+            )
+            model = data["model"]
+
         print(f"-- model.device={model.device}")
         processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
         print(f"-- processor={type(processor)}")
@@ -845,11 +848,39 @@ def test_imagetext2text_generation_gemma3_4b_it(self):
         # inputs.pop("token_type_ids", None)
         print(f"-- inputs={self.string_type(inputs)}")
 
+        # iteration 1
+        #   cache_position:T7s281,
+        #   past_key_values:StaticCache(key_cache=#0[], value_cache=#0[]),
+        #   input_ids:T7s1x281,
+        #   inputs_embeds:None,
+        #   token_type_ids:T7s1x281,
+        #   attention_mask:dict(sliding_attention:T9s1x1x281x580,
+        #                       full_attention:T9s1x1x281x580),
+        #   position_ids:None,
+        #   use_cache:bool,
+        #   logits_to_keep:None,
+        #   pixel_values:T16s1x3x896x896,
+        #   return_dict:bool)
+        # iteration 3
+        #   cache_position:T7s1,
+        #   past_key_values:StaticCache(key_cache=#34[T1s1x4x580x256,...],
+        #                               value_cache=#34[T1s1x4x580x256,...]),
+        #   input_ids:T7s1x1,
+        #   inputs_embeds:None,
+        #   token_type_ids:T7s1x1,
+        #   attention_mask:dict(sliding_attention:T9s1x1x1x580,full_attention:T9s1x1x1x580),
+        #   position_ids:None,
+        #   use_cache:bool,logits_to_keep:None,return_dict:bool)
+
         print()
         # steal forward creates a bug...
-        with steal_forward(model):  # , torch.inference_mode():
+        with steal_forward(
+            model,
+            dump_file=self.get_dump_file("test_imagetext2text_generation_gemma3_4b_it.onnx"),
+            dump_drop={"attention_mask", "past_key_values", "pixel_values"},
+        ):
             generated_ids = model.generate(
-                **inputs, max_new_tokens=300, do_sample=False, cache_implementation="hybrid"
+                **inputs, max_new_tokens=282, do_sample=False, cache_implementation="static"
             )
         output_text = processor.decode(
             generated_ids[0][inputs["input_ids"].shape[1] :], skip_special_tokens=False

onnx_diagnostic/helpers/mini_onnx_builder.py

@@ -381,6 +381,23 @@ def _mk(k):
                 else:
                     for p, o in _flatten_iterator(getattr(obj, att), sep):
                         yield f"DynamicCache_{att}{sep}{p}", o
+        elif obj.__class__.__name__ == "StaticCache":
+            # transformers
+            import transformers
+            from .cache_helper import CacheKeyValue
+
+            assert isinstance(
+                obj, transformers.cache_utils.StaticCache
+            ), f"Unexpected type {type(obj)}"
+            obj = CacheKeyValue(obj)
+            atts = ["key_cache", "value_cache"]
+            for i, att in enumerate(atts):
+                if i == len(atts) - 1:
+                    for p, o in _flatten_iterator(getattr(obj, att), sep):
+                        yield f"StaticCache._{att}{sep}{p}", o
+                else:
+                    for p, o in _flatten_iterator(getattr(obj, att), sep):
+                        yield f"StaticCache_{att}{sep}{p}", o
         else:
             raise NotImplementedError(f"Unexpected type {type(obj)}")
 

onnx_diagnostic/helpers/torch_helper.py

@@ -5,7 +5,7 @@
 import sys
 import warnings
 from collections.abc import Iterable
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import numpy as np
 import onnx
 from onnx.external_data_helper import load_external_data_for_tensor, uses_external_data
@@ -283,6 +283,7 @@ def steal_forward(
     ],
     fprint: Callable = string_type,
     dump_file: Optional[str] = None,
+    dump_drop: Optional[Set[str]] = None,
     submodules: bool = False,
     verbose: int = 0,
     storage_limit: int = 2**27,
@@ -303,6 +304,7 @@ def steal_forward(
     :param dump_file: dumps stolen inputs and outputs in an onnx model,
         they can be restored with :func:`create_input_tensors_from_onnx_model
         <onnx_diagnostic.helpers.mini_onnx_builder.create_input_tensors_from_onnx_model>`
+    :param dump_drop: to drop some inputs too big (only if dump_file is specified)
     :param submodules: if True and model is a module, the list extended with all the submodules
         the module contains
     :param verbose: verbosity
@@ -411,6 +413,9 @@ def forward(self, x, y):
             if verbose:
                 size = torch_tensor_size(storage)
                 print(f"-- gather stored {len(storage)} objects, size={size // 2 ** 20} Mb")
+            if dump_drop:
+                print(string_type(dump_drop))
+                stop
             proto = create_onnx_model_from_input_tensors(storage)
             if verbose:
                 print("-- dumps stored objects")
@@ -794,9 +799,14 @@ def torch_deepcopy(value: Any) -> Any:
         from .cache_helper import CacheKeyValue
 
         ca = CacheKeyValue(value)
+        if len(ca.key_cache) == 0:
+            # Use of deepcopy.
+            import copy
+
+            return copy.deepcopy(value)
         return make_static_cache(
             torch_deepcopy(list(zip(ca.key_cache, ca.value_cache))),
-            max_cache_len=value.max_cache_len,
+            max_cache_len=max([value.max_cache_len, *[t.shape[2] for t in ca.key_cache]]),
         )
     if value.__class__.__name__ == "HybridCache":
         from .cache_helper import CacheKeyValue

onnx_diagnostic/tasks/image_text_to_text.py

@@ -14,6 +14,27 @@
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
+    if (
+        hasattr(config, "architectures")
+        and config.architectures
+        and config.architectures[0] == "Gemma3ForConditionalGeneration"
+    ):
+        if hasattr(config, "vision_config"):
+            if hasattr(config.vision_config, "num_hidden_layers"):
+                config.vision_config.num_hidden_layers = min(
+                    config.vision_config.num_hidden_layers, nhl()
+                )
+        if hasattr(config, "text_config"):
+            if hasattr(config.text_config, "intermediate_size"):
+                config.text_config.intermediate_size = min(
+                    config.text_config.intermediate_size, 10240 // 10 * 5 // 2
+                )
+                config.text_config.hidden_size = min(
+                    config.text_config.hidden_size, 2560 // 10 * 5 // 2
+                )
+        update_config(config, kwargs)
+        return kwargs
+
     if hasattr(config, "num_hidden_layers"):
         config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     if hasattr(config, "mm_tokens_per_image"):