Expected all tensors to be on the same device, but found at least two devices
Author: ootsuka-repos
Pull request: no
Labels: name: bug

### System Info

TRLのSFTrainerが動かない。
RTX8000*2 4090D*2 num4 gpu

pip install transformers==4.51.3 NG
pip install transformers==4.49.0 OK

torch2.5.1+cuda12.4


-error text-
Expected all tensors to be on the same device, but found at least two devices, cuda:3 and cuda:0!
  File "C:\Users\admin\Desktop\llm_sft\train\SFT.py", line 185, in main
    trainer.train(
  File "C:\Users\admin\Desktop\llm_sft\train\SFT.py", line 190, in <module>
    main()
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:3 and cuda:0!

Comments:
The version of TRL is known to be irrelevant.
I get errors from Transformers version 4.49.0 and up.
I don't know why I can't do what I used to be able to do.

Hi @ootsuka-repos, you'll need to give us some minimal code to reproduce the issue!

@Rocketknight1 
code 

```
import os
import random
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'

# WANDBログイン（キー削除済み、必要なら環境変数や手動入力で対応）
import wandb
# wandb.login(key="your_wandb_api_key_here")
wandb.login()

import json
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    DataCollatorForLanguageModeling
)
from datasets import Dataset
from huggingface_hub import login
from transformers.trainer_callback import TrainerCallback
from trl import SFTTrainer, SFTConfig
from tqdm import tqdm
from datetime import datetime

# 定数（トークンとパスはサンプル/相対パスに変更）
HF_TOKEN = "your_hf_token_here"
DATA_FILE_PATH = "dataset/data.jsonl"
MODEL_NAME = "models/magnum_latest_checkpoint"

current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
models_save_path = f"models/magnum_{current_time}"
MAX_TOKEN = 3500
MIN_TOKEN = 2000

# Hugging Faceログイン
login(token=HF_TOKEN)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=True)

def count_tokens(text):
    return len(tokenizer.encode(text, add_special_tokens=False))

def load_custom_data(file_path):
    data_with_tokens = []
    max_token_count = 0
    skipped_large_count = 0
    skipped_small_count = 0

    with open(file_path, "r", encoding="utf-8") as f:
        lines = f.readlines()

    for line in tqdm(lines, desc="Processing Data", unit="line"):
        item = json.loads(line)
        text = f"""アダルトASMRのシナリオを作成してください
タイトル:{item['title']}

タグ:{item['tags']}

シナリオ:{item['transcription']}"""
        
        token_count = count_tokens(text)
        if token_count > MAX_TOKEN:
            skipped_large_count += 1
            continue
        if token_count < MIN_TOKEN:
            skipped_small_count += 1
            continue

        max_token_count = max(max_token_count, token_count)
        data_with_tokens.append((item, token_count))

    random.shuffle(data_with_tokens)
    sorted_data = [item for item, _ in data_with_tokens][:250]

    print(f"最大トークン数: {max_token_count}")
    print(f"データ数: {len(sorted_data)}")
    print(f"スキップ（大）: {skipped_large_count}")
    print(f"スキップ（小）: {skipped_small_count}")

    filtered_file_path = "dataset/data_filtered.jsonl"
    with open(filtered_file_path, "w", encoding="utf-8") as f_out:
        for item in sorted_data:
            f_out.write(json.dumps(item, ensure_ascii=False) + "\n")
    print(f"保存先: {filtered_file_path}")

    return sorted_data

class SaveModelAndTokenizerCallback(TrainerCallback):
    def __init__(self, model, tokenizer, output_dir, save_epoch_interval=1):
        self.model = model
        self.tokenizer = tokenizer
        self.output_dir = output_dir
        self.save_epoch_interval = save_epoch_interval

    def on_epoch_end(self, args, state, control, **kwargs):
        checkpoint_dir = os.path.join(self.output_dir, f"checkpoint-{state.global_step}")
        os.makedirs(checkpoint_dir, exist_ok=True)
        self.tokenizer.save_pretrained(checkpoint_dir)

def main():
    data = load_custom_data(DATA_FILE_PATH)

    dataset = Dataset.from_dict({
        "Title": [item["title"] for item in data],
        "Tags": [item["tags"] for item in data],
        "Transcription": [item["transcription"] for item in data],
    })

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        device_map="balanced",
        use_cache=False,
        torch_dtype=torch.float16,
        attn_implementation="eager",
    )

    def formatting_prompts_func(example):
        prompt = f"""アダルトASMRのシナリオを作成してください

タイトル:{example['Title']}

タグ:{example['Tags']}"""
        infe = f"""シナリオ:{example['Transcription']}"""
        messages = [
            {"role": "user", "content": prompt},
            {"role": "assistant", "content": infe}
        ]
        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False
    )

    save_model_and_tokenizer_callback = SaveModelAndTokenizerCallback(
        model, tokenizer, models_save_path, save_epoch_interval=1
    )

    training_args = SFTConfig(
        output_dir=models_save_path,
        num_train_epochs=15,
        learning_rate=1e-5,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=1,
        max_grad_norm=1.0,
        gradient_checkpointing=True,
        logging_steps=1,
        save_strategy="epoch",
        lr_scheduler_type="cosine",
        optim="adamw_bnb_8bit",
        max_seq_length=MAX_TOKEN,
        packing=False,
        neftune_noise_alpha=5,
    )

    trainer = SFTTrainer(
        model=model,
        train_dataset=dataset,
        args=training_args,
        formatting_func=formatting_prompts_func,
        data_collator=data_collator,
        callbacks=[save_model_and_tokenizer_callback],
    )
    
    trainer.train()

if __name__ == "__main__":
    main()

```

I have a similar problem.

pip install transformers==4.51.3 Error
pip install transformers==4.48.0 OK

torch2.5.1+cuda12.4

Device configuration: 3 GPUs A100 40Gb

- Error text
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:2 and cuda:0!

@Rocketknight1 @vanpelt @tmm1 
pls

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!
  File "C:\Users\admin\Desktop\ai-model-training\train\sft_llm_train.py", line 142, in train_sft
    trainer.train()
  File "C:\Users\admin\Desktop\ai-model-training\train\sft_llm_train.py", line 145, in <module>
    train_sft()
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!

transformers 4.52.3 bugs

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

cc @sunmarc maybe? The code is quite long and I don't see an obvious source for the issue, so this might need a `git bisect` or something

Could you try to find the faulty commit using git bisect @ootsuka-repos ? Also, can you show us the entire traceback to better pinpoint the issue, thanks ! 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.