重要提示

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

量化#

训练后量化 (PTQ)#

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

模型量化有两个主要优点：减少模型内存需求和提高推理吞吐量。

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

量化过程包括以下步骤

使用适当的并行策略加载模型检查点
校准模型以获得适当的算法特定缩放因子
生成包含模型配置 (json)、量化权重 (safetensors) 和分词器配置 (yaml) 的输出目录或 .qnemo tarball。

加载模型需要使用在 nemo.collections.nlp.models.language_modeling.megatron.gpt_layer_modelopt_spec 模块中定义的 ModelOpt 规范。通常，校准步骤是轻量级的，并使用小型数据集来获取用于缩放张量的适当统计信息。生成的输出目录（或 .qnemo tarball）已准备好用于使用 Nvidia TensorRT-LLM 库构建服务引擎。引擎构建步骤在 NeMo 项目的 nemo.deploy 和 nemo.export 模块中也可用。

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

支持矩阵#

下表列出了流行的 LLM 架构的已验证模型支持矩阵。每个模型条目还可选地提供指向相应 Nemo 检查点的下载链接，以用于测试目的。对其他模型系列的支持是实验性的。

模型支持矩阵#
模型系列	FP8	INT8_SQ	INT4_AWQ
Llama (1, 2, 3)	✅	✅	✅
Mistral	✅	✅	✅
GPT-3	✅	✅	✅
Nemotron-3 8b	✅	✅	✅
Nemotron-4 15b	✅	✅	✅
Nemotron-4 340b (Base, Instruct, Reward)	✅	✅	✅
StarCoder 2	✅	✅	✅
Gemma	✅	✅	✅

示例#

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

脚本必须使用等于张量并行性的进程数正确启动。这可以通过以下 torchrun 命令实现

torchrun --nproc-per-node 8 examples/nlp/language_modeling/megatron_gpt_ptq.py \
    model.restore_from_path=llama3-70b-base-bf16.nemo \
    model.tensor_model_parallel_size=8 \
    model.pipeline_model_parallel_size=1 \
    trainer.num_nodes=1 \
    trainer.devices=8 \
    trainer.precision=bf16 \
    quantization.algorithm=fp8 \
    export.decoder_type=llama \
    export.inference_tensor_parallel=2 \
    export.save_path=llama3-70b-base-fp8-qnemo

对于大型模型，该命令可以在多节点设置中使用。例如，这可以使用 NeMo 框架启动器和 Slurm 完成。

输出目录存储以下文件

llama3-70b-base-fp8-qnemo/
├── config.json
├── rank0.safetensors
├── rank1.safetensors
├── tokenizer.model
└── tokenizer_config.yaml

TensorRT-LLM 引擎可以使用 nemo.export 子模块中提供的 TensorRTLLM 类方便地构建和运行

from nemo.export.tensorrt_llm import TensorRTLLM


trt_llm_exporter = TensorRTLLM(model_dir="/path/to/trt_llm_engine_folder")
trt_llm_exporter.export(
    nemo_checkpoint_path="llama3-70b-base-fp8-qnemo",
    model_type="llama",
)
trt_llm_exporter.forward(["Hi, how are you?", "I am good, thanks, how about you?"])

或者，也可以直接使用 trtllm-build 命令构建，请参阅 TensorRT-LLM 文档

trtllm-build \
    --checkpoint_dir llama3-70b-base-fp8-qnemo \
    --output_dir /path/to/trt_llm_engine_folder \
    --max_batch_size 8 \
    --max_input_len 2048 \
    --max_output_len 512 \
    --strongly_typed

已知问题#

目前，使用 nemo.export 模块为量化的 “qnemo” 模型构建 TensorRT-LLM 引擎仅限于单节点部署。

量化感知训练 (QAT)#

QAT 是一种微调量化模型的技术，用于恢复由于量化造成的模型质量下降。在 QAT 期间，PTQ 期间计算的量化缩放因子被冻结，并且模型权重被微调。虽然 QAT 比 PTQ 需要更多的计算资源，但它在恢复模型质量方面非常有效。要在来自 PTQ 的校准模型上执行 QAT，您需要在使用小型数据集在下游任务上进一步微调模型，然后再导出到 TensorRT-LLM。您可以重用您的训练管道进行 QAT。根据经验，我们建议 QAT 持续原始训练时长的 1-10%，并使用较小的学习率，例如 Adam 优化器的 1e-5。如果您在 SFT 模型上进行 QAT，其中学习率和微调数据集大小已经很小，您可以继续使用与 SFT 相同的学习率和数据集大小作为 QAT 的起点。由于 QAT 在 PTQ 之后完成，因此支持的模型系列与 PTQ 相同。

示例#

以下示例展示了如何在 Supervised Finetuned Llama2 7B 模型上执行 PTQ 和 QAT 以达到 INT4 精度。该脚本已在 8 个 RTX 6000 Ada 48GB GPU 上使用 8 个张量并行性进行了测试。或者，单个带有 8 个 40GB GPU 的 DGX A100 节点也可用于相同的目的。对于像 Llama2 70B 这样更大的模型，您可能需要使用一个或多个带有 8 个 80GB GPU 的 DGX H100 节点。

该示例是 SFT with Llama 2 playbook 的修改版本。有关设置 BF16 NeMo 模型和 databricks-dolly-15k 指令数据集的更多详细信息，请参阅 playbook。

首先，我们将按原样运行 playbook 中的 SFT 示例命令，以训练 Llama2 7B SFT 模型 100 步。确保将 trainer.max_steps=50 更改为 trainer.max_steps=100 以用于 examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py 脚本。这将花费约 2 小时来生成验证损失约为 1.15 的模型检查点，我们将在下一步将其用于 PTQ 和 QAT。

对于量化，我们使用 SFT 脚本和配置文件的修改版本，其中包括量化和 TensorRT-LLM 导出支持。除了新参数外，请确保传递与您为 SFT 训练传递的参数相同的参数，除了模型恢复路径将是 SFT 输出 .nemo 文件。以下示例命令将在 SFT 模型检查点上执行 PTQ，然后再次执行 SFT (QAT)，然后可以导出以用于 TensorRT-LLM 推理。该脚本将花费约 2-3 小时完成。

torchrun --nproc-per-node 8 examples/nlp/language_modeling/tuning/megatron_gpt_qat.py \
    trainer.num_nodes=1 \
    trainer.devices=8 \
    trainer.precision=bf16 \
    trainer.max_steps=100 \
    model.restore_from_path=<llama2-7b-sft-nemo-path> \
    model.global_batch_size=128 \
    quantization.algorithm=int4 \
    # other parameters from sft training

正如您从日志中看到的，INT4 PTQ 模型的验证损失约为 1.31，而 QAT 模型的验证损失约为 1.17，这非常接近 BF16 模型的损失 1.15。此脚本将在实验管理器日志目录（在配置文件 yaml 文件中）中生成量化的 .nemo 检查点，该检查点可用于进一步训练。它还可以选择生成导出的 TensorRT-LLM 引擎目录或 .qnemo 文件，通过设置类似于 PTQ 示例的 export 参数，该文件可用于推理。请注意，如果需要获得更好的模型质量，您可以调整 QAT 训练器步数和学习率。

使用 NVIDIA Transformer Engine 在 FP8 中训练的 NeMo 检查点#

如果您有一个在 Transformer Engine 的预训练或微调期间生成的 FP8 量化检查点，则可以使用 nemo.export 将其直接转换为 FP8 TensorRT-LLM 引擎。API 与常规 .nemo 和 .qnemo 检查点相同

from nemo.export.tensorrt_llm import TensorRTLLM


trt_llm_exporter = TensorRTLLM(model_dir="/path/to/trt_llm_engine_folder")
trt_llm_exporter.export(
    nemo_checkpoint_path="/path/to/llama2-7b-base-fp8.nemo",
    model_type="llama",
)
trt_llm_exporter.forward(["Hi, how are you?", "I am good, thanks, how about you?"])

量化的导出设置可以通过 trt_llm_exporter.export 参数进行调整

fp8_quantized: Optional[bool] = None：手动启用/禁用 FP8 量化
fp8_kvcache: Optional[bool] = None：手动启用/禁用 KV 缓存的 FP8 量化

默认情况下，量化设置是从 NeMo 检查点自动检测的。

参考文献#

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