どんな問題が発生したか

HuggingFaceのTrainerクラスを使用してマルチGPU環境でモデルを学習させようとした際、以下のエラーに遭遇することがあります。

RuntimeError: module must have its parameters on device cuda:0 (device_ids[0]) but found one of them on device: cuda:1

発生環境の一例：

GPU: 2台以上のCUDA対応GPU
フレームワーク: PyTorch + HuggingFace Transformers
使用クラス: Trainer
現象: model_init関数内でモデルを初期化し、CUDAデバイス上手動配置後からTrainerを実行

このエラーは、Trainerが期待するデバイスと、実際にモデルが配置されているデバイスが一致しない場合に発生します。特にmodel_init関数内で独自の初期化を行い、手動でモデルを特定のGPU（cuda:1など）に移動させた場合に起こりやすい問題です。

結論

Trainerのdevice引数とmodel_initで返すモデルのデバイスを一致させるか、model_initではなくprepare_modelメソッドでデバイスマッピングを制御することで解決できます。最も簡単な方法は、Trainerを使用せず стандартный PyTorchの学習ループで устройство を明示的に管理することです。

具体的な手順

手順1: device引数を確認する

まず、Trainerのdevice引数を確認します。デフォルトではcuda:0が使用されるため、モデルをcuda:0に配置する必要があります。

from transformers import Trainer, TrainingArguments

# 明示的にdeviceを指定（推奨）
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=4,
    device=0,  # cuda:0を指定
    # 他の引数...
)

trainer = Trainer(
    model_init=model_init,
    args=training_args,
    # 他の引数...
)
trainer.train()

手順2: model_init関数でモデルを正しいデバイスに配置する

model_init関数内でモデルを初期化した後、正しいデバイスに配置します。

import torch
from transformers import AutoModelForSequenceClassification

def model_init():
    model = AutoModelForSequenceClassification.from_pretrained(
        'bert-base-uncased',
        num_labels=2
    )
    # 明示的にcuda:0に移動
    model = model.to('cuda:0')
    return model

# Trainerの初期化
trainer = Trainer(
    model_init=model_init,
    args=training_args,
    train_dataset=train_dataset,
)
trainer.train()

手順3: マルチGPUの場合はDeepSpeedまたはFSDPを使用する

複数のGPUを効率的に使用する場合は、DeepSpeedまたはFSDP（Fully Sharded Data Parallel）を使用します。

# DeepSpeed設定例 (deepspeed_config.json)
{
    "fp16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 1,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}

# トレーニング実行
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    deepspeed='deepspeed_config.json',
)
trainer = Trainer(
    model_init=model_init,
    args=training_args,
)
trainer.train()

手順4: 代替手段 – PyTorch標準ループを使用する

もしTrainerでの解決が困難な場合は、標準的なPyTorchの学習ループを使用することで、より柔軟なデバイス管理が可能になります。

import torch
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification

# デバイスの設定
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# モデルの初期化とデバイス配置
model = AutoModelForSequenceClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=2
)
model = model.to(device)

# オプティマイザーの設定
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)

# 学習ループ
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)

model.train()
for epoch in range(3):
    for batch in train_loader:
        # バッチをデバイスに移動
        batch = {k: v.to(device) for k, v in batch.items()}
        
        # 順伝播
        outputs = model(**batch)
        loss = outputs.loss
        
        # 逆伝播
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
print("Training completed!")

補足・注意点

バージョン依存の注意

Transformers 4.20以降では、Trainerのデバイス管理が改善されていますが、過去のバージョンでは異なる挙動を示す場合があります
PyTorchのバージョンも影響を与えるため、torch.__version__とtransformers.__version__を確認してください

環境別の落とし穴

単一GPU環境: 特に意識する必要はありません默认でcuda:0が使用されます
マルチGPU環境（DataParallel）: 自動的にデバイスが割り当てられる場合がありますが、意図しないデバイス配置が行われることがあります
マルチGPU環境（Distributed）: local_rankを正しく設定し、torch.distributedを使用する場合はDEVICE変数を適切に扱う必要があります
クラウド環境（Colab/AWS）: GPUの認識順序が異なる場合があるため、torch.cuda.device_count()で確認してください

よく行われるミス

model_init関数内でモデルを初期化するが、to(device)を呼ばない
DataLoaderのbatchをデバイスに移動し忘れる
損失計算後に.loss.detach().to(device)を行わず、テンソルが別のデバイスに残る
torch.cuda.set_device()を呼ばずに、複数のGPUがある環境で特定のGPUを指定してしまう

デバッグ技巧

問題が発生した場合は、以下のコマンドでデバイスの状態を確認できます。

# 現在のデバイス確認
import torch
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Device count: {torch.cuda.device_count()}")
print(f"Current device: {torch.cuda.current_device()}")

# モデルパラメータのデバイス確認
for name, param in model.named_parameters():
    print(f"{name}: {param.device}")
    break  # 最初の1つだけでOK

【HuggingFace Trainer】CUDAデバイス不一致エラーの解決法