Einführung in Cloud TPU v5e-Inferenz

Übersicht und Vorteile

Cloud TPU v5e ist ein von Google entwickelter KI-Beschleuniger, der für transformer-, Text-zu-Bild- und CNN-basiertes Training, Feinabstimmung 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 Priorität für die Bereitstellung.

In diesem Dokument wird die Bereitstellung eines Modells auf einer TPU mit einem einzelnen Host erläutert. TPU-Slices mit maximal 8 Chips haben eine TPU-VM oder einen einzigen Host und werden als Einzelhost-TPUs bezeichnet.

Mehr erfahren

Für v5e-TPUs benötigen Sie ein Kontingent. Für On-Demand-TPUs ist ein Kontingent von tpu-v5s-litepod-serving erforderlich. Reservierte TPUs benötigen ein Kontingent von tpu-v5s-litepod-serving-reserved. Weitere Informationen erhalten Sie vom Cloud-Vertrieb.

Zur Verwendung von Cloud TPU benötigen Sie ein Google Cloud-Konto und ein Projekt. Weitere Informationen finden Sie unter Cloud TPU-Umgebung einrichten.

v5e-TPUs werden mit 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.

Inferenz und Bereitstellung von Cloud TPU-Modellen

Wie Sie ein Modell für Inferenzen bereitstellen, hängt vom ML-Framework ab, mit dem Ihr Modell geschrieben wurde. TPU v5e unterstützt Bereitstellungsmodelle, die in JAX, TensorFlow und PyTorch geschrieben wurden.

Inferenz und Bereitstellung von JAX-Modellen

So stellen Sie ein Modell auf einer TPU-VM bereit:

  1. Serialisieren Sie das Modell im SavedModel von TensorFlow.
  2. Inferenzkonverter verwenden, um das gespeicherte Modell für die Bereitstellung vorzubereiten
  3. TensorFlow Serving zum Bereitstellen des Modells verwenden

Memcache-Format

Ein Literal enthält ein vollständiges TensorFlow-Programm, einschließlich der trainierten Parameter und der Berechnung. Der ursprüngliche Modellerstellungscode ist für die Ausführung nicht erforderlich.

Wenn Ihr Modell in JAX geschrieben wurde, müssen Sie jax2tf verwenden, um es im GeoJSON-Format zu serialisieren.

Inferenzkonverter

Der Cloud TPU-Inferenzkonverter bereitet ein Modell, das im SavedModel-Format exportiert wurde, für TPU-Inferenzen auf und optimiert es. 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 für den Inferenzkonverter.

Anforderungen für Inferenzkonverter

  1. Das Modell muss aus TensorFlow oder JAX im SavedModel exportiert werden.

  2. Sie müssen einen Funktionsalias für die TPU-Funktion definieren. Weitere Informationen finden Sie im Nutzerhandbuch für den Inferenzkonverter. In den Beispielen in diesem Leitfaden wird tpu_func als TPU-Funktionsalias verwendet.

  3. Achten Sie darauf, dass die CPU Ihrer Maschine Advanced Vector eXtensions (AVX) unterstützt, da die TensorFlow-Bibliothek (die Abhängigkeit des Cloud TPU Inference Converters) für die Verwendung von AVX-Anweisungen kompiliert wird. Die meisten CPUs haben die AVX-Unterstützung.

Inferenz und Bereitstellung von JAX-Modellen

In diesem Abschnitt wird beschrieben, wie JAX-Modelle mit jax2tf und TensorFlow Serving bereitgestellt werden.

  1. Verwenden Sie jax2tf, um Ihr Modell im Memcache-Format zu serialisieren
  2. Inferenzkonverter zur Vorbereitung des gespeicherten Modells für die Bereitstellung verwenden
  3. TensorFlow Serving zum Bereitstellen des Modells verwenden

Verwenden Sie jax2tf, um ein JAX-Modell in das Memcache-Format zu serialisieren

Die folgende Python-Funktion zeigt, wie jax2tf in Ihrem Modellcode verwendet wird:

# 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 Interoperabilität von JAX und Cloud TPU.

Inferenzkonverter verwenden, um das gespeicherte Modell für die Bereitstellung vorzubereiten

Eine Anleitung zur Verwendung des Inferenzkonverters finden Sie im Leitfaden zum Inferenzkonverter.

TensorFlow Serving verwenden

Anleitungen zur Verwendung von TensorFlow Serving finden Sie unter TensorFlow Serving.

Beispiele für die Bereitstellung von JAX-Modellen

Vorbereitung

  1. Richten Sie Ihre Docker-Anmeldedaten ein und rufen Sie den Inference Converter und das Cloud TPU-Bereitstellungs-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
    
  2. Stellen Sie eine SSH-Verbindung zur TPU-VM her und installieren Sie den Democode für die Inferenz:

    gsutil -m cp -r \
    "gs://cloud-tpu-inference-public/demo" \
    .
    
  3. 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.

  1. Exportieren Sie ein TPU-kompatibles gespeichertes TensorFlow-Modell aus einem Flax-BERT-Modell:

    cd demo/jax/bert
    python3 export_bert_model.py
    
  2. Starten Sie den Container des Cloud TPU-Modellservers:

    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 achten Sie darauf, dass die gRPC- und HTTP-Server aktiv sind:

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

    Wenn Sie einen Logeintrag sehen, der 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 ...
    
  3. 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?."
    
  4. bereinigen.

    Bereinigen Sie den Docker-Container, bevor Sie andere 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 zur Inferenz verwenden

Sie können ein vortrainiertes stabiles Diffusionsmodell von der Hugging Face App herunterladen.

  1. Laden Sie das stabile Diffusion-Modell in einem TPU-kompatiblen TF2-Format für gespeicherte Modelle herunter:

    cd demo/jax/stable_diffusion
    python3 export_stable_diffusion_model.py
    
  2. 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 das Containerlog des Modellservers, um sicherzugehen, dass 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 Anfragen verarbeiten können.

    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 ...
    
  3. Senden Sie eine Anfrage an den Modellserver.

    python3 stable_diffusion_request.py
    

    Dieses Skript sendet als Prompt „Gemälde eines Eichhörnchens, das in New York läuft“. Das Ausgabebild wird als stable_diffusion_images.jpg in Ihrem aktuellen Verzeichnis gespeichert.

  4. bereinigen.

    Bereinigen Sie den Docker-Container, bevor Sie andere 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

  1. TensorFlow Serving-Docker-Image für die TPU-VM herunterladen

    Beispielumgebungsvariablen 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}
    
  2. Richten Sie die Docker-Anmeldedaten ein und rufen Sie das Inference Converter und das TensorFlow 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
    
  3. Lade den Democode herunter:

    gsutil -m cp -r \
    "gs://cloud-tpu-inference-public/demo" \
    .
    
  4. Installieren Sie die TensorFlow-Demoabhängigkeiten:

    pip install -r ./demo/tf/requirements.txt
    
  5. Stellen Sie Ihr TensorFlow-Modell mit dem TensorFlow Serving-Docker-Image auf der TPU-VM bereit.

    # 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}
    
  6. Verwenden Sie die Serving Client API, um Ihr Modell abzufragen.

TensorFlow ResNet-50 Serving-Demo ausführen

  1. Exportieren Sie ein TPU-kompatibles gespeichertes TF2-Modell aus dem Keras ResNet-50-Modell.

    cd demo/tf/resnet-50
    python3 export_resnet_model.py
    
  2. Starten Sie den Container des TensorFlow-Modellservers 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 der gRPC und der 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 Anfragen verarbeiten kann. 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 ...
    
  3. Senden Sie die Anfrage an den Modellserver.

    Das Bild der Anfrage 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)]]
    
  4. bereinigen.

    Bereinigen Sie den Docker-Container, bevor Sie andere 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/
    

PyTorch-Modellinferenz und -bereitstellung

Für Modelle, die mit PyTorch geschrieben wurden, lautet der Workflow:

  1. Schreiben Sie einen Python-Modell-Handler zum Laden und Inferenzen mit TorchDynamo und PyTorch/XLA
  2. Verwenden Sie TorchModelArchiver, um ein Modellarchiv zu erstellen
  3. TorchServe zum Bereitstellen des Modells verwenden

TorchDynamo und PyTorch/XLA

TorchDynamo (Dynamo) ist ein JIT-Compiler auf Python-Ebene, der entwickelt wurde, um PyTorch-Programme schneller zu machen. Sie bietet eine saubere API, an die Compiler-Back-Ends angeschlossen werden können. Sie ändert den Python-Bytecode kurz vor der Ausführung dynamisch. In der PyTorch/XLA 2.0-Version gibt es ein experimentelles Back-End für Inferenz und Training mit Dynamo.

Dynamo stellt einen Torch FX-Diagramm (FX-Diagramm) bereit, 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:

Hier ist ein kleines Codebeispiel für das Ausführen einer 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 die 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

Als Erstes müssen Sie einen Modell-Handler bereitstellen, der den TorchServe-Modellserver-Worker anweist, Ihr Modell zu laden, die Eingabedaten zu verarbeiten und die Inferenz auszuführen. Sie können die Standard-Inferenz-Handler von TorchServe (Quelle) verwenden oder anhand von base_handler.py einen eigenen benutzerdefinierten Modell-Handler entwickeln. Möglicherweise müssen Sie auch das trainierte Modell und die Modelldefinitionsdatei angeben.

Im folgenden Densenet 161-Beispiel verwenden wir Modellartefakte und den von TorchServe bereitgestellten Standard-Bildklassifikator-Handler:

  1. Konfigurieren Sie einige Umgebungsvariablen:

    CWD="$(pwd)"
    
    WORKDIR="${CWD}/densenet_161"
    
    mkdir -p ${WORKDIR}/model-store
    mkdir -p ${WORKDIR}/logs
    
  2. Laden Sie Modellartefakte aus dem Beispiel für den TorchServe-Bildklassifikator 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}
    
  3. Laden Sie die Modellgewichtungen herunter:

    wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth
    
    mv densenet161-8d451a50.pth ${WORKDIR}
    
  4. Erstellen Sie eine TorchServe-Modellkonfigurationsdatei für das Dynamo-Back-End:

    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 das PyTorch-Modell mit Cloud TPU TorchServe bereitstellen möchten, müssen Sie den Modell-Handler und alle Modellartefakte mit dem Torch Model Archiver in einer Modellarchivdatei (*.mar) verpacken.

Erstellen Sie eine Modellarchivdatei mit „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

Sie sollten die im Verzeichnis „model-store“ generierte Modellarchivdatei sehen:

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

Inferenzanfragen verarbeiten

Mit der Modellarchivdatei können Sie den TorchServe-Modellserver starten und Inferenzanfragen verarbeiten.

  1. 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
    
  2. Serverzustand des Abfragemodells:

    curl http://localhost:8080/ping
    

    Wenn der Modellserver einsatzbereit ist, wird Folgendes angezeigt:

    {
      "status": "Healthy"
    }
    

    Verwenden Sie zum Abfragen der Standardversionen des aktuell registrierten Modells:

    curl http://localhost:8081/models
    

    Sie sollten das registrierte Modell sehen:

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

    So laden Sie ein Bild zur Inferenznutzung 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
    }
    
  3. Modellserverlogs

    Verwenden Sie die folgenden Befehle, um auf die Logs zuzugreifen:

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

    In Ihrem 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

Nach dem Einrichten der Inferenz können Sie Profiler verwenden, um die Leistung und TPU-Auslastung zu analysieren. Weitere Informationen zur Profilerstellung finden Sie unter: