Esegui il deployment di un modello nelle VM Cloud TPU

Google Cloud dà accesso ad acceleratori di machine learning progettati su misura chiamate Tensor Processing Unit (TPU). Le TPU sono ottimizzati per accelerare l'addestramento e l'inferenza dei modelli di machine learning, il che li rende ideali per una varietà di applicazioni, incluse le applicazioni di elaborazione, visione artificiale e riconoscimento vocale.

In questa pagina viene descritto come eseguire il deployment dei modelli su un host singolo Cloud TPU v5e per la previsione online in Vertex AI.

È supportata solo la versione v5e di Cloud TPU. Altro Le generazioni di Cloud TPU non sono supportate.

Importa il modello

Per il deployment sulle Cloud TPU, devi importare il modello in Vertex AI e configurarlo in modo che utilizzi uno dei seguenti container:

Container di runtime TensorFlow ottimizzato

Per importare ed eseguire un TensorFlow SavedModel su Cloud TPU, il modello deve essere ottimizzato per TPU. Se il tuo TensorFlow SavedModel non è già una TPU ottimizzato, ci sono tre modi per ottimizzare il modello:

  • Ottimizzazione manuale del modello: utilizzi il convertitore di inferenza per ottimizzare il modello e salvarlo. Quindi, devi superare --saved_model_tags='serve,tpu' e --disable_optimizer=true quando upload del modello. Ad esempio:

    model = aiplatform.Model.upload(
    display_name='Manually optimized model',
    artifact_uri="gs://model-artifact-uri",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
    serving_container_args=[
        "--saved_model_tags=serve,tpu",
        "--disable_optimizer=true",
    ]
)

  • Ottimizzazione automatica del modello con partizionamento automatico: quando importi un modello, Vertex AI tenterà di ottimizzare il modello non ottimizzato utilizzando un algoritmo di partizionamento automatico. Questa ottimizzazione non funziona su per tutti i modelli. Se l'ottimizzazione non va a buon fine, devi ottimizzare manualmente o scegli l'ottimizzazione automatica del modello con il partizionamento manuale. Ad esempio:

    model = aiplatform.Model.upload(
    display_name='TPU optimized model with automatic partitioning',
    artifact_uri="gs://model-artifact-uri",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
    serving_container_args=[
    ]
)

  • Ottimizzazione automatica del modello con il partizionamento manuale. Specifica --converter_options_string segnala e modifica ConverterOptions.TpuFunction per soddisfare le tue esigenze. Ad esempio, consulta la sezione Immagine del convertitore. Tieni presente che solo ConverterOptions.TpuFunction, che è tutto ciò che serve il partizionamento manuale è supportato. Ad esempio:

    model = aiplatform.Model.upload(
display_name='TPU optimized model with manual partitioning',
  artifact_uri="gs://model-artifact-uri",
  serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
  serving_container_args=[
      "--converter_options_string='tpu_functions { function_alias: \"partitioning function name\" }'"
  ]
)

Per saperne di più sull'importazione dei modelli, consulta Importazione di modelli in Vertex AI.

Container PyTorch predefinito

Le istruzioni per importare ed eseguire un modello PyTorch su Cloud TPU sono le stesse come istruzioni per importare ed eseguire un modello PyTorch.

Ad esempio, TorchServe per l'inferenza di Cloud TPU v5e mostra come pacchettizzare il modello Densenet 161 in artefatti del modello utilizzando Torch Model Archiver.

Quindi, carica gli artefatti del modello nella cartella Cloud Storage e modello come mostrato:

model = aiplatform.Model.upload(
    display_name='DenseNet TPU model from SDK PyTorch 2.1',
    artifact_uri="gs://model-artifact-uri",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/pytorch-tpu.2-1:latest",
    serving_container_args=[],
    serving_container_predict_route="/predictions/model",
    serving_container_health_route="/ping",
    serving_container_ports=[8080]
)

Per maggiori informazioni, vedi Esportare gli artefatti del modello per PyTorch e il blocco note Jupyter per Fornire un modello PyTorch utilizzando un container predefinito.

Container personalizzato

Per i container personalizzati, il modello non deve essere necessariamente un ma deve essere ottimizzato per TPU. Per informazioni sulla produzione di una TPU ottimizzato, consulta le seguenti guide per i framework ML comuni:

Per informazioni sulla pubblicazione di modelli addestrati con JAX, TensorFlow PyTorch su Cloud TPU v5e, vedi Inferenza di Cloud TPU v5e.

Assicurati che il container personalizzato soddisfi i requisiti dei container personalizzati.

Devi aumentare il limite di memoria bloccata in modo che il driver possa comunicare con i chip TPU tramite accesso diretto alla memoria (DMA). Ad esempio:

Riga di comando

ulimit -l 68719476736

Python

import resource

resource.setrlimit(
    resource.RLIMIT_MEMLOCK,
    (
        68_719_476_736_000,  # soft limit
        68_719_476_736_000,  # hard limit
    ),
  )

Quindi, consulta Utilizzare un container personalizzato per la previsione. per informazioni sull'importazione di un modello con un container personalizzato. Se vuoi per implementare la logica di pre o post-elaborazione, valuta la possibilità di Routine di previsione personalizzate.

Creazione di un endpoint

Le istruzioni per creare un endpoint per le Cloud TPU sono le stesse della istruzioni per la creazione di qualsiasi endpoint.

Ad esempio, il seguente comando crea una classe endpoint risorsa:

endpoint = aiplatform.Endpoint.create(display_name='My endpoint')

La risposta contiene l'ID del nuovo endpoint, che utilizzerai nei passaggi successivi.

Per ulteriori informazioni sulla creazione di un endpoint, consulta eseguire il deployment di un modello in un endpoint.

Esegui il deployment di un modello

Le istruzioni per il deployment di un modello nelle Cloud TPU sono le stesse della istruzioni per il deployment di qualsiasi modello, tranne per il fatto che devi specificare uno dei seguenti tipi di macchine Cloud TPU supportati:

Tipo di macchina Numero di chip TPU
ct5lp-hightpu-1t 1
ct5lp-hightpu-4t 4
ct5lp-hightpu-8t 8

Gli acceleratori TPU sono integrati nel tipo di macchina. Non è necessario specificare il tipo di acceleratore o il numero di acceleratori.

Ad esempio, il comando seguente esegue il deployment di un modello chiamando deployModel:

machine_type = 'ct5lp-hightpu-1t'

deployed_model = model.deploy(
    endpoint=endpoint,
    deployed_model_display_name='My deployed model',
    machine_type=machine_type,
    traffic_percentage=100,
    min_replica_count=1
    sync=True,
)

Per ulteriori informazioni, consulta Eseguire il deployment di un modello in un endpoint.

Ricevi previsioni online

L'istruzione per ottenere previsioni online da una Cloud TPU è la stessa come istruzione per ottenere previsioni online.

Ad esempio, il comando seguente invia una richiesta di previsione online chiamando predict:

deployed_model.predict(...)

Per i container personalizzati, consulta i requisiti per le richieste e le risposte di previsione per i container personalizzati.

Protezione della capacità

Per impostazione predefinita, la quota per Custom model serving TPU v5e cores per region è 0.

Per richiedere un aumento, consulta Richiedere un limite della quota superiore.

Prezzi

I tipi di macchine TPU vengono fatturati per ora, proprio come tutti gli altri tipo di macchina in Vertex Prediction. Per maggiori informazioni informazioni, consulta Prezzi delle previsioni.

Passaggi successivi