概述
大语言模型能够解决复杂的数学问题。然而，要使其在规模上高效运行，仅有一个强大的模型检查点是不够的。您需要合适的服务堆栈、量化策略和解码方法——这些通常分散在不同的工具中，无法很好地协同工作。团队最终需要管理容器、转换脚本和临时代码，才能比较BF16与FP8，或者测试推测解码设置。
本文展示了如何使用某中心NeMo-Skills库来管理某中心TensorRT-LLM，从而构建一个快速、可复现的推理流水线。这个流水线是我们在2024年赢得人工智能数学奥林匹克奖时所使用配置的简化版本，它在两个某中心H100 GPU上，通过FP8量化和ReDrafter推测解码实现了4倍的批量推理加速。相同的工作流程可以在单个工作站上运行，也可以轻松扩展到集群，只需极少改动。
关键步骤
通过本文，您将学习如何：
- 准备一个OpenMath模型并将其量化为FP8 TensorRT-LLM引擎。
- 训练并集成用于推测解码的ReDrafter草案模型。
- 启动一个优化的推理服务器，可选择通过安全的代码沙箱进行工具调用。
- 在BF16、FP8和FP8+ReDrafter配置之间对延迟和吞吐量进行基准测试。
设置环境
第一步是建立一个一致且隔离的环境。我们将使用一个某中心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-ytensorrt pip3installtensorrt_llm==1.1.0rc0# 安装NeMo-Skillspipinstallgit+https://github.com/NVIDIA/NeMo-Skills.git准备模型权重
下一步是准备我们的大语言模型。我们将下载nvidia/OpenMath-Nemotron-14B-Kaggle模型，并使用FP8量化将其转换为优化的TensorRT-LLM引擎。
关于FP8量化的说明：FP8（8位浮点）量化效率很高，但需要支持E4M3 FP8的GPU（如某中心Hopper GPU）。对于其他GPU，建议使用int8_wo（仅权重的8位整数量化），它不需要校准。
下载模型权重和数据集
生成一个Hugging Face令牌并将其导出为环境变量。然后使用Hugging Face CLI下载必要的模型和数据集。
# 导出您的Hugging Face令牌exportHF_TOKEN=hf_YOUR_HUGGING_FACE_TOKEN# 安装Hugging Face CLIpipinstall-U"huggingface_hub[cli]"# 下载14B参数的主模型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数据集的子集来创建它。提供了一个示例来生成HuggingFace格式的数学校准数据集。
转换并量化为TensorRT-LLM引擎
现在，将Hugging Face模型转换为TensorRT-LLM引擎，应用FP8量化，并使用准备好的校准数据集。此步骤生成FP8量化后的LLM推理引擎。
ns convert\--input_modelOpenMath-Nemotron-14B-kaggle\--output_modelOpenMath-Nemotron-14B-kaggle-fp8-trtllm\--convert_fromhf\--convert_totrtllm\--num_gpus2\--dtypefp8\--hf_model_namenvidia/OpenMath-Nemotron-14B-kaggle\--model_typeqwen\--max_input_len30000\--max_seq_len32000\--no-trt_reuse_tmp_engine\--calib_dataset./calibration_dataset此命令执行后，您的FP8 LLM引擎已准备就绪，可以进行部署。
使用ReDrafter加速推理
为了进一步提高推理效率，我们集成了ReDrafter。这种推测解码技术使用一个更小的“草案”模型来预测令牌，使主LLM能够更快地生成响应。ReDrafter是由某机构开发的基于RNN的推理方法。其实现与TensorRT-LLM库支持的大多数模型兼容。
安装和训练ReDrafter
首先，安装ReDrafter库。草案模型的Tokenizer和训练数据应与基础模型使用的保持一致。如果无法获取原始训练数据，也可以使用基础模型的输出来训练草案模型。
# 安装ReDrafter库pipinstall--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-mnemo_skills.training.train_redrafter\--llm_name_or_path'OpenMath-Nemotron-14B-kaggle'\--dataset"OpenMathReasoning"\--dataset_split"tir"\--bf16True\--output_dir"redrafter_output"\--num_train_epochs1\--per_device_train_batch_size1\--gradient_accumulation_steps4\--save_strategy"no"\--learning_rate0.001\--weight_decay0.\--warmup_ratio0.1\--lr_scheduler_type"cosine"\--logging_steps20\--tf32True\--model_max_length2048\--dataset_nrows50000\--drafter_predict_n_tokens3\--drafter_num_layers2\--rnnTrue\--phasetrain\--report_towandb# 如果不使用wandb则删除此行在训练过程中，观察redrafter2_top1分数。如果该分数高于0.6，则表明接近2倍的运行时性能（60%的步骤接受接下来三个起草的令牌）。
为ReDrafter模型构建TensorRT-LLM引擎
接下来，我们将把训练好的ReDrafter模型转换为TensorRT-LLM检查点，然后将其与我们的主LLM结合，创建最终的、加速的TensorRT-LLM引擎。
首先，克隆TensorRT-LLM仓库以访问其转换脚本：
gitclone https://github.com/NVIDIA/TensorRT-LLM/然后，将训练好的ReDrafter PyTorch检查点转换为TensorRT-LLM检查点。
# FP8量化步骤中生成的基础模型中间检查点exportBASE_TRTLLM_CKPT=$(pwd)/OpenMath-Nemotron-14B-kaggle-fp8-trtllm-tmp-ckpt# 训练好的草案检查点exportREDRAFTER_PYTORCH_CKPT=$(pwd)/redrafter_output/redrafter__redrafter_OpenMath-Nemotron-14B-kaggle_n_3_lr_0.001_layers_2exportREDRAFTER_TRTLLM_CKPT=$(pwd)/OpenMath-Nemotron-14B-kaggle-fp8-draft-ckptcd./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\--dtypebfloat16\--tp_size2\--redrafter_num_beams1\--redrafter_draft_len_per_beam3cd../../../最后，构建集成了草案头用于推测解码的TensorRT-LLM引擎。
trtllm-build\--checkpoint_dir$REDRAFTER_TRTLLM_CKPT\--output_dirOpenMath-Nemotron-14B-kaggle-fp8-redrafter-trtllm\--gemm_pluginfp8\--use_paged_context_fmha=enable\--max_batch_size32\--max_seq_len32000\--max_input_len32000\--max_num_tokens32000\--speculative_decoding_modeexplicit_draft_tokens\--max_beam_width1\--kv_cache_typepaged至此，您的TensorRT-LLM引擎已通过ReDrafter技术增强，可以开始服务了！
基准测试与结果
我们准备了一个配套的笔记本，您可以自己尝试完整的流程。该笔记本使用了与上述“容器设置”部分相同的容器设置和安装步骤，并使用两个H100 GPU进行推理。在该笔记本中，您可以：
- 在不同的TensorRT-LLM引擎（BF16、FP8、FP8+ReDrafter）上运行推理。
- 比较性能基准，如首个令牌时间和每设备吞吐量。
- 探索高级控制，例如在固定时间后提前停止或在完成前N次生成后终止。
- 使用工具调用运行推理。
以下是一组基准测试结果的示例：
| 指标 | BF16 | FP8 | FP8+ReDrafter |
|---|---|---|---|
| 总生成时间（秒） | 144.2 | 64.7 | 30.5 |
| 平均样本吞吐量（令牌/秒） | 34.6 | 75.2 | 138.5 |
表 1. 在两个H100 GPU上不同配置的TensorRT-LLM性能对比
完整的基准测试和代码可在笔记本中找到。更多结果请查看AIMO-2获奖方案论文。
可选：启用工具调用和代码执行沙箱
OpenMath LLM是一个强大的工具指令推理模型。这意味着它不仅生成文本，还能在安全的沙箱中编写和执行Python代码来解决问题。在配套笔记本中，我们提供了如何启动LLM服务器及其伴随的代码执行沙箱的示例。
其交互流程如下：
- LLM生成包裹在
<tool_call>和</tool_call>令牌中的Python代码。 - 推理引擎提取此代码并将其发送到沙箱。
- 沙箱执行代码并返回结果。
- 输出被反馈给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要在配套笔记本中关闭工具调用，请使用get_model而不是get_code_execution_model，如NeMo-Skills文档所示。
您可以亲自尝试。运行配套笔记本来在您的硬件上对这些性能改进进行基准测试，并体验工具调用功能。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）
