diff --git a/lib/bumblebee/layers.ex b/lib/bumblebee/layers.ex index f4bfd087..1920c1a8 100644 --- a/lib/bumblebee/layers.ex +++ b/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, diff --git a/lib/bumblebee/shared.ex b/lib/bumblebee/shared.ex index e4ff6c55..cec37411 100644 --- a/lib/bumblebee/shared.ex +++ b/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. diff --git a/test/bumblebee/layers_test.exs b/test/bumblebee/layers_test.exs index 9a07de39..008ec9d5 100644 --- a/test/bumblebee/layers_test.exs +++ b/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 diff --git a/test/bumblebee/shared_test.exs b/test/bumblebee/shared_test.exs index 3467286c..38a5f736 100644 --- a/test/bumblebee/shared_test.exs +++ b/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 diff --git a/test/bumblebee/text/llama_test.exs b/test/bumblebee/text/llama_test.exs index fb143c60..7d58e9f2 100644 --- a/test/bumblebee/text/llama_test.exs +++ b/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(