FasterTransformer: 加速Transformer模型推理的高性能库

FasterTransformer简介

FasterTransformer是由NVIDIA开发的一个开源库,旨在优化和加速Transformer模型的推理过程。它通过高度优化的CUDA实现和多种加速技术,可以显著提升Transformer模型的推理性能,尤其适用于大规模语言模型的部署。

FasterTransformer主要具有以下特点:

支持多种Transformer模型架构,包括BERT、GPT、T5等。
提供C++、TensorFlow、PyTorch等多种接口,方便集成到不同的深度学习框架中。
针对NVIDIA GPU进行了深度优化,可充分利用Tensor Core等硬件特性。
支持FP16、INT8等低精度推理,进一步提升性能。
实现了Tensor并行和流水线并行等分布式推理技术,可扩展到多GPU和多节点。

主要优化技术

FasterTransformer采用了多种优化技术来加速Transformer模型的推理:

1. 内核融合

将多个小的CUDA内核融合成一个大的内核,减少内核启动开销和数据传输。例如,将Self-Attention中的多个操作融合到一个内核中。

2. 内存优化

通过重用内存缓冲区、优化内存访问模式等方式减少内存占用和访问延迟。

3. 低精度计算

支持FP16、INT8等低精度计算,充分利用Tensor Core等硬件加速单元。

4. 并行计算

实现了模型并行和数据并行等分布式推理技术,可以扩展到多GPU和多节点。

5. 动态shape支持

支持动态输入序列长度,避免了不必要的padding开销。

性能表现

根据NVIDIA的测试结果,FasterTransformer在多个模型上都取得了显著的性能提升:

BERT模型上,相比原生PyTorch实现可获得4-6倍的加速。
GPT模型上,相比Megatron-LM实现可获得1.5-2倍的加速。
T5模型上,可获得2-3倍的加速。

以下是BERT模型在T4 GPU上的性能对比图:

BERT性能对比

可以看到,FasterTransformer在各种batch size和sequence length下都取得了显著的性能提升,尤其是在大batch size和长序列的情况下。

支持的模型和框架

FasterTransformer支持多种Transformer模型和深度学习框架:

模型:BERT、GPT、T5、BLOOM、Longformer等
框架:TensorFlow、PyTorch、Triton Inference Server
精度:FP32、FP16、INT8
并行:Tensor并行、Pipeline并行

下面是FasterTransformer支持的主要模型和特性矩阵:

模型	框架	FP16	INT8	稀疏化	Tensor并行	Pipeline并行
BERT	TensorFlow	✓	✓	-	-	-
BERT	PyTorch	✓	✓	✓	✓	✓
GPT	PyTorch	✓	-	-	✓	✓
T5	PyTorch	✓	-	-	✓	✓

使用方法

FasterTransformer提供了C++、Python等多种接口。以PyTorch为例,使用FasterTransformer加速BERT模型的基本步骤如下:

安装FasterTransformer

pip install fastertransformer

导入相关模块

import torch
from fastertransformer import BertEncoderOP

创建FasterTransformer BERT encoder

ft_encoder = BertEncoderOP(layer_num=12, 
                           head_num=12,
                           size_per_head=64,
                           int8_mode=0)

将PyTorch模型权重转换为FasterTransformer格式

ft_weights = ft_encoder.convert_pytorch_model(torch_model)

执行推理

output = ft_encoder.forward(input_ids, 
                            attention_mask,
                            ft_weights)

通过简单的几步,就可以将FasterTransformer集成到现有的PyTorch代码中,实现性能的大幅提升。

总结

FasterTransformer作为一个高性能的Transformer推理库,通过多种优化技术显著提升了BERT、GPT等模型的推理速度。它易于集成、功能丰富,是部署大规模语言模型的理想选择。虽然该项目已停止开发,但其优化思路和技术仍具有重要的参考价值。对于需要高性能Transformer推理的应用场景,FasterTransformer仍是一个值得考虑的解决方案。