add custom RMSNorm to ALL_LAYERNORM_LAYERS

[shijie-wu]

What does this PR do?

It fixed a issue discovered during discussion of PR #26152.

@ArthurZucker: the ALL_LAYERNORM_LAYERS should contain all the custom layer norm classes (from transformers modeling files) and should be updated if that is not the case

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@ArthurZucker

[ArthurZucker]

Thanks! Could you also add IdeficsRMSNorm and MegaRMSNorm? 😉 (make sure make fixup and make fix-copies also pass!)

[shijie-wu]

@ArthurZucker done! MegaRMSNorm - a part of MegaSequenceNorm - is already in ALL_LAYERNORM_LAYERS

transformers/src/transformers/models/mega/modeling_mega.py

Line 317 in 493b24b

ALL_LAYERNORM_LAYERS.append(MegaSequenceNorm)

do we want to add all these to ALL_LAYERNORM_LAYERS?

transformers/src/transformers/models/mega/modeling_mega.py

Lines 294 to 301 in 493b24b

	elif norm_type == "scalenorm":
	self.norm = MegaScaleNorm(dim=-1, eps=eps, affine=affine)
	elif norm_type == "rmsnorm":
	self.norm = MegaRMSNorm(embedding_dim, eps=eps, affine=affine)
	elif norm_type == "batchnorm":
	self.norm = nn.BatchNorm1d(embedding_dim, eps=eps, affine=affine)
	elif norm_type == "syncbatchnorm":
	self.norm = nn.SyncBatchNorm(embedding_dim, eps=eps, affine=affine)

[ArthurZucker]

Thanks again! Let's add all of the mega ones as you said!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[shijie-wu]

after discussion with @ArthurZucker, it seems best to limit to scope of this PR to adding custom RMSNorm to ALL_LAYERNORM_LAYERS. Adding nn.BatchNorm1d might create unintentional impact.