Fix MPS fused RMSNorm: do the weight multiply in fp32 to match CPU/CUDA#189617
Fix MPS fused RMSNorm: do the weight multiply in fp32 to match CPU/CUDA#189617Incheonkirin wants to merge 2 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/189617
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: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
Bot
added
ciflow/mps
|
The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:
Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows. |
|
@jhavukainen this replaces #187779. I accidentally detached my fork from pytorch's fork network while changing its visibility, which closed that PR and left it un-reopenable from my account. The patch here is line-for-line identical to the one you approved, rebased onto main. The one adverse check, macos-m2-26, was cancelled by pytorch-bot after 10h55m without ever getting a runner or starting a step; m1-14 and m2-15 passed on the same commit in ~44 min. Could you re-review and approve if it still looks good? A rerun of the cancelled job would help too, if you have the permissions. Thanks. |
albanD
added
the
triaged
|
@jhavukainen gentle re-ping: this is the resubmission of the #187779 patch you approved (identical lines). Could you re-review when you get a chance? |
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 pytorch#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).
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.
7b1a34d to
fd7d864
Compare
3 participants
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.
Validation
git range-diff: both reviewed commits are patch-equivalent after rebase