Diffuser des modèles Open Source à l'aide de TPU sur GKE avec Optimum TPU

Ce tutoriel vous explique comment diffuser des modèles de grands modèles de langage (LLM) Open Source à l'aide de Tensor Processing Units (TPU) sur Google Kubernetes Engine (GKE) avec le framework de diffusion Optimum TPU de Hugging Face. Dans ce tutoriel, vous allez télécharger des modèles Open Source depuis Hugging Face et les déployer sur un cluster GKE Standard à l'aide d'un conteneur qui exécute Optimum TPU.

Ce guide est un bon point de départ si vous avez besoin du contrôle précis, de l'évolutivité, de la résilience, de la portabilité et de la rentabilité des services Kubernetes gérés lors du déploiement et de la diffusion de vos charges de travail d'IA/de ML.

Ce tutoriel est destiné aux clients d'IA générative de l'écosystème Hugging Face, aux utilisateurs nouveaux ou existants de GKE, aux ingénieurs en ML, aux ingénieurs MLOps (DevOps) ou aux administrateurs de plate-forme qui souhaitent utiliser les fonctionnalités d'orchestration de conteneurs Kubernetes pour diffuser des LLM.

Pour rappel, vous disposez de plusieurs options pour l'inférence LLM sur Google Cloud, qui couvrent des offres telles que Vertex AI, GKE et Google Compute Engine. Vous pouvez y intégrer des bibliothèques de diffusion telles que JetStream, vLLM et d'autres offres partenaires. Par exemple, vous pouvez utiliser JetStream pour obtenir les dernières optimisations du projet. Si vous préférez les options Hugging Face, vous pouvez utiliser Optimum TPU.

Optimum TPU est compatible avec les fonctionnalités suivantes :

  • Regroupement continu
  • Streaming de jetons
  • Recherche gloutonne et échantillonnage multinomial à l'aide de transformateurs.

Accéder au modèle

Vous pouvez utiliser les modèles Gemma 2B ou Llama3 8B. Ce tutoriel se concentre sur ces deux modèles, mais Optimum TPU est compatible avec d'autres modèles.

Gemma 2B

Pour accéder aux modèles Gemma en vue du déploiement sur GKE, vous devez d'abord signer le contrat d'autorisation de licence, puis générer un jeton d'accès Hugging Face.

Vous devez signer le contrat de consentement pour utiliser Gemma. Procédez comme suit :

  1. Accédez à la page de consentement du modèle.
  2. Vérifiez l'autorisation à l'aide de votre compte Hugging Face.
  3. Acceptez les conditions du modèle.

Générer un jeton d'accès

Générez un nouveau jeton Hugging Face si vous n'en possédez pas déjà un :

  1. Cliquez sur Your Profile > Settings > Access Tokens (Votre profil > Paramètres > Jetons d'accès).
  2. Cliquez sur Nouveau jeton.
  3. Spécifiez le nom de votre choix et un rôle d'au moins Read.
  4. Cliquez sur Generate a token (Générer un jeton).
  5. Copiez le jeton dans votre presse-papiers.

Llama 3 8B

Vous devez signer le contrat de consentement pour utiliser Llama3 8b dans le dépôt Hugging Face.

Générer un jeton d'accès

Générez un nouveau jeton Hugging Face si vous n'en possédez pas déjà un :

  1. Cliquez sur Your Profile > Settings > Access Tokens (Votre profil > Paramètres > Jetons d'accès).
  2. Sélectionnez New Token (Nouveau jeton).
  3. Spécifiez le nom de votre choix et un rôle d'au moins Read.
  4. Sélectionnez Générer un jeton.
  5. Copiez le jeton dans votre presse-papiers.

Créer un cluster GKE

Créez un cluster GKE Standard avec un nœud de processeur :

gcloud container clusters create CLUSTER_NAME \
    --project=PROJECT_ID \
    --num-nodes=1 \
    --location=ZONE

Créer un pool de nœuds TPU

Créez un pool de nœuds TPU v5e avec un nœud et huit puces :

gcloud container node-pools create tpunodepool \
    --location=ZONE \
    --num-nodes=1 \
    --machine-type=ct5lp-hightpu-8t \
    --cluster=CLUSTER_NAME

Si des ressources TPU sont disponibles, GKE provisionne le pool de nœuds. Si les ressources TPU sont temporairement indisponibles, le résultat affiche un message d'erreur GCE_STOCKOUT. Pour résoudre les erreurs de rupture de stock de TPU, consultez Ressources TPU insuffisantes pour répondre à la demande de TPU.

Configurez kubectl de manière à communiquer avec votre cluster :

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${ZONE}

Créer le conteneur

Exécutez la commande suivante pour créer l'image :

cd optimum-tpu && make tpu-tgi

Transfert de l'image vers Artifact Registry

gcloud artifacts repositories create optimum-tpu --repository-format=docker --location=REGION_NAME && \
gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
docker image tag huggingface/optimum-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest && \
docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest

Créer un secret Kubernetes pour les identifiants Hugging Face

Créez un secret Kubernetes contenant le jeton Hugging Face:

kubectl create secret generic hf-secret \
  --from-literal=hf_api_token=${HF_TOKEN} \
  --dry-run=client -o yaml | kubectl apply -f -

Déployer Optimum TPU

Pour déployer Optimum TPU, ce tutoriel utilise un déploiement Kubernetes. Un déploiement est un objet de l'API Kubernetes qui vous permet d'exécuter plusieurs instances dupliquées de pods répartis entre les nœuds d'un cluster.

Gemma 2B

  1. Enregistrez le fichier manifeste de déploiement suivant sous le nom optimum-tpu-gemma-2b-2x4.yaml :

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: tgi-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: tgi-tpu
  template:
    metadata:
      labels:
        app: tgi-tpu
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      containers:
      - name: tgi-tpu
        image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
        args:
        - --model-id=google/gemma-2b
        - --max-concurrent-requests=4
        - --max-input-length=8191
        - --max-total-tokens=8192
        - --max-batch-prefill-tokens=32768
        - --max-batch-size=16
        securityContext:
            privileged: true
        env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
        ports:
        - containerPort: 80
        resources:
          limits:
            google.com/tpu: 8
        livenessProbe:
          httpGet:
            path: /health
            port: 80
          initialDelaySeconds: 300
          periodSeconds: 120

---
apiVersion: v1
kind: Service
metadata:
  name: service
spec:
  selector:
    app: tgi-tpu
  ports:
    - name: http
      protocol: TCP
      port: 8080
      targetPort: 80

    Ce fichier manifeste décrit un déploiement Optimum TPU avec un équilibreur de charge interne sur le port TCP 8080.

  2. Appliquer le fichier manifeste

    kubectl apply -f optimum-tpu-gemma-2b-2x4.yaml

Llama 3 8B

  1. Enregistrez le manifeste suivant sous le nom optimum-tpu-llama3-8b-2x4.yaml :

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: tgi-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: tgi-tpu
  template:
    metadata:
      labels:
        app: tgi-tpu
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      containers:
      - name: tgi-tpu
        image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
        args:
        - --model-id=meta-llama/Meta-Llama-3-8B
        - --max-concurrent-requests=4
        - --max-input-length=8191
        - --max-total-tokens=8192
        - --max-batch-prefill-tokens=32768
        - --max-batch-size=16
        env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
        ports:
        - containerPort: 80
        resources:
          limits:
            google.com/tpu: 8
        livenessProbe:
          httpGet:
            path: /health
            port: 80
          initialDelaySeconds: 300
          periodSeconds: 120
---
apiVersion: v1
kind: Service
metadata:
  name: service
spec:
  selector:
    app: tgi-tpu
  ports:
    - name: http
      protocol: TCP
      port: 8080
      targetPort: 80

    Ce fichier manifeste décrit un déploiement Optimum TPU avec un équilibreur de charge interne sur le port TCP 8080.

  2. Appliquer le fichier manifeste

    kubectl apply -f optimum-tpu-llama3-8b-2x4.yaml

Affichez les journaux du déploiement en cours d'exécution :

kubectl logs -f -l app=tgi-tpu

La sortie devrait ressembler à ce qui suit :

2024-07-09T22:39:34.365472Z  WARN text_generation_router: router/src/main.rs:295: no pipeline tag found for model google/gemma-2b
2024-07-09T22:40:47.851405Z  INFO text_generation_router: router/src/main.rs:314: Warming up model
2024-07-09T22:40:54.559269Z  INFO text_generation_router: router/src/main.rs:351: Setting max batch total tokens to 64
2024-07-09T22:40:54.559291Z  INFO text_generation_router: router/src/main.rs:352: Connected
2024-07-09T22:40:54.559295Z  WARN text_generation_router: router/src/main.rs:366: Invalid hostname, defaulting to 0.0.0.0

Assurez-vous que le modèle est entièrement téléchargé avant de passer à la section suivante.

Diffuser le modèle

Configurez le transfert de port vers le modèle :

kubectl port-forward svc/service 8080:8080

Interagir avec le serveur de modèles à l'aide de curl

Vérifiez vos modèles déployés :

Dans une nouvelle session de terminal, utilisez curl pour discuter avec le modèle :

curl 127.0.0.1:8080/generate     -X POST     -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":40}}'     -H 'Content-Type: application/json'

La sortie devrait ressembler à ce qui suit :

{"generated_text":"\n\nDeep learning is a subset of machine learning that uses artificial neural networks to learn from data.\n\nArtificial neural networks are inspired by the way the human brain works. They are made up of multiple layers"}