创建 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
推荐阅读
- 在 v5e Cloud TPU 虚拟机上进行 JetStream PyTorch 推断
- 在 v5e Cloud TPU 虚拟机上进行 JetStream MaxText 推理
- 使用 Pax 在单主机 TPU 上训练
更多芯擎AI开发板干货请关注芯擎AI开发板专栏。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。