ChatGLM3多卡微调中的Tensor JSON序列化问题解决方案

问题背景

在使用ChatGLM3进行多卡微调训练时，许多开发者遇到了一个常见的技术问题：当训练过程中尝试保存检查点(checkpoint)时，系统会抛出"TypeError: Object of type Tensor is not JSON serializable"的错误。这个问题通常出现在使用多GPU进行LoRA微调的场景下，特别是在第一个检查点保存完成后继续训练时。

错误现象分析

该问题的典型表现是：

1. 训练过程可以正常启动并运行初始阶段
2. 在达到第一个检查点保存步数时，系统能够成功保存临时检查点文件(如tmp-checkpoint-500)
3. 但在保存完成后继续训练时，程序会崩溃并报错
4. 错误信息明确指出Tensor对象无法被JSON序列化

根本原因

经过技术分析，发现这个问题与DeepSpeed库的版本兼容性有关。在较新版本的DeepSpeed(如0.14.0)中，某些内部数据结构处理方式发生了变化，导致在多卡训练环境下保存模型状态时，尝试将Tensor对象直接序列化为JSON格式时失败。

解决方案

解决这个问题的有效方法是降级DeepSpeed到0.13.1版本。具体操作步骤如下：

1. 首先卸载当前安装的DeepSpeed版本：

   pip uninstall deepspeed
   

2. 安装指定版本的DeepSpeed：

   pip install deepspeed==0.13.1
   

3. 确保其他相关依赖也符合要求：

   • transformers >= 4.36.2
   • torch >= 2.1.1
   • peft >= 0.6.2

完整微调命令

在正确配置环境后，可以使用以下命令启动多卡微调：

OMP_NUM_THREADS=1 torchrun --standalone --nnodes=1 --nproc_per_node=2 finetune_hf.py data/fix/ /path/to/chatglm3-6b configs/lora_multiple.yaml

重要提示：需要确保在lora配置文件中(deepspeed部分)已经取消注释，以便正确启用DeepSpeed优化。

技术建议

1. 版本控制：在深度学习项目中，特别是涉及多卡训练时，保持各组件版本的兼容性非常重要。建议使用虚拟环境管理不同项目的依赖。

2. 检查点验证：在训练开始前，可以尝试手动保存一个检查点来验证序列化功能是否正常。

3. 日志监控：训练过程中密切关注日志输出，特别是在接近检查点保存步数时的系统状态。

4. 硬件适配：虽然本文案例使用的是3090显卡，但解决方案同样适用于其他NVIDIA显卡，如4090等。

总结

多卡微调是训练大语言模型的重要手段，而版本兼容性问题常常是阻碍训练顺利进行的绊脚石。通过合理控制DeepSpeed版本，开发者可以避免Tensor JSON序列化问题，确保ChatGLM3模型在多卡环境下稳定训练。这一经验也提醒我们，在深度学习工程实践中，组件版本管理是需要特别关注的技术细节。