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

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 executa 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 através de 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 MDAs 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, novos ou utilizadores existentes do SDXL, e a 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:

Exemplo de arquitetura de publicação do MaxDiffusion com a TPU v5e no GKE.

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 serviço de modelos para o serviço ClusterIP.

Antes de começar

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

    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 identificador de um utilizador num Workload Identity Pool. Para ver detalhes, consulte Represente utilizadores do pool de força de trabalho em políticas de IAM ou contacte o seu administrador.

    5. Clique em Selecionar uma função e, de seguida, pesquise a 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, usa 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 Ícone de ativação do Cloud Shell 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 do 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.

      Crie um cluster do GKE

      Pode publicar SDXLs em TPUs num cluster padrão ou no GKE Autopilot. 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 seu 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}. \
    --location=${REGION}

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

      2. Execute o seguinte comando para criar um conjunto 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 node pool de TPUs v5e com uma topologia 1x1 e um nó.

        Para criar node pools 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 seu cluster:

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

      Crie o contentor de inferência 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ídos 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: