TensorRT 10.7 中自注意力机制精度问题分析与解决方案

问题背景

在将OpenCLIP视觉模型转换为TensorRT格式时，开发者遇到了严重的精度损失问题。该问题特别出现在使用TensorRT 10.7版本处理包含自注意力机制的ViT-SO400M-14-SigLIP-384模型时。即使在FP32精度下，相对误差的P90值高达6.266，远超出可接受范围。

问题现象

当使用标准转换流程将OpenCLIP模型导出为ONNX并转换为TensorRT引擎时，出现了以下关键现象：

1. 在FP32模式下，输出张量的相对误差显著偏高
2. 使用FP8模式时误差消失（尽管硬件不支持FP8）
3. 通过日志分析发现，FP32模式下TensorRT识别并优化了自注意力模式
4. 当使用自定义插件替换Softmax层阻止自注意力优化时，精度恢复正常

根本原因分析

经过深入调查，发现问题根源在于OpenCLIP模型中使用的torch.nn.scaled_dot_product_attention函数。TensorRT对该函数的处理存在精度问题，特别是在进行自注意力模式识别和优化时。

具体表现为：

1. TensorRT的自注意力融合优化在某些情况下会引入数值不稳定性
2. 标准的缩放点积注意力计算在TensorRT优化路径下产生了较大误差
3. 当禁用这种优化（如使用FP8或自定义插件）时，精度恢复正常

解决方案

目前有两种可行的解决方案：

临时解决方案

修改模型代码，用简化的注意力实现替换标准scaled_dot_product_attention：

import math
import torch

def naive_sdpa(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
    _, _B, _Nt, E = q.shape
    q = q * math.sqrt(1.0/float(E))
    attn = q @ k.transpose(-2,-1)
    if attn_mask is not None:
        attn += attn_mask
    attn = attn.softmax(dim=-1)
    return (attn @ v)

torch.nn.functional.scaled_dot_product_attention = naive_sdpa

这种方法可以绕过TensorRT的问题路径，但可能牺牲部分性能优化。

长期解决方案

等待NVIDIA官方修复TensorRT中对scaled_dot_product_attention的处理。该问题已被确认为TensorRT的已知问题，正在跟踪解决中。

技术建议

对于遇到类似问题的开发者，建议：

1. 首先验证问题是否确实由自注意力优化引起
2. 考虑使用--noTF32标志提高计算精度（会有轻微性能损失）
3. 在关键应用场景中，考虑使用临时解决方案确保精度
4. 关注TensorRT的后续版本更新，及时获取官方修复

总结

TensorRT 10.7版本在处理包含自注意力机制的模型时存在已知精度问题。开发者在使用OpenCLIP等基于Transformer架构的模型时需特别注意。目前可通过简化注意力实现作为临时解决方案，长期来看应等待官方修复。这个问题凸显了在模型优化过程中平衡性能与精度的重要性，特别是在使用自动优化技术时。