PyTorch Lightning 2.4.0多设备输入与模型设备不匹配问题

问题背景 #

使用PyTorch-Lightning 2.4.0进行多GPU训练，lmm是封装HuggingFace模型与处理器的nn.Module子类，核心代码如下：

class ICVModel(pl.LightningModule):
    def __init__(self, lmm, icv_encoder: torch.nn.Module) -> None:
        super().__init__()
        self.lmm = lmm
        self.lmm.requires_grad_(False)
        self.icv_encoder = icv_encoder
        self.eos_token = self.lmm.processor.tokenizer.eos_token

    def forward(self, ice_texts, query_texts, answers, images):
        query_answer = [
            query + answer + self.eos_token
            for query, answer in zip(query_texts, answers)
        ]
        query_images = [img[-setting.num_image_in_query :] for img in images]
        query_inputs = self.lmm.process_input(query_answer, query_images)
        query_outputs = self.lmm.model(
            **query_inputs,
            labels=query_inputs["input_ids"],
        )

调用lmm.model(...)时触发设备不匹配错误。

设备信息排查 #

• lmm.model.forward外部打印（对应inputs.pixel_values.device、self.device、self.lmm.device）：

rank[0]: cpu cuda:0 cuda:0
rank[1]: cpu cuda:1 cuda:1

• Idefics模型内部forward打印（对应inputs.pixel_values.device、self.device）：

rank[0]: cuda:0 cuda:0
rank[1]: cuda:0 cuda:1

手动将pixel_values移至正确设备后，后续前向传播中仍会被移至错误设备，最终报错：

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)

可能原因 #

1. lmm.process_input方法未将处理后的图像tensor分配到对应GPU，导致初始pixel_values停留在CPU；
2. Idefics模型内部存在硬编码的设备指定（如强制使用cuda:0），未适配多GPU分布式场景；
3. PyTorch-Lightning的自动设备分配逻辑与模型内部的设备操作冲突。

解决方案 #

1. 修正lmm.process_input的设备分配 #

在process_input方法末尾，将所有输出tensor统一移至模型所在设备：

class LMMWrapper(nn.Module):
    # 其他代码...
    def process_input(self, texts, images):
        # 原预处理逻辑
        inputs = self.processor(texts, images, return_tensors="pt", padding=True)
        # 新增：将所有tensor移至模型所在设备
        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
        return inputs

2. 在ICVModel的forward中强制同步设备 #

传入模型前，将所有输入tensor强制移至当前LightningModule的设备：

def forward(self, ice_texts, query_texts, answers, images):
    # 原逻辑...
    query_inputs = self.lmm.process_input(query_answer, query_images)
    # 强制将输入tensor移至当前设备
    query_inputs = {k: v.to(self.device) for k, v in query_inputs.items()}
    query_outputs = self.lmm.model(**query_inputs, labels=query_inputs["input_ids"])

3. 检查Idefics模型内部的硬编码设备 #

如果是自定义的Idefics子类，排查forward或初始化逻辑中是否存在硬编码的cuda:0，替换为self.device：

# 错误示例
x = x.to("cuda:0")
# 修正为
x = x.to(self.device)

4. 确认分布式训练配置正确性 #

确保PyTorch-Lightning的Trainer使用正确的分布式策略，例如：

trainer = pl.Trainer(
    accelerator="gpu",
    devices=2,
    strategy="ddp",
    # 其他配置...
)

初始化时验证模型设备分配：

def __init__(self, lmm, icv_encoder: torch.nn.Module) -> None:
    super().__init__()
    self.lmm = lmm
    print(f"Rank {self.global_rank}: LMM model device: {self.lmm.model.device}")
    # 其他代码...

内容的提问来源于stack exchange，提问作者Kamichanw