[Gemma4/Gemma3n] MultimodalEmbedder.forward missing dtype alignment that subclass override already applies
Author: BiggieW
Pull request: no

### Summary

`Gemma4MultimodalEmbedder.forward` and `Gemma3nMultimodalEmbedder.forward` (on the `inputs_embeds` branch) feed activations directly into an `nn.Linear` projection without aligning the activation dtype with `embedding_projection.weight.dtype`. This is the only place in either model's vision/audio→LM bridge where this defensive cast is missing — a sibling subclass already applies it, and the same convention is used elsewhere in the same model.

### Existing precedent (already in this codebase)

The strongest precedent is the same model's own subclass override:

1. **`Gemma4UnifiedMultimodalEmbedder.forward`** (`modeling_gemma4_unified.py:856`)
   — subclasses `Gemma4MultimodalEmbedder` and **only overrides `forward`** to add
   `inputs_embeds = inputs_embeds.to(self.embedding_projection.weight.dtype)`,
   with the comment `# Additional dtype casting`. The override exists solely
   because the parent class doesn't do this. Moving the cast into the parent
   removes the duplication.

2. **`Gemma4VisionPatchEmbedder.forward`** (`modeling_gemma4.py:613`)
   — `pixel_values.to(self.input_proj.weight.dtype)`. Different position in the
   pipeline (vision input rather than projection bridge), but the same
   `.to(weight.dtype)` convention applied within the same model.

### Proposed change

Move the cast into the parent classes so the four sibling layers follow the same convention:

- `Gemma4MultimodalEmbedder.forward` — add `inputs_embeds = inputs_embeds.to(self.embedding_projection.weight.dtype)` before `embedding_pre_projection_norm`.
- `Gemma3nMultimodalEmbedder.forward` — same cast on the `inputs_embeds` branch (token-ids branch unaffected).
- `Gemma4UnifiedMultimodalEmbedder` — drop the now-redundant `forward` override; behaviour preserved by inheritance.

When the dtypes already match, the cast is a no-op.

### Why this matters

`Gemma{3n,4}RMSNorm.forward` returns `normed_output.type_as(hidden_states)`, i.e. the norm preserves the **input** dtype rather than the weight dtype. So if upstream activations and the projection weight end up in different dtypes at this point in the graph, `embedding_projection` (an `nn.Linear`) raises:

`RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16`

In practice this happens under **module-level mixed precision** — wrappers that cast parameters but don't uniformly cast intermediate activations (FSDP2 `MixedPrecisionPolicy` is the canonical example, but the issue is not specific to it). `torch.autocast` users are unaffected, because autocast intercepts at the op level and the parameter dtype attribute doesn't change.

To be clear: this PR is **not** claimed to be the right layer of fix for any specific framework's wrapping configuration — incorrect wrap units or missing `cast_forward_inputs` need to be handled at the wrapper layer. The motivation here is purely consistency with `Gemma4UnifiedMultimodalEmbedder`'s existing override, which removes one possible mismatch surface.

### Tests

Plan to add a regression test in `tests/models/gemma4/test_modeling_gemma4.py` that casts only `embedding_projection` to bf16, feeds an fp32 activation, and asserts the forward succeeds and returns bf16. Pure PyTorch, no distributed setup. Red without the patch, green with it.

### Question for maintainers

Would you welcome a single PR that:
1. Adds the cast to `Gemma4MultimodalEmbedder.forward` (modular source).
2. Adds the same cast to `Gemma3nMultimodalEmbedder.forward` on the `inputs_embeds` branch (modular source).
3. Removes the now-redundant `forward` override in `Gemma4UnifiedMultimodalEmbedder` (modular source).
4. Adds the regression test described above.


Comments:
cc @zucchini-nlp and @Cyrilvallez — would appreciate your guidance on whether to proceed with the unified PR.

I want to add a short explanation on why the `to(dtype)` exist in the fist place. We usually don't force users to pass pixels/audio in model's dtype so the very first multimodal module will have to cast inputs. In Unified Gemma4 there's no encoder, thus the burden of casting dtype is on the MLP proj now

Re your issue with mixed precision training, I’m not sure we want to upstream a fix for this. In general, the model assumes input/weight dtype consistency, and adding defensive casts in the forward pass for wrapper-specific behavior can easily become hard to reason about and maintain long term

Make sense, thanks for the clear explanation!