使用 SAX 在 TPU v5e 上部署和提供 Llama 2 模型

创建 TPU

以下步骤展示了如何创建用于您的模型的 TPU 虚拟机。

创建环境变量：

 export PROJECT_ID=PROJECT_ID export ACCELERATOR_TYPE=ACCELERATOR_TYPE export ZONE=ZONE export RUNTIME_VERSION=v2-alpha-tpuv5-lite export SERVICE_ACCOUNT=SERVICE_ACCOUNT export TPU_NAME=TPU_NAME export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID

环境变量说明

`PROJECT_ID` 您的 Google Cloud 项目的 ID。 `ACCELERATOR_TYPE` 加速器类型指定 您要创建的 Cloud TPU。不同大小的 Llama 2 不同的 TPU 大小要求： <ph type="x-smartling-placeholder"> </ph>+ **7B**：`v5litepod-4` 或更大 + **13B**：`v5litepod-8` 或更大 + **70B**：`v5litepod-16` 或更大 `ZONE` 您要在哪个[可用区](https://cloud.google.com/tpu/docs/regions-zones?hl=zh-cn)中创建 Cloud TPU。 `SERVICE_ACCOUNT` 您要关联到 Cloud TPU 的服务账号。 `TPU_NAME` Cloud TPU 的名称。 `QUEUED_RESOURCE_ID` 已加入队列的资源请求的标识符。

在您的活跃 Google Cloud CLI 配置中设置项目 ID 和可用区：

gcloud config set project $PROJECT_ID && gcloud config set compute/zone $ZONE

创建 TPU 虚拟机：

gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \ --node-id ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --accelerator-type ${ACCELERATOR_TYPE} \ --runtime-version ${RUNTIME_VERSION} \ --service-account ${SERVICE_ACCOUNT}

检查 TPU 是否处于活跃状态：

gcloud compute tpus queued-resources list --project $PROJECT_ID --zone $ZONE

设置检查点转换节点

如需在 SAX 集群上运行 LLama 模型，您需要将原始 Llama 检查点设置为 SAX 兼容格式。

转换需要大量内存资源，具体取决于模型大小：

对于 7B 和 13B 模型，您可以在 TPU 虚拟机上运行转换。对于 70B 模型，则需要创建一个 Compute Engine 实例，大约 1TB 磁盘空间：

gcloud compute instances create INSTANCE_NAME --project=$PROJECT_ID --zone=$ZONE \ --machine-type=n2-highmem-128 \ --network-interface=network-tier=PREMIUM,stack-type=IPV4_ONLY,subnet=default \ --maintenance-policy=MIGRATE --provisioning-model=STANDARD \ --service-account=$SERVICE\_ACCOUNT \ --scopes=https://www.googleapis.com/auth/cloud-platform \ --tags=http-server,https-server \ --create-disk=auto-delete=yes,boot=yes,device-name=bk-workday-dlvm,image=projects/ml-images/global/images/c0-deeplearning-common-cpu-v20230925-debian-10,mode=rw,size=500,type=projects/$PROJECT_ID/zones/$ZONE/diskTypes/pd-balanced \ --no-shielded-secure-boot \ --shielded-vtpm \ --shielded-integrity-monitoring \ --labels=goog-ec-src=vm_add-gcloud \ --reservation-affinity=any

无论您使用 TPU 还是 Compute Engine 实例作为转化服务器，设置您的服务器以转换 Llama 2 检查点：

对于 7B 和 13B 模型，将将服务器名称环境变量设置为 TPU 的名称：

export CONV_SERVER_NAME=$TPU_NAME

对于 70B 模型，请将服务器名称环境变量设置为您的 Compute Engine 实例：

export CONV_SERVER_NAME=INSTANCE_NAME

使用 SSH 连接到转化节点。

如果您的转化节点是 TPU，请连接到 TPU：

gcloud compute tpus tpu-vm ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE

如果您的转化节点是 Compute Engine 实例，请连接到 Compute Engine 虚拟机：

gcloud compute ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE

在转化节点上安装所需的软件包：

sudo apt update sudo apt-get install python3-pip sudo apt-get install git-all pip3 install paxml==1.1.0 pip3 install torch pip3 install jaxlib==0.4.14

下载 Llama 检查点转换脚本：

gcloud storage cp gs://cloud-tpu-inference-public/sax-tokenizers/llama/convert_llama_ckpt.py .

下载 Llama 2 权重

在转换模型之前，您需要下载 Llama 2 权重。为此所以必须使用原始的 Llama 2 权重（例如， meta-llama/Llama-2-7b）而非已针对 Hugging Face Transformer 格式转换的权重（例如， meta-llama/Llama-2-7b-hf).

如果您已经有了 Llama 2 权重，请直接跳到将权重。

要从 Hugging Face Hub 下载权重，您需要将设置一个用户访问令牌，并且请求访问 Llama 2 模型。如需申请访问权限，请按照说明操作找到想要使用的模型，例如 meta-llama/Llama-2-7b.

为权重创建一个目录：

 sudo mkdir WEIGHTS_DIRECTORY

从 Hugging Face Hub 获取 Llama2 权重：

安装 Hugging Face Hub CLI：

pip install -U "huggingface_hub[cli]"

切换到 weights 目录：

cd WEIGHTS_DIRECTORY

下载 Llama 2 文件：

python3 from huggingface_hub import login login() from huggingface_hub import hf_hub_download, snapshot_download import os PATH=os.getcwd() snapshot_download(repo_id="meta-llama/LLAMA2_REPO", local_dir_use_symlinks=False, local_dir=PATH) 将 LLAMA2_REPO 替换为 Hugging Face 代码库的名称 您要下载的文件：`Llama-2-7b`、`Llama-2-13b`或`Llama-2-70b`。

转换权重

修改转化脚本，然后运行转化脚本来转换模型权重。

创建一个目录来保存转换后的权重：

 sudo mkdir CONVERTED_WEIGHTS

将 Saxml GitHub 代码库克隆到您已进行读取、写入、和执行权限：

 git clone https://github.com/google/saxml.git -b r1.1.0

切换到 saxml 目录：

cd saxml

打开 saxml/tools/convert_llama_ckpt.py 文件。

在 saxml/tools/convert_llama_ckpt.py 文件中，将第 169 行从：

'scale': pytorch_vars[0]['layers.%d.attention_norm.weight' % (layer_idx)].type(torch.float16).numpy()

收件人：

'scale': pytorch_vars[0]['norm.weight'].type(torch.float16).numpy()

运行 saxml/tools/init_cloud_vm.sh 脚本：

saxml/tools/init_cloud_vm.sh

仅适用于 70B：关闭测试模式：

打开 saxml/server/pax/lm/params/lm_cloud.py 文件。

在 saxml/server/pax/lm/params/lm_cloud.py 将第 344 行从：

return True

收件人：

return False

转换权重：

python3 saxml/tools/convert_llama_ckpt.py --base-model-path WEIGHTS_DIRECTORY \ --pax-model-path CONVERTED_WEIGHTS \ --model-size MODEL_SIZE

替换以下内容： + WEIGHTS_DIRECTORY：原始权重的目录。 + CONVERTED_WEIGHTS：已转换权重的目标路径。 + MODEL_SIZE：`7b`、`13b` 或 `70b`。

准备检查点目录

转换检查点后，检查点目录应包含以下结构：

checkpoint_00000000 metadata/ metadata state/ mdl_vars.params.lm*/ ... ... step/

创建一个名为 commit_success.txt 的空文件，并将其副本放入 checkpoint_00000000、metadata 和 state 目录，SAX 可通过这些信息了解确认此检查点已完全转换并可以开始加载：

切换到检查点目录：

cd CONVERTED_WEIGHTS/checkpoint_00000000

创建一个名为 commit_success.txt 的空文件：

touch commit_success.txt

切换到元数据目录，然后创建一个名为 commit_success.txt 的空文件：

cd metadata && touch commit_success.txt

切换到状态目录并创建名为 commit_success.txt 的空文件：

cd .. && cd state && touch commit_success.txt

检查点目录现在应具有以下结构：

checkpoint_00000000 commit_success.txt metadata/ commit_success.txt metadata state/ commit_success.txt mdl_vars.params.lm*/ ... ... step/

创建 Cloud Storage 存储桶

您需要将转换后 Cloud Storage 存储桶中的检查点可用。

为 Cloud Storage 存储桶的名称设置环境变量：

export GSBUCKET=BUCKET_NAME

创建存储桶：

gcloud storage buckets create gs://${GSBUCKET}

将转换后的检查点文件复制到您的存储桶：

 gcloud storage cp -r CONVERTED_WEIGHTS/checkpoint_00000000 gs://$GSBUCKET/sax_models/llama2/SAX_LLAMA2_DIR/

将 SAX_LLAMA2_DIR 替换为适当的值：

+ **70 亿**：`saxml_llama27b` + **130 亿**：`saxml_llama213b` + **700 亿**：`saxml_llama270b`

创建 SAX 集群

如需创建 SAX 集群，您需要执行以下操作：

创建管理服务器创建模型服务器将模型部署到模型服务器

在典型部署中，您需要在 Compute Engine 上运行管理服务器，实例以及 TPU 或 GPU 上的模型服务器。为了本教程中，您需要将管理服务器和模型服务器部署在同一 TPU 上 v5e 实例。

创建管理服务器

创建管理服务器 Docker 容器：

在转化服务器上，安装 Docker：

sudo apt-get update sudo apt-get install docker.io

启动管理服务器 Docker 容器：

sudo docker run --name sax-admin-server \ -it \ -d \ --rm \ --network host \ --env GSBUCKET=${GSBUCKET} us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.1.0

您可以运行不带 -d 选项的 docker run 命令，以查看日志和确保管理服务器正常启动。

创建模型服务器

以下部分介绍了如何创建模型服务器。

7b 模型

启动模型服务器 Docker 容器：

sudo docker run --privileged \ -it \ -d \ --rm \ --network host \ --name "sax-model-server" \ --env SAX_ROOT=gs://${GSBUCKET}/sax-root us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \ --sax_cell="/sax/test" \ --port=10001 \ --platform_chip=tpuv5e \ --platform_topology='4'

13b 模型

“lm_cloud.py”中缺少“LLaMA13BFP16TPUv5e”的配置。通过以下步骤显示了如何更新 lm_cloud.py 并提交新的 Docker 映像。

启动模型服务器：

sudo docker run --privileged \ -it \ -d \ --rm \ --network host \ --name "sax-model-server" \ --env SAX_ROOT=gs://${GSBUCKET}/sax-root \ us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \ --sax_cell="/sax/test" \ --port=10001 \ --platform_chip=tpuv5e \ --platform_topology='8'

使用 SSH 连接到 Docker 容器：

sudo docker exec -it sax-model-server bash

在 Docker 映像中安装 Vim：

$ apt update $ apt install vim

打开 saxml/server/pax/lm/params/lm_cloud.py 文件。搜索 LLaMA13B。您应该会看到以下代码：

@servable_model_registry.register @quantization.for_transformer(quantize_on_the_fly=False) class LLaMA13B(BaseLLaMA): """13B model on a A100-40GB. April 12, 2023 Latency = 5.06s with 128 decoded tokens. 38ms per output token. """ NUM_LAYERS = 40 VOCAB_SIZE = 32000 DIMS_PER_HEAD = 128 NUM_HEADS = 40 MODEL_DIMS = 5120 HIDDEN_DIMS = 13824 ICI_MESH_SHAPE = [1, 1, 1] @property def test_mode(self) -> bool: return True

注释或删除以 @quantization 开头的行。在此之后更改，文件应如下所示：

@servable_model_registry.register class LLaMA13B(BaseLLaMA): """13B model on a A100-40GB. April 12, 2023 Latency = 5.06s with 128 decoded tokens. 38ms per output token. """ NUM_LAYERS = 40 VOCAB_SIZE = 32000 DIMS_PER_HEAD = 128 NUM_HEADS = 40 MODEL_DIMS = 5120 HIDDEN_DIMS = 13824 ICI_MESH_SHAPE = [1, 1, 1] @property def test_mode(self) -> bool: return True

添加以下代码以支持 TPU 配置。

@servable_model_registry.register class LLaMA13BFP16TPUv5e(LLaMA13B): """13B model on TPU v5e-8. """ BATCH_SIZE = [1] BUCKET_KEYS = [128] MAX_DECODE_STEPS = [32] ENABLE_GENERATE_STREAM = False ICI_MESH_SHAPE = [1, 1, 8] @property def test_mode(self) -> bool: return False

退出 Docker 容器 SSH 会话：

exit

将更改提交至新的 Docker 映像：

sudo docker commit sax-model-server sax-model-server:v1.1.0-mod

检查是否已创建新的 Docker 映像：

sudo docker images

您可以将 Docker 映像发布到项目的 Artifact Registry 中教程将继续使用本地映像

停止模型服务器。本教程的其余部分将使用更新后的模型服务器。

```auto sudo docker stop sax-model-server ```

使用更新后的 Docker 映像启动模型服务器。请务必指定更新后的映像名称 sax-model-server:v1.1.0-mod：

```auto sudo docker run --privileged \ -it \ -d \ --rm \ --network host \ --name "sax-model-server" \ --env SAX_ROOT=gs://${GSBUCKET}/sax-root \ sax-model-server:v1.1.0-mod \ --sax_cell="/sax/test" \ --port=10001 \ --platform_chip=tpuv5e \ --platform_topology='8' ```

70B 型号

使用 SSH 连接到您的 TPU 并启动模型服务器：

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --worker=all \ --command=" gcloud auth configure-docker \ us-docker.pkg.dev # Pull SAX model server image sudo docker pull us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 # Run model server sudo docker run \ --privileged \ -it \ -d \ --rm \ --network host \ --name "sax-model-server" \ --env SAX_ROOT=gs://${GSBUCKET}/sax-root \ us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \ --sax_cell="/sax/test" \ --port=10001 \ --platform_chip=tpuv5e \ --platform_topology='16' "

检查日志

检查模型服务器日志，确保模型服务器已启动正确：

docker logs -f sax-model-server

如果模型服务器未启动，请参阅问题排查部分。

对于 70B 模型，请对每个 TPU 虚拟机重复以下步骤：

使用 SSH 连接到 TPU：

 gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --worker=WORKER_NUMBER

WORKER_NUMBER 是一个从 0 开始的索引，表示您要运行哪个 TPU 虚拟机连接。

查看日志：

sudo docker logs -f sax-model-server

三个 TPU 虚拟机应显示它们已连接到其他实例：

I1117 00:16:07.196594 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.3:10001 I1117 00:16:07.197484 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.87:10001 I1117 00:16:07.199437 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.13:10001

其中一个 TPU 虚拟机应包含显示模型服务器正在启动的日志：

I1115 04:01:29.479170 139974275995200 model_service_base.py:867] Started joining SAX cell /sax/test ERROR: logging before flag.Parse: I1115 04:01:31.479794 1 location.go:141] Calling Join due to address update ERROR: logging before flag.Parse: I1115 04:01:31.814721 1 location.go:155] Joined 10.182.0.44:10000

发布模型

SAX 附带一个名为 saxutil 的命令行工具，与 SAX 模型服务器交互。在本教程中，您将使用 saxutil 用于发布模型。如需查看 saxutil 命令的完整列表，请参阅 Saxml 自述文件文件。

切换到您克隆 Saxml GitHub 代码库的目录：

cd saxml

对于 70B 模型，请连接到转化服务器：

gcloud compute ssh ${CONV_SERVER_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE}

安装 Bazel：

sudo apt-get install bazel

设置别名，以便通过 Cloud Storage 存储桶运行 saxutil：

alias saxutil='bazel run saxml/bin:saxutil -- --sax_root=gs://${GSBUCKET}/sax-root'

使用 saxutil 发布模型。在 TPU 上大约需要 10 分钟 v5litepod-8。

 saxutil --sax_root=gs://${GSBUCKET}/sax-root publish '/sax/test/MODEL' \ saxml.server.pax.lm.params.lm_cloud.PARAMETERS \ gs://${GSBUCKET}/sax_models/llama2/SAX_LLAMA2_DIR/checkpoint_00000000/ \ 1

执行以下变量替换操作：

 | 模型大小 | 值 | | --- | --- | | 70 亿 | `MODEL`：llama27b `PARAMETERS`：saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e `SAX_LLAMA2_DIR`：saxml_llama27b | | 130 亿 | `MODEL`：llama213b `PARAMETERS`：saxml.server.pax.lm.params.lm_cloud.LLaMA13BFP16TPUv5e `SAX_LLAMA2_DIR`：saxml_llama213b | | 700 亿 | `MODEL`：llama270b `PARAMETERS`：saxml.server.pax.lm.params.lm_cloud.LLaMA70BFP16TPUv5e `SAX_LLAMA2_DIR`：saxml_llama270b |

测试部署

如需检查部署是否成功，请使用 saxutil ls 命令：

saxutil ls /sax/test/MODEL

成功的部署的副本数量应该大于零，与以下代码类似：

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root is /sax/test/1lama27b +----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+ | MODEL | MODEL PATH | CHECKPOINT PATH | # OF REPLICAS | (SELECTED) REPLICAADDRESS | +----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+ | llama27b | saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e | gs://${MODEL_BUCKET}/sax_models/llama2/7b/pax_7B/checkpoint_00000000/ | 1 | 10.182.0.28:10001 | +----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+

模型服务器的 Docker 日志将类似于以下内容：

I1114 17:31:03.586631 140003787142720 model_service_base.py:532] Successfully loaded model for key: /sax/test/llama27b INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root is /sax/test/1lama27b

问题排查

如果部署失败，请检查模型服务器日志：

sudo docker logs -f sax-model-server

为了成功部署，您应该会看到以下输出：

Successfully loaded model for key: /sax/test/llama27b

如果日志未显示模型已部署，请检查模型配置和模型检查点的路径。

生成回答

您可以使用 saxutil 工具生成对提示的回复。

生成问题回复：

saxutil lm.generate -extra=“temperature:0.2” /sax/test/MODEL “Q: Who is Harry Potter’s mother? A:”

输出应类似如下所示：

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root' lm.generate /sax/test/llama27b 'Q: Who is Harry Potter's mother? A: ` +-------------------------------+------------+ | GENERATE | SCORE | +-------------------------------+------------+ | 1. Harry Potter's mother is | -20.214787 | | Lily Evans. 2. Harry Potter's | | | mother is Petunia Evans | | | (Dursley). | | +-------------------------------+------------+

从客户端与模型交互

SAX 代码库包含可用于与 SAX 单元交互的客户端。客户端支持 C++、Python 和 Go。以下示例展示了如何构建 Python 客户端。

构建 Python 客户端：

bazel build saxml/client/python:sax.cc --compile_one_dependency

将客户端添加到 PYTHONPATH。此示例假定您的 saxml 您的主目录：

export PYTHONPATH=${PYTHONPATH}:$HOME/saxml/bazel-bin/saxml/client/python/

从 Python shell 与 SAX 交互：

$ python3 Python 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0] on linux Type "help", "copyright", "credits" or "license" for more information. >>> import sax >>>

从 HTTP 端点与模型交互

如需从 HTTP 端点与模型进行交互，请创建一个 HTTP 客户端：

创建 Compute Engine 虚拟机：

 export PROJECT_ID=PROJECT_ID export ZONE=ZONE export HTTP_SERVER_NAME=HTTP_SERVER_NAME export SERVICE_ACCOUNT=SERVICE_ACCOUNT export MACHINE_TYPE=e2-standard-8 gcloud compute instances create $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE \ --machine-type=$MACHINE_TYPE \ --network-interface=network-tier=PREMIUM,stack-type=IPV4_ONLY,subnet=default \ --maintenance-policy=MIGRATE --provisioning-model=STANDARD \ --service-account=$SERVICE_ACCOUNT \ --scopes=https://www.googleapis.com/auth/cloud-platform \ --tags=http-server,https-server \ --create-disk=auto-delete=yes,boot=yes,device-name=$HTTP_SERVER_NAME,image=projects/ml-images/global/images/c0-deeplearning-common-cpu-v20230925-debian-10,mode=rw,size=500,type=projects/$PROJECT_ID/zones/$ZONE/diskTypes/pd-balanced \ --no-shielded-secure-boot \ --shielded-vtpm \ --shielded-integrity-monitoring \ --labels=goog-ec-src=vm_add-gcloud \ --reservation-affinity=any

使用 SSH 连接到 Compute Engine 虚拟机：

gcloud compute ssh $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE

克隆 AI on GKE GitHub 代码库：

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

切换到 HTTP 服务器目录：

cd ai-on-gke/tools/saxml-on-gke/httpserver

构建 Docker 文件：

docker build -f Dockerfile -t sax-http .

运行 HTTP 服务器：

docker run -e SAX_ROOT=gs://${GSBUCKET}/sax-root -p 8888:8888 -it sax-http

从本地机器或能够访问端口的其他服务器测试端点 8888：

导出服务器的 IP 地址和端口的环境变量：

 export LB_IP=HTTP_SERVER_EXTERNAL_IP export PORT=8888

设置包含模型和查询的 JSON 载荷：

 json_payload=$(cat << EOF { "model": "/sax/test/MODEL", "query": "Example query" } EOF )

发送请求：

curl --request POST --header "Content-type: application/json" -s $LB_IP:$PORT/generate --data "$json_payload"

文章来源：google cloud