如何实现数学问题求解推理速度提升4倍

原创于 2025-12-31 12:15:23 发布 · 186 阅读

CC 4.0 BY-SA版权

文章标签：

#AI推理 #模型量化 #推测解码 #程序那些事 #AIGC #大语言模型 #推理优化

概述

大语言模型能够解决复杂的数学问题。然而，要使其在规模上高效运行，仅有一个强大的模型检查点是不够的。您需要合适的服务堆栈、量化策略和解码方法——这些通常分散在不同的工具中，无法很好地协同工作。团队最终需要管理容器、转换脚本和临时代码，才能比较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 -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 模型，并使用FP8量化将其转换为优化的TensorRT-LLM引擎。

关于FP8量化的说明：FP8（8位浮点）量化效率很高，但需要支持E4M3 FP8的GPU（如某中心Hopper GPU）。对于其他GPU，建议使用int8_wo（仅权重的8位整数量化），它不需要校准。

下载模型权重和数据集

生成一个Hugging Face令牌并将其导出为环境变量。然后使用Hugging Face CLI下载必要的模型和数据集。

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

# 安装Hugging Face CLI
pip install -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_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。这种推测解码技术使用一个更小的“草案”模型来预测令牌，使主LLM能够更快地生成响应。ReDrafter是由某机构开发的基于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 ../../../

最后，构建集成了草案头用于推测解码的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技术增强，可以开始服务了！

基准测试与结果

我们准备了一个配套的笔记本，您可以自己尝试完整的流程。该笔记本使用了与上述“容器设置”部分相同的容器设置和安装步骤，并使用两个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/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）