Fix facebookresearch#173 (facebookresearch#175)

CHANGELOG.md

@@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## TBD
 ### Fixed
 - Much faster fused dropout [#164]
+- Fused dropout repeatability [#173]
 
 ### Added
 - Embedding weight tying option [#172]

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(

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),