tl;dr：使用大型模型推理(LMI) 容器和带有 vLLM 滚动批处理的深度 Java 库(DJL) 库在 SageMaker 上部署Phi-3-mini-4k-instruct模型的代码。

import io
import sagemaker
import boto3
import json
# Change this to your role
iam_role = "arn:aws:iam::1111111111:role/service-role/AmazonSageMaker-ExecutionRole-00000000T000000"
sagemaker_session = sagemaker.session.Session()
region = sess._region_name
smr_client = boto3.client("sagemaker-runtime")
container_uri = sagemaker.image_uris.retrieve(framework="djl-lmi", version="0.28.0", region=region)
instance_type = "ml.g5.4xlarge"
endpoint_name = sagemaker.utils.name_from_base("phi3-4k-lmi-endpoint")
model = sagemaker.Model(
    image_uri=container_uri, 
    role=iam_role,
    env={
        "HF_MODEL_ID": "microsoft/Phi-3-mini-4k-instruct",
        "OPTION_ROLLING_BATCH": "vllm",
        "TENSOR_PARALLEL_DEGREE": "max",
        "OPTION_MAX_ROLLING_BATCH_SIZE": "2",
        "OPTION_DTYPE":"fp16",
    }
)
model.deploy(
    instance_type=instance_type,
    initial_instance_count=1,
    endpoint_name=endpoint_name,
)

已部署终端的流式输出演示：

在理解代码之前，我们首先要了解一些术语，如 vLLM 和大型模型推理容器。如果你还不了解这些框架，这将帮助你更好地理解代码。

vLLM

vLLM 在 Kwon 等人的论文 "Efficient Memory Management for Large Language Model Serving with PagedAttention"（ArXiv，9 月 2023 日）中提出，是虚拟大型语言模型的缩写。vLLM 解决了 GPU 内存分配的难题，尤其是当前 LLM 服务系统中管理键值（KV）缓存内存的低效率问题。这一限制导致 GPU 利用率不足、推理速度较慢以及内存使用率较高。

本文作者受到操作系统中使用的内存和分页技术的启发，引入了一种名为 PagedAttention 的注意力算法来应对这些挑战。PagedAttention 使用分页技术--一种将硬件地址映射到虚拟地址的方法--通过允许非连续存储注意力键和值 (KV)，实现了高效的内存管理。

推理请求的批处理主要有两种类型：

• 客户端（静态）--通常情况下，当客户端向服务器发送请求时，服务器会默认按顺序处理每个请求，这对吞吐量来说并不是最佳选择。为了优化吞吐量，客户端会在单个有效载荷中批量处理推理请求，服务器则会实施预处理逻辑，将批量请求分解为多个请求，并分别运行每个请求的推理。在这一方案中，客户端需要更改批处理代码，而且解决方案与批处理规模紧密相关。
• 服务器端（动态）--批处理的另一种技术是使用推理来帮助在服务器端实现批处理。当独立的推理请求到达服务器时，推理服务器可在服务器端动态地将它们分组为更大的批次。推理服务器可以管理批处理，以满足指定的延迟目标，最大限度地提高吞吐量，同时保持在所需的延迟范围内。推理服务器会自动处理，因此无需更改客户端代码。服务器端批处理包括不同的技术，以进一步优化基于自动回归解码的生成语言模型的吞吐量。这些批处理技术包括动态批处理、连续批处理和 PagedAttention（vLLM）批处理。

vLLM 也使用连续批处理，它能在模型生成输出标记时动态调整批处理大小。

连续批处理

连续批处理是一种专门针对文本生成的优化方法。它提高了吞吐量，而且不会牺牲到第一个字节的延迟时间。连续批处理（也称为迭代或滚动批处理）解决了 GPU 闲置时间的难题，并在动态批处理方法的基础上，进一步通过在批处理中不断推送新请求来实现。下图显示了请求的连续批处理。当请求 2 和请求 3 处理完毕后，会调度另一组请求。

下面的交互式图表深入介绍了连续批处理的工作原理。

vLLM 与 TensorRT-LLM、LMI-Dist、Transformers 和 NeuronX 等其他框架的不同之处在于，它专注于高效的内存管理和可扩展的服务。其他框架优先考虑模型优化和加速，而 vLLM 的 PagedAttention 算法和连续批处理使其非常适合生产环境。

尽管许多其他库（如 HuggingFace 文本生成库）也在效仿 vLLM，但它们也进行了一些内部优化，以提高性能。

在生产环境中使用 vLLM 有几个优势。vLLM 易于使用，再加上其强大的功能，使它成为希望在其应用程序中利用 LLM 的开发人员的一个极具吸引力的选择。此外，vLLM 对分布式推理和实时处理的支持可实现可扩展的高效模型服务。

SageMaker LMI 容器和 Deep Java 库

SageMaker LMI 容器是来自 Deep Java Library (DJL) 的一组预构建容器，可在 Amazon SageMaker 上高效推断大型语言模型 (LLM)。这些容器提供了一种简单、可扩展的方式来部署 LLM，使开发人员能够专注于构建 AI 应用程序，而无需担心底层基础设施。

Deep Java Library（DJL）是一个用于深度学习的开源、高级、引擎无关的 Java 框架。它提供原生的 Java 开发体验，让 Java 开发人员能够轻松上手机器学习和深度学习，而无需掌握该领域的丰富专业知识。DJL 的设计易于使用，能与 Java 应用程序很好地集成，并支持多种深度学习引擎，如 TensorFlow、PyTorch 和 MXNet。

DJL 的 SageMaker LMI 容器具有多项关键功能，包括支持 vLLM、TensorRT-LLM、LMI-Dist、NauronX 和 Hugging Face Transformers 等流行的 LLM 框架，优化大型模型的性能，以及与 SageMaker 的托管基础架构无缝集成。它们还提供模型服务、批处理和缓存等功能，使在生产中部署和管理 LLM 变得更加容易。

如果你以前使用自带容器（BYOC）方法在SageMaker上部署过自定义模型，那么你可能已经熟悉了sagemaker-inference-toolkit。

SageMaker 推论工具包

SageMaker Inference Toolkit 是 AWS SageMaker 提供的一个开源库，使开发人员能够在各种硬件平台（包括 CPU、GPU 以及 TPU 和 FPGA 等专用 AI 加速器）上优化、编译和运行机器学习模型。它提供了一套工具和 API，用于优化模型推理性能、减少延迟和提高吞吐量。

该工具包支持 TensorFlow、PyTorch 和 MXNet 等流行的深度学习框架，允许开发人员在云、边缘或内部部署环境中部署模型。它还提供了模型优化、量化和内核优化等功能，以提高推理性能。

SageMaker Inference Toolkit是一个更广泛的库，侧重于在各种硬件平台上优化和部署机器学习模型，而SageMaker LMI深度学习容器则是专门为在SageMaker和其他云平台上使用DJL库部署LLM等大型深度学习模型而设计的。

我提到推理工具包是为了与大型模型推理容器形成对比。如果你不想使用 DJL，可以使用该工具包从头开始创建自己的容器。使用 DJL 部署的模型也可以通过 serving.properties 配置进行自定义，而不是本文提到的环境变量配置。

代码演练

import io
import sagemaker
import boto3
import json

# Your IAM role that provides access to SageMaker and S3.
# See https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-ex-role.html
# if running on a SageMaker notebook or directly use
# sagemaker.get_execution_role() if running on SageMaker studio
iam_role = "arn:aws:iam::1111111111:role/service-role/AmazonSageMaker-ExecutionRole-00000000T000000"
# manages interactions with the sagemaker apis
sagemaker_session = sagemaker.session.Session()
region = sagemaker_session._region_name
# boto3 Sagemaker runtime client to invoke the endpoint
# with streaming response
smr_client = boto3.client("sagemaker-runtime")
# get the lmi image uri
# available frameworks: "djl-lmi" (for vllm, lmi-dist), "djl-tensorrtllm" (for tensorrt-llm),
# "djl-neuronx" (for transformers neuronx)
container_uri = sagemaker.image_uris.retrieve(
    framework="djl-lmi", version="0.28.0", region=region
)
# instance type you will deploy your model to
# Go for bigger instance if your model is bigger
# than 7B parameters
instance_type = "ml.g5.4xlarge"
# create a unique endpoint name
endpoint_name = sagemaker.utils.name_from_base("phi3-4k-lmi-endpoint")
# create your SageMaker Model
# phi-3-mini-4k model fits well on our instance's GPU
# as it only has 3.8B parameters
model = sagemaker.Model(
    image_uri=container_uri,
    role=iam_role,
    # specify all environment variable configs in this map
    env={
        "HF_MODEL_ID": "microsoft/Phi-3-mini-4k-instruct",
        "OPTION_ROLLING_BATCH": "vllm",
        "TENSOR_PARALLEL_DEGREE": "max",
        "OPTION_MAX_ROLLING_BATCH_SIZE": "2",
        "OPTION_DTYPE": "fp16",
        # Streaming will work without this variable
        # "OPTION_ENABLE_STREAMING":"true"
    },
)
# deploy your model
model.deploy(
    instance_type=instance_type,
    initial_instance_count=1,
    endpoint_name=endpoint_name,
)

1. IAM 角色

从AWS账户的IAM服务中复制SageMaker IAM角色的ARN，如果在SageMaker上运行笔记本，则使用sagemaker.get_execution_role() 。

这样，SageMaker 就可以在角色允许的策略范围内配置资源并与其他 AWS 服务交互。

2. SageMaker会话、区域和客户端

我们创建一个SageMaker会话实例来获取默认区域。该区域取自 $HOME/.aws/config 下的 aws cli 配置。

创建 Boto3 SageMaker 运行时客户端，以便稍后调用端点。

3. 容器 URI

我们使用 sagemaker.image_uris.retrieve 获取 SageMaker DLC（深度学习容器）的 URI。我们传递框架、区域和版本，因为它们决定了图像的 URI。

4. 实例类型

Amazon EC2 G5 实例是基于 GPU 的高性能实例，适用于图形密集型应用和机器学习推理。

G5 有两种类型的实例：单 GPU 虚拟机和多 GPU 虚拟机。

因为 phi-3-mini 有 38 亿个参数：

模型加载大小：3.8 * 2 (fp16) = 7.6 GB

KV 缓存大小：2（键和值矩阵）* 2（fp16）* 32（层数）* 4096（隐藏大小）= 每个令牌 524288 字节或每个令牌 0.00052 GB

4096 个标记 * 4（批处理大小） * 每个标记 0.00052 GB = 8.52 GB KV 缓存（4 批处理大小

用于推理的模型总大小 = 7.6 + 8.5 = ~16.1 GB

16.1GB 是我们的 phi-3-mini-4k 模型批处理规模为 4 时所需的最小 GPU vRAM。我们选择了拥有 24GB GPU 内存的 g5.4xlarge 实例。这为 vLLM 框架批量预填充和加速文本生成留出了更多空间。

对于更大的模型，这里有 DJL 提供的实例类型选择指南：

要为 SageMaker 端点使用 g5 实例，可能需要申请增加 AWS 账户的配额。

5. 端点名称

使用sagemaker.utils api创建一个唯一的端点名称。使用此 api 是可选的。

6. 创建模型

使用 sagemaker.model api 在账户的 SageMaker 模型注册表中创建一个新模型。你也可以从 SageMaker 用户界面检查该模型。

这里我们传递 container_uri、iam_role 和一些环境变量。这些环境变量决定了：

• HF_MODEL_ID - 下载哪个模型、
• OPTION_ROLLING_BATCH - 使用哪个后端（这里我们通过 vllm 来实现滚动批处理）。LMI 容器中的其他选项包括 lmi-dist（lmi-dist 用于 lmi-dist）、auto（自动）和 disabled（禁用）（hugginface accelerate），其中自动用于文本生成模型，禁用用于非文本生成模型。在 TensorRT-LLM 容器中，trllm 用于 TensorRT-LLM，也是默认选项。在 Transformers NeuronX 容器中，auto（自动）用于 TensorRT-LLM，也是默认选项。
• TENSOR_PARALLEL_DEGREE （张量并行度）- 使用张量并行度将模型分片到多少个 GPU。
• OPTION_MAX_ROLLING_BATCH_SIZE - 批量大小
• OPTION_DTYPE - 浮点数

7. 部署模型

model.deploy 会创建端点配置和端点。

端点重新进入服务状态需要几分钟时间。你可以使用"-"和"！"字符从流输出中查看进度，也可以从 AWS 账户中的 SageMaker UI 查看进度。

流式输出

要从部署的模型中生成输出，你可以使用sagemaker.Predictor api，如下面的示例，或者使用sagemaker-runtime客户端的invoke_endpoint_with_response_stream api。

Predictor api 将一次性返回所有生成的文本。

# Get a predictor for your endpoint
predictor = sagemaker.Predictor(
    endpoint_name=endpoint_name,
    sagemaker_session=sagemaker_session,
    serializer=sagemaker.serializers.JSONSerializer(),
    deserializer=sagemaker.deserializers.JSONDeserializer(),
)
# Make a prediction with your endpoint
outputs = predictor.predict(
    {
        "inputs": "The meaning of life is",
        "parameters": {"do_sample": True, "max_new_tokens": 256},
    }
)
outputs["generated_text"]

示范回答：

‘用最简单的话来说，爱因斯坦提出的相对论由两部分组成： 狭义相对论和广义相对论。狭义相对论认为，物理定律对于所有非加速观察者都是相同的，真空中的光速是恒定的，与光源或观察者的运动无关。这就产生了著名的方程 E=mc²，断言能量（E）和质量（m）是可以互换的。另一方面，广义相对论处理的是万有引力。爱因斯坦没有将其视为一种力，而是提出大质量物体会导致时空扭曲，我们将其感知为引力。这一理论预测了引力波、黑洞等现象，并解释了大质量物体对光线的弯曲。然而’

从端点流式输出

首先，我们将创建一个行迭代器类：

class LineIterator:
    """
    A helper class for parsing the byte stream input.
    The output of the model will be in the following format:
    ```
    b'{"outputs": [" a"]}\n'
    b'{"outputs": [" challenging"]}\n'
    b'{"outputs": [" problem"]}\n'
    ...
    ```
    While usually each PayloadPart event from the event stream will contain a byte array
    with a full json, this is not guaranteed and some of the json objects may be split across
    PayloadPart events. For example:
    ```
    {'PayloadPart': {'Bytes': b'{"outputs": '}}
    {'PayloadPart': {'Bytes': b'[" problem"]}\n'}}
    ```
    This class accounts for this by concatenating bytes written via the 'write' function
    and then exposing a method which will return lines (ending with a '\n' character) within
    the buffer via the 'scan_lines' function. It maintains the position of the last read
    position to ensure that previous bytes are not exposed again.
    """
    def __init__(self, stream):
        self.byte_iterator = iter(stream)
        self.buffer = io.BytesIO()
        self.read_pos = 0
    def __iter__(self):
        return self
    def __next__(self):
        while True:
            self.buffer.seek(self.read_pos)
            line = self.buffer.readline()
            if line and line[-1] == ord("\n"):
                self.read_pos += len(line)
                return line[:-1]
            try:
                chunk = next(self.byte_iterator)
            except StopIteration:
                if self.read_pos < self.buffer.getbuffer().nbytes:
                    continue
                raise
            if "PayloadPart" not in chunk:
                print("Unknown event type:" + chunk)
                continue
            self.buffer.seek(0, io.SEEK_END)
            self.buffer.write(chunk["PayloadPart"]["Bytes"])

创建一个停止标记变量，这样与之前的输出不同，我们会在正确的标记处中断。

stop_token = "\n""\n"

现在，我们将创建一个 body 对象，调用端点并解析流式响应：

# Create body object and pass 'stream' to True
body = {
    "inputs": "The meaning of life",
    "parameters": {
        "max_new_tokens": 400,
        # "return_full_text": False  # This does not work with Phi3
    },
    "stream": True,
}
# Invoke the endpoint
resp = smr_client.invoke_endpoint_with_response_stream(
    EndpointName=endpoint_name, Body=json.dumps(body), ContentType="application/json"
)
# Parse the streaming response
event_stream = resp["Body"]
start_json = b"{"
for line in LineIterator(event_stream):
    if line != b"" and start_json in line:
        data = json.loads(line[line.find(start_json) :].decode("utf-8"))
        if data["token"]["text"] != stop_token:
            print(data["token"]["text"], end="")

jupyter 笔记本中的流式响应演示：