大模型推理有妙招：Amazon SageMaker 让你事半功倍，伸缩自如

背景

在当今快速发展的人工智能领域，大型语言模型（Large Language Models，LLMs）已经成为推动自然语言处理前进的重要力量。然而在我们完成训练后，部署这些大型模型仍然面临着诸多挑战，如计算资源的限制、推理性能的瓶颈以及成本效益的权衡等。本文将以 Llama 3 模型为例，介绍如何在 Amazon SageMaker 便捷地进行模型推理，并根据设定的伸缩策略进行自动扩缩，以应对不同的工作负载需求。

关于 Amazon SageMaker

Amazon SageMaker 是一个全托管机器学习平台，让算法工程师可以快速构建、训练和部署各类模型。

全托管服务：无需管理底层基础设施，开发者可以专注于模型的开发和优化。
集成 Jupyter 环境：方便的模型训练、预处理和实验。
一键部署：通过简单的步骤即可以将训练好的模型部署为可扩展的推理端点。
监控和管理：提供全面的监控和日志功能，便于模型的维护和优化。

目前在 SageMaker 中提供了多种推理方式，以满足不同的推理需求。

以下是主要的几种推理方式：

实时推理（Real-time Inference）：端点能够实时响应请求，适用于需要低延迟和高吞吐量的应用场景。
无服务器推理（Serverless Inference）：根据请求量自动扩展和缩减资源，适用于不规则流量和成本敏感的应用场景。
异步推理（Asynchronous Inference）：允许用户提交推理请求并在稍后时间获取结果，适用于处理时间较长的推理任务。
批量转换（Batch Transform）：对大批量数据进行离线推理，适用于不需要实时响应的批处理任务。

多样化的推理方式提供了灵活的选择，用户可以根据具体的应用需求和资源限制选择最合适的推理方式。

推理框架的选择

LMI 容器是亚马逊云科技专为大语言模型推理而构建的高性能 Docker 容器。与多个推理库提供集成，包括 vLLM、LMI-Dist、TensorRT-LLM、Transformers NeuronX 和 HuggingFace Accelerate。每个推理框架都提供了一组独特的特性和优化。借助 LMI 内置的推理处理程序和统一配置，尝试不同的推理技术栈变得非常简单，只需更改几个配置即可在 Amazon SageMaker 端点上部署大语言模型。在本文中，我们将采用 vLLM 框架来支持 Llama 3 8B 模型的推理。

使用 SageMaker 进行大模型推理

使用步骤

您可以在亚马逊云科技的控制台创建 SageMaker Notebook 实例，选择合适的实例类型和存储空间。

以下内容在该 Notebook 中找到完整代码，本文只做关键代码的说明

构建 SageMaker 对象

%pip install sagemaker --upgrade  --quiet
import boto3
import sagemaker
from sagemaker import Model, image_uris, serializers, deserializers
role = sagemaker.get_execution_role()  # execution role for the endpoint
sess = sagemaker.session.Session()  # sagemaker session for interacting with different AWS APIs
region = sess._region_name  # region name of the current SageMaker Studio environment
account_id = sess.account_id()

notebook 中配置 lmi 相关的参数

%%writefile {local_code_dir}/serving.properties
engine=Python
option.model_id=S3PATH
option.dtype=bf16
option.task=text-generation
option.rolling_batch=vllm
option.tensor_parallel_degree=1
option.device_map=auto
option.gpu_memory_utilization=0.85
option.max_model_len=8192
option.max_tokens=8192
option.output_formatter = json
option.model_loading_timeout = 1200
option.enforce_eager=true

选择推理容器镜像

inference_image_uri = image_uris.retrieve(
        framework="djl-deepspeed",
        region=sess.boto_session.region_name,
        version="0.27.0"
    )

创建模型

s3_code_prefix = "large-model-lmi-dist/code"
bucket = sess.default_bucket()  # bucket to house artifacts
code_artifact = sess.upload_data("mymodel.tar.gz", bucket, s3_code_prefix)
print(f"S3 Code or Model tar ball uploaded to --- &gt; {code_artifact}")
model = Model(image_uri=image_uri, model_data=code_artifact, role=role)

创建 SageMaker 推理 Endpoint

endpoint_config_name = f"{model_name}-config"
endpoint_name = f"{model_name}-endpoint"

endpoint_config_response = sm_client.create_endpoint_config(
    EndpointConfigName=endpoint_config_name,
    ProductionVariants=[
        {
            "VariantName": "variant1",
            "ModelName": model_name,
            "InstanceType": "ml.g5.2xlarge",
            "InitialInstanceCount": 1,
            "ContainerStartupHealthCheckTimeoutInSeconds": 10*60,
        },
    ],
)

create_endpoint_response = sm_client.create_endpoint(
    EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name
)
print(f"Created Endpoint: {create_endpoint_response['EndpointArn']}")

使用 streaming 方式调用 Endpoint

import json
import boto3
input_text = """今天天气，很不错哦"""
smr_client = boto3.client("sagemaker-runtime")
response_model = smr_client.invoke_endpoint_with_response_stream(
    EndpointName=endpoint_name,
    Body=json.dumps(
    {
        "inputs": input_text,
        "parameters": parameters,
        "stream" : True
    }
    ),
    ContentType="application/json",
)

def print_response_stream(response_stream):
    event_stream = response_stream.get('Body')
    for line in LineIterator(event_stream):
        print(line, end='')

print_response_stream(response_model)

SageMaker Endpoint 伸缩

Amazon SageMaker Endpoints 支持为您托管的模型提供自动伸缩功能，根据您的工作负载的变化动态调整预置的实例数量。当工作负载增加时，自动上线更多实例。当工作负载减少时，自动移除不必要的实例，以避免您为未使用的配置实例付费。

SageMaker Endpoints 支持两种自动伸缩策略：目标跟踪策略和步进缩放策略。一般来说，建议使用目标跟踪策略。这种策略根据指标和目标值创建和管理推理集群，根据需要添加和删除实例数量，以保持指标接近或等于指定的目标值。例如，使用预定义的 InvocationsPerInstance 指标和目标值为 70 的扩展策略，可以将 InvocationsPerInstance 保持在 70 左右。

Amazon SageMaker Endpoints 自动缩放功能可以通过三种方式进行配置：SageMaker 控制台、AWS CLI 和 AWS SDK（通过 Application Auto Scaling API）。

下图为通过SageMaker控制台进行目标跟踪策略配置的示例，需要设置自动缩放的最小实例数和最大实例数，以及监控的指标和目标值。

您也可以使用 AWS SDK 编写代码进行配置。以下实现前述相同配置的参考代码。

autoscale = boto3.client("application-autoscaling")

autoscale.register_scalable_target(
    ServiceNamespace="sagemaker",
    ResourceId=f"endpoint/{endpoint_name}/variant/variant1",
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    MinCapacity=1,
    MaxCapacity=5
)

autoscale.put_scaling_policy(
    PolicyName="SageMakerEndpointInvocationScalingPolicy",
    ServiceNamespace="sagemaker",
    ResourceId="endpoint/" + endpoint_name + "/variant/variant1",
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    PolicyType="TargetTrackingScaling",
    TargetTrackingScalingPolicyConfiguration={
        "TargetValue": 70.0,
        "PredefinedMetricSpecification":
        {
            "PredefinedMetricType": "SageMakerVariantInvocationsPerInstance"
        }
    }
)

另外，SageMaker 还支持基于计划的自动缩放，实现在特定时间执行扩缩活动。您可以创建只执行一次的计划操作或按定期计划执行的计划操作。在此不再赘述。

结论

本文以 Llama 3 模型为例，介绍了在 Amazon SageMaker 上部署和推理大模型的完整流程。通过代码示例，我们展示了从配置推理参数到部署 Endpoint 的全过程。此外，还介绍了 SageMaker Endpoint 的自动伸缩能力，可以根据实际情况灵活调整推理资源，达到成本和效益的均衡。借助 SageMaker，可以轻松实现高效、弹性的大模型推理服务。

参考链接

https://thinkwithwp.com/cn/sagemaker/

https://djl.ai/

https://github.com/tsaol/llama-on-aws-sagemaker/

亚马逊AWS官方博客