在RNN模型上应用DataParallel时触发Assertion Error问题求助

数据老陈

2026-5-20

这种情况确实有点棘手——单GPU跑完全正常，一用DataParallel就报设备不匹配，而且你已经确认输入都是cuda类型，那大概率是DataParallel的特殊机制或者模型内部的细节没处理到位。我踩过不少类似的坑，给你几个针对性的排查和解决方向：

1. 先确认模型所有参数都在GPU上

有时候我们以为模型已经移到GPU了，但可能某些自定义子模块没被正确注册到nn.Module里，或者DataParallel的包装顺序错了。你可以在模型初始化后，打印所有参数的设备，排查漏网之鱼：

for name, param in model.named_parameters():
    print(f"{name}: {param.device}")

如果发现有参数在CPU上，那问题就出在这。正确的模型初始化+DataParallel姿势应该是：

# 方式1：先移GPU再包装
model = MyRNNModel(your_model_args).cuda()
model = nn.DataParallel(model)

# 方式2：一步到位
model = nn.DataParallel(MyRNNModel(your_model_args)).cuda()

注意：如果你的模型有自定义子模块，一定要让它继承nn.Module，并且用self.submodule = SubModule()的方式添加到主模型中，否则DataParallel不会处理这些子模块的参数，它们会留在CPU上。

2. RNN的Hidden状态是重灾区

RNN的hidden状态（尤其是LSTM的(h,c)元组）在多GPU环境下特别容易出问题，常见的两个坑：

Hidden的batch维度不匹配：单GPU时，hidden的形状是(num_layers*num_directions, batch_size, hidden_size)，但DataParallel会把输入batch拆分成N份（N是GPU数量），每个GPU处理batch_size/N的数据。如果你的hidden还是保持原batch_size的大小，就会和拆分后的输入维度不匹配，这时候报错信息可能会误导你说是设备问题，而不是维度问题。
解决方法：确保hidden是作为forward的输入传入（别存在模型的成员变量里），这样DataParallel会自动帮你拆分hidden到各个GPU。如果是训练循环中需要维护的hidden，每次forward后记得用hidden = tuple(h.detach() for h in hidden)分离梯度，并且确认它的设备正确。
Hidden元组里混了CPU tensor：如果hidden是元组（比如LSTM的h和c），一定要确认两个tensor都在GPU上，不能一个在GPU一个在CPU。你可以在forward开头加个打印排查：

def forward(self, input_batch, input_batch_length, hidden):
    if isinstance(hidden, tuple):
        print(f"h device: {hidden[0].device}, c device: {hidden[1].device}")
    else:
        print(f"hidden device: {hidden.device}")
    # 后续业务代码

3. 检查模型内部创建的临时Tensor

有时候我们在forward里会手动创建一些tensor（比如torch.zeros()、torch.ones()），如果没指定设备，这些tensor默认会在CPU上，和GPU上的输入运算就会触发设备不匹配错误。解决方法是用模型的设备或者输入的设备来创建：

# 不要直接写torch.zeros(shape)
temp_tensor = torch.zeros(shape, device=self.device)
# 或者跟着输入的设备走
temp_tensor = torch.zeros(shape, device=input_batch.device)

4. 留意input_batch_length的处理

如果input_batch_length是一个普通列表，一般没问题（pack_padded_sequence可以接受CPU上的长度列表），但如果它是tensor，一定要确保它在GPU上。如果你的模型里有把它转成tensor的操作，记得指定设备：

input_batch_length = torch.tensor(input_batch_length, device=input_batch.device)

你可以先从检查模型参数和hidden状态的设备开始，这两个是最常见的原因。如果还是解决不了，可以把forward函数的完整代码贴出来，这样能更精准定位问题。

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