使用新的 Amazon SageMaker 容器提升 Mixtral 和 Llama 2 模型的推理性能 |亚马逊网络服务

在2024一月， 亚马逊SageMaker 推出新版本 (0.26.0) 大型模型推理 (LMI) 深度学习容器 (DLC)。该版本提供对新模型（包括专家混合）的支持、推理后端的性能和可用性改进，以及用于增强控制和预测可解释性的新一代详细信息（例如生成完成的原因和令牌级别日志概率）。

LMI DLC 提供低代码接口，可简化最先进的推理优化技术和硬件的使用。 LMI 允许您应用张量并行性；最新的高效注意力、批处理、量化和内存管理技术；令牌流；等等，只需提供模型 ID 和可选模型参数即可。借助 SageMaker 上的 LMI DLC，您可以加快实现价值的时间 生成人工智能（AI） 应用程序、卸载与基础设施相关的繁重工作，并针对您选择的硬件优化大型语言模型 (LLM)，以实现一流的性价比。

在这篇文章中，我们将探讨此版本中引入的最新功能，检查性能基准，并提供有关以高性能部署带有 LMI DLC 的新 LLM 的详细指南。

LMI DLC 的新功能

在本节中，我们将讨论 LMI 后端的新功能，并深入研究其他一些特定于后端的功能。 LMI 目前支持以下后端：

• LMI分布式库 – 这是受 OSS 启发，使用 LLM 运行推理的 AWS 框架，以实现结果的最佳延迟和准确性
• LMI 法学硕士 – 这是内存高效的AWS后端实现 法学硕士 推理库
• LMI TensorRT-LLM 工具包 – 这是 AWS 后端实现 NVIDIA TensorRT-法学硕士，它创建 GPU 特定的引擎来优化不同 GPU 上的性能
• LMI DeepSpeed – 这是 AWS 的改编 极速，它增加了真正的连续批处理、SmoothQuant 量化以及在推理过程中动态调整内存的能力
• LMI 神经元X – 您可以使用它来部署 AWS Inferentia2 和 AWS 培训基于实例，具有真正的连续批处理和加速，基于 AWS 神经元开发工具包

下表总结了新添加的功能，包括常见功能和特定于后端的功能。

支持新型号

跨后端支持新的流行模型，例如 Mistral-7B（所有后端）、基于 MoE 的 Mixtral（除 Transformers-NeuronX 之外的所有后端）和 Llama2-70B (Transformers-NeuronX)。

上下文窗口扩展技术

基于旋转位置嵌入 (RoPE) 的上下文缩放现在可在 LMI-Dist、vLLM 和 TensorRT-LLM 后端上使用。 RoPE 缩放可以在推理过程中将模型的序列长度扩展到几乎任何大小，而无需进行微调。

以下是使用 RoPE 时的两个重要注意事项：

• 模型困惑度 – 随着序列长度的增加， 所以可以 模型的 困惑。通过对大于原始训练中使用的输入序列进行最小微调，可以部分抵消这种影响。要深入了解 RoPE 如何影响模型质量，请参阅 延长 RoPE.
• 推理性能 – 较长的序列长度将消耗更高的加速器的高带宽内存 (HBM)。内存使用量的增加可能会对加速器可以处理的并发请求数量产生不利影响。

添加了生成详细信息

您现在可以获得有关生成结果的两个细粒度详细信息：

• 完成原因 – 这给出了生成完成的原因，可以是达到最大生成长度、生成句尾（EOS）标记或生成用户定义的停止标记。它与最后一个流序列块一起返回。
• 日志概率 – 这将返回模型为流序列块中的每个标记分配的对数概率。您可以通过计算序列的联合概率作为模型置信度的粗略估计 log_probs 各个标记的值，这对于对模型输出进行评分和排名非常有用。请注意，未经校准，LLM 代币概率通常会过于自信。

您可以通过在 LMI 的输入负载中添加details=True 来启用生成结果输出，并保持所有其他参数不变：

payload = {“inputs”:“your prompt”,
“parameters”:{max_new_tokens”:256,...,“details”:True}
}

整合配置参数

最后，LMI 配置参数也得到了整合。有关所有常见和特定于后端的部署配置参数的更多信息，请参阅 大型模型推理配置.

LMI-分布式后端

在 AWS re:Invent 2023 上，LMI-Dist 添加了新的、优化的集体操作，以加速 GPU 之间的通信，从而为对于单个 GPU 来说太大的模型带来更低的延迟和更高的吞吐量。这些集合仅适用于 SageMaker、p4d 实例。

之前的迭代仅支持跨所有 8 个 GPU 的分片，而 LMI 0.26.0 以部分全部到全部的模式引入了对 4 张量并行度的支持。这可以与 SageMaker 推理组件，您可以使用它来精细配置应为部署在端点后面的每个模型分配多少个加速器。这些功能共同提供了对底层实例的资源利用率的更好控制，使您能够通过在一个端点后面托管不同的模型来增加模型多租户，或者微调部署的聚合吞吐量以匹配您的模型和流量特征。

下图比较了直接all-to-all和部分all-to-all。

TensorRT-LLM 后端

NVIDIA 的 TensorRT-LLM 是作为之前的 LMI DLC 版本 (0.25.0) 的一部分引入的，可在使用 NVIDIA GPU 时实现最先进的 GPU 性能和优化，例如 SmoothQuant、FP8 以及 LLM 的连续批处理。

TensorRT-LLM 要求在部署之前将模型编译成高效的引擎。 LMI TensorRT-LLM DLC 可以在启动服务器并加载模型进行实时推理之前，自动即时编译支持的模型列表 (JIT)。 DLC 0.26.0 版本增加了 JIT 编译支持的模型列表，引入了 Baichuan、ChatGLM、GPT2、GPT-J、InternLM、Mistral、Mixtral、Qwen、SantaCoder 和 StarCoder 模型。

JIT 编译会增加端点配置和扩展时间的几分钟开销，因此始终建议提前编译模型。有关如何执行此操作的指南以及支持的型号列表，请参阅 TensorRT-LLM 模型提前编译教程。如果您选择的型号尚不支持，请参阅 TensorRT-LLM手动编译模型教程 编译 TensorRT-LLM 支持的任何其他模型。

此外，LMI 现在公开了原生 TensorRT-LLM SmootQuant 量化，以及通过令牌或通道控制 alpha 和缩放因子的参数。相关配置的更多信息可以参考 TensorRT-法学硕士.

vLLM后端

LMI DLC 中包含的 vLLM 更新版本通过 CUDA 图形模式（而非急切模式）实现了高达 50% 的性能提升。 CUDA 图表通过一次性启动多个 GPU 操作而不是单独启动它们来加速 GPU 工作负载，从而减少开销。当使用张量并行性时，这对于小模型特别有效。

性能的提高是以 GPU 内存消耗的增加为代价的。 CUDA 图形模式现在是 vLLM 后端的默认模式，因此如果您受到可用 GPU 内存量的限制，您可以设置 option.enforce_eager=True 强制 PyTorch eager 模式。

变形金刚-NeuronX 后端

更新后的版本为 神经元X LMI NeuronX DLC 中包含的模型现在支持具有分组查询注意机制的模型，例如 Mistral-7B 和 LLama2-70B。分组查询注意力是默认变压器注意力机制的重要优化，其中模型使用比查询头更少的键和值头进行训练。这减少了 GPU 内存上 KV 缓存的大小，从而实现更大的并发性并提高性价比。

下图说明了多头、分组查询和多查询注意方法（资源).

不同的KV缓存分片策略可以适应不同类型的工作负载。有关分片策略的更多信息，请参阅 分组查询注意力 (GQA) 支持。您可以启用您想要的策略（shard-over-heads，例如）使用以下代码：

option.group_query_attention=shard-over-heads

此外，NeuronX DLC 的新实现引入了 TransformerNeuronX 的缓存 API，可以访问 KV 缓存。它允许您在处理批量推理时从新请求中插入和删除 KV 缓存行。在引入此 API 之前，会针对任何新添加的请求重新计算 KV 缓存。与 LMI V7 (0.25.0) 相比，我们将并发请求的延迟提高了 33% 以上，并支持更高的吞吐量。

选择正确的后端

要根据所选模型和任务决定使用哪个后端，请使用以下流程图。有关各个后端用户指南以及支持的型号，请参阅 LMI 后端用户指南.

使用具有附加属性的 LMI DLC 部署 Mixtral

让我们逐步了解如何使用 LMI 8 容器部署 Mixtral-7x0.26.0B 模型并生成其他详细信息，例如 log_prob 和 finish_reason 作为输出的一部分。我们还讨论了如何通过内容生成用例从这些附加属性中受益。

包含详细说明的完整笔记本可在 GitHub回购.

我们首先导入库并配置会话环境：

import boto3
import sagemaker 
import json 
import io 
import numpy as np 
from sagemaker import Model, image_uris, serializers, deserializers 

role = sagemaker.get_execution_role() # execution role for the endpoint 
session = sagemaker.session.Session() # sagemaker session for interacting with different AWS APIs 
region = session._region_name # region name of the current SageMaker Studio environment

您可以使用 SageMaker LMI 容器来托管模型，而无需任何额外的推理代码。您可以通过环境变量或 serving.properties 文件。或者，您可以有一个 model.py 文件进行任何预处理或后处理以及 requirements.txt 需要安装的任何其他软件包的文件。

在这种情况下，我们使用 serving.properties 文件来配置参数并自定义 LMI 容器行为。有关更多详细信息，请参阅 GitHub回购。该存储库解释了您可以设置的各种配置参数的详细信息。我们需要以下关键参数：

• 发动机 – 指定 DJL 使用的运行时引擎。这驱动了模型加速器中的分片和模型加载策略。
• 选项.model_id – 指定 亚马逊简单存储服务 (Amazon S3) 预训练模型的 URI 或托管在模型存储库内的预训练模型的模型 ID 拥抱脸。在本例中，我们提供 Mixtral-8x7B 模型的模型 ID。
• 选项.tensor_parallel_ Degree – 设置 Accelerate 需要对模型进行分区的 GPU 设备的数量。此参数还控制 DJL 服务运行时将启动的每个模型的工作线程数量。我们将此值设置为 max （当前机器上的最大 GPU）。
• 选项.rolling_batch – 实现连续批处理以优化加速器利用率和整体吞吐量。对于 TensorRT-LLM 容器，我们使用 auto.
• option.model_loading_timeout – 设置下载和加载模型以服务推理的超时值。
• 选项.max_rolling_batch – 设置连续批次的最大大小，定义在任何给定时间可以并行处理的序列数。

%%writefile serving.properties 
engine=MPI 
option.model_id=mistralai/Mixtral-8x7B-v0.1 
option.tensor_parallel_degree=max 
option.max_rolling_batch_size=32 
option.rolling_batch=auto 
option.model_loading_timeout = 7200

我们包装了 serving.properties tar.gz 格式的配置文件，以满足 SageMaker 托管要求。我们将 DJL LMI 容器配置为 tensorrtllm 作为后端引擎。此外，我们指定了容器的最新版本 (0.26.0)。

image_uri = image_uris.retrieve(
   framework="djl-tensorrtllm",
   region=sess.boto_session.region_name,
   version="0.26.0"
)

接下来，我们上传本地 tarball（包含 serving.properties 配置文件）到 S3 前缀。我们使用 DJL 容器的图像 URI 以及模型服务工件 tarball 上传到的 Amazon S3 位置来创建 SageMaker 模型对象。

model = Model(image_uri=image_uri, model_data=code_artifact, role=role) 

instance_type = "ml.p4d.24xlarge" 
endpoint_name = sagemaker.utils.name_from_base("mixtral-lmi-model") 

model.deploy(
   initial_instance_count=1,
   instance_type=instance_type,
   endpoint_name=endpoint_name,
   container_startup_health_check_timeout=1800
)

作为 LMI 0.26.0 的一部分，您现在可以使用有关生成输出的两个附加细粒度详细信息：

• 日志概率 – 这是模型为流序列块中的每个标记分配的对数概率。您可以通过将序列的联合概率计算为各个标记的对数概率之和，将它们用作模型置信度的粗略估计，这对于对模型输出进行评分和排名非常有用。请注意，未经校准，LLM 代币概率通常会过于自信。
• 完成原因 – 这是生成完成的原因，可以是达到最大生成长度、生成EOS代币或生成用户定义的停止代币。这是与最后一个流序列块一起返回的。

您可以通过传递来启用这些 "details"=True 作为模型输入的一部分。

让我们看看如何生成这些详细信息。我们使用内容生成示例来了解它们的应用。

我们定义一个 LineIterator 辅助类，它具有从响应流中延迟获取字节、缓冲它们并将缓冲区分解为行的功能。这个想法是从缓冲区提供字节，同时从流中异步获取更多字节。

class LineIterator:
    def __init__(self, stream):
        # Iterator to get bytes from stream 
        self.byte_iterator = iter(stream)  
        # Buffer stream bytes until we get a full line
        self.buffer = io.BytesIO()  
        # Track current reading position within buffer
        self.read_pos = 0

   def __iter__(self):
        # Make class iterable 
        return self

    def __next__(self):
        while True:
           # Seek read position within buffer
           self.buffer.seek(self.read_pos)  
           # Try reading a line from current position
           line = self.buffer.readline()
           # If we have a full line
           if line and line[-1] == ord('n'):
               # Increment reading position past this line
               self.read_pos += len(line)  
               # Return the line read without newline char
               return line[:-1] 
           # Fetch next chunk from stream  
           try:
               chunk = next(self.byte_iterator)
           # Handle end of stream 
           except StopIteration:
               # Check if we have any bytes still unread
               if self.read_pos < self.buffer.getbuffer().nbytes:
                   continue
               # If not, raise StopIteration
               raise
           # Add fetched bytes to end of buffer
           self.buffer.seek(0, io.SEEK_END)  
           self.buffer.write(chunk['PayloadPart']['Bytes'])

生成并使用令牌概率作为附加细节

考虑一个我们生成内容的用例。具体来说，我们的任务是为一个以生活方式为中心的网站撰写一段关于定期锻炼的好处的简短段落。我们想要生成内容并输出模型对生成内容的置信度的一些指示性分数。

我们使用提示调用模型端点并捕获生成的响应。我们设置 "details": True 作为模型输入中的运行时参数。由于对数概率是为每个输出标记生成的，因此我们将各个对数概率附加到列表中。我们还捕获响应中生成的完整文本。

sm_client = boto3.client("sagemaker-runtime")

# Set details: True as a runtime parameter within the input.
body = {"inputs": prompt, "parameters": {"max_new_tokens":512, "details": True}}
resp = sm_client.invoke_endpoint_with_response_stream(EndpointName=endpoint_name, Body=json.dumps(body), ContentType="application/json")
event_stream = resp['Body']

overall_log_prob = []

for line in LineIterator(event_stream):
    resp = json.loads(line)
    if resp['token'].get('text') != None:
        token_log_prob = resp['token']['log_prob']
        overall_log_prob.append(token_log_prob)
    elif resp['generated_text'] != None:
        generated_text= resp['generated_text']

为了计算总体置信度得分，我们计算所有单个标记概率的平均值，然后获得 0 到 1 之间的指数值。这是我们推断的生成文本的整体置信度得分，在本例中是一段关于好处的段落定期锻炼。

print(generated_text) 
overall_score=np.exp(np.mean(overall_log_prob)) 
print(f"nnOverall confidence score in the generated text: {overall_score}")

这是如何生成和使用的一个示例 log_prob，在内容生成用例的上下文中。同样，您可以使用 log_prob 作为分类用例的置信度得分的度量。

或者，您可以将其用于整体输出序列或句子级评分，以评估温度等参数对生成输出的影响。

生成并使用完成原因作为附加详细信息

让我们以相同的用例为基础，但这次我们的任务是写一篇更长的文章。此外，我们希望确保输出不会由于生成长度问题（最大令牌长度）或由于遇到停止令牌而被截断。

为了实现这一点，我们使用 finish_reason 输出中生成的属性，监视其值，并继续生成，直到生成整个输出。

我们定义一个推理函数，它接受负载输入并调用 SageMaker 端点、流回响应并处理响应以提取生成的文本。有效负载包含提示文本作为输入和参数，例如最大令牌和详细信息。在流中读取响应并逐行处理，以将生成的文本标记提取到列表中。我们提取细节，例如 finish_reason。我们在循环（链式请求）中调用推理函数，同时每次添加更多上下文，并跟踪生成的令牌数量和发送的请求数量，直到模型完成。

def inference(payload):
    # Call SageMaker endpoint and get response stream
    resp = sm_client.invoke_endpoint_with_response_stream(EndpointName=endpoint_name, Body=json.dumps(payload), ContentType="application/json")
    event_stream = resp['Body']
    text_output = []
    for line in LineIterator(event_stream):
        resp = json.loads(line) 
        # Extract text tokens if present
        if resp['token'].get('text') != None:
            token = resp['token']['text']
            text_output.append(token)  
            print(token, end='')
        # Get finish reason if details present
        if resp.get('details') != None:
            finish_reason = resp['details']['finish_reason']
            # Return extracted output, finish reason and token length
            return payload['inputs'] + ''.join(text_output), finish_reason, len(text_output)

# set details: True as a runtime parameter within the input.
payload = {"inputs": prompt,  "parameters": {"max_new_tokens":256, "details": True}} 

finish_reason = "length"
# Print initial output 
print(f"Output: {payload['inputs']}", end='')  
total_tokens = 0
total_requests = 0
while finish_reason == 'length':
    # Call inference and get extracts
    output_text, finish_reason, out_token_len = inference(payload)
    # Update payload for next request
    payload['inputs'] = output_text 
    total_tokens += out_token_len
    total_requests += 1
# Print metrics
print(f"nntotal tokens generated: {total_tokens} ntotal requests sent: {total_requests}")

正如我们所看到的，尽管 max_new_token 参数设置为 256，我们使用 finish_reason 详细属性作为输出的一部分，将多个请求链接到端点，直到生成整个输出。

同样，根据您的用例，您可以使用 stop_reason 检测为给定任务指定的输出序列长度不足或由于人为停止序列而意外完成。

结论

在本文中，我们介绍了 AWS LMI 容器的 v0.26.0 版本。我们强调了关键的性能改进、新模型支持和新的可用性功能。借助这些功能，您可以更好地平衡成本和性能特征，同时为最终用户提供更好的体验。

要了解有关 LMI DLC 功能的更多信息，请参阅 模型并行性和大模型推理。我们很高兴看到您如何使用 SageMaker 的这些新功能。

关于作者

若昂·莫拉 是 AWS 的高级 AI/ML 专家解决方案架构师。 João 帮助 AWS 客户（从小型初创公司到大型企业）高效地训练和部署大型模型，并在 AWS 上更广泛地构建 ML 平台。

拉胡尔·沙玛 是 AWS 的高级解决方案架构师，帮助 AWS 客户设计和构建 AI/ML 解决方案。在加入 AWS 之前，Rahul 在金融和保险领域工作了几年，帮助客户构建数据和分析平台。

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

健生 是 Amazon Web Services 的软件开发工程师，致力于机器学习系统的几个关键方面。 他是 SageMaker Neo 服务的关键贡献者，专注于深度学习编译和框架运行时优化。 最近，他致力于优化大型模型推理的机器学习系统并做出贡献。

泰勒·奥斯特伯格 是 AWS 的软件开发工程师。他擅长在 SageMaker 中打造高性能机器学习推理体验。最近，他的重点是优化 SageMaker 平台上 Inferentia 深度学习容器的性能。 Tyler 擅长为大型语言模型实施高性能托管解决方案，并使用尖端技术增强用户体验。

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

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

拉古·拉梅沙 是 Amazon SageMaker Service 团队的高级 ML 解决方案架构师。 他专注于帮助客户大规模构建、部署 ML 生产工作负载并将其迁移到 SageMaker。 他专注于机器学习、人工智能和计算机视觉领域，并拥有 UT 达拉斯分校计算机科学硕士学位。 空闲时间，他喜欢旅行和摄影。

• SEO 支持的内容和 PR 分发。 今天得到放大。
• PlatoData.Network 垂直生成人工智能。 赋予自己力量。 访问这里。
• 柏拉图爱流。 Web3 智能。 知识放大。 访问这里。
• 柏拉图ESG。 碳， 清洁科技, 能源， 环境， 太阳能， 废物管理。 访问这里。
• 柏拉图健康。 生物技术和临床试验情报。 访问这里。
• Sumber: https://aws.amazon.com/blogs/machine-learning/boost-inference-performance-for-mixtral-and-llama-2-models-with-new-amazon-sagemaker-containers/