From 553deca6f64fbd9f33bcb8517c0fb3f3493d5153 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Fri, 2 Oct 2026 22:19:06 +0200 Subject: [PATCH] Add YaRN RoPE scaling and correct dynamic scaling --- lib/bumblebee/layers.ex | 239 +++++++++++++++++++++++++++++++-- lib/bumblebee/shared.ex | 183 ++++++++++++++++++++++--- test/bumblebee/layers_test.exs | 156 +++++++++++++++++++++ test/bumblebee/shared_test.exs | 93 +++++++++++++ 4 files changed, 645 insertions(+), 26 deletions(-) create mode 100644 test/bumblebee/layers_test.exs diff --git a/lib/bumblebee/layers.ex b/lib/bumblebee/layers.ex index 27d81990..3bd830ae 100644 --- a/lib/bumblebee/layers.ex +++ b/lib/bumblebee/layers.ex @@ -1239,18 +1239,206 @@ defmodule Bumblebee.Layers do def rotary_embedding(query, key, position_ids, attention_mask, size, opts \\ []) do opts = Keyword.validate!(opts, [:name, :scaling_strategy, max_positions: 2048, base: 10_000]) + scaling_strategy = + normalize_rotary_embedding_scaling_strategy(opts[:scaling_strategy], size, opts) + output = Axon.layer( &apply_rotary_embedding/5, [query, key, position_ids, Axon.optional(attention_mask)], - [size: size] ++ opts + [size: size, scaling_strategy: scaling_strategy] ++ + Keyword.drop(opts, [:scaling_strategy]) ) unwrap_tuple(output, 2) end + defp normalize_rotary_embedding_scaling_strategy(nil, _size, _opts), do: nil + + defp normalize_rotary_embedding_scaling_strategy( + %{type: :yarn} = strategy, + size, + opts + ) do + required_keys = [ + :factor, + :original_max_positions, + :beta_fast, + :beta_slow + ] + + case Enum.find(required_keys, &(not Map.has_key?(strategy, &1))) do + nil -> + validate_factor!(strategy.factor, :yarn) + + strategy = + strategy + |> Map.put_new_lazy(:attention_factor, fn -> 0.1 * :math.log(strategy.factor) + 1.0 end) + |> Map.put_new(:truncate, true) + + validate_yarn_strategy!(strategy) + + low = + yarn_correction_dim( + strategy.beta_fast, + size, + opts[:base], + strategy.original_max_positions + ) + + high = + yarn_correction_dim( + strategy.beta_slow, + size, + opts[:base], + strategy.original_max_positions + ) + + {low, high} = + if strategy.truncate do + {:math.floor(low), :math.ceil(high)} + else + {low, high} + end + + low = max(low, 0) + high = min(high, size - 1) + high = if low == high, do: high + 0.001, else: high + + Map.merge(strategy, %{ + correction_low: low, + correction_high: high + }) + + key -> + raise ArgumentError, "YaRN rotary embedding scaling strategy is missing #{inspect(key)}" + end + end + + defp normalize_rotary_embedding_scaling_strategy( + %{type: :longrope, short_factor: short_factor, long_factor: long_factor} = strategy, + size, + _opts + ) do + expected_size = div(size, 2) + + if longrope_factors?(short_factor, expected_size) and + longrope_factors?(long_factor, expected_size) do + strategy + else + raise ArgumentError, + "LongRoPE factors must each have #{expected_size} entries for a rotary size of #{size}" + end + end + + defp normalize_rotary_embedding_scaling_strategy( + %{type: :dynamic, factor: factor} = strategy, + size, + _opts + ) do + validate_factor!(factor, :dynamic) + + if size <= 2 do + raise ArgumentError, "dynamic RoPE requires a rotary size greater than 2" + end + + strategy + end + + defp normalize_rotary_embedding_scaling_strategy( + %{type: :linear, factor: factor} = strategy, + _size, + _opts + ) do + validate_factor!(factor, :linear) + strategy + end + + defp normalize_rotary_embedding_scaling_strategy(%{type: type}, _size, _opts) + when type in [:linear, :dynamic] do + raise ArgumentError, "#{type} RoPE requires :factor to be a number >= 1" + end + + defp normalize_rotary_embedding_scaling_strategy( + %{ + type: :llama3, + factor: factor, + low_frequency_factor: low, + high_frequency_factor: high, + original_max_positions: original_max_positions + } = strategy, + _size, + _opts + ) do + validate_factor!(factor, :llama3) + validate_positive_number!(low, :low_frequency_factor, :llama3) + validate_positive_number!(high, :high_frequency_factor, :llama3) + validate_positive_integer!(original_max_positions, :original_max_positions, :llama3) + + if high <= low do + raise ArgumentError, "llama3 RoPE requires :high_frequency_factor > :low_frequency_factor" + end + + strategy + end + + defp normalize_rotary_embedding_scaling_strategy(strategy, _size, _opts) do + raise ArgumentError, "unsupported rotary embedding scaling strategy: #{inspect(strategy)}" + 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 + + defp validate_yarn_strategy!(strategy) do + validate_factor!(strategy.factor, :yarn) + validate_positive_integer!(strategy.original_max_positions, :original_max_positions, :yarn) + validate_positive_number!(strategy.beta_fast, :beta_fast, :yarn) + validate_positive_number!(strategy.beta_slow, :beta_slow, :yarn) + validate_positive_number!(strategy.attention_factor, :attention_factor, :yarn) + + unless is_boolean(strategy.truncate) do + raise ArgumentError, "YaRN RoPE requires :truncate to be a boolean" + end + + if strategy.beta_fast < strategy.beta_slow do + raise ArgumentError, "YaRN RoPE requires :beta_fast >= :beta_slow" + end + end + + defp longrope_factors?(factors, expected_size) do + is_list(factors) and length(factors) == expected_size and + Enum.all?(factors, &(is_number(&1) and &1 > 0)) + end + + defp validate_factor!(factor, type) do + if is_number(factor) and factor >= 1.0 do + :ok + else + raise ArgumentError, "#{type} RoPE requires :factor to be a number >= 1" + end + end + + defp validate_positive_number!(value, key, type) do + if is_number(value) and value > 0 do + :ok + else + raise ArgumentError, "#{type} RoPE requires #{inspect(key)} to be a positive number" + end + end + + defp validate_positive_integer!(value, key, type) do + if is_integer(value) and value > 0 do + :ok + else + raise ArgumentError, "#{type} RoPE requires #{inspect(key)} to be a positive integer" + end + end + deftransformp create_sinusoidal_positions( sequence_length, + context_length, max_positions, size, base, @@ -1266,15 +1454,48 @@ defmodule Bumblebee.Layers do position = Nx.divide(position, factor) 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)) + %{type: :dynamic, factor: factor} -> + context_length = Nx.max(context_length, max_positions) + + scale = + context_length + |> Nx.multiply(factor) + |> Nx.divide(max_positions) + |> Nx.subtract(factor - 1) + + base = Nx.multiply(base, Nx.pow(scale, size / (size - 2))) inv_frequency = inv_frequency(base, range) positions_cos_sin(position, inv_frequency) + %{ + type: :yarn, + factor: factor, + attention_factor: attention_factor, + correction_low: low, + correction_high: high + } -> + inv_frequency_extrapolation = inv_frequency(base, range) + inv_frequency_interpolation = Nx.divide(inv_frequency_extrapolation, factor) + + ramp = + Nx.iota({div(size, 2)}) + |> Nx.subtract(low) + |> Nx.divide(high - low) + |> Nx.clip(0.0, 1.0) + + inv_frequency_extrapolation_factor = Nx.subtract(1.0, ramp) + + inv_frequency = + Nx.multiply( + inv_frequency_interpolation, + Nx.subtract(1.0, inv_frequency_extrapolation_factor) + ) + |> Nx.add(Nx.multiply(inv_frequency_extrapolation, inv_frequency_extrapolation_factor)) + + {cos, sin} = positions_cos_sin(position, inv_frequency) + {Nx.multiply(cos, attention_factor), Nx.multiply(sin, attention_factor)} + %{ type: :longrope, short_factor: short_factor, @@ -1394,17 +1615,19 @@ defmodule Bumblebee.Layers do _other -> Nx.axis_size(attention_mask, 1) end + position_ids = Nx.as_type(position_ids, :s64) + context_length = Nx.add(Nx.reduce_max(position_ids), 1) + {cos, sin} = create_sinusoidal_positions( sequence_length, + context_length, opts[:max_positions], opts[:size], opts[:base], opts[:scaling_strategy] ) - position_ids = Nx.as_type(position_ids, :s64) - cos = cos |> Nx.take(position_ids) |> Nx.new_axis(2) |> Nx.as_type(Nx.type(query)) sin = sin |> Nx.take(position_ids) |> Nx.new_axis(2) |> Nx.as_type(Nx.type(query)) diff --git a/lib/bumblebee/shared.ex b/lib/bumblebee/shared.ex index e4ff6c55..06fb66ff 100644 --- a/lib/bumblebee/shared.ex +++ b/lib/bumblebee/shared.ex @@ -38,10 +38,15 @@ defmodule Bumblebee.Shared do * `%{type: :dynamic, factor: number()}` + * `%{type: :yarn, factor: number(), original_max_positions: pos_integer(), beta_fast: number(), beta_slow: number()}` + * `%{type: :llama3, factor: number(), low_frequency_factor: number(), high_frequency_factor: number(), original_max_positions: pos_integer()}` * `%{type: :longrope, short_factor: list(number()), long_factor: list(number()), original_max_positions: pos_integer()}` + YaRN also accepts `:attention_factor` (defaults to `0.1 * log(factor) + 1.0`) + and `:truncate` (defaults to `true`). + For more details see https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases """ ] @@ -168,8 +173,14 @@ defmodule Bumblebee.Shared do end defp rotary_embedding_options(params, data) do + base = params["rope_theta"] + + if base != nil and (not is_number(base) or base <= 0) do + raise "conversion failed, \"rope_theta\" must be a positive number" + end + [ - rotary_embedding_base: params["rope_theta"], + rotary_embedding_base: base, rotary_embedding_percentage: params["partial_rotary_factor"], rotary_embedding_scaling_strategy: rotary_embedding_scaling_strategy(params, data) ] @@ -182,40 +193,176 @@ defmodule Bumblebee.Shared do data["original_max_position_embeddings"] || params["original_max_position_embeddings"] || data["max_position_embeddings"] - case {params["rope_type"] || params["type"] || "default", params} do + type = rope_type!(params) + + if type == "default" and scaling_fields?(params) do + if Map.has_key?(params, "type") or Map.has_key?(params, "rope_type") do + raise "conversion failed, default rotary embedding does not accept scaling parameters" + else + raise "conversion failed, rotary embedding scaling parameters require \"type\" or \"rope_type\"" + end + end + + case {type, params} do {"default", _params} -> nil {"linear", %{"factor" => factor}} -> - %{type: :linear, factor: factor} + %{type: :linear, factor: validate_factor!(factor, "linear")} {"dynamic", %{"factor" => factor}} -> - %{type: :dynamic, factor: factor} + %{type: :dynamic, factor: validate_factor!(factor, type)} {"llama3", %{"factor" => factor, "low_freq_factor" => low, "high_freq_factor" => high}} -> - %{ - type: :llama3, - factor: factor, - low_frequency_factor: low, - high_frequency_factor: high, - original_max_positions: original_max_positions - } + llama3_scaling_strategy(factor, low, high, original_max_positions) # Old Phi-3 checkpoints use "su" and "yarn" for LongRoPE - {type, %{"short_factor" => short_factor, "long_factor" => long_factor}} + {type, %{"short_factor" => _short_factor, "long_factor" => _long_factor} = params} when type in ["longrope", "su", "yarn"] -> - %{ - type: :longrope, - short_factor: short_factor, - long_factor: long_factor, - original_max_positions: original_max_positions - } + longrope_scaling_strategy(params, original_max_positions) + + {"yarn", %{"factor" => _factor}} = config -> + yarn_scaling_strategy(config, original_max_positions) _other -> raise "conversion failed, unsupported rotary embedding parameters: #{inspect(params)}" end end + defp yarn_scaling_strategy({"yarn", params}, original_max_positions) do + factor = validate_factor!(params["factor"], "yarn") + original_max_positions = validate_original_max_positions!(original_max_positions, "yarn") + beta_fast = validate_positive_number!(params["beta_fast"] || 32.0, "beta_fast", "yarn") + beta_slow = validate_positive_number!(params["beta_slow"] || 1.0, "beta_slow", "yarn") + + unless is_boolean(Map.get(params, "truncate", true)) do + raise "conversion failed, yarn requires \"truncate\" to be a boolean" + end + + if beta_fast < beta_slow do + raise "conversion failed, YaRN requires \"beta_fast\" >= \"beta_slow\"" + end + + %{ + type: :yarn, + factor: factor, + original_max_positions: original_max_positions, + beta_fast: beta_fast, + beta_slow: beta_slow, + attention_factor: yarn_attention_factor(params, factor), + truncate: Map.get(params, "truncate", true) + } + end + + defp llama3_scaling_strategy(factor, low, high, original_max_positions) do + factor = validate_factor!(factor, "llama3") + low = validate_positive_number!(low, "low_freq_factor", "llama3") + high = validate_positive_number!(high, "high_freq_factor", "llama3") + + if high <= low do + raise "conversion failed, llama3 requires \"high_freq_factor\" > \"low_freq_factor\"" + end + + %{ + type: :llama3, + factor: factor, + low_frequency_factor: low, + high_frequency_factor: high, + original_max_positions: validate_original_max_positions!(original_max_positions, "llama3") + } + end + + defp rope_type!(%{"rope_type" => rope_type, "type" => type}) when rope_type != type do + raise "conversion failed, conflicting \"rope_type\" and \"type\" values" + end + + defp rope_type!(%{"rope_type" => rope_type}), do: rope_type + defp rope_type!(%{"type" => type}), do: type + defp rope_type!(_params), do: "default" + + defp scaling_fields?(params) do + Enum.any?( + [ + "factor", + "short_factor", + "long_factor", + "low_freq_factor", + "high_freq_factor", + "beta_fast", + "beta_slow", + "attention_factor", + "mscale", + "mscale_all_dim", + "truncate" + ], + &Map.has_key?(params, &1) + ) + end + + defp longrope_scaling_strategy(params, original_max_positions) do + %{ + type: :longrope, + short_factor: validate_number_list!(params["short_factor"], "short_factor", "longrope"), + long_factor: validate_number_list!(params["long_factor"], "long_factor", "longrope"), + original_max_positions: validate_original_max_positions!(original_max_positions, "longrope") + } + end + + defp yarn_attention_factor(params, factor) do + case params["attention_factor"] do + nil -> + case {params["mscale"], params["mscale_all_dim"]} do + {nil, nil} -> + yarn_mscale(factor) + + {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 -> + raise "conversion failed, yarn requires numeric \"mscale\" and \"mscale_all_dim\" together" + end + + attention_factor -> + validate_positive_number!(attention_factor, "attention_factor", "yarn") + 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 + + defp validate_factor!(factor, type) do + if is_number(factor) and factor >= 1.0 do + factor + else + raise "conversion failed, #{type} requires a numeric \"factor\" >= 1" + end + end + + defp validate_original_max_positions!(value, type) do + if is_integer(value) and value > 0 do + value + else + raise "conversion failed, #{type} requires a positive integer \"original_max_position_embeddings\"" + end + end + + defp validate_positive_number!(value, key, type) do + if is_number(value) and value > 0 do + value + else + raise "conversion failed, #{type} requires a positive numeric #{inspect(key)}" + end + end + + defp validate_number_list!(value, key, type) do + if is_list(value) and value != [] and Enum.all?(value, &(is_number(&1) and &1 > 0)) do + value + else + raise "conversion failed, #{type} requires a non-empty list of positive numbers for #{inspect(key)}" + end + end + @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 new file mode 100644 index 00000000..688ee7bd --- /dev/null +++ b/test/bumblebee/layers_test.exs @@ -0,0 +1,156 @@ +defmodule Bumblebee.LayersTest do + use ExUnit.Case, async: true + + import Bumblebee.TestHelpers + + alias Bumblebee.Layers + + test "documented YaRN configuration uses the default attention factor and truncation" do + strategy = %{ + type: :yarn, + factor: 4.0, + original_max_positions: 2, + beta_fast: 32.0, + beta_slow: 1.0 + } + + {query, _key} = apply_rotary_embedding(strategy) + {init, predict} = Axon.build(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]]) + } + + params = init.(inputs, Axon.ModelState.empty()) + output = predict.(params, inputs) + + assert_all_close( + output[[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) + ) + end + + test "rejects an explicit unknown rotary scaling strategy" do + assert_raise ArgumentError, ~r/unsupported rotary embedding scaling strategy/, fn -> + apply_rotary_embedding(%{type: :unsupported}) + end + end + + test "rejects an explicit malformed linear scaling strategy" do + assert_raise ArgumentError, ~r/linear RoPE requires :factor to be a number >= 1/, fn -> + apply_rotary_embedding(%{type: :linear}) + end + end + + test "rejects non-positive LongRoPE factors" do + assert_raise ArgumentError, ~r/LongRoPE factors must each have 2 entries/, fn -> + apply_rotary_embedding(%{ + type: :longrope, + short_factor: [1.0, 0.0], + long_factor: [1.0, 1.0], + original_max_positions: 2 + }) + end + end + + test "dynamic RoPE exponent applies to the multiplier, using position IDs" 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, 8}) + + {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, 0, 0, 0, 0]]) + } + + {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 + + test "rejects dynamic RoPE with a rotary size of two" do + assert_raise ArgumentError, ~r/requires a rotary size greater than 2/, fn -> + query = Axon.input("query", shape: {1, 4, 1, 2}) + key = Axon.input("key", shape: {1, 4, 1, 2}) + position_ids = Axon.input("position_ids", shape: {1, 4}) + attention_mask = Axon.input("attention_mask", shape: {1, 4}) + + Layers.rotary_embedding(query, key, position_ids, attention_mask, 2, + scaling_strategy: %{type: :dynamic, factor: 2.0} + ) + end + end + + test "YaRN scales rotary frequencies at positions beyond the original context" 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: 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.3012, 0.9975], [-1.3254, 0.9950], [-1.1311, 0.9925]]) + ) + end + + defp apply_rotary_embedding(strategy) 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}) + + Layers.rotary_embedding(query, key, position_ids, attention_mask, 4, + scaling_strategy: strategy + ) + end +end diff --git a/test/bumblebee/shared_test.exs b/test/bumblebee/shared_test.exs index 3467286c..6d54ff0b 100644 --- a/test/bumblebee/shared_test.exs +++ b/test/bumblebee/shared_test.exs @@ -3,6 +3,99 @@ defmodule Bumblebee.SharedTest do alias Bumblebee.Shared + describe "rotary_embedding_options_from_transformers/1" do + test "accepts the legacy type key" do + assert [rotary_embedding_scaling_strategy: %{type: :linear, factor: 2.0}] = + Shared.rotary_embedding_options_from_transformers(%{ + "rope_scaling" => %{"type" => "linear", "factor" => 2.0} + }) + end + + test "accepts the rope_type key for YaRN" do + assert [ + rotary_embedding_scaling_strategy: %{ + type: :yarn, + factor: 4.0, + original_max_positions: 128, + beta_fast: 32.0, + beta_slow: 1.0 + } + ] = + Shared.rotary_embedding_options_from_transformers(%{ + "max_position_embeddings" => 512, + "rope_parameters" => %{ + "rope_type" => "yarn", + "factor" => 4.0, + "original_max_position_embeddings" => 128, + "beta_fast" => 32.0, + "beta_slow" => 1.0 + } + }) + end + + test "raises for an unsupported RoPE type" do + assert_raise RuntimeError, ~r/unsupported rotary embedding parameters/, fn -> + Shared.rotary_embedding_options_from_transformers(%{ + "rope_scaling" => %{"type" => "unsupported", "factor" => 2.0} + }) + end + end + + test "raises for conflicting type aliases" do + assert_raise RuntimeError, ~r/conflicting "rope_type" and "type" values/, fn -> + Shared.rotary_embedding_options_from_transformers(%{ + "rope_scaling" => %{"rope_type" => "linear", "type" => "dynamic", "factor" => 2.0} + }) + end + end + + test "raises for invalid YaRN parameters" do + assert_raise RuntimeError, ~r/requires a numeric "factor" >= 1/, fn -> + Shared.rotary_embedding_options_from_transformers(%{ + "rope_scaling" => %{ + "type" => "yarn", + "factor" => 0.5, + "original_max_position_embeddings" => 128 + } + }) + end + end + + test "rejects scaling parameters without a type" do + assert_raise RuntimeError, ~r/scaling parameters require "type" or "rope_type"/, fn -> + Shared.rotary_embedding_options_from_transformers(%{ + "rope_scaling" => %{"factor" => 2.0} + }) + end + end + + test "rejects a non-positive rope theta" do + assert_raise RuntimeError, ~r/"rope_theta" must be a positive number/, fn -> + Shared.rotary_embedding_options_from_transformers(%{"rope_theta" => 0}) + end + end + + test "preserves legacy Phi-3 YaRN LongRoPE without a factor" do + assert [ + rotary_embedding_scaling_strategy: %{ + type: :longrope, + short_factor: [1.0, 1.0], + long_factor: [2.0, 2.0], + original_max_positions: 128 + } + ] = + Shared.rotary_embedding_options_from_transformers(%{ + "max_position_embeddings" => 512, + "rope_scaling" => %{ + "type" => "yarn", + "short_factor" => [1.0, 1.0], + "long_factor" => [2.0, 2.0], + "original_max_position_embeddings" => 128 + } + }) + end + end + describe "validate_label_options/1" do test "passes when :id_to_label is empty" do spec = %{__struct__: TestConfig, num_labels: 3, id_to_label: %{}}