Repository: facebookresearch/xformers

Issue:
Dropout uses the *memory address* of seeds instead of reading seeds from memory
# 🐛 Bug

From reading the code for `k_dropout_fw` and `k_dropout_bw`, it seems to me that the seeds are never read from memory and the code simply uses the *memory address* of the seed.
For [example](https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_dropout.py):
```def _get_4_bin_masks(seed, rand_offsets, p):
    rand1, rand2, rand3, rand4 = tl.randint4x(seed.to(tl.int32), rand_offsets)
```
Here seed is still a memory address and not an integer.

As a result, when `k_dropout_fw` is passed in two identical seed tensors with different memory addresses, it produces different results.  

## To Reproduce

Setting the Pytorch seed should produce the same `seed` used in dropout, and should produce the same dropout mask.
However, that's not the case

```
import torch
from xformers.triton.dropout import dropout

x = torch.randn(3, 5, device='cuda')
print(x)

torch.manual_seed(0)
torch.cuda.manual_seed(0)
print(dropout(x, 0.5))

torch.manual_seed(0)
torch.cuda.manual_seed(0)
print(dropout(x, 0.5))
```

```
 tensor([[ 0.4821, -1.6949,  0.8196,  1.9093, -1.0018],
        [ 0.4030, -1.5175, -0.3187, -0.0959,  2.7204],
        [ 1.0645, -0.1254,  0.3978, -2.9882,  0.2232]], device='cuda:0')

tensor([[ 0.9642,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.8059, -3.0350,  0.0000, -0.1918,  5.4409],
        [ 0.0000, -0.2507,  0.7955,  0.0000,  0.0000]], device='cuda:0')

tensor([[ 0.9642, -3.3897,  0.0000,  3.8186, -2.0037],
        [ 0.0000,  0.0000, -0.6374,  0.0000,  5.4409],
        [ 2.1290, -0.2507,  0.7955,  0.0000,  0.4464]], device='cuda:0')
```


- PyTorch Version (e.g., 1.0): 1.10.1
- OS (e.g., Linux): Ubuntu 20.04
- How you installed PyTorch (`conda`, `pip`, source): conda
- Build command you used (if compiling from source): pip install -e . (from master)
- Python version: 3.8
- CUDA/cuDNN version: 11.3
- GPU models and configuration: V100

Dropout uses the *memory address* of seeds instead of reading seeds from memory
# 🐛 Bug

From reading the code for `k_dropout_fw` and `k_dropout_bw`, it seems to me that the seeds are never read from memory and the code simply uses the *memory address* of the seed.
For [example](https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_dropout.py):
```def _get_4_bin_masks(seed, rand_offsets, p):
    rand1, rand2, rand3, rand4 = tl.randint4x(seed.to(tl.int32), rand_offsets)
```
Here seed is still a memory address and not an integer.

As a result, when `k_dropout_fw` is passed in two identical seed tensors with different memory addresses, it produces different results.  

## To Reproduce

Setting the Pytorch seed should produce the same `seed` used in dropout, and should produce the same dropout mask.
However, that's not the case

```
import torch
from xformers.triton.dropout import dropout

x = torch.randn(3, 5, device='cuda')
print(x)

torch.manual_seed(0)
torch.cuda.manual_seed(0)
print(dropout(x, 0.5))

torch.manual_seed(0)
torch.cuda.manual_seed(0)
print(dropout(x, 0.5))
```

```
 tensor([[ 0.4821, -1.6949,  0.8196,  1.9093, -1.0018],
        [ 0.4030, -1.5175, -0.3187, -0.0959,  2.7204],
        [ 1.0645, -0.1254,  0.3978, -2.9882,  0.2232]], device='cuda:0')

tensor([[ 0.9642,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.8059, -3.0350,  0.0000, -0.1918,  5.4409],
        [ 0.0000, -0.2507,  0.7955,  0.0000,  0.0000]], device='cuda:0')

tensor([[ 0.9642, -3.3897,  0.0000,  3.8186, -2.0037],
        [ 0.0000,  0.0000, -0.6374,  0.0000,  5.4409],
        [ 2.1290, -0.2507,  0.7955,  0.0000,  0.4464]], device='cuda:0')
```


- PyTorch Version (e.g., 1.0): 1.10.1
- OS (e.g., Linux): Ubuntu 20.04
- How you installed PyTorch (`conda`, `pip`, source): conda
- Build command you used (if compiling from source): pip install -e . (from master)
- Python version: 3.8
- CUDA/cuDNN version: 11.3
- GPU models and configuration: V100



Discussion:
oh just seeing that, sorry for the delay !
I thought that I had a unit test for that, but actually that was not the case ([see](https://github.com/facebookresearch/xformers/blob/main/tests/test_triton_dropout.py#L93)). You're right, thanks for the report @tridao, it should be an easy fix  
oh just seeing that, sorry for the delay !
I thought that I had a unit test for that, but actually that was not the case ([see](https://github.com/facebookresearch/xformers/blob/main/tests/test_triton_dropout.py#L93)). You're right, thanks for the report @tridao, it should be an easy fix  

Resolving patch:
diff --git a/xformers/triton/k_dropout.py b/xformers/triton/k_dropout.py
--- a/xformers/triton/k_dropout.py
+++ b/xformers/triton/k_dropout.py
@@ -21,7 +21,8 @@
 
 @triton.jit
 def _get_4_bin_masks(seed, rand_offsets, p):
-    rand1, rand2, rand3, rand4 = tl.randint4x(seed.to(tl.int32), rand_offsets)
+    seed = tl.load(seed)
+    rand1, rand2, rand3, rand4 = tl.randint4x(seed, rand_offsets)
 
     # binarize masks, save registers
     # NOTE: We keep the random numbers as is there (integers over int32),


Test patch:
diff --git a/tests/test_triton_dropout.py b/tests/test_triton_dropout.py
--- a/tests/test_triton_dropout.py
+++ b/tests/test_triton_dropout.py
@@ -113,6 +113,17 @@ def test_dropout(shape, amp, bias, p):
         drop_p = (y.numel() - y.count_nonzero()) / y.numel()
         assert abs(drop_p - p) < 0.01
 
+        # Check that the same seeds lead to the same dropout
+        torch.manual_seed(0)
+        torch.cuda.manual_seed(0)
+        y_1 = triton_dropout(x, p=0.5)
+
+        torch.manual_seed(0)
+        torch.cuda.manual_seed(0)
+        y_2 = triton_dropout(x, p=0.5)
+
+        assert torch.allclose(y_1, y_2)
+
 
 @pytest.mark.skipif(not _triton_available, reason="Triton is not available")
 @pytest.mark.skipif(