v6e TPU VM での JetStream MaxText 推論

このチュートリアルでは、JetStream を使用して TPU v6e で MaxText モデルを提供する方法について説明します。JetStream は、XLA デバイス（TPU）での大規模言語モデル（LLM）推論用にスループットとメモリが最適化されたエンジンです。このチュートリアルでは、Llama2-7B モデルの推論ベンチマークを実行します。

始める前に

4 個のチップを含む TPU v6e をプロビジョニングする準備を行います。

  1. Cloud TPU 環境を設定するガイドに沿って、プロジェクトを設定し、Google Cloud CLI を構成し、Cloud TPU API を有効にして、Cloud TPU を使用するためのアクセス権があることを確認します。 Google Cloud

  2. 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 環境をプロビジョニングする

TPU VM は、GKE、GKE と XPK、またはキューに入れられたリソースとしてプロビジョニングできます。

前提条件

  • プロジェクトに十分な TPUS_PER_TPU_FAMILY 割り当てがあることを確認します。これは、Google Cloud プロジェクト内でアクセスできるチップの最大数を指定します。
  • プロジェクトに次の TPU 割り当てがあることを確認します。
    • TPU VM の割り当て
    • IP アドレスの割り当て
    • Hyperdisk Balanced の割り当て
  • ユーザー プロジェクトの権限

環境変数を作成する

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

例: 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 環境を設定する

  1. 推論ベンチマークを実行するディレクトリを作成します。

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

  2. Python 仮想環境を設定します。

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

  3. Git Large File Storage（LFS）をインストールします（OpenOrca データ用）。

    sudo apt-get install git-lfs
git lfs install

  4. 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

  5. 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

  6. Llama モデルへのアクセス権をリクエストして、Llama 2 モデルのダウンロード キーを Meta から取得します。

  7. Llama リポジトリのクローンを作成します。

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

  8. bash download.sh を実行します。プロンプトが表示されたら、ダウンロード キーを入力します。このスクリプトにより、llama ディレクトリ内に llama-2-7b ディレクトリが作成されます。

    bash download.sh

  9. ストレージ バケットを作成します。

    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}

チェックポイントの変換を行う

  1. スキャンされたチェックポイントへの変換を実行します。

    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}

  2. 未スキャンのチェックポイントに変換します。

    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

推論を実行する

  1. 検証テストを実行します。

    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

  2. 現在のターミナルでサーバーを実行します。

    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}

  3. 新しいターミナル ウィンドウを開き、TPU に接続して、最初のターミナル ウィンドウで使用した仮想環境に切り替えます。

    source venv/bin/activate

  4. 次のコマンドを実行して、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

クリーンアップ

  1. TPU との接続を解除します。

    $ (vm) exit

  2. TPU を削除します。

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

  3. バケットとその内容を削除します。

    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}