Device error on 0.0.27.dev844

[Cospui]

🐛 Bug

[rank1]: ValueError: Attention bias and Query/Key/Value should be on the same device
[rank1]:   query.device: cuda:1
[rank1]:   attn_bias   : cuda:0

Command

To Reproduce

Steps to reproduce the behavior:

out = xformers.ops.fmha.memory_efficient_attention(q, k, v, attn_bias=attn_bias)

Code works fine on version 0.0.27.dev840 and 0.0.26.post1. After upgrading to new version, the code fails with the above error.

Expected behavior

Environment

Please copy and paste the output from the
environment collection script from PyTorch
(or fill out the checklist below manually).

You can run the script with:

# For security purposes, please check the contents of collect_env.py before running it.
python -m torch.utils.collect_env

• PyTorch Version (e.g., 1.0):
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, source):
• Build command you used (if compiling from source):
• Python version:
• CUDA/cuDNN version:
• GPU models and configuration:
• Any other relevant information:

Additional context

[danthe3rd]

Hi,
Thanks for reporting this bug. Which attention bias are you using? (type(attn_bias))
There are 2 solutions to unblock you immediately:
(1) Set the correct cuda device for each rank (eg torch.set_default_device(rank))
(2) Or transfer the bias to the right device (eg attn_bias = attn_bias.to(query.device)) - but this is not backward compatible

[Cospui]

Thanks for your reply. As for the attention bias, I use

attn_bias, x = xformers.ops.fmha.BlockDiagonalMask.from_tensor_list(x)
attn_bias = attn_bias.make_causal()

I set model.to(device) in each rank, but I did not set device for attn_bias. Maybe attn_bias changes the default device to cuda:0? maybe I can set the default coda device. but currently, I just downgraded xformers to 0.0.26.post1 for now, and the error disappeared.

[danthe3rd]

maybe I can set the default coda device

Yes this would most likely solve your issue

[shizhediao]

Same error. Can I just use 0.0.26.post1 as a workaround?

[Navanit-git]

Same error still. Is there a way to set device in xformers??

[ChristopheYe]

Same error, I tried using "0.0.26.post1" but then vllm library doesn't work anymore, is there any solution to this please ?

[function2-llx]

I think device should be specified here according to the input tensor list:

xformers/xformers/ops/fmha/attn_bias.py

Line 760 in 916611d

block_diag = cls.from_seqlens(seqlens)

Or it will use the default CUDA device, which may fail when multiple CUDA devices are available:

xformers/xformers/ops/fmha/attn_bias.py

Lines 90 to 94 in 916611d

	def _get_default_bias_device(device: Optional[torch.device] = None) -> torch.device:
	if device is None:
	if torch.cuda.is_available():
	return torch.device("cuda")
	return torch.device("cpu")

[francescotaioli]

I got the same issue using trl==0.15.0.dev0, xformers==0.0.28.post3 and vllm==0.7.1