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

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 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

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 ⚙️

Model CI Report

❌ Failed tests

• llama:
  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:

[[ 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:

[-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:

[-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

[-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

[-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

[-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.