[release/2.12] [ROCm] Backport LayerNorm gamma/beta backward fix for wave32 (RDNA)#3425
[release/2.12] [ROCm] Backport LayerNorm gamma/beta backward fix for wave32 (RDNA)#3425amd-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 |
|
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/ |
|
FYI I tried to get my fix included in the PyTorch 2.12.1 patch release. I requested a review but nobody responded, so unfortunately it wasn't included in the end. PR: pytorch#185415 |
|
@bozhou2021 Great, thanks. Closing this one then |
|
@amd-maradosa why close this PR? 2.12 branch still has this bug. |
|
I thought if it ends up inside upstream repo, it doesn't need to be inside the ROCm one, reopening then. |
No problems. Thanks for this backport PR. I'm looking forward to it being merged. If anything I can be helpful, please let me know. |
|
Jenkins build for 06aff093f7ff2c725c8ab2cf45becdd1d90f773c commit finished as FAILURE |
2 participants
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)