Skip to content

Commit 0d707a6

Browse files
edenamramedenamram
committed
Refactor rotary embedding options and improve parameter mapping in Gemma4Text
1 parent 288ed5f commit 0d707a6

2 files changed

Lines changed: 73 additions & 65 deletions

File tree

‎lib/bumblebee/layers.ex‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1237,7 +1237,14 @@ defmodule Bumblebee.Layers do
12371237
Adds a rotary embedding layer to the network.
12381238
"""
12391239
def rotary_embedding(query, key, position_ids, attention_mask, size, opts \\ []) do
1240-
opts = Keyword.validate!(opts, [:name, :scaling_strategy, :rotary_dim, max_positions: 2048, base: 10_000])
1240+
opts =
1241+
Keyword.validate!(opts, [
1242+
:name,
1243+
:scaling_strategy,
1244+
:rotary_dim,
1245+
max_positions: 2048,
1246+
base: 10_000
1247+
])
12411248

12421249
output =
12431250
Axon.layer(

‎lib/bumblebee/text/gemma4_text.ex‎

Lines changed: 65 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -452,8 +452,7 @@ defmodule Bumblebee.Text.Gemma4Text do
452452
max_positions: spec.max_positions,
453453
base: spec.rotary_embedding_base,
454454
percentage: 1.0,
455-
rotary_dim:
456-
trunc(spec.global_attention_head_size * spec.partial_rotary_factor)
455+
rotary_dim: trunc(spec.global_attention_head_size * spec.partial_rotary_factor)
457456
]
458457

459458
:sliding_attention ->
@@ -588,14 +587,13 @@ defmodule Bumblebee.Text.Gemma4Text do
588587

589588
shortcut_ple = hidden_state
590589

591-
# Gate: project hidden_state DOWN to PLE dimension
590+
# Gate: project hidden_state DOWN to PLE dimension, then activation
592591
gated =
593592
Axon.dense(hidden_state, spec.hidden_size_per_layer_input,
594593
name: join(name, "per_layer_input_gate"),
595594
use_bias: false
596595
)
597596

598-
# Activation (gelu_approx_tanh, same as FFN)
599597
gated = Layers.activation(gated, spec.activation)
600598

601599
# Element-wise multiply with PLE signal
@@ -622,18 +620,19 @@ defmodule Bumblebee.Text.Gemma4Text do
622620
end
623621

624622
# 4. Layer scalar: multiply output by per-layer learned scalar
625-
layer_scalar =
626-
Axon.param("layer_scalar", fn _ -> {1} end, initializer: :ones)
627-
628623
hidden_state =
629624
Axon.layer(
630625
fn hidden_state, scalar, _opts ->
631-
Nx.multiply(hidden_state, scalar)
626+
Nx.multiply(hidden_state, Nx.reshape(scalar, {}))
632627
end,
633-
[hidden_state, layer_scalar],
634-
name: join(name, "layer_scalar")
628+
[
629+
hidden_state,
630+
Axon.param("layer_scalar", fn _ -> {1} end, initializer: Axon.Initializers.ones())
631+
],
632+
name: join(name, "layer_scalar_op")
635633
)
636634

635+
# Handle cross-attention (required by block interface but not used by Gemma 4)
637636
{_hidden_state, cross_attention_info} =
638637
steps.cross_attention_maybe.(hidden_state, fn _ ->
639638
raise "cross attention not supported"
@@ -756,59 +755,61 @@ defmodule Bumblebee.Text.Gemma4Text do
756755
end
757756

758757
defimpl Bumblebee.HuggingFace.Transformers.Model do
759-
def params_mapping(spec) do
760-
%{
761-
"embedder.token_embedding" => "model.language_model.embed_tokens",
762-
# PLE (Per-Layer Embeddings) global weights
763-
"embedder.token_embedding_per_layer" => "model.language_model.embed_tokens_per_layer",
764-
"per_layer_model_projection" => "model.language_model.per_layer_model_projection",
765-
"per_layer_projection_norm" => "model.language_model.per_layer_projection_norm",
766-
"decoder.blocks.{n}.per_layer_input_gate" =>
767-
"model.language_model.layers.{n}.per_layer_input_gate",
768-
"decoder.blocks.{n}.per_layer_projection" =>
769-
"model.language_model.layers.{n}.per_layer_projection",
770-
"decoder.blocks.{n}.post_per_layer_input_norm" =>
771-
"model.language_model.layers.{n}.post_per_layer_input_norm",
772-
# Per-layer scalar
773-
"decoder.blocks.{n}.layer_scalar" =>
774-
"model.language_model.layers.{n}",
775-
# Attention projections
776-
"decoder.blocks.{n}.self_attention.query" =>
777-
"model.language_model.layers.{n}.self_attn.q_proj",
778-
"decoder.blocks.{n}.self_attention.key" =>
779-
"model.language_model.layers.{n}.self_attn.k_proj",
780-
"decoder.blocks.{n}.self_attention.value" =>
781-
"model.language_model.layers.{n}.self_attn.v_proj",
782-
"decoder.blocks.{n}.self_attention.output" =>
783-
"model.language_model.layers.{n}.self_attn.o_proj",
784-
# QK-norm
785-
"decoder.blocks.{n}.self_attention.query_norm" =>
786-
"model.language_model.layers.{n}.self_attn.q_norm",
787-
"decoder.blocks.{n}.self_attention.key_norm" =>
788-
"model.language_model.layers.{n}.self_attn.k_norm",
789-
# Layer norms
790-
"decoder.blocks.{n}.self_attention_norm" =>
791-
"model.language_model.layers.{n}.input_layernorm",
792-
"decoder.blocks.{n}.post_attention_norm" =>
793-
"model.language_model.layers.{n}.post_attention_layernorm",
794-
# FFN layer norms
795-
"decoder.blocks.{n}.pre_ffn_norm" =>
796-
"model.language_model.layers.{n}.pre_feedforward_layernorm",
797-
"decoder.blocks.{n}.post_ffn_norm" =>
798-
"model.language_model.layers.{n}.post_feedforward_layernorm",
799-
# FFN projections
800-
"decoder.blocks.{n}.ffn.gate" => "model.language_model.layers.{n}.mlp.gate_proj",
801-
"decoder.blocks.{n}.ffn.intermediate" => "model.language_model.layers.{n}.mlp.up_proj",
802-
"decoder.blocks.{n}.ffn.output" => "model.language_model.layers.{n}.mlp.down_proj",
803-
# Output
804-
"output_norm" => "model.language_model.norm",
805-
"language_modeling_head.output" =>
806-
if(spec.tie_word_embeddings,
807-
do: "model.language_model.embed_tokens",
808-
else: "lm_head"
809-
),
810-
"sequence_classification_head.output" => "score"
811-
}
812-
end
758+
def params_mapping(spec) do
759+
%{
760+
"embedder.token_embedding" => "model.language_model.embed_tokens",
761+
# PLE global weights
762+
"embedder.token_embedding_per_layer" => "model.language_model.embed_tokens_per_layer",
763+
"per_layer_model_projection" => "model.language_model.per_layer_model_projection",
764+
"per_layer_projection_norm" => "model.language_model.per_layer_projection_norm",
765+
# PLE per-layer weights
766+
"decoder.blocks.{n}.per_layer_input_gate" =>
767+
"model.language_model.layers.{n}.per_layer_input_gate",
768+
"decoder.blocks.{n}.per_layer_projection" =>
769+
"model.language_model.layers.{n}.per_layer_projection",
770+
"decoder.blocks.{n}.post_per_layer_input_norm" =>
771+
"model.language_model.layers.{n}.post_per_layer_input_norm",
772+
# Layer scalar
773+
"decoder.blocks.{n}.layer_scalar_op" => "model.language_model.layers.{n}",
774+
"decoder.blocks.{n}.layer_scalar_op.layer_scalar" =>
775+
"model.language_model.layers.{n}.layer_scalar",
776+
# Attention projections
777+
"decoder.blocks.{n}.self_attention.query" =>
778+
"model.language_model.layers.{n}.self_attn.q_proj",
779+
"decoder.blocks.{n}.self_attention.key" =>
780+
"model.language_model.layers.{n}.self_attn.k_proj",
781+
"decoder.blocks.{n}.self_attention.value" =>
782+
"model.language_model.layers.{n}.self_attn.v_proj",
783+
"decoder.blocks.{n}.self_attention.output" =>
784+
"model.language_model.layers.{n}.self_attn.o_proj",
785+
# QK-norm
786+
"decoder.blocks.{n}.self_attention.query_norm" =>
787+
"model.language_model.layers.{n}.self_attn.q_norm",
788+
"decoder.blocks.{n}.self_attention.key_norm" =>
789+
"model.language_model.layers.{n}.self_attn.k_norm",
790+
# Layer norms
791+
"decoder.blocks.{n}.self_attention_norm" =>
792+
"model.language_model.layers.{n}.input_layernorm",
793+
"decoder.blocks.{n}.post_attention_norm" =>
794+
"model.language_model.layers.{n}.post_attention_layernorm",
795+
# FFN layer norms
796+
"decoder.blocks.{n}.pre_ffn_norm" =>
797+
"model.language_model.layers.{n}.pre_feedforward_layernorm",
798+
"decoder.blocks.{n}.post_ffn_norm" =>
799+
"model.language_model.layers.{n}.post_feedforward_layernorm",
800+
# FFN projections
801+
"decoder.blocks.{n}.ffn.gate" => "model.language_model.layers.{n}.mlp.gate_proj",
802+
"decoder.blocks.{n}.ffn.intermediate" => "model.language_model.layers.{n}.mlp.up_proj",
803+
"decoder.blocks.{n}.ffn.output" => "model.language_model.layers.{n}.mlp.down_proj",
804+
# Output
805+
"output_norm" => "model.language_model.norm",
806+
"language_modeling_head.output" =>
807+
if(spec.tie_word_embeddings,
808+
do: "model.language_model.embed_tokens",
809+
else: "lm_head"
810+
),
811+
"sequence_classification_head.output" => "score"
812+
}
813+
end
813814
end
814815
end

0 commit comments

Comments
 (0)