LLaMA-Factory项目中训练GLM-4-9B-0414模型时的梯度检查点问题分析

在LLaMA-Factory项目中使用GLM-4-9B-0414模型进行全参数微调训练时，开发者可能会遇到一个与梯度检查点相关的错误。这个问题主要出现在使用DeepSpeed进行分布式训练的场景下，错误信息表明在尝试访问梯度检查点函数的__self__属性时失败。

问题现象

当开发者尝试使用LLaMA-Factory框架训练GLM-4-9B-0414模型时，训练过程会在初始化阶段抛出AttributeError异常，提示'functools.partial' object has no attribute '__self__'。这个错误发生在梯度检查点功能的实现部分，具体是在custom_gradient_checkpointing_func函数中尝试访问函数对象的__self__属性时。

技术背景

梯度检查点(Gradient Checkpointing)是一种用于减少深度学习模型训练时显存占用的技术。它通过在前向传播过程中不保存所有中间激活值，而是在反向传播时重新计算部分激活值，从而以计算时间为代价换取显存空间的节省。

在PyTorch中，梯度检查点通常通过torch.utils.checkpoint.checkpoint函数实现。LLaMA-Factory项目为了支持更灵活的配置，实现了自定义的梯度检查点功能。

问题原因分析

这个问题的根本原因在于DeepSpeed框架对模型前向传播函数的封装方式。DeepSpeed在封装模型时使用了functools.partial来创建部分应用函数，这导致原本期望的函数对象变成了一个部分应用函数对象，而部分应用函数没有__self__属性。

具体来说：

1. 原始代码假设传入的函数是一个绑定方法(bound method)，具有__self__属性指向所属的模块实例
2. 但DeepSpeed的封装使得传入的函数变成了functools.partial对象
3. 当代码尝试访问__self__属性时，由于部分应用函数没有这个属性而抛出异常

解决方案

这个问题已经在LLaMA-Factory的最新代码中得到修复。开发者可以通过以下方式解决：

1. 更新LLaMA-Factory到最新版本
2. 如果无法立即更新，可以临时修改checkpointing.py文件中的相关代码，使用更健壮的方式来获取模块实例

修复后的代码应该能够正确处理DeepSpeed封装后的函数对象，确保梯度检查点功能正常工作。

最佳实践建议

在使用LLaMA-Factory训练大型语言模型时，建议：

1. 保持框架和依赖库的最新版本
2. 在使用DeepSpeed等分布式训练框架时，注意检查梯度检查点功能的兼容性
3. 对于GLM系列模型，确保使用适配的transformers版本(如4.51.3)
4. 在遇到类似问题时，检查错误堆栈中涉及的框架封装层次

这个问题展示了深度学习框架在多层封装时可能出现的兼容性问题，也提醒开发者在实现高级功能时需要考虑到框架可能对基础对象进行的修改。