TensorRT中LayerNorm精度设置问题分析与解决方案

问题背景

在使用NVIDIA TensorRT v8.6.11进行模型转换时，用户遇到了一个关于LayerNorm层精度控制的问题。具体表现为：当尝试将LayerNorm层强制设置为FP32精度时，TensorRT优化器将这些层与其他操作一起封装到Myelin层中，导致最终精度仍为FP16。

问题现象

用户在使用trtexec工具转换ONNX模型时，通过以下参数指定LayerNorm层使用FP32精度：

--layerPrecisions=LayerNormalization_*:fp32
--layerOutputTypes=LayerNormalization_*:fp32

然而，转换后的引擎文件中，LayerNorm层被包含在Myelin层中，实际运行精度仍为FP16。这种现象在A100 GPU上尤为明显，导致模型精度下降超过20%。

技术分析

1. Myelin层特性：Myelin是TensorRT内部的一个优化层，它会将多个操作融合为一个高效的执行单元。这种融合虽然能提高性能，但会覆盖用户指定的精度设置。

2. 精度控制机制：TensorRT的精度控制主要通过三种方式实现：

   • 全局精度模式（如--fp16）
   • 层级别精度设置（--layerPrecisions）
   • 输出类型设置（--layerOutputTypes）

3. 通配符支持问题：TensorRT可能不完全支持使用通配符(*)来匹配层名，这可能导致精度设置未能正确应用。

解决方案

1. 明确指定层名： 避免使用通配符，直接从日志中获取完整的LayerNorm层名称，然后精确指定：

   --layerPrecisions=LayerNormalization_123:fp32
   

2. 输出层设置技巧： 由于TensorRT要求输出层必须为FP32，可以将关键层设置为输出层来间接强制其使用FP32精度：

   --outputs=Add_3244
   

3. 使用自定义插件： 对于关键操作如LayerNorm，可以开发或使用现有的FP32精度插件来替代默认实现。

4. 模型结构调整： 在导出ONNX模型时，将需要FP32精度的层设置为模型输出，确保它们在转换过程中保持所需精度。

最佳实践建议

1. 在转换前仔细检查模型结构，识别所有需要特殊精度处理的层。

2. 使用--verbose参数获取详细的转换日志，确认每层的实际精度设置。

3. 对于Transformer类模型，特别注意LayerNorm、Softmax等对精度敏感的操作。

4. 在性能与精度之间权衡，仅对确实需要FP32的层进行特殊设置。

5. 考虑使用TensorRT的精度分析工具来验证各层的实际运行精度。

总结

TensorRT的自动优化机制虽然强大，但有时会与用户的精度需求产生冲突。通过理解TensorRT的内部工作机制，采用精确指定层名、利用输出层特性等方法，可以有效控制关键操作的精度，在保持性能的同时确保模型精度不受损失。对于特别敏感的操作，开发自定义插件是最可靠的解决方案。