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

  • 猎鹰

  • 百川-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=)构造函数互换使用。 以下部分描述了如何为LLM API使用这些不同的格式。

Hugging Face Hub

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

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

本地Hugging Face模型

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

git lfs install
git clone https://huggingface.co/meta-llama/Meta-Llama-3.1-8B

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

llm = LLM(model=<path_to_meta_llama_from_hf>)

使用此模型需遵守特定许可协议。同意条款并通过Hugging Face进行身份验证以开始下载。

本地 TensorRT-LLM 引擎

LLM API 可以使用 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用户：

  • 运行时错误：只有等级为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_COMM_WORLD 中，排名为 1 的进程调用了 MPI_ABORT，错误代码为 1。

    因为LLM API依赖于mpi4py库，所以将LLM类放在一个函数中，并在__main__命名空间下保护程序的主入口点，以避免在mpi4py中出现递归生成进程。

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

  • 生成后无法退出

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

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

    2. 使用LLM作为上下文管理器，使用以下代码：with LLM(...) as llm: ...，一旦它离开with语句块，shutdown方法将自动调用。