Refactor RMSNorm implementations to use torch.nn.functional.rms_norm
Author: mstojkovicTT
Pull request: yes

# What does this PR do?
Fixes #42398

This PR replaces custom `RMSNorm/T5-style` norm implementations (e.g. in Llama) that manually compute variance and scaling with the built-in `torch.nn.functional.rms_norm`. For example, code like:

``` python
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)
```

is simplified to:

``` python
return F.rms_norm(hidden_states, hidden_states.shape[-1:], self.weight, self.variance_epsilon)
```

This keeps the behavior and epsilon handling the same while reducing the number of ops, this should improve performance for users without requiring any additional changes on their side.

To verify the performance and the numerical stability, i have wrote the following test

``` python
import timeit
import torch
import torch.nn as nn

# Original implementation
class LlamaRMSNormHF(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)


# New implementation using torch.nn.functional.rms_norm
class LlamaRMSNormTorch(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        return nn.functional.rms_norm(
            hidden_states,
            hidden_states.shape[-1:],
            self.weight,
            self.variance_epsilon,
        )

def bench(module, x, iters=1000):
    # Warmup
    for _ in range(10):
        module(x)

    if x.is_cuda:
        torch.cuda.synchronize()
    start = timeit.default_timer()
    for _ in range(iters):
        module(x)
    if x.is_cuda:
        torch.cuda.synchronize()
    end = timeit.default_timer()

    return (end - start) / iters


def test_llama_rms_norm_equivalence():
    torch.manual_seed(0)

    hidden_size = 64
    batch_size = 2
    seq_len = 3

    dtype = torch.bfloat16
    device = "cpu" 

    x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device)

    hf_module = LlamaRMSNormHF(hidden_size, eps=1e-6).to(device)
    new_module = LlamaRMSNormTorch(hidden_size, eps=1e-6).to(device)

    # make sure they have the same weights
    with torch.no_grad():
        new_module.weight.copy_(hf_module.weight)

    y_hf = hf_module(x)
    y_new = new_module(x)
    y_new = y_new.to(y_hf.dtype) # torch.allclose needs same dtype

    # Check numerically close
    print(torch.allclose(y_hf, y_new, atol=1e-5, rtol=1e-5))

    # speed benchmark
    t_hf = bench(hf_module, x)
    t_new = bench(new_module, x)

    print(f"HF   RMSNorm: {t_hf * 1e6:.2f} µs / call")
    print(f"F.rms_norm  : {t_new * 1e6:.2f} µs / call")

test_llama_rms_norm_equivalence()
```

The results show the following:
- for cpu device:
```
True
HF   RMSNorm: 86.75 µs / call
F.rms_norm  : 47.25 µs / call
```
- for cuda device:
```
True
HF   RMSNorm: 112.27 µs / call
F.rms_norm  : 83.05 µs / call
```

note: I have encountered that when I try `dtypes` that are lower then `float32`, old implementation will keep it at `float32`, but my new one will have the `dtype` of the input tensor. Thats why i have to cast to `y_hf.dtype` (trying `float64` for example will make both implementation output `float64`). This can be changed, depending on what we want to accomplish.

## Who can review?
@Rocketknight1 

Comments:
Hey @mstojkovicTT, thanks for the PR! We definitely want the functions to be a drop-in replacement, so they should return exactly the same dtype as the old functions did.

Also, in your tests you're initializing `self.weight = nn.Parameter(torch.ones(hidden_size))` which means that the final scaling step is just multiplying by 1 and having no effect. Can you try with randomly initialized weights with a mean of 1 but a little bit of variance instead, so we can see if everything is equivalent with the original?

> We definitely want the functions to be a drop-in replacement, so they should return exactly the same dtype as the old functions did.

This may be problematic, because of the following scenario:
- Lets say that we have input that is `bfloat16`
- When i run the huggingface implementation i get a return dtype is `float32`
- I dont think this is intended, considering that we do
``` python
return self.weight * hidden_states.to(input_dtype)
```
- Here, one of the following should be our goal: ether we just want to upcast the input to `float32` because of precision in accumulating, in which case we dont care about the `input_dtype` and have it for no reason, or we want to accumulate and preserve the dtype (which is not the case).
- One case i see `hidden_states.to(input_dtype)` being usefull is when we have "larger" dtypes then `float32` like `float64`, but in that case, we did downcast to `float32` for no reason

Here is the [pytorch default implementation](https://github.com/pytorch/pytorch/blob/d3f61c1b1da0530839547a69602f83f9797ee15a/aten/src/ATen/native/layer_norm.cpp#L264) of `rms_norm` when it is not using custom `CUDA` or `MPS` kernels. Here we can see that they do pretty much the same thing, except that they just cast the output of the op to the input type, not the `hidden_states` that is later multiplied with weight (and that causes this wierd `dtype` change)

> Can you try with randomly initialized weights with a mean of 1 but a little bit of variance instead, so we can see if everything is equivalent with the original?

I did the same testing just with the additional

``` python
with torch.no_grad():
        random_weight = torch.randn(hidden_size, device=device) * 0.05 + 1.0
        hf_module.weight.copy_(random_weight)
        new_module.weight.copy_(random_weight)
```

and everything still works.

And also @Rocketknight1, thank you for taking a time to review this!


Yeah, I understand! My guess is that the original code was made to **follow the original model implementation, even if it seems weird**. Downcasting `hidden_states` to `bfloat16` and then back to `float32` seems to just pointlessly lose precision for no purpose, but probably if we "fix" it for pretrained models their accuracy will get a bit worse because we've created a shift from how they were trained. You need to be extremely careful when dealing with pretrained models.

Ah, that makes sense. What do you think about trying CI run against the changes for now, and if it breaks something I will take a closer look? @Rocketknight1 

Sure, let's see how it goes

run-slow: llama

This comment contains `run-slow`, running the specified jobs: 

models: ["models/llama"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/19896728031)

### Model CI Report

#### ❌ Failed tests

- [llama](https://github.com/huggingface/transformers/actions/runs/19896728031/job/57029807492):
    tests/models/llama/test_modeling_llama.py::LlamaIntegrationTest::test_llama_3_1_hard
    tests/models/llama/test_modeling_llama.py::LlamaIntegrationTest::test_model_7b_logits_bf16



Hmn, it does seem like this might shift logits slightly

Thanks for running this @Rocketknight1! 
I also reproduced the `logit` shift locally and wanted to share what I am seeing.

For `tests/models/llama/test_modeling_llama.py::LlamaIntegrationTest::test_model_7b_logits_bf16` on my machine:

### With my change (new behavior)

I see a small drift in the expected mean/slice:
- mean diff: `[[ 0.0090, 0.0019, -0.0504, 0.0004, 0.0058, -0.0122, 0.0050, 0.0407]]`
- expected mean: `[[-6.5208, -4.1218, -4.9377, -3.2536, 0.8127, -2.9811, 1.2918, -3.3848]]
`
- actual mean: `[[-6.5118, -4.1199, -4.9881, -3.2532, 0.8185, -2.9933, 1.2968, -3.3441]]`

And for the slice (`logits[0,0,0:15]`), I get:
- slice diff: 
``` python
[[ 6.2500e-02,  3.1250e-02, -7.8188e-03,  3.1250e-02,  5.0068e-05,
          0.0000e+00,  0.0000e+00,  0.0000e+00, -5.0068e-05,  3.1300e-02,
         -5.0068e-05, -5.0068e-05, -1.5638e-02, -2.3438e-02,  0.0000e+00]]
```
- expected slice: 
``` python
[-12.5625,  -7.1250,  -0.6289,  -7.8750,  -6.9688,  -7.8125,  -6.5000,
          -7.4375,  -7.6562,  -6.9688,  -6.0312,  -7.0312,  -1.8203,   1.8750,
          -8.5000]
```
- actual slice: 
``` python
[-12.5000,  -7.0938,  -0.6367,  -7.8438,  -6.9688,  -7.8125,  -6.5000,
         -7.4375,  -7.6562,  -6.9375,  -6.0312,  -7.0312,  -1.8359,   1.8516,
         -8.5000]
```

### Baseline / original implementation (before my change)

For comparison, with the original implementation I see different drift:
- mean diff: `[[-0.0539, -0.0090, 0.0018, 0.0743, 0.0040, -0.0363, -0.0077, 0.0440]]`
- expected mean: `[[-6.5208, -4.1218, -4.9377, -3.2536,  0.8127, -2.9811,  1.2918, -3.3848]]`
- actual mean: `[[-6.5747, -4.1308, -4.9359, -3.1793, 0.8167, -3.0174, 1.2841, -3.3408]]`

And for the slice (`logits[0,0,0:15]`)

- slice diff tensor
``` python
[-6.2500e-02,  6.2500e-02,  7.8119e-02, -9.3750e-02, -1.2495e-01,
         -6.2500e-02,  3.1250e-02, -1.2500e-01, -3.1300e-02,  5.0068e-05,
         -6.2550e-02, -3.1300e-02, -1.5638e-02,  2.3438e-02,  0.0000e+00]
```

- expected slice
``` python
[-12.5625,  -7.1250,  -0.6289,  -7.8750,  -6.9688,  -7.8125,  -6.5000,
          -7.4375,  -7.6562,  -6.9688,  -6.0312,  -7.0312,  -1.8203,   1.8750,
          -8.5000]
```

- actual slice
``` python
[-12.6250,  -7.0625,  -0.5508,  -7.9688,  -7.0938,  -7.8750,  -6.4688,
         -7.5625,  -7.6875,  -6.9688,  -6.0938,  -7.0625,  -1.8359,   1.8984,
         -8.5000]
```

So from my side it looks like the change does alter numerics, but it’s not obvious that it strictly makes things “worse”, it seems to move the error pattern relative to the golden expectations. I dont know if you have some kind of way to check the actual performance of the model?

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: aimv2, apertus, arcee, aria, bamba, bitnet, blt, chameleon, clvp, csm, cwm, deepseek_v2, deepseek_v3, dia, diffllama, doge

Hi @mstojkovicTT, just realized we might actually need to pause this PR! The reason is that I don't think Torch 2.3 (our current minimum supported version) has `functional.rms_norm()`. It was added in 2.4, so we might not be able to merge this for a few more months even if the numerics are okay.

That makes sense. Thanks for the responses, putting this on a pause till then :)

Yeah, and thanks for investigating! My suspicion is the numerics are similar enough that we might be able to just switch everything to using `F.rms_norm()` for speed/memory efficiency, but I might have to discuss internally. In some cases, it probably matches what the model was trained with more closely than our eager implementation does.