如何将数学问题的推理速度提升 4 倍

大语言模型能够解决具有挑战性的数学问题。然而，若要实现其大规模高效运行，仅依靠一个强大的模型检查点还远远不够。高效的部署需要匹配合适的服务架构、量化策略和解码方法。这些关键技术通常分散在不同的工具中，难以无缝协同。结果，团队往往不得不处理复杂的容器配置、模型转换脚本以及临时编写的集成代码，以对比 BF16 与 FP8 的性能，或验证预测性解码的配置效果。

本文将介绍如何利用 NVIDIA NeMo-Skills 库构建高效且可复现的推理工作流，用于管理 NVIDIA TensorRT-LLM。我们采用的这一精简配置在 2024 年 AI 数学奥林匹克竞赛中荣获奖项，通过在两块 NVIDIA H100 GPU 上应用 FP8 量化和 ReDrafter 预测解码技术，实现了高达 4 倍的批量推理加速。该工作流不仅可在单台工作站上运行，还能轻松扩展至集群环境，仅需极少的配置调整即可实现横向扩展。

在本博文的最后，您将学习如何：

准备 OpenMath 模型，并将其量化为 FP8 格式以构建 TensorRT-LLM 引擎。
训练并集成 ReDrafter 草稿模型，支持预测性解码。
启动优化的推理服务器，支持安全代码沙盒及可选的工具调用功能。
对 BF16、FP8 以及结合 ReDrafter 的 FP8 配置进行延迟和吞吐量的基准测试。

建议使用配备两块 H100 GPU（或具备类似 FP8 功能的其他 GPU）的机器，或使用具有类似配置节点的 Slurm 集群，并按照以下步骤操作。

设置您的环境

第一步是构建统一的隔离环境。我们将采用 NVIDIA PyTorch NGC 容器，并安装必要的基础库：用于模型优化的 TensorRT-LLM，以及用于整体工作流管理的 NeMo 技能套件。FP8 推理需要支持 FP8 的 NVIDIA GPU，例如基于 NVIDIA Ada Lovelace、NVIDIA Hopper、NVIDIA Blackwell 或 NVIDIA Rubin 架构的显卡。本示例假设系统中配备有两块可用 GPU。

容器设置和库安装

进入 nvcr.io/nvidia/pytorch:25.05-py3 容器后，请执行以下命令来安装 TensorRT-LLM 和 NeMo-Skill。

# Ensure no conflicting TensorRT installations and install TensorRT-LLM
[ -f /etc/pip/constraint.txt ] && : > /etc/pip/constraint.txt
pip uninstall -y tensorrt
pip3 install tensorrt_llm==1.1.0rc0

# Install NeMo-Skills
pip install git+https://github.com/NVIDIA/NeMo-Skills.git

准备模型权重

下一步是准备大语言模型（LLM）。我们将下载 NVIDIA 的 OpenMath-Nemotron-14B-Kaggle 模型，并通过 FP8 量化将其转换为优化后的 TensorRT-LLM 引擎。

FP8 量化说明：FP8（8 位浮点）量化具有很高的效率，但需要支持 E4M3 FP8 格式的 GPU（例如 NVIDIA Hopper 架构的 GPU）。对于不支持该格式的 GPU，建议采用 int8_wo（仅使用权重量化的 8 位整数）方案，该方法无需校准，使用便捷。

下载模型权重和数据集

生成 Hugging Face Token 并将其设置为环境变量，然后通过 Hugging Face CLI 下载所需的模型和数据集。

# Export your Hugging Face token
export HF_TOKEN=hf_YOUR_HUGGING_FACE_TOKEN 

# Install Hugging Face CLI
pip install -U "huggingface_hub[cli]"

# Download the 14B parameter main model
huggingface-cli download nvidia/OpenMath-Nemotron-14B-kaggle --local-dir OpenMath-Nemotron-14B-kaggle

# Download the OpenMathReasoning dataset for calibration
huggingface-cli download nvidia/OpenMathReasoning --repo-type dataset --local-dir OpenMathReasoning

用于 FP8 量化的校准数据集准备

对于 FP8 量化，一个用于推理的小型校准数据集至关重要。我们将采用 OpenMathReasoning 数据集的一个子集来构建该数据集，并提供一个示例，展示如何以 HuggingFace 格式生成数学校准数据集。

转换并量化为 TensorRT-LLM 引擎

应用 FP8 量化技术，并结合准备好的校准数据集，将 Hugging Face 模型转换为 TensorRT-LLM 引擎。该步骤将生成支持 FP8 量化的 LLM 推理引擎。

ns convert \
    --input_model OpenMath-Nemotron-14B-kaggle \
    --output_model OpenMath-Nemotron-14B-kaggle-fp8-trtllm \
    --convert_from hf \
    --convert_to trtllm \
    --num_gpus 2 \
    --dtype fp8 \
    --hf_model_name nvidia/OpenMath-Nemotron-14B-kaggle \
    --model_type qwen \
    --max_input_len 30000 \
    --max_seq_len 32000 \
    --no-trt_reuse_tmp_engine \
    --calib_dataset ./calibration_dataset

执行该命令后，您的 FP8 LLM 引擎即可启动部署。

使用 ReDrafter 加速推理

为进一步提升推理效率，我们集成了 ReDrafter。该预测性解码技术通过一个较小的“草稿”模型来预测 token，从而加速主大语言模型的响应生成。ReDrafter 是由 Apple 开发的一种基于 RNN 的推理方法，其实现可与 TensorRT-LLM 库中支持的多数模型兼容。

安装和训练 ReDrafter

首先，安装 ReDrafter 库。草稿模型的分词器和训练数据应与基础模型保持一致。若原始训练数据不可获取，也可基于基础模型来训练草稿模型。

# Install the ReDrafter library
pip install --no-binary=protobuf --ignore-requires-python \     
"git+https://github.com/apple/ml-recurrent-drafter.git#egg=recurrent-drafting[dev,train]"

# Train the ReDrafter model
ns run_cmd --log_dir ./logs/ \
torchrun --nproc_per_node=2 -m nemo_skills.training.train_redrafter \
    --llm_name_or_path 'OpenMath-Nemotron-14B-kaggle' \
    --dataset "OpenMathReasoning" \
    --dataset_split "tir" \
    --bf16 True \
    --output_dir "redrafter_output" \
    --num_train_epochs 1 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 4 \
    --save_strategy "no" \
    --learning_rate 0.001 \
    --weight_decay 0. \
    --warmup_ratio 0.1 \
    --lr_scheduler_type "cosine" \
    --logging_steps 20 \
    --tf32 True \
    --model_max_length 2048 \
    --dataset_nrows 50000 \
    --drafter_predict_n_tokens 3 \
    --drafter_num_layers 2 \
    --rnn True \
    --phase train \
    --report_to wandb # Remove if not using wandb

在训练过程中，关注 redrafter2_top1 分数。若该分数达到 0.6 以上，表明运行时性能接近 2 倍（即 60% 的步骤能够接受接下来连续生成的三个 token）。

为 ReDrafter 模型构建 TensorRT-LLM 推理引擎

将训练好的 ReDrafter 模型转换为 TensorRT-LLM 检查点后，再与主 LLM 集成，从而构建出加速的 TensorRT-LLM 引擎。

首先，克隆 TensorRT-LLM 仓库以获取其中的模型转换脚本。

git clone https://github.com/NVIDIA/TensorRT-LLM/

接下来，将已训练好的 ReDrafter PyTorch 检查点转换为 TensorRT-LLM 检查点。

# Base model intermediate checkpoint from FP8 quantization step
export BASE_TRTLLM_CKPT=$(pwd)/OpenMath-Nemotron-14B-kaggle-fp8-trtllm-tmp-ckpt
# Trained draft checkpoint
export REDRAFTER_PYTORCH_CKPT=$(pwd)/redrafter_output/redrafter__redrafter_OpenMath-Nemotron-14B-kaggle_n_3_lr_0.001_layers_2
export REDRAFTER_TRTLLM_CKPT=$(pwd)/OpenMath-Nemotron-14B-kaggle-fp8-draft-ckpt

cd ./TensorRT-LLM/examples/redrafter
python convert_checkpoint.py \
    --base_model_checkpoint_dir $BASE_TRTLLM_CKPT \
    --drafter_model_dir $REDRAFTER_PYTORCH_CKPT \
    --output_dir $REDRAFTER_TRTLLM_CKPT \
    --dtype bfloat16 \
    --tp_size 2 \
    --redrafter_num_beams 1 \
    --redrafter_draft_len_per_beam 3
cd ../../../

最后，利用草稿头构建组合的 TensorRT-LLM 引擎基础模型，以实现预测性解码。

trtllm-build \
    --checkpoint_dir $REDRAFTER_TRTLLM_CKPT \
    --output_dir OpenMath-Nemotron-14B-kaggle-fp8-redrafter-trtllm \
    --gemm_plugin fp8 \
    --use_paged_context_fmha=enable \
    --max_batch_size 32 \
    --max_seq_len 32000 \
    --max_input_len 32000 \
    --max_num_tokens 32000 \
    --speculative_decoding_mode explicit_draft_tokens \
    --max_beam_width 1 \
    --kv_cache_type paged

您的 TensorRT-LLM 引擎现已由 ReDrafter 强力支持，随时可用！

基准测试和结果

我们准备了一个配套的 Notebook，您可以在其中亲手实践完整的制作流程。该 Notebook 的容器配置与安装方式同前述容器设置一致，推理过程使用两块 H100 GPU。在 Notebook 中，您将能够：

在不同的 TensorRT-LLM 引擎（BF16、FP8、FP8 + ReDrafter）上执行推理任务。
对比各项性能指标，例如首个 token 的生成时间以及各设备的吞吐量。
探索高级控制策略，例如在达到指定时间后提前终止生成，或在完成前 N 个生成步骤后停止。
同时支持使用工具调用进行推理执行。

以下是您将看到的基准测试结果示例：

指标	BF16	FP8	FP8+ReDrafter
总生成时间（秒）	144.2	64.7	30.5
平均样本吞吐量（Token/s）	34.6	75.2	138.5

表 1：TensorRT-LLM 在两块 H100 GPU 上不同配置的性能对比

Notebook 中包含了完整的基准测试与代码。如需了解更多结果，请参阅 AIMO-2 冠军解决方案论文。

可选：启用工具调用及代码执行沙盒功能

OpenMath LLM 是一个功能强大的工具型指令推理模型。它不仅能够生成文本，还能在安全的沙盒环境中编写并执行 Python 代码以解决复杂问题。在配套的 Notebook 中，我们提供了启动 LLM 服务器及其关联代码执行沙盒的具体示例。

交互的工作原理如下：

LLM 生成的 Python 代码会被封装在 <tool_call> 和 </tool_call> 标记之间。
推理引擎将提取该代码并发送至沙盒环境。
沙盒执行代码后，返回执行结果。
系统再将结果反馈给 LLM，以支持其继续生成内容或最终确定答案。

以下是此类交互的示例：

<tool_call>
# Initialize a list to store valid bases
valid_bases = []


# Check bases from 10 upwards
for b in range(10, 10000):  # Arbitrary large upper limit
    num1 = 9 * b + 7
    num2 = b + 7
    if num1 % num2 == 0:
        valid_bases.append(b)
        print(f"Found base: {b}")


# Sum the valid bases
sum_bases = sum(valid_bases)
print(f"Sum: {sum_bases}")


# If sum is over 1000, take modulo 1000
if sum_bases > 1000:
    result = sum_bases % 1000
else:
    result = sum_bases


print(f"Final Result: {result}")
</tool_call>
```output
Found base: 21
Found base: 49
Sum: 70
Final Result: 70
```

要在配套的 Notebook 中关闭工具调用，请使用 get_model，而非 NeMo-Skills 文档中提到的 get_code_execution_model。

动手试一试吧。运行配套的 Notebook，在硬件上对这些性能优化进行基准测试，同时尝试使用工具调用功能。

如何将数学问题的推理速度提升 4 倍

设置您的环境

容器设置和库安装

准备模型权重

下载模型权重和数据集

用于 FP8 量化的校准数据集准备

转换并量化为 TensorRT-LLM 引擎

使用 ReDrafter 加速推理

安装和训练 ReDrafter

为 ReDrafter 模型构建 TensorRT-LLM 推理引擎

基准测试和结果

可选：启用工具调用及代码执行沙盒功能

标签

关于作者

如何将数学问题的推理速度提升 4 倍

设置您的环境

容器设置和库安装

准备模型权重

下载模型权重和数据集

用于 FP8 量化的校准数据集准备

转换并量化为 TensorRT-LLM 引擎

使用 ReDrafter 加速推理

安装和训练 ReDrafter

为 ReDrafter 模型构建 TensorRT-LLM 推理引擎

基准测试和结果

可选：启用工具调用及代码执行沙盒功能

标签

关于作者

相关文章

NVIDIA Blackwell 架构席卷 MLPerf Training v5.1 基准测试

在 Kubernetes 上启用 NVIDIA GB200 NVL72 及以上型号的多节点 NVLink

生成式 AI 超分辨率利用可扩展的低计算模型加速天气预报

使用 GPU 加速的 Polars DataFrame 训练 XGBoost 模型

构建可扩展且容错的 NCCL 应用