Add RMS Norm layer

[PraNavKumAr01]

🚀 The feature, motivation and pitch

Hey team, i love building things from scratch, and as i was implementing the LLaMa paper by meta obviously using pytorch i saw that pytorch did not have a nn.rmsnorm function for RMS Normalization layer . I did end up implementing it on my own, but i still feel like RMS Normalization being used in mutliple other new language models, LLaMa being the most popler one, needs to be implemented in Pytorch and should be a part of it.

Alternatives

There are many other normalization layer functions already in pytorch, but as results show, RMS Norm has been really impactful in the success of LLaMa, so i think the community might benefit with it being in pytorch

Additional context

Also, ive been using pytorch for a long time and ive always build things from scratch, and it has always been on my todo list to contribute to pytorch, i feel like this might be it.

[vadimkantorov]

Related:

• upstream apex.normalization.FusedRMSNorm #72643

I think, this layer was recently added in PyTorch, but not in a fused/optimized version yet

[PraNavKumAr01]

@vadimkantorov Oh gotcha, maybe ill read through that thread and work on the optimization