外观
FasterTransformer 部署加速详解
内容整理自学习笔记,仅供面试备考参考;不构成录用、培训或考试承诺。
1. 为什么需要 FasterTransformer?
大模型推理面临以下挑战:
- 自回归推理过程中产生大量的 key 和 value 缓存,需要重复计算
- 大模型参数量极大,即使量化到 INT4 也是不小的内存占用(如 GPT-3 175B 半精度存储需要 350GB)
- 传统的 TensorRT 加速方法无法充分利用多 GPU 加载 Transformer 的不同块
2. FasterTransformer 介绍
NVIDIA FasterTransformer (FT) 是一个库,用于实现基于 Transformer 的神经网络推理的加速引擎,特别强调大型模型,以分布式方式跨越许多 GPU 和节点。
核心特性
| 特性 | 说明 |
|---|---|
| 实现语言 | C++/CUDA 编写 |
| 依赖库 | 高度优化的 cuBLAS、cuBLASLt 和 cuSPARSELt 库 |
| 架构支持 | 完整的编码器-解码器架构(如 T5)、仅编码器模型(如 BERT)、仅解码器模型(如 GPT) |
| 包含模块 | Transformer 块的高度优化版本(编码器和解码器) |
使用 FasterTransformer 可以在 GPU 上构建最快的 Transformer 推理流程。
3. FasterTransformer 核心是什么?
3.1 张量并行(TP)和流水线并行(PP)
| 并行方式 | 说明 |
|---|---|
| 张量并行(TP) | 每个张量被分成多个块,每个块放置在单独的 GPU 上,计算时各块并行处理,最后组合结果 |
| 流水线并行(PP) | 模型被深度拆分,将不同的完整层放置到不同的 GPU/节点上 |
3.2 通信底层
| 通信方式 | 用途 |
|---|---|
| MPI | 节点间通信 |
| NVIDIA NCCL | 节点内/节点间 GPU 通信 |
TP 和 PP 可以结合在一起,在多 GPU 和多节点环境中运行具有数十亿和数万亿个参数的大型 Transformer 模型。
3.3 框架集成
| 集成方式 | 多 GPU 支持 | 多节点支持 |
|---|---|---|
| TensorFlow 集成 | ❌ 仅单 GPU | ❌ |
| PyTorch 集成 | ✅ | ✅ |
| Triton 后端 | ✅ | ✅ |
为了避免为模型并行性而拆分模型的额外工作,FasterTransformer 提供了工具将模型从不同格式拆分和转换为 FasterTransformer 二进制文件格式,可直接以二进制格式加载模型。
4. FasterTransformer 优化方法
4.1 推理缓存优化
| 方面 | 说明 |
|---|---|
| 动机 | 自回归推理时会产生非常多的 key 和 value 值,每次都需要重复计算 |
| 优化策略 | 对产生的缓存分块存储,避免重复计算 |
4.2 内存优化
| 方面 | 说明 |
|---|---|
| 动机 | 大模型参数量极大,GPT-3 175B 半精度存储需要 350GB |
| 优化策略 | 缓存激活值和输出,进行新 sentence 推理时重新利用,避免多层反复计算 |
例如 GPT-3 中层数为 96,因此只需要 1/96 的内存量用于激活。
4.3 使用 MPI 和 NCCL 通信优化
张量并行性优化:
- FasterTransformer 遵循了 Megatron 的思想
- 对于自注意力块和前馈网络块,FT 按行拆分第一个矩阵的权重,按列拆分第二个矩阵的权重
- 通过优化,FT 可以将每个 Transformer 块的归约操作减少到两倍
流水线并行性优化:
- FasterTransformer 将整批请求拆分为多个微批,隐藏了通信的泡沫
- FasterTransformer 会针对不同情况自动调整微批量大小
4.4 MatMul 内核自动调整(GEMM 自动调整)
矩阵乘法是基于 Transformer 的神经网络中最繁重的操作。FT 使用来自 CuBLAS 和 CuTLASS 库的功能来执行这些操作。
GEMM 自动调整流程:
├── GemmBatchedEx 函数实现 MatMul 操作
├── 以 "cublasGemmAlgo_t" 作为输入参数选择底层算法
├── 对所有底层算法进行实时基准测试
├── 根据模型参数和输入数据选择最佳算法
│ ├── 注意层的大小
│ ├── 注意头的数量
│ └── 隐藏层的大小
└── 使用硬件加速的底层函数(如 __expf、__shfl_xor_sync)MatMul 操作可以在"硬件"级别使用不同的低级算法以数十种不同的方式执行,FT 会自动选择最优算法。
4.5 量化推理
| 精度 | 说明 |
|---|---|
| FP16 | 半精度浮点数,较少的数据传输量和所需内存 |
| INT8 | 8位整数,进一步减少内存占用 |
INT8 和 FP16 计算可以在特殊硬件上执行,如 Tensor Core(适用于从 Volta 开始的所有 GPU 架构),以及 Hopper GPU 中的 Transformer 引擎。
5. FasterTransformer 优化总结
| 优化方向 | 核心思想 | 效果 |
|---|---|---|
| 推理缓存优化 | 分块存储 KV 缓存 | 避免重复计算 |
| 内存优化 | 缓存激活值和输出 | 减少 1/N 层的内存量 |
| 通信优化(TP) | Megatron 思想拆分矩阵 | 归约操作减少到两倍 |
| 通信优化(PP) | 微批处理隐藏通信泡沫 | 自动调整微批量大小 |
| GEMM 自动调整 | 实时基准测试选最优算法 | 最大化矩阵乘法效率 |
| 量化推理 | FP16/INT8 低精度 | 减少内存和加速计算 |
FasterTransformer 通过不修改模型架构而在计算加速层面进行深度优化,是 NVIDIA 推出的高性能 Transformer 推理引擎,特别适合多 GPU 和多节点分布式部署场景。