[release/2.12] [ROCm] Backport LayerNorm gamma/beta backward fix for wave32 (RDNA)#3425
Open
amd-maradosa wants to merge 1 commit into
Open
[release/2.12] [ROCm] Backport LayerNorm gamma/beta backward fix for wave32 (RDNA)#3425amd-maradosa wants to merge 1 commit into
amd-maradosa wants to merge 1 commit into
Conversation
…#183864) FIXES pytorch#183861 LayerNorm/RMSNorm's gamma beta backward kernel produced wrong results for AMD Strix Halo and other RDNA GPUs. It was caused by a hardcoded warp size (incorrect for RDNA/Strix Halo), and mixed usage of thread block and warp coordinates. This caused half of the values (along the x-axis/features) being dropped when saving the gradients after partial sum reduction. The fix is to query and use correct warp size instead of the hardcoded values. - No semantic change to the existing AMD and CUDA code paths. - This will bring Strix Halo/RNDA to the proven/performant CUDA code path. Issue reproduced and tested on local Strix Halo hardware using: $python test/test_nn.py TestNN.test_layer_norm_backwards_eps -v See issue pytorch#183861 Pull Request resolved: pytorch#183864 Approved by: https://github.com/Skylion007, https://github.com/cyyever, https://github.com/jeffdaily (cherry picked from commit c3534e4)
|
Jenkins build for 06aff093f7ff2c725c8ab2cf45becdd1d90f773c commit finished as FAILURE |
Author
|
Jenkins tests that are failing are unrelated to changes in this PR. Here are the exact same tests failing on another open PR: https://ml-ci-internal.amd.com/job/pytorch/job/pytorch-ci-pipeline/job/PR-3327/5/testReport/ |
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
2 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.
Issue #6173
Cherry-pick of pytorch#183864 (c3534e4)
release/2.12 carries a hardcoded wave64 launch geometry in the LayerNorm gamma/beta backward kernel (block_dim_x = 64), which is only correct on CDNA. On wave32 RDNA hardware the launch geometry no longer matches the 32-lane wavefront, so the warp-shuffle reduction drops partial sums and produces wrong dgamma/dbeta.
The upstream fix makes block_dim_x a template parameter and dispatches on the runtime warp size (at::cuda::warp_size()), selecting the matching launch shape per GPU. CDNA and CUDA code paths are unchanged.
Tested on gfx1201 (RDNA4, wave32): test_nn.py::TestNN::test_layer_norm_backwards_eps passes with this change.
(cherry picked from commit c3534e4)