Gemma RMSNorm weight init should init with 0
Author: albertz
Pull request: no
Labels: name: bug

### System Info

Looking at current master code.

### Who can help?

@ArthurZucker 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

The current code for weight initialization:
```python
    def _init_weights(self, module):
        ...
        elif (
            isinstance(
                module, (nn.LayerNorm, nn.RMSNorm, nn.GroupNorm, nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)
            )
            or "LayerNorm" in module.__class__.__name__
            or "RMSNorm" in module.__class__.__name__
        ):
            # Norms can exist without weights (in which case they are None from torch primitives)
            if hasattr(module, "weight") and module.weight is not None:
                module.weight.data.fill_(1.0)
        ...
```

That was recently changed in b16688e96a3b3e1e7a701cd5284a850c696c108e. Before, e.g. Gemma3 had this custom code:
```python
        elif isinstance(module, Gemma3RMSNorm):
            module.weight.data.fill_(1.0)
```

So the behavior has not changed (in this case).

However, this initialization is wrong (both now and before b16688e96), at least for `Gemma3RMSNorm`. This is because `Gemma3RMSNorm.forward` uses this code:
```python
        output = output * (1.0 + self.weight.float())
```
The `weight` should be initialized to 0, not 1.

`Gemma3RMSNorm.__init__` actually does it correctly:
```python
        self.weight = nn.Parameter(torch.zeros(dim))
```

Gemma2 and Gemma have the same issue. (But Gemma3n does not.) RecurrentGemma as well.
I browsed through a couple other model implementations but have not seen this issue elsewhere.



### Expected behavior

Param init of `Gemma3RMSNorm.weight` (and also other Gemmas except Gemma3n) should init with 0.

Comments:
Btw, I was wondering, why to do `output * (1 + weight)`, and not just `output * weight`: Weight decay probably works better with `output * (1 + weight)`.

Hey! Do you want to open a PR for a fix? 

I can. But I wanted to discuss the possible options how to fix this first. I see many different ways. For example:

1. An easy, but maybe not so generic / good way would be to just extend the generic base `_init_weights`, and have a special rule for all the Gemma*RMSNorm (except of Gemma3n). This code would need to be extended then for all further cases which do it like that. People might potentially forget about this.

2. We could rename `weight` in all the Gemma*RMSNorm (except Gemma3n) to sth else. But that might break some other code? Also needs probably some fixes (remappings) in model import, etc.

3. We could rename `weight` in those RMSNorms (to `add_weight`?) and add a `weight` `property` which returns `1 + add_weight`), and `_init_weights` adds another check if `weight` is a `torch.nn.Parameter` and only applies this `fill_` in that case, and `_init_weights` adds another check for `add_weight`. And this also needs the remappings for model import.

4. There could be an optional attrib like `weight_initial` or so and if available, `_init_weights` would use that instead of 1. We would set `weight_initial = 0` for all the relevant RMSNorms.


Let's go with 1.! 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Yes I'm still on it. I will update #40449 once I find some time. (Now there is ICLR deadline coming up...)

This is now fixed by #40796.