Pengantar Inferensi Cloud TPU v5e

Ringkasan dan manfaat

Cloud TPU v5e adalah akselerator AI yang dikembangkan Google dan dioptimalkan untuk pelatihan, fine-tuning, dan inferensi berbasis transformasi berbasis teks ke gambar dan CNN (inferensi). Slice TPU v5e dapat berisi hingga 256 chip.

Penyajian mengacu pada proses men-deploy model machine learning terlatih ke lingkungan produksi, yang dapat digunakan untuk inferensi. SLO Latensi adalah prioritas untuk penayangan.

Dokumen ini membahas cara menayangkan model di TPU host tunggal. Slice TPU dengan 8 chip atau kurang memiliki satu VM TPU atau host dan disebut TPU host tunggal.

Mulai

Anda memerlukan kuota untuk TPU v5e. TPU on-demand memerlukan kuota tpu-v5s-litepod-serving. TPU yang dipesan memerlukan kuota tpu-v5s-litepod-serving-reserved. Untuk mengetahui informasi selengkapnya, hubungi Cloud Sales.

Anda memerlukan akun dan project Google Cloud untuk menggunakan Cloud TPU. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan lingkungan Cloud TPU

Anda menyediakan TPU v5e menggunakan Resource yang diantrekan. Untuk mengetahui informasi selengkapnya tentang konfigurasi v5e yang tersedia untuk penayangan, lihat Jenis Cloud TPU v5e untuk inferensi.

Inferensi dan inferensi model Cloud TPU

Cara Anda menyajikan model untuk inferensi bergantung pada framework ML yang digunakan untuk menulis model Anda. TPU v5e mendukung model inferensi yang ditulis dalam JAX, TensorFlow, dan PyTorch.

Inferensi dan inferensi model JAX

Untuk menayangkan model di VM TPU, Anda harus:

Membuat serialisasi model dalam format SavedModel TensorFlow
Menggunakan Pengonversi Inferensi untuk menyiapkan model tersimpan untuk penayangan
Menggunakan TensorFlow Serve untuk menyalurkan model

FormatSavedModel

SavedModel berisi program TensorFlow lengkap, termasuk parameter dan komputasi terlatih. Metode ini tidak memerlukan kode bangunan model asli untuk dijalankan.

Jika model ditulis dalam JAX, Anda harus menggunakan jax2tf untuk membuat serialisasi model dalam formatSavedModel.

Konverter Inferensi

Cloud TPU Inference Converter menyiapkan dan mengoptimalkan model yang diekspor dalam format SavedModel untuk inferensi TPU. Anda dapat menjalankan konverter inferensi di shell lokal atau VM TPU Anda. Sebaiknya gunakan shell VM TPU karena memiliki semua alat command line yang diperlukan untuk menjalankan konverter. Untuk informasi selengkapnya tentang Pengonversi Inferensi, lihat Panduan Pengguna Pengubah Inferensi.

Persyaratan Konverter Inferensi

Model Anda harus diekspor dari TensorFlow atau JAX dalam format SavedModel.
Anda harus menentukan alias fungsi untuk fungsi TPU. Untuk mengetahui informasi selengkapnya, lihat Panduan Pengguna Pengonversi Inferensi. Contoh dalam panduan ini menggunakan tpu_func sebagai alias fungsi TPU.
Pastikan CPU mesin Anda mendukung petunjuk Advanced Vector eXtensions (AVX), karena library TensorFlow (dependensi Cloud TPU Inference Converter) dikompilasi untuk menggunakan instruksi AVX. Sebagian besar CPU memiliki dukungan AVX.

Catatan: Anda dapat menjalankan lscpu | grep avx untuk memeriksa apakah kumpulan petunjuk AVX didukung.

Inferensi dan inferensi model JAX

Bagian ini menjelaskan cara menayangkan model JAX menggunakan jax2tf dan TensorFlow Serve.

Menggunakan jax2tf untuk membuat serialisasi model ke dalam formatSavedModel
Gunakan Pengonversi Inferensi untuk menyiapkan model tersimpan untuk penayangan
Menggunakan TensorFlow Serve untuk menyalurkan model

Menggunakan `jax2tf` untuk membuat serialisasi model JAX ke formatSavedModel

Fungsi Python berikut menunjukkan cara menggunakan jax2tf dalam kode model Anda:

# Inference function
def model_jax(params, inputs):
  return params[0] + params[1] * inputs

# Wrap the parameter constants as tf.Variables; this will signal to the model
# saving code to save those constants as variables, separate from the
# computation graph.
params_vars = tf.nest.map_structure(tf.Variable, params)

# Build the prediction function by closing over the `params_vars`. If you
# instead were to close over `params` your SavedModel would have no variables
# and the parameters will be included in the function graph.
prediction_tf = lambda inputs: jax2tf.convert(model_jax)(params_vars, inputs)

my_model = tf.Module()
# Tell the model saver what the variables are.
my_model._variables = tf.nest.flatten(params_vars)
my_model.f = tf.function(prediction_tf, jit_compile=True, autograph=False)
tf.saved_model.save(my_model)

Untuk mengetahui informasi selengkapnya tentang jax2tf, lihat interoperasi JAX dan Cloud TPU.

Menggunakan Pengonversi Inferensi untuk menyiapkan model tersimpan untuk penayangan

Petunjuk untuk menggunakan Pengonversi Inferensi dijelaskan dalam Panduan pengonversi inferensi.

Menggunakan TensorFlow Serve

Petunjuk untuk menggunakan TensorFlow Serve dijelaskan dalam layanan TensorFlow.

Contoh penayangan model JAX

Prasyarat

Siapkan kredensial Docker dan ambil image Docker Inference Converter dan Cloud TPU:

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
   us-docker.pkg.dev
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Hubungkan ke VM TPU Anda dengan SSH, lalu instal kode demo inferensi:

gsutil -m cp -r \
"gs://cloud-tpu-inference-public/demo" \
.

Instal dependensi demo JAX:

pip install -r ./demo/jax/requirements.txt

Menyajikan model JAX BERT untuk inferensi

Anda dapat mendownload model BERT terlatih dari Wajah Memeluk.

Ekspor model tersimpan TensorFlow yang kompatibel dengan TPU dari model Flax BERT:
```
cd demo/jax/bert
python3 export_bert_model.py
```

Mulai penampung server model Cloud TPU:

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/jax/bert_tpu,target=/models/bert \
  -e MODEL_NAME=bert \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Sekitar 30 detik setelah penampung dimulai, periksa log penampung server model dan pastikan server gRPC dan HTTP aktif:

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

Jika Anda melihat entri log yang diakhiri dengan informasi berikut, berarti server siap melayani permintaan.

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

Kirim permintaan inferensi ke server model.

python3 bert_request.py

Outputnya akan mirip dengan berikut ini:

For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."

Jalankan pembersihan.

Pastikan untuk membersihkan container Docker sebelum menjalankan demo lainnya.

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}

Bersihkan artefak model:

sudo rm -rf /tmp/jax/

Menayangkan Difusi Stabil JAX untuk inferensi

Anda dapat mendownload model Difusi Stabil terlatih dari Wajah Memeluk.

Download model Diffusion Stabil dalam format model tersimpan TF2 yang kompatibel dengan TPU:

Catatan: Proses download memerlukan waktu sekitar 3 menit.
```
cd demo/jax/stable_diffusion
python3 export_stable_diffusion_model.py
```

Mulai penampung server model Cloud TPU untuk model:

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/jax/stable_diffusion_tpu,target=/models/stable_diffusion \
  -e MODEL_NAME=stable_diffusion \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Setelah sekitar dua menit, periksa log penampung server model untuk memastikan server gRPC dan HTTP berjalan:

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

Jika Anda melihat log yang diakhiri dengan informasi berikut, berarti server siap melayani permintaan.

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

Kirim permintaan ke server model.
```
python3 stable_diffusion_request.py
```
Skrip ini mengirimkan "Painting of a squirrel boarding in New York" sebagai prompt. Gambar output akan disimpan sebagai stable_diffusion_images.jpg di direktori Anda saat ini.

Jalankan pembersihan.

Pastikan untuk membersihkan container Docker sebelum menjalankan demo lainnya.

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}

Membersihkan artefak model

sudo rm -rf /tmp/jax/

TensorFlow Serving

Petunjuk berikut menunjukkan cara menyalurkan model TensorFlow di VM TPU.

Alur kerja inferensi TensorFlow

Download image Docker TensorFlow Serve untuk VM TPU Anda.

Menetapkan contoh variabel lingkungan

export YOUR_LOCAL_MODEL_PATH=model-path
export MODEL_NAME=model-name
# Note: this image name may change later.
export IMAGE_NAME=us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Mendownload image Docker

docker pull ${IMAGE_NAME}

Siapkan kredensial Docker dan ambil Inference Converter serta image Docker TensorFlow Serve.

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
   us-docker.pkg.dev
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Download kode demo:

gsutil -m cp -r \
"gs://cloud-tpu-inference-public/demo" \
.

Instal dependensi demo TensorFlow:

pip install -r ./demo/tf/requirements.txt

Tayangkan model TensorFlow Anda menggunakan image TensorFlow Serve Docker di VM TPU.

# PORT 8500 is for gRPC model server and 8501 is for HTTP/REST model server.
docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=${YOUR_LOCAL_MODEL_PATH},target=/models/${MODEL_NAME} \
  -e MODEL_NAME=${MODEL_NAME} \
  ${IMAGE_NAME}

Gunakan Serve Client API untuk membuat kueri model Anda.
- REST Client API
- gRPC Client API

Jalankan demo TensorFlow ResNet-50 Serve

Ekspor model tersimpan TF2 yang kompatibel dengan TPU dari model Keras ResNet-50.
```
cd demo/tf/resnet-50
python3 export_resnet_model.py
```

Luncurkan penampung server model TensorFlow untuk model tersebut.

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/tf/resnet_tpu,target=/models/resnet \
  -e MODEL_NAME=resnet \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Periksa log penampung server model dan pastikan gRPC dan Server HTTP sudah aktif:

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

Jika Anda melihat log yang diakhiri dengan informasi berikut, berarti server siap melayani permintaan. Prosesnya perlu waktu sekitar 30 detik.

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

Kirim permintaan ke server model.

Gambar permintaan adalah pisang dari https://i.imgur.com/j9xCCzn.jpeg .

python3 resnet_request.py

Outputnya akan mirip dengan berikut ini:

Predict result: [[('n07753592', 'banana', 0.94921875), ('n03532672', 'hook', 0.022338867), ('n07749582', 'lemon', 0.005126953)]]

Jalankan pembersihan.

Pastikan untuk membersihkan container Docker sebelum menjalankan demo lainnya.

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}

Bersihkan artefak model:

sudo rm -rf /tmp/tf/

Inferensi dan inferensi model PyTorch

Untuk model yang ditulis dengan PyTorch, alur kerjanya adalah:

Menulis pengendali model Python untuk memuat dan melakukan inferensi menggunakan TorchDynamo dan PyTorch/XLA
Gunakan TorchModelArchiver untuk membuat arsip model
Menggunakan TorchServe untuk menyalurkan model

TorchDynamo dan PyTorch/XLA

TorchDynamo (Dynamo) adalah compiler JIT level Python yang dirancang untuk membuat program PyTorch lebih cepat. Menyediakan API yang bersih untuk backend compiler yang akan dikaitkan. Kode ini secara dinamis mengubah bytecode Python sebelum dieksekusi. Dalam rilis PyTorch/XLA 2.0, terdapat backend eksperimental untuk inferensi dan pelatihan menggunakan Dynamo.

Dynamo menyediakan grafik Torch FX (FX) saat mengenali pola model dan PyTorch/XLA menggunakan pendekatan tensor lambat untuk mengompilasi grafik FX dan menampilkan fungsi yang dikompilasi. Untuk mengetahui informasi selengkapnya tentang Dynamo, lihat:

Postingan Diskusi Developer Pytorch
Dokumentasi TorchDynamo
PyTorch 2.0 & XLA untuk detail selengkapnya

Berikut adalah contoh kode kecil untuk menjalankan inferensi padatnet161 dengan torch.compile.

import torch
import torchvision
import torch_xla.core.xla_model as xm

def eval_model(loader):
  device = xm.xla_device()
  xla_densenet161 = torchvision.models.densenet161().to(device)
  xla_densenet161.eval()
  dynamo_densenet161 = torch.compile(
      xla_densenet161, backend='torchxla_trace_once')
  for data, _ in loader:
    output = dynamo_densenet161(data)

TorchServe

Anda dapat menggunakan image Docker torchserve-tpu yang disediakan untuk menyalurkan model pytorch yang diarsipkan di VM TPU.

Siapkan autentikasi untuk Docker:

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
    us-docker.pkg.dev

Tarik image Docker Cloud TPU TorchServe ke VM TPU Anda:

CLOUD_TPU_TORCHSERVE_IMAGE_URL=us-docker.pkg.dev/cloud-tpu-images/inference/torchserve-tpu:v0.9.0-2.1
docker pull ${CLOUD_TPU_TORCHSERVE_IMAGE_URL}

Kumpulkan artefak model

Untuk memulai, Anda harus menyediakan pengendali model, yang menginstruksikan pekerja server model TorchServe untuk memuat model, memproses data input, dan menjalankan inferensi. Anda dapat menggunakan pengendali inferensi default TorchServe (sumber), atau mengembangkan pengendali model kustom Anda sendiri dengan mengikuti base_handler.py. Anda mungkin juga perlu menyediakan model yang dilatih dan file definisi model.

Dalam contoh Densenet 161 berikut, kami menggunakan artefak model dan pengendali pengklasifikasi gambar default yang disediakan oleh TorchServe:

Konfigurasikan beberapa variabel lingkungan:

CWD="$(pwd)"

WORKDIR="${CWD}/densenet_161"

mkdir -p ${WORKDIR}/model-store
mkdir -p ${WORKDIR}/logs

Download dan salin artefak model dari contoh pengklasifikasi gambar TorchServe:

git clone https://github.com/pytorch/serve.git

cp ${CWD}/serve/examples/image_classifier/densenet_161/model.py ${WORKDIR}
cp ${CWD}/serve/examples/image_classifier/index_to_name.json ${WORKDIR}

Download bobot model:

wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth

mv densenet161-8d451a50.pth ${WORKDIR}

Buat file konfigurasi model TorchServe untuk menggunakan backend Dynamo:

echo 'pt2: "torchxla_trace_once"' >> ${WORKDIR}/model_config.yaml

Anda akan melihat file dan direktori berikut:

>> ls ${WORKDIR}
model_config.yaml
index_to_name.json
logs
model.py
densenet161-8d451a50.pth
model-store

Membuat file arsip model

Untuk menayangkan model PyTorch dengan Cloud TPU TorchServe, Anda harus mengemas pengendali model dan semua artefak model ke dalam file arsip model (*.mar) menggunakan Torch Model Archiver.

Buat file arsip model dengan torch-model-archiver:

MODEL_NAME=Densenet161

docker run \
    --privileged  \
    --shm-size 16G \
    --name torch-model-archiver \
    -it \
    -d \
    --rm \
    --mount type=bind,source=${WORKDIR},target=/home/model-server/ \
    ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
    torch-model-archiver \
        --model-name ${MODEL_NAME} \
        --version 1.0 \
        --model-file model.py \
        --serialized-file densenet161-8d451a50.pth \
        --handler image_classifier \
        --export-path model-store \
        --extra-files index_to_name.json \
        --config-file model_config.yaml

Anda akan melihat file arsip model yang dibuat di direktori model-store:

>> ls ${WORKDIR}/model-store
Densenet161.mar

Menayangkan permintaan inferensi

Setelah memiliki file arsip model, Anda dapat memulai server model TorchServe dan menayangkan permintaan inferensi.

Mulai server model TorchServe:

docker run \
    --privileged  \
    --shm-size 16G \
    --name torchserve-tpu \
    -it \
    -d \
    --rm \
    -p 7070:7070 \
    -p 7071:7071 \
    -p 8080:8080 \
    -p 8081:8081 \
    -p 8082:8082 \
    -p 9001:9001 \
    -p 9012:9012 \
    --mount type=bind,source=${WORKDIR}/model-store,target=/home/model-server/model-store \
    --mount type=bind,source=${WORKDIR}/logs,target=/home/model-server/logs \
    ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
    torchserve \
        --start \
        --ncs \
        --models ${MODEL_NAME}.mar \
        --ts-config /home/model-server/config.properties

Kondisi server model kueri:

curl http://localhost:8080/ping

Jika server model sudah aktif dan berjalan, Anda akan melihat:

{
  "status": "Healthy"
}

Untuk membuat kueri versi default dari model yang terdaftar saat ini, gunakan:

curl http://localhost:8081/models

Anda akan melihat model yang terdaftar:

{
  "models": [
    {
      "modelName": "Densenet161",
      "modelUrl": "Densenet161.mar"
    }
  ]
}

Agar dapat mendownload gambar untuk digunakan inferensi:

curl -O https://raw.githubusercontent.com/pytorch/serve/master/docs/images/kitten_small.jpg

mv kitten_small.jpg ${WORKDIR}

Untuk mengirim permintaan inferensi ke server model, gunakan:

curl http://localhost:8080/predictions/${MODEL_NAME} -T ${WORKDIR}/kitten_small.jpg

Anda akan melihat respons seperti berikut:

{
  "tabby": 0.47878125309944153,
  "lynx": 0.20393909513950348,
  "tiger_cat": 0.16572578251361847,
  "tiger": 0.061157409101724625,
  "Egyptian_cat": 0.04997897148132324
}

Log server model

Gunakan perintah berikut untuk mengakses log:

ls ${WORKDIR}/logs/
cat ${WORKDIR}/logs/model_log.log

Anda akan melihat pesan berikut di log Anda:

"Compiled model with backend torchxla\_trace\_once"

Pembersihan

Hentikan container Docker:

rm -rf serve
rm -rf ${WORKDIR}

docker stop torch-model-archiver
docker stop torchserve-tpu

Pembuatan profil

Setelah menyiapkan inferensi, Anda dapat menggunakan profiler untuk menganalisis performa dan penggunaan TPU. Untuk informasi selengkapnya tentang pembuatan profil, lihat:

Pengantar Inferensi Cloud TPU v5e

Ringkasan dan manfaat

Mulai

Inferensi dan inferensi model Cloud TPU

Inferensi dan inferensi model JAX

FormatSavedModel

Konverter Inferensi

Persyaratan Konverter Inferensi

Inferensi dan inferensi model JAX

Menggunakan jax2tf untuk membuat serialisasi model JAX ke formatSavedModel

Menggunakan Pengonversi Inferensi untuk menyiapkan model tersimpan untuk penayangan

Menggunakan TensorFlow Serve

Contoh penayangan model JAX

Prasyarat

Menyajikan model JAX BERT untuk inferensi

Menayangkan Difusi Stabil JAX untuk inferensi

TensorFlow Serving

Alur kerja inferensi TensorFlow

Jalankan demo TensorFlow ResNet-50 Serve

Inferensi dan inferensi model PyTorch

TorchDynamo dan PyTorch/XLA

TorchServe

Kumpulkan artefak model

Membuat file arsip model

Menayangkan permintaan inferensi

Pembersihan

Pembuatan profil

Menggunakan `jax2tf` untuk membuat serialisasi model JAX ke formatSavedModel