Add support for YaRN scaling

lib/bumblebee/layers.ex

@@ -1321,12 +1321,58 @@ defmodule Bumblebee.Layers do
 
         positions_cos_sin(position, inv_frequency)
 
+      %{
+        type: :yarn,
+        factor: factor,
+        original_max_positions: original_max_positions,
+        beta_fast: beta_fast,
+        beta_slow: beta_slow,
+        attention_factor: attention_factor,
+        truncate: truncate
+      } ->
+        inv_frequency_extrapolation = inv_frequency(base, range)
+        inv_frequency_interpolation = Nx.divide(inv_frequency_extrapolation, factor)
+
+        low = yarn_correction_dim(beta_fast, size, base, original_max_positions)
+        high = yarn_correction_dim(beta_slow, size, base, original_max_positions)
+
+        {low, high} =
+          if truncate do
+            {:math.floor(low), :math.ceil(high)}
+          else
+            {low, high}
+          end
+
+        low = max(low, 0.0)
+        high = min(high, size - 1.0)
+        high = if low == high, do: high + 0.001, else: high
+
+        ramp =
+          Nx.iota({div(size, 2)})
+          |> Nx.subtract(low)
+          |> Nx.divide(high - low)
+          |> Nx.clip(0.0, 1.0)
+
+        inv_frequency =
+          Nx.add(
+            Nx.multiply(inv_frequency_interpolation, ramp),
+            Nx.multiply(inv_frequency_extrapolation, Nx.subtract(1.0, ramp))
+          )
+
+        {cos, sin} = positions_cos_sin(position, inv_frequency)
+        {Nx.multiply(cos, attention_factor), Nx.multiply(sin, attention_factor)}
+
       _other ->
         inv_frequency = inv_frequency(base, range)
         positions_cos_sin(position, inv_frequency)
     end
   end
 
+  defp yarn_correction_dim(num_rotations, size, base, original_max_positions) do
+    size * :math.log(original_max_positions / (num_rotations * 2 * :math.pi())) /
+      (2 * :math.log(base))
+  end
+
   defnp llama3_inv_frequency(
           inv_frequency,
           factor,

lib/bumblebee/shared.ex

@@ -42,6 +42,8 @@ defmodule Bumblebee.Shared do
 
           * `%{type: :longrope, short_factor: list(number()), long_factor: list(number()), original_max_positions: pos_integer()}`
 
+          * `%{type: :yarn, factor: number(), original_max_positions: pos_integer(), beta_fast: number(), beta_slow: number(), attention_factor: number(), truncate: boolean()}`
+
         For more details see https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases
         """
       ]
@@ -211,11 +213,39 @@ defmodule Bumblebee.Shared do
           original_max_positions: original_max_positions
         }
 
+      {"yarn", %{"factor" => factor} = params} ->
+        attention_factor =
+          params["attention_factor"] ||
+            case {params["mscale"], params["mscale_all_dim"]} do
+              # Zero means unset, same as in hf/transformers
+              {mscale, mscale_all_dim}
+              when is_number(mscale) and mscale != 0 and is_number(mscale_all_dim) and
+                     mscale_all_dim != 0 ->
+                yarn_mscale(factor, mscale) / yarn_mscale(factor, mscale_all_dim)
+
+              _other ->
+                yarn_mscale(factor)
+            end
+
+        %{
+          type: :yarn,
+          factor: factor,
+          original_max_positions: original_max_positions,
+          beta_fast: params["beta_fast"] || 32.0,
+          beta_slow: params["beta_slow"] || 1.0,
+          attention_factor: attention_factor,
+          truncate: Map.get(params, "truncate", true)
+        }
+
       _other ->
         raise "conversion failed, unsupported rotary embedding parameters: #{inspect(params)}"
     end
   end
 
+  defp yarn_mscale(factor, mscale \\ 1.0)
+  defp yarn_mscale(factor, _mscale) when factor <= 1.0, do: 1.0
+  defp yarn_mscale(factor, mscale), do: 0.1 * mscale * :math.log(factor) + 1.0
+
   @doc """
   Merges the given list of attributes into a configuration struct.
 

test/bumblebee/layers_test.exs

@@ -34,5 +34,45 @@ defmodule Bumblebee.LayersTest do
 
       assert_all_close(outputs.query[[0, 1, 0, 1]], Nx.tensor(0.9967))
     end
+
+    test "applies YaRN scaling" do
+      query = Axon.input("query", shape: {1, 4, 1, 4})
+      key = Axon.input("key", shape: {1, 4, 1, 4})
+      position_ids = Axon.input("position_ids", shape: {1, 4})
+      attention_mask = Axon.input("attention_mask", shape: {1, 4})
+
+      strategy = %{
+        type: :yarn,
+        factor: 4.0,
+        original_max_positions: 2,
+        beta_fast: 32.0,
+        beta_slow: 1.0,
+        attention_factor: 0.1 * :math.log(4.0) + 1.0,
+        truncate: true
+      }
+
+      {query, _key} =
+        Layers.rotary_embedding(query, key, position_ids, attention_mask, 4,
+          scaling_strategy: strategy
+        )
+
+      model = Axon.container(%{query: query})
+
+      inputs = %{
+        "query" => Nx.broadcast(1.0, {1, 4, 1, 4}),
+        "key" => Nx.broadcast(1.0, {1, 4, 1, 4}),
+        "position_ids" => Nx.tensor([[0, 1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 1, 1]])
+      }
+
+      {init, predict} = Axon.build(model)
+      params = init.(inputs, Axon.ModelState.empty())
+      outputs = predict.(params, inputs)
+
+      assert_all_close(
+        outputs.query[[0, 1..3, 0, 0..1]],
+        Nx.tensor([[-0.3429, 1.1358], [-1.5092, 1.1329], [-1.2879, 1.1301]])
+      )
+    end
   end
 end

test/bumblebee/shared_test.exs

@@ -28,4 +28,30 @@ defmodule Bumblebee.SharedTest do
                    end
     end
   end
+
+  describe "rotary_embedding_options_from_transformers/1" do
+    test "loads YaRN configuration" do
+      data = %{
+        "rope_theta" => 10_000.0,
+        "rope_scaling" => %{
+          "rope_type" => "yarn",
+          "factor" => 4.0,
+          "original_max_position_embeddings" => 2048
+        }
+      }
+
+      assert Shared.rotary_embedding_options_from_transformers(data) == [
+               rotary_embedding_base: 10_000.0,
+               rotary_embedding_scaling_strategy: %{
+                 type: :yarn,
+                 factor: 4.0,
+                 original_max_positions: 2048,
+                 beta_fast: 32.0,
+                 beta_slow: 1.0,
+                 attention_factor: 0.1 * :math.log(4.0) + 1.0,
+                 truncate: true
+               }
+             ]
+    end
+  end
 end

test/bumblebee/text/llama_test.exs

@@ -54,6 +54,44 @@ defmodule Bumblebee.Text.LlamaTest do
     )
   end
 
+  test ":base rotary embedding scaling strategy :yarn" do
+    assert {:ok, %{model: model, params: params, spec: spec}} =
+             Bumblebee.load_model({:hf, "bumblebee-testing/tiny-random-LlamaModel"},
+               spec_overrides: [
+                 rotary_embedding_scaling_strategy: %{
+                   type: :yarn,
+                   factor: 4.0,
+                   original_max_positions: 16,
+                   beta_fast: 32.0,
+                   beta_slow: 1.0,
+                   attention_factor: 1.138629436111989,
+                   truncate: true
+                 }
+               ]
+             )
+
+    assert %Bumblebee.Text.Llama{
+             architecture: :base,
+             rotary_embedding_scaling_strategy: %{type: :yarn}
+           } = spec
+
+    inputs = %{
+      "input_ids" => Nx.tensor([[10, 20, 30, 40, 50, 60, 70, 80, 0, 0]]),
+      "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
+    }
+
+    outputs = Axon.predict(model, params, inputs)
+
+    assert Nx.shape(outputs.hidden_state) == {1, 10, 32}
+
+    assert_all_close(
+      outputs.hidden_state[[.., 1..3, 1..3]],
+      Nx.tensor([
+        [[1.4805, -2.0318, 0.4768], [2.3732, -0.8378, -0.0220], [0.5740, -0.0519, -1.1806]]
+      ])
+    )
+  end
+
   test ":for_sequence_classification" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model(