TensorRT-LLM:英伟达推出的专为提升大模型推理速度优化的全新框架
随着大型语言模型(LLM)如 GPT-3 和 BERT 在 AI 领域的崛起,如何在实际应用中高效地进行模型推断成为了一个关键问题。为此,英伟达推出了全新的大模型推理提速框架TensorRT-LM,可以将现有的大模型推理速度提升4倍!2016年,英伟达已经推出了TensorRT,此次发布的TensorRT-LM是在TensorRT基础上针对大模型进一步优化的加速推理库。

TensorRT简介
TensorRT是英伟达的一个深度学习模型优化器和运行时库,它可以将深度学习模型转换为优化的格式,从而在英伟达GPU上实现更快的推断速度。TensorRT的第一个版本是在2016年11月发布的,当时叫做GPU Inference Engine (GIE)。后来在2017年3月,英伟达将其改名为TensorRT,并发布了TensorRT 2.0版本。从那以后,英伟达一直不断更新和改进TensorRT。

上图是英伟达官方针对TensorRT的示意图,可以看到,TensorRT定位的是将训练结束的模型优化达到加速目的,因此与你使用的训练框架和训练过程不强相关。TensorRT通过自动识别可以合并的连续层,并将它们融合成一个操作。这减少了在 GPU 上的操作数量,从而提高了执行速度。
TensorRT-LLM简介
TensorRT 是 NVIDIA 的一个深度学习模型优化器和运行时库,旨在为深度学习模型在 NVIDIA GPU 上提供快速、高效的推断。而 TensorRT-LLM 则是其针对大型语言模型的扩展,提供了一系列专门的优化和功能。
随着 LLM 的模型参数数量不断增加,传统的推断方法在性能和成本上都面临挑战。TensorRT-LLM 提供了一种方法,可以在保持模型准确性的同时,大大提高推断速度并降低成本。
TensorRT-LLM 首先解析模型结构,然后应用一系列优化技术,如层融合、精度校准和内核选择。它还利用了并行化技术,如张量并行性,以在多个 GPU 之间分配模型的不同部分。
根据官方的说明,TensorRT-LLM的主要特点:
专为 LLM 设计: 与标准的 TensorRT 不同,TensorRT-LLM 针对大型语言模型的特定需求和挑战进行了优化。
集成优化: NVIDIA 与多家领先公司合作,将这些优化集成到了 TensorRT-LLM 中,以确保 LLM 在 NVIDIA GPU 上的最佳性能。
模块化 Python API: TensorRT-LLM 提供了一个开源的模块化 Python API,使开发者能够轻松定义、优化和执行新的 LLM 架构和增强功能。
飞行批处理(In-flight batching): 这是一种优化的调度技术,可以更有效地处理动态负载。它允许 TensorRT-LLM 在其他请求仍在进行时开始执行新请求,从而提高 GPU 利用率。
支持新的 FP8 数据格式: 在 H100 GPU 上,TensorRT-LLM 支持新的 FP8 数据格式,这可以大大减少内存消耗,同时保持模型的准确性。
广泛的模型支持: TensorRT-LLM 包括了许多在今天生产中广泛使用的 LLM 的完全优化、即用版本,如 Meta Llama 2、OpenAI GPT-2 和 GPT-3 等。
并行化和分布式推断: TensorRT-LLM 利用张量并行性进行模型并行化,这使得模型可以在多个 GPU 之间并行运行,从而实现大型模型的高效推断。
优化的内核和操作: TensorRT-LLM 包括了针对 LLM 的优化内核和操作,如 FlashAttention 和遮蔽多头注意力等。
简化的开发流程: TensorRT-LLM 旨在简化 LLM 的开发和部署过程,使开发者无需深入了解底层的技术细节。
TensorRT、TensorRT-LLM与PyTorch之间的关系和差异
TensorRT、TensorRT-LLM 和 PyTorch 之间的关系可以从以下几个方面来描述:
PyTorch是一个开源的深度学习框架,广泛用于研究和开发。开发者使用 PyTorch 进行模型的设计、训练和初步推断。
而TensorRT 是 NVIDIA 的一个深度学习模型优化器和运行时库。它的主要目的是将深度学习模型(如从 PyTorch 导出的模型)转换为优化的格式,以在 NVIDIA GPU 上实现更快的推断速度。
开发者通常首先使用 PyTorch 训练他们的模型,然后使用 TensorRT 进行优化,以便在生产环境中实现更快的模型推断。
TensorRT-LLM 是 TensorRT 的一个扩展,专门针对大型语言模型(LLM)的推断进行优化。它提供了一系列工具和功能,使得大型语言模型在 NVIDIA GPU 上的推断更为高效。
TensorRT-LLM与TensorRT的关系:TensorRT-LLM基于TensorRT,但增加了针对大型语言模型的特定优化和功能。
TensorRT-LLM与PyTorch的关系: 就像TensorRT一样,开发者可以首先使用 PyTorch 训练大型语言模型,然后使用 TensorRT-LLM 进行优化,以便在生产环境中实现更快的推断。
总的来说,PyTorch、TensorRT 和 TensorRT-LLM 之间的关系是一个连续的流程:开发者首先使用 PyTorch 进行模型的设计和训练,然后使用 TensorRT 或其扩展 TensorRT-LLM 进行优化,以在 NVIDIA GPU 上实现高效的模型推断。
TensorRT-LLM可以直接解析某些框架(如PyTorch、TF)的模型或通过 ONNX 格式导入模型,然后通过一系列优化操作来提升模型推理速度。
TensorRT-LLM的加速结果测试
官方给出了GPT-J 6B在TensorRT-LLM加持下的模型推理速度提升结果,如下图所示:

GPT-J 6B是一个由EleutherAI研究小组创建的开源自回归语言模型。它是OpenAI的GPT-3的最先进替代品之一,在各种自然语言任务(如聊天、摘要和问答等)方面表现良好。
上图使用A100作为GPT-J 6B的推理速度基准,H100的推理速度是A100的4倍,而使用了TensorRT-LLM之后的H100推理速度是A100的8倍!提升速度惊人!
而LLaMA2也有很高的速度提升,如下图所示:

使用A100作为LLaMA2 7B的推理速度(文本摘要)基准,H100的推理速度是A100的2.6倍,而使用了TensorRT-LLM之后的H100推理速度是A100的4.6倍!提升速度也是非常惊人!
TensorRT-LLM支持的硬件和模型
目前官方的说明与TensorRT类似,支持大多数的开源框架的训练结果模型,包括PyTorch、TF等。同时,关于它支持的硬件的信息如下:
NVIDIA H100 GPUs: TensorRT-LLM 在 NVIDIA H100 GPU 上提供了显著的性能提升。H100 GPU 支持新的 FP8 数据格式,这使得模型可以在更低的精度下运行,从而减少内存消耗,同时保持模型的准确性。
NVIDIA A100 GPUs: 在某些性能比较中,TensorRT-LLM 也在 NVIDIA A100 GPU 上进行了测试。但从结果来看,H100 在性能上明显优于 A100。
除了上述明确提到的 GPU,TensorRT-LLM 可能还支持其他 NVIDIA GPU,但具体支持的硬件列表可能需要查阅 NVIDIA 的官方文档或相关资料。
总的来说,TensorRT-LLM 主要针对 NVIDIA 的高端 GPU 进行了优化,以实现大型语言模型的高效推断。
目前TensorRT-LLM属于早期预览,只要注册成为NVIDIA开发者即可申请使用~
欢迎大家关注DataLearner官方微信,接受最新的AI技术推送
