Disponibilizar um LLM usando TPU Trillium no GKE com o vLLM

Neste tutorial, mostramos como disponibilizar modelos de linguagem grandes (LLMs) usando Unidades de Processamento de Tensor (TPUs) no Google Kubernetes Engine (GKE) com o framework de disponibilização vLLM. Neste tutorial, você vai servir Llama 3.1 70b, usar TPU Trillium e configurar o escalonamento automático horizontal de pods usando métricas do servidor vLLM.

Este documento é um bom ponto de partida se você precisar do controle granular, da escalonabilidade, da resiliência, da portabilidade e da economia do Kubernetes gerenciado ao implantar e disponibilizar cargas de trabalho de IA/ML.

Contexto

Ao usar o TPU Trillium no GKE, é possível implementar uma solução de veiculação robusta e pronta para produção com todos os benefícios do Kubernetes gerenciado, incluindo escalonabilidade eficiente e maior disponibilidade. Esta seção descreve as principais tecnologias usadas neste guia.

TPU Trillium

TPUs são circuitos integrados de aplicação específica (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. As TPUs são usadas para acelerar modelos de machine learning e de IA criados com o uso de frameworks como TensorFlow, PyTorch e JAX. Este tutorial usa a TPU Trillium, que é a TPU de sexta geração do Google.

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

  1. Saiba mais sobre a arquitetura do sistema da TPU Trillium.
  2. Saiba mais sobre TPUs no GKE.

vLLM

O vLLM é um framework de código aberto altamente otimizado para veiculação de LLMs. Ele pode aumentar a capacidade de exibição em GPUs com recursos como:

  • Otimização da implementação do transformador com PagedAttention.
  • Lotes contínuos para melhorar a capacidade de processamento geral de disponibilização.
  • Paralelismo de tensor e exibição distribuída em várias TPUs.

Para saber mais, consulte a documentação do vLLM.

Cloud Storage FUSE

O Cloud Storage FUSE fornece acesso do cluster do GKE ao Cloud Storage para pesos de modelo que residem em buckets de armazenamento de objetos. Neste tutorial, o bucket do Cloud Storage criado vai estar vazio. Quando o vLLM é iniciado, o GKE faz o download do modelo do Hugging Face e armazena os pesos em cache no bucket do Cloud Storage. Na reinicialização do pod ou no aumento de escala da implantação, as cargas de modelo posteriores vão fazer o download de dados em cache do bucket do Cloud Storage, aproveitando downloads paralelos para um desempenho ideal.

Para saber mais, consulte a documentação do driver CSI do Cloud Storage FUSE.

Objetivos

Este tutorial é destinado a engenheiros de MLOps ou DevOps ou administradores de plataforma que querem usar os recursos de orquestração do GKE para exibir LLMs.

Este tutorial inclui as etapas a seguir:

  1. Crie um cluster do GKE com a topologia de TPU Trillium recomendada com base nas características do modelo.
  2. Implante o framework vLLM em um pool de nós no cluster.
  3. Use o framework vLLM para disponibilizar o modelo Llama 3.1 70b usando um balanceador de carga.
  4. Configure o escalonamento automático horizontal de pods usando métricas de servidor vLLM.
  5. Disponibilize o modelo.

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.

    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, roles/iam.securityAdmin, roles/artifactregistry.writer, roles/container.clusterAdmin

    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.

Prepare o ambiente

Nesta seção, você provisiona os recursos necessários para implantar o vLLM e o modelo.

Receber acesso ao modelo

É necessário assinar o contrato de consentimento para usar o Llama 3.1 70b no repositório do Hugging Face.

Gerar um token de acesso

Gere um token do Hugging Face, se você ainda não tiver um:

  1. Clique em Seu perfil > Configurações > Tokens de acesso.
  2. Selecione Novo token.
  3. Especifique um Nome de sua escolha e um Papel de pelo menos Read.
  4. Selecione Gerar um token.

Iniciar o Cloud Shell

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 a gcloud CLI.

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

  1. No console Google Cloud , inicie uma sessão do Cloud Shell clicando em Ícone de ativação do Cloud Shell Ativar o Cloud Shell no console . 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 && \
export PROJECT_ID=$(gcloud config get project) && \
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)") && \
export CLUSTER_NAME=CLUSTER_NAME && \
export ZONE=ZONE && \
export HF_TOKEN=HUGGING_FACE_TOKEN && \
export CLUSTER_VERSION=CLUSTER_VERSION && \
export GSBUCKET=GSBUCKET && \
export KSA_NAME=KSA_NAME && \
export NAMESPACE=NAMESPACE \
export IMAGE_NAME=IMAGE_NAME

    Substitua os seguintes valores:

    • PROJECT_ID: é o ID do projeto do Google Cloud .
    • CLUSTER_NAME: o nome do cluster do GKE.
    • ZONE: uma zona compatível com a TPU Trillium (v6e).
    • CLUSTER_VERSION: a versão do GKE, que precisa ser compatível com o tipo de máquina que você quer usar. A versão padrão do GKE pode não ter disponibilidade para sua TPU de destino. A TPU Trillium tem suporte nas versões do GKE 1.31.2-gke.1115000 ou mais recentes.
    • GSBUCKET: o nome do bucket do Cloud Storage a ser usado para o Cloud Storage FUSE.
    • KSA_NAME: o nome da conta de serviço do Kubernetes usada para acessar buckets do Cloud Storage. O acesso ao bucket é necessário para que o Cloud Storage FUSE funcione.
    • NAMESPACE: o namespace do Kubernetes em que você quer implantar os recursos de vLLM.
    • IMAGE_NAME: a imagem da TPU do vLLM. Você pode usar a imagem pública docker.io/vllm/vllm-tpu:2e33fe419186c65a18da6668972d61d7bbc31564 ou criar a sua própria.

Criar um cluster do GKE

É possível disponibilizar LLMs 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. Crie um cluster do GKE Autopilot:

    gcloud container clusters create-auto CLUSTER_NAME \
    --cluster-version=CLUSTER_VERSION \
    --release-channel=rapid

    Substitua PROJECT_ID pelo ID do projeto do Google Cloud do seu projeto.

Padrão

  1. Crie um cluster do GKE Standard:

    gcloud container clusters create CLUSTER_NAME \
    --project=PROJECT_ID \
    --zone=ZONE \
    --cluster-version=CLUSTER_VERSION \
    --release-channel=rapid \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --addons GcsFuseCsiDriver

  2. Crie um pool de nós de fração da TPU:

    gcloud container node-pools create tpunodepool \
    --zone=ZONE \
    --num-nodes=1 \
    --machine-type=ct6e-standard-8t \
    --cluster=CLUSTER_NAME \
    --enable-autoscaling --total-min-nodes=1 --total-max-nodes=2

    O GKE cria os seguintes recursos para o LLM:

Configurar o kubectl para se comunicar com o cluster

Para configurar o kubectl para se comunicar com o cluster, execute o seguinte comando:

  gcloud container clusters get-credentials CLUSTER_NAME --location=ZONE

Criar um Secret do Kubernetes para as credenciais do Hugging Face

  1. Crie um namespace Pule esta etapa se estiver usando o namespace default:

    kubectl create namespace NAMESPACE

  2. Crie um secret do Kubernetes que contenha o token do Hugging Face e execute o seguinte comando:

    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=HUGGING_FACE_TOKEN \
    --namespace NAMESPACE

Criar um bucket do Cloud Storage

No Cloud Shell, execute este comando:

gcloud storage buckets create gs://GSBUCKET \
    --uniform-bucket-level-access

Isso cria um bucket do Cloud Storage para armazenar os arquivos de modelo transferidos por download do Hugging Face.

Configurar uma conta de serviço do Kubernetes para acessar o bucket

  1. Crie a conta de serviço do Kubernetes:

    kubectl create serviceaccount KSA_NAME --namespace NAMESPACE

  2. Conceda acesso de leitura e gravação à ServiceAccount do Kubernetes para acessar o bucket do Cloud Storage:

    gcloud storage buckets add-iam-policy-binding gs://GSBUCKET \
  --member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/NAMESPACE/sa/KSA_NAME" \
  --role "roles/storage.objectUser"

  3. Como alternativa, conceda acesso de leitura e gravação a todos os buckets do Cloud Storage no projeto:

    gcloud projects add-iam-policy-binding PROJECT_ID \
--member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/NAMESPACE/sa/KSA_NAME" \
--role "roles/storage.objectUser"

    O GKE cria os seguintes recursos para o LLM:

    1. Um bucket do Cloud Storage para armazenar o modelo transferido por download e o cache de compilação. Um driver CSI do Cloud Storage FUSE lê o conteúdo do bucket.
    2. Volumes com armazenamento em cache de arquivos ativado e o recurso de download paralelo do Cloud Storage FUSE.
    Prática recomendada:

    Use um cache de arquivos com suporte de tmpfs ou Hyperdisk / Persistent Disk, dependendo do tamanho esperado do conteúdo do modelo, por exemplo, arquivos de peso. Neste tutorial, você vai usar o armazenamento em cache de arquivos do Cloud Storage FUSE com suporte da RAM.

(Opcional) Criar e implantar a imagem da TPU

Se você precisar de um controle mais preciso sobre o conteúdo da imagem do Docker, escolha essa opção.

Conteinerize seu servidor vLLM:

  1. Clone o repositório vLLM e crie a imagem:

    git clone https://github.com/vllm-project/vllm && cd vllm && git reset --hard 2e33fe419186c65a18da6668972d61d7bbc31564 && docker build -f Dockerfile.tpu . -t vllm-tpu

  2. Envie a imagem para o Artifact Registry:

    gcloud artifacts repositories create vllm-tpu --repository-format=docker --location=REGION_NAME && \
gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
docker image tag vllm-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-tpu/vllm-tpu:latest && \
docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-tpu/vllm-tpu:latest

Implantar o servidor de modelos do vLLM

Para implantar o servidor de modelos vLLM, siga estas etapas:

  1. Inspecione o manifesto de implantação salvo como vllm-llama3-70b.yaml. Uma implantação é um objeto da API do Kubernetes que permite executar várias réplicas de pods distribuídos entre os nós de um cluster:

    
apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: vllm-tpu
  template:
    metadata:
      labels:
        app: vllm-tpu
      annotations:
        gke-gcsfuse/volumes: "true"
        gke-gcsfuse/cpu-limit: "0"
        gke-gcsfuse/memory-limit: "0"
        gke-gcsfuse/ephemeral-storage-limit: "0"
    spec:
      serviceAccountName: KSA_NAME
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
      containers:
      - name: vllm-tpu
        image: IMAGE_NAME
        command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
        args:
        - --host=0.0.0.0
        - --port=8000
        - --tensor-parallel-size=8
        - --max-model-len=8192
        - --model=meta-llama/Meta-Llama-3.1-70B
        - --download-dir=/data
        env: 
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        - name: VLLM_XLA_CACHE_PATH
          value: "/data"
        ports:
        - containerPort: 8000
        resources:
          limits:
            google.com/tpu: 8
        readinessProbe:
          tcpSocket:
            port: 8000
          initialDelaySeconds: 15
          periodSeconds: 10
        volumeMounts:
        - name: gcs-fuse-csi-ephemeral
          mountPath: /data
        - name: dshm
          mountPath: /dev/shm
      volumes:
      - name: gke-gcsfuse-cache
        emptyDir:
          medium: Memory
      - name: dshm
        emptyDir:
          medium: Memory
      - name: gcs-fuse-csi-ephemeral
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: GSBUCKET
            mountOptions: "implicit-dirs,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:100,file-cache:max-parallel-downloads:-1,file-cache:download-chunk-size-mb:10,file-cache:max-size-mb:-1"
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-service
spec:
  selector:
    app: vllm-tpu
  type: LoadBalancer	
  ports:
    - name: http
      protocol: TCP
      port: 8000  
      targetPort: 8000

  2. Execute este comando para aplicar o manifesto:

    kubectl apply -f vllm-llama3-70b.yaml -n NAMESPACE

  3. Confira os registros do servidor de modelos em execução:

    kubectl logs -f -l app=vllm-tpu -n NAMESPACE

    A saída será parecida com esta:

    INFO:     Started server process [1]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)

Disponibilizar o modelo

  1. Execute o comando a seguir para conferir o endereço IP externo do serviço:

    export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')

  2. Em um novo terminal, interaja com o modelo usando curl:

    curl http://$vllm_service:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
    "model": "meta-llama/Meta-Llama-3.1-70B",
    "prompt": "San Francisco is a",
    "max_tokens": 7,
    "temperature": 0
}'

    A saída será semelhante a esta:

    {"id":"cmpl-6b4bb29482494ab88408d537da1e608f","object":"text_completion","created":1727822657,"model":"meta-llama/Meta-Llama-3-8B","choices":[{"index":0,"text":" top holiday destination featuring scenic beauty and","logprobs":null,"finish_reason":"length","stop_reason":null,"prompt_logprobs":null}],"usage":{"prompt_tokens":5,"total_tokens":12,"completion_tokens":7}}

Configurar o escalonador automático personalizado

Nesta seção, você vai configurar o escalonamento automático horizontal de pods usando métricas personalizadas do Prometheus. Você usa as métricas do Google Cloud Managed Service para Prometheus no servidor vLLM.

Para saber mais, consulte Google Cloud Managed Service para Prometheus. Essa opção precisa estar ativada por padrão no cluster do GKE.

  1. Configure o adaptador de métricas personalizadas do Stackdriver no cluster:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml

  2. Adicione a função "Visualizador de monitoramento" à conta de serviço usada pelo adaptador do Stackdriver para métricas personalizadas:

    gcloud projects add-iam-policy-binding projects/PROJECT_ID \
    --role roles/monitoring.viewer \
    --member=principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/custom-metrics/sa/custom-metrics-stackdriver-adapter

  3. Salve o seguinte manifesto como vllm_pod_monitor.yaml:

    
apiVersion: monitoring.googleapis.com/v1
kind: PodMonitoring
metadata:
 name: vllm-pod-monitoring
spec:
 selector:
   matchLabels:
     app: vllm-tpu
 endpoints:
 - path: /metrics
   port: 8000
   interval: 15s

  4. Aplique ao cluster:

    kubectl apply -f vllm_pod_monitor.yaml

Criar carga no endpoint vLLM

Crie carga no servidor vLLM para testar como o GKE é escalonado automaticamente com uma métrica vLLM personalizada.

  1. Execute um script bash (load.sh) para enviar o número N de solicitações paralelas ao endpoint do vLLM:

    #!/bin/bash
N=PARALLEL_PROCESSES
export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
for i in $(seq 1 $N); do
  while true; do
    curl http://$vllm_service:8000/v1/completions -H "Content-Type: application/json" -d '{"model": "meta-llama/Meta-Llama-3.1-70B", "prompt": "Write a story about san francisco", "max_tokens": 100, "temperature": 0}'
  done &  # Run in the background
done
wait

    Substitua PARALLEL_PROCESSES pelo número de processos em paralelo que você quer executar.

  2. Execute o script bash:

    nohup ./load.sh &

Verificar se o Google Cloud Managed Service para Prometheus processa as métricas

Depois que o Google Cloud Managed Service para Prometheus extrai as métricas e você adiciona carga ao endpoint do vLLM, é possível conferir as métricas no Cloud Monitoring.

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

    Acessar o Metrics Explorer

  2. Clique em < > PromQL.

  3. Insira a seguinte consulta para observar as métricas de tráfego:

    vllm:avg_generation_throughput_toks_per_s{cluster='CLUSTER_NAME_HERE'}

No gráfico de linhas, a métrica vLLM aumenta de 0 (pré-carregamento) para um valor (pós-carregamento). Este gráfico confirma que as métricas do vLLM estão sendo ingeridas no Google Cloud Managed Service para Prometheus.

A imagem a seguir é um exemplo de gráfico após a execução do script de carga. Nesse caso, o servidor de modelo está veiculando cerca de 2.000 tokens de geração por segundo.

teste

Implantar a configuração do escalonador automático horizontal de pods

Ao decidir qual métrica usar para o escalonamento automático, recomendamos as seguintes métricas para o TPU vLLM:

  • num_requests_waiting: essa métrica está relacionada ao número de solicitações aguardando na fila do servidor do modelo. Esse número começa a crescer significativamente quando o cache kv está cheio.

  • gpu_cache_usage_perc: essa métrica está relacionada à utilização do cache kv, que está diretamente relacionada ao número de solicitações processadas para um determinado ciclo de inferência no servidor de modelos. Essa métrica funciona da mesma forma em GPUs e TPUs, mas está vinculada ao esquema de nomenclatura da GPU.

Recomendamos o uso de num_requests_waiting ao otimizar a capacidade de processamento e o custo e quando as metas de latência forem alcançáveis com a capacidade de processamento máxima do servidor de modelo.

Recomendamos o uso de gpu_cache_usage_perc quando você tiver cargas de trabalho sensíveis à latência em que o escalonamento baseado em fila não é rápido o suficiente para atender aos seus requisitos.

Para mais explicações, consulte Práticas recomendadas para o escalonamento automático de cargas de trabalho de inferência de modelos de linguagem grandes (LLM) com TPUs.

Ao selecionar um alvo averageValue para a configuração da HPA, você vai precisar determinar isso experimentalmente. Confira a postagem do blog Economizar em GPUs: escalonamento automático mais inteligente para suas cargas de trabalho de inferência do GKE para mais ideias sobre como otimizar essa parte. O profile-generator usado nesta postagem do blog também funciona para a TPU do vLLM.

Para implantar a configuração do escalonador automático horizontal de pods usando num_requests_waiting, siga estas etapas:

  1. Salve o seguinte manifesto como vllm-hpa.yaml:

    
apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
 name: vllm-hpa
spec:
 scaleTargetRef:
   apiVersion: apps/v1
   kind: Deployment
   name: vllm-tpu
 minReplicas: 1
 maxReplicas: 2
 metrics:
   - type: Pods
     pods:
       metric:
         name: prometheus.googleapis.com|vllm:num_requests_waiting|gauge
       target:
         type: AverageValue
         averageValue: 10

    As métricas do vLLM no Google Cloud Managed Service para Prometheus seguem o formato vllm:metric_name.

    Prática recomendada:

    Use num_requests_waiting para dimensionar a taxa de transferência. Use gpu_cache_usage_perc para casos de uso de TPU sensíveis à latência.

  2. Implante a configuração do escalonador automático horizontal de pods:

    kubectl apply -f vllm-hpa.yaml

    O GKE programa outro pod para implantação, o que aciona o escalonador automático do pool de nós para adicionar um segundo nó antes de implantar a segunda réplica do vLLM.

  3. Acompanhe o progresso do escalonamento automático de pods:

    kubectl get hpa --watch

    O resultado será assim:

    NAME       REFERENCE             TARGETS       MINPODS   MAXPODS   REPLICAS   AGE
vllm-hpa   Deployment/vllm-tpu   <unknown>/5   1         2         0          6s
vllm-hpa   Deployment/vllm-tpu   34972m/5      1         2         1          16s
vllm-hpa   Deployment/vllm-tpu   25112m/5      1         2         2          31s
vllm-hpa   Deployment/vllm-tpu   35301m/5      1         2         2          46s
vllm-hpa   Deployment/vllm-tpu   25098m/5      1         2         2          62s
vllm-hpa   Deployment/vllm-tpu   35348m/5      1         2         2          77s

  4. Aguarde 10 minutos e repita as etapas na seção Verificar se o Google Cloud Managed Service para Prometheus consome as métricas. O Google Cloud Managed Service para Prometheus processa as métricas dos dois endpoints do vLLM.

test2

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 os recursos implantados

Para evitar cobranças na sua conta do pelos recursos criados neste guia, execute os seguintes comandos:

ps -ef | grep load.sh | awk '{print $2}' | xargs -n1 kill -9

gcloud container clusters delete CLUSTER_NAME \
  --location=ZONE

A seguir