Exécuter l'inférence sur Cloud TPU

La mise en service est le processus de déploiement d'un modèle de machine learning entraîné dans un environnement de production pour exécuter l'inférence. Les TPU v5e et versions ultérieures sont compatibles avec l'inférence. Le respect des objectifs de niveau de service (SLO) de latence est une priorité pour les charges de travail de diffusion.

Ce document explique comment mettre en service un modèle sur un TPU à hôte unique. Les tranches TPU comportant huit puces ou moins disposent d'une VM TPU ou d'un hôte et sont appelées TPU à hôte unique. Pour en savoir plus sur l'inférence multi-hôte avec PyTorch, consultez vLLM-TorchTPU et SGLang-TorchTPU. Pour en savoir plus sur l'inférence multi-hôte avec JAX, consultez Effectuer une inférence multi-hôte à l'aide de Pathways.

Configurer votre TPU

Quelle que soit la pile que vous choisissez, vous avez besoin d'un projet Google Cloud avec un TPU à hôte unique avant de pouvoir diffuser un modèle. Les étapes suivantes s'appliquent à toutes les combinaisons de bibliothèques et de frameworks :

  1. Créer un projet Vous avez besoin d'un compte et d'un projet Google Cloud pour utiliser Cloud TPU. Pour en savoir plus, consultez Configurer un projetGoogle Cloud pour les TPU.

  2. Vérifiez votre quota. Assurez-vous de disposer d'un quota suffisant pour le nombre de cœurs TPU que vous prévoyez d'utiliser pour l'inférence. Les TPU v5e utilisent des quotas distincts pour l'entraînement et la mise en service. Dans l'API Cloud TPU, les quotas spécifiques à la mise en service pour les TPU v5e sont les suivants :

    • Ressources v5e à la demande : TPU v5 lite pod cores for serving per project per zone
    • Ressources v5e préemptives : Preemptible TPU v5 lite pod cores for serving per project per zone

    Pour les autres versions de TPU, les charges de travail d'entraînement et de mise en service utilisent le même quota. Pour en savoir plus, consultez Quotas Cloud TPU.

  3. Créez une VM TPU. Choisissez une version de TPU et une zone disposant d'une capacité de diffusion, puis créez une VM TPU à hôte unique. Pour en savoir plus, consultez Régions et zones TPU et Gérer les ressources TPU.

Une fois la VM TPU en cours d'exécution, suivez le guide de configuration de votre bibliothèque et de votre framework dans Choisir une pile d'inférence.

Choisir une pile d'inférence

Cloud TPU est compatible avec l'inférence de grands modèles de langage (LLM) via deux bibliothèques de diffusion Open Source : vLLM et SGLang. Les deux bibliothèques s'exécutent sur Cloud TPU avec un backend JAX ou PyTorch. Vous pouvez donc choisir la bibliothèque qui convient le mieux à votre charge de travail, quel que soit le framework dans lequel votre modèle est écrit.

Choisissez d'abord la bibliothèque, puis sélectionnez le backend correspondant au framework de votre modèle, comme indiqué dans le tableau suivant :

Bibliothèque de diffusion Backend JAX Backend PyTorch
vLLM vllm-jax vllm-torchtpu
SGLang sglang-jax sglang-torchtpu

Inférence avec vLLM

vLLM est une bibliothèque Open Source pour l'inférence et la mise en service rapides de LLM, basée sur le traitement par lot continu et PagedAttention. Cloud TPU fonctionne avec vLLM via le plug-in vllm-jax pour les modèles JAX et le plug-in vllm-torchtpu pour les modèles PyTorch.

Pour obtenir d'autres exemples d'utilisation de vLLM pour diffuser un modèle sur des TPU, consultez les ressources suivantes :

vLLM avec JAX

La documentation de vLLM TPU couvre l'installation, un guide de démarrage rapide, les modèles compatibles et les options de réglage. Pour en savoir plus, consultez Diffuser des modèles JAX avec vLLM sur TPU.

vLLM avec PyTorch

Le plug-in vllm-torchtpu ajoute la compatibilité du noyau et de l'environnement d'exécution TorchTPU à vLLM. Pour en savoir plus, consultez Mettre en service des modèles PyTorch avec vLLM sur TPU.

Inférence avec SGLang

SGLang est un framework de diffusion Open Source pour les grands modèles de langage et multimodaux qui associe un environnement d'exécution hautes performances à une interface de génération structurée. Sur Cloud TPU, SGLang s'exécute via le backend SGLang-JAX (sglang-jax) pour les modèles JAX et le backend SGLang-TorchTPU (sglang-torchtpu) pour les modèles PyTorch.

SGLang avec PyTorch

Pour en savoir plus sur l'installation, le lancement du serveur et les configurations compatibles, consultez Diffuser des modèles PyTorch avec l'environnement d'exécution SGLang-TorchTPU sur TPU.

SGLang avec JAX

Pour en savoir plus sur l'installation, le lancement du serveur et les configurations compatibles, consultez Diffuser des modèles JAX avec l'environnement d'exécution SGLang-JAX sur TPU.

Profiler les performances d'inférence

Une fois l'inférence configurée, vous pouvez utiliser des profileurs pour analyser les performances et l'utilisation des TPU. Pour en savoir plus sur le profilage, consultez les ressources suivantes :