Fix gemma2 accuracy through the correct softcapping logic

[AllentDan]

NOTE
There is still another place that is not aligned with gemma2, but I don't think that is important.
Gemma2 implementation of RMS Norm:

class DefaultRMSNormImpl(RMSNormImpl):
    """RMS norm implementation api."""

    def __init__(self, hidden_size: int, eps: float = 1e-6):
        self.hidden_size = hidden_size
        self.eps = eps

    def forward(self,
                x: torch.Tensor,
                weight: torch.Tensor,
                residual: torch.Tensor = None):
        """forward."""
        input_dtype = x.dtype
        if residual is not None:
            x = x + residual
            residual = x
        x = x.to(torch.float32)
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        x = (weight.float() * x).to(input_dtype)
        if residual is None:
            return x
        return x, residual

Llama does x.to(float16) * w whilst Gemma is (x * w).to(float16)
See huggingface/transformers#29402