Align weight dtype in DISABLED grad_input path under autocast - #4689
Open
guptaishaan wants to merge 1 commit into
Open
Align weight dtype in DISABLED grad_input path under autocast#4689guptaishaan wants to merge 1 commit into
guptaishaan wants to merge 1 commit into
Conversation
Float8Linear.forward casts only the input to the autocast dtype, so save_for_backward stores a bf16 input next to an fp32 weight. The forward gemm survives because autocast is still active there and torch.mm is autocast-eligible. Backward runs with autocast off, so when cast_config_weight_for_grad_input is ScalingType.DISABLED the raw fp32 weight reaches torch.mm against a bf16 grad_output and it raises "expected mat1 and mat2 to have the same dtype". Cast the saved weight to the dtype of the grad_output tensor it is multiplied against. This is a no-op without autocast, where both are already fp32. The grad_weight gemm needs no change, input_hp was already cast to the autocast dtype in forward. Adds test_autocast_backward_scaling_disabled, which uses an all-DISABLED config so it runs without sm_89 and checks both gradients against a plain nn.Linear reference under the same autocast. Fixes pytorch#4616
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4689
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
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
1 participant
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 #4616
Float8Linear.forwardcasts only the input to the autocast dtype, soctx.save_for_backwardholds a (bf16 input, fp32 weight) pair. The forward gemm gets away with this because autocast is still active there andtorch.mmis autocast-eligible, but backward runs with autocast off, so theScalingType.DISABLEDbranch forcast_config_weight_for_grad_inputhands the raw fp32 weight totorch.mmagainst a bf16grad_outputand raisesexpected mat1 and mat2 to have the same dtype.Cast the saved weight to the dtype of the grad_output tensor it is multiplied against. No-op without autocast, since both are already fp32. The
grad_weightgemm is unaffected,input_hpwas already cast to the autocast dtype in forward.Verified on 1x A40 (sm_86), torch 2.13.0+cu126:
float8_linear.py:143, both with the issue's config plusemulate=Trueand with all six cast configs set toDISABLED. Both pass after the change.TestFloat8Linear::test_autocast_backward_scaling_disabledfails before the change with the reportedRuntimeErrorand passes after. It uses the all-DISABLED config so it needs no fp8 hardware, and asserts both gradients match a plainnn.Linearreference bitwise under the same autocast.pytest test/float8/test_base.py66 passed, 60 skipped.pytest test/float8/minus the distributed modules, 95 passed, 76 skipped. Ruff lint and format clean.Not verified: sm_86 cannot run
torch._scaled_mm, so the real fp8 path (DYNAMIC forward with DISABLEDweight_for_grad_input,emulate=False) was not executed, and neither were the FSDP/DTensor tests.Float8TrainingTensor.dtypeis_orig_dtype, so the new.to()is well defined on that path, but it was not run on hardware. Only bf16 autocast was tested, not fp16.Thanks to @TangSiLiYang for the report and the diagnosis, which pinned the exact line and the correct fix.