Apresente o Stable Diffusion XL (SDXL) com TPUs no GKE com o MaxDiffusion

Autopilot Padrão

Este tutorial mostra como publicar um modelo de geração de imagens SDXL usando unidades de processamento tensor (TPUs) no Google Kubernetes Engine (GKE) com o MaxDiffusion. Neste tutorial, vai transferir o modelo do Hugging Face e implementá-lo num cluster Autopilot ou Standard usando um contentor que execute o MaxDiffusion.

Este guia é um bom ponto de partida se precisar do controlo detalhado, da personalização, da escalabilidade, da resiliência, da portabilidade e da rentabilidade do Kubernetes gerido ao implementar e publicar as suas cargas de trabalho de IA/ML. Se precisar de uma plataforma de IA gerida unificada para criar e apresentar rapidamente modelos de ML de forma rentável, recomendamos que experimente a nossa solução de implementação do Vertex AI.

Contexto

Ao publicar o SDXL com TPUs no GKE com o MaxDiffusion, pode criar uma solução de publicação robusta e pronta para produção com todas as vantagens do Kubernetes gerido, incluindo rentabilidade, escalabilidade e maior disponibilidade. Esta secção descreve as tecnologias principais usadas neste tutorial.

Stable Diffusion XL (SDXL)

O Stable Diffusion XL (SDXL) é um tipo de modelo de difusão latente (LDM) suportado pelo MaxDiffusion para inferência. Para a IA generativa, pode usar os MDAs para gerar imagens de alta qualidade a partir de descrições de texto. Os MDGs são úteis para aplicações como a pesquisa de imagens e a geração de legendas de imagens.

O SDXL suporta a inferência de anotações de fragmentação com um ou vários anfitriões. Isto permite que o SDXL seja preparado e executado em várias máquinas, o que pode melhorar a eficiência.

Para saber mais, consulte o repositório de modelos generativos da Stability AI e o artigo SDXL.

TPUs

As TPUs são circuitos integrados específicos da aplicação (ASICs) desenvolvidos pela Google e usados para acelerar a aprendizagem automática e os modelos de IA criados com frameworks como o TensorFlow, o PyTorch e o JAX.

Antes de usar as TPUs no GKE, recomendamos que conclua o seguinte percurso de aprendizagem:

1. Saiba mais sobre a disponibilidade da versão atual da TPU com a arquitetura do sistema Cloud TPU.
2. Saiba mais sobre as TPUs no GKE.

Este tutorial aborda a publicação do modelo SDXL. O GKE implementa o modelo em nós da TPU v5e de anfitrião único com topologias de TPU configuradas com base nos requisitos do modelo para publicar comandos com baixa latência. Neste guia, o modelo usa um chip TPU v5e com uma topologia 1x1.

MaxDiffusion

O MaxDiffusion é uma coleção de implementações de referência, escritas em Python e Jax, de vários modelos de difusão latente que são executados em dispositivos XLA, incluindo TPUs e GPUs. O MaxDiffusion é um ponto de partida para projetos de difusão para investigação e produção.

Para saber mais, consulte o repositório MaxDiffusion.

Objetivos

Este tutorial destina-se a clientes de IA generativa que usam o JAX, utilizadores novos ou existentes do SDXL, e engenheiros de AA, engenheiros de MLOps (DevOps) ou administradores de plataformas que tenham interesse em usar capacidades de orquestração de contentores do Kubernetes para publicar MDIs.

Este tutorial abrange os seguintes passos:

1. Crie um cluster do GKE Autopilot ou Standard com a topologia de TPU recomendada, com base nas caraterísticas do modelo.
2. Crie uma imagem de contentor de inferência do SDXL.
3. Implemente o servidor de inferência SDXL no GKE.
4. Publicar uma interação com o modelo através de uma app Web.

Arquitetura

Esta secção descreve a arquitetura do GKE usada neste tutorial. A arquitetura consiste num cluster do GKE Autopilot ou Standard que aprovisiona as TPUs e aloja os componentes do MaxDiffusion. O GKE usa estes componentes para implementar e publicar os modelos.

O diagrama seguinte mostra os componentes desta arquitetura:

Esta arquitetura inclui os seguintes componentes:

• Um cluster regional do GKE Autopilot ou Standard.
• Um único conjunto de nós de fatia de TPU de anfitrião único que aloja o modelo SDXL na implementação do MaxDiffusion.
• O componente Service com um balanceador de carga do tipo ClusterIP. Este serviço distribui o tráfego de entrada por todas as réplicas MaxDiffusion HTTP.
• O servidor WebApp HTTP com um serviço LoadBalancer externo que distribui o tráfego de entrada e redireciona o tráfego de publicação de modelos para o serviço ClusterIP.

Antes de começar

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.

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 (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 API.

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 API

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 (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 API.

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 API

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

  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 column 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.

     Aceder ao IAM
  2. Selecione o projeto.
  3. Clique em Conceder acesso.

  4. No campo Novos responsáveis, introduza o identificador do utilizador. Normalmente, este é o endereço de email de uma Conta Google.

  5. Na lista Selecionar uma função, selecione uma função.
  6. Para conceder funções adicionais, clique em Adicionar outra função e adicione cada função adicional.
  7. Clique em Guardar.
  • Verifique se tem quota suficiente para chips TPU v5e PodSlice Lite. Neste tutorial, vai usar instâncias a pedido.

  Prepare o ambiente

  Neste tutorial, vai usar o Cloud Shell para gerir recursos alojados no Google Cloud. O Cloud Shell vem pré-instalado com o software de que precisa para este tutorial, incluindo o kubectl e a CLI gcloud.

  Para configurar o seu ambiente com o Cloud Shell, siga estes passos:

  1. Na Google Cloud consola, inicie uma sessão do Cloud Shell clicando em Ativar Cloud Shell na Google Cloud consola. Esta ação inicia uma sessão no painel inferior da consola. Google Cloud

  2. Defina as variáveis de ambiente predefinidas:

     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 REGION=REGION_NAME
     export ZONE=ZONE
     

     Substitua os seguintes valores:

     • PROJECT_ID: o seu Google Cloud ID do projeto.
     • CLUSTER_NAME: o nome do seu cluster do GKE.
     • REGION_NAME: a região onde o cluster do GKE, o contentor do Cloud Storage e os nós da TPU estão localizados. A região contém zonas onde os tipos de máquinas TPU v5e estão disponíveis (por exemplo, us-west1, us-west4, us-central1, us-east1, us-east5 ou europe-west4).
     • (Apenas para clusters padrão) ZONE: a zona onde os recursos de TPU estão disponíveis (por exemplo, us-west4-a). Para clusters do Autopilot, não precisa de especificar a zona, apenas a região.

  3. Clone o repositório de exemplo e abra o diretório do 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
     

  Crie e configure Google Cloud recursos

  Siga estas instruções para criar os recursos necessários.

  Nota: pode ter de criar uma reserva de capacidade para a utilização de alguns aceleradores. Para saber como reservar e consumir recursos reservados, consulte o artigo Consumir recursos zonais reservados.

  Crie um cluster do GKE

  Pode publicar SDXLs em TPUs num cluster do GKE Autopilot ou Standard. Recomendamos que use um cluster do Autopilot para uma experiência do Kubernetes totalmente gerida. Para escolher o modo de funcionamento do GKE mais adequado às suas cargas de trabalho, consulte o artigo Escolha um modo de funcionamento do GKE.

  Piloto automático

  1. No Cloud Shell, execute o seguinte comando:

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

     O GKE cria um cluster do Autopilot com nós de CPU e TPU, conforme solicitado pelas cargas de trabalho implementadas.

  2. Configure kubectl para comunicar com o cluster:

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

  Standard

  1. Crie um cluster padrão do GKE regional que use a Workload Identity Federation para o 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}
     

     A criação do cluster pode demorar vários minutos.

  2. Execute o seguinte comando para criar um node pool para o 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
     

     O GKE cria um node pool de TPUs v5e com uma topologia 1x1 e um nó.

     Para criar pools de nós com topologias diferentes, saiba como planear a configuração da TPU. Certifique-se de que atualiza os valores de exemplo neste tutorial, como cloud.google.com/gke-tpu-topology e google.com/tpu.

  3. Configure kubectl para comunicar com o cluster:

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

  Crie o contentor de inferência do SDXL

  Siga estas instruções para criar uma imagem de contentor para o servidor de inferência SDXL.

  1. Abra o manifesto 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. Execute a compilação e crie a imagem do contentor de inferência.

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

     A saída contém o caminho da imagem do contentor.

  Implemente o servidor de inferência SDXL

  Nesta secção, implementa o servidor de inferência SDXL. Para implementar o servidor, este tutorial usa uma implementação do Kubernetes. Uma implementação é um objeto da API Kubernetes que lhe permite executar várias réplicas de pods distribuídas entre os nós num cluster.

  1. Explore o serve_sdxl_v5e.yaml manifesto.

     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. Atualize o ID do projeto no manifesto.

     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. Aplique o manifesto:

     kubectl apply -f serve_sdxl_v5e.yaml
     

     O resultado é semelhante ao seguinte:

     deployment.apps/max-diffusion-server created
     

  4. Valide o estado do modelo:

     kubectl get deploy --watch
     

     O resultado é semelhante ao seguinte:

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

  5. Recupere o endereço ClusterIP:

     kubectl get service max-diffusion-server
     

     O resultado contém um campo ClusterIP. Tome nota do valor CLUSTER-IP.

  6. Valide a implementação:

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

     Substitua CLUSTER_IP pelo valor CLUSTER-IP que anotou anteriormente. O resultado é semelhante ao seguinte:

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

  7. Veja os registos da implementação:

     kubectl logs -l app=max-diffusion-server
     

     Quando a implementação terminar, o resultado é semelhante ao seguinte:

     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.
     

  Implemente o cliente da app Web

  Nesta secção, implementa o cliente da app Web para publicar o modelo SDXL.

  1. Explore o build/webapp/cloudbuild.yaml manifesto.

     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. Execute a compilação e crie a imagem do contentor do cliente no diretório build/webapp.

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

     A saída contém o caminho da imagem do contentor.

  3. Abra o manifesto 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. Edite o ID do projeto no manifesto:

     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. Aplique o manifesto:

     kubectl apply -f serve_sdxl_client.yaml
     

  6. Recupere o endereço IP LoadBalancer:

     kubectl get service max-diffusion-client-service
     

     O resultado contém um campo LoadBalancer. Tome nota do valor EXTERNAL-IP.

  Êxito: implementou com êxito o modelo SDXL através do MaxDiffusion e da TPU v5e no GKE.

  Interaja com o modelo através da página Web

  1. Acesso ao seguinte URL a partir de um navegador de Internet:

     http://EXTERNAL_IP:8080
     

     Substitua EXTERNAL_IP pelo valor EXTERNAL_IP que anotou anteriormente.

  2. Interaja com o SDXL através da interface de chat. Adicione um comando e clique em Enviar. Por exemplo:

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

  O resultado é uma imagem gerada pelo modelo semelhante ao seguinte exemplo:

  

  Limpar

  Para evitar incorrer em custos na sua conta do Google Cloud pelos recursos usados neste tutorial, elimine o projeto que contém os recursos ou mantenha o projeto e elimine os recursos individuais.

  Elimine o projeto

  Atenção: a eliminação de um projeto tem os seguintes efeitos:

  • Tudo no projeto é eliminado. Se usou um projeto existente para as tarefas neste documento, quando o elimina, também elimina qualquer outro trabalho que tenha feito no projeto.
  • Os IDs dos projetos personalizados são perdidos. Quando criou este projeto, pode ter criado um ID do projeto personalizado que quer usar no futuro. Para preservar os URLs que usam o ID do projeto, como um URL appspot.com, elimine os recursos selecionados no projeto em vez de eliminar o projeto inteiro.

  Se planeia explorar várias arquiteturas, tutoriais ou inícios rápidos, a reutilização de projetos pode ajudar a evitar exceder os limites de quota de projetos.

  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.

  Elimine os recursos individuais

  Manter o projeto e eliminar os recursos individuais, conforme descrito na secção seguinte. Execute os seguintes comandos e siga as instruções:

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

  O que se segue?

  • Configure o tutorial com outras topologias de TPU. Para saber mais sobre outras topologias de TPU, consulte o artigo Planeie a configuração da TPU.
  • Explore o código de exemplo do servidor de inferência MaxDiffusion no repositório de exemplo que clonou neste tutorial.
  • Saiba mais acerca das TPUs no GKE.
  • Explore o repositório do GitHub do JetStream.
  • Explore o Vertex AI Model Garden.