Gestisci Stable Diffusion XL (SDXL) utilizzando le TPU su GKE con MaxDiffusion

Autopilot Standard

Questo tutorial mostra come pubblicare un modello di generazione di immagini SDXL utilizzando le Tensor Processing Unit (TPU) su Google Kubernetes Engine (GKE) con MaxDiffusion. In questo tutorial scaricherai il modello da Hugging Face ed eseguirai il deployment su un cluster Autopilot o Standard utilizzando un container che esegue MaxDiffusion.

Questa guida è un buon punto di partenza se hai bisogno del controllo granulare, della personalizzazione, della scalabilità, della resilienza, della portabilità e dell'economicità di Kubernetes gestito durante il deployment e la gestione dei tuoi carichi di lavoro AI/ML. Se hai bisogno di una piattaforma di AI gestita unificata per creare e pubblicare rapidamente modelli ML in modo economico, ti consigliamo di provare la nostra soluzione di deployment Vertex AI.

Sfondo

Se utilizzi SDXL con le TPU su GKE con MaxDiffusion, puoi creare una soluzione di gestione solida e pronta per la produzione con tutti i vantaggi di Kubernetes gestito, tra cui efficienza dei costi, scalabilità e maggiore disponibilità. Questa sezione descrive le tecnologie chiave utilizzate in questo tutorial.

Stable Diffusion XL (SDXL)

Stable Diffusion XL (SDXL) è un tipo di modello di diffusione latente (LDM) supportato da MaxDiffusion per l'inferenza. Per l'AI generativa, puoi utilizzare i modelli LDM per generare immagini di alta qualità a partire da descrizioni testuali. I modelli LDM sono utili per applicazioni come la ricerca immagini e la creazione di didascalie per le immagini.

SDXL supporta l'inferenza su uno o più host con annotazioni di sharding. In questo modo, SDXL può essere addestrato ed eseguito su più macchine, il che può migliorare l'efficienza.

Per saperne di più, consulta il repository Generative Models di Stability AI e l'articolo SDXL.

TPU

Le TPU sono circuiti integrati specifici per le applicazioni (ASIC) sviluppati da Google e utilizzati per accelerare i modelli di machine learning e AI creati utilizzando framework come TensorFlow, PyTorch e JAX.

Prima di utilizzare le TPU in GKE, ti consigliamo di completare il seguente percorso di apprendimento:

1. Scopri di più sulla disponibilità della versione attuale della TPU con l'architettura di sistema di Cloud TPU.
2. Scopri di più sulle TPU in GKE.

Questo tutorial illustra l'erogazione del modello SDXL. GKE esegue il deployment del modello su nodi TPU v5e a host singolo con topologie TPU configurate in base ai requisiti del modello per la gestione dei prompt con bassa latenza. In questa guida, il modello utilizza un chip TPU v5e con una topologia 1x1.

MaxDiffusion

MaxDiffusion è una raccolta di implementazioni di riferimento, scritte in Python e Jax, di vari modelli di diffusione latente che vengono eseguiti su dispositivi XLA, tra cui TPU e GPU. MaxDiffusion è un punto di partenza per i progetti di diffusione sia per la ricerca che per la produzione.

Per saperne di più, consulta il repository MaxDiffusion.

Obiettivi

Questo tutorial è rivolto ai clienti dell'AI generativa che utilizzano JAX, agli utenti nuovi o esistenti di SDXL e a qualsiasi ML Engineer, MLOps (DevOps) Engineer o amministratore della piattaforma interessato a utilizzare le funzionalità di orchestrazione dei container Kubernetes per il servizio di LLM.

Questo tutorial illustra i seguenti passaggi:

1. Crea un cluster GKE Autopilot o Standard con la topologia TPU consigliata, in base alle caratteristiche del modello.
2. Crea un'immagine container di inferenza SDXL.
3. Esegui il deployment del server di inferenza SDXL su GKE.
4. Pubblica e interagisci con il modello tramite un'app web.

Architettura

Questa sezione descrive l'architettura GKE utilizzata in questo tutorial. L'architettura è costituita da un cluster GKE Autopilot o Standard che esegue il provisioning delle TPU e ospita i componenti MaxDiffusion. GKE utilizza questi componenti per eseguire il deployment e pubblicare i modelli.

Il seguente diagramma mostra i componenti di questa architettura:

Questa architettura include i seguenti componenti:

• Un cluster regionale GKE Autopilot o Standard.
• Un singolo pool di nodi di slice TPU single-host che ospita il modello SDXL nel deployment MaxDiffusion.
• Il componente Service con un bilanciatore del carico di tipo ClusterIP. Questo servizio distribuisce il traffico in entrata a tutte le repliche MaxDiffusion HTTP.
• Il server WebApp HTTP con un servizio LoadBalancer esterno che distribuisce il traffico in entrata e reindirizza il traffico di erogazione del modello a ClusterIP Service.

Prima di iniziare

Accedi al tuo account Google Cloud . Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti senza costi per l'esecuzione, il test e il deployment dei workload.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required APIs.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the APIs

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required APIs.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the APIs

• Assicurati di disporre dei seguenti ruoli nel progetto: roles/container.admin, roles/iam.serviceAccountAdmin, roles/artifactregistry.admin, roles/cloudbuild.builds.editor

  Controlla i ruoli

  1. Nella console Google Cloud vai alla pagina IAM.

     Vai a IAM
  2. Seleziona il progetto.

  3. Nella colonna Entità, trova tutte le righe che identificano te o un gruppo di cui fai parte. Per scoprire a quali gruppi appartieni, contatta il tuo amministratore.

  4. Per tutte le righe che ti specificano o ti includono, controlla la colonna Ruolo per verificare se l'elenco dei ruoli include i ruoli richiesti.

  Concedi i ruoli

  1. Nella console Google Cloud vai alla pagina IAM.

     Vai a IAM
  2. Seleziona il progetto.
  3. Fai clic su Concedi l'accesso.

  4. Nel campo Nuove entità, inserisci il tuo identificatore dell'utente. In genere si tratta dell'indirizzo email di un Account Google.

  5. Fai clic su Seleziona un ruolo, quindi cerca il ruolo.
  6. Per concedere altri ruoli, fai clic su Aggiungi un altro ruolo e aggiungi ogni ruolo successivo.
  7. Fai clic su Salva.

• Verifica di disporre di una quota sufficiente per i chip TPU v5e PodSlice Lite. In questo tutorial utilizzi le istanze on demand.

Prepara l'ambiente

In questo tutorial utilizzerai Cloud Shell per gestire le risorse ospitate su Google Cloud. Cloud Shell include il software preinstallato necessario per questo tutorial, tra cui kubectl e gcloud CLI.

Per configurare l'ambiente con Cloud Shell:

1. Nella console Google Cloud , avvia una sessione di Cloud Shell facendo clic su Attiva Cloud Shell nella Google Cloud console. Viene avviata una sessione nel riquadro inferiore della console Google Cloud .

2. Imposta le variabili di ambiente predefinite:

   gcloud config set project PROJECT_ID
   gcloud config set billing/quota_project PROJECT_ID
   export PROJECT_ID=$(gcloud config get project)
   export CLUSTER_NAME=CLUSTER_NAME
   export CLUSTER_VERSION=CLUSTER_VERSION
   export REGION=REGION_NAME
   export ZONE=ZONE
   

   Sostituisci i seguenti valori:

   • PROJECT_ID: il tuo Google Cloud ID progetto.
   • CLUSTER_NAME: il nome del cluster GKE.
   • CLUSTER_VERSION : la versione di GKE. Devi specificare una versione GKE che supporti TPU Trillium (v6e). Per maggiori informazioni, consulta Verificare la disponibilità delle TPU in GKE.
   • REGION_NAME: la regione in cui si trovano il cluster GKE, il bucket Cloud Storage e i nodi TPU. La regione contiene zone in cui sono disponibili i tipi di macchine TPU v5e (ad esempio us-west1, us-west4, us-central1, us-east1, us-east5 o europe-west4).
   • (Solo cluster Standard) ZONE: la zona in cui sono disponibili le risorse TPU (ad esempio, us-west4-a). Per i cluster Autopilot, non è necessario specificare la zona, ma solo la regione.

3. Clona il repository di esempio e apri la directory del tutorial:

   git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
   cd kubernetes-engine-samples/ai-ml/maxdiffusion-tpu
   WORK_DIR=$(pwd)
   gcloud artifacts repositories create gke-llm --repository-format=docker --location=$REGION
   gcloud auth configure-docker $REGION-docker.pkg.dev
   

Creare e configurare risorse Google Cloud

Segui queste istruzioni per creare le risorse richieste.

Nota: potresti dover creare una prenotazione di capacità per l'utilizzo di alcuni acceleratori. Per scoprire come prenotare e utilizzare le risorse prenotate, consulta Utilizzo delle risorse di zona prenotate.

Crea un cluster GKE

Puoi gestire SDXL su TPU in un cluster GKE Autopilot o Standard. Ti consigliamo di utilizzare un cluster Autopilot per un'esperienza Kubernetes completamente gestita. Per scegliere la modalità operativa GKE più adatta ai tuoi workload, consulta Scegliere una modalità operativa GKE.

Autopilot

1. In Cloud Shell, esegui questo comando:

   gcloud container clusters create-auto ${CLUSTER_NAME} \
     --project=${PROJECT_ID} \
     --location=${REGION} \
     --release-channel=rapid \
     --cluster-version=${CLUSTER_VERSION}
   

   GKE crea un cluster Autopilot con nodi CPU e TPU come richiesto dai carichi di lavoro di cui è stato eseguito il deployment.

2. Configura kubectl per comunicare con il cluster:

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

Standard

1. Crea un cluster GKE Standard regionale che utilizzi Workload Identity Federation for GKE.

   gcloud container clusters create ${CLUSTER_NAME} \
       --enable-ip-alias \
       --machine-type=n2-standard-4 \
       --num-nodes=2 \
       --workload-pool=${PROJECT_ID}.svc.id.goog \
       --location=${REGION}
   

   La creazione del cluster potrebbe richiedere diversi minuti.

2. Esegui questo comando per creare un pool di nodi per il tuo cluster:

   gcloud container node-pools create maxdiffusion-tpu-nodepool \
     --cluster=${CLUSTER_NAME} \
     --machine-type=ct5lp-hightpu-1t \
     --num-nodes=1 \
     --location=${REGION} \
     --node-locations=${ZONE} \
     --spot
   

   GKE crea un pool di nodi TPU v5e con una topologia 1x1 e un nodo.

   Per creare node pool con topologie diverse, scopri come pianificare la configurazione TPU. Assicurati di aggiornare i valori di esempio in questo tutorial, ad esempio cloud.google.com/gke-tpu-topology e google.com/tpu.

3. Configura kubectl per comunicare con il cluster:

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

Crea il container di inferenza SDXL

Segui queste istruzioni per creare un'immagine container per il server di inferenza SDXL.

1. Apri il file manifest build/server/cloudbuild.yaml:

   steps:
   - name: 'gcr.io/cloud-builders/docker'
     args: [ 'build', '-t', '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion:latest', '.' ]
   images:
   - '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion:latest'

2. Esegui la build e crea l'immagine container di inferenza.

   cd $WORK_DIR/build/server
   gcloud builds submit . --region=$REGION
   

   L'output contiene il percorso dell'immagine container.

Esegui il deployment del server di inferenza SDXL

In questa sezione, esegui il deployment del server di inferenza SDXL. Per eseguire il deployment del server, questo tutorial utilizza un deployment Kubernetes. Un deployment è un oggetto API Kubernetes che ti consente di eseguire più repliche di pod distribuite tra i nodi di un cluster.

1. Esplora il manifest di serve_sdxl_v5e.yaml.

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: stable-diffusion-deployment
   spec:
     selector:
       matchLabels:
         app: max-diffusion-server
     replicas: 1  # number of nodes in node-pool
     template:
       metadata:
         labels:
           app: max-diffusion-server
       spec:
         nodeSelector:
           cloud.google.com/gke-tpu-topology: 1x1 #  target topology
           cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
           #cloud.google.com/gke-spot: "true"
         volumes:
         - name: dshm
           emptyDir:
                 medium: Memory
         containers:
         - name: serve-stable-diffusion
           image: REGION-docker.pkg.dev/PROJECT_ID/gke-llm/max-diffusion:latest
           env:
           - name: MODEL_NAME
             value: 'stable_diffusion'
           ports:
           - containerPort: 8000
           resources:
             requests:
               google.com/tpu: 1  # TPU chip request
             limits:
               google.com/tpu: 1  # TPU chip request
           volumeMounts:
               - mountPath: /dev/shm
                 name: dshm
   
   ---
   apiVersion: v1
   kind: Service
   metadata:
     name: max-diffusion-server
     labels:
       app: max-diffusion-server
   spec:
     type: ClusterIP
     ports:
       - port: 8000
         targetPort: 8000
         name: http-max-diffusion-server
         protocol: TCP
     selector:
       app: max-diffusion-server

2. Aggiorna l'ID progetto nel file manifest.

   cd $WORK_DIR
   perl -pi -e 's|PROJECT_ID|PROJECT_ID|g' serve_sdxl_v5e.yaml
   perl -pi -e 's|REGION|REGION_NAME|g' serve_sdxl_v5e.yaml
   

3. Applica il manifest:

   kubectl apply -f serve_sdxl_v5e.yaml
   

   L'output è simile al seguente:

   deployment.apps/max-diffusion-server created
   

4. Verifica lo stato del modello:

   kubectl get deploy --watch
   

   L'output è simile al seguente:

   NAME                          READY   UP-TO-DATE   AVAILABLE   AGE
   stable-diffusion-deployment   1/1     1            1           8m21s
   

5. Recupera l'indirizzo ClusterIP:

   kubectl get service max-diffusion-server
   

   L'output contiene un campo ClusterIP. Prendi nota del valore di CLUSTER-IP.

6. Convalida il deployment:

    export ClusterIP=CLUSTER_IP
    kubectl run curl --image=curlimages/curl \
       -it --rm --restart=Never \
       -- "$ClusterIP:8000"
   

   Sostituisci CLUSTER_IP con il valore CLUSTER-IP che hai annotato in precedenza. L'output è simile al seguente:

   {"message":"Hello world! From FastAPI running on Uvicorn with Gunicorn."}
   pod "curl" deleted
   

7. Visualizza i log del deployment:

   kubectl logs -l app=max-diffusion-server
   

   Al termine del deployment, l'output è simile al seguente:

   2024-06-12 15:45:45,459 [INFO] __main__: replicate params:
   2024-06-12 15:45:46,175 [INFO] __main__: start initialized compiling
   2024-06-12 15:45:46,175 [INFO] __main__: Compiling ...
   2024-06-12 15:45:46,175 [INFO] __main__: aot compiling:
   2024-06-12 15:45:46,176 [INFO] __main__: tokenize prompts:2024-06-12 15:48:49,093 [INFO] __main__: Compiled in 182.91802048683167
   INFO:     Started server process [1]
   INFO:     Waiting for application startup.
   INFO:     Application startup complete.
   

Esegui il deployment del client dell'app web

In questa sezione, esegui il deployment del client dell'app web per pubblicare il modello SDXL.

1. Esplora il manifest di build/webapp/cloudbuild.yaml.

   steps:
   - name: 'gcr.io/cloud-builders/docker'
     args: [ 'build', '-t', '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion-web:latest', '.' ]
   images:
   - '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion-web:latest'

2. Esegui la build e crea l'immagine container del client nella directory build/webapp.

   cd $WORK_DIR/build/webapp
   gcloud builds submit . --region=$REGION
   

   L'output contiene il percorso dell'immagine container.

3. Apri il file manifest serve_sdxl_client.yaml:

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: max-diffusion-client
   spec:
     selector:
       matchLabels:
         app: max-diffusion-client
     template:
       metadata:
         labels:
           app: max-diffusion-client
       spec:
         containers:
         - name: webclient
           image: REGION-docker.pkg.dev/PROJECT_ID/gke-llm/max-diffusion-web:latest
           env:
             - name: SERVER_URL
               value: "http://ClusterIP:8000"
           resources:
             requests:
               memory: "128Mi"
               cpu: "250m"
             limits:
               memory: "256Mi"
               cpu: "500m"
           ports:
           - containerPort: 5000
   ---
   apiVersion: v1
   kind: Service
   metadata:
     name: max-diffusion-client-service
   spec:
     type: LoadBalancer
     selector:
       app: max-diffusion-client
     ports:
     - port: 8080
       targetPort: 5000

4. Modifica l'ID progetto nel manifest:

   cd $WORK_DIR
   perl -pi -e 's|PROJECT_ID|PROJECT_ID|g' serve_sdxl_client.yaml
   perl -pi -e 's|ClusterIP|CLUSTER_IP|g' serve_sdxl_client.yaml
   perl -pi -e 's|REGION|REGION_NAME|g' serve_sdxl_client.yaml
   

5. Applica il manifest:

   kubectl apply -f serve_sdxl_client.yaml
   

6. Recupera l'indirizzo IP LoadBalancer:

   kubectl get service max-diffusion-client-service
   

   L'output contiene un campo LoadBalancer. Prendi nota del valore di EXTERNAL-IP.

Operazione riuscita: hai eseguito il deployment del modello SDXL utilizzando MaxDiffusion e TPU v5e in GKE.

Interagisci con il modello utilizzando la pagina web

1. Accesso al seguente URL da un browser web:

   http://EXTERNAL_IP:8080
   

   Sostituisci EXTERNAL_IP con il valore EXTERNAL_IP che hai annotato in precedenza.

2. Interagisci con SDXL utilizzando l'interfaccia di chat. Aggiungi un prompt e fai clic su Invia. Ad esempio:

   Create a detailed image of a fictional historical site, capturing its unique architecture and cultural significance
   

L'output è un'immagine generata dal modello simile al seguente esempio:

Esegui la pulizia

Per evitare che al tuo Account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.

Elimina il progetto

Attenzione: l'eliminazione di un progetto ha i seguenti effetti:

• L'intero contenuto del progetto viene eliminato. Se hai utilizzato un progetto esistente per, le attività descritte in questo documento, eliminerai anche tutto il lavoro che hai svolto nel progetto.
• Gli ID progetto personalizzati non sono più disponibili. Quando hai creato questo progetto, potresti aver creato un ID progetto personalizzato che vuoi utilizzare in futuro. Per conservare gli URL che utilizzano l'ID progetto, ad esempio un URL appspot.com, elimina le risorse selezionate all'interno del progetto anziché eliminare l'intero progetto.

Se intendi esplorare più architetture, tutorial o guide rapide, puoi riutilizzare i progetti ed evitare così di superare i limiti di quota.

1. Nella console Google Cloud , vai alla pagina Gestisci risorse.

   Vai a Gestisci risorse

2. Nell'elenco dei progetti, seleziona quello che vuoi eliminare, quindi fai clic su Elimina.
3. Nella finestra di dialogo, digita l'ID del progetto e fai clic su Chiudi per eliminare il progetto.

Elimina le singole risorse

Mantieni il progetto ed elimina le singole risorse, come descritto nella sezione seguente. Esegui questi comandi e segui le istruzioni:

gcloud container clusters delete ${CLUSTER_NAME} --location=${REGION}

Passaggi successivi

• Configura il tutorial con altre topologie TPU. Per scoprire di più su altre topologie TPU, consulta la sezione Pianificare la configurazione TPU.
• Esplora il codice campione del server di inferenza MaxDiffusion nel repository di esempio che hai clonato in questo tutorial.
• Scopri di più sulle TPU in GKE.
• Esplora il repository GitHub di JetStream.
• Esplora Vertex AI Model Garden.