重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

量化#

NeMo 提供训练后量化 (PTQ)，以将 FP16/BF16 模型后处理为较低精度格式，从而实现高效部署。以下章节详细介绍了如何使用它。

训练后量化#

PTQ 支持以低精度格式（FP8、INT4 或 INT8）部署模型，以实现高效服务。不同的 量化方法 可用，包括 FP8 量化、INT8 SmoothQuant 和 INT4 AWQ。

模型量化有三个主要优点：减少模型内存需求、降低内存带宽压力以及提高推理吞吐量。

在 NeMo 中，量化由 NVIDIA TensorRT 模型优化器 (ModelOpt) 启用 – 该库用于量化和压缩深度学习模型，以在 GPU 上实现优化的推理。

量化过程包括以下步骤

  1. 使用适当的并行策略加载模型检查点。

  2. 校准模型以获得较低精度 GEMM 的缩放因子。

  3. 生成包含模型配置 (json) 和量化权重 (safetensors) 的 TensorRT-LLM 检查点。此外，还保存了设置模型分词器所需的上下文。

加载模型需要使用在 nemo.collections.nlp.models.language_modeling.megatron.gpt_layer_modelopt_spec 模块中定义的自定义 ModelOpt 规范。通常，校准步骤是轻量级的，并使用小型数据集来获取用于缩放张量的适当统计信息。生成的输出目录已准备好用于使用 NVIDIA TensorRT-LLM 库构建服务引擎（请参阅 通过导出 TensorRT-LLM 部署 NeMo 模型）。我们将此检查点称为 qnemo 检查点。

量化算法也可以方便地设置为 "no_quant"，以便仅使用 TensorRT-LLM 部署的默认精度执行权重导出步骤。这对于获得用于比较的基线性能和准确性结果非常有用。

支持矩阵#

下表列出了常用 LLM 架构的经过验证的模型支持矩阵。对其他模型系列的支持是实验性的。

模型名称

模型参数

解码器类型

FP8

INT8 SQ

INT4 AWQ

GPT

2B、8B、43B

gptnext

✓

✓

✓

Nemotron-3

8B、22B

gptnext

✓

✓

✓

Nemotron-4

15B、340B

gptnext

✓

✓

✓

Llama 2

7B、13B、70B

llama

✓

✓

✓

Llama 3

8B、70B

llama

✓

✓

✓

Llama 3.1

8B、70B、405B

llama

✓

✓

✓

Falcon

7B、40B

falcon

✗

✗

✗

Gemma 1

2B、7B

gemma

✓

✓

✓

StarCoder 1

15B

gpt2

✓

✓

✓

StarCoder 2

3B、7B、15B

gptnext

✓

✓

✓

Mistral

7B

llama

✓

✓

✓

Mixtral

8x7B

llama

✓

✗

✗

运行 PTQ 时，将根据使用的模型自动检测用于导出 TensorRT-LLM 检查点的解码器类型。如有必要，可以使用 decoder_type 参数覆盖它。

示例#

以下示例展示了如何使用单个 DGX H100 节点上的 8 个张量并行度将 Llama 3 70b 模型量化为 FP8 精度。量化模型旨在用于使用 export.inference_tp 参数指定的 2 个 H100 GPU 进行服务。

可以使用 NeMo CLI 或使用带有 torchrun 或 Slurm 的 PTQ 脚本启动量化工作流程。如下所示。

使用 NeMo CLI#

以下命令可以在 NeMo 容器内启动（仅支持单节点用例）

CALIB_TP=8
INFER_TP=2

nemo llm ptq \
    nemo_checkpoint=/opt/checkpoints/llama3-70b-base \
    calibration_tp=$CALIB_TP \
    quantization_config.algorithm=fp8 \
    export_config.inference_tp=$INFER_TP \
    export_config.path=/opt/checkpoints/llama3-70b-base-fp8-qnemo \
    run.executor=torchrun \
    run.executor.ntasks_per_node=$CALIB_TP

将 PTQ 脚本与 torchrun 或 Slurm 结合使用#

或者，可以直接使用 torchrun 命令和 scripts/llm/ptq.py。必须使用等于张量并行度的进程数正确启动脚本

CALIB_TP=8
CALIB_PP=1
INFER_TP=2

torchrun --nproc_per_node $CALIB_TP /opt/NeMo/scripts/llm/ptq.py \
    --nemo_checkpoint=/opt/checkpoints/llama3-70b-base \
    --calibration_tp=$CALIB_TP \
    --calibration_pp=$CALIB_PP \
    --algorithm=fp8 \
    --inference_tp=$INFER_TP \
    --export_path=/opt/checkpoints/llama3-70b-base-fp8-qnemo

对于大型模型，可以通过设置 --calibration_tp 和 --calibration_pp 以及相应的 Slurm --ntasks-per-node 和 --nodes 参数，在 Slurm 上为多节点用例启动此脚本

CALIB_TP=8
CALIB_PP=2
INFER_TP=8

srun --nodes $CALIB_PP --ntasks-per-node $CALIB_TP ... \
    python /opt/NeMo/scripts/llm/ptq.py \
        --nemo_checkpoint=/opt/checkpoints/nemotron4-340b-base \
        --calibration_tp=$CALIB_TP \
        --calibration_pp=$CALIB_PP \
        ...

对于 Llama 3 70b 示例，输出目录具有以下结构

llama3-70b-base-fp8-qnemo/
├── config.json
├── nemo_context/
├── rank0.safetensors
└── rank1.safetensors

下一步是为生成的检查点构建 TensorRT-LLM 引擎。这可以使用 nemo.export 模块中提供的 TensorRTLLM 类方便地实现和运行。有关详细信息，请参阅 通过导出 TensorRT-LLM 部署 NeMo 模型。或者，您可以直接使用 TensorRT-LLM trtllm-build 命令。

参考资料#

有关量化技术的更多详细信息，请参阅以下论文

本页内容