GLM-4V多卡推理实现与设备一致性问题的解决方案

背景介绍

GLM-4V作为新一代多模态大模型，在实际部署时经常需要利用多GPU进行推理加速。然而，开发者在尝试多卡部署时遇到了设备不一致的问题，表现为"Expected all tensors to be on the same device"错误。本文将深入分析问题原因并提供完整的解决方案。

问题现象分析

在多GPU环境下运行GLM-4V时，模型会抛出设备不一致的错误，具体表现为：

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:2 and cuda:3!

这种错误通常发生在视觉模块(visual.py)中的Transformer层前向传播过程中，特别是在残差连接操作时。根本原因是模型的不同部分被分配到了不同的GPU设备上，而PyTorch要求参与运算的张量必须位于同一设备上。

技术原理探究

在多卡并行推理场景下，GLM-4V的视觉模块存在以下技术特点：

1. 模型并行性：模型的不同层可能被分配到不同的GPU设备上
2. 残差连接：Transformer架构中的残差连接要求输入和输出必须在同一设备上
3. 特殊标记处理：视觉模块中的BOI(开始标记)和EOI(结束标记)需要与主计算流保持设备一致

解决方案实现

针对上述问题，我们通过对visual.py文件进行两处关键修改实现了多卡推理支持：

1. Transformer层设备一致性修复

在TransformerLayer类的forward方法中，我们显式确保MLP输出与输入位于同一设备：

class TransformerLayer(nn.Module):
    def forward(self, hidden_states):
        # ...原有代码...
        mlp_output = self.post_attention_layernorm(self.mlp(mlp_input))
        mlp_output = mlp_output.to(mlp_input.device)  # 新增设备同步
        output = mlp_input + mlp_output
        return output

2. 视觉标记设备同步

在视觉模块的前向传播中，确保BOI和EOI标记与输入张量设备一致：

def forward(self, images: "tensor(B, C, H, W)") -> "tensor(B, L, D)":
    # ...原有代码...
    boi = self.boi.expand(x.shape[0], -1, -1).to(x.device)  # 设备同步
    eoi = self.eoi.expand(x.shape[0], -1, -1).to(x.device)  # 设备同步
    x = torch.cat((boi, x, eoi), dim=1)

部署建议

1. 硬件要求：每张GPU建议至少有16GB显存
2. 软件版本：
   • transformers ≥ 4.42.4
   • torch ≥ 2.3.1
3. 最佳实践：对于6卡及以上配置，建议检查模型拆分是否合理，部分模块可能不适合跨卡拆分

总结

通过显式管理张量设备位置，我们成功解决了GLM-4V在多卡环境下的设备不一致问题。这一解决方案已在官方代码库中合并，为大规模部署GLM-4V提供了可靠的技术支持。开发者现在可以充分利用多GPU的计算能力，实现高效的视觉-语言多模态推理。