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


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 imparerai a scarica il modello da Hugging Face ed eseguine il deployment su Autopilot o Standard utilizzando un container che esegue MaxDiffusion.

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

Sfondo

Mediante la gestione di SDXL mediante TPU su GKE con MaxDiffusion, puoi creare una soluzione di pubblicazione solida e pronta per la produzione con tutte i vantaggi di Kubernetes gestito, ad esempio efficienza in termini di costi, scalabilità e maggiore disponibilità. Questa sezione descrive le tecnologie chiave usate 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'IA generativa, puoi utilizzare i modelli LDM per generare immagini di alta qualità dalle descrizioni di testo. Gli LDM sono utili per applicazioni come la ricerca di immagini e la creazione di didascalie.

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

Per saperne di più, consulta il repository Generative Models by Stability AI e il documento SDXL.

TPU

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

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

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

Questo tutorial illustra la pubblicazione 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 pubblicazione di prompt con bassa latenza. In questo guida, il modello utilizza un chip TPU v5e con una topologia 1x1.

MaxDiffusion

MaxDiffusion è una raccolta di riferimenti implementazioni, scritte in Python e Jax, di vari modelli di diffusione latente 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 scoprire di più, consulta il repository MaxDiffusion.

Obiettivi

Questo tutorial è rivolto ai clienti di IA generativa che utilizzano JAX, agli utenti nuovi o esistenti di SDXL e a qualsiasi esperto di ML, MLOps (DevOps) o amministratore della piattaforma interessato a utilizzare le funzionalità di orchestrazione dei contenitori Kubernetes per la pubblicazione 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 un'interazione con il modello tramite un'app web.

Architettura

Questa sezione descrive l'architettura GKE utilizzata in questo tutorial. L'architettura è composta da un sistema GKE Autopilot Cluster standard che esegue il provisioning di TPU e ospita i componenti MaxDiffusion. GKE utilizza questi componenti per eseguire il deployment e il servizio dei modelli.

Il seguente diagramma mostra i componenti di questa architettura:

Architettura di esempio per la pubblicazione di MaxDiffusion con TPU v5e su GKE.

Questa architettura include i seguenti componenti:

  • Un cluster GKE Autopilot o Standard a livello di regione.
  • Un pool di nodi di sezione TPU con un solo host che ospita il modello SDXL nel deployment di MaxDiffusion.
  • Il componente del servizio 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 pubblicazione del modello al servizio ClusterIP.

Prima di iniziare

  • Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

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

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

  • Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin

    Check for the roles

    1. In the Google Cloud console, go to the IAM page.

      Go to IAM
    2. Select the project.
    3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

    4. For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.

    Grant the roles

    1. In the Google Cloud console, go to the IAM page.

      Vai a IAM
    2. Seleziona il progetto.
    3. Fai clic su Concedi l'accesso.
    4. Nel campo Nuove entità, inserisci il tuo identificatore utente. In genere si tratta dell'indirizzo email di un Account Google.

    5. Nell'elenco Seleziona un ruolo, seleziona un ruolo.
    6. Per concedere altri ruoli, fai clic su Aggiungi un altro ruolo e aggiungi ogni altro ruolo.
    7. Fai clic su Salva.
  • Assicurati di avere una quota sufficiente per TPU v5e PodSlice Lite chip. In questo tutorial vengono utilizzate le istanze on demand.

prepara l'ambiente

In questo tutorial utilizzerai Cloud Shell per gestire le risorse ospitate in Google Cloud. Cloud Shell è preinstallato con il software necessario per questo tutorial, inclusi kubectl e con gcloud CLI.

Per configurare l'ambiente con Cloud Shell, segui questi passaggi:

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

  2. Imposta le variabili di ambiente predefinite:

    gcloud config set project PROJECT_ID
    export PROJECT_ID=$(gcloud config get project)
    export CLUSTER_NAME=CLUSTER_NAME
    export REGION=REGION_NAME
    export ZONE=ZONE
    

    Sostituisci i seguenti valori:

    • PROJECT_ID: il tuo ID progetto Google Cloud.
    • CLUSTER_NAME: il nome del tuo cluster GKE.
    • REGION_NAME: la regione in cui GKE cluster, il bucket Cloud Storage e i nodi TPU. La regione contiene zone in cui sono disponibili 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
    

Creazione e configurazione delle risorse Google Cloud

Segui queste istruzioni per creare le risorse richieste.

Crea un cluster GKE

Puoi pubblicare 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 di GKE più adatta ai tuoi carichi di lavoro, consulta Scegliere una modalità operativa GKE.

Autopilot

  1. In Cloud Shell, esegui questo comando:

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

    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 in modo che comunichi con il tuo cluster:

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

Standard

  1. Crea un cluster GKE Standard regionale che utilizza 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 il comando seguente per creare un pool di nodi per il cluster:

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

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

    Per creare pool di nodi con topologie diverse, scopri come Pianifica la configurazione di 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 in modo che comunichi con il tuo cluster:

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

Crea il contenitore di inferenza SDXL

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

  1. Apri il manifest 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

  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
    sed -i "s|PROJECT_ID|$PROJECT_ID|g" serve_sdxl_v5e.yaml
    sed -i "s|REGION|$REGION|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:

    watch kubectl get deploy
    

    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 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 di cui hai parlato 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 della web app 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 del container 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
    sed -i "s|PROJECT_ID|$PROJECT_ID|g" serve_sdxl_client.yaml
    sed -i "s|ClusterIP|$ClusterIP|g" serve_sdxl_client.yaml
    sed -i "s|REGION|$REGION|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 EXTERNAL-IP.

Interagire con il modello utilizzando la pagina web

  1. Accedi al seguente URL da un browser web:

    http://EXTERNAL_IP:8080
    

    Sostituisci EXTERNAL_IP con il valore EXTERNAL_IP che di cui hai parlato in precedenza.

  2. Interagisci con SDXL utilizzando l'interfaccia di chat. Aggiungi un prompt 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 all'esempio seguente:

Immagine generata da SDXL

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

  1. In the Google Cloud console, go to the Manage resources page.

    Go to Manage resources

  2. In the project list, select the project that you want to delete, and then click Delete.
  3. In the dialog, type the project ID, and then click Shut down to delete the project.

Elimina le singole risorse

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

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

Passaggi successivi