用于实时推理的基于 Transformer 的神经网络的 GPU 加速优化

本文介绍了使用 NVIDIA TensorRT 进行混合精度优化的 Transformer 模型的 GPU 加速推理管道的设计和评估。我们对 BERT-base（110M 参数）和 GPT-2（124M 参数）进行了评估，批量大小为 1 到 32，序列长度为 32 到 512。该系统的 CPU 基线速度提高了 64.4 倍，单样本推理延迟低于 10 毫秒，内存使用量减少了 63%。我们引入了一种混合精度策略，该策略保留 FP32 用于数字敏感操作，例如 softmax 和层归一化，同时将 FP16 应用于线性层。这种方法保持了高数值保真度（相对于基线输出的余弦相似度 >= 0.9998）并消除了 NaN 的不稳定性。该管道作为模块化、容器化系统实现，可在 360 多种配置中实现可重复的基准测试。 NVIDIA A100 上的跨 GPU 验证显示 FP16 加速比在 1.84 倍到 2.00 倍之间保持一致，并且数值行为稳定。 SST-2 的下游评估表明在混合精度下精度没有下降。 WikiText-2 上的验证表明，对于完整的 FP16，随机输入低估了 NaN 不稳定性高达 6 倍，同时确认了混合方法的稳健性（0.0% NaN，余弦相似度 >= 0.9998）。这些结果提供了跨 GPU 架构的性能和准确性权衡的详细表征，并为在延迟关键环境中部署 Transformer 模型提供了实用指导。