PEFT项目中DeepSpeed Zero3与LoRA训练中embedding层保存问题分析

问题背景

在PEFT(Parameter-Efficient Fine-Tuning)项目的最新版本中，当使用DeepSpeed Zero3优化策略进行LoRA(Low-Rank Adaptation)训练时，开发人员发现一个关键问题：当配置modules_to_save参数来保存embedding层(embed_tokens)和语言模型头部(lm_head)时，这些层在保存的适配器模型文件中变成了空张量，而非预期的实际权重值。

问题现象

具体表现为：

1. 训练过程可以正常完成，没有报错
2. 但在尝试合并模型时，会出现张量形状不匹配的错误
3. 检查保存的adapter_model.safetensors文件，发现embed_tokens和lm_head对应的张量确实存在，但形状为0，即空张量

技术分析

根本原因

通过代码调试发现，这一问题源于PEFT项目中get_peft_model_state_dict函数的变更。在之前的版本中，该函数包含专门处理modules_to_save的逻辑，能够正确保存这些模块的权重。但在最新版本中，这部分代码被移除，替换为使用AuxiliaryTrainingWrapper的处理逻辑。

DeepSpeed Zero3的影响

DeepSpeed Zero3是一种内存优化技术，它会将模型参数、梯度和优化器状态分区到不同的GPU上。这种分布式策略可能导致：

1. 参数在不同设备间的分布
2. 需要特殊的处理来收集和保存完整的参数
3. 对modules_to_save这类特殊参数的处理可能出现边缘情况

代码变更对比

旧版本处理方式：

if getattr(model, "modules_to_save", None) is not None:
    for key, value in state_dict.items():
        if any(f"{module_name}.modules_to_save.{adapter_name}" in key for module_name in model.modules_to_save):
            to_return[key.replace("modules_to_save.", "")] = value

新版本处理方式：

for name, module in model.named_modules():
    if isinstance(module, AuxiliaryTrainingWrapper):
        to_return.update({f"{name}.{k}": v for k, v in module.adapter_state_dict(adapter_name).items()})

解决方案

临时解决方案是恢复旧版本中处理modules_to_save的代码逻辑。但更完善的解决方案应该考虑：

1. 确保与DeepSpeed Zero3的兼容性
2. 统一参数保存的处理方式
3. 添加针对此情况的测试用例
4. 可能需要修改AuxiliaryTrainingWrapper的实现，使其能够正确处理modules_to_save参数

最佳实践建议

对于使用PEFT+DeepSpeed Zero3进行模型训练的用户，建议：

1. 在配置modules_to_save时进行验证性测试
2. 检查保存的适配器模型文件，确认关键参数是否正确保存
3. 关注PEFT项目的更新，及时获取官方修复
4. 在关键训练前，先进行小规模测试验证整个流程

总结

这一问题揭示了参数高效微调技术与分布式训练框架集成时的复杂性。特别是在处理需要特殊保存的模块时，需要考虑框架间的交互和参数分发机制。随着PEFT技术的普及，这类集成问题将越来越受到关注，需要开发者和用户共同注意。