optimum-rbln 0.1.12__py3-none-any.whl → 0.1.15__py3-none-any.whl

optimum/rbln/modeling_diffusers.py

@@ -0,0 +1,329 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+import copy
+import importlib
+from os import PathLike
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import torch
+
+from .modeling import RBLNModel
+from .modeling_config import RUNTIME_KEYWORDS, ContextRblnConfig, use_rbln_config
+from .utils.decorator_utils import remove_compile_time_kwargs
+
+
+if TYPE_CHECKING:
+    from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
+
+
+class RBLNDiffusionMixin:
+    """
+    RBLNDiffusionMixin provides essential functionalities for compiling Stable Diffusion pipeline components to run on RBLN NPUs.
+    This mixin class serves as a base for implementing RBLN-compatible Stable Diffusion pipelines. It contains shared logic for
+    handling the core components of Stable Diffusion.
+
+    To use this mixin:
+
+    1. Create a new pipeline class that inherits from both this mixin and the original StableDiffusionPipeline.
+    2. Define the required _submodules class variable listing the components to be compiled.
+    3. If needed, implement get_default_rbln_config for custom configuration of submodules.
+
+    Example:
+        ```python
+        class RBLNStableDiffusionPipeline(RBLNDiffusionMixin, StableDiffusionPipeline):
+            _submodules = ["text_encoder", "unet", "vae"]
+
+            @classmethod
+            def get_default_rbln_config(cls, model, submodule_name, rbln_config):
+                # Configuration for other submodules...
+                pass
+        ```
+
+    Class Variables:
+        _submodules: List of submodule names that should be compiled (typically ["text_encoder", "unet", "vae"])
+
+    Methods:
+        from_pretrained: Creates and optionally compiles a model from a pretrained checkpoint
+
+    Notes:
+        - When `export=True`, all compatible submodules will be compiled for NPU inference
+        - The compilation config can be customized per submodule by including submodule names
+          as keys in rbln_config
+    """
+
+    _submodules = []
+
+    @classmethod
+    @property
+    def img2img_pipeline(cls):
+        return "Img2Img" in cls.__name__
+
+    @classmethod
+    @property
+    def inpaint_pipeline(cls):
+        return "Inpaint" in cls.__name__
+
+    @classmethod
+    def get_submodule_rbln_config(
+        cls, model: torch.nn.Module, submodule_name: str, rbln_config: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        submodule = getattr(model, submodule_name)
+        submodule_class_name = submodule.__class__.__name__
+
+        if submodule_class_name == "MultiControlNetModel":
+            submodule_class_name = "ControlNetModel"
+
+        submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), f"RBLN{submodule_class_name}")
+
+        submodule_config = rbln_config.get(submodule_name, {})
+        submodule_config = copy.deepcopy(submodule_config)
+
+        pipe_global_config = {k: v for k, v in rbln_config.items() if k not in cls._submodules}
+
+        submodule_config.update({k: v for k, v in pipe_global_config.items() if k not in submodule_config})
+        submodule_config.update(
+            {
+                "img2img_pipeline": cls.img2img_pipeline,
+                "inpaint_pipeline": cls.inpaint_pipeline,
+            }
+        )
+        submodule_config = submodule_cls.update_rbln_config_using_pipe(model, submodule_config)
+        return submodule_config
+
+    @staticmethod
+    def _maybe_apply_and_fuse_lora(
+        model: torch.nn.Module,
+        lora_ids: Optional[Union[str, List[str]]] = None,
+        lora_weights_names: Optional[Union[str, List[str]]] = None,
+        lora_scales: Optional[Union[float, List[float]]] = None,
+    ) -> torch.nn.Module:
+        lora_ids = [lora_ids] if isinstance(lora_ids, str) else lora_ids
+        lora_weights_names = [lora_weights_names] if isinstance(lora_weights_names, str) else lora_weights_names
+        lora_scales = [lora_scales] if isinstance(lora_scales, float) else lora_scales
+
+        # adapt lora weight into pipeline before compilation
+        if lora_ids and lora_weights_names:
+            if len(lora_ids) == 1:
+                if len(lora_ids) != len(lora_weights_names):
+                    raise ValueError(
+                        f"You must define the same number of lora ids ({len(lora_ids)} and lora weights ({len(lora_weights_names)}))"
+                    )
+                else:
+                    model.load_lora_weights(lora_ids[0], weight_name=lora_weights_names[0])
+                    model.fuse_lora(lora_scale=lora_scales[0] if lora_scales else 1.0)
+            elif len(lora_ids) > 1:
+                if not len(lora_ids) == len(lora_weights_names):
+                    raise ValueError(
+                        f"If you fuse {len(lora_ids)} lora models, but you must define the same number for lora weights and adapters."
+                    )
+
+                adapter_names = [f"adapter_{i}" for i in range(len(lora_ids))]
+
+                for lora_id, lora_weight, adapter_name in zip(lora_ids, lora_weights_names, adapter_names):
+                    model.load_lora_weights(lora_id, weight_name=lora_weight, adapter_name=adapter_name)
+
+                if lora_scales:
+                    model.set_adapters(adapter_names, adapter_weights=lora_scales)
+
+                model.fuse_lora()
+        return model
+
+    @classmethod
+    @use_rbln_config
+    def from_pretrained(
+        cls,
+        model_id: str,
+        *,
+        export: bool = False,
+        model_save_dir: Optional[PathLike] = None,
+        rbln_config: Dict[str, Any] = {},
+        lora_ids: Optional[Union[str, List[str]]] = None,
+        lora_weights_names: Optional[Union[str, List[str]]] = None,
+        lora_scales: Optional[Union[float, List[float]]] = None,
+        **kwargs,
+    ) -> RBLNModel:
+        if export:
+            # keep submodules if user passed any of them.
+            passed_submodules = {
+                name: kwargs.pop(name) for name in cls._submodules if isinstance(kwargs.get(name), RBLNModel)
+            }
+
+        else:
+            # raise error if any of submodules are torch module.
+            model_index_config = None
+            for submodule_name in cls._submodules:
+                if isinstance(kwargs.get(submodule_name), torch.nn.Module):
+                    raise AssertionError(
+                        f"{submodule_name} is not compiled torch module. If you want to compile, set `export=True`."
+                    )
+
+                # Load submodule outside if runtime kwargs(e.g. device) is specified.
+                if submodule_config := rbln_config.get(submodule_name):
+                    if any(kwd in submodule_config for kwd in RUNTIME_KEYWORDS):
+                        if model_index_config is None:
+                            model_index_config = cls.load_config(pretrained_model_name_or_path=model_id)
+
+                        module_name, class_name = model_index_config[submodule_name]
+                        if module_name != "optimum.rbln":
+                            raise ValueError(
+                                f"Invalid module_name '{module_name}' found in model_index.json for "
+                                f"submodule '{submodule_name}'. "
+                                "Expected 'optimum.rbln'. Please check the model_index.json configuration."
+                            )
+                        submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), class_name)
+                        submodule = submodule_cls.from_pretrained(
+                            model_id, export=False, subfolder=submodule_name, rbln_config=submodule_config
+                        )
+                        kwargs[submodule_name] = submodule
+
+        with ContextRblnConfig(
+            device=rbln_config.get("device"),
+            device_map=rbln_config.get("device_map"),
+            create_runtimes=rbln_config.get("create_runtimes"),
+            optimize_host_mem=rbln_config.get("optimize_host_memory"),
+        ):
+            model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
+
+        if not export:
+            return model
+
+        model = cls._maybe_apply_and_fuse_lora(
+            model,
+            lora_ids=lora_ids,
+            lora_weights_names=lora_weights_names,
+            lora_scales=lora_scales,
+        )
+
+        compiled_submodules = cls._compile_submodules(model, passed_submodules, model_save_dir, rbln_config)
+        return cls._construct_pipe(model, compiled_submodules, model_save_dir, rbln_config)
+
+    @classmethod
+    def _compile_submodules(
+        cls,
+        model: torch.nn.Module,
+        passed_submodules: Dict[str, RBLNModel],
+        model_save_dir: Optional[PathLike],
+        rbln_config: Dict[str, Any],
+    ) -> Dict[str, RBLNModel]:
+        compiled_submodules = {}
+
+        for submodule_name in cls._submodules:
+            submodule = passed_submodules.get(submodule_name) or getattr(model, submodule_name, None)
+            submodule_rbln_config = cls.get_submodule_rbln_config(model, submodule_name, rbln_config)
+
+            if submodule is None:
+                raise ValueError(f"submodule ({submodule_name}) cannot be accessed since it is not provided.")
+            elif isinstance(submodule, RBLNModel):
+                pass
+            elif submodule_name == "controlnet" and hasattr(submodule, "nets"):
+                # In case of multicontrolnet
+                submodule = cls._compile_multicontrolnet(
+                    controlnets=submodule,
+                    model_save_dir=model_save_dir,
+                    controlnet_rbln_config=submodule_rbln_config,
+                )
+            elif isinstance(submodule, torch.nn.Module):
+                submodule_cls: RBLNModel = getattr(
+                    importlib.import_module("optimum.rbln"), f"RBLN{submodule.__class__.__name__}"
+                )
+                submodule = submodule_cls.from_model(
+                    model=submodule,
+                    subfolder=submodule_name,
+                    model_save_dir=model_save_dir,
+                    rbln_config=submodule_rbln_config,
+                )
+            else:
+                raise ValueError(f"Unknown class of submodule({submodule_name}) : {submodule.__class__.__name__} ")
+
+            compiled_submodules[submodule_name] = submodule
+        return compiled_submodules
+
+    @classmethod
+    def _compile_multicontrolnet(
+        cls,
+        controlnets: "MultiControlNetModel",
+        model_save_dir: Optional[PathLike],
+        controlnet_rbln_config: Dict[str, Any],
+    ):
+        # Compile multiple ControlNet models for a MultiControlNet setup
+        from .diffusers.models.controlnet import RBLNControlNetModel
+        from .diffusers.pipelines.controlnet import RBLNMultiControlNetModel
+
+        compiled_controlnets = [
+            RBLNControlNetModel.from_model(
+                model=controlnet,
+                subfolder="controlnet" if i == 0 else f"controlnet_{i}",
+                model_save_dir=model_save_dir,
+                rbln_config=controlnet_rbln_config,
+            )
+            for i, controlnet in enumerate(controlnets.nets)
+        ]
+        return RBLNMultiControlNetModel(compiled_controlnets, config=controlnets.nets[0].config)
+
+    @classmethod
+    def _construct_pipe(cls, model, submodules, model_save_dir, rbln_config):
+        # Construct finalize pipe setup with compiled submodules and configurations
+
+        if model_save_dir is not None:
+            # To skip saving original pytorch modules
+            for submodule_name in cls._submodules:
+                delattr(model, submodule_name)
+
+            # Direct calling of `save_pretrained` causes config.unet = (None, None).
+            # So config must be saved again, later.
+            model.save_pretrained(model_save_dir)
+            # FIXME: Here, model touches its submodules such as model.unet,
+            # Causing warning messeages.
+
+        update_dict = {}
+        for submodule_name in cls._submodules:
+            # replace submodule
+            setattr(model, submodule_name, submodules[submodule_name])
+            update_dict[submodule_name] = ("optimum.rbln", submodules[submodule_name].__class__.__name__)
+
+        # Update config to be able to load from model directory.
+        #
+        # e.g)
+        # update_dict = {
+        #     "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
+        #     "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
+        #     "unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
+        # }
+        model.register_to_config(**update_dict)
+
+        if model_save_dir:
+            # overwrite to replace incorrect config
+            model.save_config(model_save_dir)
+
+        if rbln_config.get("optimize_host_memory") is False:
+            # Keep compiled_model objs to further analysis. -> TODO: remove soon...
+            model.compiled_models = []
+            for name in cls._submodules:
+                submodule = getattr(model, name)
+                model.compiled_models.extend(submodule.compiled_models)
+
+        return model
+
+    @remove_compile_time_kwargs
+    def __call__(self, *args, **kwargs):
+        return super().__call__(*args, **kwargs)

optimum/rbln/transformers/__init__.py

@@ -28,7 +28,6 @@ from transformers.utils import _LazyModule
 
 _import_structure = {
     "cache_utils": ["RebelDynamicCache"],
-    "generation": ["BatchTextIteratorStreamer"],
     "models": [
         "RBLNAutoModel",
         "RBLNAutoModelForAudioClassification",
@@ -57,6 +56,7 @@ _import_structure = {
         "RBLNWhisperForConditionalGeneration",
         "RBLNLlamaForCausalLM",
         "RBLNPhiForCausalLM",
+        "RBLNT5EncoderModel",
         "RBLNT5ForConditionalGeneration",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNMidmLMHeadModel",
@@ -67,7 +67,6 @@ _import_structure = {
 
 if TYPE_CHECKING:
     from .cache_utils import RebelDynamicCache
-    from .generation import BatchTextIteratorStreamer
     from .models import (
         RBLNAutoModel,
         RBLNAutoModelForAudioClassification,
@@ -97,6 +96,7 @@ if TYPE_CHECKING:
         RBLNMistralForCausalLM,
         RBLNPhiForCausalLM,
         RBLNQwen2ForCausalLM,
+        RBLNT5EncoderModel,
         RBLNT5ForConditionalGeneration,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,

optimum/rbln/transformers/cache_utils.py

@@ -12,9 +12,11 @@ class RebelDynamicCache(DynamicCache):
     `[batch_size, num_heads, seq_len, head_dim]`.
     """
 
-    def __init__(self, current_steps) -> None:
+    def __init__(self, position_ids) -> None:
         super().__init__()
-        self.current_steps = current_steps
+        # batch, _ = position_ids.shape
+        # current_steps = [position_ids[b][0] for b in range(batch)]
+        self.current_steps = position_ids[:, 0]
 
     def assign(
         self,
@@ -58,13 +60,7 @@ class RebelDynamicCache(DynamicCache):
     @classmethod
     def from_input_format(cls, position_ids, num_hidden_layer, *past_key_values) -> "DynamicCache":
         """Converts a cache in the rbln cache format (list of past_kv) into an equivalent `DynamicCache`."""
-
-        batch, _ = position_ids.shape
-        current_steps = [position_ids[b][0] for b in range(batch)]
-
-        assert len(current_steps) == batch
-        cache = cls(current_steps)
-
+        cache = cls(position_ids)
         for layer_idx in range(num_hidden_layer):
             key_states = past_key_values[layer_idx * 2]
             value_states = past_key_values[layer_idx * 2 + 1]

optimum/rbln/transformers/modeling_rope_utils.py

@@ -0,0 +1,283 @@
+import math
+from typing import Optional, Tuple
+
+import torch
+from transformers import PretrainedConfig
+
+
+def _compute_default_rope_parameters(
+    config: Optional[PretrainedConfig] = None,
+    seq_len: Optional[int] = None,
+) -> Tuple["torch.Tensor", float]:
+    """
+    Computes the inverse frequencies according to the original RoPE implementation
+    Args:
+        config ([`~transformers.PretrainedConfig`]):
+            The model configuration.
+        seq_len (`int`, *optional*):
+            The current sequence length. Unused for this type of RoPE.
+    Returns:
+        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+    """
+
+    base = config.rope_theta
+    partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    dim = int(head_dim * partial_rotary_factor)
+
+    attention_factor = 1.0  # Unused in this type of RoPE
+
+    # Compute the inverse frequencies
+    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
+    return inv_freq, attention_factor
+
+
+def _compute_linear_scaling_rope_parameters(
+    config: Optional[PretrainedConfig] = None,
+    seq_len: Optional[int] = None,
+) -> Tuple["torch.Tensor", float]:
+    """
+    Computes the inverse frequencies with linear scaling. Credits to the Reddit user /u/kaiokendev
+    Args:
+        config ([`~transformers.PretrainedConfig`]):
+            The model configuration.
+        seq_len (`int`, *optional*):
+            The current sequence length. Unused for this type of RoPE.
+    Returns:
+        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+    """
+
+    factor = config.rope_scaling["factor"]
+
+    # Gets the default RoPE parameters
+    inv_freq, attention_factor = _compute_default_rope_parameters(config, seq_len)
+
+    # Then applies linear scaling to the frequencies.
+    # NOTE: originally, scaling was applied to the position_ids. However, we get `embs = inv_freq @ position_ids`, so
+    # applying scaling to the inverse frequencies is equivalent.
+    inv_freq /= factor
+    return inv_freq, attention_factor
+
+
+def _compute_dynamic_ntk_parameters(
+    config: Optional[PretrainedConfig] = None,
+    seq_len: Optional[int] = None,
+) -> Tuple["torch.Tensor", float]:
+    """
+    Computes the inverse frequencies with NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla
+    Args:
+        config ([`~transformers.PretrainedConfig`]):
+            The model configuration.
+        seq_len (`int`, *optional*):
+            The current sequence length, used to update the dynamic RoPE at inference time.
+    Returns:
+        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+    """
+
+    base = config.rope_theta
+    partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    dim = int(head_dim * partial_rotary_factor)
+    max_position_embeddings = config.max_position_embeddings
+    factor = config.rope_scaling["factor"]
+
+    attention_factor = 1.0  # Unused in this type of RoPE
+
+    # Process with chunk_size to reduce precesion error
+    chunk_size = 4096
+    chunks = (seq_len + chunk_size - 1) // chunk_size
+
+    inv_freq_list = []
+    for i in range(chunks):
+        start = i * chunk_size
+        end = min((i + 1) * chunk_size, seq_len)
+
+        seq_lens = torch.arange(start, end, dtype=torch.float32).view(-1, 1) + 1.0
+        seq_lens = torch.where(seq_lens > max_position_embeddings, seq_lens, max_position_embeddings)
+
+        # Compute the inverse frequencies for each chunk
+        scaled_base = base * ((factor * seq_lens / max_position_embeddings) - (factor - 1)) ** (dim / (dim - 2))
+        inv_freq = 1.0 / (scaled_base ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
+
+        inv_freq_list.append(inv_freq)
+
+    final_inv_freq = torch.cat(inv_freq_list, dim=0)
+
+    return final_inv_freq, attention_factor
+
+
+def _compute_yarn_parameters(config: PretrainedConfig, seq_len: Optional[int] = None) -> Tuple["torch.Tensor", float]:
+    """
+    Computes the inverse frequencies with NTK scaling. Please refer to the
+    [original paper](https://arxiv.org/abs/2309.00071)
+    Args:
+        config ([`~transformers.PretrainedConfig`]):
+            The model configuration.
+        seq_len (`int`, *optional*):
+            The current sequence length. Unused for this type of RoPE.
+    Returns:
+        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+        post-processing scaling factor applied to the computed cos/sin.
+    """
+
+    base = config.rope_theta
+    partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    dim = int(head_dim * partial_rotary_factor)
+    max_position_embeddings = config.max_position_embeddings
+    factor = config.rope_scaling["factor"]
+
+    # Sets the attention factor as suggested in the paper
+    attention_factor = config.rope_scaling.get("attention_factor")
+    if attention_factor is None:
+        attention_factor = 0.1 * math.log(factor) + 1.0
+
+    # Optional config options
+    # beta_fast/beta_slow: as suggested in the paper, default to 32/1 (correspondingly)
+    beta_fast = config.rope_scaling.get("beta_fast") or 32
+    beta_slow = config.rope_scaling.get("beta_slow") or 1
+
+    # Compute the inverse frequencies
+    def find_correction_dim(num_rotations, dim, base, max_position_embeddings):
+        """Inverse dimension formula to find the dimension based on the number of rotations"""
+        return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))
+
+    def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings):
+        """Find dimension range bounds based on rotations"""
+        low = math.floor(find_correction_dim(low_rot, dim, base, max_position_embeddings))
+        high = math.ceil(find_correction_dim(high_rot, dim, base, max_position_embeddings))
+        return max(low, 0), min(high, dim - 1)
+
+    def linear_ramp_factor(min, max, dim):
+        if min == max:
+            max += 0.001  # Prevent singularity
+
+        linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
+        ramp_func = torch.clamp(linear_func, 0, 1)
+        return ramp_func
+
+    # Note on variable naming: "interpolation" comes from the original technique, where we interpolate the position IDs
+    # to expand the possible context length. In other words, interpolation = apply scaling factor.
+    pos_freqs = base ** (torch.arange(0, dim, 2).float() / dim)
+    inv_freq_extrapolation = 1.0 / pos_freqs
+    inv_freq_interpolation = 1.0 / (factor * pos_freqs)
+
+    low, high = find_correction_range(beta_fast, beta_slow, dim, base, max_position_embeddings)
+
+    # Get n-dimensional rotational scaling corrected for extrapolation
+    inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2).float()
+    inv_freq = (
+        inv_freq_interpolation * (1 - inv_freq_extrapolation_factor)
+        + inv_freq_extrapolation * inv_freq_extrapolation_factor
+    )
+
+    return inv_freq, attention_factor
+
+
+def _compute_longrope_parameters(
+    config: PretrainedConfig, seq_len: Optional[int] = None
+) -> Tuple["torch.Tensor", float]:
+    """
+    Computes the inverse frequencies with LongRoPE scaling. Please refer to the
+    [original implementation](https://github.com/microsoft/LongRoPE)
+    Args:
+        config ([`~transformers.PretrainedConfig`]):
+            The model configuration.
+        seq_len (`int`, *optional*):
+            The current sequence length. Unused for this type of RoPE.
+    Returns:
+        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+        post-processing scaling factor applied to the computed cos/sin.
+    """
+
+    base = config.rope_theta
+    partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    dim = int(head_dim * partial_rotary_factor)
+    long_factor = config.rope_scaling["long_factor"]
+    short_factor = config.rope_scaling["short_factor"]
+    factor = config.rope_scaling.get("factor")
+    attention_factor = config.rope_scaling.get("attention_factor")
+
+    # NOTE: Phi3 (and potentially other models) modify `max_position_embeddings` and have a
+    # `original_max_position_embeddings` field containing the pretrained value. They use the ratio between these two
+    # values to compute the default attention scaling factor, instead of using `factor`.
+    if hasattr(config, "original_max_position_embeddings"):
+        max_position_embeddings = config.original_max_position_embeddings
+        expanded_max_position_embeddings = config.max_position_embeddings
+        factor = expanded_max_position_embeddings / max_position_embeddings
+    else:
+        max_position_embeddings = config.max_position_embeddings
+        expanded_max_position_embeddings = max_position_embeddings * factor
+
+    # Sets the attention factor as suggested in the paper
+    if attention_factor is None:
+        if factor <= 1.0:
+            attention_factor = 1.0
+        else:
+            attention_factor = math.sqrt(1 + math.log(factor) / math.log(max_position_embeddings))
+
+    # Compute the inverse frequencies -- scaled based on the target sequence length
+    if expanded_max_position_embeddings > max_position_embeddings:
+        ext_factors = torch.tensor(long_factor, dtype=torch.float32)
+    else:
+        ext_factors = torch.tensor(short_factor, dtype=torch.float32)
+    inv_freq_shape = torch.arange(0, dim, 2, dtype=torch.int64).float() / dim
+    inv_freq = 1.0 / (ext_factors * base**inv_freq_shape)
+
+    return inv_freq, attention_factor
+
+
+def _compute_llama3_parameters(
+    config: PretrainedConfig, seq_len: Optional[int] = None
+) -> Tuple["torch.Tensor", float]:
+    """
+    Computes the inverse frequencies for llama 3.1.
+
+    Args:
+        config ([`~transformers.PretrainedConfig`]):
+            The model configuration.
+        seq_len (`int`, *optional*):
+            The current sequence length. Unused for this type of RoPE.
+    Returns:
+        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+        post-processing scaling factor applied to the computed cos/sin.
+    """
+    # Gets the default RoPE parameters
+    inv_freq, attention_factor = _compute_default_rope_parameters(config, seq_len)
+
+    factor = config.rope_scaling["factor"]  # `8` in the original implementation
+    low_freq_factor = config.rope_scaling["low_freq_factor"]  # `1` in the original implementation
+    high_freq_factor = config.rope_scaling["high_freq_factor"]  # `4` in the original implementation
+    old_context_len = config.rope_scaling["original_max_position_embeddings"]  # `8192` in the original implementation
+
+    low_freq_wavelen = old_context_len / low_freq_factor
+    high_freq_wavelen = old_context_len / high_freq_factor
+
+    wavelen = 2 * math.pi / inv_freq
+    # wavelen < high_freq_wavelen: do nothing
+    # wavelen > low_freq_wavelen: divide by factor
+    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq)
+    # otherwise: interpolate between the two, using a smooth factor
+    smooth_factor = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
+    smoothed_inv_freq = (1 - smooth_factor) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
+    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
+    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
+
+    return inv_freq_llama, attention_factor
+
+
+# This maps the "rope_type" string field in rope config to the corresponding function to compute the RoPE parameters
+# from the model config. You can append new {'rope_type': callable} pairs to this dictionary to enable custom RoPE
+# parameterizations, as long as the callable has the same signature.
+ROPE_INIT_FUNCTIONS = {
+    "default": _compute_default_rope_parameters,
+    "linear": _compute_linear_scaling_rope_parameters,
+    "dynamic": _compute_dynamic_ntk_parameters,
+    "yarn": _compute_yarn_parameters,
+    "longrope": _compute_longrope_parameters,
+    "llama3": _compute_llama3_parameters,
+}