Replace RMSNorm by nn.RMSNorm

[manuelcandales]

In this PR we replace torchchat's own RMSNorm implementation by nn.RMSNorm, and we bump the PyTorch pin to capture the massive speed up (30x-40x) to RMSNorm on MPS backend introduced in pytorch/pytorch#145301

Preliminary benchmarks on an M1 Pro with 16GB RAM, show a 33% speed up on token generation when running Llama 3.2 1B with 4-bit quantization

Motivation: Token generation on MPS backend is currently CPU bound, because of MPSGraph overhead. Surprisingly, the ops that are impacting performance the most are simple ones: mul, copy_, add, where, mean, rsqrt, sub, cat, stack. Experiments on an M1 Pro show that each of those op calls on the MPS backend, has at least 20us of CPU overhead. Also, these ops dominate the graph. For example, in aggregate, these ops are called 770 times for each token, when running Llama 3.2 1B. Compare that to SDPA which is called only 33 times, and linear which is called 113 times.

• mul is called 275 times per token
• copy_ is called 202 times per token
• add is called 97 times per token
• where is called 34 times per token
• mean is called 33 times per token
• rsqrt is called 33 times per token
• sub is called 32 times per token
• cat is called 32 times per token
• stack is called 32 times per token

Currently, torchchat's own RMSNorm operation is basically implemented like this:

norm = x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
output = norm(x.float()).type_as(x) * weight

This means that a single call to torchchat's RMSNorm involves 3 calls to aten::mul and calls to aten::rsqrt, aten::mean and aten::add. RMSNorm is called 33 times for each token. Hence, RMSNorm contributes 5 * 33 = 165 of those 770 op calls.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1464

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit b8801f5 with merge base f4ae60f ():

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / test-build-runner-et-android / linux-job (gh)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Jack-Khuu]

Thanks for the fix!!

Can you add a short snippet of your internal post to the description of this PR? It'll help with OSS visibility