TensorRT引擎缓存中哈希计算性能问题的分析与优化

在PyTorch/TensorRT项目中，当使用引擎缓存功能时，开发人员发现重用缓存的引擎比从头开始训练一个不可重构的引擎还要慢。经过深入分析，发现问题出在FxGraphCachePickler.get_hash(new_gm)这个哈希计算函数上，它占据了总编译时间的很大一部分。

问题背景

在TensorRT引擎缓存机制中，为了能够重用之前编译好的引擎，系统需要为当前的模型计算一个唯一的哈希值作为缓存键。这个哈希值是通过FxGraphCachePickler.get_hash(new_gm)函数生成的，它首先将模型序列化，然后计算序列化数据的SHA256哈希。

性能瓶颈分析

通过实际测试Llama2-7b模型，发现哈希计算阶段消耗了惊人的时间：

• 在启用缓存的REFIT模式下，第一次运行（需要构建引擎并保存到缓存）耗时约727660ms，其中哈希计算就占了323009ms
• 后续运行（直接重用缓存的引擎）耗时约410300ms，哈希计算仍然占了323275ms
• 相比之下，从头编译一个不可重构的引擎仅需约267709ms

进一步分析发现，FxGraphCachePickler.get_hash(new_gm)函数中，serialized_data = cls.dumps(obj)这一行序列化操作消耗了绝大部分时间，而实际的SHA256哈希计算耗时很少。

问题根源

经过与PyTorch团队的沟通，确认这个性能问题的根源在于Python的pickle序列化机制。对于大型模型，pickle序列化会非常耗时，而PyTorch团队表示这是pickle本身的限制，无法在他们的层面进行优化。

解决方案

为了解决这个问题，开发团队设计并实现了一个新的哈希函数，完全绕过了pickle序列化步骤。这个新方案通过直接处理模型的关键特征来生成哈希值，避免了昂贵的序列化操作。

新方案的主要特点包括：

1. 直接提取模型图中的关键信息（如操作类型、参数形状等）
2. 使用更高效的哈希计算方法
3. 保持与原有哈希相同的唯一性保证

实施效果

通过PR #3293提交的这个优化，显著提升了引擎缓存重用的性能：

• 哈希计算时间从原来的300多秒降低到几乎可以忽略不计
• 重用缓存的引擎现在比从头编译要快得多
• 整体编译时间大幅缩短，特别是在需要频繁重用缓存的场景下

这个优化对于大型模型特别有益，使得TensorRT的引擎缓存功能真正发挥出了其应有的价值，为模型部署和推理提供了显著的性能提升。