[RFC] Fix Gemma 3 FP16 with activation scaling
Author: gau-nernst
Pull request: yes

# What does this PR do?

## The problem

Gemma 3 doesn't work with FP16 inference. The cause of it seems to be

1. Embedding output is scaled by sqrt(hidden_size) -> large activations to the transfromer trunk

https://github.com/huggingface/transformers/blob/199d7adf1037e1e3a0a7a2ceb155e468abef5e26/src/transformers/models/gemma3/modular_gemma3.py#L575-L578

2. Hence, the weight of post norms are huge, to counteract the huge activations in the identity stream

https://github.com/huggingface/transformers/blob/199d7adf1037e1e3a0a7a2ceb155e468abef5e26/src/transformers/models/gemma3/modular_gemma3.py#L524-L528

![image](https://github.com/user-attachments/assets/5c386e7b-b53f-4531-9305-4d626dda3d27)

This makes activations go beyond the range of FP16 (~65k)

## The fix

We want to scale the "main" activations down, roughly by sqrt(hidden_size). To do this, we only need to scale the 2 following:
- Embedding layer
- Weight of `post_xx_layernorm`

This is because any scaling before `pre_xx_layernorm` will not change the math due to normalization. We need to scale the weight of `post_xx_layernorm` so that it is consistent with input to `pre_xx_layernorm`.

The final norm `model.norm` will remove the scaling, before LM head.

## Sanity check

```python
import torch
from transformers import pipeline, Gemma3ForCausalLM, AutoTokenizer

model_id = "google/gemma-3-4b-it"
model = Gemma3ForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    activation_scale=1/64,
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device="cuda")

messages = [
    {"role": "user", "content": "Write a poem on Hugging Face, the company"},
]
print(pipe(messages, max_new_tokens=100))
```

- BF16, no activation scale (default to 1) -> normal
- FP16, no activation scale (default to 1) -> CUDA error due to inf/nan
- FP16, activation scale=1/64 -> normal

Thank you to X users for drawing my attention to this and coming up with the fix: https://x.com/SeunghyunSEO7/status/1907350826940805266

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@ArthurZucker
- Feel free to test with larger models. I tested with 4B and 12B. For larger model, you may need to use smaller activation scale (1 / sqrt(hidden_size))
- This is an RFC since you may want a different design to introduce scaling
- Technically the scaling to `post_xx_norm` can be absorbed into the weight, but it's risky, since it needs to modify weight loading code. And things like `load_state_dict()` will not work correctly (or whatever method people might use to load weights)

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @eustlb
- graph models: @clefourrier

Library:

- flax: @gante and @Rocketknight1
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @zach-huggingface and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc @zach-huggingface
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @Rocketknight1
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page). This will assign reviewers and trigger CI.

I think we can support this, but we should be careful that there is no impact on performance/quality in float32/bfloat16, which are much more important. In general, I see float16 as a "legacy" format right now that is rarely used in practice, with most large models training in bfloat16 instead.

> there is no impact on performance/quality in float32/bfloat16, which are much more important

I totally agree.

By default, `activation_scale=1`, so there "shouldn't" be any difference. To be more certain, we can run it on a few evals. Or compare the outputs given fixed inputs (can be random)

Will you be able to help me verify it? I can try to run the smaller ones (1B, 4B), but I definitely don't have the hardware to run 12B and 27B in BF16.

cc @arthurzucker @cyrilvallez as core maintainers first - do we want to change the code to support `float16` when `bfloat16` and `float32` already work well, and I don't think Gemma3 was ever intended to support `float16` inference?

Some input. Apart from supporting older hardware, many quant kernels, like AutoAWQ, Marlin, only support FP16, hence being able to run in FP16 is invaluable.

> I don't think Gemma3 was ever intended to support float16 inference?

If that's the case, Gemma3 should error out if users try to load it in FP16. Because right now, it will give gibberish outputs (or simply doesn't work at all)

@gau-nernst thanks for this PR! Unfortunately I will not have a chance to review it this week. @ArthurZucker @Rocketknight1  can you please hold off on accepting until I, @osanseviero, or another Googler LGTM; we want to make sure whatever happens here is aligned across our other Gemma implementations. 

> missing some docs on the page to mention why or when you should set this and what value you should give

Where should this be?

Hey! Sorry I'm just seeing this right now! If I understand correctly, we have issues when loading the model in fp16 instead of bf16 right? In which case, I would suggest to simply use the `_keep_in_fp32` flag for that particular weight, which will make it always be loaded in fp32, so precision should not be lost. This is the reason why it was originally introduced, see the comment I added [here](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L4427-L4430). 

Would that solution fix the issue? It would be simpler, and would probably be more aligned with the Google's team code no?

The issue is in the activations, not the weights. `x + postnorm(...)` will explode. If you want to use FP32 to fix the issue, the activation flow must stay in FP32 until the next pre norm, which scales the activations back to the small range, representable by FP16. It won't be exactly a clean/simple fix. Iirc, Unsloth uses this fix cc @danielhanchen.

(hence, llama.cpp also doesnt have the issue, though it primarily uses FP16, because all activations in llama.cpp uses FP32 iirc)

Which approach is better is for you to decide. 

I see, then if we cannot solve it simply by forcing a given precision for a few weights, I'm happy with your solution given that the default does not change. Let's just add a bit more details in the docstring of the new config arg though, as @ArthurZucker mentioned 🤗

We still need the green light from @RyanMullins as well, as per his request 🤗

cc @RyanMullins !

Thanks for the ping @ArthurZucker and sorry for the prolonged delay @gau-nernst.

The principle of this solution is sensible, but I want to minimize the DX disruption as much as possible by recommending exactly one `activation_scale` value for all models. 

Given the goal is to...

> scale the "main" activations down, roughly by sqrt(hidden_size).

...and we have four different hidden sizes for the 1B (1152, sqrt ~= 33.94), 4B (2560, sqrt ~= 50.60), 12B (3840, sqrt ~= 61.97), and 27B (5376, sqrt ~= 73.32), can we recommend `activation_scale=1/74` for all models when using `float16`?