NVIDIA 开源 TensorRT-LLM:H100 推理吞吐翻 8 倍
NVIDIA 推出开源库 TensorRT-LLM,专为加速大语言模型推理设计。它引入动态组批与 FP8 精度,在 H100 上相比 A100 带来最高 8 倍吞吐提升,Llama 2 推理提速 4.6 倍。开发者通过 Python API 即可调用底层并行优化,大幅降低算力与电费支出。

TensorRT-LLM 是 NVIDIA 开源的大语言模型推理加速库,专门解决大模型在 GPU 上跑起来又慢又贵的问题。它把底层复杂的并行计算和显存优化打包起来,让开发者用简单的 Python API 就能在 Ampere、Ada Lovelace 和 Hopper 架构 GPU 上榨干硬件性能。
它能做什么
在实际部署中,大模型推理最耗费算力也最占显存。TensorRT-LLM 针对大模型推理常见的性能瓶颈做了深度优化,最直接的效果就是大幅拉高吞吐量并压低成本。
在主流的 CNN/Daily Mail 文本摘要基准测试中,结合了 FP8 精度和动态组批(in-flight batching)的 H100 GPU,整体吞吐量达到了 A100 基础 PyTorch 模式的 8 倍。在热门的开源模型 Llama 2 上,推理速度也能提速 4.6 倍。
吞吐提升直接换来了真金白银的成本下降。根据官方测试,在跑类似 GPT-J 6B 这样的小参数模型时,相比 A100 基线,总体拥有成本(TCO)能降低 5.3 倍,电费能耗降低 5.6 倍。跑 70B 参数的 Llama 2 时,TCO 也能缩减 3 倍,能耗减少 3.2 倍。
它开箱支持当前几乎所有主流模型架构,包括 Meta 的 Llama 2、OpenAI 的 GPT-2 和 GPT-3、Falcon、Mosaic MPT 以及 BLOOM 等数十个模型。
怎么做到的
以前要把 70B 这样的大模型跑起来,开发者往往得手动拆解模型权重、手写底层代码来协调多张显卡。TensorRT-LLM 内置了张量并行机制,利用 NVLink 自动把权重矩阵拆分到多张 GPU 或多个节点上并行计算,完全不需要人工重写模型结构。
另一个关键改进是被称为 in-flight batching 的动态组批调度。大模型生成文本有长有短,传统的批处理必须等最长的那条请求生成完才能开始下一批,很多算力在干等。TensorRT-LLM 能在每个计算迭代中实时把已经完成的请求移出显存,并立刻塞进新的请求,实测能让 H100 上的吞吐直接翻倍。
在显存优化上,它结合了 Hopper 架构的 Transformer Engine,开发者不需要改动模型代码,就能直接把权重转成 FP8 格式运行。和传统的 INT8 或 INT4 相比,FP8 在大幅节省显存的同时几乎不损失模型精度。此外,它还集成了 PagedAttention 以及 FlashAttention 等最新的底层算子融合。
Databricks 工程副总裁 Naveen Rao 表示,MosaicML 已经在自家的推理服务里无缝接入了这套工具,集成的过程非常轻松。
怎么跑起来
硬件方面,TensorRT-LLM 支持 NVIDIA Ampere、Lovelace 和 Hopper 架构显卡。对于百亿参数以上的大模型,建议搭配具备高速 NVLink 的多卡环境使用。
项目目前已经完全开源,代码托管在 GitHub 的 NVIDIA/TensorRT-LLM 仓库中,同时也整合进了企业级的 NVIDIA NeMo 框架。开发者只要熟悉 Python,就可以通过模块化的 Python API 自定义模型结构、定义优化策略并启动服务,不需要深入折腾底层的 C++ 或 CUDA 细节。
为什么重要
我们认为,TensorRT-LLM 最大的价值在于大幅降低了企业自建大模型推理集群的技术门槛与运行成本。过去只有少数头部大厂有能力手写底层算子榨干硬件性能,现在任何有部署需求的技术团队都能直接拿来即用,省下了昂贵的高性能计算工程师人力支出。
对于行业来说,这不仅是 NVIDIA 在给自家新一代显卡带货,更是在大模型推理领域筑深自家的 CUDA 护城河。不过需要注意的是,这套优化深度绑定了 NVIDIA 较新的 GPU 架构,特别是 FP8 相关红利基本由 Hopper 独享。如果你正在运维基于 A100 或 H100 的推理集群,现在就非常值得把现有服务迁移上去试试。
来源 (1)
- developer.nvidia.com(在新标签页打开)原文developer.nvidia.com