gemma3 fp16 fix

[mobicham]

What does this PR do?
Fixes float16 inference with Gemma 3 models by simply clipping the activations. The residual addition step should also be clipped for more accurate outputs. Without this fix, the output is empty because of NaN in the activations.

Fixes #36822

Test it here: https://gist.github.com/mobicham/2be18274476bbb32fe6c7aab9d50b316

Outputs:

google/gemma-3-12b-it - BFP16 (HQQ 8-bit, gs=64)
#################################################
Here's a detailed description of the image:

**Overall Impression:**

The image is a close-up, vibrant photograph of a cluster of pink cosmos flowers with a bumblebee actively foraging on one of them. The scene is set in a garden or natural setting, with a blurred background of foliage.

**Main Elements:**

*   **Cosmos Flowers:** The primary focus is on the cosmos flowers. They are a soft, delicate pink color with broad, slightly ruffled petals. Several flowers are in various stages of bloom, with some fully open and others still budding or past their prime.
*   **Bumblebee:** A bumblebee is prominently positioned on one of the cosmos flowers. It appears to be actively collecting pollen, with its body covered in yellow and black stripes.
*   **Foliage:** The background is filled with green foliage, including large leaves and stems of other plants. The foliage is out of focus, creating a soft, blurred backdrop that highlights the flowers and bee.
*   **Dried Flowers:** There are also some dried, seed-filled flower heads visible among the cosmos, adding a touch of natural decay and the cycle of life to the scene.

**Color Palette:**

The dominant colors are pink (from the cosmos), green (from the foliage), yellow and black (from the bumblebee), and hints of red from another flower in the lower right corner.

**Composition:**

The composition is well-balanced, with the cosmos flowers and bumblebee taking center stage. The blurred background adds depth and draws the viewer's eye to the main subject. The arrangement of the flowers and foliage creates a natural, organic feel.

**Overall Tone:**

The image evokes a sense of natural beauty, tranquility, and the importance of pollinators in the ecosystem.


google/gemma-3-12b-it - FP16 (HQQ 8-bit, gs=64)
##################################################
Here' bee! Let's describe this image!

Here's a detailed description of the image:

**Overall Impression:**

The image is a close-up, medium-angle shot of flowers in a garden setting. It has a soft, natural lighting that suggests a bright day. The overall feeling conveyed is one of natural beauty and a busy ecosystem.

**Subject Details:**

*   **Flowers:** The main focus is on several pink cosmos flowers. They are fully bloomed and appear to be the stars of the show. The petals are a delicate shade of pink, and their shape is typical of cosmos flowers – star-shaped.
*   **Bee:** A bumblebee is captured in action, busily collecting nectar from one of the cosmos flowers. It's a key element of the image, conveying a sense of life and the interconnectedness of nature.
*   **Background Foliage:** The background is filled with a variety of green foliage, creating a lush and natural backdrop for the flowers. There are also some spent flower heads and stems, indicating the cycle of plant life.
*   **Other Flowers:** There are also a few red flowers visible in the bottom-right corner of the frame, which adds some color variation to the overall image.

**Compositional Elements:**

*   **Shots Framing:** The frame is tightly framed on the flowers, with a shallow depth of field. This draws the viewer's eye directly to the main subject matter.
*   **Shots Framing:** The composition is fairly balanced, with the main mass of flowers in the center of the frame.
*   **Color Palette:** The color palette is dominated by pinks, greens, and touches of red, creating a soft and harmonious color scheme.

**Overall Impression:**

The image is beautifully composed and conveys a strong sense of natural beauty and the busy activity of nature. The bee collecting nectar from the cosmos flower is a focal element that adds to the overall charm of the image.

Who can review?
@ArthurZucker @Rocketknight1

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the Ready for review button (at the bottom of the PR page).

[Rocketknight1]

This looks okay to me, but I'm a bit worried about the clamp_ - it implies that the model is getting overflows in float16, and after clamping the values may be significantly incorrect depending on how large the overflow was.

As a result, the model quality might significantly decline in float16 instead of bfloat16. There's not much we can do about this, but we should be careful not to advertise that we fully support it!

cc @ArthurZucker do you think it's okay?

[mobicham]

@Rocketknight1 True! But it's better than not outputting anything. The few tests I ran looked quite fine. I can run some evaluations to compare bfloat16 vs. fp16 with this fix.

[mobicham]

Meanwhile, I am also looking into more accurate ways of doing this...so let's not merge yet when approved.

[ArthurZucker]

My main inquire is to put this in the layer norm layer directly instead of the decoder layer if possible. Now I get that's not super feasible if you want to keep float32 for the residuals, but if you return hiddenstates as f32 but residuals are fp16 mixed precision should be alright no?

[mobicham]

My main inquire is to put this in the layer norm layer directly instead of the decoder layer if possible. Now I get that's not super feasible if you want to keep float32 for the residuals, but if you return hiddenstates as f32 but residuals are fp16 mixed precision should be alright no?

Yeah that was my first version, then I noticed that the range expands when the residual is added, so they should be clipped before entering the layernorm.

Let me run some lm-eval benchmarks to see the impact between different settings, it's difficult to judge with a single output.

[mobicham]

@ArthurZucker lm-eval results

google/gemma-3-12b-it - bfloat16
mmlu: 0.730
winogrande: 0.766
gsm8k: 0.874
truthfulqa_mc2: 0.580
arc_challenge: 0.7235
hellaswag: 0.839
--------------------------
Average: 0.752

google/gemma-3-12b-it - float16 fix
mmlu: 0.703
Winogrande: 0.701
gsm8k: 0.647
truthfulqa_mc2: 0.535
arc_challenge: 0.6186
hellaswag: 0.739
--------------------------
Average: 0.6573

[ArthurZucker]

Thanks for the update, let's keep this open for reference, unless the google team wants to merge it!