onnx-diagnostic 0.7.12__py3-none-any.whl → 0.7.14__py3-none-any.whl

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1,13 +1,20 @@
 import inspect
 import math
+import os
 from dataclasses import dataclass
 from functools import wraps
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, List, Optional, Tuple, Union
 import packaging.version as pv
 import torch
 import transformers
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
 from transformers.cache_utils import StaticCache, Cache
+from transformers.generation.utils import (
+    GenerateNonBeamOutput,
+    GenerationConfig,
+    StoppingCriteriaList,
+    LogitsProcessorList,
+)
 
 try:
     from transformers.cache_utils import parse_processor_args  # noqa: F401
@@ -114,6 +121,7 @@ if patch_masking_utils:
         """manual patch for function ``transformers.masking_utils.eager_mask``."""
         # The masks for eager attention are simply boolean mask from sdpa, casted to 0 and -inf
         _ = kwargs.pop("allow_is_causal_skip", None)
+        # PATCHED: this line called the patched version of sdpa_mask
         mask = patched_sdpa_mask_recent_torch(
             batch_size=batch_size,
             cache_position=cache_position,
@@ -126,7 +134,7 @@ if patch_masking_utils:
             **kwargs,
         )
         min_dtype = torch.finfo(dtype).min
-        # The patched line.
+        # PATCHED: the following line
         # we need 0s where the tokens should be taken into account,
         # and -inf otherwise (mask is already of boolean type)
         # mask =
@@ -158,6 +166,7 @@ if patch_masking_utils:
             mask_function = and_masks(mask_function, padding_mask_function(padding_mask))
         batch_arange = torch.arange(batch_size, device=cache_position.device)
         head_arange = torch.arange(1, device=cache_position.device)
+        # PATCHED: this line calls the patched version of vmap_for_bhqkv
         causal_mask = patched__vmap_for_bhqkv(mask_function)(
             batch_arange, head_arange, cache_position, kv_arange
         )
@@ -214,6 +223,7 @@ if patch_DynamicLayer:
             self.dtype, self.device = key_states.dtype, key_states.device
             new_shape = list(key_states.shape)
             new_shape[-2] = 0
+            # PATCHED: used a tensor with an empty shape and not en empty list to initialize
             self.keys = torch.empty(new_shape, dtype=self.dtype, device=self.device)
             self.values = torch.empty(new_shape, dtype=self.dtype, device=self.device)
             if patch_is_initialized:
@@ -248,6 +258,8 @@ def _patch_make_causal_mask(
         diagonal = past_key_values_length - sliding_window - 1
 
         context_mask = torch.tril(torch.ones_like(mask, dtype=torch.bool), diagonal=diagonal)
+        # PATCHED: removed if is_torchdynamo_compiling(): mask = mask.clone()
+        # and used masked_fill instead of masked_fill_
         # In this case, the current implementation of torch fails (17/12/2024).
         # Try model Phi-3.5-Mini-Instruct.
         mask = mask.masked_fill(context_mask, torch.finfo(dtype).min)
@@ -455,7 +467,16 @@ class patched_GenerationMixin:
     _PATCHES_ = [
         "_cache_dependant_input_preparation",
         "_cache_dependant_input_preparation_exporting",
-        "prepare_inputs_for_generation",
+        (
+            None
+            if pv.Version(transformers.__version__) >= pv.Version("4.56")
+            else "prepare_inputs_for_generation"
+        ),
+        (
+            "_sample"
+            if pv.Version(transformers.__version__) == pv.Version("4.57.0.dev0")
+            else None
+        ),
     ]
     _PATCHED_CLASS_ = transformers.generation.utils.GenerationMixin
 
@@ -588,7 +609,7 @@ class patched_GenerationMixin:
         model_inputs = {}
         # - some models don't have `Cache` support
         # (which implies they don't expect `cache_position` in `forward`)
-        if self._supports_cache_class:
+        if getattr(self, "_supports_cache_class", False):
             model_inputs["cache_position"] = cache_position
         # - `cache_position` was not a mandatory input in
         # `prepare_inputs_for_generation` for those models, and this
@@ -728,6 +749,192 @@ class patched_GenerationMixin:
         model_inputs.pop("labels", None)
         return model_inputs
 
+    def _sample(
+        self,
+        input_ids: torch.LongTensor,
+        logits_processor: "LogitsProcessorList",  # noqa: F821
+        stopping_criteria: "StoppingCriteriaList",  # noqa: F821
+        generation_config: "GenerationConfig",  # noqa: F821
+        synced_gpus: bool = False,
+        streamer: Optional["BaseStreamer"] = None,  # noqa: F821
+        **model_kwargs,
+    ) -> Union["GenerateNonBeamOutput", torch.LongTensor]:  # noqa: F821
+        """
+        2025/09/29: updates for Gemma3 models, fix for eager mode as well as the export.
+        """
+        # init values
+        pad_token_id = generation_config._pad_token_tensor
+        output_attentions = generation_config.output_attentions
+        output_hidden_states = generation_config.output_hidden_states
+        output_scores = generation_config.output_scores
+        output_logits = generation_config.output_logits
+        return_dict_in_generate = generation_config.return_dict_in_generate
+        has_eos_stopping_criteria = any(
+            hasattr(criteria, "eos_token_id") for criteria in stopping_criteria
+        )
+        do_sample = generation_config.do_sample
+
+        # init attention / hidden states / scores tuples
+        scores = () if (return_dict_in_generate and output_scores) else None
+        raw_logits = () if (return_dict_in_generate and output_logits) else None
+        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+        cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+        decoder_hidden_states = (
+            () if (return_dict_in_generate and output_hidden_states) else None
+        )
+
+        # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+        if return_dict_in_generate and self.config.is_encoder_decoder:
+            encoder_attentions = (
+                model_kwargs["encoder_outputs"].get("attentions")
+                if output_attentions
+                else None
+            )
+            encoder_hidden_states = (
+                model_kwargs["encoder_outputs"].get("hidden_states")
+                if output_hidden_states
+                else None
+            )
+
+        # keep track of which sequences are already finished
+        batch_size, cur_len = input_ids.shape[:2]
+        this_peer_finished = False
+        unfinished_sequences = torch.ones(
+            batch_size, dtype=torch.long, device=input_ids.device
+        )
+        model_kwargs = self._get_initial_cache_position(
+            cur_len, input_ids.device, model_kwargs
+        )
+
+        model_forward = self.__call__
+        compile_forward = self._valid_auto_compile_criteria(model_kwargs, generation_config)
+        if compile_forward:
+            os.environ["TOKENIZERS_PARALLELISM"] = "0"
+            # If we use FA2 and a static cache, we cannot compile with fullgraph
+            if self.config._attn_implementation == "flash_attention_2":
+                # only raise warning if the user passed an explicit compile-config
+                if (
+                    generation_config.compile_config is not None
+                    and generation_config.compile_config.fullgraph
+                ):
+                    generation_config.compile_config.fullgraph = False
+            model_forward = self.get_compiled_call(generation_config.compile_config)
+
+        if generation_config.prefill_chunk_size is not None:
+            model_kwargs = self._prefill_chunking(input_ids, generation_config, **model_kwargs)
+            is_prefill = False
+        else:
+            is_prefill = True
+
+        while self._has_unfinished_sequences(
+            this_peer_finished, synced_gpus, device=input_ids.device
+        ):
+            # prepare model inputs
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+
+            if is_prefill:
+                outputs = self(**model_inputs, return_dict=True)
+                is_prefill = False
+            else:
+                outputs = model_forward(**model_inputs, return_dict=True)
+
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs,
+                model_kwargs,
+                is_encoder_decoder=self.config.is_encoder_decoder,
+            )
+            if synced_gpus and this_peer_finished:
+                continue
+
+            next_token_logits = outputs.logits[:, -1, :].to(
+                copy=True, dtype=torch.float32, device=input_ids.device
+            )
+
+            # pre-process distribution
+            next_token_scores = logits_processor(input_ids, next_token_logits)
+
+            # Store scores, attentions and hidden_states when required
+            if return_dict_in_generate:
+                if output_scores:
+                    scores += (next_token_scores,)
+                if output_logits:
+                    raw_logits += (next_token_logits,)
+                if output_attentions:
+                    decoder_attentions += (
+                        (outputs.decoder_attentions,)
+                        if self.config.is_encoder_decoder
+                        else (outputs.attentions,)
+                    )
+                    if self.config.is_encoder_decoder:
+                        cross_attentions += (outputs.cross_attentions,)
+
+                if output_hidden_states:
+                    decoder_hidden_states += (
+                        (outputs.decoder_hidden_states,)
+                        if self.config.is_encoder_decoder
+                        else (outputs.hidden_states,)
+                    )
+
+            # token selection
+            if do_sample:
+                probs = torch.nn.functional.softmax(next_token_scores, dim=-1)
+                next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+            else:
+                next_tokens = torch.argmax(next_token_scores, dim=-1)
+
+            # finished sentences should have their next token be a padding token
+            if has_eos_stopping_criteria:
+                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (
+                    1 - unfinished_sequences
+                )
+
+            # update generated ids, model inputs, and length for next step
+            # PATCHED: the two following lines, next_tokens can 2D already for this model
+            next_tokens_2d = (
+                next_tokens if len(next_tokens.shape) == 2 else next_tokens[:, None]
+            )
+            input_ids = torch.cat([input_ids, next_tokens_2d], dim=-1)
+            if streamer is not None:
+                streamer.put(next_tokens.cpu())
+
+            unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores)
+            this_peer_finished = unfinished_sequences.max() == 0
+            cur_len += 1
+
+            # This is needed to properly delete outputs.logits which may be very large
+            # for first iteration
+            # Otherwise a reference to outputs is kept which keeps
+            # the logits alive in the next iteration
+            del outputs
+
+        if streamer is not None:
+            streamer.end()
+
+        if return_dict_in_generate:
+            if self.config.is_encoder_decoder:
+                return transformers.generation.utils.GenerateEncoderDecoderOutput(
+                    sequences=input_ids,
+                    scores=scores,
+                    logits=raw_logits,
+                    encoder_attentions=encoder_attentions,
+                    encoder_hidden_states=encoder_hidden_states,
+                    decoder_attentions=decoder_attentions,
+                    cross_attentions=cross_attentions,
+                    decoder_hidden_states=decoder_hidden_states,
+                    past_key_values=model_kwargs.get("past_key_values"),
+                )
+            else:
+                return transformers.generation.utils.GenerateDecoderOnlyOutput(
+                    sequences=input_ids,
+                    scores=scores,
+                    logits=raw_logits,
+                    attentions=decoder_attentions,
+                    hidden_states=decoder_hidden_states,
+                    past_key_values=model_kwargs.get("past_key_values"),
+                )
+        else:
+            return input_ids
+
 
 def patched__compute_dynamic_ntk_parameters(
     config: Optional[transformers.PretrainedConfig] = None,
@@ -791,6 +998,7 @@ def patched__compute_dynamic_ntk_parameters(
     if seq_len is None:
         seq_len = max_position_embeddings
     else:
+        # PATCHED: remove the line using max
         torch._check(isinstance(seq_len, torch.Tensor))
         seq_len = torch.maximum(
             seq_len,
@@ -896,6 +1104,7 @@ def patched_dynamic_rope_update(rope_forward):
         )
         original_inv_freq = self.original_inv_freq.to(device)
 
+        # PATCHED: uses torch.cond instead of a test
         cond = (seq_len > original_max_position_embeddings).item()
         inv_freq = torch.cond(
             cond,
@@ -967,6 +1176,7 @@ def patched_dynamic_rope_update(rope_forward):
 
         original_inv_freq = self.original_inv_freq.to(device)
         cond = (seq_len >= self.original_max_seq_len).item()
+        # PATCHED: uses torch.cond instead of a test
         inv_freq = torch.cond(
             cond,
             (lambda x, y: x.clone()),
@@ -1002,6 +1212,7 @@ def common_eager_attention_forward(
 
     attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
     if attention_mask is not None:
+        # PATCHED
         # The two following lines were added.
         if attention_mask is not None and attention_mask.ndim == 4:
             attention_mask = attention_mask[:, :, :, : key.shape[-2]]
@@ -1074,6 +1285,7 @@ def patched_modeling_marian_eager_attention_forward(
 class common_RotaryEmbedding(torch.nn.Module):
     # This may cause some issues.
     # @torch.no_grad()
+    # PATCHED: the decorator
     @patched_dynamic_rope_update
     def forward(self, x, position_ids):
         inv_freq_expanded = (
@@ -1629,3 +1841,56 @@ if patch_qwen3:
                 batch_size, sequence_length, hidden_dim
             )
             return final_hidden_states, router_logits
+
+
+try:
+    from transformers.models.gemma3.modeling_gemma3 import Gemma3Model  # noqa: F401
+
+    patch_gemma3 = True
+except ImportError:
+    patch_gemma3 = False
+
+
+if patch_gemma3:
+
+    class patched_Gemma3Model(torch.nn.Module):
+        _PATCHES_ = ["get_placeholder_mask"]
+        _PATCHED_CLASS_ = transformers.models.gemma3.modeling_gemma3.Gemma3Model
+        _PATCHED_PR_ = "https://github.com/huggingface/transformers/pull/41319"
+
+        def get_placeholder_mask(
+            self,
+            input_ids: torch.LongTensor,
+            inputs_embeds: torch.FloatTensor,
+            image_features: torch.FloatTensor,
+        ):
+            if input_ids is None:
+                special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                    torch.tensor(
+                        self.config.image_token_id,
+                        dtype=torch.long,
+                        device=inputs_embeds.device,
+                    )
+                )
+                special_image_mask = special_image_mask.all(-1)
+            else:
+                special_image_mask = input_ids == self.config.image_token_id
+
+            n_image_tokens = special_image_mask.sum()
+            special_image_mask = (
+                special_image_mask.unsqueeze(-1)
+                .expand_as(inputs_embeds)
+                .to(inputs_embeds.device)
+            )
+            n_image_features = image_features.shape[0] * image_features.shape[1]
+            # PATCHED: torch._check
+            # if inputs_embeds[special_image_mask].numel() != image_features.numel():
+            #    raise ValueError( ... )
+            torch._check(
+                inputs_embeds[special_image_mask].numel() == image_features.numel(),
+                lambda: (
+                    f"Image features and image tokens do not match: tokens: "
+                    f"{n_image_tokens}, features {n_image_features}"
+                ),
+            )
+            return special_image_mask

onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py

@@ -4829,3 +4829,39 @@ def _ccached_microsoft_phi3_mini_128k_instruct():
             "vocab_size": 32064,
         }
     )
+
+
+def _ccached_google_gemma_3_4b_it_like():
+    "google/gemma-3-4b-it"
+    return transformers.Gemma3Config(
+        **{
+            "architectures": ["Gemma3ForConditionalGeneration"],
+            "boi_token_index": 255999,
+            "eoi_token_index": 256000,
+            "eos_token_id": [1, 106],
+            "image_token_index": 262144,
+            "initializer_range": 0.02,
+            "mm_tokens_per_image": 256,
+            "model_type": "gemma3",
+            "text_config": {
+                "hidden_size": 2560,
+                "intermediate_size": 10240,
+                "model_type": "gemma3_text",
+                "num_hidden_layers": 34,
+                "rope_scaling": {"factor": 8.0, "rope_type": "linear"},
+                "sliding_window": 1024,
+            },
+            "torch_dtype": "bfloat16",
+            "transformers_version": "4.50.0.dev0",
+            "vision_config": {
+                "hidden_size": 1152,
+                "image_size": 896,
+                "intermediate_size": 4304,
+                "model_type": "siglip_vision_model",
+                "num_attention_heads": 16,
+                "num_hidden_layers": 27,
+                "patch_size": 14,
+                "vision_use_head": false,
+            },
+        }
+    )

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -25,6 +25,20 @@ def _code_needing_rewriting(model: Any) -> Any:
     return code_needing_rewriting(model)
 
 
+def _preprocess_model_id(
+    model_id: str, subfolder: Optional[str], same_as_pretrained: bool, use_pretrained: bool
+) -> Tuple[str, Optional[str], bool, bool]:
+    if subfolder or "//" not in model_id:
+        return model_id, subfolder, same_as_pretrained, use_pretrained
+    spl = model_id.split("//")
+    if spl[-1] == "pretrained":
+        return _preprocess_model_id("//".join(spl[:-1]), "", True, True)
+    if spl[-1] in {"transformer", "vae"}:
+        # known subfolder
+        return "//".join(spl[:-1]), spl[-1], same_as_pretrained, use_pretrained
+    return model_id, subfolder, same_as_pretrained, use_pretrained
+
+
 def get_untrained_model_with_inputs(
     model_id: str,
     config: Optional[Any] = None,
@@ -57,7 +71,7 @@ def get_untrained_model_with_inputs(
         to get a smaller model
     :param use_pretrained: download the pretrained weights as well
     :param use_preinstalled: use preinstalled configurations
-    :param add_second_input: provides a second inputs to check a model
+    :param add_second_input: provides others inputs to check a model
         supports different shapes
     :param subfolder: subfolder to use for this model id
     :param use_only_preinstalled: use only preinstalled version
@@ -85,8 +99,16 @@ def get_untrained_model_with_inputs(
         f"model_id={model_id!r}, preinstalled model is only available "
         f"if use_only_preinstalled is False."
     )
+    model_id, subfolder, same_as_pretrained, use_pretrained = _preprocess_model_id(
+        model_id,
+        subfolder,
+        same_as_pretrained=same_as_pretrained,
+        use_pretrained=use_pretrained,
+    )
     if verbose:
-        print(f"[get_untrained_model_with_inputs] model_id={model_id!r}")
+        print(
+            f"[get_untrained_model_with_inputs] model_id={model_id!r}, subfolder={subfolder!r}"
+        )
         if use_preinstalled:
             print(f"[get_untrained_model_with_inputs] use preinstalled {model_id!r}")
     if config is None:
@@ -178,7 +200,7 @@ def get_untrained_model_with_inputs(
 
         if verbose:
             print(
-                f"[get_untrained_model_with_inputs] package_source={package_source.__name__} é"
+                f"[get_untrained_model_with_inputs] package_source={package_source.__name__} "
                 f"from {package_source.__file__}"
             )
         if use_pretrained:
@@ -193,7 +215,7 @@ def get_untrained_model_with_inputs(
             )
             if verbose:
                 print(
-                    f"[get_untrained_model_with_inputs] -- done in "
+                    f"[get_untrained_model_with_inputs] -- done(1) in "
                     f"{time.perf_counter() - begin}s"
                 )
         else:
@@ -250,14 +272,36 @@ def get_untrained_model_with_inputs(
                 )
             if verbose:
                 print(
-                    f"[get_untrained_model_with_inputs] -- done in "
+                    f"[get_untrained_model_with_inputs] -- done(2) in "
                     f"{time.perf_counter() - begin}s"
                 )
 
             seed = int(os.environ.get("SEED", "17"))
             torch.manual_seed(seed)
+
+            if verbose:
+                begin = time.perf_counter()
+                print(
+                    f"[get_untrained_model_with_inputs] "
+                    f"instantiate_specific_model {cls_model}"
+                )
+
             model = instantiate_specific_model(cls_model, config)
+
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] -- done(3) in "
+                    f"{time.perf_counter() - begin}s (model is {type(model)})"
+                )
+
             if model is None:
+
+                if verbose:
+                    print(
+                        f"[get_untrained_model_with_inputs] "
+                        f"instantiate_specific_model(2) {cls_model}"
+                    )
+
                 try:
                     if type(config) is dict:
                         model = cls_model(**config)
@@ -268,6 +312,12 @@ def get_untrained_model_with_inputs(
                         f"Unable to instantiate class {cls_model.__name__} with\n{config}"
                     ) from e
 
+                if verbose:
+                    print(
+                        f"[get_untrained_model_with_inputs] -- done(4) in "
+                        f"{time.perf_counter() - begin}s (model is {type(model)})"
+                    )
+
     # input kwargs
     seed = int(os.environ.get("SEED", "17")) + 1
     torch.manual_seed(seed)