Halaman ini diterjemahkan oleh Cloud Translation API.

Inferensi MaxText JetStream pada VM TPU v6e

Tutorial ini menunjukkan cara menggunakan JetStream untuk menayangkan model MaxText di TPU v6e. JetStream adalah mesin yang dioptimalkan untuk throughput dan memori untuk inferensi model bahasa besar (LLM) di perangkat XLA (TPU). Dalam tutorial ini, Anda akan menjalankan benchmark inferensi untuk model Llama2-7B.

Sebelum memulai

Bersiap untuk menyediakan TPU v6e dengan 4 chip:

Ikuti panduan Menyiapkan lingkungan Cloud TPU untuk menyiapkan Google Cloud project, mengonfigurasi Google Cloud CLI, mengaktifkan Cloud TPU API, dan memastikan Anda memiliki akses untuk menggunakan Cloud TPU.
Lakukan autentikasi dengan Google Cloud dan konfigurasikan project dan zona default untuk Google Cloud CLI.
```
gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE
```

Kapasitas aman

Jika Anda siap untuk mengamankan kapasitas TPU, lihat Kuota Cloud TPU untuk mengetahui informasi selengkapnya tentang kuota Cloud TPU. Jika ada pertanyaan tambahan tentang cara mengamankan kapasitas, hubungi tim akun atau tim penjualan Cloud TPU Anda.

Menyediakan lingkungan Cloud TPU

Anda dapat menyediakan VM TPU dengan GKE, dengan GKE dan XPK, atau sebagai resource dalam antrean.

Prasyarat

Pastikan project Anda memiliki kuota TPUS_PER_TPU_FAMILY yang cukup, yang menentukan jumlah maksimum chip yang dapat Anda akses dalam projectGoogle Cloud .
Pastikan project Anda memiliki cukup kuota TPU untuk:
- Kuota VM TPU
- Kuota alamat IP
- Kuota Hyperdisk Balanced
Izin project pengguna
- Jika Anda menggunakan GKE dengan XPK, lihat Izin Konsol Cloud di akun pengguna atau layanan untuk mengetahui izin yang diperlukan untuk menjalankan XPK.

Membuat variabel lingkungan

Di Cloud Shell, buat variabel lingkungan berikut:

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

Deskripsi flag perintah

Variabel	Deskripsi
`PROJECT_ID`	Google Cloud nama project. Gunakan project yang ada atau buat project baru.
`TPU_NAME`	Nama TPU.
`ZONE`	Lihat dokumen Region dan zona TPU untuk zona yang didukung.
`ACCELERATOR_TYPE`	Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
`RUNTIME_VERSION`	Versi software Cloud TPU.
`SERVICE_ACCOUNT`	Alamat email untuk akun layanan Anda . Anda dapat menemukannya dengan membuka halaman Akun Layanan di Google Cloud konsol. Contoh: `tpu-service-account@PROJECT_ID.iam.gserviceaccount.com`
`QUEUED_RESOURCE_ID`	ID teks yang ditetapkan pengguna untuk permintaan resource yang diantrekan.

Menyediakan TPU v6e

Gunakan perintah berikut untuk menyediakan 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}

Gunakan perintah list atau describe untuk membuat kueri status resource yang diantrekan.

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

Untuk mengetahui informasi selengkapnya tentang status permintaan resource yang diantrekan, lihat Mengelola resource yang diantrekan.

Menghubungkan ke TPU menggunakan SSH

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}

Setelah terhubung ke TPU, Anda dapat menjalankan benchmark inferensi.

Menyiapkan lingkungan VM TPU

Buat direktori untuk menjalankan benchmark inferensi:

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

Siapkan lingkungan virtual Python:

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

Instal Penyimpanan File Besar (LFS) Git (untuk data OpenOrca):
```
sudo apt-get install git-lfs
git lfs install
```

Clone dan instal 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

Siapkan 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

Minta Akses ke Model Llama untuk mendapatkan kunci download dari Meta untuk model Llama 2.

Buat clone repositori Llama:

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

Jalankan bash download.sh. Jika diminta, berikan kunci download Anda. Skrip ini membuat direktori llama-2-7b di dalam direktori llama Anda.
```
bash download.sh
```

Buat bucket penyimpanan:

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}

Melakukan konversi checkpoint

Melakukan konversi ke checkpoint yang dipindai:

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}

Mengonversi ke checkpoint yang tidak dipindai:

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

Melakukan inferensi

Jalankan pengujian validasi:

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

Jalankan server di terminal saat ini:

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}

Buka jendela terminal baru, hubungkan ke TPU, dan beralih ke lingkungan virtual yang sama dengan yang Anda gunakan di jendela terminal pertama:
```
source venv/bin/activate
```

Jalankan perintah berikut untuk menjalankan benchmark 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

Hasil

Output berikut dihasilkan saat menjalankan benchmark menggunakan v6e-8. Hasilnya akan bervariasi berdasarkan hardware, software, model, dan jaringan.

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

Pembersihan

Putuskan sambungan dari TPU:
```
$ (vm) exit
```

Hapus TPU:

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

Hapus bucket dan isinya:

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}