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

Faça login na sua conta do Google Cloud . Se você começou a usar o Google Cloud, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

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

• Verifique se você tem os seguintes papéis no projeto: roles/container.admin, roles/iam.serviceAccountAdmin, roles/artifactregistry.admin, roles/cloudbuild.builds.editor

  Verificar os papéis

  1. No console do Google Cloud , acesse a página IAM.

     Acessar IAM
  2. Selecione o projeto.

  3. Na coluna Principal, encontre todas as linhas que identificam você ou um grupo no qual você está incluído. Para saber em quais grupos você está incluído, entre em contato com o administrador.

  4. Em todas as linhas que especificam ou incluem você, verifique a coluna Papel para ver se a lista de papéis inclui os papéis necessários.

  Conceder os papéis

  1. No console do Google Cloud , acesse a página IAM.

     Acessar IAM
  2. Selecione o projeto.
  3. Clique em Conceder acesso.

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

  5. Clique em Selecionar um papel e pesquise o papel.
  6. Para conceder outros papéis, adicione-os clicando em Adicionar outro papel.
  7. Clique em Salvar.

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

Prepare o ambiente

Neste tutorial, você vai 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 CLI gcloud.

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 doGoogle Cloud . Isso inicia uma sessão no painel inferior do console Google Cloud .

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

   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
   

   Substitua os seguintes valores:

   • PROJECT_ID: o ID do projeto do Google Cloud .
   • CLUSTER_NAME: o nome do cluster do GKE.
   • CLUSTER_VERSION : a versão do GKE. Você precisa especificar uma versão do GKE que seja compatível com a TPU Trillium (v6e). Para mais informações, consulte Validar a disponibilidade da TPU no 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 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} \
     --location=${REGION} \
     --release-channel=rapid \
     --cluster-version=${CLUSTER_VERSION}
   

   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 \
     --location=${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 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 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

Nesta seção, você vai implantar o servidor de inferência SDXL. Para implantar o servidor, este tutorial usa uma implantação do Kubernetes. Uma implantação é um objeto da API Kubernetes que permite executar várias réplicas de pods distribuídos entre os nós de um cluster.

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
   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 será assim:

   deployment.apps/max-diffusion-server created
   

4. Verifique o status do modelo:

   kubectl get deploy --watch
   

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

   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.

Excluir o projeto

Cuidado: excluir um projeto tem os seguintes efeitos:

• Tudo no projeto é excluído. Se você tiver usado um projeto atual para as tarefas neste documento, a exclusão dele incluirá a exclusão de quaisquer outros trabalhos feitos nele.
• Os IDs do projeto personalizados são perdidos. Ao criar o projeto, você pode ter criado um ID 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 explorar várias arquiteturas, tutoriais e guias de início rápido, a reutilização de projetos pode evitar que você exceda limites da cota.

1. No console Google Cloud , acesse a página Gerenciar recursos.

   Acessar "Gerenciar recursos"

2. Na lista de projetos, selecione o projeto que você quer excluir e clique em Excluir .
3. Na caixa de diálogo, digite o ID do projeto e clique em Encerrar para excluí-lo.

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} --location=${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.