英伟达大模型推理提速实操：60分钟搞定4倍加速

各位跨境电商的实战精英们，大家好！在当前这个瞬息万变的数字化时代，效率就是生命线，尤其是在大模型（LLM）的应用中。咱们都知道，大模型在解决复杂数学问题上展现出了惊人的潜力，但要把它们真正部署到实际业务中，实现规模化的高效运行，可不仅仅是训练一个强大的模型那么简单。

很多时候，大家会发现，要流畅地实现BF16与FP8精度切换，或是部署一个推测解码方案，得在各种工具、容器和转换脚本之间来回折腾，费时费力，就像是“东拼西凑”，效率自然上不去。

今天，新媒网跨境获悉，咱们将深入探讨如何借助英伟达（NVIDIA）的NeMo-Skills库，来管理英伟达（NVIDIA）TensorRT-LLM，构建一个既快速又可复现的推理流水线。这套精简的方案，正是咱们在2024年AI数学奥林匹克竞赛中夺冠时使用的核心技术，它能够在两块英伟达H100 GPU上，通过FP8量化和ReDrafter推测解码，实现高达4倍的批量推理速度提升。更妙的是，这套工作流无论是跑在单台工作站，还是扩展到集群，只需做极少的改动。

通过今天的实战演练，各位将手把手学会：

• 如何将OpenMath模型准备好，并量化为FP8精度的TensorRT-LLM引擎。
• 如何训练并整合一个ReDrafter草稿模型，来实现高效的推测解码。
• 如何启动一个优化过的推理服务器，甚至还可以选择通过安全的代码沙盒进行工具调用。
• 如何对比BF16、FP8以及FP8+ReDrafter配置下的延迟和吞吐量，用数据说话。

如果您准备跟着一起操作，我们建议您准备一台配备两块H100（或其他支持FP8计算）GPU的机器，或者一个拥有类似节点的Slurm集群。工欲善其事，必先利其器，咱们先把环境搭建好。

环境搭建：基础先行，事半功倍

首先，咱们需要建立一个统一且隔离的运行环境。这就像盖房子要打好地基一样。我们将使用英伟达PyTorch NGC容器，并在其中安装必要的库：TensorRT-LLM负责模型优化，NeMo-Skills则负责整个流水线的管理。这里要特别提醒一句，FP8推理功能需要支持FP8的英伟达GPU，比如英伟达Ada Lovelace、英伟达Hopper、英伟达Blackwell或英伟达Rubin架构的显卡。咱们的例子假设有两块GPU可用。

容器设置与库安装

进入nvcr.io/nvidia/pytorch:25.05-py3容器后，请按照下面的指令安装TensorRT-LLM和NeMo-Skills：

# 确保没有冲突的TensorRT安装，然后安装TensorRT-LLM
[ -f /etc/pip/constraint.txt ] && : > /etc/pip/constraint.txt
pip uninstall -y tensorrt
pip3 install tensorrt_llm==1.1.0rc0

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

模型权重：精打细算，优化先行

接下来，咱们要准备核心的大语言模型。我们将下载nvidia/OpenMath-Nemotron-14B-Kaggle模型，并将其转换为优化过的TensorRT-LLM引擎，同时应用FP8量化技术。

关于FP8量化的小提示： FP8（8位浮点）量化效率极高，但它对GPU有要求，需要支持E4M3 FP8的显卡（比如英伟达Hopper架构的GPU）。如果您的GPU不满足这个条件，那么推荐使用int8_wo（8位整型，仅权重部分量化），这种方式不需要校准数据集，兼容性更广。

下载模型权重与数据集

首先，生成您的Hugging Face令牌（token），并将其导出为环境变量。然后，使用Hugging Face命令行工具下载所需的模型和数据集。

# 导出您的Hugging Face令牌
export HF_TOKEN=hf_YOUR_HUGGING_FACE_TOKEN

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

# 下载140亿参数的主模型
huggingface-cli download nvidia/OpenMath-Nemotron-14B-kaggle --local-dir OpenMath-Nemotron-14B-kaggle

# 下载OpenMathReasoning数据集用于校准
huggingface-cli download nvidia/OpenMathReasoning --repo-type dataset --local-dir OpenMathReasoning

为FP8量化准备校准数据集

对于FP8量化来说，一个能够代表推理数据的小型校准数据集至关重要。我们将利用OpenMathReasoning数据集的一个子集来创建它。这里提供了一个示例，用来在Hugging Face格式下生成数学校准数据集。

转换为TensorRT-LLM引擎并进行量化

现在，咱们将Hugging Face模型转换为TensorRT-LLM引擎，同时应用FP8量化并使用咱们刚才准备好的校准数据集。这一步将生成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技术。这是一种推测解码（speculative decoding）方法，它通过一个更小、更轻量的“草稿模型”来预测潜在的令牌序列，这样一来，主LLM就能更快地生成响应。新媒网了解到，ReDrafter是一种由苹果公司（Apple Inc.）开发的基于RNN的推理方法。它与TensorRT-LLM库中大多数支持的模型都兼容。

安装与训练ReDrafter

首先，咱们来安装ReDrafter库。请注意，草稿模型的tokenizer和训练数据，应该与基础模型使用的保持一致。如果获取不到原始的训练数据，咱们也可以用基础模型生成的数据来训练草稿模型。

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

# 训练ReDrafter模型
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 # 如果不使用wandb，请移除此行

在训练过程中，请关注redrafter2_top1这个指标。目标是让它高于0.6，这意味着接近2倍的运行时性能提升（60%的步骤能够接受接下来草稿模型预测的三个令牌）。

为ReDrafter模型构建TensorRT-LLM引擎

现在，咱们将训练好的ReDrafter模型转换为TensorRT-LLM的检查点，然后将其与咱们的主LLM结合起来，创建最终的、加速后的TensorRT-LLM引擎。

首先，克隆TensorRT-LLM仓库，以便获取其转换脚本：

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

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

# 来自FP8量化步骤的基础模型中间检查点
export BASE_TRTLLM_CKPT=$(pwd)/OpenMath-Nemotron-14B-kaggle-fp8-trtllm-tmp-ckpt

# 训练好的草稿模型检查点
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 ../../../

最后，构建结合了草稿头（draft head）的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的加持，已经准备好提供服务了！

实测为证：性能说话，效率之选

为了方便大家亲身体验这套完整的流水线，我们还准备了一个配套的Jupyter Notebook。这个Notebook是在与上述容器设置和库安装相同的环境下，以及两块H100 GPU上运行推理的。

在这个Notebook里，您可以：

• 针对不同的TensorRT-LLM引擎（BF16、FP8、FP8+ReDrafter）运行推理任务。
• 比较各种性能基准，比如首个令牌的生成时间（time to first token）和每设备吞吐量。
• 探索一些高级控制选项，例如在固定时间后提前停止，或者在完成N次生成后终止。
• 通过工具调用（tool-calling）的方式进行推理。

下面是您可能看到的基准测试结果示例：

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

完整的基准测试数据和代码都可以在配套的Notebook中找到。更多详细结果，可以参考AI数学奥林匹克竞赛的获奖方案论文。新媒网跨境认为，这样的性能飞跃，无疑能为咱们的跨境业务带来更低的成本、更高的效率，在激烈的市场竞争中占据优势。

拓展能力：工具调用与代码执行沙盒（选学）

OpenMath LLM是一个强大的工具指令推理模型。这意味着它不仅仅能生成文本，还能在安全的沙盒环境中编写并执行Python代码来解决问题。在配套的Notebook中，我们提供了一个如何同时启动LLM服务器及其配套代码执行沙盒的示例。

它的交互流程是这样的：

1. LLM生成被<tool_call>和</tool_call>标签包裹的Python代码。
2. 推理引擎提取这些代码，并将其发送到沙盒。
3. 沙盒执行代码并返回结果。
4. 输出结果被反馈回LLM，用于继续生成或最终确定其答案。

以下是一个交互示例：

<tool_call>
# 初始化一个列表来存储有效基数
valid_bases = []
# 从10开始检查基数
for b in range(10, 10000): # 任意设定的较大上限
num1 = 9 * b + 7
num2 = b + 7
if num1 % num2 == 0:
valid_bases.append(b)
print(f"Found base: {b}")

# 求和所有有效基数
sum_bases = sum(valid_bases)
print(f"Sum: {sum_bases}")

# 如果和超过1000，取模1000
if sum_bases > 1000:
result = sum_bases % 1000
else:
result = sum_bases
print(f"Final Result: {result}")
</tool_call>

Found base: 21
Found base: 49
Sum: 70
Final Result: 70

如果要在配套Notebook中关闭工具调用功能，您可以使用get_model而不是get_code_execution_model，具体操作可以参考NeMo-Skills的文档。

现在就去尝试吧！运行配套的Notebook，亲自在您的硬件上检验这些性能提升，并体验工具调用的强大功能。

风险前瞻与时效提醒

在咱们享受技术进步带来便利的同时，作为跨境实战专家，更要对潜在风险和技术时效性保持清醒的认识。

1. 合规性与数据安全： 跨境业务涉及多国法规，模型训练和推理过程中使用的数据务必符合当地的数据隐私法律法规。例如，个人信息保护、数据出境等，都需要严格遵守。同时，开源工具和模型的引入，也要求咱们在生产环境中仔细审计代码，确保没有潜在的安全漏洞。

2. 技术迭代速度： 大模型领域技术日新月异，今天的高效方案可能明天就有更优解。本教程基于2025年的技术发展前沿，旨在提供当前最优实践。但请大家保持持续学习的心态，关注英伟达及其他科技巨头的最新动向，及时更新您的知识和工具链。

3. 硬件依赖与成本： 本文提到的FP8量化等技术，对GPU硬件有较高要求，特别是需要支持FP8计算的英伟达Hopper等架构显卡。这可能意味着前期较高的硬件投入。在部署前，务必对成本效益进行详细评估，并考虑到全球芯片供应链的潜在波动。

4. 工具链的成熟度： NeMo-Skills和TensorRT-LLM等工具虽然强大，但随着技术发展，其API或功能可能也会更新。在实际应用中，建议大家关注官方文档的最新版本，确保使用的库版本兼容稳定。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/nvidia-llm-infer-boost-guide-1hr-get-4x-accel.html