Fix rms_norm in fp16/bf16#147203
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/147203
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 5a0c8cf with merge base 81847d0 ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchbot label "topic: not user facing" |
colesbury
added
the
triaged
eqy
left a comment
eqy
left a comment
There was a problem hiding this comment.
Would it make sense to add a test that checks for whether the expected tolerances are met?
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 1 mandatory check(s) failed. The first few are: Dig deeper by viewing the failures on hud |
|
@pytorchbot merge |
|
Pull workflow has not been scheduled for the PR yet. It could be because author doesn't have permissions to run those or skip-checks keywords were added to PR/commits, aborting merge. Please get/give approval for the workflows and/or remove skip ci decorators before next merge attempt. If you think this is a mistake, please contact PyTorch Dev Infra. |
|
@pytorchbot rebase |
|
You don't have permissions to rebase this PR since you are a first time contributor. If you think this is a mistake, please contact PyTorch Dev Infra. |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 1 mandatory check(s) failed. The first few are: Dig deeper by viewing the failures on hud |
|
@pytorchbot merge -i |
|
|
|
@pytorchmergebot rebase |
|
@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here |
|
Successfully rebased |
b15c62b to
5a0c8cf
Compare
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
|
@riccardofelluga @eqy This PR seems to break torchtitan float8 training with rowwise scales, when RMS norm is used.
Can we either revert this change or look into a fix asap please? This NaN issue is currently blocking the release of a blog post on float8 rowwise training, so we are eager to resolve it as soon as possible. Thanks! cc @vkuzo @lessw2020 |
|
fyi @drisspg on #147203 (comment) as well since Vasiliy is OOO |
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).
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).
6 participants
Fixes #134106. This PR moves the
upcasted_resultdown-casting after all computation is done.Since the multiplication with the weight_opt input is not done in half precision, the current code path is doing the following: fp16 -> fp32 -> fp16 -> fp32 -> fp16. What we want tho is to avoid down-casting and this PR proposes: fp16 -> fp32 -> fp16. This results in better accuracy as it avoids truncating.