自定义Transformer分类器部署CUDA时设备不匹配报错排查

可能的原因 #

• 模型内部非参数类张量（如位置编码、注意力掩码）未迁移到CUDA：从零实现Transformer时，常手动生成位置编码或掩码，这类张量若默认在CPU创建，前向传播时会与CUDA上的模型参数、输入张量冲突。
• DataLoader多进程加载导致张量留存在CPU：当num_workers>0时，子进程加载的数据可能未正确迁移到CUDA，或子进程内创建的张量未处理。
• 模型自定义层的缓冲区未迁移：部分自定义层会使用缓冲区（buffer）而非可训练参数，若初始化时未注册或未移至目标设备，会导致设备不匹配。
• 前向传播中动态生成的张量未指定设备：在forward函数内新建的张量（如临时计算的掩码）默认在CPU，未同步到模型/输入的设备。

修复方案 #

• 迁移所有静态张量：模型初始化时将位置编码等静态张量移至目标设备，或在forward函数中基于输入张量的设备动态创建：

  def forward(self, input_ids, labels=None):
      # 获取输入所在设备
      device = input_ids.device
      # 基于设备动态生成注意力掩码
      attn_mask = torch.ones(input_ids.shape, device=device)
      # 后续计算逻辑
  

• 排查DataLoader设置：先将num_workers设为0验证是否是多进程问题，若需保留多进程，在collate_fn中统一将所有张量移至目标设备：

  def collate_fn(batch):
      input_ids = torch.stack([item['input_ids'] for item in batch]).to('cuda')
      labels = torch.tensor([item['labels'] for item in batch]).to('cuda')
      return {'input_ids': input_ids, 'labels': labels}
  

• 检查并迁移模型缓冲区：通过model.buffers()查看所有缓冲区，确保它们已同步到CUDA；若自定义层使用缓冲区，需用self.register_buffer()注册，保证调用model.to(device)时自动迁移。
• 显式指定所有张量的设备：创建任何张量时都通过device参数指定目标设备，避免依赖默认CPU设备。
• 打印设备信息定位问题：在forward函数开头添加设备打印，快速定位异常张量：

  def forward(self, input_ids, labels=None):
      print("input_ids 设备:", input_ids.device)
      print("模型参数设备:", next(self.parameters()).device)
      print("位置编码设备:", self.pos_emb.device)
      # 后续逻辑
  

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