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 pour la prédiction en ligne dans Vertex AI.

Seule la version v5e de Cloud TPU est compatible. Les autres générations Cloud TPU ne sont pas compatibles.

Importer votre modèle

Pour le déploiement sur Cloud TPU, vous devez importer votre modèle 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, trois options s'offrent à vous pour optimiser votre modèle :

  • Optimisation manuelle du modèle : vous utilisez le convertisseur d'inférence pour optimiser votre modèle et l'enregistrer. Vous devez ensuite transmettre les options --saved_model_tags='serve,tpu' et --disable_optimizer=true lorsque vous importez (upload) votre modèle. Exemple :

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

  • Optimisation automatique du modèle avec partitionnement automatique : lorsque vous importez un modèle, Vertex AI tente d'optimiser 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 ou choisir l'optimisation automatique du modèle avec partitionnement manuel. Exemple :

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

  • Optimisation automatique du modèle avec partitionnement manuel. Spécifiez l'option --converter_options_string et ajustez le paramètre ConverterOptions.TpuFunction en fonction de vos besoins. Pour obtenir un exemple, consultez la page Image de convertisseur. Notez que seul ConverterOptions.TpuFunction, soit tout ce qui est nécessaire pour le partitionnement manuel, est accepté. Exemple :

    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\" }'"
  ]
)

Pour en savoir plus sur l'importation de modèles, consultez la page 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 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.

Ensuite, importez les artefacts de modèle dans votre dossier Cloud Storage et importez le 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 Jupyter 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 la page Inférence Cloud TPU v5e.

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

Vous devez augmenter la limite de mémoire verrouillée pour que le pilote puisse communiquer avec les puces TPU via l'accès direct à la mémoire (DM). 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 la page Utiliser un conteneur personnalisé pour la prédiction. Si vous souhaitez mettre en œuvre une logique de pré ou post-traitement, envisagez d'utiliser des routines de prédiction personnalisées.

Créer un point de terminaison

Les instructions de création d'un point de terminaison pour Cloud TPU sont identiques à celles de création d'un 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 la section Déployer un modèle sur un point de terminaison.

Déployer un modèle

Les instructions de déploiement d'un modèle sur des Cloud TPU sont identiques à celles du déploiement de n'importe quel modèle, à la différence que vous spécifiez l'un des types de machines Cloud TPU compatibles suivants :

Type de machine Nombre de puces TPU
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 la section Déployer un modèle sur un point de terminaison.

Obtenir des prédictions en ligne

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

Par exemple, la commande suivante envoie une requête de prédiction en ligne en appelant predict :

deployed_model.predict(...)

Pour les conteneurs personnalisés, consultez les exigences concernant les requêtes et les réponses de prédiction pour les conteneurs personnalisés.

Sécuriser la capacité

Par défaut, le quota pour Custom model serving TPU v5e cores per region est de 0.

Pour demander une augmentation de quota, consultez la section Demander une augmentation de limite 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 la page Tarifs des prédictions.

Étapes suivantes