今天，Amazon SageMaker推出了新版本（0.25.0）的Large Model Inference (LMI) Deep Learning Containers (DLCs)，并且增加了对NVIDIA的TensorRT-LLM Library的支持。通过这些升级，您可以轻松地使用最先进的工具在SageMaker上优化大型语言模型（LLMs），并获得价格性能方面的好处 – 相较于以往版本，Amazon SageMaker LMI TensorRT-LLM DLC将Llama2-70B、Falcon-40B和CodeLlama-34B模型的延迟平均降低了33%，吞吐量平均提高了60%。
LLMs在各种应用中出现了前所未有的增长。然而，这些模型往往过大，无法适应单个加速器或GPU设备，使得实现低延迟推理和扩展性变得困难。SageMaker提供LMI DLCs来帮助您最大程度地利用可用资源并提高性能。最新的LMI DLCs提供连续批处理支持推理请求以提高吞吐量，高效推理集合操作以提高延迟，Paged Attention V2（提高具有较长序列长度的工作负载性能）以及来自NVIDIA的最新TensorRT-LLM库以实现GPU上的最大性能。LMI DLCs提供了一个低代码界面，通过仅需要模型ID和可选模型参数即可简化TensorRT-LLM的编译；所有构建TensorRT-LLM优化模型和创建模型存储库所需的繁重工作都由LMI DLC管理。此外，您可以使用LMI DLCs提供的最新量化技术（GPTQ，AWQ和SmoothQuant）。因此，借助SageMaker上的LMI DLCs，您可以加速生成式AI应用的价值实现时间，并为您选择的硬件优化LLMs，实现最佳的性价比。
在本文中，我们深入探讨了最新发布的LMI DLCs的新功能，讨论了性能基准，并概述了将LLMs与LMI DLCs一起部署以最大限度提高性能并降低成本所需的步骤。
SageMaker LMI DLCs的新功能
在本节中，我们将讨论三个SageMaker LMI DLCs的新功能。
SageMaker LMI现在支持TensorRT-LLM
SageMaker现在作为最新LMI DLC版本（0.25.0）的一部分提供了NVIDIA的TensorRT-LLM，使得在使用NVIDIA GPU时能够实现像SmoothQuant，FP8和连续批处理等最先进的优化。TensorRT-LLM为实现超低延迟的体验打开了大门，极大地提高了性能。TensorRT-LLM SDK支持从单GPU到多GPU配置的部署，通过张量并行等技术可以获得额外的性能提升。要使用TensorRT-LLM库，选择可用的LMI DLCs中的TensorRT-LLM DLC，并设置engine=MPI和其他设置，如option.model_id。以下是TensorRT-LLM技术堆栈的示意图。
高效推理集合操作
在常规的LLMs部署中，模型参数分布在多个加速器上，以满足无法安装在单个加速器上的大型模型的要求。通过启用每个加速器并行进行部分计算，可以提高推理速度。然后，在这些过程的末尾引入一个集合操作来合并这些部分结果，并在加速器之间重新分配。
对于P4D实例类型，SageMaker实现了一种新的集合操作，加快了GPU之间的通信。因此，与以前的版本相比，最新的LMI DLCs在延迟和吞吐量方面都具有更低的性能。此外，此功能已经内置于SageMaker LMI DLCs中，无需任何配置即可使用该功能，该功能专为Amazon SageMaker独家提供。
量化支持
SageMaker LMI DLC现在支持最新的量化技术，包括使用GPTQ预量化模型、激活感知权重量化（AWQ）和SmoothQuant的即时量化。
GPTQ允许LMI运行来自Hugging Face的流行INT3和INT4模型。它提供了在单个GPU/多个GPU上容纳的最小模型权重。LMI DLC还支持AWQ推理，这可以加快推理速度。最后，LMI DLC现在支持SmoothQuant，它通过INT8量化来减少模型的内存占用和计算成本，同时几乎不会降低模型的准确性。目前，我们允许您对SmoothQuant模型进行即时转换，无需任何额外步骤。但是，GPTQ和AWQ需要使用数据集进行量化才能与LMI DLC一起使用。您也可以选择使用受欢迎的预量化的GPTQ和AWQ模型在LMI DLC上使用。要使用SmoothQuant，请在serving.properties中设置option.quantize=smoothquan以及engine=DeepSpeed。一份使用SmoothQuant在ml.g5.12xlarge上托管GPT-Neox的示例笔记本位于GitHub上。
使用SageMaker LMI DLCs
您可以使用新的LMI DLCs 0.25.0在SageMaker上部署您的LLMs，而无需对代码进行任何更改。SageMaker LMI DLCs使用DJL serving为您的模型提供推理服务。要开始使用，请创建一个配置文件，指定诸如模型并行化和推理优化库等设置。有关使用SageMaker LMI DLC的说明和教程，请参阅模型并行处理和大型模型推理以及我们的可用SageMaker LMI DLCs清单。
DeepSpeed容器包含一个名为LMI分布式推理库（LMI-Dist）的库。LMI-Dist是一个推理库，用于在vLLM、文本生成推理（最高版本为0.9.4）、FasterTransformer和DeepSpeed框架中运行大型模型推理的最佳优化。该库采用了FlashAttention、PagedAttention、FusedKernel和高效的GPU通信内核等开源流行技术，以加速模型并减少内存消耗。
TensorRT LLM是NVIDIA于2023年10月发布的一个开源库。我们对TensorRT-LLM库进行了推理加速的优化，并创建了一个工具包，以支持即时模型转换。该工具包使用户能够提供Hugging Face模型ID并端到端地部署模型。它还支持连续批处理和流式处理。对于Llama-2 7B和13B模型，编译时间约为1-2分钟，70B模型约为7分钟。如果您希望在SageMaker端点设置和实例扩展时避免此编译开销，建议使用我们的教程来进行提前（AOT）编译以准备好模型。我们还接受任何可与LMI DLC一起使用的专为Triton Server构建的TensorRT LLM模型。
性能基准结果
我们将最新版本（0.25.0）的SageMaker LMI DLC与上一个版本（0.23.0）的性能进行了比较。我们对Llama-2 70B、Falcon 40B和CodeLlama 34B模型进行了实验，以展示使用TensorRT-LLM和高效推理集体操作（在SageMaker上可用）的性能提升。
SageMaker LMI容器附带了一个默认的处理程序脚本，用于加载和托管模型，提供了一个低代码选项。如果需要对模型加载步骤进行任何自定义，您还可以选择使用自己的脚本。您需要在serving.properties文件中传递所需的参数。该文件包含了Deep Java Library (DJL)模型服务器下载和托管模型所需的配置。下面的代码是用于我们的部署和基准测试的serving.properties：
engine=MPIoption.use_custom_all_reduce=true option.model_id={{s3url}}option.tensor_parallel_degree=8option.output_formatter=jsonoption.max_rolling_batch_size=64option.model_loading_timeout=3600engine参数用于定义DJL模型服务器的运行时引擎。我们可以使用model_id参数指定Hugging Face模型ID或Amazon Simple Storage Service (Amazon S3)模型的位置。任务参数用于定义自然语言处理（NLP）任务。tensor_parallel_degree参数设置张量并行模块分布的设备数。对于启用了NVLink的GPU实例，将use_custom_all_reduce参数设置为true以加速模型推理。您可以为P4D、P4de、P5和其他连接了NVLink的GPU设置此参数。output_formatter参数设置输出格式。max_rolling_batch_size参数设置最大并发请求限制。model_loading_timeout设置下载和加载模型用于服务推理的超时值。有关配置选项的详细信息，请参阅配置和设置。
Llama-2 70B
以下是Llama-2 70B的性能比较结果。并发数为16时，延迟减少了28%，吞吐量增加了44%，使用了新的LMI TensorRT LLM DLC。
Falcon 40B
以下是Falcon 40B的性能比较结果。并发数为16时，延迟减少了36%，吞吐量增加了59%，使用了新的LMI TensorRT LLM DLC。
CodeLlama 34B
以下是CodeLlama 34B的性能比较结果。并发数为16时，延迟减少了36%，吞吐量增加了77%，使用了新的LMI TensorRT LLM DLC。
建议的配置和容器用于托管LLMs
在最新版本中，SageMaker提供了两个容器：0.25.0-deepspeed和0.25.0-tensorrtllm。DeepSpeed容器包含DeepSpeed，即LMI分布式推理库。TensorRT-LLM容器包含NVIDIA的TensorRT-LLM库用于加速LLM推理。
我们建议使用下面的图示配置进行部署。
要开始使用，请参考示例笔记本：
结论
在本文中，我们展示了如何使用SageMaker LMI DLCs优化LLMs，以符合您的业务用例并实现价格性能优势。要了解更多关于LMI DLC功能的信息，请参考模型并行性和大型模型推理。我们非常期待看到您如何使用来自Amazon SageMaker的这些新功能。