Self-LLM项目多卡微调GLM-4模型时的设备一致性解决方案

问题背景

在使用Self-LLM项目进行GLM-4模型微调时，当尝试使用多张GPU进行训练时，系统会报出设备不一致的错误。具体表现为模型计算损失函数时，部分张量分布在不同的GPU设备上（如cuda:3和cuda:1），导致无法正常完成交叉熵损失计算。

错误分析

该问题的根本原因是模型在分布式训练环境下，不同部分的计算被自动分配到了不同的GPU设备上，但损失计算时未能确保所有参与计算的张量位于同一设备。这种情况在使用device_map="auto"参数时尤为常见，因为该参数会自动优化模型各层在不同设备间的分布。

解决方案

经过深入分析，发现需要在计算损失函数前显式确保所有张量位于同一设备。具体修改方案如下：

1. 定位到模型文件中的损失计算部分（通常在modeling_chatglm.py文件中，约1020行附近）

2. 原始代码如下：

if labels is not None:
    lm_logits = lm_logits.to(torch.float32)
    shift_logits = lm_logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    loss_fct = CrossEntropyLoss(ignore_index=-100)
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    lm_logits = lm_logits.to(hidden_states.dtype)
    loss = loss.to(hidden_states.dtype)

3. 修改后的代码：

if labels is not None:
    lm_logits = lm_logits.to(torch.float32)
    shift_logits = lm_logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    
    # 新增设备同步代码
    device = shift_labels.device
    shift_logits = shift_logits.to(device)
    shift_labels = shift_labels.to(device)
    
    loss_fct = CrossEntropyLoss(ignore_index=-100)
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    lm_logits = lm_logits.to(hidden_states.dtype)
    loss = loss.to(hidden_states.dtype)

技术原理

这种修改确保了在计算交叉熵损失前：

1. 获取标签张量所在的设备
2. 将预测logits显式移动到同一设备
3. 确保标签张量也在同一设备

这种处理方式在多GPU环境下尤为重要，因为：

• 模型的不同层可能被自动分配到不同GPU上
• 损失计算需要所有输入张量位于同一设备
• 显式设备同步可以避免隐式的跨设备数据传输

其他注意事项

1. 虽然使用device_map="cuda"可以暂时规避这个问题，但这会限制模型只能使用单卡，无法发挥多GPU的并行计算优势。

2. 对于大规模模型训练，推荐使用这种显式设备同步的方式，因为它：

   • 保持多卡并行能力
   • 明确控制张量位置
   • 提高代码可读性和可维护性

3. 在实际应用中，还可以考虑使用PyTorch的分布式训练工具（如DistributedDataParallel）来更好地管理多GPU训练过程。

总结

通过这种简单的代码修改，我们成功解决了Self-LLM项目中GLM-4模型多卡微调时的设备不一致问题。这种解决方案不仅适用于GLM-4模型，对于其他需要多GPU训练的大语言模型也具有参考价值，是进行分布式模型训练时的一个实用技巧。