大语言模型能够解决具有挑战性的数学问题。然而，若要实现其大规模高效运行，仅依靠一个强大的模型检查点还远远不够。高效的部署需要匹配合适的服务架构、量化策略和解码方法。这些关键技术通常分散在不同的工具中，难以无缝协同。结果，团队往往不得不处理复杂的容器配置、模型转换脚本以及临时编写的集成代码，以对比 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.txtpip uninstall -y tensorrtpip3 install tensorrt_llm==1.1.0rc0# Install NeMo-Skillspip 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 tokenexport HF_TOKEN=hf_YOUR_HUGGING_FACE_TOKEN# Install Hugging Face CLIpip install -U "huggingface_hub[cli]"# Download the 14B parameter main modelhuggingface-cli download nvidia/OpenMath-Nemotron-14B-kaggle --local-dir OpenMath-Nemotron-14B-kaggle# Download the OpenMathReasoning dataset for calibrationhuggingface-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 librarypip install --no-binary=protobuf --ignore-requires-python \"git+https://github.com/apple/ml-recurrent-drafter.git#egg=recurrent-drafting[dev,train]"# Train the ReDrafter modelns 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 stepexport BASE_TRTLLM_CKPT=$(pwd)/OpenMath-Nemotron-14B-kaggle-fp8-trtllm-tmp-ckpt# Trained draft checkpointexport REDRAFTER_PYTORCH_CKPT=$(pwd)/redrafter_output/redrafter__redrafter_OpenMath-Nemotron-14B-kaggle_n_3_lr_0.001_layers_2export REDRAFTER_TRTLLM_CKPT=$(pwd)/OpenMath-Nemotron-14B-kaggle-fp8-draft-ckptcd ./TensorRT-LLM/examples/redrafterpython 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 3cd ../../../
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 |
总生成时间（秒） |
|
|
|
平均样本吞吐量（Token/s） |
|
|
|
Notebook 中包含了完整的基准测试与代码。如需了解更多结果，请参阅 AIMO-2 冠军解决方案 论文。
可选：启用工具调用及代码执行沙盒功能
OpenMath LLM 是一个功能强大的工具型指令推理模型。它不仅能够生成文本，还能在安全的沙盒环境中编写并执行 Python 代码以解决复杂问题。在配套的 Notebook 中，我们提供了启动 LLM 服务器及其关联代码执行沙盒的具体示例。
交互的工作原理如下：
-
LLM 生成的 Python 代码会被封装在 <tool_call> 和 </tool_call> 标记之间。 -
推理引擎将提取该代码并发送至沙盒环境。 -
沙盒执行代码后，返回执行结果。 -
系统再将结果反馈给 LLM，以支持其继续生成内容或最终确定答案。
以下是此类交互的示例：
<tool_call># Initialize a list to store valid basesvalid_bases = []# Check bases from 10 upwardsfor b in range(10, 10000): # Arbitrary large upper limitnum1 = 9 * b + 7num2 = b + 7if num1 % num2 == 0:valid_bases.append(b)print(f"Found base: {b}")# Sum the valid basessum_bases = sum(valid_bases)print(f"Sum: {sum_bases}")# If sum is over 1000, take modulo 1000if sum_bases > 1000:result = sum_bases % 1000else:result = sum_basesprint(f"Final Result: {result}")</tool_call>```outputFound base: 21Found base: 49Sum: 70Final Result: 70```
get_model，而非 NeMo-Skills 文档中提到的 get_code_execution_model。
动手试一试吧。运行配套的 Notebook，在硬件上对这些性能优化进行基准测试，同时尝试使用工具调用功能。
— 关于作者 —
Igor Gitman
于 2017 年 12 月毕业于卡内基梅隆大学，获得机器学习硕士学位。此后，他一直在 NVIDIA 深度学习软件团队工作，负责开发新算法和工具包，以便在 GPU 上快速训练深度神经网络
Rakib Hasan
是 NVIDIA 的高级 AI 开发者技术工程师，专门优化深度学习工作负载，包括大语言模型 (LLM) 推理。他通过添加对 Llama 模型的支持并实现 RoPE 缩放和推理解码等功能，为 TensorRT-LLM 做出了贡献。Rakib 拥有路易斯安那州立大学 (LSU) 的博士学位，专注于在 x64 和 ARM CPU 上优化数学库。
Alessio Devoto
是罗马大学数据科学博士生，负责开发高效且可解释的机器学习模型。他的研究涵盖计算机视觉和自然语言处理，重点关注高效语言和视觉模型、AI 可解释性和自适应计算方法。他曾在慕尼黑的 NVIDIA 实习，并担任爱丁堡 NLP 的客座研究员。
Darragh Hanley
是 NVIDIA 应用智能体研究团队的高级 LLM 技术人员，主要研究自主推理和多智能体协作框架。他是 NVIDIA Kaggle Grandmaster 团队的成员，多次在竞赛中获胜 (峰值排名：全球排名前 4) ，并在加州大学伯克利分校的数据科学硕士课程中教授 DeepLearning。