使用 AWS Inferentia2 最大限度地提高稳定扩散性能并降低推理成本 | 亚马逊网络服务

由于其在创建真实文本、图像、代码和音频方面令人印象深刻的能力，生成式人工智能模型近几个月来一直在快速增长。 在这些模型中，稳定扩散模型因其在根据文本提示创建高质量图像方面的独特优势而脱颖而出。 稳定扩散可以生成各种高质量图像，包括逼真的肖像、风景，甚至抽象艺术。 而且，与其他生成式 AI 模型一样，稳定扩散模型需要强大的计算能力来提供低延迟推理。

在这篇文章中，我们将展示如何运行稳定扩散模型并以最低的成本实现高性能 亚马逊弹性计算云 （亚马逊 EC2）使用 Amazon EC2 Inf2实例 powered by AWS Inferentia2。 我们研究稳定扩散模型的架构，并逐步完成使用以下命令编译稳定扩散模型的步骤： AWS 神经元 并将其部署到 Inf2 实例。 我们还讨论了 Neuron SDK 自动进行的优化以提高性能。 您可以经济高效地在 AWS Inferentia2.1 上运行 Stable Diffusion 1.5 和 2 版本。 最后，我们展示了如何将稳定扩散模型部署到 Inf2 实例 亚马逊SageMaker.

浮点 2.1 (FP32) 中的稳定扩散 32 模型大小为 5 GB，bfoat2.5 (BF16) 中的稳定扩散 16 模型大小为 2 GB。 单个 inf2.xlarge 实例具有一个 AWS Inferentia32 加速器和 2.1 GB HBM 内存。 Stable Diffusion 2 模型可以适合单个 infXNUMX.xlarge 实例。 稳定扩散是一种文本到图像模型，您只需提供文本提示作为输入，即可使用它来创建不同样式和内容的图像。 要了解有关稳定扩散模型架构的更多信息，请参阅 使用 Stable Diffusion 模型创建高质量图像，并使用 Amazon SageMaker 经济高效地部署它们.

Neuron SDK 如何优化稳定扩散性能

在 AWS Inferentia2.1 实例上部署 Stable Diffusion 2 模型之前，我们需要使用以下命令编译模型组件 神经元SDK。 Neuron SDK 包括深度学习编译器、运行时和工具，可编译并自动优化深度学习模型，以便它们可以在 Inf2 实例上高效运行并充分发挥 AWS Inferentia2 加速器的性能。 我们在以下网站上提供了稳定扩散 2.1 模型的示例 GitHub回购。 本笔记本提供了一个端到端示例，说明如何编译稳定扩散模型、保存编译的神经元模型并将其加载到运行时以进行推理。

我们使用 StableDiffusionPipeline 来自拥抱的脸 diffusers 用于加载和编译模型的库。 然后，我们使用以下命令编译 Neuron 模型的所有组件 torch_neuronx.trace() 并将优化后的模型保存为TorchScript。 编译过程可能会占用大量内存，需要大量 RAM。 为了避免这个问题，在跟踪每个模型之前，我们创建一个 deepcopy 正在跟踪的管道部分。 接下来，我们使用以下命令从内存中删除管道对象 del pipe。 在 RAM 较低的实例上进行编译时，此技术特别有用。

此外，我们还对稳定扩散模型进行优化。 UNet 拥有推理中计算最密集的方面。 UNet 组件对批量大小为 XNUMX 的输入张量进行操作，生成批量大小也为 XNUMX 的相应输出张量，以生成单个图像。 这些批次中的元素完全彼此独立。 我们可以利用这种行为，通过在每个 Neuron 核心上运行一批来获得最佳延迟。 我们编译一批 UNet（通过使用一批输入张量），然后使用 torch_neuronx.DataParallel 用于将此单批模型加载到每个核心上的 API。 该API的输出是一个无缝的两批模块：我们可以将两批的输入传递给UNet，并返回两批输出，但在内部，两个单批模型在两个Neuron核心上运行。 该策略优化了资源利用率并减少了延迟。

在 Inf2 EC2 实例上编译并部署稳定扩散模型

要在 Inf2 EC2 实例上编译和部署稳定扩散模型，请登录 AWS管理控制台 并创建一个 inf2.8xlarge 实例。 请注意，仅编译模型时需要 inf2.8xlarge 实例，因为编译需要更高的主机内存。 稳定扩散模型可以托管在 inf2.xlarge 实例上。 您可以使用以下命令找到带有 Neuron 库的最新 AMI AWS命令行界面 （AWS CLI）命令：

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

在此示例中，我们使用深度学习 AMI Neuron PyTorch 2 (Ubuntu 1.13) 创建了一个 EC20.04 实例。 然后，您可以通过连接到实例并运行以下步骤来创建 JupyterLab 实验室环境：

run source /opt/aws_neuron_venv_pytorch/bin/activate
pip install jupyterlab
jupyter-lab

包含编译和托管模型的所有步骤的笔记本位于 GitHub上.

让我们看一下其中一个文本编码器块的编译步骤。 作为稳定扩散管道一部分的其他块可以类似地编译。

第一步是从 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。 它存储在文件名 model.pt 下 text_encoder 我们的编译器工作区目录：

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。 编译完所有模型后，您可以按照以下步骤加载并运行模型：

1. 定义编译模型的路径。
2. 加载预训练的 StableDiffusionPipeline 模型，其配置指定为使用 bfloat16 数据类型。
3. 使用以下命令将 UNet 模型加载到两个神经元核心上 torch_neuronx.DataParallel API。 这允许执行数据并行推理，从而可以显着加快模型性能。
4. 加载模型的其余部分（text_encoder, decoder及 post_quant_conv）到单个神经元核心上。

然后，您可以通过提供输入文本作为提示来运行管道。 以下是模型生成的一些提示图片：

• 雷诺·塞尚肖像，钢笔和墨水，复杂的线条图，作者：Craig Mullins、阮嘉、kentaro miura、greg rutkowski、loundraw

• 19 世纪老煤矿工人肖像，美丽的画作，由格雷格·鲁特科斯基 (greg rutkowski) 绘制的极其细致的面部彩绘

• 森林中央的一座城堡

在 AWS Inferentia2.1 和 SageMaker 上托管 Stable Diffusion 2

使用 SageMaker 托管稳定扩散模型还需要使用 Neuron SDK 进行编译。 您可以使用大型模型推理 (LMI) 容器提前或在运行时完成编译。 提前编译可以加快模型加载时间，是首选选项。

SageMaker LMI 容器提供两种部署模型的方法：

• 无代码选项，我们只需提供 serving.properties 具有所需配置的文件
• 自带推理脚本

我们查看这两种解决方案并检查配置和推理脚本（model.py）。 在这篇文章中，我们使用存储在 亚马逊简单存储服务 （亚马逊 S3）存储桶。 您可以使用此预编译模型进行部署。

使用提供的脚本配置模型

在本节中，我们将展示如何配置 LMI 容器来托管稳定扩散模型。 SD2.1 笔记本电脑可在 GitHub上。 第一步是根据以下目录结构创建模型配置包。 我们的目标是使用托管模型所需的最小模型配置。 需要的目录结构如下：

<config-root-directory> / 
    ├── serving.properties
    │   
    └── model.py [OPTIONAL]

接下来，我们创建 服务.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

参数指定以下内容：

• 选项.model_id – LMI 容器使用 s5cmd 从 S3 位置加载模型，因此我们需要指定编译权重的位置。
• 选项.入口点 – 要使用内置处理程序，我们指定 Transformers-neuronx 类。 如果您有自定义推理脚本，则需要提供该脚本。
• 选项.dtype – 这指定加载特定大小的权重。 在这篇文章中，我们使用 BF16，与 FP32 相比，它进一步降低了我们的内存需求，并因此降低了我们的延迟。
• 选项.tensor_parallel_ Degree – 此参数指定我们用于此模型的加速器数量。 AWS Inferentia2 芯片加速器有两个 Neuron 核心，因此指定值 2 意味着我们使用一个加速器（两个核心）。 这意味着我们现在可以创建多个工作线程来增加端点的吞吐量。
• 选项.引擎 – 设置为 Python 表示我们不会为此托管使用 DeepSpeed 或 Faster Transformer 等其他编译器。

带上你自己的剧本

如果您想自带自定义推理脚本，则需要删除 option.entryPoint 止 serving.properties。 在这种情况下，LMI 容器将寻找 model.py 文件位于与 serving.properties 并用它来运行推理。

创建您自己的推理脚本 (model.py)

使用 LMI 容器创建您自己的推理脚本相对简单。 该容器需要您的 model.py 文件具有以下方法的实现：

def handle(inputs: Input) which returns an object of type Outputs

让我们来看看该计划的一些关键领域 附笔记本，演示了自带脚本功能。

更换 cross_attention 具有优化版本的模块：

# Replace original cross-attention module with custom cross-attention module for better performance
    CrossAttention.get_attention_scores = get_attention_scores
Load the compiled weights for the following
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 优化权重时，请注意，我们还指定了需要加载这些权重的神经元核心的数量。 在这里，我们加载到具有两个核心的单个加速器：

# Load the compiled UNet onto two neuron cores.
    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)
   
 
    # Load other compiled models onto a single neuron core.
 
    # - load encoders
    pipe.text_encoder = NeuronTextEncoder(pipe.text_encoder)
    clip_compiled = torch.jit.load(text_encoder_filename)
    pipe.text_encoder.neuron_text_encoder = clip_compiled
    #- load decoders
    pipe.vae.decoder = torch.jit.load(decoder_filename)
    pipe.vae.post_quant_conv = torch.jit.load(post_quant_conv_filename)

通过提示运行推理会调用管道对象来生成图像。

创建 SageMaker 端点

我们使用 Boto3 API 创建 SageMaker 端点。 完成以下步骤：

1. 创建仅包含服务和可选选项的 tarball model.py 文件并将其上传到 Amazon S3。
2. 使用图像容器和之前上传的模型 tarball 创建模型。
3. 使用以下关键参数创建端点配置：
   1. 使用一个 ml.inf2.xlarge 实例。
   2. 在 ContainerStartupHealthCheckTimeoutInSeconds 更改为 240 以确保模型部署后启动健康检查。
   3. 在 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商店 (Amazon EBS) 卷附件。 我们可以将模型下载超时和容器启动运行状况检查设置为更高的值，这将为容器提供足够的时间从 Amazon S3 中提取权重并将其加载到 AWS Inferentia2 加速器中。 欲了解更多详情，请参阅 创建端点配置.

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 上对稳定扩散模型进行基准测试

我们在 Inf16 上运行了一些测试，对使用 BF 2 数据类型的稳定扩散模型进行基准测试，并且我们能够得出可与稳定扩散的其他一些加速器相媲美或超过的延迟数。 再加上 AWS Inferentia2 芯片的较低成本，使其成为一个极其有价值的提议。

以下数字来自部署在 inf2.xl 实例上的稳定扩散模型。 有关成本的更多信息，请参阅 Amazon EC2 Inf2实例.

结论

在这篇文章中，我们深入研究了使用 Inf2.1 实例编译、优化和部署 Stable Diffusion 2 模型。 我们还演示了使用 SageMaker 部署稳定扩散模型。 Inf2 实例还为 Stable Diffusion 1.5 提供了极高的性价比。 要了解有关 Inf2 实例为何非常适合生成式 AI 和大型语言模型的更多信息，请参阅 用于低成本、高性能生成式 AI 推理的 Amazon EC2 Inf2 实例现已全面上市。 有关性能详细信息，请参阅 Inf2性能。 查看更多示例 GitHub回购.

特别感谢 Matthew Mcclain、Beni Hegedus、Kamran Khan、Shruti Koparkar 和 Qing Lan 的审核和提供的宝贵意见。

作者简介

维韦克甘加萨尼 是 Amazon Web Services 的高级机器学习解决方案架构师。 他与机器学习初创公司合作，在 AWS 上构建和部署 AI/ML 应用程序。 他目前专注于为 MLOps、ML 推理和低代码 ML 提供解决方案。 他从事过不同领域的项目，包括自然语言处理和计算机视觉。

董家强 是 AWS Annapurna Labs 的高级解决方案架构师。 他专门从事大型深度学习模型在云中的大规模训练和部署。 他拥有博士学位。 在达拉斯德克萨斯大学西南医学中心获得分子生物物理学博士学位。 他曾在 AWS 峰会和 AWS Reinvent 上发表过演讲。 现在，他帮助客户在 AWS 云中训练和部署大型 PyTorch 和 TensorFlow 模型。 他是两本书的作者： 学习 TensorFlow 企业版 和 TensorFlow 2 袖珍参考.

鲁宾德·格鲁瓦尔 是 AWS 的高级人工智能/机器学习专家解决方案架构师。 他目前专注于在 SageMaker 上提供模型和 MLOps。 在担任此职务之前，他曾担任机器学习工程师构建和托管模型。 工作之余，他喜欢打网球和在山路上骑自行车。

• SEO 支持的内容和 PR 分发。 今天得到放大。
• PlatoData.Network 垂直生成人工智能。 赋予自己力量。 访问这里。
• 柏拉图爱流。 Web3 智能。 知识放大。 访问这里。
• 柏拉图ESG。 汽车/电动汽车， 碳， 清洁科技, 能源， 环境， 太阳能， 废物管理。 访问这里。
• 块偏移量。 现代化环境抵消所有权。 访问这里。
• Sumber: https://aws.amazon.com/blogs/machine-learning/maximize-stable-diffusion-performance-and-lower-inference-costs-with-aws-inferentia2/