使用 NVIDIA NeMo 和 NVIDIA TensorRT 模型优化器对 LLM 进行训练后量化

在本文中，我们概述了如何在 NVIDIA NeMo 中实施 PTQ。这是通过使用 NVIDIA TensorRT Model Optimizer 库实现的，该库可量化和压缩深度学习模型，以在 GPU 上优化推理。它还使用 NVIDIA TensorRT-LLM，这是一个开源库，用于优化大语言模型（LLM）推理。我们展示了量化模型的准确性和性能结果。在整个示例中，我们使用了 Llama 3 模型。

PTQ 是 NeMo LLM 构建和自定义功能的自然扩展，使用 NVIDIA TensorRT 模型优化器和 NVIDIA TensorRT-LLM 实现无缝、高效的部署路径。例如，NVIDIA NIM 受益于 NeMo 中的 PTQ 工作流程。

从技术角度来看，量化具有以下优势：

• 它减小了模型大小，因此适合使用较少的GPU和较低的总设备可用内存进行部署。
• 它通过使用更少的位数据类型来减轻内存带宽压力
• 例如，它可显著加快 NVIDIA 架构上的矩阵乘法（GEMM）运算，与微基准测试中的 FP16/BF16 数据类型相比，FP8 的运算速度最高可提升 2 倍。

NeMo 功能概述

NVIDIA NeMo 是一个端到端平台，可随时随地开发自定义生成式 AI。它包含用于训练、微调、检索增强生成、护栏和工具包的工具、数据管护工具和预训练模型，为企业提供了一种简单、经济高效且快速的生成式 AI 采用方式。

在使用 NeMo 工具包提供的各种选项构建模型后，可以使用 NeMo 导出和部署工具应用 PTQ 方法，并为优化模型提供服务。

最近发布的 NeMo 容器版本是一个独立工具包，包含应用 PTQ 和部署量化 LLM 所需的所有依赖项。

NeMo 和 TensorRT 模型优化器提供了适用于量化的各种模型，包括以下系列：

• GPT
• Llama
• Gemma
• StarCoder
• Nemotron（包括最近发布的Nemotron 4-340b）

PTQ 支持还提供多节点支持，用于使用适当的张量和流水线并行设置校准最大的 LLMs。

量化和部署 NeMo 模型

总的来说，PTQ 工作流程包括以下步骤：

1. 加载模型。
2. 校准模型以获取较低精度的 GEMMs 缩放系数，并将量化模型导出到 TensorRT-LLM 检查点。
3. 构建 TensorRT-LLM 引擎。
4. 部署模型 (例如，使用 PyTriton)。

加载模型

典型的 PTQ 用例始于以高精度格式（例如 FP16 或 BF16）训练的模型，该模型应以较低精度的数据类型（例如 FP8）提供。输入模型可以是通过之前的工作流步骤获得的基础 LLM 或指令调优 LLM。

NeMo 还为各种模型提供社区模型转换器，这些模型可用于生成相应的 NeMo 检查点。

校准和导出量化模型

在 PTQ 中，校准是一种获取模型层中执行的矩阵乘法运算的缩放系数的过程，以便使用低于训练精度的格式进行计算。

您可以直接从 NeMo 容器（例如使用 torchrun）轻松启动此步骤，也可以在多节点用例的 Slurm 集群上使用 NeMo 框架启动器。

简言之，量化代码可以归结为以下代码示例：

from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
from nemo.export.quantize import Quantizer

# Set quantization and export configs appropriately, see https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/conf/megatron_gpt_ptq.yaml

quantizer = Quantizer(quantization_config, export_config)

model = MegatronGPTModel.restore_from(...)

dataloader = ...  # A dataloader that yields lists of strings

def forward_loop(model):
    # Model forward pass for collecting activation statistics for calibration
    ...

model = quantizer.quantize(model, forward_loop)

quantizer.export(model)

完整的脚本 megatron_gpt_ptq.py 是校准工作流程的入口点。重要的量化参数在 megatron_gpt_ptq.yaml 配置文件中指定，建议使用默认设置。最重要的是，可用的低精度格式和量化算法是 FP8、INT4 AWQ 和 INT8 SQ。

通常，数据集的选择不会对准确性产生重大影响。然而，对于高度特定领域的应用程序（例如 StarCoder2 等代码完成模型），建议使用代码数据集以准确估计校正统计数据。

校正步骤的最后一步是以适合在下一步构建TensorRT-LLM引擎的TensorRT-LLM检查点格式保存模型。

总体而言，使用 NVIDIA DGX H100 GPU 节点对具有 70B 参数的中等大小模型使用 tensor 并行，校正过程只需几分钟。

构建 TensorRT-LLM 引擎

在运行 TensorRT-LLM 之前，您需要编译一组二进制文件，以构建推理引擎，这些二进制文件考虑了对特定 GPU 硬件、模型架构和推理设置的优化。

使用与常规 NeMo 模型相同的 API，为在校准步骤中获得的量化检查点构建引擎。若要构建 FP8 引擎，必须使用具有必要 FP8 支持的计算资源（例如 NVIDIA H100 Hopper 或 NVIDIA L40 Ada Lovelace GPU）运行此步骤。

以下 Python 命令展示了如何构建 TensorRT-LLM 引擎，以及如何在模型中传递示例提示。

from nemo.export.tensorrt_llm import TensorRTLLM

trt_llm_exporter = TensorRTLLM(model_dir=”path/to/trt_llm_engine”)
trt_llm_exporter.export(
    nemo_checkpoint_path=”path/to/model_qnemo”,
    max_batch_size=8,
    max_input_len=2048,
    max_output_len=512,
)
trt_llm_exporter.forward(["How does PTQ work?"])

导出命令通常需要几分钟才能完成构建或导出 TensorRT-LLM 引擎，并将其保存到 model_dir 参数中。

部署模型

可以使用 PyTriton 轻松部署给定的 TensorRT-LLM 引擎。

from nemo.deploy import DeployPyTriton
from nemo.export.tensorrt_llm import TensorRTLLM


trt_llm_exporter = TensorRTLLM(model_dir="path/to/trt_llm_engine")

nm = DeployPyTriton(
    model=trt_llm_exporter,
    triton_model_name="llama3_70b_fp8",
    port=8000,
)
nm.deploy()
nm.serve()

最后，在客户端上，NeMo Framework 提供了一个专用类，用于向服务器发送查询。以下代码示例展示了如何使用它。

from nemo.deploy.nlp import NemoQueryLLM


nq = NemoQueryLLM(
    url="localhost:8000",
    model_name="llama3_70b_fp8",
)

nq.query_llm(
    prompts=["How does PTQ work?"],
    top_k=1,
)

Llama 3 PTQ 示例和结果

出于演示目的，我们展示了两个预训练 Llama 3 模型变体（8B 和 70B）的 Llama 3 PTQ 吞吐量和准确性结果。我们分别使用 benchmark.py 和 mmlu.py 脚本评估了 TensorRT-LLM 引擎的性能和准确性。

以下结果是使用 TensorRT-LLM 0.12.0 和 TensorRT Model Optimizer 0.17.0 的 NVIDIA H100 80GB GPU 的结果。软件栈随附最新的 NeMo 框架容器，为您提供完整的环境。

准确性结果

图 2 显示了跨不同量化方法的两个 Llama 3 模型大小的 MMLU 结果，以及基准 FP16 结果。

值得注意的是，FP8 量化在很大程度上保持了准确度。在 INT8 SQ 和 Llama 3 模型大小的情况下，我们发现 SmoothQuant alpha 参数可以提高准确性。此参数用于管理量化焦点，从权重仅到激活仅，可以在量化配置中方便地进行设置。在两个 Llama 3 模型大小的情况下，alpha=0.8 的中间值都会产生最佳的 MMLU 结果。

在表 1 中，括号中的百分比数字是基准 FP16 分数的一小部分，用于测量给定场景保持准确性的程度。

性能结果

图 3 显示了定义为不同量化方法和两个 Llama 3 模型大小的量化模型与 FP16 基准的吞吐量比的推理加速。本文稍后将详细介绍实现的确切吞吐量结果。

在所有实验中，我们分别使用 2048 和 512 的输入和输出长度来构建 TensorRT-LLM 引擎并收集性能数据。这些值可视为文本摘要场景的代表性参数。

下表显示了用于为给定量化格式构建引擎的GPU数量，以及两种批量大小（32和1）的FP16基准结果。吞吐量相对于使用的GPU数量进行了标准化。

我们观察到，在较小的 Llama 3 变体中，FP8、INT8 SQ 和 INT4 AWQ 的速度分别提升了 1.45 倍、1.40 倍和 1.08 倍。在较大的模型中，FP8 和 INT8 SQ 的速度提升高达 1.81 倍。

INT4 AWQ 是一种仅权重的量化方法，推荐用于小批量大小。它主要提高了内存带宽，但对于更大的批量大小，会受到计算限制。

我们展示批量大小为 1 的结果以供比较。在本例中，对于 Llama 3 8B 和 Llama 3 70B 模型，我们分别获得了比 FP16 基准高 1.56 倍和 2.66 倍的性能优势。所有 Llama 3 70B 模型的量化变体都可以使用仅一个 NVIDIA H100 GPU 提供服务，而基准 FP16 精度至少需要两个 GPUs。

报告的吞吐量数据不应视为峰值性能，因为可以使用 TensorRT-LLM 的其他功能，例如在飞批处理，以进一步改进这些数据。

我们还使用 TensorRT-LLM 的 gptManagerBenchmark 工具检查了性能统计数据，重点关注 batch size = 32 的 FP16 基准和 FP8 量化引擎。

在 Llama 3 8B 模型中，第一个令牌（TTFT）的时间得到改善，令牌间延迟（ITL）的加速大致相当于本文前面报告的基于吞吐量的加速。

对于更大的 Llama 3 70B 模型，量化引擎在少 2 倍的 GPUs 上运行所实现的 TTFT 和 ITL 结果与基准 FP16 结果相似。这直接转化为 2 倍的资源节省。借助 PTQ，可以使用更少的 GPUs 更高效地服务模型。

总结

本文向您展示了如何在 NeMo 中使用 PTQ 为 LLM 部署构建高效的 TensorRT-LLM 引擎。在未来迭代中，用于模型的位数将大幅减少，因为新一代 NVIDIA B100 Blackwell 架构将支持 FP4。

还值得一提的是，对于某些应用程序，PTQ 可能已经足够，而其他应用程序可能需要量化感知训练（QAT）技术来微调量化权重以保持模型准确性。NeMo 中也提供了 QAT 来满足这些需求。

有关更多信息，请参阅训练后量化（Post-Training Quantization，PTQ）。PTQ 的切入点是 megatron_gpt_ptq.py 脚本。你可能还发现 NeMo 框架训练后量化（PTQ）手册很有用。它使用两个示例模型（Llama 3 和 Nemotron-340b）指导你完成整个部署过程。

至于 QAT，入门指南是 megatron_gpt_qat.py 脚本，相应的手册是适用于 Llama2 SFT 模型的 NeMo 框架量化感知训练（QAT）。有关更多信息，请参阅调整 TensorRT-LLM 性能的最佳实践。

致谢

非常感谢 NVIDIA 各个团队的许多敬业工程师的帮助，他们为成功集成 NeMo 和 TensorRT 模型优化器做出了贡献，包括 Asma Kuriparambil Thekkumpate、Keval Morabia、Wei-Ming Chen、Huizi Mao、Ao Tang、Dong Hyuk Chang、Alexandros Koumparoulis、Enwei Zhu 和 Simon Layton。