API 介绍#

LLM API 是一个高级 Python API，专为 LLM 工作流程设计。 此 API 正在开发中，将来可能会有重大更改。

支持的模型#

  • Llama（包括变体 Mistral、Mixtral、InternLM）

  • GPT（包括变体 Starcoder-1/2、Santacoder）

  • Gemma-1/2

  • Phi-1/2/3

  • ChatGLM（包括变体 glm-10b、chatglm、chatglm2、chatglm3、glm4）

  • QWen-1/1.5/2

  • Falcon

  • Baichuan-1/2

  • GPT-J

  • Mamba-1/2

模型准备#

LLM 类支持来自以下任何来源的输入

  1. Hugging Face Hub：触发从 Hugging Face 模型中心下载，例如 TinyLlama/TinyLlama-1.1B-Chat-v1.0。

  2. 本地 Hugging Face 模型：使用本地存储的 Hugging Face 模型。

  3. 本地 TensorRT-LLM 引擎：由 trtllm-build 工具构建或由 Python LLM API 保存。

这些格式中的任何一种都可以与 LLM(model=<any-model-path>) 构造函数互换使用。

以下部分介绍如何为 LLM API 使用这些不同的格式。

Hugging Face Hub#

使用 Hugging Face Hub 就像在 LLM 构造函数中指定 repo 名称一样简单

llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")

您还可以直接在 Hugging Face Hub 上加载 TensorRT 模型优化器的 量化检查点，方法相同。

本地 Hugging Face 模型#

鉴于 Hugging Face 模型中心的普及，API 支持 Hugging Face 格式作为起点之一。 要将 API 与 Llama 3.1 模型一起使用，请使用以下命令从 Meta Llama 3.1 8B 模型页面下载模型

git lfs install
git clone https://hugging-face.cn/meta-llama/Meta-Llama-3.1-8B

模型下载完成后，您可以加载模型

llm = LLM(model=<path_to_meta_llama_from_hf>)

使用此模型受制于特定的 许可。 同意条款并使用 Hugging Face 进行身份验证以开始下载。

本地 TensorRT-LLM 引擎#

有两种方法可以构建 TensorRT-LLM 引擎

  1. 您可以直接使用 trtllm-build 工具从 Hugging Face 模型构建 TensorRT-LLM 引擎，然后将引擎保存到磁盘以供以后使用。 请参阅 GitHub 上 examples/llama 存储库中的 README。

    引擎构建完成后，我们可以加载模型

    llm = LLM(model=<path_to_trt_engine>)

  2. 使用 LLM 实例创建引擎并持久化到本地磁盘

    llm = LLM(<model-path>)

# Save engine to local disk
llm.save(<engine-dir>)

    可以使用 model 参数加载引擎，如第一种方法所示。

提示和故障排除#

以下提示通常可以帮助熟悉 TensorRT-LLM 其他 API 的新 LLM API 用户

  • RuntimeError：只有 rank 0 才能启动多节点会话，得到 1

    对于使用 LLM API 启动单节点多 GPU 推理，无需添加 mpirun 前缀。

    例如，您可以运行 python llm_inference_distributed.py 以在单个节点上执行多 GPU。

  • Slurm 节点上的挂起问题

    如果您在使用 Slurm 管理的节点上遇到挂起或其他问题，请将前缀 mpirun -n 1 --oversubscribe --allow-run-as-root 添加到您的启动脚本。

    例如，尝试 mpirun -n 1 --oversubscribe --allow-run-as-root python llm_inference_distributed.py。

  • MPI_ABORT 在 communicator MPI_COMM_WORLD 中的 rank 1 上被调用，错误代码为 1。

    由于 LLM API 依赖于 mpi4py 库，请将 LLM 类放在一个函数中，并在 __main__ 命名空间下保护程序的主要入口点，以避免 mpi4py 中的 递归派生过程。

    此限制仅适用于多 GPU 推理。

  • 生成后无法退出

    LLM 实例管理线程和进程，这可能会阻止其引用计数达到零。 为了解决这个问题，有两种常见的解决方案

    1. 如快速入门指南中所示，将 LLM 实例包装在一个函数中。 这将减少引用计数并触发关闭过程。

    2. 将 LLM 用作上下文管理器，代码如下：with LLM(...) as llm: ...，一旦退出 with 语句块，shutdown 方法将自动调用。