OpenRLHF项目中Zero3 SFT训练模型加载问题分析与解决方案

问题背景

在OpenRLHF项目中进行Zero3模式的监督式微调(SFT)训练后，用户遇到了模型无法正常加载的问题。具体表现为尝试加载训练好的PeftModel时出现大量参数形状不匹配的错误，特别是lora_B.default.weight参数形状从预期的torch.Size([14336, 8])变成了torch.Size([0])。

错误现象分析

当使用PeftModel.from_pretrained()方法加载模型时，系统报出RuntimeError，显示多个层的mlp模块中的lora_B.default.weight参数形状不匹配。错误信息表明：

1. 检查点(checkpoint)中的参数形状为torch.Size([0])（空张量）
2. 当前模型期望的形状为torch.Size([14336, 8])

这种错误在模型的所有32个层中重复出现，涉及gate_proj、up_proj和down_proj三个投影层的LoRA参数。

根本原因

经过项目维护者的分析，这个问题主要由两个因素导致：

1. safetensors文件保存问题：在Zero3模式下保存的模型文件可能包含空张量
2. 字典键名不一致：保存的模型状态字典中的键名与加载时期望的键名不匹配

特别值得注意的是，这个问题主要出现在Zero3模式下，而在Zero2模式下通常不会出现类似问题。这是因为Zero3采用了更复杂的分片策略，对模型参数的存储和加载提出了更高要求。

解决方案

项目维护者已经在主分支中修复了这个问题。修复方案主要涉及以下几个方面：

1. 保存逻辑优化：修正了模型保存时的参数处理流程，确保所有必要参数都被正确保存
2. 键名统一：标准化了状态字典中的键名，保证保存和加载时的一致性
3. LoRA参数处理：特别处理了LoRA相关参数的保存机制

对于遇到类似问题的用户，建议：

1. 更新到最新版本的OpenRLHF代码
2. 在保存模型前确保所有参数都已正确聚合
3. 检查模型保存和加载时的配置一致性

技术启示

这个问题揭示了分布式训练中模型保存和加载的几个重要原则：

1. 分布式策略的影响：不同的分布式策略(Zero2/Zero3)对模型序列化有不同要求
2. 参数聚合时机：在分布式环境下，参数的保存需要特别考虑聚合状态
3. 兼容性设计：模型保存格式需要同时考虑训练和推理场景的需求

通过这个案例，我们可以更好地理解大规模语言模型训练中参数管理的复杂性，以及在不同并行策略下确保模型可重现性的重要性。