在 Google Cloud Vertex AI 上部署 Meta Llama 3.1 405B

发布日期:2024 年 8 月 19 日
在 GitHub 上更新

Meta Llama 3.1 是 Meta 于 2024 年 7 月发布的最新开源 LLM。Meta Llama 3.1 有三种大小:8B 用于在消费级 GPU 上高效部署和开发,70B 用于大规模 AI 原生应用,以及 405B 用于合成数据、作为法官的 LLM 或蒸馏等用例。其主要功能包括:128K token 的大上下文长度(与原始的 8K 相比),多语言能力,工具使用能力,以及更宽松的许可。

在本博客中,你将学习如何在 Google Cloud Vertex AI 的 Google Cloud A3 节点上以编程方式部署 meta-llama/Meta-Llama-3.1-405B-Instruct-FP8,即 meta-llama/Meta-Llama-3.1-405B-Instruct 的 FP8 量化变体,该节点配备 8 个 H100 NVIDIA GPU,并使用 文本生成推理 (TGI) 和 Hugging Face 专门为 Google Cloud 构建的深度学习容器 (DLC)。

另外,你也可以直接从 HubVertex 模型花园 部署 meta-llama/Meta-Llama-3.1-405B-Instruct-FP8,无需编写任何代码!

本博客将涵盖

Vertex AI 简介

  1. Google Cloud 上 Meta Llama 3.1 模型的要求
  2. 为 Vertex AI 设置 Google Cloud
  3. 在 Vertex AI 上注册 Meta Llama 3.1 405B 模型
  4. 在 Vertex AI 上部署 Meta Llama 3.1 405B
  5. 使用 Meta Llama 3.1 405B 运行在线预测
    1. 通过 Python
      1. 在同一会话中
      2. 从不同会话中
    2. 通过 Vertex AI 在线预测用户界面
  6. 清理资源

结论

让我们开始吧!🚀 或者,你也可以从 此 Jupyter Notebook 开始学习。

Vertex AI 简介

Vertex AI 是一个机器学习 (ML) 平台,可让你训练和部署 ML 模型和 AI 应用,并自定义大型语言模型 (LLM) 以用于你的 AI 驱动应用。Vertex AI 结合了数据工程、数据科学和 ML 工程工作流程,使你的团队能够使用通用工具集进行协作,并利用 Google Cloud 的优势扩展你的应用。

本博客将重点介绍如何从 Hugging Face Hub 部署一个已微调的模型,使用预构建的容器获取实时在线预测。因此,我们将演示 Vertex AI 在推理方面的使用。

更多信息请参见 Vertex AI - 文档 - Vertex AI 简介

1. Google Cloud 上 Meta Llama 3.1 模型的要求

Meta Llama 3.1 带来了令人兴奋的进步。然而,运行这些模型需要仔细考虑你的硬件资源。对于推理,内存需求取决于模型大小和权重精度。下表显示了不同配置所需的近似内存:

模型大小 FP16 FP8 INT4
8B 16 GB 8 GB 4 GB
70B 140 GB 70 GB 35 GB
405B 810 GB 405 GB 203 GB

注意:上述引用数字表示仅加载模型检查点所需的 GPU 显存。它们不包括用于内核或 CUDA 图的 torch 保留空间。

例如,一个 H100 节点(8 个 H100,每个 80GB)总共有约 640GB 的显存,因此 405B 模型需要在多节点设置中运行,或者以较低精度(例如 FP8)运行,这将是推荐的方法。有关更多信息,请参阅 Hugging Face 博客上的 Meta Llama 3.1

Google Cloud 中的 A3 加速器优化机器系列配备 8 个 H100 80GB NVIDIA GPU、208 个 vCPU 和 1872 GB 内存。该机器系列针对计算和内存密集型、网络受限的 ML 训练和 HPC 工作负载进行了优化。有关 A3 机器可用性公告的更多信息,请参阅 推出专为 AI 构建的配备 NVIDIA H100 GPU 的 A3 超级计算机,有关 A3 机器系列的更多信息,请参阅 Compute Engine - 加速器优化机器系列

即使 A3 机器在 Google Cloud 中可用,你仍然需要请求自定义配额增加,因为这些需要特定批准。请注意,A3 机器仅在某些区域可用,因此请务必在 Compute Engine - GPU 区域和可用区 中检查 A3 High 或 A3 Mega 在每个区域的可用性。

在这种情况下,要请求增加配额以使用 A3 High GPU 机器类型,你需要增加以下配额:

  • Service: Vertex AI APIName: Custom model serving Nvidia H100 80GB GPUs per region 设置为 8
  • Service: Vertex AI APIName: Custom model serving A3 CPUs per region 设置为 208

A3 Quota Request in Google Cloud

有关如何请求配额增加的更多信息,请参阅 Google Cloud 文档 - 查看和管理配额

2. 为 Vertex AI 设置 Google Cloud

在继续之前,我们将为方便起见设置以下环境变量

%env PROJECT_ID=your-project-id
%env LOCATION=your-region

首先,你需要根据 Cloud SDK - 安装 gcloud CLI 中的说明在你的机器上安装 gcloud;并登录你的 Google Cloud 账户,设置你的项目和首选的 Google Compute Engine 区域。

gcloud auth login
gcloud config set project $PROJECT_ID
gcloud config set compute/region $LOCATION

安装 Google Cloud SDK 后,你需要启用使用其 Artifact Registry for Docker 中的深度学习容器 (DLC) 来使用 Vertex AI 所需的 Google Cloud API。

gcloud services enable aiplatform.googleapis.com
gcloud services enable compute.googleapis.com
gcloud services enable container.googleapis.com
gcloud services enable containerregistry.googleapis.com
gcloud services enable containerfilesystem.googleapis.com

然后,你还需要安装 google-cloud-aiplatform,这是从 Python 以编程方式与 Google Cloud Vertex AI 交互所必需的。

pip install --upgrade --quiet google-cloud-aiplatform

然后通过 Python 如下初始化它

import os
from google.cloud import aiplatform

aiplatform.init(project=os.getenv("PROJECT_ID"), location=os.getenv("LOCATION"))

最后,由于 Meta Llama 3.1 模型在 Hugging Face Hub 的 meta-llama 组织 下受限,你需要请求访问并等待批准,这通常不会超过 24 小时。然后,你需要安装 huggingface_hub Python SDK 以使用 huggingface-cli 登录 Hugging Face Hub 以下载这些模型。

pip install --upgrade --quiet huggingface_hub

或者,你也可以跳过 huggingface_hub 的安装,直接生成一个 Hugging Face 细粒度令牌,该令牌对受限仓库 meta-llama/Meta-Llama-3.1-405B-Instruct-FP8meta-llama 组织 下的任何其他模型具有只读权限,选择例如 仓库权限 -> meta-llama/Meta-Llama-3.1-405B-Instruct-FP8 -> 对所选仓库内容的读取权限。然后将该令牌设置为 HF_TOKEN 环境变量,或者直接手动提供给 notebook_login 方法,如下所示

from huggingface_hub import notebook_login

notebook_login()

3. 在 Vertex AI 上注册 Meta Llama 3.1 405B 模型

要在 Vertex AI 上注册 Meta Llama 3.1 405B 模型,你需要使用 google-cloud-aiplatform Python SDK。但在继续之前,你需要首先定义要使用的 DLC,在这种情况下,它将是 Hugging Face TGI 最新版 GPU DLC。

截至目前(2024 年 8 月),最新的 Hugging Face TGI DLC,即 us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-2.ubuntu2204.py310 使用 TGI v2.2。此版本支持 Meta Llama 3.1 架构,该架构需要与之前的 Meta Llama 3 不同的 RoPE 缩放方法。

要检查 Google Cloud 中有哪些 Hugging Face DLC 可用,你可以导航到 Google Cloud Artifact Registry 并按“huggingface-text-generation-inference”过滤,或者使用以下 gcloud 命令

gcloud container images list --repository="us-docker.pkg.dev/deeplearning-platform-release/gcr.io" | grep "huggingface-text-generation-inference"

然后,你需要为容器定义配置,即 text-generation-launcher 预期作为参数的环境变量(根据 官方文档),在这种情况下,如下所示:

  • MODEL_ID Hugging Face Hub 上的模型 ID,即 meta-llama/Meta-Llama-3.1-405B-Instruct-FP8
  • HUGGING_FACE_HUB_TOKEN 对受限仓库 meta-llama/Meta-Llama-3.1-405B-Instruct-FP8 的只读访问令牌,下载 Hugging Face Hub 中的权重所需。
  • NUM_SHARD 要使用的分片数,即要使用的 GPU 数量,在这种情况下设置为 8,因为将使用具有 8 个 H100 NVIDIA GPU 的 A3 实例。

此外,作为建议,您还应该定义 HF_HUB_ENABLE_HF_TRANSFER=1 以通过 hf_transfer 工具实现更快的下载速度,因为 Meta Llama 3.1 405B 大约 400 GiB,否则下载权重可能需要更长时间。

然后,您可以通过 google-cloud-aiplatform Python SDK 在 Vertex AI 的模型注册表中注册模型,如下所示:

from huggingface_hub import get_token

model = aiplatform.Model.upload(
    display_name="meta-llama--Meta-Llama-3.1-405B-Instruct-FP8",
    serving_container_image_uri="us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-2.ubuntu2204.py310",
    serving_container_environment_variables={
        "MODEL_ID": "meta-llama/Meta-Llama-3.1-405B-Instruct-FP8",
        "HUGGING_FACE_HUB_TOKEN": get_token(),
        "HF_HUB_ENABLE_HF_TRANSFER": "1",
        "NUM_SHARD": "8",
    },
)
model.wait()

Meta Llama 3.1 405B FP8 registered on Vertex AI

4. 在 Vertex AI 上部署 Meta Llama 3.1 405B

Meta Llama 3.1 405B 在 Vertex AI 模型注册表中注册后,您可以创建 Vertex AI 端点并将模型部署到该端点,并使用 Hugging Face TGI 的 DLC 作为服务容器。

如前所述,由于 FP8 中的 Meta Llama 3.1 405B 占用约 400 GiB 的磁盘空间,这意味着我们需要至少 400 GiB 的 GPU VRAM 来加载模型,并且节点内的 GPU 需要支持 FP8 数据类型。在这种情况下,将使用一台配备 8 个 NVIDIA H100 80GB 的 A3 实例,总共约 640 GiB VRAM,以加载模型,同时为 KV 缓存和 CUDA 图形留出一些空闲 VRAM。

endpoint = aiplatform.Endpoint.create(display_name="Meta-Llama-3.1-405B-FP8-Endpoint")

deployed_model = model.deploy(
    endpoint=endpoint,
    machine_type="a3-highgpu-8g",
    accelerator_type="NVIDIA_H100_80GB",
    accelerator_count=8,
)

请注意,在 Vertex AI 上部署 Meta Llama 3.1 405B 可能需要大约 25-30 分钟,因为它需要分配 Google Cloud 上的资源、从 Hugging Face Hub 下载权重(约 10 分钟)以及将其加载到 TGI 中进行推理(约 2 分钟)。

Meta Llama 3.1 405B Instruct FP8 deployed on Vertex AI

恭喜,您已成功在 Google Cloud 账户中部署了 Meta Llama 3.1 405B!🔥 现在是时候对模型进行测试了。

5. 使用 Meta Llama 3.1 405B 运行在线预测

Vertex AI 将在 /predict 路由中公开一个在线预测端点,该端点正在提供来自文本生成推理 (TGI) DLC 的文本生成,确保 I/O 数据符合 Vertex AI 有效负载(有关 Vertex AI I/O 有效负载的更多信息,请参阅 Vertex AI 文档 - 从自定义训练模型获取在线预测)。

由于公开的端点是 /generate,在向 Vertex AI 发送请求之前,你需要使用聊天模板格式化消息,因此建议安装 🤗transformers 以使用 PreTrainedTokenizerFast 分词器实例的 apply_chat_template 方法。

pip install --upgrade --quiet transformers

然后使用分词器将聊天模板应用于对话,如下所示

import os
from huggingface_hub import get_token
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    "meta-llama/Meta-Llama-3.1-405B-Instruct-FP8",
    token=get_token(),
)

messages = [
    {"role": "system", "content": "You are an assistant that responds as a pirate."},
    {"role": "user", "content": "What's the Theory of Relativity?"},
]
inputs = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)

现在你已经从初始对话消息中得到一个字符串,使用 Meta Llama 3.1 的默认聊天模板进行了格式化

<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are an assistant that responds as a pirate.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nWhat's the Theory of Relativity?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n

这就是你将在有效负载中发送到已部署的 Vertex AI 端点的内容,以及生成参数,如 使用文本生成推理 (TGI) -> 生成 中所述。

5.1 通过 Python

5.1.1 在同一会话中

如果你想在当前会话中(即用于部署模型的会话)运行在线预测,你可以通过 aiplatform.Model.deploy 方法返回的 aiplatform.Endpoint 以编程方式发送请求,如下面的代码片段所示。

output = deployed_model.predict(
    instances=[
        {
            "inputs": inputs,
            "parameters": {
                "max_new_tokens": 128,
                "do_sample": True,
                "top_p": 0.95,
                "temperature": 0.7,
            },
        },
    ]
)

生成以下 output

预测(预测结果=[“你想了解那些花哨的科学玩意儿,嗯?那好,伙伴,喝杯朗姆酒坐下来,仔细听着。我要告诉你的是那个大胆的天才阿尔伯特·爱因斯坦提出的相对论。\n\n你看,爱因斯坦说时间和空间就像大海和风一样相互关联。你不能只有其中一个而没有另一个,明白吗?他还提出,你如何看待时间和空间取决于你移动的速度以及你所处的位置。这就叫做相对论,我的”], deployed_model_id='', metadata=None, model_version_id='1', model_resource_name='projects//locations//models/', explanations=None)

5.1.2 从不同会话中

如果 Vertex AI 端点是在不同会话中部署的,而您只是想使用它,但无法访问 aiplatform.Model.deploy 方法返回的 deployed_model 变量,那么您也可以运行以下代码片段来通过其资源名称实例化已部署的 aiplatform.Endpoint,该资源名称可以在 Vertex AI 在线预测 UI 中找到,也可以从上面实例化的 aiplatform.Endpoint 中获取,或者直接替换 projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID} 中的值。

import os
from google.cloud import aiplatform

aiplatform.init(project=os.getenv("PROJECT_ID"), location=os.getenv("LOCATION"))

endpoint = aiplatform.Endpoint(f"projects/{os.getenv('PROJECT_ID')}/locations/{os.getenv('LOCATION')}/endpoints/{ENDPOINT_ID}")
output = endpoint.predict(
    instances=[
        {
            "inputs": inputs,
            "parameters": {
                "max_new_tokens": 128,
                "do_sample": True,
                "top_p": 0.95,
                "temperature": 0.7,
            },
        },
    ],
)

生成以下 output

预测 (predictions=["你正在寻找关于那些高级物理学的宝藏知识,嗯?好的,伙伴,坐下来,喝杯朗姆酒,仔细听着,我将为你讲述爱因斯坦的相对论的故事。\n\n它分为两部分,我的朋友:狭义相对论和广义相对论。现在,我知道你在想什么:有什么区别呢?好吧,伙伴,让我为你分解一下。\n\n狭义相对论是这样一种观点:时间和空间像大海和天空一样相互连接。"], deployed_model_id='', metadata=None, model_version_id='1', model_resource_name='projects//locations//models/', explanations=None)

5.2 通过 Vertex AI 在线预测用户界面

或者,出于测试目的,你也可以使用 Vertex AI 在线预测用户界面,该界面提供了一个字段,用于输入根据 Vertex AI 规范(如上例所示)格式化的 JSON 有效负载:

{
    "instances": [
        {
            "inputs": "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are an assistant that responds as a pirate.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nWhat's the Theory of Relativity?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n",
            "parameters": {
                "max_new_tokens": 128,
                "do_sample": true,
                "top_p": 0.95,
                "temperature": 0.7
            }
        }
    ]
}

这样就可以在用户界面中生成并打印输出。

Meta Llama 3.1 405B Instruct FP8 online prediction on Vertex AI

6. 清理资源

完成后,您可以按如下方式释放您创建的资源,以避免不必要的成本。

  • deployed_model.undeploy_all 用于从所有端点取消部署模型。
  • deployed_model.delete 用于在 undeploy_all 方法后,从部署模型的端点中优雅地删除模型。
  • model.delete 用于从注册表中删除模型。
deployed_model.undeploy_all()
deployed_model.delete()
model.delete()

或者,你也可以按照以下步骤从 Google Cloud Console 中删除这些资源:

  • 前往 Google Cloud 中的 Vertex AI
  • 前往“部署和使用”->“在线预测”
  • 点击端点,然后点击已部署的模型以“从端点取消部署模型”
  • 然后返回端点列表并删除该端点
  • 最后,前往“部署和使用”->“模型注册表”,然后删除模型

结论

就是这样!您已成功在 Google Cloud Vertex AI 上注册并部署了 Meta Llama 3.1 405B Instruct FP8,然后通过编程方式和 Google Cloud 控制台运行了在线预测,最后清理了使用的资源以避免不必要的成本。

多亏了 Hugging Face 的 Text Generation Inference (TGI) DLC 和 Google Cloud Vertex AI,部署用于服务大型语言模型 (LLM) 的高性能文本生成容器从未如此简单。我们不会止步于此——敬请期待,我们将在 Google Cloud 上实现更多使用开放模型构建 AI 的体验!

社区

注册登录 发表评论