[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 |
|
Given the just announced AMD Helios system on MI455X (gfx1250) is using wave32, it would be great to have this merged asap to give users/researchers a smooth experience. And even better if we can help the official PyTorch upgrade from ROCm 7.2 to the latest. (FYI I hit on an MIOpen bug recently using the official PyTorch2.13+ROCm7.2 build, which seems to have been fixed in the latest TheRock version of ROCm. However TheRock doesn't provide PyTorch2.13, which is the only version has this LayerNorm fix. This forced me to disable MIOpen and use PyTorch fallback instead). MI455X (gfx1250) is using wave32 as documented here: |
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)