From 8b72ff6b74e27f3a52147aaa5cf39ac483731a02 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Sun, 4 Oct 2026 10:19:16 +0200 Subject: [PATCH 1/3] Fix dynamic RoPE scaling formula --- lib/bumblebee/layers.ex | 6 ++--- test/bumblebee/layers_test.exs | 45 ++++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 4 deletions(-) create mode 100644 test/bumblebee/layers_test.exs diff --git a/lib/bumblebee/layers.ex b/lib/bumblebee/layers.ex index 27d81990..f4bfd087 100644 --- a/lib/bumblebee/layers.ex +++ b/lib/bumblebee/layers.ex @@ -1267,10 +1267,8 @@ defmodule Bumblebee.Layers do positions_cos_sin(position, inv_frequency) %{type: :dynamic, factor: factor} when sequence_length > max_positions -> - base = - base - |> Nx.multiply(factor * sequence_length / max_positions - (factor - 1)) - |> Nx.pow(size / (size - 2)) + scale = factor * sequence_length / max_positions - (factor - 1) + base = Nx.multiply(base, Nx.pow(scale, size / (size - 2))) inv_frequency = inv_frequency(base, range) positions_cos_sin(position, inv_frequency) diff --git a/test/bumblebee/layers_test.exs b/test/bumblebee/layers_test.exs new file mode 100644 index 00000000..11d84382 --- /dev/null +++ b/test/bumblebee/layers_test.exs @@ -0,0 +1,45 @@ +defmodule Bumblebee.LayersTest do + use ExUnit.Case, async: true + + import Bumblebee.TestHelpers + + alias Bumblebee.Layers + + describe "rotary_embedding/6" do + test "applies dynamic scaling when sequence length exceeds max_positions" 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}) + + {query, _key} = + Layers.rotary_embedding(query, key, position_ids, attention_mask, 4, + max_positions: 2, + base: 10_000, + scaling_strategy: %{type: :dynamic, factor: 2.0} + ) + + 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) + + scale = 2.0 * 4 / 2 - (2.0 - 1.0) + base = 10_000 * :math.pow(scale, 4 / (4 - 2)) + angle = 1 / :math.sqrt(base) + + assert_all_close( + outputs.query[[0, 1, 0, 1]], + Nx.tensor(:math.cos(angle) - :math.sin(angle)) + ) + end + end +end From 4bb22bb7cb197957f0d1bc82d001e1d3bf0759be Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Sun, 4 Oct 2026 10:21:11 +0200 Subject: [PATCH 2/3] Add YaRN RoPE scaling support --- lib/bumblebee/layers.ex | 46 ++++++++++++++++++++++++++++++ lib/bumblebee/shared.ex | 26 +++++++++++++++++ test/bumblebee/layers_test.exs | 41 ++++++++++++++++++++++++++ test/bumblebee/shared_test.exs | 25 ++++++++++++++++ test/bumblebee/text/llama_test.exs | 37 ++++++++++++++++++++++++ 5 files changed, 175 insertions(+) diff --git a/lib/bumblebee/layers.ex b/lib/bumblebee/layers.ex index f4bfd087..5a1f9fc5 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 + } = yarn -> + beta_fast = yarn[:beta_fast] || 32.0 + beta_slow = yarn[:beta_slow] || 1.0 + attention_factor = yarn[:attention_factor] || 1.0 + + 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 Map.get(yarn, :truncate, true) 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..6cefa76a 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()}` + For more details see https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases """ ] @@ -211,11 +213,35 @@ 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 + {mscale, mscale_all_dim} when is_number(mscale) and is_number(mscale_all_dim) -> + 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 + } + _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 11d84382..ebc65091 100644 --- a/test/bumblebee/layers_test.exs +++ b/test/bumblebee/layers_test.exs @@ -41,5 +41,46 @@ defmodule Bumblebee.LayersTest do Nx.tensor(:math.cos(angle) - :math.sin(angle)) ) 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 + } + + {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.3012, 0.9975], [-1.3254, 0.9950], [-1.1311, 0.9925]]) + |> Nx.multiply(0.1 * :math.log(4.0) + 1.0), + atol: 1.0e-4 + ) + end end end diff --git a/test/bumblebee/shared_test.exs b/test/bumblebee/shared_test.exs index 3467286c..c609e284 100644 --- a/test/bumblebee/shared_test.exs +++ b/test/bumblebee/shared_test.exs @@ -28,4 +28,29 @@ 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 + } + ] + end + end end diff --git a/test/bumblebee/text/llama_test.exs b/test/bumblebee/text/llama_test.exs index fb143c60..c380810f 100644 --- a/test/bumblebee/text/llama_test.exs +++ b/test/bumblebee/text/llama_test.exs @@ -54,6 +54,43 @@ 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 + } + ] + ) + + 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( From 41fde273ef26bc865e5bbd917504f115fccf491a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonatan=20K=C5=82osko?= Date: Mon, 5 Oct 2026 14:11:31 +0200 Subject: [PATCH 3/3] Parse YaRN truncate option and use reference values in tests Assisted-By: Claude Opus 5.5 --- lib/bumblebee/layers.ex | 14 +++++++------- lib/bumblebee/shared.ex | 10 +++++++--- test/bumblebee/layers_test.exs | 7 +++---- test/bumblebee/shared_test.exs | 3 ++- test/bumblebee/text/llama_test.exs | 3 ++- 5 files changed, 21 insertions(+), 16 deletions(-) diff --git a/lib/bumblebee/layers.ex b/lib/bumblebee/layers.ex index 5a1f9fc5..1920c1a8 100644 --- a/lib/bumblebee/layers.ex +++ b/lib/bumblebee/layers.ex @@ -1324,12 +1324,12 @@ defmodule Bumblebee.Layers do %{ type: :yarn, factor: factor, - original_max_positions: original_max_positions - } = yarn -> - beta_fast = yarn[:beta_fast] || 32.0 - beta_slow = yarn[:beta_slow] || 1.0 - attention_factor = yarn[:attention_factor] || 1.0 - + 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) @@ -1337,7 +1337,7 @@ defmodule Bumblebee.Layers do high = yarn_correction_dim(beta_slow, size, base, original_max_positions) {low, high} = - if Map.get(yarn, :truncate, true) do + if truncate do {:math.floor(low), :math.ceil(high)} else {low, high} diff --git a/lib/bumblebee/shared.ex b/lib/bumblebee/shared.ex index 6cefa76a..cec37411 100644 --- a/lib/bumblebee/shared.ex +++ b/lib/bumblebee/shared.ex @@ -42,7 +42,7 @@ 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()}` + * `%{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 """ @@ -217,7 +217,10 @@ defmodule Bumblebee.Shared do attention_factor = params["attention_factor"] || case {params["mscale"], params["mscale_all_dim"]} do - {mscale, mscale_all_dim} when is_number(mscale) and is_number(mscale_all_dim) -> + # 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 -> @@ -230,7 +233,8 @@ defmodule Bumblebee.Shared do original_max_positions: original_max_positions, beta_fast: params["beta_fast"] || 32.0, beta_slow: params["beta_slow"] || 1.0, - attention_factor: attention_factor + attention_factor: attention_factor, + truncate: Map.get(params, "truncate", true) } _other -> diff --git a/test/bumblebee/layers_test.exs b/test/bumblebee/layers_test.exs index ebc65091..13075773 100644 --- a/test/bumblebee/layers_test.exs +++ b/test/bumblebee/layers_test.exs @@ -54,7 +54,8 @@ defmodule Bumblebee.LayersTest do original_max_positions: 2, beta_fast: 32.0, beta_slow: 1.0, - attention_factor: 0.1 * :math.log(4.0) + 1.0 + attention_factor: 0.1 * :math.log(4.0) + 1.0, + truncate: true } {query, _key} = @@ -77,9 +78,7 @@ defmodule Bumblebee.LayersTest do assert_all_close( outputs.query[[0, 1..3, 0, 0..1]], - Nx.tensor([[-0.3012, 0.9975], [-1.3254, 0.9950], [-1.1311, 0.9925]]) - |> Nx.multiply(0.1 * :math.log(4.0) + 1.0), - atol: 1.0e-4 + Nx.tensor([[-0.3429, 1.1358], [-1.5092, 1.1329], [-1.2879, 1.1301]]) ) end end diff --git a/test/bumblebee/shared_test.exs b/test/bumblebee/shared_test.exs index c609e284..38a5f736 100644 --- a/test/bumblebee/shared_test.exs +++ b/test/bumblebee/shared_test.exs @@ -48,7 +48,8 @@ defmodule Bumblebee.SharedTest do original_max_positions: 2048, beta_fast: 32.0, beta_slow: 1.0, - attention_factor: 0.1 * :math.log(4.0) + 1.0 + attention_factor: 0.1 * :math.log(4.0) + 1.0, + truncate: true } ] end diff --git a/test/bumblebee/text/llama_test.exs b/test/bumblebee/text/llama_test.exs index c380810f..7d58e9f2 100644 --- a/test/bumblebee/text/llama_test.exs +++ b/test/bumblebee/text/llama_test.exs @@ -64,7 +64,8 @@ defmodule Bumblebee.Text.LlamaTest do original_max_positions: 16, beta_fast: 32.0, beta_slow: 1.0, - attention_factor: 1.138629436111989 + attention_factor: 1.138629436111989, + truncate: true } ] )