重要提示

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

通过导出 vLLM 部署 NeMo 模型#

本节介绍如何使用脚本和 API 将 NeMo LLM 导出到 vLLM，并使用 NVIDIA Triton 推理服务器进行部署。

快速示例#

  1. 从 NVIDIA NGC 上的 NeMo 框架 下载 NeMo 格式的支持模型检查点。在此，我们将使用 Gemma 7B Base 模型作为示例。为了运行以下命令，需要一个 NGC 帐户。请访问 NGC 创建一个帐户并运行下面的 ngc 命令。

    ngc registry model download-version "nvidia/nemo/gemma_7b_base:1.1"

  2. 使用下面显示的命令拉取并运行 Docker 容器镜像。将 :vr 标签更改为您要使用的容器版本

    docker pull nvcr.io/nvidia/nemo:vr

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

  3. 在容器中，激活包含 vLLM 安装的虚拟环境 (venv)。

    source /opt/venv/bin/activate

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

    python scripts/deploy/nlp/deploy_vllm_triton.py \
    --nemo_checkpoint /opt/checkpoints/gemma_7b_base_v1.1/pytorch-7b-pt.nemo \
    --model_type gemma \
    --triton_model_name gemma \
    --tensor_parallelism_size 1

  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 gemma -p "The capital of Canada is" -mol 50

  8. 要导出和部署不同的模型（如 Llama3、Mixtral 或 Starcoder），请更改 scripts/deploy/nlp/deploy_vllm_triton.py 脚本中的 model_type。请查看下面以查看模型类型列表。

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

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

导出和部署 LLM 模型#

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

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

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

    python scripts/deploy/nlp/deploy_vllm_triton.py \
    --nemo_checkpoint /opt/checkpoints/gemma_7b_base_v1.1/pytorch-7b-pt.nemo \
    --model_type gemma \
    --triton_model_name gemma \
    --tensor_parallelism_size 1

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

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

    • --model_type: 模型类型。可以是“llama”、“mistral”、“mixtral”、“starcoder2”、“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: 用于权重转换的临时文件夹。默认为 /tmp/ 中的新文件夹。

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

    • --dtype: 部署模型的数据类型。默认为“bfloat16”。

    • --max_model_len: 模型最大输入 + 输出长度。默认为 512。

    • --max_batch_size: 模型最大批大小。默认为 8。

    • --debug_mode: 启用更详细的输出。

    • --weight_storage: 用于存储 vLLM 转换权重的策略。可以是“auto”、“cache”、“file”、“memory”。使用 --help 获取更多信息。

    注意

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

  3. 要导出和部署不同的模型（如 Llama3、Mixtral 和 Starcoder），请更改 scripts/deploy/nlp/deploy_vllm_triton.py 脚本中的 model_type 参数。请参阅下表以了解有关 LLM 模型使用的 model_type 的更多信息。

    模型名称

    model_type

    Llama 2

    llama

    Llama 3

    llama

    Gemma

    gemma

    StarCoder2

    starcoder2

    Mistral

    mistral

    Mixtral

    mixtral

  4. 通过缓存权重更快地导出。

    每当执行部署脚本时，它都会通过将 NeMo 检查点导出到 vLLM 来启动服务，其中包括将权重转换为兼容的格式。默认情况下，对于单 GPU 用例，转换在内存中进行，速度很快。对于多个 GPU，转换通过临时文件进行，并且可以选择在运行之间保留该文件以加快部署速度。为此，您需要创建一个空目录并使其可用于部署脚本。

    停止正在运行的容器，然后运行以下命令以指定一个空目录

    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_vllm_triton.py \
   --nemo_checkpoint /opt/checkpoints/gemma_7b_base_v1.1/pytorch-7b-pt.nemo \
   --model_type gemma \
   --triton_model_name gemma \
   --triton_model_repository /opt/checkpoints/tmp_triton_model_repository \
   --weight_storage cache \
   --tensor_parallelism_size 1

    --weight_storage cache 设置表明权重将通过 --triton_model_repository 指定的目录中的文件进行转换。仅当此文件比输入 .nemo 文件旧时，才会被覆盖。

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

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

    • 登录到 Hugging Face

    huggingface-cli login

    • 或者，设置 HF_TOKEN 环境变量

    export HF_TOKEN=your_token_here

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

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

将 LLM 模型导出到 vLLM#

您可以使用导出模块中的 API 将 NeMo 检查点导出到 vLLM。以下代码示例假定 gemma_7b_base_v1.1/pytorch-7b-pt.nemo 检查点已下载并挂载到 /opt/checkpoints/ 路径。

import os
from nemo.export.vllm_exporter import vLLMExporter

checkpoint_file = "/opt/checkpoints/gemma_7b_base_v1.1/pytorch-7b-pt.nemo"
model_dir = "/opt/checkpoints/gemma_7b_base_v1.1/vllm_export"

# Ensure that the temporary directory exists
if not os.path.exsts(model_dir):
   os.mkdir(model_dir)

# Export the checkpoint to vLLM, prepare for inference
exporter = vLLMExporter()
exporter.export(
   nemo_checkpoint=checkpoint_file,
   model_dir=model_dir,
   model_type="gemma",
)

# Run inference and print the output
output = exporter.forward(["What is the best city in the world?"], max_output_len=50, top_k=1, top_p=0.0, temperature=1.0)
print("output: ", output)

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

使用 vLLM 在 Triton 服务器上部署 LLM 模型#

您可以使用部署模块中的 API 将 vLLM 模型部署到 Triton。首先使用上面的导出示例将模型导出到 vLLM，只需删除末尾的 forward 和 print 调用即可。然后初始化 Triton 服务器并提供模型服务

from nemo.deploy import DeployPyTriton

nm = DeployPyTriton(model=exporter, triton_model_name="gemma", port=8000)
nm.deploy()
nm.serve()
本页内容