本頁面由 Cloud Translation API 翻譯而成。

在 v6e TPU VM 上執行 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 VM。

必要條件

請確認專案具備足夠的 TPUS_PER_TPU_FAMILY 配額，該配額會指定您在Google Cloud 專案中可存取的晶片數量上限。
確認專案具備足夠的 TPU 配額，可供下列項目使用：
- TPU VM 配額
- IP 位址配額
- Hyperdisk Balanced 配額
使用者專案權限
- 如果您要搭配 XPK 使用 GKE，請參閱「使用者或服務帳戶的 Cloud Console 權限」，瞭解執行 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 專案 ID。使用現有專案或建立新專案。
`TPU_NAME`	TPU 的名稱。
`ZONE`	建立 TPU VM 的區域。如要進一步瞭解支援的區域，請參閱「TPU 地區和區域」一文。
`ACCELERATOR_TYPE`	加速器類型會指定您要建立的 Cloud TPU 版本和大小。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。
`RUNTIME_VERSION`	Cloud TPU 軟體版本。
`SERVICE_ACCOUNT`	服務帳戶的電子郵件地址。您可以前往 Google Cloud 控制台的「Service Accounts」頁面查看。例如： `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 VM 環境

建立用於執行推論基準測試的目錄：

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 Large File Storage (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}