重要提示

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

在框架中部署 NeMo 模型#

本节演示如何使用 NVIDIA Triton Inference Server 在框架内（称为“框架内”）部署 PyTorch 级别的 NeMo LLM。

快速示例#

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

  2. 使用下面显示的命令拉取并运行 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. 运行以下部署脚本以验证一切正常。该脚本直接在 Triton 服务器上提供 .nemo 模型

    python scripts/deploy/nlp/deploy_inframework_triton.py --nemo_checkpoint /opt/checkpoints/Nemotron-3-8B-Base-4k.nemo --triton_model_name nemotron

  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_inframework.py -mn nemotron -p "What is the color of a banana?" -mol 5

任何与 collections/nlp/models/language_modeling 中的 MegatronGPTModel 类兼容的模型都应与此脚本一起使用。

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

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

部署 NeMo LLM 模型#

执行该脚本将直接部署框架内 (.nemo) 模型并在 Triton 上启动服务。

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

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

    python scripts/deploy/nlp/deploy_inframework_triton.py --nemo_checkpoint /opt/checkpoints/Nemotron-3-8B-Base-4k.nemo --triton_model_name nemotron

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

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

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

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

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

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

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

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

    • --debug_mode：启用来自脚本的其他调试日志消息

  3. 要部署不同的模型，只需在 scripts/deploy/nlp/deploy_inframework_triton.py 脚本中更改 --nemo_checkpoint。任何与 NeMo 的 MegatronLLMDeployable 类兼容的模型都可以在此处使用。Nemotron、Llama2 和 Mistral 已经过测试并确认可以工作。

  4. 使用 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 的部署模块为将模型部署到 Triton 提供了简单的 API。

将 LLM 模型部署到 TensorRT-LLM#

您可以使用 deploy 模块中的 API 将框架内模型部署到 Triton。以下代码示例假定 Nemotron-3-8B-Base-4k.nemo 检查点已下载并挂载到 /opt/checkpoints/ 路径。

  1. 运行以下命令

    from nemo.deploy.nlp import MegatronLLMDeployable
from nemo.deploy import DeployPyTriton

megatron_deployable = MegatronLLMDeployable("/opt/checkpoints/Nemotron-3-8B-Base-4k.nemo")
nm = DeployPyTriton(model=megatron_deployable, triton_model_name="nemotron", port=8000)

nm.deploy()
nm.serve()
在本页中