PyTorch TensorRT 中 FP8 精度支持的技术解析

概述

在深度学习模型部署领域，PyTorch TensorRT 作为 NVIDIA 推出的高性能推理引擎，一直致力于提供更高效的模型优化方案。近期，关于 FP8（8位浮点数）精度的支持问题引起了开发者社区的关注。

FP8 精度支持现状

FP8 数据类型是 NVIDIA 近年来推出的新型计算格式，旨在进一步降低计算和存储开销，同时保持合理的模型精度。然而，在 PyTorch TensorRT 24.06 版本中，开发者发现尝试使用 torch.float8_e4m3fn 数据类型时会遇到错误提示："Provided an unsupported data type as an input data type (support: bool, int32, long, half, float), got: torch.float8_e4m3fn"。

技术背景

FP8 格式有两种主要变体：

• E4M3（4位指数，3位尾数）
• E5M2（5位指数，2位尾数）

这种格式特别适合新一代 GPU 架构，能够在保持合理模型精度的同时显著提升计算效率和降低内存占用。

解决方案与建议

根据官方反馈，FP8 支持将在后续版本中提供。开发者可以采取以下策略：

1. 对于当前项目，可以考虑使用已支持的 FP16 或 INT8 量化方案作为临时解决方案
2. 关注 PyTorch TensorRT 24.07 版本的发布，该版本将正式包含 FP8 支持
3. 对于急于测试 FP8 功能的开发者，可以尝试使用 2.3 版本或 nightly 构建版本

技术展望

FP8 支持的引入将为深度学习推理带来新的优化空间：

• 更低的显存占用，支持更大 batch size
• 更高的计算吞吐量
• 适用于新一代 GPU 架构的优化

开发者应当密切关注 PyTorch TensorRT 的版本更新，及时获取最新的功能支持，以充分利用硬件加速能力提升模型推理性能。