Attention bias and Query/Key/Value should be on the same device

[calledice]

using GRPO

[rank0]: ValueError: Error in model execution (input dumped to /tmp/err_execute_model_input_20250211-033220.pkl): Attention bias and Query/Key/Value should be on the same device
[rank0]: query.device: cuda:2
[rank0]: attn_bias : cuda:0

[ashenwitch]

same here

[anirudhb11]

I'm also facing a similar issue here, I'm running on 2 GPUs (1 for training and 1 for inference). My configuration is as follows:

PyTorch Version: 2.5.1+cu124
CUDA Available: True
CUDA Version: 12.4
GPU Details: [{'GPU Name': 'Quadro RTX 8000', 'Total Memory (MB)': 45539.0, 'CUDA Capability': '7.5'}, {'GPU Name': 'Quadro RTX 8000', 'Total Memory (MB)': 45539.0, 'CUDA Capability': '7.5'}]

Just to highlight I'm using a RTX 8000 which doesn't support bfloat16 - so I'm not sure if this is causing the issue. To enable running in float16 I change the config file to include --vllm-dtype float16 and change the torch_dtype in the config file to float16

This is the traceback I see;

[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/src/open_r1/grpo.py", line 279, in <module>
[rank0]:     main(script_args, training_args, model_args)
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/src/open_r1/grpo.py", line 214, in main
[rank0]:     trainer = GRPOTrainer(
[rank0]:               ^^^^^^^^^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/trl/trainer/grpo_trainer.py", line 404, in __init__
[rank0]:     self.llm = LLM(
[rank0]:                ^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/utils.py", line 1039, in inner
[rank0]:     return fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/entrypoints/llm.py", line 240, in __init__
[rank0]:     self.llm_engine = self.engine_class.from_engine_args(
[rank0]:                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/engine/llm_engine.py", line 482, in from_engine_args
[rank0]:     engine = cls(
[rank0]:              ^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/engine/llm_engine.py", line 274, in __init__
[rank0]:     self._initialize_kv_caches()
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/engine/llm_engine.py", line 414, in _initialize_kv_caches
[rank0]:     self.model_executor.determine_num_available_blocks())
[rank0]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/executor/executor_base.py", line 99, in determine_num_available_blocks
[rank0]:     results = self.collective_rpc("determine_num_available_blocks")
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/executor/uniproc_executor.py", line 49, in collective_rpc
[rank0]:     answer = run_method(self.driver_worker, method, args, kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/utils.py", line 2208, in run_method
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/worker/worker.py", line 228, in determine_num_available_blocks
[rank0]:     self.model_runner.profile_run()
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/worker/model_runner.py", line 1236, in profile_run
[rank0]:     self._dummy_run(max_num_batched_tokens, max_num_seqs)
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/worker/model_runner.py", line 1347, in _dummy_run
[rank0]:     self.execute_model(model_input, kv_caches, intermediate_tensors)
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/mila/b/buvanesa/code/open-r1/openr1/lib/python3.11/site-packages/vllm/worker/model_runner_base.py", line 152, in _wrapper
[rank0]:     raise type(err)(
[rank0]: ValueError: Error in model execution (input dumped to /tmp/err_execute_model_input_20250211-144354.pkl): Attention bias and Query/Key/Value should be on the same device
[rank0]:   query.device: cuda:1
[rank0]:   attn_bias   : cuda:0

[rank0]:[W211 14:43:54.574278762 ProcessGroupNCCL.cpp:1250] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present,  but this warning has only been added since PyTorch 2.4 (function operator())

[PascalPolygon]

Same problem

[calledice]

same here

in config_demo.yaml you can set
vllm_device: cuda:0

[byhteng]

this problem can be resolved in this way:

1. find vllm/attention/backends/xformers.py in our conda env
2. at line 733, change to attn_bias=attn_bias[0].to(query.device),
   this problem seems to result from the initialization af attention bias without specifying the device of the vllm process, but using the device of the main process