onnx-diagnostic 0.5.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -107,7 +107,7 @@ def torch_export_patches(
 ) -> Callable:
     """
     Tries to bypass some situations :func:`torch.export.export` does not support.
-    See also :ref:`l-patches-explained`.
+    See also :ref:`l-patches-explained` and :ref:`l-patch-coverage`.
 
     :param patch_sympy: fix missing method ``name`` for IntegerConstant
     :param patch_torch: patches :epkg:`torch` with supported implementation
@@ -140,6 +140,7 @@ def torch_export_patches(
     * ``torch.jit.isinstance``
     * ``torch._dynamo.mark_static_address``
     * ``torch._subclasses.fake_impls.infer_size``
+    * ``torch.vmap``
     * fix missing method ``name`` for ``sympy.S.IntegerConstant``
     * ``AttentionMaskConverter._make_causal_mask``
     * Serialization of ``MambaCache`` (in :epkg:`transformers`)
@@ -251,6 +252,7 @@ def torch_export_patches(
         if patch_torch:
             from .patches.patch_torch import (
                 patched_infer_size,
+                patched_vmap,
                 patched__broadcast_shapes,
                 _catch_produce_guards_and_solve_constraints,
                 patch__check_input_constraints_for_graph,
@@ -261,6 +263,10 @@ def torch_export_patches(
                 print(f"[torch_export_patches] stop_if_static={stop_if_static!r}")
                 print("[torch_export_patches] patch pytorch")
 
+            # torch.vmap
+            f_vmap = torch.vmap
+            torch.vmap = patched_vmap
+
             # torch.jit.isinstance
             f_jit_isinstance = torch.jit.isinstance
             torch.jit.isinstance = isinstance
@@ -328,6 +334,11 @@ def torch_export_patches(
         ####################
 
         if patch_transformers:
+            try:
+                import transformers.masking_utils as masking_utils
+            except ImportError:
+                masking_utils = None
+
             if verbose:
                 import transformers
 
@@ -339,6 +350,15 @@ def torch_export_patches(
                 patch_transformers_list, verbose=verbose
             )
 
+            if masking_utils and hasattr(masking_utils, "_vmap_for_bhqkv"):
+                if verbose:
+                    print(
+                        "[torch_export_patches] patches "
+                        "transformers.masking_utils._vmap_for_bhqkv"
+                    )
+                f_transformers__vmap_for_bhqkv = masking_utils._vmap_for_bhqkv
+                masking_utils._vmap_for_bhqkv = patch_transformers_list.patched__vmap_for_bhqkv
+
         if custom_patches:
             if verbose:
                 print("[torch_export_patches] applies custom patches")
@@ -381,6 +401,7 @@ def torch_export_patches(
 
             if patch_torch:
                 # this should disappear when torch.jit is removed
+                torch.vmap = f_vmap
                 torch.jit.isinstance = f_jit_isinstance
                 torch._dynamo.mark_static_address = f_mark_static_address
                 # tracked by https://github.com/pytorch/pytorch/issues/143495
@@ -430,12 +451,24 @@ def torch_export_patches(
             ##############
 
             if patch_transformers:
+                try:
+                    import transformers.masking_utils as masking_utils
+                except ImportError:
+                    masking_utils = None
                 if verbose:
                     print("[torch_export_patches] unpatch transformers")
                 unpatch_module_or_classes(
                     patch_transformers_list, revert_patches_info, verbose=verbose
                 )
 
+                if masking_utils and hasattr(masking_utils, "_vmap_for_bhqkv"):
+                    if verbose:
+                        print(
+                            "[torch_export_patches] unpatch "
+                            "transformers.masking_utils._vmap_for_bhqkv"
+                        )
+                    masking_utils._vmap_for_bhqkv = f_transformers__vmap_for_bhqkv
+
             ########
             # caches
             ########

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -1,5 +1,5 @@
 import pprint
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import packaging.version as pv
 import optree
 import torch
@@ -133,6 +133,11 @@ def register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
         # To avoid doing it multiple times.
         PATCH_OF_PATCHES.add(BaseModelOutput)
 
+    return serialization_functions(verbose=verbose)
+
+
+def serialization_functions(verbose: int = 0) -> Dict[str, Union[Callable, int]]:
+    """Returns the list of serialization functions."""
     return dict(
         DynamicCache=register_class_serialization(
             DynamicCache,

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,148 @@ def ast_or_into_bitor(node: "ast.Node") -> "ast.Node":
     return new_node
 
 
+@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)
+        ),
+        pre_rewriter=ast_or_into_bitor,
+    )
+
+    def _add(f):
+        g = bd.copy()
+        g["function"] = f
+        return g
+
+    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
@@ -30,34 +169,15 @@ def code_needing_rewriting(cls_name: str) -> Optional[List[Any]]:
     .. runpython::
         :showcode:
 
+        import pprint
         from onnx_diagnostic.torch_export_patches.patch_module_helper import (
             code_needing_rewriting,
         )
 
-        print(code_needing_rewriting("BartForConditionalGeneration"))
+        pprint.pprint(code_needing_rewriting("BartForConditionalGeneration"))
     """
-    if cls_name in {
-        "BartEncoderLayer",
-        "BartForConditionalGeneration",
-        "PLBartEncoderLayer",
-        "PLBartForConditionalGeneration",
-    }:
-        import transformers
-
-        bd = dict(
-            filter_node=(
-                lambda node: isinstance(node, ast.If) and not isinstance(node.test, ast.Name)
-            ),
-            pre_rewriter=ast_or_into_bitor,
-        )
-
-        def _add(f):
-            g = bd.copy()
-            g["function"] = f
-            return g
-
-        return [
-            _add(transformers.models.bart.modeling_bart.BartEncoderLayer.forward),
-            _add(transformers.models.plbart.modeling_plbart.PLBartEncoderLayer.forward),
-        ]
+    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_torch.py

@@ -370,3 +370,94 @@ class patched_ShapeEnv:
         #     RuntimeWarning,
         #     stacklevel=0,
         # )
+
+
+def patched_vmap(func, in_dims=0, out_dims=0):
+    """
+    Python implementation of :func:`torch.vmap`.
+    The implementation raises an issue when it is being exported with
+    :func:`torch.export.export` when the function is called with
+    non tensors arguments and the batch size is dynamic.
+    """
+    from ...helpers import string_type
+
+    def wrapped(*args):
+        assert all(not isinstance(a, dict) for a in args), (
+            f"dictionaries are not implemented in "
+            f"args={string_type(args, with_shape=True)}"
+        )
+
+        in_dims_ = (
+            ([in_dims] * len(args))
+            if not isinstance(in_dims, (list, tuple))
+            else list(in_dims)
+        )
+        assert len(in_dims_) == len(args), (
+            f"Mismtch between in_dims={in_dims_} and "
+            f"args={string_type(args, with_shape=True)}"
+        )
+
+        batch_size = None
+        batched_args = []
+        for arg, in_dim in zip(args, in_dims_):
+            if in_dim is None:
+                batched_args.append(arg)
+                continue
+
+            assert batch_size is None or batch_size == arg.size(in_dim), (
+                f"Unable to continue, batch_size={batch_size}, in_dim={in_dim}, "
+                f"arg.size(in_dim)={arg.size(in_dim)}"
+            )
+            if batch_size is None:
+                batch_size = arg.size(in_dim)
+            arg = arg.movedim(in_dim, 0)
+            batched_args.append(arg)
+
+        if all(isinstance(a, torch.Tensor) for a in args) and isinstance(
+            batch_size, torch.SymInt
+        ):
+            batched_tensors = [
+                (
+                    arg
+                    if (isinstance(arg, torch.Tensor) and in_dim is not None)
+                    else arg.unsqueeze(0).expand((batch_size, *arg.shape))
+                )
+                for arg, in_dim in zip(batched_args, in_dims_)
+            ]
+            results = torch.ops.higher_order.scan(func, [], batched_tensors, [])
+            stacked = results[0]
+            if out_dims != 0:
+                return stacked.movedim(0, out_dims)
+            return stacked
+
+        else:
+            torch._check(
+                not isinstance(batch_size, torch.SymInt),
+                lambda: (
+                    f"patched_vmap supports dynamic batch_size only if all argument "
+                    f"are tensors but types are {[type(a) for a in args]}"
+                ),
+            )
+            batched_tensors = [
+                (
+                    (None, arg)
+                    if (isinstance(arg, torch.Tensor) and in_dim is not None)
+                    else (arg, arg)
+                )
+                for arg, in_dim in zip(batched_args, in_dims_)
+            ]
+
+            results = []
+            for i in range(batch_size):
+                input_slice = [v if v is not None else arg[i] for v, arg in batched_tensors]
+                result = func(*input_slice)
+                results.append(result)
+
+            if isinstance(results[0], torch.Tensor):
+                stacked = torch.stack(results)
+                if out_dims != 0:
+                    return stacked.movedim(0, out_dims)
+                return stacked
+            return results
+
+    return wrapped

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1,6 +1,7 @@
 import inspect
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional, Tuple
 import torch
 import transformers
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
@@ -9,6 +10,34 @@ from ...ext_test_case import has_transformers
 from ...helpers.torch_helper import is_torchdynamo_exporting
 
 
+def patched__vmap_for_bhqkv(mask_function: Callable, bh_indices: bool = True) -> Callable:
+    """Patch for function ``transformers.masking_utils._vmap_for_bhqkv``."""
+    from ...helpers import string_type
+
+    dimensions: List[Tuple[Optional[int], ...]] = [
+        (None, None, None, 0),
+        (None, None, 0, None),
+    ]
+    if bh_indices:
+        dimensions.extend([(None, 0, None, None), (0, None, None, None)])
+    dimensions = [tuple(1 if d is None else -1 for d in shape) for shape in dimensions]
+    dimensions = tuple(reversed(dimensions))
+    indices = tuple(shape.index(-1) for shape in dimensions)
+
+    def vector_mask_function(
+        *args, mask_function=mask_function, dimensions=dimensions, indices=indices
+    ):
+        assert len(args) == len(
+            dimensions
+        ), f"Mismatch between args={string_type(args)} and dimensions={dimensions}"
+        new_args = [a.reshape(shape) for a, shape in zip(args, dimensions)]
+        max_shape = tuple(args[i].shape[0] for i in indices)
+        expanded_args = [a.expand(max_shape) for a in new_args]
+        return mask_function(*expanded_args)
+
+    return vector_mask_function
+
+
 def _patch_make_causal_mask(
     input_ids_shape: torch.Size,
     dtype: torch.dtype,
@@ -503,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,
+        }
+    )