[Bugfix] Gemma4: fix multimodal embedder norm order to match HF reference (#40411)

Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
This commit is contained in:
Luciano Martins
2026-04-21 02:28:26 +00:00
committed by GitHub
co-authored by Luciano Martins
parent 18563f2072
commit 20d3743491
+9 -8
View File
@@ -849,22 +849,23 @@ class Gemma4MultimodalEmbedder(nn.Module):
or multimodal_config.hidden_size
)
self.embedding_pre_projection_norm = RMSNorm(
embedding_dim,
eps=self.eps,
has_weight=False,
)
self.embedding_projection = ReplicatedLinear(
embedding_dim,
self.text_hidden_size,
bias=False,
)
self.embedding_post_projection_norm = RMSNorm(
self.text_hidden_size,
eps=self.eps,
has_weight=False,
)
def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
"""Project soft tokens from a multimodal tower into LM space."""
embs_proj, _ = self.embedding_projection(inputs_embeds)
return self.embedding_post_projection_norm(embs_proj)
embs_normed = self.embedding_pre_projection_norm(inputs_embeds)
embs_proj, _ = self.embedding_projection(embs_normed)
return embs_proj
# ---------------------------------------------------------------------------