通过使用AWS Inferentia2，最大化稳定扩散性能并降低推理成本
通过AWS Inferentia2，最大化扩散性能，降低推理成本
由于其在创建逼真的文本、图像、代码和音频方面的出色能力，生成式AI模型近几个月来得到了快速发展。在这些模型中，稳定扩散模型以其在基于文本提示创建高质量图像方面的独特优势而脱颖而出。稳定扩散可以生成各种高质量的图像，包括逼真的肖像、风景以及抽象艺术。与其他生成式AI模型一样，稳定扩散模型需要强大的计算能力以提供低延迟推理。
在本文中，我们将展示如何在亚马逊弹性计算云（Amazon EC2）上以最低成本运行稳定扩散模型并实现高性能，使用由AWS Inferentia2提供动力的Amazon EC2 Inf2实例。我们将查看稳定扩散模型的架构，并演示使用AWS Neuron编译稳定扩散模型并将其部署到Inf2实例的步骤。我们还将讨论Neuron SDK自动进行的优化以提高性能。您可以在AWS Inferentia2上高效地运行稳定扩散2.1和1.5版本。最后，我们将展示如何使用Amazon SageMaker将稳定扩散模型部署到Inf2实例。
稳定扩散2.1模型在32位浮点数（FP32）中的大小为5 GB，在bfoat16（BF16）中的大小为2.5 GB。单个inf2.xlarge实例具有一个带有32 GB HBM内存的AWS Inferentia2加速器。稳定扩散2.1模型可以适应单个inf2.xlarge实例。稳定扩散是一种文本到图像的模型，您可以通过提供文本提示作为输入来创建不同风格和内容的图像。要了解有关稳定扩散模型架构的更多信息，请参阅使用稳定扩散模型创建高质量图像，并通过Amazon SageMaker以低成本进行部署。
Neuron SDK如何优化稳定扩散性能
在我们可以在AWS Inferentia2实例上部署稳定扩散2.1模型之前，我们需要使用Neuron SDK编译模型组件。Neuron SDK包括深度学习编译器、运行时和工具，可以编译并自动优化深度学习模型，使其能够在Inf2实例上高效运行，并充分发挥AWS Inferentia2加速器的性能。我们在GitHub仓库上提供了稳定扩散2.1模型的示例。本文档提供了一个端到端的示例，演示了如何编译稳定扩散模型、保存编译后的Neuron模型，并将其加载到运行时进行推理。
我们使用来自Hugging Face的diffusers库的StableDiffusionPipeline来加载和编译模型。然后，我们使用torch_neuronx.trace()将模型的所有组件编译为Neuron，并将优化后的模型保存为TorchScript。编译过程可能需要相当多的内存，因此对于每个要追踪的模型，我们在追踪之前创建一个正在被追踪的管道的deepcopy。之后，我们使用del pipe从内存中删除管道对象。在内存较低的实例上编译时，这种技术特别有用。
此外，我们还对稳定扩散模型进行了优化。UNet是推理中计算密集型的部分。UNet组件在批次大小为2的输入张量上操作，并生成一个相应的输出张量，批次大小也为2，以生成一张图像。这些批次中的元素彼此完全独立。我们可以利用这种行为，在每个Neuron核心上运行一个批次以获得最佳的延迟。我们将UNet编译为一个批次（使用一个批次的输入张量），然后使用torch_neuronx.DataParallel API将此单批次模型加载到每个核心上。此API的输出是一个无缝的两批次模块：我们可以将两个批次的输入传递给UNet，并返回两个批次的输出，但在内部，两个单批次模型在两个Neuron核心上运行。这种策略优化了资源利用率并降低了延迟。
在Inf2 EC2实例上编译和部署稳定扩散模型
要在Inf2 EC2实例上编译和部署稳定扩散模型，请登录AWS管理控制台并创建一个inf2.8xlarge实例。请注意，仅在编译模型时需要inf2.8xlarge实例，因为编译需要更高的主机内存。稳定扩散模型可以托管在inf2.xlarge实例上。您可以使用以下AWS命令行界面（AWS CLI）命令查找带有Neuron库的最新AMI：
aws ec2 describe-images --region us-east-1 --owners amazon \
--filters 'Name=name,Values=Deep Learning AMI Neuron PyTorch 1.13.? (Amazon Linux 2) ????????' 'Name=state,Values=available' \
--query 'reverse(sort_by(Images, &CreationDate))[:1].ImageId' \
--output text对于这个例子，我们使用Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04)创建了一个EC2实例。然后，您可以通过连接到实例并运行以下步骤来创建一个JupyterLab实验环境：
运行 source /opt/aws_neuron_venv_pytorch/bin/activate
pip install jupyterlab
jupyter-lab包含编译和托管模型的所有步骤的笔记本位于GitHub上。
让我们来看看编译文本编码器块的步骤。Stable Diffusion管道的其他块可以类似地进行编译。
第一步是从Hugging Face加载预训练模型。`StableDiffusionPipeline.from_pretrained`方法将预训练模型加载到我们的管道对象`pipe`中。然后，我们使用`deepcopy`创建文本编码器的副本，有效地克隆它。然后使用`del pipe`命令删除原始管道对象，释放其占用的内存。这里，我们将模型量化为BF16权重：
model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
text_encoder = copy.deepcopy(pipe.text_encoder)
del pipe这一步涉及使用`NeuronTextEncoder`包装我们的文本编码器。编译文本编码器模块的输出将是`dict`类型。我们使用这个包装器将其转换为`list`类型：
text_encoder = NeuronTextEncoder(text_encoder)我们用一些值初始化PyTorch张量`emb`。张量`emb`用作`torch_neuronx.trace`函数的示例输入。该函数会跟踪我们的文本编码器并将其编译为针对Neuron进行优化的格式。编译模型的目录路径由将`COMPILER_WORKDIR_ROOT`与子目录`text_encoder`连接构成：
emb = torch.tensor([...])
text_encoder_neuron = torch_neuronx.trace(
text_encoder.neuron_text_encoder,
emb,
compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder'),
)使用`torch.jit.save`保存已编译的文本编码器。它存储在我们编译器工作区的`text_encoder`目录下的文件名为`model.pt`的文件中：
text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder/model.pt')
torch.jit.save(text_encoder_neuron, text_encoder_filename)该笔记本包括编译模型的类似步骤：UNet、VAE解码器和VAE`post_quant_conv`。在编译完所有模型之后，您可以按照以下步骤加载和运行模型：
- 定义编译模型的路径。
- 加载预训练的`StableDiffusionPipeline`模型，并指定其配置以使用bfloat16数据类型。
- 使用`torch_neuronx.DataParallel`API将UNet模型加载到两个Neuron核上。这允许进行数据并行推断，可以显著提高模型的性能。
- 将模型的其他部分（`text_encoder`、`decoder`和`post_quant_conv`）加载到单个Neuron核上。
然后，您可以通过提供输入文本作为提示来运行流水线。以下是模型为这些提示生成的一些图片：
- 雷诺·塞尚的肖像，钢笔和墨水，精细线条的绘画，作者：克雷格·穆林斯、鲁安·贾、肯塔罗·米乌拉、格雷格·鲁特科夫斯基、loundraw
- 19世纪煤矿工人的肖像，美丽的绘画，由Greg Rutkowski绘制的高度详细人脸绘画
- 森林中的一座城堡
在AWS Inferentia2和SageMaker上托管稳定扩散2.1
使用SageMaker托管稳定扩散模型还需要使用Neuron SDK进行编译。您可以提前完成编译，也可以在运行时使用Large Model Inference（LMI）容器进行编译。提前编译可以加快模型加载速度，是首选选项。
SageMaker LMI容器提供两种部署模型的方式：
- 无代码选项，我们只需提供一个包含所需配置的
serving.properties文件 - 自定义推理脚本
我们将介绍这两种解决方案，并介绍配置和推理脚本（model.py）。在本文中，我们演示了使用存储在Amazon Simple Storage Service（Amazon S3）存储桶中的预编译模型进行部署。您可以将此预编译模型用于您的部署。
使用提供的脚本配置模型
在本节中，我们将展示如何配置LMI容器以托管稳定扩散模型。SD2.1笔记本在GitHub上可用。第一步是根据以下目录结构创建模型配置包。我们的目标是使用托管模型所需的最小模型配置。所需的目录结构如下：
<config-root-directory> /
├── serving.properties
│
└── model.py [OPTIONAL]接下来，我们使用以下参数创建serving.properties文件：
%%writefile code_sd/serving.properties
engine=Python
option.entryPoint=djl_python.transformers-neuronx
option.use_stable_diffusion=True
option.model_id=s3url
option.tensor_parallel_degree=2
option.dtype=bf16这些参数指定了以下内容：
- option.model_id – LMI容器使用s5cmd从S3位置加载模型，因此我们需要指定编译权重所在的位置。
- option.entryPoint – 为了使用内置处理程序，我们指定transformers-neuronx类。如果您有自定义推理脚本，则需要提供相应的脚本。
- option.dtype – 这指定以特定大小加载权重。在本文中，我们使用BF16，这进一步降低了内存需求和延迟。
- option.tensor_parallel_degree – 此参数指定我们为此模型使用的加速器数量。AWS Inferentia2芯片加速器有两个Neuron核心，因此指定值为2表示我们使用一个加速器（两个核心）。这意味着我们现在可以创建多个工作进程，以增加端点的吞吐量。
- option.engine – 设置为Python，表示我们将不使用DeepSpeed或Faster Transformer等其他编译器进行托管。
自定义推理脚本
如果您想使用自定义的推理脚本，需要从serving.properties中删除option.entryPoint。在这种情况下，LMI容器将在与serving.properties相同的位置查找model.py文件，并使用该文件运行推理。
创建自己的推理脚本（model.py）
使用LMI容器创建自己的推理脚本相对简单。容器要求您的model.py文件中有以下方法的实现：
def handle(inputs: Input) -> Outputs让我们来查看一下附带的笔记本中的一些关键部分，它演示了自定义脚本功能。
将cross_attention模块替换为优化版本：
# 用自定义的交叉注意力模块替换原始的交叉注意力模块，以提高性能
CrossAttention.get_attention_scores = get_attention_scores
加载以下编译后的权重文件
text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder.pt')
decoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'vae_decoder.pt')
unet_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'unet.pt')
post_quant_conv_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'vae_post_quant_conv.pt')这些是我们在创建编译时使用的编译后的权重文件的名称。可以自由更改文件名，但请确保您的权重文件名与您在此处指定的名称匹配。
然后，我们需要使用Neuron SDK加载它们，并将它们设置为实际模型的权重。当加载优化后的UNet权重时，请注意我们还指定了要加载到的Neuron核心数量。在这里，我们将其加载到具有两个核心的单个加速器上：
# 将编译的UNet加载到两个Neuron核心上。
pipe.unet = NeuronUNet(UNetWrap(pipe.unet))
logging.info(f"Loading model: unet:created")
device_ids = [idx for idx in range(tensor_parallel_degree)]
pipe.unet.unetwrap = torch_neuronx.DataParallel(torch.jit.load(unet_filename), device_ids, set_dynamic_batching=False)
# 将其他编译的模型加载到单个Neuron核心上。
# - 加载编码器
pipe.text_encoder = NeuronTextEncoder(pipe.text_encoder)
clip_compiled = torch.jit.load(text_encoder_filename)
pipe.text_encoder.neuron_text_encoder = clip_compiled
#- 加载解码器
pipe.vae.decoder = torch.jit.load(decoder_filename)
pipe.vae.post_quant_conv = torch.jit.load(post_quant_conv_filename)使用提示运行推理将调用pipe对象生成一张图片。
创建SageMaker端点
我们使用Boto3 API创建SageMaker端点。完成以下步骤：
- 使用仅包含serving和可选
model.py文件的tar包，并将其上传到Amazon S3。 - 使用镜像容器和先前上传的模型tar包创建模型。
- 使用以下重要参数创建端点配置：
- 使用
ml.inf2.xlarge实例。 - 将
ContainerStartupHealthCheckTimeoutInSeconds设置为240，以确保在部署模型后进行健康检查。 - 将
VolumeInGB设置为较大的值，以便用于加载大小为32 GB的模型权重。
- 使用
创建SageMaker模型
在创建model.tar.gz文件并将其上传到Amazon S3之后，我们需要创建一个SageMaker模型。我们使用LMI容器和上一步的模型工件来创建SageMaker模型。SageMaker允许我们自定义和注入各种环境变量。对于此工作流程，我们可以将所有内容保持默认设置。请参阅以下代码：
inference_image_uri = (
f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0 djl-serving-inf2"
)创建模型对象，实际上创建了一个锁定的容器，该容器加载到实例上并用于推理：
model_name = name_from_base(f"inf2-sd")
create_model_response = boto3_sm_client.create_model(
ModelName=model_name,
ExecutionRoleArn=role,
PrimaryContainer={"Image": inference_image_uri, "ModelDataUrl": s3_code_artifact},
)创建SageMaker端点
在此示例中，我们使用一个ml.inf2.xlarge实例。我们需要设置VolumeSizeInGB参数来提供加载模型和权重所需的磁盘空间。此参数适用于支持Amazon Elastic Block Store（Amazon EBS）卷附加的实例。我们可以将模型下载超时和容器启动健康检查设置为较高的值，这将为容器从Amazon S3拉取权重并加载到AWS Inferentia2加速器提供足够的时间。有关更多详细信息，请参阅CreateEndpointConfig。
endpoint_config_response = boto3_sm_client.create_endpoint_config(
EndpointConfigName=endpoint_config_name,
ProductionVariants=[
{
"VariantName": "variant1",
"ModelName": model_name,
"InstanceType": "ml.inf2.xlarge", # -
"InitialInstanceCount": 1,
"ContainerStartupHealthCheckTimeoutInSeconds": 360,
"VolumeSizeInGB": 400
},
],
)最后，我们创建一个SageMaker端点：
create_endpoint_response = boto3_sm_client.create_endpoint(
EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name
)调用模型端点
这是一个生成模型，所以我们传入模型用于生成图像的提示。有效载荷的类型为JSON：
response_model = boto3_sm_run_client.invoke_endpoint(
EndpointName=endpoint_name,
Body=json.dumps(
{
"prompt": "Mountain Landscape",
"parameters": {} #
}
),
ContentType="application/json",
)在Inf2上对稳定扩散模型进行基准测试
我们进行了一些基准测试，使用BF 16数据类型在Inf2上对稳定扩散模型进行测试，并能够得到与其他加速器相媲美或超过的延迟数据。再加上AWS Inferentia2芯片的低成本，使其成为一个非常有价值的选择。
以下数据来自部署在inf2.xl实例上的稳定扩散模型。有关成本的更多信息，请参阅Amazon EC2 Inf2实例。
| 模型 稳定扩散 1.5 稳定扩散 1.5 稳定扩散 1.5 稳定扩散 1.5 | 分辨率 512×512 768×768 512×512 768×768 | 数据类型 bf16 bf16 bf16 bf16 | 迭代次数 50 50 30 30 | P95延迟（毫秒） 2,427.4 8,235.9 1,456.5 4,941.6 | Inf2.xl按需每小时成本 $0.76 $0.76 $0.76 $0.76 | Inf2.xl（每张图像的成本） $0.0005125 $0.0017387 $0.0003075 $0.0010432 |
| 稳定扩散 2.1 稳定扩散 2.1 稳定扩散 2.1 稳定扩散 2.1 | 512×512 768×768 512×512 768×768 | bf16 bf16 bf16 bf16 | 50 50 30 30 | 1,976.9 6,836.3 1,186.2 4,101.8 | $0.76 $0.76 $0.76 $0.76 | $0.0004174 $0.0014432 $0.0002504 $0.0008659 |
结论
在本文中，我们深入探讨了使用Inf2实例对Stable Diffusion 2.1模型进行编译、优化和部署的过程。我们还演示了使用SageMaker部署Stable Diffusion模型的方法。Inf2实例还为Stable Diffusion 1.5提供了出色的性价比。要了解更多关于Inf2实例为生成式人工智能和大型语言模型的优势，请参阅Amazon EC2 Inf2实例现已正式推出，可提供低成本、高性能的生成式人工智能推理。有关性能细节，请参阅Inf2性能。在GitHub仓库中还有更多示例可供参考。
特别感谢Matthew Mcclain、Beni Hegedus、Kamran Khan、Shruti Koparkar和Qing Lan对审阅和提供宝贵意见。
