🐛 Describe the bug

based on this issue: #156012 and this PR: #156903

The fix patched forward but did not patch backward.
To patch backward, add
logsumexp /= 0.6931471805599453
at
https://github.com/ROCm/pytorch/blob/cfa0de7c5151cfd4d036b2b4ee6d35a37bd7a983/torch/distributed/tensor/experimental/_attention.py#L498

Versions

before patching gradient diff is 1e-1, after patching is 1e-7

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @naromero77amd