随着能够创建类人文本、图像、代码和音频的生成式人工智能模型的兴起，人工智能（AI）和机器学习（ML）的世界见证了模式的转变。与传统的机器学习模型相比，生成式AI模型要大得多，也复杂得多。然而，这种模型增加了复杂性，也带来了高昂的推理成本，以及对强大计算资源日益增长的需求。对于资源有限的企业和研究人员来说，生成式AI模型的高推理成本可能会成为进入市场的障碍，因此需要更高效、更具成本效益的解决方案。此外，大多数生成式AI使用案例都涉及人机交互或真实世界场景，因此需要能提供低延迟性能的硬件。亚马逊云科技一直在利用专用芯片进行创新，以满足对功能强大、高效且经济实惠的计算硬件的日益增长的需求。
近日，亚马逊云科技宣布Amazon SageMaker支持基于AWS Inferentia2（ml.inf2）和AWS Trainium（ml.trn1）的SageMaker实例，以托管用于实时推理和异步推理的生成式AI模型。ml.inf2实例可在美国东部（俄亥俄）的SageMaker上部署模型，ml.trn1实例可在美国东部（弗吉尼亚北部）的SageMaker上部署模型。
现在可以在SageMaker上使用这些实例，以低成本实现生成式AI模型的高性能，包括大型语言模型（LLM）、Stable Diffusion和Vision Transformers。此外，还可以使用Amazon SageMaker Inference Recommender来协助您运行负载测试，并评估在这些实例上部署模型的性价比优势。并且可以使用ml.inf2和ml.trn1实例在SageMaker上运行机器学习应用程序，以实现文本摘要、代码生成、视频和图像生成、语音识别、个性化、欺诈侦测等。要轻松上手，可以在配置SageMaker端点时首先指定ml.trn1或ml.inf2实例。然后将兼容ml.trn1和ml.inf2的AWS Deep Learning Containers（DLC）用于PyTorch、TensorFlow、Hugging Face和大型模型推理（LMI）。
此篇文章将展示利用LMI容器，在不需要任何额外编码的情况下，使用SageMaker在AWS Inferentia2上部署大型语言模型的过程。使用GPT4ALL-J，这是一种经过微调的GPT-J 7B模型，可提供聊天机器人风格的互动。
ml.trn1和ml.inf2实例概述
ml.trn1实例由Trainium加速器提供支持，该加速器主要用于生成式AI模型（包括LLM）的高性能深度学习训练。不过，这些实例也支持比Inf2所适合的模型更大的推理工作负载。最大的实例（trn1.32xlarge实例）在单个实例中配备了16个Trainium加速器和512 GB加速器内存，可提供高达3.4 petaflops的FP16/BF16计算能力。16个Trainium加速器通过超高速NeuronLinkv2进行连接，这样可简化集体通信。
ml.Inf2实例由AWS Inferentia2加速器提供支持，这是一款专为推理而构建的加速器。与第一代AWS Inferentia相比，这款加速器的计算性能提高了三倍，吞吐量提高了四倍，延迟降低了多达10倍。最大的实例Inf2.48xlarge在单个实例中配备了12个AWS Inferentia2加速器和384 GB加速器内存，BF16/FP16的综合计算能力为2.3 petaflops。这使您能够在单个实例中部署包含多达1750亿个参数的模型。Inf2是唯一提供这种互连功能的推理优化实例，而这种功能只有在更昂贵的训练实例中才有。对于单个加速器无法容纳的超大型模型，数据可通过NeuronLink直接在加速器之间流动，完全绕过CPU。借助NeuronLink，Inf2支持更快的分布式推理，并提高吞吐量、降低延迟。
AWS Inferentia2和Trainium加速器都有两个NeuronCores-v2、32 GB HBM内存堆栈和专用的集体计算引擎，在进行多加速器推理时，通过重叠计算和通信自动优化运行时系统。
AWS Neuron SDK
AWS Neuron是用于在基于AWS Inferentia和Trainium的实例上运行深度学习工作负载的SDK。AWS Neuron包括深度学习编译器、运行时系统和原生集成到TensorFlow和PyTorch中的工具。利用Neuron，可以在ml.trn1和ml.inf2上开发、分析和部署高性能机器学习工作负载。
Neuron编译器接受各种格式的机器学习模型（TensorFlow、PyTorch、XLA HLO），并对这些模型进行优化，以便在Neuron设备上运行。在机器学习框架内调用Neuron编译器，而在该框架内，机器学习模型由Neuron框架插件发送给编译器。由此产生的编译器构件称为NEFF文件（Neuron可执行文件格式），Neuron运行时系统会将该文件加载到Neuron设备中。
Neuron运行时系统由内核驱动程序和C/C++库组成，这些库提供API来访问AWS Inferentia和Trainium Neuron设备。适用于TensorFlow和PyTorch的Neuron机器学习框架插件使用Neuron运行时系统在NeuronCore上加载和运行模型。Neuron运行时系统将编译好的深度学习模型（NEFF）加载到Neuron设备上，并针对高吞吐量和低延迟进行优化。
使用SageMaker ml.inf2实例托管NLP模型
transformers-neuronx是一个开源库，可将模型的大型权重矩阵分片到多个NeuronCore上，在深入研究如何使用该库为LLM提供服务之前，先简单了解一下可用于单个NeuronCore的模型的典型部署流程。
查看支持的模型列表，确保AWS Inferentia2支持该模型。接下来，需要使用Neuron编译器对模型进行预编译。可以使用SageMaker notebook或Amazon Elastic Compute Cloud（Amazon EC2）实例来编译模型。可以借助SageMaker Python SDK，使用PyTorch等流行的深度学习框架部署模型。可以将模型部署到SageMaker托管服务，并获得可用于推理的端点。这些端点是完全托管的，支持自动扩缩。
使用SageMaker ml.inf2实例托管LLM
大型语言模型通常具有数十亿个参数，规模太大，单个加速器无法容纳。这就需要使用模型并行技术，在多个加速器上托管LLM。托管LLM的另一个关键要求是实施高性能的模型服务解决方案。该解决方案应该能够高效地加载模型、管理分区，并通过HTTP端点无缝地处理请求。
SageMaker包括专门的Deep Learning Containers（DLC）、库和工具，用于模型并行化和大型模型推理。SageMaker使用流行的开源库维护DLC，以便在亚马逊云科技基础设施上托管GPT、T5、OPT、BLOOM和Stable Diffusion等大型模型。这些专用的DLC称为SageMaker LMI容器。
SageMaker LMI容器使用DJLServing，这是一种与transformers-neuronx库集成的模型服务器，可支持NeuronCore之间的张量并行。DJL模型服务器和transformers-neuronx库是容器的核心组件，其中还包括Neuron SDK。这种设置便于将模型加载到AWS Inferentia2加速器上，在多个NeuronCore上并行处理模型，并通过HTTP端点提供服务。
LMI容器支持从Amazon Simple Storage Service（Amazon S3）存储桶或Hugging Face Hub加载模型。默认的处理程序脚本会加载模型，将模型编译并转换为Neuron优化格式，然后再次加载模型。要使用LMI容器托管LLM，有两种选择：
无代码（首选）——这是使用LMI容器部署LLM的最简单方法。在这种方法中，您可以使用提供的默认处理程序，只需传递模型名称和serving.properties文件中所需的参数，即可加载和托管模型。要使用默认处理程序，我们需要将entryPoint参数设置为djl_python.transformers-neuronx。
自带脚本——在这种方法中，您可以选择创建自己的model.py文件，其中包含加载和服务模型所需的代码。该文件充当DJLServing API 和transformers-neuronx API之间的中介。要自定义模型加载过程，可以为serving.properties提供可配置的参数。
运行时系统架构
tensor_parallel_degree属性值决定了张量并行模块在多个NeuronCore上的分布。例如，inf2.24xlarge有六个AWS Inferentia2加速器。每个AWS Inferentia2加速器都有两个NeuronCore。每个NeuronCore都有一个16 GB的专用高带宽内存（HBM），用于存储张量并行模块。当张量并行度为4时，LMI将为同一模型分配三个模型副本，每个副本使用4个NeuronCore。如下图所示，当LMI容器启动时，将首先在CPU可寻址内存中加载和跟踪模型。跟踪完成后，根据张量并行度跨NeuronCore对模型进行分区。
LMI使用DJLServing作为其模型服务堆栈。在SageMaker中通过容器的运行状况检查后，容器即可处理推理请求。DJLServing启动相当于TOTAL NUMBER OF NEURON CORES/TENSOR_PARALLEL_DEGREE的多个Python进程。每个Python进程都包含相当于TENSOR_PARALLEL_DEGREE的C++线程。每个C++线程在一个NeuronCore上保存一个模型分片。
当使用多个独立请求调用服务器时，许多实践者（Python进程）倾向于按顺序运行推理。尽管设置起来更容易，但利用加速器的计算能力通常不是最佳实践。为解决这一问题，DJLServing提供了动态批处理的内置优化功能，可在服务器端将这些独立的推理请求动态合并成一个更大的批处理，以提高吞吐量。所有请求都先到达动态批处理器，然后才进入实际的作业队列等待推理。您可以使用serving.properties中的batch_size设置，为动态批处理设置首选的批处理大小。您还可以配置max_batch_delay，根据延迟要求指定批处理器中等待其他请求加入批处理的最长延迟时间。吞吐量还取决于模型副本的数量和容器中启动的Python进程组。如下图所示，当张量并行度设置为4时，LMI容器会启动三个Python进程组，每个进程组都包含模型的完整副本。这使您可以增加批处理大小，获得更高的吞吐量。
用于部署LLM的SageMaker notebook
在本节中，逐步演示如何部署GPT4All-J，这是一种包含60亿个参数的模型，采用FP32时有24 GB。GPT4All-J是一款流行的聊天机器人，接受过的训练包括单词问题、对话、代码、诗歌、歌曲和故事等各种交互内容。GPT4all-J是一种经过微调的GPT-J模型，可产生与人类互动类似的响应。
GitHub上提供完整的notebook示例。可以使用SageMaker Python SDK将模型部署到Inf2实例。我们使用提供的默认处理程序来加载模型。这样，我们只需提供一个servings.properties文件。此文件具有DJL模型服务器下载和托管模型所需的配置。我们可以使用model_id参数指定Hugging Face模型的名称，以便直接从Hugging Face存储库下载模型。或者，您也可以通过提供s3url参数从Amazon S3下载模型。entryPoint参数配置为指向用于加载模型的库。
tensor_parallel_degree属性值决定了张量并行模块在多个设备上的分布。例如，如果有12个NeuronCore，张量并行度为4，那么LMI将分配3个模型副本，每个副本使用4个NeuronCore。您还可以使用属性dtype定义精度类型。n_position参数定义了模型的最大输入和输出序列长度之和。
总结
综上所述，此文展示了SageMaker新推出的功能，它现在支持ml.inf2和ml.trn1实例来托管生成式AI模型。并且演示了如何在不编写任何代码的情况下，使用SageMaker和LMI容器在AWS Inferentia2上部署生成式AI模型GPT4ALL-J。还展示了如何使用DJLServing和transformers-neuronx加载模型、对模型进行分区和提供服务。