使用大型模型推理深度学习容器和 DeepSpeed 在 Amazon SageMaker 上部署 BLOOM-176B 和 OPT-30B

过去几年，深度学习领域发展迅速。 尽管硬件有所改进，例如 NVIDIA 和亚马逊的最新一代加速器，但高级机器学习 (ML) 从业者在为自然语言处理 (NLP) 等应用程序部署大型深度学习模型时仍然经常遇到问题。

在之前的一篇文章中，我们讨论了 功能和可配置设置 in Amazon SageMaker 模型部署 这可以使这些大型模型的推理变得更容易。 今天，我们宣布一个新的 亚马逊SageMaker 深度学习容器 (DLC)，您可以使用它在几分钟内开始进行大型模型推理。 这个 DLC 打包了一些最流行的模型并行推理开源库，例如 DeepSpeed 和 Hugging Face Accelerate。

在这篇文章中，我们使用新的 SageMaker 大型模型推理 DLC 来部署两个最流行的大型 NLP 模型：BigScience 的 布卢姆-176B 和 Meta 的 OPT-30B 来自拥抱脸存储库。 特别是，我们使用 DeepSpeed 的深度 Java 库 (DJL) 服务和张量并行技术，在文本生成用例中实现每个令牌 0.1 秒的延迟。

您可以在我们的 GitHub存储库.

大型模型推理技术

语言模型最近在规模和流行度上都呈爆炸式增长。 随着 Hugging Face 等模型库的轻松访问以及分类和文本生成等 NLP 任务的准确性和性能的提高，从业者越来越多地接触这些大型模型。 但是，大型模型通常太大而无法容纳在单个加速器的内存中。 例如，BLOOM-176B 型号可能需要超过 350 GB 的加速器内存，这远远超过了当今可用的硬件加速器的容量。 这需要使用 DeepSpeed 和 Hugging Face Accelerate 等库中的模型并行技术，将模型分布在多个加速器上进行推理。 在这篇文章中，我们使用 SageMaker 大型模型推理容器 使用这两个开源库生成和比较延迟和吞吐量性能。

DeepSpeed 和 Accelerate 使用不同的技术来优化大型语言模型以进行推理。 关键区别在于 DeepSpeed 的 使用优化内核. 这些内核可以通过减少模型计算图中的瓶颈来显着改善推理延迟。 优化的内核可能难以开发，并且通常特定于特定的模型架构； DeepSpeed 使用这些优化的内核支持流行的大型模型，例如 OPT 和 BLOOM。 相比之下，Hugging Face 的 Accelerate 库在撰写本文时不包括优化内核。 正如我们在结果部分中讨论的那样，这种差异是 DeepSpeed 超过 Accelerate 的大部分性能优势的原因。

DeepSpeed 和 Accelerate 之间的第二个区别是模型并行的类型。 Accelerate 使用管道并行性在模型的隐藏层之间划分模型，而 DeepSpeed 使用张量并行性来划分层本身。 流水线并行是一种灵活的方法，它支持更多的模型类型，并且可以在使用更大的批量大小时提高吞吐量。 张量并行性需要 GPU 之间的更多通信，因为模型层可以分布在多个设备上，但可以通过同时使用多个 GPU 来改善推理延迟。 您可以在以下位置了解有关并行技术的更多信息 模型并行性介绍 和 模型并行.

解决方案概述

为了有效地托管大型语言模型，我们需要以下关键领域的功能和支持：

• 构建和测试解决方案 – 鉴于 ML 开发的迭代性质，我们需要能够构建、快速迭代和测试托管这些模型时推理端点的行为，包括快速失败的能力。 这些模型通常只能托管在 p4dn 或 g5 等较大的实例上，并且考虑到模型的大小，启动推理实例并运行任何测试迭代可能需要一段时间。 本地测试通常有限制，因为您需要一个大小相似的实例来测试，而这些模型并不容易获得。
• 大规模部署和运行 – 模型文件需要加载到推理实例上，考虑到大小，这本身就是一个挑战。 以 Bloom-176B 的 Tar / Un-Tar 为例，创建大约需要 1 小时，加载需要另外一个小时。 我们需要一种替代机制来允许轻松访问模型文件。
• 将模型加载为单例 – 对于多工作进程，我们需要确保模型只加载一次，这样我们就不会遇到竞争条件并进一步消耗不必要的资源。 在这篇文章中，我们展示了一种直接从 亚马逊简单存储服务 （亚马逊 S3）。 但是，这仅在我们使用 DJL 的默认设置时才有效。 此外，端点的任何扩展都需要能够在几分钟内启动，这需要重新考虑如何加载和分发模型。
• 分片框架 – 这些模型通常需要，通常通过张量并行机制或通过管道分片作为典型的分片技术，并且我们有先进的概念，例如建立在张量分片之上的零分片。 有关分片技术的更多信息，请参阅 模型并行. 为了实现这一点，我们可以有各种组合并使用来自 NIVIDIA、DeepSpeed 等的框架。 这需要能够测试 BYOC 或使用 1P 容器并迭代解决方案并运行基准测试。 您可能还想测试各种托管选项，例如异步、无服务器等。
• 硬件选择 – 您对硬件的选择取决于上述所有要点以及进一步的流量模式、用例需求和模型大小。

在这篇文章中，我们使用 DeepSpeed 的优化内核和张量并行技术在 SageMaker 上托管 BLOOM-176B 和 OPT-30B。 我们还比较了 Accelerate 的结果，以展示优化内核和张量并行的性能优势。 有关 DeepSpeed 和 Accelerate 的更多信息，请参阅 DeepSpeed Inference：以前所未有的规模实现 Transformer 模型的高效推理 和 使用 DeepSpeed 和 Accelerate 实现令人难以置信的快速 BLOOM 推理.

在本例中，我们使用 DJLServing 作为模型服务解决方案。 DJLServing 是一种高性能通用模型服务解决方案，由与编程语言无关的深度 Java 库 (DJL) 提供支持。 要了解有关 DJL 和 DJLServing 的更多信息，请参阅 使用 DJLServing 和 DeepSpeed 模型并行推理在 Amazon SageMaker 上部署大型模型.

值得注意的是，优化的内核可能会导致精度变化和修改的计算图，理论上这可能会导致模型行为发生变化。 尽管这偶尔会改变推理结果，但我们预计这些差异不会对模型的基本评估指标产生重大影响。 尽管如此，建议从业者在使用这些内核时确认模型输出是否符合预期。

以下步骤演示了如何使用 DJLServing 和 SageMaker 大型模型推理容器在 SageMaker 中部署 BLOOM-176B 模型。 完整的例子也可以在我们的 GitHub存储库.

使用 DJLServing SageMaker DLC 映像

将区域替换为您正在运行笔记本的特定区域后，使用以下代码使用 DJLServing SageMaker DLC 图像：

763104351884.dkr.ecr..amazonaws.com/djl-inference:0.19.0-deepspeed0.7.3-cu113
# example uri might be like 763104351884.dkr.ecr.us-east-1.amazonaws.com/djl-inference:0.19.0-deepspeed0.7.3-cu113

创建我们的模型文件

首先，我们创建一个名为 serving.properties 只包含一行代码。 这告诉 DJL 模型服务器使用 DeepSpeed 引擎。 该文件包含以下代码：

engine=DeepSpeed

serving.properties 是由 DJLServing 定义的文件，用于配置每个模型的配置。

接下来，我们创建我们的 model.py 文件，它定义了加载和服务模型所需的代码。 在我们的代码中，我们读到 TENSOR_PARALLEL_DEGREE 环境变量（默认值为 1）。 这设置了张量并行模块分布的设备数量。 请注意，DeepSpeed 提供了一些内置分区定义，包括用于 BLOOM 模型的分区定义。 我们通过指定使用它 replace_method 和 relpace_with_kernel_inject. 如果您有自定义模型并且需要 DeepSpeed 进行有效分区，则需要更改 relpace_with_kernel_inject 至 false 并添加 injection_policy 使运行时分区工作。 有关详细信息，请参阅 初始化推理. 对于我们的示例，我们在 DeepSpeed 上使用了预分区的 BLOOM 模型。

其次，在 model.py 文件，我们还在端点启动后从 Amazon S3 加载模型。 模型被加载到 /tmp 容器上的空间，因为 SageMaker 映射 /tmp 到 Amazon Elastic Block商店 当我们指定终端节点创建参数时挂载的 (Amazon EBS) 卷 VolumeSizeInGB. 对于像 p4dn 这样预构建卷实例的实例，我们可以继续利用 /tmp 在容器上。 请参阅以下代码：

from djl_python import Input, Output
import os
import deepspeed
import torch
import torch.distributed as dist
import sys
import subprocess
import time
from glob import glob
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from transformers.models.opt.modeling_opt import OPTDecoderLayer

predictor = None

def check_config():
    local_rank = os.getenv('LOCAL_RANK')
    
    if not local_rank:
        return False
    return True
    
def get_model():

    if not check_config():
        raise Exception("DJL:DeepSpeed configurations are not default. This code does not support non default configurations") 
    
    tensor_parallel = int(os.getenv('TENSOR_PARALLEL_DEGREE', '1'))
    local_rank = int(os.getenv('LOCAL_RANK', '0'))
    model_dir = "/tmp/model"
    bucket = os.environ.get("MODEL_S3_BUCKET")
    key_prefix = os.environ.get("MODEL_S3_PREFIX")
    print(f"rank: {local_rank}")
    if local_rank == 0:
        if f"{model_dir}/DONE" not in glob(f"{model_dir}/*"):
            print("Starting Model downloading files")
            try:
                proc_run = subprocess.run(
                    ["aws", "s3", "cp", "--recursive", f"s3://{bucket}/{key_prefix}", model_dir]
                )
                print("Model downloading finished")
                # write file when download complete. Could use dist.barrier() but this makes it easier to check if model is downloaded in case of retry
                with open(f"{model_dir}/DONE", "w") as f:
                    f.write("download_complete")
                    
                proc_run.check_returncode() # to throw the error in case there was one
                
            except subprocess.CalledProcessError as e:
                print ( "Model download failed: Error:nreturn code: ", e.returncode, "nOutput: ", e.stderr )
                raise # FAIL FAST  
                               
    dist.barrier()
                
    
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    
    # has to be FP16 as Int8 model loading not yet supported
    with deepspeed.OnDevice(dtype=torch.float16, device="meta"):
        model = AutoModelForCausalLM.from_config(
            AutoConfig.from_pretrained(model_dir), torch_dtype=torch.bfloat16
        )
    model = model.eval()
    
    model = deepspeed.init_inference(
        model,
        mp_size=tensor_parallel,
        dtype=torch.int8,
        base_dir = model_dir,
        checkpoint=os.path.join(model_dir, "ds_inference_config.json"),
        replace_method='auto',
        replace_with_kernel_inject=True
    )

    model = model.module
    dist.barrier()
    return model, tokenizer

DJLServing 管理在定义的任何 pip 包上的运行时安装 requirement.txt. 该文件将具有：

awscli
boto3

我们创建了一个名为 code 和 model.py, serving.properties及 requirements.txt 文件已在此目录中创建。 要查看文件，您可以从终端运行以下代码：

mkdir -p code
cat code/model.py 
cat code/serving.properties 
cat code/requirements.txt 

下图显示了结构 model.tar.gz.

最后，我们创建模型文件并将其上传到 Amazon S3：

tar cvfz model.tar.gz code
s3_code_artifact = sess.upload_data("model.tar.gz", bucket, s3_code_prefix)

从 Hugging Face 下载并存储模型（可选）

如果您想将模型下载到 Amazon S3 并从那里使用它，我们提供了本节中的步骤。 GitHub 上的 Jupyter 文件中提供了这些步骤。 以下屏幕截图显示了这些步骤的快照。

创建 SageMaker 模型

我们现在创建一个 SageMaker 模型。 我们使用 Amazon Elastic Container注册 (Amazon ECR) 图像和上一步中创建 SageMaker 模型的模型工件。 在模型设置中，我们配置 TENSOR_PARALLEL_DEGREE=8，这意味着模型沿 8 个 GPU 划分。 请参阅以下代码：

PrimaryContainer={
        "Image": inference_image_uri,
        "ModelDataUrl": s3_code_artifact,
        "Environment": {
            "MODEL_S3_BUCKET": bucket,
            "MODEL_S3_PREFIX": s3_model_prefix,
            "TENSOR_PARALLEL_DEGREE": "8",
},

在 Jupyter 文件中运行上述单元后，您会看到类似于以下内容的输出：

{
    "ModelArn": "arn:aws:sagemaker:us-east-1::model/bloom-djl-ds-"
}

创建 SageMaker 端点

您可以使用具有多个 GPU 的任何实例进行测试。 在此演示中，我们使用 p4d.24xlarge 实例。 在下面的代码中，注意我们如何设置 ModelDataDownloadTimeoutInSeconds, ContainerStartupHealthCheckTimeoutInSeconds及 VolumeSizeInGB 参数以适应大模型尺寸。 这 VolumeSizeInGB 参数适用于支持 EBS 卷附件的 GPU 实例。

endpoint_config_response = sm_client.create_endpoint_config(
    EndpointConfigName=endpoint_config_name,
    ProductionVariants=[
        {
            "VariantName": "variant1",
            "ModelName": model_name,
            "InstanceType": "ml.p4d.24xlarge",
            "InitialInstanceCount": 1,
            #"VolumeSizeInGB" : 200,
            "ModelDataDownloadTimeoutInSeconds": 2400,
            "ContainerStartupHealthCheckTimeoutInSeconds": 2400,
        },
    ],
)'

最后，我们创建一个 SageMaker 端点：

create_endpoint_response = sm_client.create_endpoint(
    EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name
)

您会看到它在以下代码中打印出来：

{
    "EndpointArn": "arn:aws:sagemaker:us-east-1::endpoint/bloom-djl-ds-"
}

启动端点可能需要一段时间。 遇到问题可以多试几次 InsufficientInstanceCapacity 错误，或者您可以向 AWS 提出请求以增加您账户中的限制。

性能调优

如果您打算将这篇文章和随附的笔记本用于不同的模型，您可能需要探索 SageMaker、DeepSpeed 和 DJL 提供的一些可调参数。 反复试验这些参数可能会对您托管的大型模型的延迟、吞吐量和成本产生重大影响。 要了解有关调整参数（例如工作人员数量、张量并行度、作业队列大小等）的更多信息，请参阅 DJL 服务配置 和 使用 DJLServing 和 DeepSpeed 模型并行推理在 Amazon SageMaker 上部署大型模型.

成果

在这篇文章中，我们使用 DeepSpeed 在 SageMaker ML 实例上托管 BLOOM-176B 和 OPT-30B。 下表总结了我们的性能结果，包括与 Hugging Face 的 Accelerate 的比较。 延迟反映了产生 256 个令牌的字符串四次所需的毫秒数（batch_size=4) 从模型中。 吞吐量反映了每个测试每秒生成的令牌数。 对于 Hugging Face Accelerate，我们使用库的默认加载和 GPU 内存映射。 对于 DeepSpeed，我们使用了其更快的检查点加载机制。

从延迟的角度来看，BLOOM-3.4B 的 DeepSpeed 比 Accelerate 快 176 倍，OPT-5.6B 的 30 倍。 DeepSpeed 的优化内核是造成这种延迟差异的主要原因。 鉴于这些结果，如果您选择的模型受支持，我们建议使用 DeepSpeed 而不是 Accelerate。

还值得注意的是，使用 DeepSpeed 的模型加载时间要短得多，如果您预计需要快速扩展端点数量，它是一个更好的选择。 如果您拥有 DeepSpeed 不支持的模型或模型精度，Accelerate 更灵活的管道并行技术可能是更好的选择。

这些结果还证明了不同模型大小的延迟和吞吐量的差异。 在我们的测试中，OPT-30B 在比 BLOOM-2.4B 便宜三倍以上的实例类型上每单位时间生成的令牌数量是 BLOOM-176B 的 30 倍。 在每单位吞吐量的价格基础上，g5.24xl 实例上的 OPT-8.9B 比 p176d.4xl 实例上的 BLOOM-24B 好 XNUMX 倍。 如果您有严格的延迟、吞吐量或成本限制，请考虑使用仍能满足功能要求的最小模型。

清理

作为最佳实践的一部分，始终建议删除空闲实例。 下面的代码向您展示了如何删除实例。

# - Delete the end point
sm_client.delete_endpoint(EndpointName=endpoint_name)

# - In case the end point failed we still want to delete the model
sm_client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)
sm_client.delete_model(ModelName=model_name)

可选择从 S3 中删除模型检查点

!aws s3 rm --recursive s3:///{s3_model_prefix}

结论

在这篇文章中，我们演示了如何使用 SageMaker 大型模型推理容器来托管两种大型语言模型 BLOOM-176B 和 OPT-30B。 我们在单个 SageMaker ML 实例上使用了 DeepSpeed 的模型并行技术和多个 GPU。

有关 Amazon SageMaker 及其大型模型推理功能的更多详细信息，请参阅 Amazon SageMaker 现在支持通过可配置的卷大小和超时配额部署大型模型 和 实时推断.

关于作者

西蒙·扎马林 是一位AI / ML解决方案架构师，其主要重点是帮助客户从其数据资产中提取价值。 在业余时间，Simon喜欢与家人共度时光，阅读科幻小说，并从事各种DIY房屋项目。

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

刘 是AWS深度学习的软件工程师。 他专注于为软件工程师和科学家构建创新的深度学习工具。 在业余时间，他喜欢与朋友和家人一起远足。

艾伦谭 是 SageMaker 的高级产品经理，负责大型模型推理方面的工作。 他热衷于将机器学习应用于分析领域。 工作之余，他喜欢户外活动。

达瓦尔·帕特尔 是 AWS 的首席机器学习架构师。 他曾与从大型企业到中型初创公司的组织合作，解决与分布式计算和人工智能相关的问题。 他专注于深度学习，包括 NLP 和计算机视觉领域。 他帮助客户在 SageMaker 上实现高性能模型推理。

青岚 是 AWS 的一名软件开发工程师。 他一直在亚马逊开发几个具有挑战性的产品，包括高性能 ML 推理解决方案和高性能日志记录系统。 清的团队以极低的延迟成功推出了亚马逊广告中的第一个十亿参数模型。 青对基础设施优化和深度学习加速有深入的了解。

李庆伟 是Amazon Web Services的机器学习专家。 他获得了博士学位。 在他打破了顾问的研究补助金帐户并未能兑现他所承诺的诺贝尔奖之后，他便在运筹学获得博士学位。 目前，他帮助金融服务和保险行业的客户在AWS上构建机器学习解决方案。 在业余时间，他喜欢阅读和教学。

罗伯特·范杜森 是 Amazon SageMaker 的高级产品经理。 他领导大型模型推理等应用的深度学习模型优化。

悉达多文卡特桑 是 AWS 深度学习的软件工程师。 他目前专注于构建大型模型推理的解决方案。 在加入 AWS 之前，他曾在 Amazon Grocery 组织工作，为全球客户构建新的支付功能。 工作之余，他喜欢滑雪、户外活动和观看体育赛事。