chore(root): Updates source files with RC versions.

config.json

@@ -125,7 +125,7 @@
       2.849999999999998,
       2.9499999999999975
     ],
-    "type": "longrope"
+    "type": "su"
   },
   "rope_theta": 10000.0,
   "sliding_window": 262144,

configuration_phi3.py

@@ -83,8 +83,8 @@ class Phi3Config(PretrainedConfig):
         rope_theta (`float`, *optional*, defaults to 10000.0):
             The base period of the RoPE embeddings.
         rope_scaling (`dict`, *optional*):
-            The scaling factor for the RoPE embeddings. If `None`, no scaling is applied. If a dictionary, it must
-            contain the following keys: `type`, `short_factor` and `long_factor`. The `type` must be `longrope` and
+            The scaling strategy for the RoPE embeddings. If `None`, no scaling is applied. If a dictionary, it must
+            contain the following keys: `type`, `short_factor` and `long_factor`. The `type` must be either `su` or `yarn` and
             the `short_factor` and `long_factor` must be lists of numbers with the same length as the hidden size
             divided by the number of attention heads divided by 2.
         eos_token_id (`int`, *optional*, defaults to 32000):
@@ -158,6 +158,7 @@ class Phi3Config(PretrainedConfig):
         self.use_cache = use_cache
         self.rope_theta = rope_theta
         self.rope_scaling = rope_scaling
+        self._rope_scaling_validation()
         self.sliding_window = sliding_window
 
         super().__init__(
@@ -168,33 +169,41 @@ class Phi3Config(PretrainedConfig):
         )
 
     def _rope_scaling_validation(self):
+        """
+        Validate the `rope_scaling` configuration.
+        """
         if self.rope_scaling is None:
             return
 
-        assert (
-            (isinstance(self.rope_scaling, dict))
-            and ("type" in self.rope_scaling)
-            and ("short_factor" in self.rope_scaling)
-            and ("long_factor" in self.rope_scaling)
-        ), (
-            "`rope_scaling` must be a dictionary with three keys: `type`, `short_factor` and `long_factor`, "
-            f"got {self.rope_scaling}."
-        )
-
-        assert self.rope_scaling["type"].lower() == "longrope", "RoPE scaling type must be `longrope`."
-
-        short_factor = self.rope_scaling["short_factor"]
-        assert isinstance(short_factor, list) and all(
-            isinstance(x, (int, float)) for x in short_factor
-        ), f"RoPE scaling factor must be a list of numbers, got {short_factor}."
-        assert (
-            len(short_factor) == self.hidden_size // self.num_attention_heads // 2
-        ), f"Length of RoPE scaling factor must be half of the attention head, got {short_factor}."
-
-        long_factor = self.rope_scaling["long_factor"]
-        assert isinstance(long_factor, list) and all(
-            isinstance(x, (int, float)) for x in long_factor
-        ), f"RoPE scaling factor must be a list of numbers, got {long_factor}."
-        assert (
-            len(long_factor) == self.hidden_size // self.num_attention_heads // 2
-        ), f"Length of RoPE scaling factor must be half of the attention head, got {long_factor}."
+        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 3:
+            raise ValueError(
+                "`rope_scaling` must be a dictionary with three fields, `type`, `short_factor` and `long_factor`, "
+                f"got {self.rope_scaling}"
+            )
+        rope_scaling_type = self.rope_scaling.get("type", None)
+        rope_scaling_short_factor = self.rope_scaling.get("short_factor", None)
+        rope_scaling_long_factor = self.rope_scaling.get("long_factor", None)
+        if rope_scaling_type is None or rope_scaling_type not in ["su", "yarn"]:
+            raise ValueError(f"`rope_scaling`'s type field must be one of ['su', 'yarn'], got {rope_scaling_type}")
+        if not (
+            isinstance(rope_scaling_short_factor, list)
+            and all(isinstance(x, (int, float)) for x in rope_scaling_short_factor)
+        ):
+            raise ValueError(
+                f"`rope_scaling`'s short_factor field must be a list of numbers, got {rope_scaling_short_factor}"
+            )
+        if not len(rope_scaling_short_factor) == self.hidden_size // self.num_attention_heads // 2:
+            raise ValueError(
+                f"`rope_scaling`'s short_factor field must have length {self.hidden_size // self.num_attention_heads // 2}, got {len(rope_scaling_short_factor)}"
+            )
+        if not (
+            isinstance(rope_scaling_long_factor, list)
+            and all(isinstance(x, (int, float)) for x in rope_scaling_long_factor)
+        ):
+            raise ValueError(
+                f"`rope_scaling`'s long_factor field must be a list of numbers, got {rope_scaling_long_factor}"
+            )
+        if not len(rope_scaling_long_factor) == self.hidden_size // self.num_attention_heads // 2:
+            raise ValueError(
+                f"`rope_scaling`'s long_factor field must have length {self.hidden_size // self.num_attention_heads // 2}, got {len(rope_scaling_long_factor)}"
+            )

modeling_phi3.py

@@ -54,26 +54,17 @@ logger = logging.get_logger(__name__)
 _flash_supports_window_size = False
 try:
     from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
 
     _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
-
-    if not _flash_supports_window_size:
-        raise ValueError("Please update flash-attention to support window size.")
-
-    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
-    from flash_attn.ops.activations import swiglu
-    from flash_attn.ops.rms_norm import RMSNorm as Phi3FlashRMSNorm
-# else:
 except ImportError as error:
     logger.warning(
-        f"Flash Attention or Flash Attention Submodules not found, consider installing for better performance: {error}."
+        f"`flash-attention` package not found, consider installing for better performance: {error}."
     )
     if not _flash_supports_window_size:
         logger.warning(
-            "This version of flash does not support window size. Please use `attn_implementation='eager'` or upgrade flash-attn library."
+            "Current `flash-attenton` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`."
         )
-    swiglu = None
-    Phi3FlashRMSNorm = None
 
 _CHECKPOINT_FOR_DOC = "microsoft/Phi-3-mini-4k-instruct"
 _CONFIG_FOR_DOC = "Phi3Config"
@@ -103,9 +94,6 @@ class Phi3RMSNorm(nn.Module):
         return self.weight * hidden_states.to(input_dtype)
 
 
-PHI3_NORM_CLASS = Phi3RMSNorm if Phi3FlashRMSNorm is None else Phi3FlashRMSNorm
-
-
 # Copied from transformers.models.llama.modeling_llama._get_unpad_data
 def _get_unpad_data(attention_mask):
     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
@@ -156,43 +144,27 @@ class Phi3RotaryEmbedding(nn.Module):
         )
 
 
-class Phi3LongScaledRotaryEmbedding(nn.Module):
+class _Phi3ScaledRotaryEmbedding(nn.Module):
     def __init__(
         self,
         dim,
         short_factor,
         long_factor,
-        max_position_embeddings=4096,
-        original_max_position_embeddings=4096,
+        max_position_embeddings=2048,
+        original_max_position_embeddings=2048,
         base=10000,
-        magnitude_scaling_policy="su",
     ):
         super().__init__()
 
         self.dim = dim
+        self.short_factor = short_factor
+        self.long_factor = long_factor
         self.max_position_embeddings = max_position_embeddings
         self.original_max_position_embeddings = original_max_position_embeddings
         self.base = base
 
-        if magnitude_scaling_policy == "su":
-            self._calc_mscale = self._calc_mscale_su
-        elif magnitude_scaling_policy == "yarn":
-            self._calc_mscale = self._calc_mscale_yarn
-        else:
-            self._calc_mscale = lambda scale: float(scale)
-
-        self.short_factor = short_factor
-        self.long_factor = long_factor
-
-    def _calc_mscale_su(self, scale):
-        if scale <= 1.0:
-            return 1.0
-        return math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings))
-
-    def _calc_mscale_yarn(self, scale):
-        if scale <= 1.0:
-            return 1.0
-        return 0.1 * math.log(scale) + 1.0
+    def _calc_mscale(self, scale):
+        raise NotImplementedError("`_calc_mscale` should be implemented in subclasses")
 
     @torch.no_grad()
     def forward(self, x, seq_len=None):
@@ -206,9 +178,6 @@ class Phi3LongScaledRotaryEmbedding(nn.Module):
         else:
             t = torch.arange(self.original_max_position_embeddings, device=x.device, dtype=torch.float32)
             rescale_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device)
-        assert rescale_factors.shape == (
-            self.dim // 2,
-        ), f"misaligned shape for LongRoPE rescale factors: {rescale_factors.shape}"
 
         inv_freq = 1.0 / (
             rescale_factors * (self.base ** (torch.arange(0, self.dim, 2).float().to(x.device) / self.dim))
@@ -221,6 +190,20 @@ class Phi3LongScaledRotaryEmbedding(nn.Module):
         return (emb.cos() * mscale).to(x.dtype), (emb.sin() * mscale).to(x.dtype)
 
 
+class Phi3SuScaledRotaryEmbedding(_Phi3ScaledRotaryEmbedding):
+    def _calc_mscale(self, scale):
+        if scale <= 1.0:
+            return 1.0
+        return math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings))
+
+
+class Phi3YarnScaledRotaryEmbedding(_Phi3ScaledRotaryEmbedding):
+    def _calc_mscale(self, scale):
+        if scale <= 1.0:
+            return 1.0
+        return 0.1 * math.log(scale) + 1.0
+
+
 # Copied from transformers.models.llama.modeling_llama.rotate_half
 def rotate_half(x):
     """Rotates half the hidden dims of the input."""
@@ -253,24 +236,12 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
     cos = cos[position_ids].unsqueeze(unsqueeze_dim)
     sin = sin[position_ids].unsqueeze(unsqueeze_dim)
     # Need fp32 here to match logits
-    q_embed = (q.to(dtype=torch.float32) * cos.to(dtype=torch.float32)) + (
-        rotate_half(q).to(dtype=torch.float32) * sin.to(dtype=torch.float32)
-    )
-    k_embed = (k.to(dtype=torch.float32) * cos.to(dtype=torch.float32)) + (
-        rotate_half(k).to(dtype=torch.float32) * sin.to(dtype=torch.float32)
-    )
+    q_embed = (q.float() * cos.float()) + (rotate_half(q).float() * sin.float())
+    k_embed = (k.float() * cos.float()) + (rotate_half(k).float() * sin.float())
     return q_embed.to(q.dtype), k_embed.to(k.dtype)
 
 
 class Phi3MLP(nn.Module):
-    """Gated Linear Unit.
-
-    Reference:
-        Language Modeling with Gated Convolutional Networks.
-        https://arxiv.org/pdf/1612.08083v3.pdf.
-
-    """
-
     def __init__(self, config):
         super().__init__()
 
@@ -283,13 +254,8 @@ class Phi3MLP(nn.Module):
     def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
         y = self.gate_up_proj(hidden_states)
 
-        # Special case for SwiGLU
-        if self.config.hidden_act == "silu" and swiglu is not None:
-            gate, y = y.chunk(2, dim=-1)
-            y = swiglu(gate, y)
-        else:
-            gate, y = y.chunk(2, dim=-1)
-            y = y * self.activation_fn(gate)
+        gate, y = y.chunk(2, dim=-1)
+        y = y * self.activation_fn(gate)
 
         return self.down_proj(y)
 
@@ -330,7 +296,6 @@ class Phi3Attention(nn.Module):
         self.max_position_embeddings = config.max_position_embeddings
         self.original_max_position_embeddings = config.original_max_position_embeddings
         self.rope_theta = config.rope_theta
-        self.rope_scaling = config.rope_scaling
         self.is_causal = True
 
         if (self.head_dim * self.num_heads) != self.hidden_size:
@@ -341,24 +306,38 @@ class Phi3Attention(nn.Module):
 
         op_size = self.num_heads * self.head_dim + 2 * (self.num_key_value_heads * self.head_dim)
         self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
-
         self.qkv_proj = nn.Linear(self.hidden_size, op_size, bias=False)
+        self._init_rope()
 
-        if self.rope_scaling is None:
+    def _init_rope(self):
+        if self.config.rope_scaling is None:
             self.rotary_emb = Phi3RotaryEmbedding(
                 self.head_dim,
                 max_position_embeddings=self.max_position_embeddings,
                 base=self.rope_theta,
             )
         else:
-            self.rotary_emb = Phi3LongScaledRotaryEmbedding(
-                self.head_dim,
-                self.config.rope_scaling["short_factor"],
-                self.config.rope_scaling["long_factor"],
-                max_position_embeddings=self.config.max_position_embeddings,
-                original_max_position_embeddings=self.config.original_max_position_embeddings,
-                base=self.config.rope_theta,
-            )
+            scaling_type = self.config.rope_scaling["type"]
+            if scaling_type == "su":
+                self.rotary_emb = Phi3SuScaledRotaryEmbedding(
+                    self.head_dim,
+                    self.config.rope_scaling["short_factor"],
+                    self.config.rope_scaling["long_factor"],
+                    max_position_embeddings=self.config.max_position_embeddings,
+                    original_max_position_embeddings=self.config.original_max_position_embeddings,
+                    base=self.config.rope_theta,
+                )
+            elif scaling_type == "yarn":
+                self.rotary_emb = Phi3YarnScaledRotaryEmbedding(
+                    self.head_dim,
+                    self.config.rope_scaling["short_factor"],
+                    self.config.rope_scaling["long_factor"],
+                    max_position_embeddings=self.config.max_position_embeddings,
+                    original_max_position_embeddings=self.config.original_max_position_embeddings,
+                    base=self.config.rope_theta,
+                )
+            else:
+                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
 
     def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
         return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
@@ -859,11 +838,11 @@ class Phi3DecoderLayer(nn.Module):
         self.self_attn = PHI3_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)
 
         self.mlp = Phi3MLP(config)
-        self.input_layernorm = PHI3_NORM_CLASS(config.hidden_size, eps=config.rms_norm_eps)
+        self.input_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
         self.resid_attn_dropout = nn.Dropout(config.resid_pdrop)
         self.resid_mlp_dropout = nn.Dropout(config.resid_pdrop)
-        self.post_attention_layernorm = PHI3_NORM_CLASS(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(
         self,
@@ -1066,9 +1045,8 @@ class Phi3Model(Phi3PreTrainedModel):
         self.layers = nn.ModuleList(
             [Phi3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
         )
-        self.norm = PHI3_NORM_CLASS(config.hidden_size, eps=config.rms_norm_eps)
-
         self._attn_implementation = config._attn_implementation
+        self.norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
         self.gradient_checkpointing = False
         # Initialize weights and apply final processing
@@ -1255,6 +1233,7 @@ class Phi3ForCausalLM(Phi3PreTrainedModel):
     def get_decoder(self):
         return self.model
 
+    # Ignore copy
     @add_start_docstrings_to_model_forward(PHI3_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(
@@ -1284,8 +1263,8 @@ class Phi3ForCausalLM(Phi3PreTrainedModel):
         ```python
         >>> from transformers import AutoTokenizer, Phi3ForCausalLM
 
-        >>> model = Phi3ForCausalLM.from_pretrained("microsoft/phi-3")
-        >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-3")
+        >>> model = Phi3ForCausalLM.from_pretrained("microsoft/phi-3-mini-4k-instruct")
+        >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-3-mini-4k-instruct")
 
         >>> prompt = "This is an example script ."
         >>> inputs = tokenizer(prompt, return_tensors="pt")
@@ -1293,7 +1272,7 @@ class Phi3ForCausalLM(Phi3PreTrainedModel):
         >>> # Generate
         >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
         >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-        'This is an example script .\n\n\n\nfrom typing import List\n\ndef find_most_common_letter(words: List[str'
+        'This is an example script .\n Certainly! Below is a sample script that demonstrates a simple task, such as calculating the sum'
         ```"""
 
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions