PyTorch 后端#

注意

注意：此功能目前为实验性功能，相关 API 在未来版本中可能会发生更改。

为了提高系统的可用性并提高开发人员的效率，TensorRT-LLM 推出了一种基于 PyTorch 的新型实验性后端。

TensorRT-LLM 的 PyTorch 后端在 0.17 及更高版本中可用。您可以通过导入 tensorrt_llm._torch 来尝试它。

快速开始#

这是一个简单的示例，展示了如何将 tensorrt_llm._torch.LLM API 与 Llama 模型一起使用。

 1from tensorrt_llm import SamplingParams
 2from tensorrt_llm._torch import LLM
 3
 4
 5def main():
 6    prompts = [
 7        "Hello, my name is",
 8        "The president of the United States is",
 9        "The capital of France is",
10        "The future of AI is",
11    ]
12    sampling_params = SamplingParams(max_tokens=32)
13
14    llm = LLM(model='TinyLlama/TinyLlama-1.1B-Chat-v1.0')
15    outputs = llm.generate(prompts, sampling_params)
16
17    for i, output in enumerate(outputs):
18        prompt = output.prompt
19        generated_text = output.outputs[0].text
20        print(f"[{i}] Prompt: {prompt!r}, Generated text: {generated_text!r}")
21
22
23if __name__ == '__main__':
24    main()

量化#

PyTorch 后端支持 FP8 和 NVFP4 量化。您可以传递 HF 模型 hub 中量化的模型，这些模型由 TensorRT 模型优化器生成。

from tensorrt_llm._torch import LLM
llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8')
llm.generate("Hello, my name is")

或者您可以尝试以下命令自行获取量化模型

git clone https://github.com/NVIDIA/TensorRT-Model-Optimizer.git
cd TensorRT-Model-Optimizer/examples/llm_ptq
scripts/huggingface_example.sh --model <huggingface_model_card> --quant fp8 --export_fmt hf

开发者指南#

关键组件#

已知问题#

  • SBSA 上的 PyTorch 工作流程与 Ubuntu 24.04 等裸机环境不兼容。 请使用 [PyTorch NGC 容器 (https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) 以获得 SBSA 平台的最佳支持。