Description
Feature request
LayerNorm and RMSNorm are the top two normalization methods in Transformers. We unify them in Pre-Normalization Transformers in our paper https://arxiv.org/abs/2305.14858. The arithmetic equivalence allows us to convert Pre-LN Transformers into Pre-RMSNorm models without impact on the model functionality. Since RMSNorm offers superior efficiency compared to LayerNorm, our method enables faster equivalent inference and training for any Pre-LN Transformers, e.g., GPT, ViT.
Our implementation is at https://github.com/ZixuanJiang/pre-rmsnorm-transformer. As the first step, we can start by accelerating the deployment of the existing Pre-LN Transformers.
Motivation
Related GitHub issue: pytorch/pytorch#72643 (comment)
Your contribution
We have provided our reference implementation at https://github.com/ZixuanJiang/pre-rmsnorm-transformer. We are open to submitting a related PR in the future.