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

Autopilot Standard

---

Neste tutorial, mostramos como disponibilizar Modelo de geração de imagem SDXL usando Unidades de Processamento de Tensor (TPUs) no Google Kubernetes Engine (GKE) com MaxDiffusion (link em inglês). Neste tutorial, baixe o modelo do Hugging Face e implante-o no Autopilot ou Padrão cluster usando um contêiner que executa o MaxDiffusion.

Este guia é um bom ponto de partida se você precisar do controle granular, da customização, da escalonabilidade, da resiliência, da portabilidade e da economia do Kubernetes gerenciado ao implantar e disponibilizar suas cargas de trabalho de IA/ML. Se você precisa de uma plataforma de IA gerenciada unificada para criar e disponibilizar modelos de ML rapidamente e de maneira econômica, recomendamos testar nossa solução de implantação da Vertex AI.

Contexto

Ao disponibilizar SDXL usando TPUs no GKE com o MaxDiffusion, é possível criar uma solução de disponibilização robusta e pronta para produção com todos Os benefícios do Kubernetes gerenciado, incluindo economia, escalonabilidade e maior disponibilidade. Esta seção descreve as principais tecnologias usadas neste tutorial.

Stable Diffusion XL (SDXL)

O Stable Diffusion XL (SDXL) é um tipo de modelo de difusão latente (LDM) compatível com o MaxDiffusion para inferência. Na IA generativa, é possível usar LDMs para gerar imagens de alta qualidade a partir de descrições de texto. Os LDMs são úteis para aplicativos, como pesquisa e legendagem de imagens.

A SDXL oferece suporte à inferência de um ou vários hosts com anotações de fragmentação. Isso permite treinar e executar a SDXL em várias máquinas, o que pode melhorar a eficiência.

Para saber mais, consulte a Repositório de modelos generativos da Stability AI e o Papel SDXL.

TPUs

TPUs são circuitos integrados de aplicação específica (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Eles são usados para acelerar modelos de machine learning e de IA criados com o uso de frameworks comoTensorFlow , PyTorch eJAX.

Antes de usar TPUs no GKE, recomendamos que você conclua o seguinte programa de aprendizado:

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

Este tutorial aborda a exibição do modelo SDXL. O GKE implanta o modelo em nós TPU de host único com topologias de TPU configuradas com base nos requisitos do modelo para exibir prompts com baixa latência. Neste o modelo usa um chip TPU v5e com topologia 1x1.

MaxDiffusion

O MaxDiffusion é uma coleção de objetos implementações escritas em Python e Jax de vários modelos de difusão latente executados em dispositivos XLA, incluindo TPUs e GPUs. O MaxDiffusion é um ponto de partida muito importante para projetos de difusão para pesquisa e produção.

Para saber mais, consulte o repositório MaxDiffusion (em inglês).

Objetivos

Este tutorial é destinado a clientes de IA generativa que usam o JAX, usuários novos ou atuais do SDXL, engenheiros de ML, engenheiros de MLOps (DevOps) ou administradores de plataformas interessados em usar os recursos de orquestração de contêineres do Kubernetes para veiculação de LLMs.

Este tutorial inclui as etapas a seguir:

1. Crie um cluster do Autopilot ou do GKE Standard com a topologia de TPU recomendada com base nas características do modelo.
2. Crie uma imagem de contêiner de inferência SDXL.
3. Implantar o servidor de inferência SDXL no GKE.
4. Exibir e interagir com o modelo usando um app da Web.

Arquitetura

Nesta seção, descrevemos a arquitetura do GKE usada neste tutorial. A arquitetura consiste em um Autopilot do GKE ou Cluster padrão que provisiona TPUs e hospeda componentes do MaxDiffusion. O GKE usa esses componentes para implantar e disponibilizar os modelos.

O diagrama a seguir mostra os componentes dessa arquitetura:

Essa arquitetura inclui os seguintes componentes:

• Um cluster regional ou Autopilot do GKE.
• Um pool de nós de fração de TPU de host único que hospeda o modelo SDXL na implantação do MaxDiffusion.
• O componente Serviço com um balanceador de carga do tipo ClusterIP. Este serviço distribui o tráfego de entrada para todas as MaxDiffusion HTTP réplicas.
• O servidor WebApp HTTP com um serviço LoadBalancer externo que distribui o tráfego de entrada e redireciona o tráfego que veicula o modelo para ClusterIP Serviço.

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.

  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

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

  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

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

     Acessar o IAM
  2. Selecionar um projeto.
  3. Clique em CONCEDER ACESSO.

  4. No campo Novos principais, insira seu identificador de usuário. Normalmente, é o endereço de e-mail de uma Conta do Google.

  5. Na lista Selecionar um papel, escolha um.
  6. Para conceder outros papéis, clique em Adicionar outro papel e adicione cada papel adicional.
  7. Clique em Salvar.

• Verifique se você tem cota suficiente para TPU v5e PodSlice Lite ícones. Neste tutorial, você usará as instâncias sob demanda.

Prepare o ambiente

Neste tutorial, você usará o Cloud Shell para gerenciar recursos hospedados no Google Cloud. O Cloud Shell vem pré-instalado com o software necessário para este tutorial, incluindo kubectl e gcloud CLI.

Para configurar o ambiente com o Cloud Shell, siga estas etapas:

1. No console do Google Cloud, inicie uma sessão do Cloud Shell clicando em Ativar o Cloud Shell no console do Google Cloud. Isso inicia uma sessão no painel inferior do console do Google Cloud.

2. Defina as variáveis de ambiente padrão:

   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
   

   Substitua os seguintes valores:

   • PROJECT_ID: é o ID do projeto do Google Cloud.
   • CLUSTER_NAME: o nome do cluster do GKE.
   • REGION_NAME: a região em que o cluster do GKE, o bucket do Cloud Storage e os nós da TPU estão localizados. A região contém zonas em que os tipos de máquina TPU v5e estão disponíveis (por exemplo, us-west1, us-west4, us-central1, us-east1, us-east5 ou europe-west4).
   • (Apenas cluster padrão) ZONE: a zona em que os recursos da TPU estão disponíveis, por exemplo, us-west4-a. Para clusters do Autopilot, não é necessário 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
   

Criar e configurar recursos do Google Cloud

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

Observação: pode ser necessário criar uma reserva de capacidade para o uso de alguns aceleradores. Para saber como reservar e consumir recursos reservados, consulte Como consumir recursos por zona reservados.

Crie um cluster do GKE

É possível disponibilizar o SDXL em TPUs em um cluster do GKE Autopilot ou Standard. Recomendamos que você use um cluster do Autopilot para ter uma experiência totalmente gerenciada do Kubernetes. Para escolher o modo de operação do GKE mais adequado para suas cargas de trabalho, consulte Escolher um modo de operação do GKE.

Piloto automático

1. No Cloud Shell, execute este comando:

   gcloud container clusters create-auto ${CLUSTER_NAME} \
     --project=${PROJECT_ID} \
     --region=${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 implantadas.

2. Configure kubectl para se comunicar com o cluster:

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

Padrão

1. Criar um cluster regional do GKE Standard que use a federação de identidade da carga de trabalho 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 levar vários minutos.

2. Execute o seguinte comando para criar um pool de nós para o 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
   

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

   Para criar pools de nós com topologias diferentes, aprenda a Planeje a configuração da TPU. Atualize os valores dos exemplos neste tutorial, como como cloud.google.com/gke-tpu-topology e google.com/tpu.

3. Configure kubectl para se comunicar com o cluster:

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

Criar o contêiner de inferência da SDXL

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

1. Abra o manifesto 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 o build e crie a imagem de contêiner de inferência.

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

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

Implantar o servidor de inferência SDXL

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

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

   kubectl apply -f serve_sdxl_v5e.yaml
   

   O resultado será assim:

   deployment.apps/max-diffusion-server created
   

4. Verifique o status do modelo:

   watch kubectl get deploy
   

   O resultado será assim:

   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
   

   A resposta contém um campo ClusterIP. Anote o valor de CLUSTER-IP

6. Validar a implantaçã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 que você anotou anteriormente. O resultado será assim:

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

7. Veja os registros da implantação:

   kubectl logs -l app=max-diffusion-server
   

   Quando a implantação é concluída, a saída é semelhante a esta:

   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.
   

Implantar o cliente do webapp

Nesta seção, você vai implantar o cliente webapp para exibir a SDXL um modelo de machine learning.

1. Confira o manifesto 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. Execute o build e crie a imagem de contêiner do cliente no build/webapp.

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

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

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

   kubectl apply -f serve_sdxl_client.yaml
   

6. Recupere o endereço IP LoadBalancer:

   kubectl get service max-diffusion-client-service
   

   A resposta contém um campo LoadBalancer. Anote o valor de EXTERNAL-IP.

.

Sucesso: você implantou com sucesso o modelo SDXL usando MaxDiffusion e TPU v5e no GKE.

Interaja com o modelo usando a página da Web

1. Acesse o seguinte URL em um navegador da Web:

   http://EXTERNAL_IP:8080
   

   Substitua o EXTERNAL_IP pelo valor EXTERNAL_IP que que você anotou anteriormente.

2. Interaja com a SDXL usando a interface de chat. Adicione uma solicitação e clique em Enviar. Por exemplo:

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

A saída é uma imagem gerada pelo modelo semelhante ao exemplo a seguir:

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

Exclua o projeto

Cuidado: excluir um projeto tem os seguintes efeitos:

• Tudo no projeto é excluído. Se você tiver usado um projeto existente para as tarefas neste documento, a exclusão dele incluirá a exclusão de quaisquer outros trabalhos feitos no projeto.
• Os IDs do projeto personalizados são perdidos. Ao criar o projeto, você pode ter criado um código do projeto personalizado para ser usado no futuro. Para preservar os URLs que usam o ID do projeto, como um URL appspot.com, exclua recursos específicos do projeto, em vez de excluir o projeto inteiro.

Se você planeja passar por várias arquiteturas, tutoriais ou guias de início rápido, a reutilização de projetos pode evitar que você exceda os limites da cota do projeto.

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.

Excluir recursos individuais

Mantenha o projeto e exclua os recursos individuais, conforme descrito na próxima seção. Execute os comandos a seguir e siga as instruções:

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

A seguir

• Configure o tutorial com outras topologias de TPU. Para saber mais sobre outras topologias de TPU, consulte Planejar a configuração de TPU.
• Explore o código de amostra do servidor de inferência MaxDiffusion no repositório de exemplos que você clonou neste tutorial.
• Saiba mais sobre TPUs no GKE.
• Explore o repositório do JetStream no GitHub (link em inglês).
• Conhecer o Model Garden da Vertex AI.