[Bugfix][Model] Gemma4: register aliased embedding scalars as buffers (#54213)
Signed-off-by: Yannick Schnider <Yannick.Schnider1@ibm.com>
Yannick Schnider committed
Sep 28, 2026 at 12:56 UTC
8cc9aa5ad350f6f80cf3bebb298d410803ff75a9
1 file changed
+13
-4
vllm/model_executor/models/gemma4.py
+13
-4
@@ -846,13 +846,22 @@ class Gemma4SelfDecoderLayers(nn.Module):
846
# Shared references to modules owned by Gemma4Model — must be
847
# inside this nn.Module so torch.compile captures them.
848
self.embed_tokens = embed_tokens
849
- self.normalizer = normalizer
849
self.embed_tokens_per_layer = embed_tokens_per_layer
851
- self.embed_scale_per_layer = embed_scale_per_layer
850
self.per_layer_model_projection = per_layer_model_projection
851
self.per_layer_projection_norm = per_layer_projection_norm
854
- self.per_layer_input_scale = per_layer_input_scale
855
- self.per_layer_projection_scale = per_layer_projection_scale
852
+
853
+ # Register the shared scalars as buffers so they follow the model on
854
+ # .to(device); plain attributes stay on CPU.
855
+ self.register_buffer("normalizer", normalizer, persistent=False)
856
+ self.register_buffer(
857
+ "embed_scale_per_layer", embed_scale_per_layer, persistent=False
858
+ )
859
+ self.register_buffer(
860
+ "per_layer_input_scale", per_layer_input_scale, persistent=False
861
+ )
862
+ self.register_buffer(
863
+ "per_layer_projection_scale", per_layer_projection_scale, persistent=False
864
+ )
865
866
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
867
return self.embed_tokens(input_ids) * self.normalizer