Add GraniteRMSNorm

[NielsRogge]

What does this PR do?

This PR is a follow-up of #31502 which broke Transformers for PyTorch < 2.4.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[mayank31398]

Hi this makes sense.
I had added the nn.RMSNorm class
but looks like there is some issue with fp16 as well: pytorch/pytorch#134106

[ArthurZucker]

thanks

[NielsRogge]

Should I run slow tests or can this be merged as-is?

[mayank31398]

might be better to run slow tests on the granite class @NielsRogge

[NielsRogge]

Seems like the slow tests are failing (cc @ydshieh), but I assume it's safe to merge this PR since the following passes

from torch import nn
import torch


class GraniteRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        GraniteRMSNorm is equivalent to T5LayerNorm
        """
        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)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
    

a = torch.nn.RMSNorm(10)

b = GraniteRMSNorm(10)

assert a.weight.shape == b.weight.shape

c = torch.randn(1, 10)

assert torch.allclose(a(c), b(c))

[iamsaurabhgupt]

when will this release to fix transformers?
creating lot of conflicts as pytorch==2.4 is not supported by many packages like vllm, flash_attn, etc.

[mayank31398]

@ArthurZucker can you merge this?
I will create a new PR for fixing the slow tests.
Looks like its blocking the DeepSpeed team.

[ArthurZucker]

Okay

[ArthurZucker]

cc @ydshieh if you see failures on this, it's expected!

[mayank31398]

Thanks @ArthurZucker , Ill fix this test in a new PR