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 hôte unique Cloud TPU v5e ou v6e pour l'inférence en ligne dans Vertex AI.

Seules les versions v5e et v6e des Cloud TPU sont compatibles. Les autres générations de Cloud TPU ne sont pas compatibles.

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

Importer votre modèle

Pour déployer votre modèle sur des instances Cloud TPU, vous devez l'importer dans Vertex AI 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 SavedModel TensorFlow 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 optimiser votre modèle automatiquement. Pour ce faire, importez votre modèle. Vertex AI optimisera ensuite votre modèle non optimisé à 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 Importer des modèles dans Vertex AI.

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 montre comment empaqueter le modèle Densenet 161 dans des artefacts de modèle à l'aide de Torch Model Archiver.

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 Exporter des artefacts de modèle pour PyTorch et le notebook du tutoriel 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
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é

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

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.

Étapes suivantes