Einführung in Cloud TPU v5e-Inferenzen
Übersicht und Vorteile
Cloud TPU v5e ist ein von Google entwickelter KI-Beschleuniger, der für transformerbasiertes, Text-zu-Image- und CNN-basiertes Training, Abstimmung und Bereitstellung (Inferenz) optimiert ist. TPU v5e-Slices können bis zu 256 Chips enthalten.
Bereitstellung bezieht sich auf die Bereitstellung eines trainierten Modells für maschinelles Lernen in einer Produktionsumgebung, in der es für Inferenzen verwendet werden kann. Latenz-SLOs haben bei der Bereitstellung Priorität.
In diesem Dokument wird die Bereitstellung eines Modells auf einer TPU mit einem einzelnen Host erläutert. TPU-Slices mit 8 oder weniger Chips haben eine TPU-VM oder einen TPU-Host und werden als TPUs mit einem einzelnen Host bezeichnet.
Jetzt starten
Sie benötigen ein Kontingent für v5e-TPUs. On-Demand-TPUs erfordern ein Kontingent von tpu-v5s-litepod-serving
. Reservierte TPUs erfordern ein Kontingent von tpu-v5s-litepod-serving-reserved
. Weitere Informationen erhalten Sie vom Cloud-Vertrieb.
Sie benötigen ein Google Cloud-Konto und ein Projekt, um Cloud TPU verwenden zu können. Weitere Informationen finden Sie unter Cloud TPU-Umgebung einrichten.
v5e-TPUs werden mithilfe von Ressourcen in der Warteschlange bereitgestellt. Weitere Informationen zu verfügbaren v5e-Konfigurationen für die Bereitstellung finden Sie unter Cloud TPU v5e-Typen für die Bereitstellung.
Cloud TPU-Modellinferenz und -bereitstellung
Wie Sie ein Modell für die Inferenz bereitstellen, hängt vom ML-Framework ab, mit dem Ihr Modell geschrieben wurde. TPU v5e unterstützt die Bereitstellung von Modellen, die in JAX, TensorFlow und PyTorch geschrieben wurden.
Inferenz und Bereitstellung von JAX-Modellen
So stellen Sie ein Modell auf einer TPU-VM bereit:
- Modell im SavedModel-Format von TensorFlow serialisieren
- Das gespeicherte Modell mit dem Inferenzkonverter für die Bereitstellung vorbereiten
- Modell mit TensorFlow Serving bereitstellen
SavedModel-Format
Ein SavedModel enthält ein vollständiges TensorFlow-Programm, einschließlich trainierter Parameter und Berechnungen. Der ursprüngliche Modellbaucode ist für die Ausführung nicht erforderlich.
Wenn Ihr Modell in JAX geschrieben wurde, müssen Sie jax2tf
verwenden, um es im SavedModel-Format zu serialisieren.
Inferenzkonverter
Der Cloud TPU Inference Converter bereitet und optimiert ein Modell, das im SavedModel-Format exportiert wurde, für TPU-Inferenzen vor. Sie können den Inferenzkonverter in einer lokalen Shell oder auf Ihrer TPU-VM ausführen. Wir empfehlen die Verwendung Ihrer TPU-VM-Shell, da sie alle Befehlszeilentools enthält, die zum Ausführen des Converters erforderlich sind. Weitere Informationen zum Inferenzkonverter finden Sie im Nutzerhandbuch zum Inferenzkonverter.
Anforderungen für Inferenzkonverter
Das Modell muss aus TensorFlow oder JAX im SavedModel-Format exportiert werden.
Sie müssen einen Funktionsalias für die TPU-Funktion definieren. Weitere Informationen finden Sie im Nutzerhandbuch zum Inferenzkonverter. In den Beispielen in diesem Leitfaden wird
tpu_func
als TPU-Funktionsalias verwendet.Achten Sie darauf, dass die CPU Ihrer Maschine AVX-Anweisungen (Advanced Vector eXtensions) unterstützt, da die TensorFlow-Bibliothek (die Abhängigkeit des Cloud TPU Inference Converters) für die Verwendung von AVX-Anweisungen kompiliert ist. Die meisten CPUs bieten AVX-Unterstützung.
Inferenz und Bereitstellung von JAX-Modellen
In diesem Abschnitt wird beschrieben, wie JAX-Modelle mithilfe von jax2tf
und TensorFlow Serving bereitgestellt werden.
- Verwenden Sie
jax2tf
, um Ihr Modell im SavedModel-Format zu serialisieren - Gespeichertes Modell mit dem Inferenzkonverter für die Bereitstellung vorbereiten
- Modell mit TensorFlow Serving bereitstellen
jax2tf
verwenden, um ein JAX-Modell im SavedModel-Format zu serialisieren
Die folgende Python-Funktion zeigt, wie Sie jax2tf
in Ihrem Modellcode verwenden:
# 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)
Weitere Informationen zu jax2tf
finden Sie unter Interaktion zwischen JAX und Cloud TPU.
Das gespeicherte Modell mit dem Inferenzkonverter für die Bereitstellung vorbereiten
Eine Anleitung zur Verwendung des Inferenzkonverters finden Sie im Leitfaden zum Inferenzkonverter.
TensorFlow Serving verwenden
Eine Anleitung zur Verwendung von TensorFlow Serving finden Sie unter TensorFlow Serving.
Beispiele für die Bereitstellung von JAX-Modellen
Vorbereitung
Richten Sie Ihre Docker-Anmeldedaten ein und rufen Sie das Inference Converter- und das Cloud TPU Serving-Docker-Image ab:
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
Stellen Sie mit SSH eine Verbindung zu Ihrer TPU-VM her und installieren Sie den Inferenz-Democode:
gsutil -m cp -r \ "gs://cloud-tpu-inference-public/demo" \ .
Installieren Sie die JAX-Demoabhängigkeiten:
pip install -r ./demo/jax/requirements.txt
Das JAX BERT-Modell für die Inferenz bereitstellen
Sie können das vortrainierte BERT-Modell von Hugging Face herunterladen.
Exportieren Sie ein TPU-kompatibles TensorFlow-Modell aus einem Flax-BERT-Modell:
cd demo/jax/bert python3 export_bert_model.py
Starten Sie den Containercontainer für das Cloud TPU-Modell:
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
Prüfen Sie etwa 30 Sekunden nach dem Start des Containers das Containerlog des Modellservers und vergewissern Sie sich, dass die gRPC- und HTTP-Server aktiv sind:
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker logs ${CONTAINER_ID}
Wenn ein Logeintrag mit den folgenden Informationen endet, kann der Server Anfragen verarbeiten.
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 ...
Inferenzanfrage an den Modellserver senden.
python3 bert_request.py
Die Ausgabe sollte in etwa so aussehen:
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?."
bereinigen.
Bereinigen Sie den Docker-Container, bevor Sie weitere Demos ausführen.
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker stop ${CONTAINER_ID}
Bereinigen Sie die Modellartefakte:
sudo rm -rf /tmp/jax/
JAX Stable Diffusion für die Inferenz bereitstellen
Sie können das vortrainierte Stable Diffusion-Modell von Hugging Face herunterladen.
Laden Sie das Stable Diffusion-Modell in einem TPU-kompatiblen, mit TF2 gespeicherten Modellformat herunter:
cd demo/jax/stable_diffusion python3 export_stable_diffusion_model.py
Starten Sie den Cloud TPU-Modellservercontainer für das Modell:
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
Prüfen Sie nach etwa zwei Minuten im Containerlog des Modellservers, ob die gRPC- und HTTP-Server ausgeführt werden:
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker logs ${CONTAINER_ID}
Wenn das Log mit den folgenden Informationen endet, bedeutet dies, dass die Server zum Verarbeiten von Anfragen bereit sind.
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 ...
Senden Sie eine Anfrage an den Modellserver.
python3 stable_diffusion_request.py
Dieses Skript sendet als Prompt „Gemälde eines Eichhörnchens auf Skateboard in New York“. Das Ausgabebild wird als
stable_diffusion_images.jpg
im aktuellen Verzeichnis gespeichert.bereinigen.
Bereinigen Sie den Docker-Container, bevor Sie weitere Demos ausführen.
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker stop ${CONTAINER_ID}
Modellartefakte bereinigen
sudo rm -rf /tmp/jax/
TensorFlow bereitstellen
Die folgende Anleitung zeigt, wie Sie Ihr TensorFlow-Modell auf TPU-VMs bereitstellen können.
TensorFlow-Bereitstellungsworkflow
Laden Sie das TensorFlow Serving-Docker-Image für Ihre TPU-VM herunter.
Beispiel-Umgebungsvariablen festlegen
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
Docker-Image herunterladen
docker pull ${IMAGE_NAME}
Richten Sie die Docker-Anmeldedaten ein und rufen Sie das Docker-Image für Inference Converter und TensorFlow Serving ab.
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
Laden Sie den Democode herunter:
gsutil -m cp -r \ "gs://cloud-tpu-inference-public/demo" \ .
Installieren Sie die TensorFlow-Demoabhängigkeiten:
pip install -r ./demo/tf/requirements.txt
Ihr TensorFlow-Modell mithilfe des Docker-Images von TensorFlow auf Ihrer TPU-VM bereitstellen.
# 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}
Verwenden Sie die Serving Client API, um Ihr Modell abzufragen.
Demo zu TensorFlow ResNet-50 Serving ausführen
Exportieren Sie ein TPU-kompatibles TF2-Modell aus dem Keras ResNet-50-Modell.
cd demo/tf/resnet-50 python3 export_resnet_model.py
Starten Sie den TensorFlow-Modellservercontainer für das Modell.
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
Prüfen Sie das Containerlog des Modellservers und achten Sie darauf, dass gRPC und HTTP-Server aktiv sind:
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker logs ${CONTAINER_ID}
Wenn das Protokoll mit den folgenden Informationen endet, bedeutet dies, dass der Server zum Verarbeiten von Anfragen bereit ist. Das dauert etwa 30 Sekunden.
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 ...
Senden Sie die Anfrage an den Modellserver.
Das Anforderungsbild ist eine Banane von https://i.imgur.com/j9xCCzn.jpeg .
python3 resnet_request.py
Die Ausgabe sollte in etwa so aussehen:
Predict result: [[('n07753592', 'banana', 0.94921875), ('n03532672', 'hook', 0.022338867), ('n07749582', 'lemon', 0.005126953)]]
bereinigen.
Bereinigen Sie den Docker-Container, bevor Sie weitere Demos ausführen.
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}') docker stop ${CONTAINER_ID}
Bereinigen Sie die Modellartefakte:
sudo rm -rf /tmp/tf/
Inferenz und Bereitstellung von PyTorch-Modellen
Für Modelle, die mit PyTorch geschrieben wurden, gilt folgender Workflow:
- Python-Modell-Handler zum Laden und Ableiten mit
TorchDynamo
und PyTorch/XLA schreiben - Mit
TorchModelArchiver
ein Modellarchiv erstellen TorchServe
verwenden, um das Modell bereitzustellen
TorchDynamo und PyTorch/XLA
TorchDynamo (Dynamo) ist ein JIT-Compiler auf Python-Ebene, der entwickelt wurde, um PyTorch-Programme zu beschleunigen. Es bietet eine saubere API, an die Compiler-Back-Ends eingebunden werden können. Der Python-Bytecode wird unmittelbar vor der Ausführung dynamisch geändert. Im PyTorch/XLA 2.0-Release gibt es ein experimentelles Back-End für Inferenz und Training mit Dynamo.
Dynamo liefert eine Torch FX-Grafik (FX), wenn ein Modellmuster erkannt wird. PyTorch/XLA verwendet einen Lazy Tensor-Ansatz, um den FX-Graphen zu kompilieren und die kompilierte Funktion zurückzugeben. Weitere Informationen zu Dynamo finden Sie unter:
- Beitrag in den Pytorch-Entwicklerdiskussionen
- TorchDynamo-Dokumentation
- PyTorch 2.0 und XLA für weitere Informationen
Hier ist ein kleines Codebeispiel für die Ausführung von Densenet161-Inferenz mit 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
Sie können das bereitgestellte Docker-Image torchserve-tpu
zum Bereitstellen Ihres archivierten Pytorch-Modells auf einer TPU-VM verwenden.
Richten Sie die Authentifizierung für Docker ein:
sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
us-docker.pkg.dev
Rufen Sie das Cloud TPU TorchServe-Docker-Image auf Ihre TPU-VM ab:
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}
Modellartefakte erfassen
Zuerst müssen Sie einen Modell-Handler bereitstellen, der den TorchServe-Modellserver-Worker anweist, Ihr Modell zu laden, die Eingabedaten zu verarbeiten und eine Inferenz auszuführen. Sie können die Standard-Inferenz-Handler von TorchServe (Quelle) verwenden oder einen eigenen benutzerdefinierten Modell-Handler anhand von base_handler.py entwickeln. Möglicherweise müssen Sie auch das trainierte Modell und die Modelldefinitionsdatei bereitstellen.
Im folgenden Densenet 161-Beispiel werden Modellartefakte und der von TorchServe bereitgestellte Standard-Bildklassifikator-Handler verwendet:
Konfigurieren Sie einige Umgebungsvariablen:
CWD="$(pwd)" WORKDIR="${CWD}/densenet_161" mkdir -p ${WORKDIR}/model-store mkdir -p ${WORKDIR}/logs
Laden Sie Modellartefakte aus dem Bildklassifikatorbeispiel von TorchServe herunter und kopieren Sie sie:
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}
Laden Sie die Modellgewichtungen herunter:
wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth mv densenet161-8d451a50.pth ${WORKDIR}
Erstellen Sie eine Konfigurationsdatei für das TorchServe-Modell, um das Dynamo-Back-End zu verwenden:
echo 'pt2: "torchxla_trace_once"' >> ${WORKDIR}/model_config.yaml
Sie sollten die folgenden Dateien und Verzeichnisse sehen:
>> ls ${WORKDIR} model_config.yaml index_to_name.json logs model.py densenet161-8d451a50.pth model-store
Modellarchivdatei generieren
Wenn Sie Ihr PyTorch-Modell mit Cloud TPU TorchServe bereitstellen möchten, müssen Sie Ihren Modell-Handler und alle Modellartefakte mit dem Torch Model Archiver in einer Modellarchivdatei (*.mar)
verpacken.
Generieren Sie mit torch-model-archiver eine Modellarchivdatei:
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
Die Modellarchivdatei sollte im Modellspeicherverzeichnis generiert werden:
>> ls ${WORKDIR}/model-store
Densenet161.mar
Inferenzanfragen bereitstellen
Jetzt haben Sie die Modellarchivdatei und können den TorchServe-Modellserver starten und Inferenzanfragen bereitstellen.
Starten Sie den TorchServe-Modellserver:
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
Serverstatus des Abfragemodells:
curl http://localhost:8080/ping
Wenn der Modellserver aktiv ist, wird Folgendes angezeigt:
{ "status": "Healthy" }
Verwenden Sie Folgendes, um die Standardversionen des aktuell registrierten Modells abzufragen:
curl http://localhost:8081/models
Sie sollten das registrierte Modell sehen:
{ "models": [ { "modelName": "Densenet161", "modelUrl": "Densenet161.mar" } ] }
So laden Sie ein Bild für die Inferenz herunter:
curl -O https://raw.githubusercontent.com/pytorch/serve/master/docs/images/kitten_small.jpg mv kitten_small.jpg ${WORKDIR}
So senden Sie eine Inferenzanfrage an den Modellserver:
curl http://localhost:8080/predictions/${MODEL_NAME} -T ${WORKDIR}/kitten_small.jpg
Die Antwort sieht ungefähr so aus:
{ "tabby": 0.47878125309944153, "lynx": 0.20393909513950348, "tiger_cat": 0.16572578251361847, "tiger": 0.061157409101724625, "Egyptian_cat": 0.04997897148132324 }
Modellserverlogs
Greifen Sie mit den folgenden Befehlen auf die Logs zu:
ls ${WORKDIR}/logs/ cat ${WORKDIR}/logs/model_log.log
Im Protokoll sollte die folgende Meldung angezeigt werden:
"Compiled model with backend torchxla\_trace\_once"
Bereinigen
Beenden Sie den Docker-Container:
rm -rf serve
rm -rf ${WORKDIR}
docker stop torch-model-archiver
docker stop torchserve-tpu
Profilerstellung
Nachdem Sie die Inferenz eingerichtet haben, können Sie mit Profilern die Leistung und die TPU-Auslastung analysieren. Weitere Informationen zur Profilerstellung finden Sie unter: