Repository: pytorch
Issue: Gradient value calculation error in MultiLabelMarginLoss
Labels: module: loss,module: cuda,triaged,module: correctness (silent)

### 🐛 Describe the bug

We make sure that the gradient value of the backward of the multilabel_margin_loss operator is computed wrongly under some input cases. The situation is clear on some hardware, such as NVIDIA T4 / 1080, but it may occur randomly on V100 (currently we just have these several GPUs to test). We found this problem for the first time when comparing the same data input of float32 and float64, and there was a big difference in output data in one value (but not a numerical accuracy error). As shown in the following Python code, we provide a simplified case for reproducing the bug. More importantly, we dissect the CUDA code implementation in aten, we finally found that it was a BUG caused by thread synchronization. As this code shows: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/MultiLabelMarginCriterion.cu#L194-#L201 

When thread zero reads and writes to grad_input_k[target_idx], but the writing is not completed yet, the non-zero thread continues to execute asynchronously, and when threadIdx.x == target_idx, the thread reads and writes to grad_input_k[target_idx] at the same time, the conflicts occurred. This is the origin of the bug. The solution is to add the necessary thread synchronization between these.

The Aten CUDA code snippet is as follows:

```cpp
  for (int dt = 0; dt < dim; dt++) {
    ...
    if (threadIdx.x == 0) {
          if (threadIdx.x == 0) {
      grad_input_k[target_idx] += static_cast<scalar_t>(total_sum);
    }
  }

  for (int d = threadIdx.x; d < dim; d += blockDim.x) {
    grad_input_k[d] *= *grad_output_k;
  }
```

By the way, here is a google colab URL with test code, you can easily reproduce this bug with the NVIDIA T4 provided by the google colab platform:
https://colab.research.google.com/drive/170CHg9BBvoc9wiHXmHuhZ3EZjpnFPeHV?usp=sharing

A simplified case for reproducing:

```python
import torch
import numpy as np

n = 1
c = 33
index = 32

# make test case
x = np.zeros((n, c), dtype=np.float32)
y = np.zeros((n, c), dtype=np.int64) + -1
y[0][0] = index

x_fp32 = torch.tensor(x, dtype=torch.float32, requires_grad=True).cuda()
x_fp32.retain_grad()
x_fp64 = torch.tensor(x, dtype=torch.float64, requires_grad=True).cuda()
x_fp64.retain_grad()

y = torch.tensor(y, dtype=torch.int64).cuda()

# Compute with data input in float32 and float64 formats respectively.
out_fp32 = torch.nn.functional.multilabel_margin_loss(x_fp32, y)
out_fp32.backward()
out_fp64 = torch.nn.functional.multilabel_margin_loss(x_fp64, y)
out_fp64.backward()

# Print grad of input.
print('x_fp32.grad=', x_fp32.grad)  # The value of the last element in the output should be -0.9697, not 0.0, it's easy to make a check by hand.
print('x_fp64.grad=', x_fp64.grad)
```

the result tested on NVIDIA T4 with PyTorch 1.12.1:

```text
x_fp32.grad= tensor([[0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303,
         0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303,
         0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0303,
         0.0303, 0.0303, 0.0303, 0.0303, 0.0303, 0.0000]], device='cuda:0')
x_fp64.grad= tensor([[ 0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,
          0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,
          0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,
          0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,  0.0303,
         -0.9697]], device='cuda:0', dtype=torch.float64)
```

### Versions

PyTorch version: 1.12.1+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.22.6
Libc version: glibc-2.26

Python version: 3.7.13 (default, Apr 24 2022, 01:04:09)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 460.32.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.6
[pip3] torch==1.12.1+cu113
[pip3] torchaudio==0.12.1+cu113
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.13.1
[pip3] torchvision==0.13.1+cu113
[conda] Could not collect

cc @ngimel