Linux服务器SentenceTransformer训练时nvmlDeviceGetNvLinkRemoteDeviceType报错

问题背景 #

我正在编写Python代码训练分类器，对包含10个句子的样本进行分类。该分类器基于SentenceTransformer框架，并集成了GASCOM项目的自定义额外层，部署在Linux服务器上运行训练。

核心训练代码 #

# 原代码块内容

触发的RuntimeError栈 #

# 原错误栈代码块内容

可能原因及解决方法 #

1. 张量维度不匹配 #

自定义层的输出维度与SentenceTransformer的输出/下游分类头的输入维度不兼容。比如SentenceTransformer输出的句子嵌入为768维，但自定义层处理后维度改变，而分类头仍按768维初始化，导致运算时维度不匹配。

解决方法：

• 打印SentenceTransformer的输出形状：print(model.encode("test").shape)，再打印自定义层处理后的张量形状，确认两者与后续分类头的输入维度一致。
• 若维度不一致，修改自定义层的输出逻辑（如调整线性层的out_features参数），或重新初始化分类头以匹配新的输入维度。

2. 自定义层前向传播逻辑错误 #

GASCOM自定义层的forward方法中存在张量形状变换错误、未处理batch维度等问题，导致输出张量形状不符合后续层的要求。

解决方法：

• 在自定义层的forward方法中添加形状打印语句，定位异常步骤：

  class GASCOM_CustomLayer(nn.Module):
      def forward(self, x):
          print("Input to custom layer:", x.shape)
          # 原有操作
          x = self.transform(x)
          print("Output from custom layer:", x.shape)
          return x
  

• 根据打印结果修正形状变换逻辑，比如确保batch维度被正确保留，避免误删或修改。

3. 输入数据形状不匹配 #

训练数据经过编码后的形状与模型（尤其是自定义层）的输入要求不一致，比如batch输入的形状应为(batch_size, hidden_dim)，但实际传入的是(batch_size, seq_len, hidden_dim)。

解决方法：

• 检查数据加载后的张量形状，确认每个样本的嵌入形状符合自定义层的输入要求。
• 若需要，在数据预处理阶段添加维度压缩或变换步骤，比如用x = x.mean(dim=1)将序列维度压缩为句子嵌入。

4. 设备不匹配 #

模型部分层部署在CPU，部分在GPU（若使用GPU训练），导致张量跨设备运算时出错。

解决方法：

• 模型初始化后，统一移至目标设备：model.to(device)（device为torch.device("cuda")或torch.device("cpu")）。
• 确保输入数据也移至对应设备：inputs = {k: v.to(device) for k, v in inputs.items()}。

5. 损失函数与输出不匹配 #

若分类器为多分类任务，损失函数的输入格式不符合要求，比如CrossEntropyLoss要求输入为(batch_size, num_classes)，但模型输出为其他形状。

解决方法：

• 检查损失函数的输入形状，确保模型输出的维度与损失函数的要求一致。
• 若输出维度不符，在模型最后添加线性层调整至num_classes维度。

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