Repository: facebookresearch/xformers

Issue:
[microGPT] Crash on sampling "sample not on the right device"
# 🐛 Bug
Maybe something which changed in lightning, but the sampling in the end of microGPT fails on the sample not being on the right device now

## Command
`python3 example/microGPT.py`

## Additional context
Callstack
```
Traceback (most recent call last):
  File "/home/lefaudeux/Git/xformers/examples/microGPT.py", line 325, in <module>
    y = sample(model, x, steps=1000, temperature=1.0, sample=True, top_k=10)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "/home/lefaudeux/Git/xformers/examples/microGPT.py", line 246, in sample
    logits = model(x_cond)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/examples/microGPT.py", line 158, in forward
    prediction = self.model(src)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/factory/model_factory.py", line 159, in forward
    memory = encoder(memory, input_mask=encoder_input_mask)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/factory/block_factory.py", line 315, in forward
    x = self.wrap_ff(x)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/components/residual.py", line 48, in forward
    return inputs[0] + self.layer(*inputs, *args, **kwargs)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/components/residual.py", line 69, in forward
    return self.sublayer(*x_norm, *args, **kwargs)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/components/feedforward/fused_mlp.py", line 64, in forward
    return self.mlp(inputs)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/container.py", line 141, in forward
    input = module(input)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/triton/dropout.py", line 156, in forward
    return _dropout.apply(x, p, self.bias, self.activation, self.activation_grad)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/cuda/amp/autocast_mode.py", line 94, in decorate_fwd
    return fwd(*args, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/triton/dropout.py", line 48, in forward
    k_dropout_fw[grid](
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/triton/code_gen.py", line 676, in __call__
    return self.kernel(*wargs, **kwargs, grid=self.grid)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/triton/code_gen.py", line 724, in __call__
    return self.kernel(*args, num_warps=config.num_warps, num_stages=config.num_stages, **meta, **config.meta)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/triton/code_gen.py", line 583, in __call__
    raise ValueError("Arguments at index {invalid_args} are on the wrong device.".format(invalid_args=invalid_args) +
ValueError: Arguments at index [0, 1, 2, 3] are on the wrong device. Only CUDA is supported at the moment`
```
[microGPT] Crash on sampling "sample not on the right device"
# 🐛 Bug
Maybe something which changed in lightning, but the sampling in the end of microGPT fails on the sample not being on the right device now

## Command
`python3 example/microGPT.py`

## Additional context
Callstack
```
Traceback (most recent call last):
  File "/home/lefaudeux/Git/xformers/examples/microGPT.py", line 325, in <module>
    y = sample(model, x, steps=1000, temperature=1.0, sample=True, top_k=10)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "/home/lefaudeux/Git/xformers/examples/microGPT.py", line 246, in sample
    logits = model(x_cond)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/examples/microGPT.py", line 158, in forward
    prediction = self.model(src)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/factory/model_factory.py", line 159, in forward
    memory = encoder(memory, input_mask=encoder_input_mask)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/factory/block_factory.py", line 315, in forward
    x = self.wrap_ff(x)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/components/residual.py", line 48, in forward
    return inputs[0] + self.layer(*inputs, *args, **kwargs)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/components/residual.py", line 69, in forward
    return self.sublayer(*x_norm, *args, **kwargs)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/components/feedforward/fused_mlp.py", line 64, in forward
    return self.mlp(inputs)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/container.py", line 141, in forward
    input = module(input)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/triton/dropout.py", line 156, in forward
    return _dropout.apply(x, p, self.bias, self.activation, self.activation_grad)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/torch/cuda/amp/autocast_mode.py", line 94, in decorate_fwd
    return fwd(*args, **kwargs)
  File "/home/lefaudeux/Git/xformers/xformers/triton/dropout.py", line 48, in forward
    k_dropout_fw[grid](
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/triton/code_gen.py", line 676, in __call__
    return self.kernel(*wargs, **kwargs, grid=self.grid)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/triton/code_gen.py", line 724, in __call__
    return self.kernel(*args, num_warps=config.num_warps, num_stages=config.num_stages, **meta, **config.meta)
  File "/home/lefaudeux/.conda/envs/xformers_dev/lib/python3.9/site-packages/triton/code_gen.py", line 583, in __call__
    raise ValueError("Arguments at index {invalid_args} are on the wrong device.".format(invalid_args=invalid_args) +
ValueError: Arguments at index [0, 1, 2, 3] are on the wrong device. Only CUDA is supported at the moment`
```


Discussion:

Resolving patch:
diff --git a/xformers/triton/dropout.py b/xformers/triton/dropout.py
--- a/xformers/triton/dropout.py
+++ b/xformers/triton/dropout.py
@@ -13,7 +13,7 @@
 import triton
 from torch.cuda.amp import custom_bwd, custom_fwd
 
-from xformers.components.activations import Activation
+from xformers.components.activations import Activation, build_activation
 from xformers.triton.k_activations import (
     get_triton_activation_bwd_kernel,
     get_triton_activation_kernel,
@@ -140,7 +140,7 @@ def __init__(
     ) -> None:
         super().__init__()
         self.p = p
-        self.activation = activation
+        self.activation_type = activation
         self.register_buffer(
             "bias", torch.zeros(bias_shape) if bias_shape is not None else None
         )
@@ -152,5 +152,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.bias is not None:  # type: ignore
             self.bias = self.bias.to(dtype=x.dtype, device=x.device)  # type: ignore
 
+        # Catch a non-cuda setup, fallback to pytorch
+        if not x.is_cuda:
+            activation = build_activation(self.activation_type)
+            x = x + self.bias if self.bias is not None else x
+            x = activation(x)
+            return torch.nn.functional.dropout(x, self.p)
+
         p = self.p if self.training else 0.0
         return _dropout.apply(x, p, self.bias, self.activation, self.activation_grad)


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
@@ -11,6 +11,7 @@
 from torch.cuda.amp.autocast_mode import autocast
 
 from xformers.components import Activation, build_activation
+from xformers.triton.dropout import FusedDropoutBias
 
 _triton_available = torch.cuda.is_available()
 
@@ -38,6 +39,12 @@
 ]
 
 
+def test_dropout_cpu():
+    triton_dropout = FusedDropoutBias(p=0.1, bias_shape=None)
+    x = torch.normal(0, 1, size=(16, 16), device="cpu")
+    _ = triton_dropout(x)
+
+
 @pytest.mark.skipif(not _triton_available, reason="Triton is not available")
 @pytest.mark.skipif(
     not _triton_available or gpu_capabilities_older_than_70(),