此页面由 Cloud Translation API 翻译。

在 v6e TPU 虚拟机上进行 JetStream MaxText 推理

本教程介绍了如何使用 JetStream 在 TPU v6e 上提供 MaxText 模型。JetStream 是一款针对 XLA 设备 (TPU) 上的大语言模型 (LLM) 推理进行了吞吐量和内存优化的引擎。在本教程中，您将针对 Llama2-7B 模型运行推理基准测试。

准备工作

准备预配具有 4 个芯片的 TPU v6e：

按照设置 Cloud TPU 环境指南设置 Google Cloud 项目、配置 Google Cloud CLI、启用 Cloud TPU API，并确保您有权使用 Cloud TPU。

使用 Google Cloud 进行身份验证，并为 Google Cloud CLI 配置默认项目和区域。

gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE

保障容量

当您准备好预订 TPU 容量时，请参阅 Cloud TPU 配额，详细了解 Cloud TPU 配额。如果您对如何确保容量还有其他疑问，请与您的 Cloud TPU 销售团队或客户支持团队联系。

预配 Cloud TPU 环境

您可以使用 GKE、GKE 和 XPK 预配 TPU 虚拟机，也可以将其作为队列化资源预配。

前提条件

验证您的项目是否有足够的 TPUS_PER_TPU_FAMILY 配额，该配额指定您可以在Google Cloud 项目中访问的芯片数量上限。
验证您的项目是否有足够的 TPU 配额：
- TPU 虚拟机配额
- IP 地址配额
- Hyperdisk Balanced 配额
用户项目权限
- 如果您将 GKE 与 XPK 搭配使用，请参阅用户账号或服务账号的 Cloud 控制台权限，了解运行 XPK 所需的权限。

创建环境变量

在 Cloud Shell 中，创建以下环境变量：

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east5-b
export ACCELERATOR_TYPE=v6e-4
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

命令标志说明

变量	说明
`PROJECT_ID`	Google Cloud 项目名称。使用现有项目或创建新项目。
`TPU_NAME`	TPU 的名称。
`ZONE`	如需了解支持的区域，请参阅 TPU 区域和可用区文档。
`ACCELERATOR_TYPE`	加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。
`RUNTIME_VERSION`	Cloud TPU 软件版本。
`SERVICE_ACCOUNT`	您的服务账号的电子邮件地址。您可以前往 Google Cloud 控制台中的“服务账号”页面找到该 ID。例如：`tpu-service-account@PROJECT_ID.iam.gserviceaccount.com`
`QUEUED_RESOURCE_ID`	已加入队列的资源请求的用户分配的文本 ID。

预配 TPU v6e

使用以下命令预配 TPU v6e：

gcloud alpha 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}

使用 list 或 describe 命令查询队列中资源的状态。

gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
    --project ${PROJECT_ID} --zone ${ZONE}

如需详细了解已排队的资源请求状态，请参阅管理已排队的资源。

使用 SSH 连接到 TPU

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}

连接到 TPU 后，您就可以运行推理基准测试。

设置 TPU 虚拟机环境

创建一个用于运行推理基准测试的目录：

export MAIN_DIR=your-main-directory
mkdir -p ${MAIN_DIR}

设置 Python 虚拟环境：

cd ${MAIN_DIR}
sudo apt update
sudo apt install python3.10 python3.10-venv
python3.10 -m venv venv
source venv/bin/activate

安装 Git 大型文件存储 (LFS)（适用于 OpenOrca 数据）：
```
sudo apt-get install git-lfs
git lfs install
```

克隆并安装 JetStream：

cd $MAIN_DIR
git clone https://github.com/google/JetStream.git
cd JetStream
git checkout main
pip install -e .
cd benchmarks
pip install -r requirements.in

设置 MaxText：

cd $MAIN_DIR
git clone https://github.com/google/maxtext.git
cd maxtext
git checkout main
bash setup.sh
pip install torch --index-url https://download.pytorch.org/whl/cpu

申请访问 Llama 模型，以从 Meta 获取 Llama 2 模型的下载密钥。

克隆 Llama 代码库：

cd $MAIN_DIR
git clone https://github.com/meta-llama/llama
cd llama

运行 bash download.sh。出现提示时，提供您的下载密钥。此脚本会在 llama 目录中创建一个 llama-2-7b 目录。
```
bash download.sh
```

创建存储分区：

export CHKPT_BUCKET=gs://your-checkpoint-bucket
export BASE_OUTPUT_DIRECTORY=gs://your-output-dir
export CONVERTED_CHECKPOINT_PATH=gs://bucket-to-store-converted-checkpoints
export MAXTEXT_BUCKET_UNSCANNED=gs://bucket-to-store-unscanned-data
gcloud storage buckets create ${CHKPT_BUCKET}
gcloud storage buckets create ${BASE_OUTPUT_DIRECTORY}
gcloud storage buckets create ${CONVERTED_CHECKPOINT_PATH}
gcloud storage buckets create ${MAXTEXT_BUCKET_UNSCANNED}
gcloud storage cp --recursive llama-2-7b/* ${CHKPT_BUCKET}

执行检查点转换

执行转换为扫描的检查点：

cd $MAIN_DIR/maxtext
python3 -m MaxText.llama_or_mistral_ckpt \
    --base-model-path $MAIN_DIR/llama/llama-2-7b \
    --model-size llama2-7b \
    --maxtext-model-path ${CONVERTED_CHECKPOINT_PATH}

转换为未扫描的检查点：

export CONVERTED_CHECKPOINT=${CONVERTED_CHECKPOINT_PATH}/0/items
export DIRECT_PARAMETER_CHECKPOINT_RUN=direct_generate_param_only_checkpoint
python3 -m MaxText.generate_param_only_checkpoint \
    MaxText/configs/base.yml \
    base_output_directory=${MAXTEXT_BUCKET_UNSCANNED} \
    load_parameters_path=${CONVERTED_CHECKPOINT} \
    run_name=${DIRECT_PARAMETER_CHECKPOINT_RUN} \
    model_name='llama2-7b' \
    force_unroll=true

执行推理

运行验证测试：

export UNSCANNED_CKPT_PATH=${MAXTEXT_BUCKET_UNSCANNED}/${DIRECT_PARAMETER_CHECKPOINT_RUN}/checkpoints/0/items
python3 -m MaxText.decode \
    MaxText/configs/base.yml \
    load_parameters_path=${UNSCANNED_CKPT_PATH} \
    run_name=runner_decode_unscanned_${idx} \
    base_output_directory=${BASE_OUTPUT_DIRECTORY} \
    per_device_batch_size=1 \
    model_name='llama2-7b' \
    ici_autoregressive_parallelism=4 \
    max_prefill_predict_length=4 \
    max_target_length=16 \
    prompt="I love to" \
    attention=dot_product \
    scan_layers=false

在当前终端中运行服务器：

export TOKENIZER_PATH=assets/tokenizer.llama2
export LOAD_PARAMETERS_PATH=${UNSCANNED_CKPT_PATH}
export MAX_PREFILL_PREDICT_LENGTH=1024
export MAX_TARGET_LENGTH=2048
export MODEL_NAME=llama2-7b
export ICI_FSDP_PARALLELISM=1
export ICI_AUTOREGRESSIVE_PARALLELISM=1
export ICI_TENSOR_PARALLELISM=-1
export SCAN_LAYERS=false
export WEIGHT_DTYPE=bfloat16
export PER_DEVICE_BATCH_SIZE=11

cd $MAIN_DIR/maxtext
python3 -m MaxText.maxengine_server \
    MaxText/configs/base.yml \
    tokenizer_path=${TOKENIZER_PATH} \
    load_parameters_path=${LOAD_PARAMETERS_PATH} \
    max_prefill_predict_length=${MAX_PREFILL_PREDICT_LENGTH} \
    max_target_length=${MAX_TARGET_LENGTH} \
    model_name=${MODEL_NAME} \
    ici_fsdp_parallelism=${ICI_FSDP_PARALLELISM} \
    ici_autoregressive_parallelism=${ICI_AUTOREGRESSIVE_PARALLELISM} \
    ici_tensor_parallelism=${ICI_TENSOR_PARALLELISM} \
    scan_layers=${SCAN_LAYERS} \
    weight_dtype=${WEIGHT_DTYPE} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE}

打开一个新终端窗口，连接到 TPU，然后切换到您在第一个终端窗口中使用的虚拟环境：
```
source venv/bin/activate
```

运行以下命令以运行 JetStream 基准测试。

export MAIN_DIR=your-main-directory
cd $MAIN_DIR

python JetStream/benchmarks/benchmark_serving.py \
    --tokenizer $MAIN_DIR/maxtext/assets/tokenizer.llama2 \
    --warmup-mode sampled \
    --save-result \
    --save-request-outputs \
    --request-outputs-file-path outputs.json \
    --num-prompts 1000 \
    --max-output-length 1024 \
    --dataset openorca \
    --dataset-path $MAIN_DIR/JetStream/benchmarks/open_orca_gpt4_tokenized_llama.calibration_1000.pkl

结果

使用 v6e-8 运行基准测试时，生成了以下输出。具体结果因硬件、软件、型号和网络而异。

Mean output size: 929.5959798994975
Median output size: 1026.0
P99 output size: 1026.0
Successful requests: 995
Benchmark duration: 195.533269 s
Total input tokens: 217011
Total generated tokens: 924948
Request throughput: 5.09 requests/s
Input token throughput: 1109.84 tokens/s
Output token throughput: 4730.39 tokens/s
Overall token throughput: 5840.23 tokens/s
Mean ttft: 538.49 ms
Median ttft: 95.66 ms
P99 ttft: 13937.86 ms
Mean ttst: 1218.72 ms
Median ttst: 152.57 ms
P99 ttst: 14241.30 ms
Mean TPOT: 91.83 ms
Median TPOT: 16.63 ms
P99 TPOT: 363.37 ms

清理

断开与 TPU 的连接：
```
$ (vm) exit
```

删除 TPU：

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
    --project ${PROJECT_ID} \
    --zone ${ZONE} \
    --force \
    --async

删除存储分区及其内容：

export CHKPT_BUCKET=gs://your-checkpoint-bucket
export BASE_OUTPUT_DIRECTORY=gs://your-output-dir
export CONVERTED_CHECKPOINT_PATH=gs://bucket-to-store-converted-checkpoints
export MAXTEXT_BUCKET_UNSCANNED=gs://bucket-to-store-unscanned-data
gcloud storage rm -r ${CHKPT_BUCKET}
gcloud storage rm -r ${BASE_OUTPUT_DIRECTORY}
gcloud storage rm -r ${CONVERTED_CHECKPOINT_PATH}
gcloud storage rm -r ${MAXTEXT_BUCKET_UNSCANNED}