TensorRT项目中关于权重剥离引擎在torch.compile下的问题分析

问题背景

在深度学习模型部署过程中，TensorRT提供了权重剥离引擎(weight-stripped engine)的功能，这一特性允许开发者将模型权重从引擎文件中分离出来，从而减小引擎文件体积并支持运行时权重更新。然而，在PyTorch/TensorRT项目中，开发者发现这一功能在通过torch.compile()接口使用时出现了异常。

三种编译方式的差异

PyTorch/TensorRT提供了三种主要方式来编译模型：

1. torch export + torch_trt.dynamo.compile：通过导出模型后使用TensorRT的动态编译接口
2. torch_trt.compile(ir="dynamo")：直接使用TensorRT的Dynamo后端编译
3. torch.compile(backend="tensorrt")：使用PyTorch原生编译接口指定TensorRT后端

前两种方式能够正确支持权重剥离功能，但第三种方式却出现了问题。测试表明，当使用torch.compile()时，即使设置了strip_engine_weights=True，生成的引擎仍然包含了权重信息，导致输出结果与预期不符。

技术分析

通过深入分析引擎属性，我们发现：

• 使用torch.compile()时，无论是否启用权重剥离，引擎大小几乎相同
• 而使用torch_trt.compile()时，启用权重剥离后引擎大小显著减小
• torch.compile()生成的引擎层数明显少于torch_trt.compile()

进一步研究发现，torch.compile()的内部处理流程与其它两种方式不同：

1. 它调用了pretraced_backend()而非直接调用编译函数
2. 在最新的PyTorch版本中，所有模型权重都被注册为输入参数
3. 这种改变导致权重剥离机制无法按预期工作

解决方案

鉴于当前的技术限制，项目团队决定：

1. 在使用torch.compile()时发出警告，提示权重剥离功能不受支持
2. 在这种情况下返回包含权重的编译后模块
3. 建议需要权重剥离功能的用户使用前两种编译方式

对开发者的建议

对于需要使用权重剥离引擎的场景，建议：

1. 优先考虑使用torch_trt.dynamo.compile或torch_trt.compile(ir="dynamo")
2. 如果必须使用torch.compile()，需要了解其权重处理方式的差异
3. 关注项目更新，未来可能会提供统一的解决方案

这一问题的发现和解决过程展示了深度学习编译工具链中不同接口间的微妙差异，也提醒开发者在选择编译方式时需要充分考虑功能需求和技术限制。