Training with 2 GPUs, without deepspeed causes tensor GPU mismatch error

[sunildkumar]

I need to fix this, but unfortunately it isn't my top priority right now.

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument tensors in method wrapper_CUDA_cat)
[rank0]: Traceback (most recent call last):
[rank0]:   File "/millcreek/home/sunil/r1_vlm/src/r1_vlm/environments/digit_recognition_env/train.py", line 90, in <module>
[rank0]:     trainer.train()
[rank0]:   File "/millcreek/home/sunil/r1_vlm/.venv/lib/python3.12/site-packages/transformers/trainer.py", line 2241, in train
[rank0]:     return inner_training_loop(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/millcreek/home/sunil/r1_vlm/.venv/lib/python3.12/site-packages/transformers/trainer.py", line 2548, in _inner_training_loop
[rank0]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank0]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/millcreek/home/sunil/r1_vlm/.venv/lib/python3.12/site-packages/transformers/trainer.py", line 3692, in training_step
[rank0]:     inputs = self._prepare_inputs(inputs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/millcreek/home/sunil/trl/trl/trainer/qwen_grpo_trainer.py", line 650, in _prepare_inputs
[rank0]:     prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1)
[rank0]:                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument tensors in method wrapper_CUDA_cat)

[Anirudh257]

Hi, I was able to solve this issue by arbitrarily adding many ".to(device)" steps in trl repository. You can view the changes here: https://github.com/Anirudh257/r1_vlm_fixed_gpus_2

[sunildkumar]

Awesome! I'll try to take a look when I have time. However, if you have time to write a PR against our TRL fork, that would be amazing. I'm pretty occupied on building multi turn tool use.

[sunildkumar]

@EivindKjosbakken - my email says you replied to this saying it was working for you (just needed to change which GPU ran vllm). Is this right?

[EivindKjosbakken]

Hi, yeah I first left the comment saying it worked with vllm_device="cuda:0". When you do this, you don't get the device errors as previously, but I then get out of memory errors which I wasn't able to solve (I am running on 2 A100 80 Gb), so I removed the comment until I was able to figure out the memory issue

[sunildkumar]

I can help with that! Some questions:

1. Which dataset are you training with?
2. Can you post your training config here please?
3. How large are your images?

[EivindKjosbakken]

Thank you!

I am trying to run Task 2: Digit Recognition by running the command:
CUDA_VISIBLE_DEVICES=0,1 uv run src/r1_vlm/environments/digit_recognition_env/train.py

I am using the training  config in GRPO/r1_vlm/src/r1_vlm/environments/digit_recognition_env/train.py, but I lowered the batch size and number generations to 2:

import os

from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
from trl import GRPOConfig, ModelConfig
from trl.trainer.qwen_grpo_trainer import QwenGRPOTrainer

from r1_vlm.environments.digit_recognition_env.digit_recognition_env import (
    DigitRecognitionEnv,
)

os.environ["WANDB_ENTITY"] = "groundlightai"
os.environ["WANDB_PROJECT"] = "digit-recognition"


vf_env = DigitRecognitionEnv()
dataset = vf_env.get_dataset()
rubric = vf_env.get_rubric()


# Flag that determines if gradient checkpointing is used. If it is, we need to set use_cache to False.
gradient_checkpointing = False


model_config = ModelConfig(
    model_name_or_path="Qwen/Qwen2.5-VL-3B-Instruct",
    torch_dtype="bfloat16",
    use_peft=False,
)

model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    pretrained_model_name_or_path=model_config.model_name_or_path,
    torch_dtype=model_config.torch_dtype,
    use_cache=False,
)

# use cache if not gradient checkpointing
if gradient_checkpointing:
    model.config.use_cache = False
elif not gradient_checkpointing:
    model.config.use_cache = True
else:
    raise ValueError("Invalid gradient checkpointing value")


processor = AutoProcessor.from_pretrained(
    model_config.model_name_or_path, padding_side="left"
)

training_args = GRPOConfig(
    model_init_kwargs=model_config,
    output_dir="vlm-r1-digit-recognition",
    learning_rate=1e-6,
    adam_beta2=0.98,
    lr_scheduler_type="cosine",
    warmup_steps=0,
    logging_steps=1,
    save_steps=20,
    save_total_limit=50,
    num_train_epochs=1,
    per_device_train_batch_size=2,
    num_generations=2,
    gradient_accumulation_steps=1,
    gradient_checkpointing=gradient_checkpointing,
    bf16=True,
    # GRPO specific parameters
    max_prompt_length=None,  # must be None for vllm + verifiers
    max_completion_length=128,
    beta=0.001,
    temperature=1.0,
    sync_ref_model=True,
    ref_model_sync_steps=64,
    eval_strategy="no",
    log_completions=True,
    use_vllm=True,
    vllm_gpu_memory_utilization=0.95,
    report_to="wandb",
    vllm_device="cuda:1",
)


trainer = QwenGRPOTrainer(
    model=model,
    processing_class=processor,
    reward_funcs=rubric,
    args=training_args,
    train_dataset=dataset,
    env=vf_env,
)

trainer.train()

I wanted to reduce the batch size to 1 to save memory, but was not allowed to do that since I would have no allowed values for "num_generations"

3. I am using the dataset in https://huggingface.co/datasets/sunildkumar/digit-recognition-r1 so they are 100x100.

[sunildkumar]

thanks for this info. One more question - which branch of our TRL fork are you on?

[EivindKjosbakken]

I'm on the "release_2025_03_06" branch (I followed the README, so everything will be according to that)

[sunildkumar]

Got it. Nothing you shared seems wrong. I'm waiting for some time on our A100 cluster to repro and I'll get back to you then.

[lsbuschoff]

Hey I'm experiencing a similar error. I was initially running the 4 GPU setup but ran into NCCL watchdog timeouts which I couldn't fix. I tried the 2 GPU setup just now and it produces the following trace (note that this is using my own environment which just swaps out the dataset for lsbuschoff/interact_single_r1 with its own reward function; this runs on the 4 GPU setup for the first model completion until some process gets hung up). Do you have any idea how to solve this device error for the 2 GPU setting? I can't see the exact changes made to TRL in #42 (comment).

Traceback (most recent call last):
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/src/r1_vlm/environments/block_stacking_env/train_2.py", line 89, in <module>
    trainer.train()
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/transformers/trainer.py", line 2241, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/transformers/trainer.py", line 2548, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/transformers/trainer.py", line 3692, in training_step
    inputs = self._prepare_inputs(inputs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/trl/trl/trainer/qwen_grpo_trainer.py", line 671, in _prepare_inputs
    ref_per_token_logps = self._get_per_token_logps(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/trl/trl/trainer/qwen_grpo_trainer.py", line 525, in _get_per_token_logps
    logits = model(
             ^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/accelerate/utils/operations.py", line 819, in forward
    return model_forward(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/accelerate/utils/operations.py", line 807, in __call__
    return convert_to_fp32(self.model_forward(*args, **kwargs))
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py", line 1795, in forward
    image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py", line 518, in forward
    hidden_states = self.patch_embed(hidden_states)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py", line 111, in forward
    hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/nn/modules/conv.py", line 725, in forward
    return self._conv_forward(input, self.weight, self.bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/hkfs/home/haicore/hgf_hmgu/hgf_roz4189/interact/r1_vlm/.venv/lib/python3.12/site-packages/torch/nn/modules/conv.py", line 720, in _conv_forward
    return F.conv3d(
           ^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument weight in method wrapper_CUDA__cudnn_convolution)

[hongge831]

@Anirudh257 Hi, have you solved the device mismatch issues? I didn't see any related modification of your link, could you please double check you link, and post your link again?

[hongge831]

@sunildkumar Hi, have you solved the device mismatch issues?

[Anirudh257]

Hi I just added many to(device) in https://github.com/Anirudh257/r1_vlm_fixed_gpus_2/trl/