onnx-diagnostic 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

onnx_diagnostic/torch_export_patches/patch_module_helper.py

@@ -1,5 +1,6 @@
 import ast
-from typing import Any, List, Optional
+import functools
+from typing import Any, Dict, List, Optional
 
 
 class OrToBitOrTransformer(ast.NodeTransformer):
@@ -19,10 +20,129 @@ def ast_or_into_bitor(node: "ast.Node") -> "ast.Node":
     return new_node
 
 
-def _rewrite_bart_encoder_layer():
-    "BartEncoderLayer, PLBartEncoderLayer"
+@functools.lru_cache
+def _rewrite_forward_clamp_float16() -> Dict[str, List[type]]:
+
     import transformers
 
+    _known = {
+        "AutoformerEncoderLayer": [
+            transformers.models.autoformer.modeling_autoformer.AutoformerEncoderLayer
+        ],
+        "BartEncoderLayer": [
+            transformers.models.bart.modeling_bart.BartEncoderLayer,
+            transformers.models.plbart.modeling_plbart.PLBartEncoderLayer,
+        ],
+        "BigBirdPegasusEncoderLayer": [
+            transformers.models.bigbird_pegasus.modeling_bigbird_pegasus.BigBirdPegasusEncoderLayer
+        ],
+        "BlenderbotSmallEncoderLayer": [
+            transformers.models.blenderbot_small.modeling_blenderbot_small.BlenderbotSmallEncoderLayer
+        ],
+        "InformerEncoderLayer": [
+            transformers.models.informer.modeling_informer.InformerEncoderLayer
+        ],
+        "LEDEncoderLayer": [transformers.models.led.modeling_led.LEDEncoderLayer],
+        "MarianEncoderLayer": [transformers.models.marian.modeling_marian.MarianEncoderLayer],
+        "MvpEncoderLayer": [transformers.models.mvp.modeling_mvp.MvpEncoderLayer],
+        "NllbMoeEncoderLayer": [
+            transformers.models.nllb_moe.modeling_nllb_moe.NllbMoeEncoderLayer
+        ],
+        "TimeSeriesTransformerEncoderLayer": [
+            transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesTransformerEncoderLayer
+        ],
+    }
+    return _known
+
+
+@functools.lru_cache
+def known_transformers_rewritings_clamp_float16() -> Dict[str, str]:
+    """
+    This functions returns the list of known classes to be rewritten.
+    in :epkg:`transformers`. Each class is mapped to an alias,
+    this alias is then given to :func:`rewritings_transformers_clamp_float16`
+    to rewrite the encoder layers because of a specific control flow.
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        from onnx_diagnostic.torch_export_patches.patch_module_helper import (
+            known_transformers_rewritings_clamp_float16,
+        )
+
+        pprint.pprint(known_transformers_rewritings_clamp_float16())
+    """
+    _alias = {
+        "AutoformerEncoder": "AutoformerEncoderLayer",
+        "AutoformerEncoderLayer": "AutoformerEncoderLayer",
+        "AutoformerForPrediction": "AutoformerEncoderLayer",
+        "AutoformerModel": "AutoformerEncoderLayer",
+        "BartEncoderLayer": "BartEncoderLayer",
+        "BartForConditionalGeneration": "BartEncoderLayer",
+        "BigBirdPegasusForConditionalGeneration": "BigBirdPegasusEncoderLayer",
+        "BigBirdPegasusForQuestionAnswering": "BigBirdPegasusEncoderLayer",
+        "BigBirdPegasusForCausalLM": "BigBirdPegasusEncoderLayer",
+        "BlenderbotSmallEncoderLayer": "BlenderbotSmallEncoderLayer",
+        "BlenderbotSmallForConditionalGeneration": "BlenderbotSmallEncoderLayer",
+        "BlenderbotSmallForCausalLM": "BlenderbotSmallEncoderLayer",
+        "InformerEncoderLayer": "InformerEncoderLayer",
+        "InformerForPrediction": "InformerEncoderLayer",
+        "LEDEncoderLayer": "LEDEncoderLayer",
+        "LEDClassificationHead": "LEDEncoderLayer",
+        "LEDForConditionalGeneration": "LEDEncoderLayer",
+        "MarianEncoderLayer": "MarianEncoderLayer",
+        "MarianEncoder": "MarianEncoderLayer",
+        "MarianModel": "MarianEncoderLayer",
+        "MarianMTModel": "MarianEncoderLayer",
+        "MvpEncoderLayer": "MvpEncoderLayer",
+        "MvpPrompt": "MvpEncoderLayer",
+        "MvpForConditionalGeneration": "MvpEncoderLayer",
+        "MvpForSequenceClassification": "MvpEncoderLayer",
+        "MvpForQuestionAnswering": "MvpEncoderLayer",
+        "MvpForCausalLM": "MvpEncoderLayer",
+        "NllbMoeEncoderLayer": "NllbMoeEncoderLayer",
+        "NllbMoeForConditionalGeneration": "NllbMoeEncoderLayer",
+        "PLBartEncoderLayer": "BartEncoderLayer",
+        "PLBartForConditionalGeneration": "BartEncoderLayer",
+        "TimeSeriesTransformerEncoderLayer": "TimeSeriesTransformerEncoderLayer",
+        "TimeSeriesTransformerForPrediction": "TimeSeriesTransformerEncoderLayer",
+    }
+    return _alias
+
+
+def rewritings_transformers_clamp_float16(cls_name) -> List[type]:
+    """
+    Rewrites known control flows equal to this:
+
+    .. code-block:: python
+
+        if hidden_states.dtype == torch.float16 and (
+            torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()
+        ):
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+    *cls_name* is the class name. It is mapped with a list of other class names
+    to rename. Here is the known list:
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        from onnx_diagnostic.torch_export_patches.patch_module_helper import (
+            _rewrite_forward_clamp_float16,
+        )
+
+        pprint.pprint(_rewrite_forward_clamp_float16())
+
+    Function `_rewrite_forward_clamp_float16` collects
+    all model classes using those layers.
+    """
+    _known = _rewrite_forward_clamp_float16()
+
+    assert cls_name in _known, f"cls_name={cls_name!r} unknown in {sorted(_known)}."
+
     bd = dict(
         filter_node=(
             lambda node: isinstance(node, ast.If) and not isinstance(node.test, ast.Name)
@@ -35,16 +155,13 @@ def _rewrite_bart_encoder_layer():
         g["function"] = f
         return g
 
-    return [
-        _add(transformers.models.bart.modeling_bart.BartEncoderLayer.forward),
-        _add(transformers.models.plbart.modeling_plbart.PLBartEncoderLayer.forward),
-    ]
+    return [_add(cls.forward) for cls in _known[cls_name]]
 
 
 def code_needing_rewriting(cls_name: str) -> Optional[List[Any]]:
     """
-    Returns a known list of methods or functions to rewrite because of control flow
-    for a specific model class.
+    Returns a known list of classes mapped to a known rewritings
+    because of control flow. See :func:`known_transformers_rewritings_clamp_float16`.
 
     :param cls_name: name of the class
     :return: a list of rewriting
@@ -59,11 +176,8 @@ def code_needing_rewriting(cls_name: str) -> Optional[List[Any]]:
 
         pprint.pprint(code_needing_rewriting("BartForConditionalGeneration"))
     """
-    if cls_name in {
-        "BartEncoderLayer",
-        "BartForConditionalGeneration",
-        "PLBartEncoderLayer",
-        "PLBartForConditionalGeneration",
-    }:
-        return _rewrite_bart_encoder_layer()
+    aliases = known_transformers_rewritings_clamp_float16()
+    if cls_name in aliases:
+        alias = aliases[cls_name]
+        return rewritings_transformers_clamp_float16(alias)
     return None

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1,5 +1,6 @@
 import inspect
 from dataclasses import dataclass
+from functools import wraps
 from typing import Any, Callable, Dict, List, Optional, Tuple
 import torch
 import transformers
@@ -531,3 +532,90 @@ class patched_GenerationMixin:
         # 8. Remove unexpected `generate` inputs (TODO @joao: fix trainer and examples)
         model_inputs.pop("labels", None)
         return model_inputs
+
+
+def patched_dynamic_rope_update(rope_forward):
+    """
+    patch:transformers.modeling_rope_utils.dynamic_rope_update
+    """
+
+    def longrope_frequency_update(self, position_ids, device):
+        seq_len = torch.max(position_ids) + 1
+        if hasattr(self.config, "original_max_position_embeddings"):
+            original_max_position_embeddings = self.config.original_max_position_embeddings
+        else:
+            original_max_position_embeddings = self.config.max_position_embeddings
+        # At export time, seq_len is unknown.
+        long_inv_freq, _ = self.rope_init_fn(
+            self.config, device, seq_len=original_max_position_embeddings + 1
+        )
+        original_inv_freq = self.original_inv_freq.to(device)
+
+        cond = (seq_len > original_max_position_embeddings).item()
+        inv_freq = torch.cond(
+            cond,
+            (lambda x, y: x.clone()),
+            (lambda x, y: y.clone()),
+            [long_inv_freq, original_inv_freq],
+        )
+        self.inv_freq = inv_freq
+        # if seq_len > original_max_position_embeddings:
+        #    self.inv_freq = self.long_inv_freq
+        # else:
+        #    self.inv_freq = self.original_inv_freq
+
+    def dynamic_frequency_update(self, position_ids, device):
+        seq_len = torch.max(position_ids) + 1
+        if seq_len > self.max_seq_len_cached:  # growth
+            inv_freq, self.attention_scaling = self.rope_init_fn(
+                self.config, device, seq_len=seq_len
+            )
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
+            self.max_seq_len_cached = seq_len
+
+        if (
+            seq_len < self.original_max_seq_len
+            and self.max_seq_len_cached > self.original_max_seq_len
+        ):
+            self.original_inv_freq = self.original_inv_freq.to(device)
+            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
+            self.max_seq_len_cached = self.original_max_seq_len
+
+    @wraps(rope_forward)
+    def wrapper(self, x, position_ids):
+        if "dynamic" in self.rope_type:
+            dynamic_frequency_update(self, position_ids, device=x.device)
+        elif self.rope_type == "longrope":
+            longrope_frequency_update(self, position_ids, device=x.device)
+        return rope_forward(self, x, position_ids)
+
+    return wrapper
+
+
+class patched_Phi3RotaryEmbedding(torch.nn.Module):
+    _PATCHES_ = ["forward"]
+    _PATCHED_CLASS_ = transformers.models.phi3.modeling_phi3.Phi3RotaryEmbedding
+
+    @torch.no_grad()
+    @patched_dynamic_rope_update
+    def forward(self, x, position_ids):
+        inv_freq_expanded = (
+            self.inv_freq[None, :, None]
+            .float()
+            .expand(position_ids.shape[0], -1, 1)
+            .to(x.device)
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = (
+            x.device.type
+            if isinstance(x.device.type, str) and x.device.type != "mps"
+            else "cpu"
+        )
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)

onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py

@@ -3951,3 +3951,145 @@ def _ccached_facebook_bart_large_cnn():
             "vocab_size": 50264,
         }
     )
+
+
+def _ccached_microsoft_phi4_reasoning():
+    "microsoft/Phi-4-mini-reasoning"
+    return transformers.Phi3Config(
+        **{
+            "architectures": ["Phi3ForCausalLM"],
+            "attention_bias": false,
+            "attention_dropout": 0.0,
+            "bos_token_id": 199999,
+            "embd_pdrop": 0.0,
+            "eos_token_id": 199999,
+            "full_attn_mod": 1,
+            "hidden_act": "silu",
+            "hidden_size": 3072,
+            "initializer_range": 0.02,
+            "intermediate_size": 8192,
+            "interpolate_factor": 1,
+            "lm_head_bias": false,
+            "max_position_embeddings": 131072,
+            "mlp_bias": false,
+            "model_type": "phi3",
+            "num_attention_heads": 24,
+            "num_hidden_layers": 32,
+            "num_key_value_heads": 8,
+            "original_max_position_embeddings": 4096,
+            "pad_token_id": 199999,
+            "partial_rotary_factor": 0.75,
+            "resid_pdrop": 0.0,
+            "rms_norm_eps": 1e-05,
+            "rope_scaling": {
+                "long_factor": [
+                    1,
+                    1.118320672,
+                    1.250641126,
+                    1.398617824,
+                    1.564103225,
+                    1.74916897,
+                    1.956131817,
+                    2.187582649,
+                    2.446418898,
+                    2.735880826,
+                    3.059592084,
+                    3.421605075,
+                    3.826451687,
+                    4.279200023,
+                    4.785517845,
+                    5.351743533,
+                    5.984965424,
+                    6.693110555,
+                    7.485043894,
+                    8.370679318,
+                    9.36110372,
+                    10.4687158,
+                    11.70738129,
+                    13.09260651,
+                    14.64173252,
+                    16.37415215,
+                    18.31155283,
+                    20.47818807,
+                    22.90118105,
+                    25.61086418,
+                    28.64115884,
+                    32.03,
+                    32.1,
+                    32.13,
+                    32.23,
+                    32.6,
+                    32.61,
+                    32.64,
+                    32.66,
+                    32.7,
+                    32.71,
+                    32.93,
+                    32.97,
+                    33.28,
+                    33.49,
+                    33.5,
+                    44.16,
+                    47.77,
+                ],
+                "short_factor": [
+                    1,
+                    1.118320672,
+                    1.250641126,
+                    1.398617824,
+                    1.564103225,
+                    1.74916897,
+                    1.956131817,
+                    2.187582649,
+                    2.446418898,
+                    2.735880826,
+                    3.059592084,
+                    3.421605075,
+                    3.826451687,
+                    4.279200023,
+                    4.785517845,
+                    5.351743533,
+                    5.984965424,
+                    6.693110555,
+                    7.485043894,
+                    8.370679318,
+                    9.36110372,
+                    10.4687158,
+                    11.70738129,
+                    13.09260651,
+                    14.64173252,
+                    16.37415215,
+                    18.31155283,
+                    20.47818807,
+                    22.90118105,
+                    25.61086418,
+                    28.64115884,
+                    32.03,
+                    32.1,
+                    32.13,
+                    32.23,
+                    32.6,
+                    32.61,
+                    32.64,
+                    32.66,
+                    32.7,
+                    32.71,
+                    32.93,
+                    32.97,
+                    33.28,
+                    33.49,
+                    33.5,
+                    44.16,
+                    47.77,
+                ],
+                "type": "longrope",
+            },
+            "rope_theta": 10000.0,
+            "sliding_window": 262144,
+            "tie_word_embeddings": true,
+            "torch_dtype": "bfloat16",
+            "transformers_version": "4.50.0",
+            "use_cache": true,
+            "vocab_size": 200064,
+        }
+    )