Hugging Face's logo

AWS Trainium 和 Inferentia 文档

在 AWS Inferentia2 上部署 Mixtral 8x7B

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

在 AWS Inferentia2 上部署 Mixtral 8x7B

Mixtral 8x7B 是 Mistral AI 的开源 LLM。它是一个稀疏专家混合模型，架构与 Mistral 7B 类似，但有一个特点：它实际上是将 8 个“专家”模型合而为一。如果您想了解更多关于 MoE 的信息，请查看 专家混合模型详解。

在本教程中，您将学习如何在 Amazon SageMaker 上使用 Hugging Face Optimum Neuron 在 AWS Inferentia2 上部署 mistralai/Mixtral-8x7B-Instruct-v0.1 模型。我们将使用 Hugging Face TGI Neuron 容器，这是一个专门构建的推理容器，可以轻松地在由 Text Generation Inference 和 Optimum Neuron 提供支持的 AWS Inferentia2 上部署 LLM。

我们将介绍如何：

  1. 设置开发环境
  2. 获取最新的 Hugging Face TGI Neuron DLC
  3. 将 Mixtral 8x7B 部署到 Inferentia2
  4. 清理

让我们开始吧！🚀

AWS Inferentia (Inf2) 是专为深度学习 (DL) 推理工作负载构建的 EC2。以下是 Inferentia2 系列的不同实例。

实例大小 加速器 Neuron 核心 加速器内存 vCPU CPU 内存 按需价格 ($/小时)
inf2.xlarge 1 2 32 4 16 0.76
inf2.8xlarge 1 2 32 32 128 1.97
inf2.24xlarge 6 12 192 96 384 6.49
inf2.48xlarge 12 24 384 192 768 12.98

1. 设置开发环境

在本教程中，我们将使用 Amazon SageMaker 中的 Notebook 实例，并结合 Python 3 (ipykernel) 和 sagemaker python SDK，将 Mixtral 8x7B 部署到 SageMaker 推理端点。

确保您已安装最新版本的 SageMaker SDK。

!pip install sagemaker --upgrade --quiet

然后，实例化 sagemaker 角色和会话。

import sagemaker
import boto3

sess = sagemaker.Session()
# sagemaker session bucket -> used for uploading data, models and logs
# sagemaker will automatically create this bucket if it not exists
sagemaker_session_bucket = None
if sagemaker_session_bucket is None and sess is not None:
    # set to default bucket if a bucket name is not given
    sagemaker_session_bucket = sess.default_bucket()

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client("iam")
    role = iam.get_role(RoleName="sagemaker_execution_role")["Role"]["Arn"]

sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

print(f"sagemaker role arn: {role}")
print(f"sagemaker session region: {sess.boto_region_name}")

2. 获取最新的 Hugging Face TGI Neuron DLC

最新的 Hugging Face TGI Neuron DLC 可用于在 AWS Inferentia2 上运行推理。您可以使用 sagemaker SDK 的 get_huggingface_llm_image_uri 方法，根据您所需的 backend、session、region 和 version，检索合适的 Hugging Face TGI Neuron DLC URI。如果尚未添加到 SageMaker SDK，您可以在此处找到容器的最新版本。

from sagemaker.huggingface import get_huggingface_llm_image_uri

# retrieve the llm image uri
llm_image = get_huggingface_llm_image_uri(
    "huggingface-neuronx",
)

# print ecr image uri
print(f"llm image uri: {llm_image}")

3. 将 Mixtral 8x7B 部署到 Inferentia2

在撰写本文时，AWS Inferentia2 不支持推理的动态形状，这意味着我们需要提前指定序列长度和批次大小。为了方便客户充分利用 Inferentia2 的强大功能，我们创建了一个 neuron 模型缓存，其中包含最流行的 LLM（包括 Mixtral 8x7B）的预编译配置。

这意味着我们不需要自己编译模型，而是可以使用缓存中的预编译模型。您可以在 Hugging Face Hub 上找到已编译/缓存的配置。如果您所需的配置尚未缓存，您可以使用 Optimum CLI 自己编译，或在 Cache 仓库 中提出请求。

让我们检查缓存中不同的配置。为此，您首先需要使用具有读取权限的 用户访问令牌 登录 Hugging Face Hub。

确保您具有访问该模型的必要权限。您可以在此处请求访问该模型。

from huggingface_hub import notebook_login

notebook_login()

然后，我们需要安装最新版本的 Optimum Neuron。

!pip install optimum-neuron --upgrade --quiet

最后，我们可以查询缓存并检索我们维护了模型编译版本的现有配置集。

HF_MODEL_ID = "mistralai/Mixtral-8x7B-Instruct-v0.1"

!optimum-cli neuron cache lookup $HF_MODEL_ID

您应该在缓存中检索到两个条目

*** 2 entrie(s) found in cache for mistralai/Mixtral-8x7B-Instruct-v0.1 for inference.***

auto_cast_type: bf16
batch_size: 1
checkpoint_id: mistralai/Mixtral-8x7B-Instruct-v0.1
checkpoint_revision: 41bd4c9e7e4fb318ca40e721131d4933966c2cc1
compiler_type: neuronx-cc
compiler_version: 2.16.372.0+4a9b2326
num_cores: 24
sequence_length: 4096
task: text-generation

auto_cast_type: bf16
batch_size: 4
checkpoint_id: mistralai/Mixtral-8x7B-Instruct-v0.1
checkpoint_revision: 41bd4c9e7e4fb318ca40e721131d4933966c2cc1
compiler_type: neuronx-cc
compiler_version: 2.16.372.0+4a9b2326
num_cores: 24
sequence_length: 4096
task: text-generation

将 Mixtral 8x7B 部署到 SageMaker 端点

在将模型部署到 Amazon SageMaker 之前，我们必须定义 TGI Neuron 端点配置。我们需要确保定义以下附加参数：

  • HF_NUM_CORES：用于编译的 Neuron 核心数。
  • HF_BATCH_SIZE：用于编译模型的批次大小。
  • HF_SEQUENCE_LENGTH：用于编译模型的序列长度。
  • HF_AUTO_CAST_TYPE：用于编译模型的自动转换类型。

我们仍然需要使用以下参数定义传统的 TGI 参数：

  • HF_MODEL_ID：Hugging Face 模型 ID。
  • HF_TOKEN：用于访问 gated 模型的 Hugging Face API 令牌。
  • MAX_BATCH_SIZE：模型可以处理的最大批次大小，等于用于编译的批次大小。
  • MAX_INPUT_TOKEN：模型可以处理的最大输入长度。
  • MAX_TOTAL_TOKENS：模型可以生成的最大令牌总数，等于用于编译的序列长度。

可选地，您可以配置端点以支持聊天模板

  • MESSAGES_API_ENABLED：启用 Messages API

选择正确的实例类型

Mixtral 8x7B 是一个大型模型，需要大量内存。我们将使用 inf2.48xlarge 实例类型，它具有 192 个 vCPU 和 384 GB 的加速器内存。inf2.48xlarge 实例配备了 12 个 Inferentia2 加速器，其中包括 24 个 Neuron 核心。在我们的例子中，我们将使用批次大小为 4，序列长度为 4096。

之后，我们可以创建端点配置并将模型部署到 Amazon SageMaker。我们将启用 Messages API 部署端点，以便它与 OpenAI Chat Completion API 完全兼容。

from sagemaker.huggingface import HuggingFaceModel

# sagemaker config
instance_type = "ml.inf2.48xlarge"
health_check_timeout = 2400  # additional time to load the model
volume_size = 512  # size in GB of the EBS volume

# Define Model and Endpoint configuration parameter
config = {
    "HF_MODEL_ID": "mistralai/Mixtral-8x7B-Instruct-v0.1",
    "HF_NUM_CORES": "24",  # number of neuron cores
    "HF_AUTO_CAST_TYPE": "bf16",  # dtype of the model
    "MAX_BATCH_SIZE": "4",  # max batch size for the model
    "MAX_INPUT_TOKENS": "4000",  # max length of input text
    "MAX_TOTAL_TOKENS": "4096",  # max length of generated text
    "MESSAGES_API_ENABLED": "true",  # Enable the messages API
    "HF_TOKEN": "<REPLACE WITH YOUR TOKEN>",
}

assert (
    config["HF_TOKEN"] != "<REPLACE WITH YOUR TOKEN>"
), "Please replace '<REPLACE WITH YOUR TOKEN>' with your Hugging Face Hub API token"


# create HuggingFaceModel with the image uri
llm_model = HuggingFaceModel(role=role, image_uri=llm_image, env=config)

创建 HuggingFaceModel 后，我们可以使用 deploy 方法将其部署到 Amazon SageMaker。我们将使用 ml.inf2.48xlarge 实例类型部署模型。TGI 将自动在所有 Inferentia 设备上分发和分片模型。

# Deploy model to an endpoint
# https://sagemaker.readthedocs.io/en/stable/api/inference/model.html#sagemaker.model.Model.deploy
llm_model._is_compiled_model = True

llm = llm_model.deploy(
    initial_instance_count=1,
    instance_type=instance_type,
    container_startup_health_check_timeout=health_check_timeout,
    volume_size=volume_size,
)

SageMaker 现在将创建我们的端点并将模型部署到该端点。部署大约需要 15 分钟。

在我们的端点部署完成后，我们可以在其上运行推理。我们将使用 predictor 中的 predict 方法在我们的端点上运行推理。

该端点支持 Messages API，它与 OpenAI Chat Completion API 完全兼容。Messages API 允许我们以对话方式与模型进行交互。我们可以定义消息的角色和内容。角色可以是 system、assistant 或 user。system 角色用于向模型提供上下文，而 user 角色用于向模型提问或提供输入。

参数可以像在 payload 的 parameters 属性中定义的那样进行定义。查看聊天完成 文档 以查找支持的参数。

{
  "messages": [
    { "role": "system", "content": "You are a helpful assistant." },
    { "role": "user", "content": "What is deep learning?" }
  ]
}
# Prompt to generate
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "What is deep learning in one sentence?"},
]

# Generation arguments https://platform.openai.com/docs/api-reference/chat/create
parameters = {
    "max_tokens": 100,
}

好的，让我们测试一下。

chat = llm.predict({"messages": messages, **parameters, "steam": True})

print(chat["choices"][0]["message"]["content"].strip())

4. 清理

要清理，我们可以删除模型和端点。

llm.delete_model()
llm.delete_endpoint()

←在 SageMaker 上使用 AWS Inferentia2 部署 Llama 3.3 70B Amazon EC2 入门→