Fix MPS fused RMSNorm: do the weight multiply in fp32 to match CPU/CUDA#187779
Fix MPS fused RMSNorm: do the weight multiply in fp32 to match CPU/CUDA#187779Incheonkirin wants to merge 2 commits into
Conversation
The MPS fused RMSNorm kernel cast the normalized value to the half dtype before the weight multiply, diverging from the CPU composite and CUDA fused/reference paths which keep x*inv*weight in fp32 and cast once (CPU aligned in #147203). Do the weight multiply in fp32 in the reference order ((x*inv)*weight) and cast once. Adds an MPS regression test (fp16/bf16, N=4095/4097).
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/187779
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ✅ No FailuresAs of commit 3b06ac9 with merge base fd03d74 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
|
|
|
jhavukainen
left a comment
There was a problem hiding this comment.
I think the idea is correct, we shouldn't truncate mid computation and aligning with the other backend implementations makes sense here.
Left a few comments on how this might be tidied up a bit.
Per review: the four write sites repeated the same upcast, so move it into one inline helper and use opmath_t<T> instead of hardcoding float (needs c10/metal/utils.h for opmath_t). No behavior change.
jhavukainen
left a comment
There was a problem hiding this comment.
Looks good to me after the cleanup.
|
I accidentally closed this PR while changing the visibility of my local upstream mirror. The fix is still absent from main, and the existing MPS review and CI were green. I have restored the source repository and branch, but GitHub no longer lets me reopen the PR from my account after the fork relationship was detached (the API still reports the patch as mergeable).\n\n@jhavukainen, could you please reopen it? Sorry for the noise. |
|
Replacement PR is now #189617 from a proper fork; no reopen is needed. It preserves the reviewed patch and is rebased onto current main. |
What does this PR do?
The MPS fused RMSNorm kernel casts the normalized value to the half dtype before multiplying by the weight (
out = w * static_cast<T>(x * inv)), so its output diverges from PyTorch's CPU composite and CUDA fused/reference behavior, which keep thex * inv * weightproduct in fp32 and cast once at the end (the CPU side was aligned in #147203). This affects half/bfloat16 inference on MPS wheneverRMSNorm/rms_normdispatches to the fused affine kernel (weight present, matching dtype, no grad).Fix: do the weight multiply in fp32 in the same order as the reference (
(x * inv) * weight) and cast once. Adds an MPS regression test that runs the fused kernel in fp16/bf16 and in fp32 with the same inputs (identical reduction) and asserts they agree, so only the weight-multiply precision is checked (N = 4095, 4097, covering both shader paths).Reproduced on Apple Silicon; the kernel recompile (a full source build) is left to CI.