重要提示

您正在查看 NeMo 2.0 文档。此版本引入了对 API 的重大更改和一个新的库 NeMo Run。我们目前正在将所有功能从 NeMo 1.0 移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

通过导出 TensorRT-LLM 部署 NeMo 模型#

本节介绍如何使用脚本和 API 将 NeMo LLM (.nemo) 或量化 LLM (.qnemo) 导出到 TensorRT-LLM，并使用 NVIDIA Triton Inference Server 部署它。

快速示例#

  1. 按照 部署 NeMo LLM 主页 中的步骤下载 nemotron-3-8b-base-4k 模型。

  2. 在终端中，转到下载 Nemotron-3-8B-Base-4k.nemo 文件的文件夹。使用下面显示的命令拉取并运行 Docker 容器镜像。将 :vr 标签更改为您要使用的容器版本

    docker pull nvcr.io/nvidia/nemo:vr

docker run --gpus all -it --rm --shm-size=4g -p 8000:8000 \
    -v ${PWD}/Nemotron-3-8B-Base-4k.nemo:/opt/checkpoints/Nemotron-3-8B-Base-4k.nemo \
    -w /opt/NeMo \
    nvcr.io/nvidia/nemo:vr

  3. 运行以下部署脚本以验证一切正常。该脚本将 Nemotron NeMo 检查点导出到 TensorRT-LLM，然后在 Triton 服务器上提供服务

    python scripts/deploy/nlp/deploy_triton.py \
    --nemo_checkpoint /opt/checkpoints/Nemotron-3-8B-Base-4k.nemo \
    --model_type gptnext \
    --triton_model_name nemotron \
    --tensor_parallelism_size 1

  4. 如果测试产生与共享内存相关的错误，请使用 --shm-size 增加共享内存大小（例如，逐渐增加 50%）。

  5. 在单独的终端中，运行以下命令以获取正在运行的容器的容器 ID。请访问 nvcr.io/nvidia/nemo:vr 镜像以获取容器 ID。

    docker ps

  6. 访问正在运行的容器，并将 container_id 替换为实际的容器 ID，如下所示

    docker exec -it container_id bash

  7. 要向 Triton 服务器发送查询，请运行以下脚本

    python scripts/deploy/nlp/query.py -mn nemotron -p "What is the color of a banana?" -mol 5

  8. 要导出和部署不同的模型（例如 Llama3、Mixtral 或 Starcoder），请更改 deploy_triton.py 脚本中的 model_type。请查看下面以了解模型类型的列表。对于量化的 qnemo 模型，指定 model_type 不是必需的，因为此信息包含在 TensorRT-LLM 检查点中。

使用脚本在 Triton 服务器上部署 NeMo LLM#

您可以使用提供的脚本从 Triton 上的 NeMo 检查点部署 LLM。

导出和部署 LLM 模型#

执行脚本后，它会将模型导出到 TensorRT-LLM，然后在 Triton 上启动服务。

  1. 使用快速示例部分中描述的步骤启动容器。

  2. 要开始为 下载的模型 或 量化 模型提供服务，请运行以下脚本

    python scripts/deploy/nlp/deploy_triton.py \
    --nemo_checkpoint /opt/checkpoints/Nemotron-3-8B-Base-4k.nemo \
    --model_type gptnext \
    --triton_model_name nemotron \
    --tensor_parallelism_size 1

    以下参数在 deploy_triton.py 脚本中定义

    • --nemo_checkpoint：.nemo 或 .qnemo 检查点文件的路径。

    • --model_type：模型类型。choices=[“gptnext”, “gpt”, “llama”, “falcon”, “starcoder”, “mixtral”, “gemma”]。

    • --triton_model_name：Triton 上的模型名称。

    • --triton_model_version：模型版本。默认为 1。

    • --triton_port：Triton 服务器侦听请求的端口。默认为 8000。

    • --triton_http_address：Triton 服务器的 HTTP 地址。默认为 0.0.0.0。

    • --triton_model_repository：TensorRT 临时文件夹。默认为 /tmp/trt_llm_model_dir/。

    • --num_gpus：用于推理的 GPU 数量。大型模型需要多 GPU 导出。此参数已弃用。

    • --tensor_parallelism_size：用于张量并行拆分张量的 GPU 数量。默认为 1。

    • --pipeline_parallelism_size：用于流水线并行拆分模型的 GPU 数量。默认为 1。

    • --dtype：TensorRT-LLM 上模型的数据类型。默认为 “bfloat16”。目前，仅支持 “bfloat16”。

    • --max_input_len：模型的最大输入长度。默认为 256。

    • --max_output_len：模型的最大输出长度。默认为 256。

    • --max_batch_size：模型的最大批处理大小。默认为 8。

    • --max_num_tokens：最大令牌数。默认为 None。

    • --opt_num_tokens：最佳令牌数。默认为 None。

    • --ptuning_nemo_checkpoint：提示嵌入表的源 .nemo 文件。

    • --task_ids：提示嵌入的唯一任务名称。

    • --max_prompt_embedding_table_size：最大提示嵌入表大小。

    • --lora_ckpt：LoRA 权重的检查点列表。

    • --use_lora_plugin：激活 LoRA 插件，该插件启用嵌入共享。

    • --lora_target_modules：将 LoRA 添加到指定的模块，但仅当启用 --use_lora_plugin 时才激活。

    • --max_lora_rank：不同 LoRA 模块的最大 LoRA 秩。它用于计算 LoRA 插件的工作区大小。

    • --no_paged_kv_cache：禁用 TensorRT-LLM 中的分页 kv 缓存。

    • --disable_remove_input_padding：禁用 TensorRT-LLM 的删除输入填充选项。

    • --use_parallel_embedding：启用 TensorRT-LLM 的并行嵌入功能。

    • --export_fp8_quantized：手动覆盖 FP8 量化设置。

    • --use_fp8_kv_cache：手动覆盖 FP8 KV 缓存量化设置。

    弃用警告：num_gpus 参数已弃用，将在下一个版本发布后删除。

    注意

    此处描述的参数是通用的，应与任何 NeMo 检查点兼容。但是，重要的是您要查看主 部署 NeMo LLM 页面 中的 LLM 模型表，以了解优化的推理模型兼容性。我们正在积极努力扩展对其他检查点的支持。

  3. 要导出和部署不同的模型（例如 Llama3、Mixtral 或 Starcoder），请更改 deploy_triton.py 脚本中的 model_type。请查看下表以了解更多关于哪个 model_type 用于 LLM 模型的信息（对于 .qnemo 模型不是必需的）。

    模型名称

    model_type

    GPT

    gpt

    Nemotron

    gpt

    Llama 2

    llama

    Llama 3

    llama

    Llama 3.1

    llama

    Gemma

    gemma

    StarCoder1

    starcoder

    StarCoder2

    starcoder

    Mistral

    llama

    Mixtral

    mixtral

  4. 每当执行脚本时，它都会通过将 NeMo 检查点导出到 TensorRT-LLM 来启动服务。如果您想跳过优化推理选项中的导出步骤，您可以指定一个空目录来保存生成的 TensorRT-LLM 引擎。停止正在运行的容器，然后运行以下命令以指定一个空目录

    mkdir tmp_triton_model_repository

docker run --gpus all -it --rm --shm-size=4g -p 8000:8000 \
    -v ${PWD}:/opt/checkpoints/ \
    -w /opt/NeMo \
    nvcr.io/nvidia/nemo:vr

python scripts/deploy/nlp/deploy_triton.py \
    --nemo_checkpoint /opt/checkpoints/Nemotron-3-8B-Base-4k.nemo \
    --model_type gptnext \
    --triton_model_name nemotron \
    --triton_model_repository /opt/checkpoints/tmp_triton_model_repository \
    --tensor_parallelism_size 1

    检查点将在执行上述脚本后导出到指定文件夹，以便稍后可以重复使用。

  5. 要直接加载导出的模型，请在容器内运行以下脚本

    python scripts/deploy/nlp/deploy_triton.py \
    --triton_model_name nemotron \
    --triton_model_repository /opt/checkpoints/tmp_triton_model_repository \
    --model_type gptnext

  6. 使用 Hugging Face 令牌访问模型。

    如果您想使用 StarCoder1、StarCoder2 或 LLama3 模型运行推理，您需要生成一个有权访问这些模型的 Hugging Face 令牌。访问 Hugging Face 了解更多信息。获得令牌后，执行以下步骤之一。

    • 登录到 Hugging Face

    huggingface-cli login

    • 或者，设置 HF_TOKEN 环境变量

    export HF_TOKEN=your_token_here

使用提示嵌入表#

您可以使用学习到的虚拟令牌在推理期间执行下游流任务。一旦使用 NeMo 框架训练容器学习了虚拟令牌，所有令牌都将保存在 .nemo 文件中。您可以将此文件输入到脚本中，如下面的命令所示。由于 NVIDIA NGC 或 Hugging Face 上没有专门用于虚拟令牌的 NeMo 检查点，因此您需要查找或生成一个检查点。

  1. 假设存在提示嵌入表的检查点，请运行以下命令

    python scripts/deploy/nlp/deploy_triton.py \
    --nemo_checkpoint /opt/checkpoints/Nemotron-3-8B-Base-4k.nemo \
    --model_type gptnext \
    --triton_model_name nemotron \
    --triton_model_repository /opt/checkpoints/tmp_triton_model_repository \
    --max_prompt_embedding_table_size 1024 \
    --ptuning_nemo_checkpoint /opt/checkpoints/my_ptuning_table.nemo \
    --task_ids "task 1" \
    --tensor_parallelism_size 1

    max_prompt_embedding_table_size 参数应设置为所有下游任务的虚拟令牌总数。

  2. 要传递多个 NeMo 检查点，请运行以下命令

    python scripts/deploy/nlp/deploy_triton.py \
    --nemo_checkpoint /opt/checkpoints/Nemotron-3-8B-Base-4k.nemo \
    --model_type gptnext \
    --triton_model_name nemotron \
    --triton_model_repository /opt/checkpoints/tmp_triton_model_repository \
    --max_prompt_embedding_table_size 1024 \
    --ptuning_nemo_checkpoint /opt/checkpoints/my_ptuning_table-1.nemo /opt/checkpoints/my_ptuning_table-2.nemo \
    --task_ids "task 1" "task 2" \
    --tensor_parallelism_size 1

    请确保 my_ptuning_table-1.nemo 和 my_ptuning_table-2.nemo 中的虚拟令牌总数不超过 max_prompt_embedding_table_size 参数。

使用 NeMo 导出和部署模块 API 运行推理#

到目前为止，我们已经使用脚本来导出和部署 LLM 模型。但是，NeMo 的部署和导出模块提供了简单的 API，用于将模型部署到 Triton 并将 NeMo 检查点导出到 TensorRT-LLM。

将 LLM 模型导出到 TensorRT-LLM#

您可以使用导出模块中的 API 将 NeMo 检查点导出到 TensorRT-LLM。以下代码示例假定 Nemotron-3-8B-Base-4k.nemo 检查点已下载并挂载到 /opt/checkpoints/ 路径。此外，还假定 /opt/checkpoints/tmp_trt_llm 路径存在。

  1. 运行以下命令

    from nemo.export.tensorrt_llm import TensorRTLLM


trt_llm_exporter = TensorRTLLM(model_dir="/opt/checkpoints/tmp_trt_llm/")
trt_llm_exporter.export(
    nemo_checkpoint_path="/opt/checkpoints/Nemotron-3-8B-Base-4k.nemo",
    model_type="gptnext",
    n_gpus=1,
)
trt_llm_exporter.forward(
    ["What is the best city in the world?"],
    max_output_token=15,
    top_k=1,
    top_p=0.0,
    temperature=1.0,
)

  2. 请务必查看 TensorRTLLM 类文档字符串以了解详细信息。

  3. 对于量化的 qnemo 模型 – 请参阅 量化 – 可以使用相同的 TensorRTLLM 类。在这种情况下，指定 model_type 不是必需的。或者，高级用户可以使用 trtllm-build 命令直接构建引擎，请参阅 TensorRT-LLM 文档

    trtllm-build \
   --checkpoint_dir /opt/checkpoints/llama3-70b-base-fp8-qnemo \
   --output_dir /path/to/trt_llm_engine_folder \
   --gemm_plugin fp8 \
   --max_batch_size 8 \
   --max_input_len 2048 \
   --max_seq_len 2560 \
   --workers 2

将 LLM 模型部署到 TensorRT-LLM#

您可以使用部署模块中的 API 将 TensorRT-LLM 模型部署到 Triton。以下代码示例假定 Nemotron-3-8B-Base-4k.nemo 检查点已下载并挂载到 /opt/checkpoints/ 路径。此外，还假定 /opt/checkpoints/tmp_trt_llm 路径存在。

  1. 运行以下命令

    from nemo.export.tensorrt_llm import TensorRTLLM
from nemo.deploy import DeployPyTriton


trt_llm_exporter = TensorRTLLM(model_dir="/opt/checkpoints/tmp_trt_llm/")
trt_llm_exporter.export(
    nemo_checkpoint_path="/opt/checkpoints/Nemotron-3-8B-Base-4k.nemo",
    model_type="gptnext",
    n_gpus=1,
)

nm = DeployPyTriton(model=trt_llm_exporter, triton_model_name="nemotron", port=8000)
nm.deploy()
nm.serve()

  2. 相同的说明适用于量化检查点。

FP8 训练模型的直接 TensorRT-LLM 导出#

如果您有一个 FP8 训练的检查点，该检查点是在使用 NVIDIA Transformer Engine 进行预训练或微调期间生成的，则可以使用 nemo.export 将其直接转换为 FP8 TensorRT-LLM 引擎。入口点与常规 .nemo 和 .qnemo 检查点相同

from nemo.export.tensorrt_llm import TensorRTLLM


trt_llm_exporter = TensorRTLLM(model_dir="/opt/checkpoints/tmp_trt_llm/")
trt_llm_exporter.export(
    nemo_checkpoint_path="/opt/checkpoints/llama2-7b-base-fp8.nemo",
    model_type="llama",
)
trt_llm_exporter.forward(["Hi, how are you?", "I am good, thanks, how about you?"])

量化的导出设置可以通过 trt_llm_exporter.export 参数进行调整

  • fp8_quantized: Optional[bool] = None：启用/禁用 FP8 量化

  • fp8_kvcache: Optional[bool] = None：启用/禁用 KV 缓存的 FP8 量化

默认情况下，量化设置是从 NeMo 检查点自动检测的。

在本页中