FP8 量化#

量化是一种技术,它允许模型以 int8 和 fp8 等较低精度运行,同时保持可接受的输出质量。以较低精度运行可以显著提升性能,大幅增加吞吐量并降低延迟。权衡是输出质量的下降,但在许多情况下,输出质量仍然可以接受,并且许多实际部署都利用了量化。如果您想了解更多关于量化的信息,请参考 掌握 LLM 技术 - 推理优化

本节将介绍如何启用 FP8 量化,并重点介绍一些针对 FP8 量化的配置选项,以提升性能。它还将继续以通过张量并行性分割到 4 个 H100-sxm-80GB GPU 上的 Llama-3.3-70B 作为案例研究,展示启用这些配置选项对性能的影响。

免责声明:此处显示的性能数据是真实的,但仅用于演示目的。环境、SKU、互连和工作负载的差异都可能显著影响性能,导致您的结果与此处所示不同。

启用量化#

要启用量化,您需要配置 QuantConfig 类并将其传递给 LLM 类的 quant_config 参数。至少必须指定 quant_algo 参数,该参数设置量化算法(fp8、每令牌 fp8、int8awq 等)。您可以在文档的 LLM-API->参考 部分找到 QuantConfig 的所有支持的量化算法和其他可配置选项。如果您使用的权重/检查点已经量化,则无需此步骤,但如果您使用的是 fp16 检查点,则还需要通过 CalibConfig 指定用于确定量化比例的校准数据集。CalibConfig 提供了设置校准数据集的多种选项,这些选项也可以在文档的 LLM-API->参考 部分进行参考。尽管 TensorRT-LLM 支持其他几种类型的量化,但本指南重点关注 fp8。

以下是从 bf16 检查点构建并保存 fp8 引擎的示例(请注意,fp8 仅在计算能力 > 8.9 的设备上受支持 - Ada、Hopper、Blackwell 及更高版本)

from tensorrt_llm import LLM, BuildConfig
from tensorrt_llm.llmapi import QuantConfig, QuantAlgo, CalibConfig

def main():

    quant_config = QuantConfig(quant_algo=QuantAlgo.FP8)

    calib_config = CalibConfig(
        calib_batches=512,
        calib_batch_size=1,
        calib_max_seq_length=2048,
        tokenizer_max_seq_length=4096
    )

    build_config = BuildConfig(
        max_num_tokens=2048,
        max_batch_size=512,
    )

    build_config.plugin_config.use_paged_context_fmha = True
    build_config.plugin_config.multiple_profiles = True

    llm = LLM(
        model="/path/to/Llama-3.3-70B",
        tensor_parallel_size=4,
        pipeline_parallel_size=1,
        build_config=build_config,
        quant_config=quant_config,
        calib_config=calib_config
    )

    llm.save("baseline_fp8_engine")

if __name__ == '__main__':
    main()

关于如何使用 TensorRT-LLM CLI 工作流程 构建 fp8 引擎的示例,请参阅 TensorRT-LLM LLaMA 示例。简而言之,您首先运行 examples/quantization/quantize.py 对模型检查点进行量化并转换为 TensorRT-LLM 格式,然后使用 trtllm-build

注意:虽然量化旨在保持模型精度,但这不能保证,因此检查量化后输出质量是否仍然足够非常重要。

FP8 “基准”性能#

对上述示例生成的引擎进行基准测试,得出以下性能结果。请注意,我们启用了之前提及的一些构建标志(多个配置、分页上下文 fmha)并调优了最大批量大小和最大令牌数。这样做是为了让您了解调优 fp8 引擎在排除为量化量身定制的选项后可以实现的性能。我们建议对量化引擎禁用 gemm 插件,因此此处未包含它(默认情况下是关闭的)。Reduce fusion 有一个量化特定的优化,将在稍后介绍。对于本页的其余部分,我们将此设置称为 fp8 的“基准”数据。

指标

令牌吞吐量 (tokens/秒)

3389.5305

请求吞吐量 (请求/秒)

1.6550

平均首令牌时间 (ms)

96.1597

平均令牌间隔延迟 (ms)

12.4248

量化 KV-Cache#

默认情况下 KV-Cache 不会量化,但 TensorRT-LLM 支持量化 KV-Cache 以进一步提升性能。然而,更激进地量化模型也会增加模型输出质量下降的风险,因此在使用此功能时务必进行检查。

启用量化 KV Cache#

LLM-API 通过 QuantConfig 中的 kv_cache_quant_algo 字段暴露用于 KV Cache 的量化算法。要启用 fp8 KV Cache,您可以按如下方式修改 QuantConfig

quant_config = QuantConfig(quant_algo=QuantAlgo.FP8,
                           kv_cache_quant_algo=QuantAlgo.FP8)

如果您使用 CLI 工作流程构建引擎,请将 --kv_cache_dtype fp8 参数传递给 examples/quantization/quantize.py

量化 KV Cache 的性能#

指标

基准线

FP8 KV-Cache 开启

令牌吞吐量 (tokens/秒)

3389.5305

5299.6372

请求吞吐量 (请求/秒)

1.6550

2.5877

平均首令牌时间 (ms)

96.1597

97.1287

平均令牌间隔延迟 (ms)

12.4248

12.5496

用于 Llama 模型的 Reduce Norm Fusion with User Buffers#

Reduce Norm Fusion 功能支持 fp8。对于 fp8 模型,还支持一个名为“User Buffers”的额外优化。用户缓冲区功能旨在消除通信内核中从本地缓冲区到共享缓冲区的额外复制,从而提升端到端性能。

启用 Reduce Norm Fusion with User Buffers#

要启用带有用户缓冲区的 Reduce Norm Fusion,请在 BuildConfig 初始化下方添加以下行

build_config.plugin_config.reduce_fusion = True
build_config.plugin_config.user_buffer = True

如果您使用 CLI 工作流程构建引擎,请将 --reduce_fusion enable--user_buffer enable 参数传递给 trtllm-build 以启用此功能。

注意:您必须先启用 reduce_fusion 才能启用 user_buffer

Reduce Norm Fusion + User Buffers 的性能:#

Reduce Norm Fusion + User Buffer ON: 与之前称为 FP8 KV-Cache ON 的引擎相同。

Reduce Norm Fusion + User Buffer ON: 之前的示例,启用了 reduce fusion 和 user buffers。调优后最大令牌数设置为 16384,最大批量大小设置为 512。

指标

Reduce Norm Fusion + User Buffer OFF

Reduce Norm Fusion + User Buffer ON

令牌吞吐量 (tokens/秒)

5299.6372

5980.7842

请求吞吐量 (请求/秒)

2.5877

2.9203

平均首令牌时间 (ms)

97.1287

82.2679

平均令牌间隔延迟 (ms)

12.5496

12.6975

Gated-MLP 中的 GEMM + SwiGLU Fusion#

Gated-MLP 中的 GEMM + SwiGLU Fusion 将两个 Matmul 操作和一个 SwiGLU 操作组合到单个内核中。目前仅在 Hopper 上支持 FP8 精度。虽然此融合可以提高性能,但由于丢弃了一个量化比例因子,可能会略微降低 FP8 PTQ 的精度。

我们建议在 Hopper 上运行大型 FP8 精度模型时启用此功能。对于非常小的工作负载或精度损失无法接受的情况,我们不建议启用此功能。

启用 GEMM + SwiGLU Fusion#

要启用 GEMM + SwiGLU Fusion,请在 BuildConfig 初始化下方添加以下行

build_config.plugin_config.gemm_swiglu_plugin = 'fp8'

对于延迟敏感的小批量大小情况,您可以将上述行替换为

build_config.plugin_config.low_latency_gemm_swiglu_plugin = 'fp8'

如果您使用 CLI 工作流程构建引擎,请将 --gemm_swiglu_plugin=fp8--low_latency_gemm_swiglu_plugin=fp8 参数(低延迟情况,二者选其一)传递给 trtllm-build

GEMM + SwiGLU Fusion 的性能#

指标

GEMM + SwiGLU fusion OFF

GEMM + SwiGLU fusion ON

令牌吞吐量 (tokens/秒)

5980.7842

5976.7977

请求吞吐量 (请求/秒)

2.9203

2.9184

平均首令牌时间 (ms)

82.2679

81.8841

平均令牌间隔延迟 (ms)

12.6975

11.7031

在此案例中,GEMM + SwiGLU 插件的表现与禁用时几乎相同。吞吐量的下降在运行间的差异范围内,TTFT 和 ITL 的改进微小。然而,我们发现,当与接下来讨论的低延迟 gemm 插件配合使用时,启用此功能对于获得最大吞吐量是必要的。

低延迟 GEMM 插件#

之前我们提到了 GEMM 插件 功能。虽然它支持 fp8,但我们建议禁用它(默认情况下是禁用的)。然而,对于 fp8 中的低延迟场景,我们建议尝试低延迟 GEMM 插件,看看它是否对您的工作负载有效。

启用低延迟 GEMM 插件#

要启用低延迟 GEMM 插件,请在 BuildConfig 初始化下方添加以下行

build_config.plugin_config.low_latency_gemm_plugin = 'fp8'

如果您使用 CLI 工作流程构建引擎,请将 --low_latency_gemm_plugin=fp8 参数传递给 trtllm-build 以启用此功能。再次强调,我们建议对 fp8 禁用 gemm 插件,因此如果您向 trtllm-build 传递了 --gemm_plugin=fp8 参数,我们建议将其移除。

低延迟 GEMM 插件的性能#

Low Latency GEMM ON: 配置与之前示例相同,但启用了低延迟 GEMM 插件。调优后最大令牌数设置为 16384,最大批量大小设置为 512。

指标

Low Latency GEMM OFF

Low Latency GEMM ON

令牌吞吐量 (tokens/秒)

5976.7977

6049.1625

请求吞吐量 (请求/秒)

2.9184

2.9537

平均首令牌时间 (ms)

81.8841

88.0162

平均令牌间隔延迟 (ms)

11.7031

10.8225

在此案例中,启用低延迟 gemm 插件实际上对吞吐量提供了显著提升。此外,它还改进了 ITL,但以牺牲 TTFT 为代价。更有甚者,当与 gemm+swiglu fusion 一起使用时,性能实际上比不启用插件时要差。这表明,对于此工作负载,低延迟 gemm 插件在 swiglu 之前的 gemm 中选择了较差的内核,但一旦由 gemm+swiglu fusion 自定义内核处理,低延迟 gemm 插件选择的其余内核则优于基准线,从而提高了性能。这强调了对不同设置进行基准测试的重要性,因为此插件的影响高度依赖于工作负载。如果可能,对于对性能要求极高的工作负载,进行一些网格搜索可能会很有用。

结论#

总体而言,利用量化可以显著提升性能。以下是我们的调优 fp8 模型与在指南前一页中达到的调优 fp16 数据相比的性能提升:

指标

调优 FP16 模型

调优 FP8 模型

提升百分比

令牌吞吐量 (tokens/秒)

2474.2581

6049.1625

144.48

请求吞吐量 (请求/秒)

1.2081

2.9537

144.49

平均首令牌时间 (ms)

147.5742

88.0162

40.36

平均令牌间隔延迟 (ms)

14.6852

10.8225

26.30

此外,与 fp8 基准数据(基准数据经过一定程度的调优,详见基准性能)相比,通过启用上面讨论的标志,我们获得了以下性能提升

指标

基准 FP8 模型

调优 FP8 模型

提升百分比

令牌吞吐量 (tokens/秒)

3389.5305

6049.1625

78.47

请求吞吐量 (请求/秒)

1.6550

2.9537

78.47

平均首令牌时间 (ms)

96.1597

88.0162

8.47

平均令牌间隔延迟 (ms)

12.4248

10.8225

12.90

如前所述,利用量化的注意事项是精度可能下降,我们强烈建议在使用量化之前先测试模型输出质量是否可接受。话虽如此,许多实际应用成功使用了量化,它带来的显著性能提升通常值得付出努力来确定是否适合。

配置选项建议总结:#

  1. 量化 KV-Cache:通常能显著提升吞吐量。只要启用此功能后输出质量仍然可接受,我们就建议开启它。

  2. Reduce fusion + user buffers:此功能仅在 fp8 Llama 和 Mistral/Mixtral 模型上受支持。其有效性取决于工作负载,因此我们建议开启并进行基准测试以检查效果。

  3. Gemm + Swiglu Plugin:此功能仅在带有 Swiglu 运算符的 fp8 模型(如 Llama、Mixtral 等)上受支持。与 reduce fusion 类似,其有效性取决于工作负载,我们建议进行有效性检查。由于它丢弃了一个量化比例,因此增加了影响精度的风险。

  4. 低延迟 GEMM 插件:其有效性取决于工作负载,因此我们建议开启并进行基准测试。正如我们在案例研究中看到的,其有效性可能受到其他标志的影响,因此如果可能,对各种配置选项组合进行基准测试是理想的选择。