PyTorch/TensorRT 中保存大模型时 pickle_protocol 问题的解决方案

在深度学习模型部署过程中，我们经常需要将优化后的模型保存到磁盘。然而，当处理大型模型时，比如 CLIP 这样的视觉语言模型，可能会遇到模型序列化的问题。本文将详细介绍在 PyTorch/TensorRT 项目中保存大型模型时遇到的 pickle 协议限制问题及其解决方案。

问题背景

当使用 torch_tensorrt.save 方法保存优化后的大型模型时，如果模型大小超过 4GB，会遇到序列化错误。这是因为 Python 的 pickle 模块在默认协议版本（protocol 2）下无法处理超过 4GB 的数据块。错误信息通常会显示："serializing a string larger than 4 GiB requires pickle protocol 4 or higher"。

技术分析

PyTorch 的序列化机制底层依赖于 pickle 模块。在 torch.save 方法中，已经提供了 pickle_protocol 参数来支持更高版本的协议。然而，torch_tensorrt.save 方法内部使用的是 torch.export.save，这个接口最初并没有暴露 pickle_protocol 参数。

具体的技术栈调用关系如下：

1. torch_tensorrt.save 调用 torch.export.save
2. torch.export.save 调用 serialize_torch_artifact
3. serialize_torch_artifact 内部使用 torch.save

解决方案

经过社区与 PyTorch 核心团队的协作，这个问题已经在最新版本中得到解决。现在用户可以直接在 torch_tensorrt.save 方法中指定 pickle_protocol 参数：

torch_tensorrt.save(model, "large_model.ep", inputs=[input_tensor], pickle_protocol=4)

对于使用较旧版本的用户，可以暂时使用以下变通方案：

exported = torch_tensorrt.dynamo._exporter.export(compiled_model)
torch.export.save(exported, save_path, pickle_protocol=4)

最佳实践

1. 对于大型模型，总是使用 pickle_protocol=4 或更高版本
2. 确保使用的 PyTorch 和 TensorRT 版本是最新的
3. 在保存模型前，先进行小批量数据的前向传播以确保模型正确性
4. 考虑模型量化等优化技术来减小模型体积

总结

随着深度学习模型规模的不断增大，序列化大模型已成为部署流程中的常见需求。PyTorch/TensorRT 生态通过不断完善其序列化机制，为开发者提供了更强大的工具支持。了解并正确使用 pickle_protocol 参数，可以避免在模型保存阶段遇到不必要的障碍，使大型模型的部署更加顺畅。