Déployer un modèle sur des VM Cloud TPU

Google Cloud donne accès à des accélérateurs de machine learning conçus sur mesure, qui sont appelés Tensor Processing Units (TPU). Les TPU sont optimisés pour accélérer l'entraînement et l'inférence des modèles de machine learning. Ils sont donc parfaitement adaptés à diverses applications, telles que le traitement du langage naturel, la vision par ordinateur et la reconnaissance vocale.

Cette page explique comment déployer vos modèles sur un seul hôte Cloud TPU v5e, v6e ou TPU7x pour l'inférence en ligne dans Gemini Enterprise Agent Platform.

Seules les versions v5e, v6e, et TPU7x des Cloud TPU (tpu7x-standard-4t type de machine) sont compatibles. Les autres générations Cloud TPU ne sont pas compatibles.

Pour savoir à quel emplacement les versions v5e, v6e et TPU7x des Cloud TPU sont disponibles, consultez les emplacements.

Importer votre modèle

Pour le déploiement sur Cloud TPU, vous devez importer votre modèle dans Agent Platform et le configurer pour utiliser l'un des conteneurs suivants :

Conteneur d'exécution TensorFlow prédéfini optimisé

Pour importer et exécuter un TensorFlow SavedModel sur un Cloud TPU, le modèle doit être optimisé pour les TPU. Si votre SavedModel TensorFlow n'est pas déjà optimisé pour les TPU, vous pouvez l'optimiser automatiquement. Pour ce faire, importez votre modèle, puis Gemini Enterprise Agent Platform l'optimise à l'aide d'un algorithme de partitionnement automatique. Cette optimisation ne fonctionne pas sur tous les modèles. Si l'optimisation échoue, vous devez optimiser manuellement votre modèle.

L'exemple de code suivant montre comment utiliser l'optimisation automatique des modèles avec le partitionnement automatique :

  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=[
      ]
  )

Pour en savoir plus sur l'importation de modèles, consultez la page Importer des modèles dans Agent Platform.

Conteneur PyTorch prédéfini

Les instructions permettant d'importer et d'exécuter un modèle PyTorch sur un Cloud TPU sont les mêmes que pour importer et exécuter un modèle PyTorch.

Par exemple, Inférence TorchServe pour Cloud TPU v5e Ensuite, importez les artefacts de modèle dans votre dossier Cloud Storage et importez votre modèle comme indiqué ci-dessous :

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]
)

Pour en savoir plus, consultez les sections Exporter des artefacts de modèle pour PyTorch et le notebook de tutoriel pour Diffuser un modèle PyTorch à l'aide d'un conteneur prédéfini.

Conteneur personnalisé

Pour les conteneurs personnalisés, votre modèle ne doit pas nécessairement être un modèle TensorFlow, mais il doit être optimisé pour les TPU. Pour en savoir plus sur la production d'un modèle optimisé pour les TPU, consultez les guides suivants pour les frameworks de ML courants :

Pour en savoir plus sur la diffusion de modèles entraînés avec JAX, TensorFlow ou PyTorch sur Cloud TPU v5e, consultez Inférence sur Cloud TPU v5e.

Assurez-vous que votre conteneur personnalisé répond aux exigences concernant les conteneurs personnalisés.

Vous devez augmenter la limite de mémoire verrouillée pour que le pilote puisse communiquer avec les puces TPU par le biais de l'accès direct à la mémoire. Exemple :

Ligne de commande

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
    ),
  )

Pour en savoir plus sur l'importation d'un modèle avec un conteneur personnalisé, consultez Utiliser un conteneur personnalisé pour l'inférence. Si vous souhaitez appliquer une logique de prétraitement ou post-traitement, envisagez d'utiliser des routines d'inférence personnalisées.

Créer un point de terminaison

La procédure à suivre pour créer un point de terminaison pour des Cloud TPU est identiques à celle permettant de créer n'importe quel point de terminaison.

Par exemple, la commande suivante crée une ressource endpoint :

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

La réponse contient l'ID du nouveau point de terminaison, que vous utiliserez dans les étapes suivantes.

Pour en savoir plus sur la création d'un point de terminaison, consultez Déployer un modèle sur un point de terminaison.

Déployer un modèle

La procédure à suivre pour déployer un modèle sur des Cloud TPU est identique à celle permettant de déployer n'importe quel modèle, à la différence que vous devez spécifier l'un des types de machines Cloud TPU compatibles suivants :

Type de machine Nombre de puces TPU
tpu7x-standard-4t 4
ct6e-standard-1t 1
ct6e-standard-4t 4
ct6e-standard-8t 8
ct5lp-hightpu-1t 1
ct5lp-hightpu-4t 4
ct5lp-hightpu-8t 8

Les accélérateurs TPU sont intégrés au type de machine. Vous n'avez pas besoin de spécifier le type ou le nombre d'accélérateurs.

Par exemple, la commande suivante déploie un modèle en appelant 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,
)

Pour en savoir plus, consultez Déployer un modèle sur un point de terminaison.

Obtenir des inférences en ligne

La procédure à suivre pour obtenir des inférences en ligne à partir d'un Cloud TPU est identique à celle permettant d'obtenir des inférences en ligne.

Par exemple, la commande suivante envoie une requête d'inférence en ligne en appelant predict :

deployed_model.predict(...)

Pour les conteneurs personnalisés, consultez les exigences concernant les requêtes et les réponses d'inférence pour les conteneurs personnalisés.

Sécuriser la capacité

Pour la plupart des régions, le TPU v5e, v6e, and TPU7x cores per region quota pour la diffusion de modèles personnalisés est de 0. Dans certaines régions, il est limité.

Pour demander une augmentation de quota, consultez Demander un ajustement de quota.

Tarifs

Les types de machines TPU sont facturés à l'heure, comme tous les autres types de machines Vertex Prediction. Pour en savoir plus, consultez les tarifs des prédictions.

Étape suivante