Esegui il deployment di un modello nelle VM Cloud TPU

Google Cloud fornisce accesso ad acceleratori di machine learning progettati su misura chiamati Tensor Processing Unit (TPU). Le TPU sono ottimizzate per accelerare l'addestramento e l'inferenza dei modelli di machine learning, caratteristica che le rende ideali per numerose applicazioni, tra cui l'elaborazione del linguaggio naturale, la visione artificiale e il riconoscimento vocale.

Questa pagina descrive come eseguire il deployment dei modelli su una singola macchina Cloud TPU v5e o v6e per l'inferenza online in Vertex AI.

Sono supportate solo le versioni v5e e v6e di Cloud TPU. Non sono supportate altre generazioni di Cloud TPU.

Per scoprire in quali località sono disponibili le versioni v5e e v6e di Cloud TPU, consulta la sezione Località.

Importare il modello

Per il deployment su Cloud TPU, devi importare il modello in Vertex AI e configurarlo per utilizzare uno dei seguenti container:

Contenitore del runtime TensorFlow ottimizzato predefinito

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

  • Ottimizzazione manuale del modello: utilizzi Inference Converter per ottimizzare il modello e salvarlo. Poi, devi passare i flag --saved_model_tags='serve,tpu' e --disable_optimizer=true quando upload il 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 tenta di ottimizzare il modello non ottimizzato utilizzando un algoritmo di partizionamento automatico. Questa ottimizzazione non funziona su tutti i modelli. Se l'ottimizzazione non va a buon fine, devi ottimizzare manualmente il modello o scegliere 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 partizionamento manuale. Specifica il flag --converter_options_string e regola ConverterOptions.TpuFunction in base alle tue esigenze. Per un esempio, vedi Immagine del convertitore. Tieni presente che è supportato solo ConverterOptions.TpuFunction, che è tutto ciò che serve per il partizionamento manuale. 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 ulteriori informazioni sull'importazione di 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 di quelle per importare ed eseguire un modello PyTorch.

Ad esempio, TorchServe per l'inferenza 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 carica il modello come mostrato di seguito:

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 ulteriori informazioni, consulta Esportare gli artefatti del modello per PyTorch e il notebook del tutorial Pubblica un modello PyTorch utilizzando un container predefinito.

Container personalizzato

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

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

Assicurati che il container personalizzato soddisfi i requisiti per i 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
    ),
  )

Per informazioni sull'importazione di un modello con un container personalizzato, consulta Utilizzo di un container personalizzato per l'inferenza. Se vuoi implementare una logica di pre-elaborazione o post-elaborazione, valuta la possibilità di utilizzare routine di inferenza personalizzate.

Creazione di un endpoint

Le istruzioni per creare un endpoint per le Cloud TPU sono le stesse di quelle per creare qualsiasi endpoint.

Ad esempio, il seguente comando crea una risorsa endpoint:

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 su Cloud TPU sono le stesse di quelle 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
ct6e-standard-1t 1
ct6e-standard-4t 4
ct6e-standard-8t 8
ct5lp-hightpu-1t 1
ct5lp-hightpu-4t 4
ct5lp-hightpu-8t 8

Gli acceleratori TPU sono integrati nel tipo di macchina. Non devi specificare il tipo o il numero di acceleratori.

Ad esempio, il seguente comando 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 saperne di più, consulta Eseguire il deployment di un modello su un endpoint.

Ottenere inferenze online

Le istruzioni per ottenere inferenze online da una Cloud TPU sono le stesse delle istruzioni per ottenere inferenze online.

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

deployed_model.predict(...)

Per i container personalizzati, consulta i requisiti di richiesta e risposta di inferenza per i container personalizzati.

Protezione della capacità

Per la maggior parte delle regioni, la TPU v5e and v6e cores per region quota per il servizio del modello personalizzato è 0. In alcune regioni, è limitato.

Per richiedere un aumento della quota, consulta Richiedi un aggiustamento delle quote.

Prezzi

I tipi di macchine TPU vengono fatturati all'ora, proprio come tutti gli altri tipo di macchina in Vertex Prediction. Per ulteriori informazioni, consulta la sezione Prezzi delle previsioni.

Passaggi successivi