[AMP] [BUGFIX] Fixes RMSNorm mixed precision - #174824
Closed
AaronWang04 wants to merge 2 commits into
Closed
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/174824
Note: Links to docs will display an error until the docs builds have been completed.
|
This PR needs a
|
Contributor
Author
|
@pytorchbot label "topic: not user facing" |
AaronWang04
changed the title
AaronWang04
changed the title
AaronWang04
marked this pull request as ready for review
February 12, 2026 02:02
Collaborator
|
Please add regression test |
malfet
added
the
triaged
eqy
approved these changes
Feb 18, 2026
Collaborator
|
@pytorchmergebot merge |
Collaborator
Merge failedReason: 1 mandatory check(s) are pending/not yet run. The first few are:
Dig deeper by viewing the pending checks on hud |
AaronWang04
force-pushed
the
add-rms-norm-autocast-fp32
branch
from
February 18, 2026 21:35
de89115 to
ec1521b
Compare
Collaborator
|
@pytorchmergebot merge |
Collaborator
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 |
Collaborator
Merge failedReason: 1 mandatory check(s) failed. The first few are: Dig deeper by viewing the failures on hud |
Collaborator
|
@pytorchmergebot rebase |
Collaborator
|
@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here |
Collaborator
|
Successfully rebased |
pytorchmergebot
force-pushed
the
add-rms-norm-autocast-fp32
branch
from
February 19, 2026 19:29
ec1521b to
c29b8e9
Compare
Contributor
Author
|
@pytorchmergebot merge |
Collaborator
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 |
norx1991
pushed a commit
that referenced
this pull request
Feb 24, 2026
Fixes #167308 Running the repro script from the issue Before: ``` Testing MinimalRMS with apply_linear=False, layernorm=RMSNorm <FusedRmsNormBackward0 object at 0xea647dfdb130> Testing MinimalRMS with apply_linear=True, layernorm=RMSNorm /opt/pytorch/pytorch/torch/nn/functional.py:2954: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /opt/pytorch/pytorch/aten/src/ATen/native/layer_norm.cpp:344.) return torch.rms_norm(input, normalized_shape, weight, eps) <ToCopyBackward0 object at 0xea647db65a50> Testing MinimalRMS with apply_linear=False, layernorm=LayerNorm <NativeLayerNormBackward0 object at 0xea647dfdb130> Testing MinimalRMS with apply_linear=True, layernorm=LayerNorm <NativeLayerNormBackward0 object at 0xea647db65a50> ``` After ``` Testing MinimalRMS with apply_linear=False, layernorm=RMSNorm <FusedRmsNormBackward0 object at 0xe08d32d630a0> Testing MinimalRMS with apply_linear=True, layernorm=RMSNorm <FusedRmsNormBackward0 object at 0xe08d3458b7f0> Testing MinimalRMS with apply_linear=False, layernorm=LayerNorm <NativeLayerNormBackward0 object at 0xe08d32d630a0> Testing MinimalRMS with apply_linear=True, layernorm=LayerNorm <NativeLayerNormBackward0 object at 0xe08d3458b7f0> ``` Pull Request resolved: #174824 Approved by: https://github.com/eqy
EmanueleCoradin
pushed a commit
to EmanueleCoradin/pytorch
that referenced
this pull request
Mar 30, 2026
Fixes pytorch#167308 Running the repro script from the issue Before: ``` Testing MinimalRMS with apply_linear=False, layernorm=RMSNorm <FusedRmsNormBackward0 object at 0xea647dfdb130> Testing MinimalRMS with apply_linear=True, layernorm=RMSNorm /opt/pytorch/pytorch/torch/nn/functional.py:2954: UserWarning: Mismatch dtype between input and weight: input dtype = c10::BFloat16, weight dtype = float, Cannot dispatch to fused implementation. (Triggered internally at /opt/pytorch/pytorch/aten/src/ATen/native/layer_norm.cpp:344.) return torch.rms_norm(input, normalized_shape, weight, eps) <ToCopyBackward0 object at 0xea647db65a50> Testing MinimalRMS with apply_linear=False, layernorm=LayerNorm <NativeLayerNormBackward0 object at 0xea647dfdb130> Testing MinimalRMS with apply_linear=True, layernorm=LayerNorm <NativeLayerNormBackward0 object at 0xea647db65a50> ``` After ``` Testing MinimalRMS with apply_linear=False, layernorm=RMSNorm <FusedRmsNormBackward0 object at 0xe08d32d630a0> Testing MinimalRMS with apply_linear=True, layernorm=RMSNorm <FusedRmsNormBackward0 object at 0xe08d3458b7f0> Testing MinimalRMS with apply_linear=False, layernorm=LayerNorm <NativeLayerNormBackward0 object at 0xe08d32d630a0> Testing MinimalRMS with apply_linear=True, layernorm=LayerNorm <NativeLayerNormBackward0 object at 0xe08d3458b7f0> ``` Pull Request resolved: pytorch#174824 Approved by: https://github.com/eqy
yaoyi1222
added a commit
to yaoyi1222/pytorch
that referenced
this pull request
Jul 26, 2026
Fixes pytorch#188955. pytorch#174824 put `rms_norm` on the fp32 cast policy to fix pytorch#167308, where a bf16 activation meeting an fp32 weight tripped the dtype-mismatch guard in `rms_norm_symint` and fell back to the composite path with a warning. The fp32 policy fixed that by making the two dtypes agree, but it applies unconditionally, including when input and weight are already both bf16. In that configuration the promotion cannot recover any information: `rms_norm` upcasts internally regardless, so the only effect is a full-size fp32 copy of the activation and, under torch.compile, standalone `triton_poi_fused__to_copy` kernels for casts that did not fuse into the norm. The promote policy already expresses the intended rule: run in the widest dtype among the arguments. Matched bf16 stays bf16, and a bf16 activation with an fp32 weight still resolves to fp32, so pytorch#167308 stays fixed. This depends on the `prioritize()` fixes in the preceding commit, without which the optional `weight` argument would not participate in the decision at all and pytorch#167308 would regress. Note that `layer_norm` and `native_layer_norm` pay the same cost and are deliberately left alone here. Their fp32 policy predates 2023 and covers far more code, so changing it is a much larger compatibility question than restoring the 2.11 behavior of an op whose policy is five months old. Test Plan: Autocast coverage, on an RTX 3080 with a CUDA build of this commit. The `rms_norm` entry moves lists, and its new entry passes an fp32 weight so that the optional-argument path is exercised: ```bash python test_cuda.py -k autocast python test_cuda.py -v TestCudaAutocast.test_autocast_torch_fp32 python test_cuda.py -v TestCudaAutocast.test_autocast_torch_need_autocast_promote ``` Resulting dtypes under `torch.autocast("cuda", dtype=torch.bfloat16)`: ``` input weight out bfloat16 bfloat16 bfloat16 (was float32) bfloat16 float32 float32 float32 float32 float32 float16 float32 float32 ``` pytorch#167308's repro shape still reaches the fused kernel, with no mismatch warning: a bf16 activation from a preceding `Linear` into an fp32-weight `RMSNorm` yields `FusedRmsNormBackward0`, not `ToCopyBackward0`. For the memory regression itself, the repro from pytorch#188955 (Linear -> RMSNorm, bf16, compiled, forward+backward, input (8, 4096, 4096)): ```bash python repro.py # norm inside autocast python repro.py exempt # norm wrapped in autocast(enabled=False) ``` Before, the two modes differ; after, they are identical, which is the point: ``` 2.12.1 autocast peak 1.88 GiB 3 cast kernels 2.12.1 exempt peak 1.38 GiB 0 cast kernels patched autocast peak 1.38 GiB 0 cast kernels patched exempt peak 1.38 GiB 0 cast kernels ``` Step time is unchanged (53.1 ms before, 54.0 ms after, within run-to-run noise), so this is a memory fix and not a throughput one. The saving is one fp32 activation copy. Under torch.compile it does not accumulate with the number of norm layers, since the partitioner recomputes the cast rather than saving it; measured at 1/2/4/8 norms the delta stays flat at 0.12 GiB for a (2, 4096, 4096) activation. In eager it does accumulate, at roughly 0.06 GiB per norm for the same shape. Numerically the promotion was buying nothing in this configuration. Against an fp64 reference on identical bf16 inputs, backward error is the same to four significant figures either way (dx 3.099e-03, dw 1.961e-03), and the forward result is bit-identical once the fp32 output is rounded to bf16, which is what the next lower-precision op does anyway. Authored with assistance from Claude Code. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #167308
Running the repro script from the issue
Before:
After
cc @mcarilli @ptrblck @leslie-fang-intel @jgong5