Acelerar o carregamento de modelos no GKE com o Run:ai Model Streamer

Este documento mostra como acelerar o carregamento dos pesos de modelos de IA grandes do Cloud Storage usando o Run:ai Model Streamer com o servidor de inferência vLLM no Google Kubernetes Engine (GKE).

A solução neste documento pressupõe que você já tenha o modelo e os pesos de IA no formato safetensors carregados em um bucket do Cloud Storage.

Ao adicionar a flag --load-format=runai_streamer à implantação do vLLM, você pode usar o Run:ai Model Streamer para melhorar a eficiência do download de modelos para cargas de trabalho de IA no GKE.

Este documento é destinado aos seguintes usuários:

• Engenheiros de machine learning (ML) que precisam carregar modelos de IA grandes do armazenamento de objetos para nós de GPU/TPU o mais rápido possível.
• Administradores e operadores de plataforma que automatizam e otimizam a infraestrutura de disponibilização de modelos no GKE.
• Arquitetos de nuvem que avaliam ferramentas especializadas de carregamento de dados para cargas de trabalho de IA/ML.

Para saber mais sobre papéis comuns e tarefas de exemplo referenciados em Cloud de Confiance by S3NS content, consulte Funções e tarefas comuns do usuário do GKE.

Visão geral

A solução descrita neste documento usa três componentes principais (Run:ai Model Streamer, vLLM e o formato de arquivo safetensors) para acelerar o processo de carregamento de pesos de modelos do Cloud Storage para nós de GPU ou TPU.

Run:ai Model Streamer

Run:ai Model Streamer é um SDK Python de código aberto que acelera o carregamento de modelos de IA grandes em aceleradores. Ele transmite pesos de modelos diretamente do armazenamento, como buckets do Cloud Storage, para a memória da GPU ou TPU. O streamer de modelos é especialmente adequado para acessar arquivos safetensors localizados no Cloud Storage.

safetensors

safetensors é um formato de arquivo para armazenar tensores, as estruturas de dados principais em modelos de IA, de uma forma que melhora a segurança e a velocidade. safetensors foi projetado como uma alternativa ao formato pickle do Python e permite tempos de carregamento rápidos usando uma abordagem de cópia zero. Essa abordagem permite que os tensores sejam acessados diretamente da origem sem precisar carregar o arquivo inteiro na memória local primeiro.

vLLM

vLLM é uma biblioteca de código aberto para inferência e veiculação de LLM. É um servidor de inferência de alto desempenho otimizado para carregar rapidamente modelos de IA grandes. Neste documento, o vLLM é o mecanismo principal que executa o modelo de IA no GKE e processa solicitações de inferência recebidas. O suporte de autenticação integrado do Run:ai Model Streamer para o Cloud Storage exige a versão 0.11.1 ou mais recente do vLLM para GPUs e 0.18.0 ou mais recente para TPUs.

Como o Run:ai Model Streamer acelera o carregamento de modelos

Ao iniciar um aplicativo de IA baseado em LLM para inferência, um atraso significativo geralmente ocorre antes que o modelo esteja pronto para uso. Esse atraso, conhecido como inicialização a frio, ocorre porque o arquivo modelo multigigabyte inteiro precisa ser baixado de um local de armazenamento, como um bucket do Cloud Storage, para o disco local da máquina. Em seguida, o arquivo é carregado na memória do acelerador. Durante esse período de carregamento, o acelerador caro fica inativo, o que é ineficiente e caro.

Em vez do processo de download e carregamento, o streamer de modelos transmite o modelo diretamente do Cloud Storage para a memória da GPU ou TPU. O streamer usa um back-end de alto desempenho para ler várias partes do modelo, chamadas de tensores, em paralelo. A leitura de tensores simultaneamente é significativamente mais rápida do que carregar o arquivo sequencialmente.

Informações gerais da arquitetura

O Run:ai Model Streamer se integra ao vLLM no GKE para acelerar o carregamento de modelos transmitindo pesos de modelos diretamente do Cloud Storage para a memória do acelerador, ignorando o disco local.

O diagrama a seguir mostra essa arquitetura:

Essa arquitetura inclui os seguintes componentes e fluxo de trabalho:

• Bucket do Cloud Storage: armazena os pesos do modelo de IA no formato safetensors.
• Pod do GKE com GPUs ou TPUs: executa o servidor de inferência vLLM.
• Servidor de inferência vLLM: configurado com a flag --load-format=runai_streamer, que ativa a funcionalidade do streamer de modelos.
• Run:ai Model Streamer: quando o vLLM é iniciado, o streamer de modelos lê pesos de modelos do caminho gs:// especificado no bucket do Cloud Storage. Em vez de baixar arquivos para o disco, ele transmite dados de tensor diretamente para a memória do acelerador do pod do GKE, onde fica imediatamente disponível para o vLLM para inferência.
• Rapid Cache do Cloud Storage (opcional): se ativado, o Rapid Cache armazena em cache os dados do bucket na mesma zona dos nós do GKE, acelerando ainda mais o acesso a dados para o streamer.

Benefícios

• Tempos de inicialização a frio reduzidos:o streamer de modelos reduz significativamente o tempo necessário para que os modelos sejam iniciados. Ele carrega os pesos do modelo até seis vezes mais rápido em comparação com os métodos convencionais. Para mais informações, consulte Comparativos do Run:ai Model Streamer.
• Utilização aprimorada do acelerador:ao minimizar os atrasos no carregamento do modelo, aceleradores como GPUs e TPUs podem dedicar mais tempo às tarefas de inferência reais, aumentando a eficiência geral e a capacidade de processamento.
• Fluxo de trabalho simplificado:a solução descrita neste documento se integra ao GKE, permitindo que servidores de inferência como vLLM ou SGLang acessem diretamente modelos em buckets do Cloud Storage.

Antes de começar

Siga os seguintes pré-requisitos:

Selecionar ou criar um projeto e ativar as APIs

Configurar o Cloud Shell

Este documento usa a Google Cloud CLI e comandos kubectl para criar e gerenciar os recursos necessários para essa solução. É possível executar esses comandos no Cloud Shell clicando em Ativar o Cloud Shell na parte de cima da o Cloud de Confiance console.

No Cloud de Confiance console, ative o Cloud Shell.

Ativar o Cloud Shell

Como alternativa, você pode instalar e inicializar a CLI gcloud no ambiente shell local para executar os comandos. Se você quiser usar um terminal de shell local, execute o comando gcloud auth login para fazer a autenticação com Cloud de Confiance by S3NS.

Conceder papéis do IAM

Verifique se sua Cloud de Confiance conta do tem os seguintes papéis do IAM no projeto para que você possa criar um cluster do GKE e gerenciar o Cloud Storage:

• roles/container.admin
• roles/storage.admin

Para conceder esses papéis, execute os seguintes comandos:

gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="user:$(gcloud config get-value account)" \
    --role="roles/container.admin"
gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="user:$(gcloud config get-value account)" \
    --role="roles/storage.admin"

Substitua PROJECT_ID pelo ID do seu projeto.

Preparar o ambiente

Esta seção orienta você na configuração do cluster do GKE e na configuração de permissões para acessar o modelo no Cloud Storage.

Criar um cluster do GKE

O Run:ai Model Streamer pode ser usado com clusters do GKE Autopilot e Standard. Escolha o modo de cluster que melhor atenda às suas necessidades.

1. Defina variáveis para o projeto e o nome do cluster:

   export PROJECT_ID=PROJECT_ID
   export CLUSTER_NAME=CLUSTER_NAME
   

   Substitua:

   • PROJECT_ID: o ID do Cloud de Confiance by S3NS projeto do. Para encontrar o ID do projeto, execute o comando gcloud config get-value project.
   • CLUSTER_NAME: o nome do cluster. Por exemplo, run-ai-test.

2. Crie um cluster do Autopilot ou Standard:

   Autopilot

   Siga estas etapas para criar um cluster do GKE Autopilot:

   1. Defina a região do cluster:

      export REGION=REGION
      

      Substitua REGION pela região em que você quer criar o cluster. Para ter um desempenho ideal, use a mesma região do bucket do Cloud Storage.

   2. Crie o cluster:

      gcloud container clusters create-auto $CLUSTER_NAME \
          --project=$PROJECT_ID \
          --location=$REGION
      

   Os clusters do Autopilot provisionam nós automaticamente com base nos requisitos da carga de trabalho. Ao implantar o servidor vLLM em uma etapa posterior, o Autopilot vai provisionar os nós de GPU ou TPU, se necessário. Para mais informações, consulte Sobre a criação automática de pools de nós.

   Padrão

   Siga estas etapas para criar um cluster do GKE Standard:

   1. Defina a zona do cluster:

      export ZONE=ZONE
      

      Substitua ZONE pela zona em que você quer criar o cluster. Para ter um desempenho ideal, use uma zona na mesma região do bucket do Cloud Storage.

   2. Crie o cluster:

      gcloud container clusters create $CLUSTER_NAME \
          --project=$PROJECT_ID \
          --zone=$ZONE \
          --workload-pool=$PROJECT_ID.s3ns.svc.id.goog \
          --num-nodes=1
      

Crie um pool de nós.

Se você criou um cluster Standard, é necessário criar um pool de nós com GPUs ou TPUs. Os clusters do Autopilot provisionam nós automaticamente com base nos requisitos da carga de trabalho. Crie um pool de nós com base no acelerador que você quer usar:

*   {GPU}

    Create a node pool with one G2 machine (NVIDIA L4 GPU):

    ```sh
    gcloud container node-pools create g2-gpu-pool \
        --cluster=$CLUSTER_NAME \
        --zone=$ZONE \
        --machine-type=g2-standard-16 \
        --num-nodes=1 \
        --accelerator=type=nvidia-l4
    ```

*   {TPU}

    Create a node pool with TPU v7x nodes:

    ```sh
    gcloud container node-pools create tpu7x-pool \
        --cluster=$CLUSTER_NAME \
        --zone=$ZONE \
        --machine-type=tpu7x-standard-4t \
        --num-nodes=1
    ```

Configurar a Federação de Identidade da Carga de Trabalho para GKE

Configure a Federação de Identidade da Carga de Trabalho para GKE para permitir que as cargas de trabalho do GKE acessem com segurança o modelo no bucket do Cloud Storage.

1. Defina variáveis para a conta de serviço e o namespace do Kubernetes:

   export KSA_NAME=KSA_NAME
   export NAMESPACE=NAMESPACE
   

   Substitua:

   • NAMESPACE: o namespace em que você quer que as cargas de trabalho sejam executadas. Use o mesmo namespace para criar todos os recursos neste documento.
   • KSA_NAME: o nome da conta de serviço do Kubernetes que o pod pode usar para fazer a autenticação em Cloud de Confiance by S3NS APIs.

2. Crie um namespace do Kubernetes:

   kubectl create namespace $NAMESPACE
   

3. Crie uma conta de serviço do Kubernetes (KSA):

   kubectl create serviceaccount $KSA_NAME \
       --namespace=$NAMESPACE
   

4. Conceda à KSA as permissões necessárias:

   1. Defina as variáveis de ambiente:

      export BUCKET_NAME=BUCKET_NAME
      export PROJECT_ID=PROJECT_ID
      export PROJECT_NUMBER=$(gcloud projects describe $PROJECT_ID \
          --format 'get(projectNumber)')
      

      Substitua:

      • BUCKET_NAME: o nome do bucket do Cloud Storage que contém os arquivos safetensors.
      • PROJECT_ID: o ID do Cloud de Confiance by S3NS projeto do.

      O PROJECT_NUMBER, PROJECT_ID, NAMESPACE e KSA_NAME serão usados para criar o identificador principal da Federação de Identidade da Carga de Trabalho para GKE do seu projeto nas etapas a seguir.

   2. Conceda o papel roles/storage.bucketViewer à KSA para visualizar objetos no bucket do Cloud Storage:

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

   3. Conceda o papel roles/storage.objectUser à KSA para ler, gravar e excluir objetos no bucket do Cloud Storage:

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

Agora você configurou um cluster do GKE com GPUs ou TPUs e configurou a Federação de Identidade da Carga de Trabalho para GKE, concedendo a uma conta de serviço do Kubernetes as permissões necessárias para acessar o modelo de IA no Cloud Storage. Com o cluster e as permissões no lugar, você está pronto para implantar o servidor de inferência vLLM, que usará essa conta de serviço para transmitir pesos de modelos com o Run:ai Model Streamer.

Implantar o vLLM com o Run:ai Model Streamer

Implante um pod que executa o servidor compatível com o vLLM OpenAI e que está configurado com a flag --load-format=runai_streamer para usar o Run:ai Model Streamer. A versão do vLLM precisa ser 0.11.1 ou mais recente para GPUs e 0.18.0 ou mais recente para TPUs.

Os exemplos de manifestos a seguir mostram como configurar o vLLM com o streamer de modelos ativado para um modelo de tamanho pequeno, como gemma-2-9b-it.

A flag --model-loader-extra-config={"distributed":true} ativa o carregamento distribuído de pesos de modelos e é uma configuração recomendada para melhorar o desempenho do carregamento de modelos do armazenamento de objetos.

Para mais informações, consulte Paralelismo de tensores e Parâmetros ajustáveis.

1. Selecione o manifesto de exemplo com base no acelerador que você quer usar:

   GPU

   O exemplo de manifesto a seguir usa uma única GPU NVIDIA L4. Se você estiver usando um modelo grande que exige várias GPUs, aumente o valor --tensor-parallel-size para o número necessário de GPUs.

   Salve o seguinte manifesto como vllm-deployment.yaml. O manifesto foi projetado para flexibilidade em clusters do Autopilot e Standard.

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: vllm-streamer-deployment
     namespace: NAMESPACE
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: vllm-streamer
     template:
       metadata:
         labels:
           app: vllm-streamer
       spec:
         serviceAccountName: KSA_NAME
         containers:
           - name: vllm-container
             image: vllm/vllm-openai:v0.11.1
             command:
               - python3
               - -m
               - vllm.entrypoints.openai.api_server
             args:
               - --model=gs://BUCKET_NAME/PATH_TO_MODEL
               - --load-format=runai_streamer
               - --model-loader-extra-config={"distributed":true}
               - --host=0.0.0.0
               - --port=8000
               - --disable-log-requests
               - --tensor-parallel-size=1
             ports:
               - containerPort: 8000
                 name: api
             # startupProbe allows for longer startup times for large models
             startupProbe:
               httpGet:
                 path: /health
                 port: 8000
               failureThreshold: 60  # 60 * 10s = 10 minutes timeout
               periodSeconds: 10
               initialDelaySeconds: 30
             readinessProbe:
               httpGet:
                 path: /health
                 port: 8000
               failureThreshold: 3
               periodSeconds: 10
             resources:
               limits:
                 nvidia.com/gpu: "1"
               requests:
                 nvidia.com/gpu: "1"
             volumeMounts:
             - mountPath: /dev/shm
               name: dshm
         nodeSelector:
           cloud.google.com/gke-accelerator: nvidia-l4
         volumes:
         - emptyDir:
             medium: Memory
           name: dshm
   

   Substitua:

   • NAMESPACE: o namespace do Kubernetes.
   • KSA_NAME: o nome da conta de serviço do Kubernetes.
   • BUCKET_NAME: o nome do bucket do Cloud Storage.
   • PATH_TO_MODEL: o caminho para o diretório do modelo no bucket, por exemplo, models/my-llama.

   TPU

   O exemplo de manifesto a seguir usa nós da TPU v7x.

   Salve o seguinte manifesto como vllm-deployment.yaml.

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: tpu-vllm
     namespace: NAMESPACE
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: tpu-vllm
     template:
       metadata:
         labels:
           app: tpu-vllm
       spec:
         serviceAccountName: KSA_NAME
         containers:
         - name: vllm-container
           image: vllm/vllm-tpu:v0.18.0
           resources:
             limits:
               google.com/tpu: "4"
           command: ["sh", "-c"]
           args:
           - >-
             python3 -m vllm.entrypoints.openai.api_server
             --model=gs://BUCKET_NAME/PATH_TO_MODEL
             --load-format=runai_streamer
             --tensor-parallel-size=8
             --port=8000
           ports:
           - containerPort: 8000
           env:
           - name: VLLM_XLA_CACHE_PATH
             value: "gs://BUCKET_NAME/PATH_TO_CACHE"
         nodeSelector:
           cloud.google.com/gke-tpu-accelerator: tpu7x
           cloud.google.com/gke-tpu-topology: 2x2x1
   

   Substitua:

   • NAMESPACE: o namespace do Kubernetes.
   • KSA_NAME: o nome da conta de serviço do Kubernetes.
   • BUCKET_NAME: o nome do bucket do Cloud Storage.
   • PATH_TO_MODEL: o caminho para o diretório do modelo no bucket, por exemplo, models/my-llama.
   • PATH_TO_CACHE: o caminho para o diretório de cache de compilação XLA no bucket, por exemplo, models/xla-cache.

2. Aplique o manifesto para criar a implantação:

   kubectl create -f vllm-deployment.yaml
   

É possível gerar mais manifestos do vLLM usando a ferramenta Guia de início rápido do GKE Inference.

Verificar a implantação

1. Verifique o status da implantação:

   kubectl get deployments -n NAMESPACE
   

2. Consiga o nome do pod:

   kubectl get pods -n NAMESPACE | grep vllm-streamer
   

   Anote o nome do pod que começa com vllm-streamer-deployment.

3. Para verificar se o streamer de modelos baixa o modelo e os pesos, consulte os registros do pod:

   kubectl logs -f POD_NAME -n NAMESPACE
   

   Substitua o POD_NAME pelo nome do pod da etapa anterior. Os registros de streaming bem-sucedidos são semelhantes a este:

   [RunAI Streamer] Overall time to stream 15.0 GiB of all files: 13.4s, 1.1 GiB/s
   

Opcional: melhorar o desempenho com o Rapid Cache

O Rapid Cache do Cloud Storage pode acelerar ainda mais o carregamento de modelos armazenando dados em cache mais perto dos nós do GKE. O armazenamento em cache é especialmente benéfico ao escalonar vários nós na mesma zona.

Você ativa o Rapid Cache para um bucket específico do Cloud Storage em uma zona específica Cloud de Confiance by S3NS . Para melhorar o desempenho, a zona do cache precisa corresponder à zona em que os pods de inferência do GKE são executados. Sua abordagem depende de os pods serem executados em zonas previsíveis.

• Para clusters zonais do GKE Standard, em que você sabe em qual zona os pods serão executados, ative o Rapid Cache para essa zona específica.

• Para clusters regionais do GKE (Autopilot e Standard), em que os pods podem ser programados em várias zonas, você tem as seguintes opções:

  • Ativar o armazenamento em cache em todas as zonas: ative o Rapid Cache em todas as zonas na região do cluster. Isso garante que um cache esteja disponível, independentemente de onde o GKE programa os pods. Observe que você incorre em custos para cada zona em que o armazenamento em cache está ativado. Para mais informações, consulte Preços do Rapid Cache.
  • Colocar pods em uma zona específica: use uma nodeSelector ou nodeAffinity no manifesto da carga de trabalho para restringir os pods a uma única zona. Em seguida, é possível ativar o Rapid Cache apenas nessa zona. Essa é uma abordagem mais econômica se a carga de trabalho tolerar a restrição a uma única zona.

Para ativar o Rapid Cache na zona em que o cluster do GKE reside, execute os seguintes comandos:

# Enable the cache
gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE

# Check the status of the cache
gcloud storage buckets anywhere-caches describe $BUCKET_NAME/$ZONE

Limpar

Para evitar cobranças na sua Cloud de Confiance by S3NS conta do pelos recursos usados neste documento, exclua o projeto que contém os recursos ou mantenha o projeto e exclua os recursos individuais.

Para excluir os recursos individuais, siga estas etapas:

1. Exclua o cluster do GKE. Essa ação remove todos os nós e cargas de trabalho.

   gcloud container clusters delete CLUSTER_NAME --location=ZONE_OR_REGION
   

   Substitua:

   • CLUSTER_NAME: o nome do cluster.
   • ZONE_OR_REGION: a zona ou região do cluster.

2. Desative o Rapid Cache, se você o ativou, para evitar custos contínuos. Para mais informações, consulte Desativar um cache.

A seguir

• Saiba como carregar modelos do Hugging Face no Cloud Storage (experimental)
• Saiba como disponibilizar um LLM com várias GPUs no GKE.
• Saiba mais sobre as GPUs disponíveis no GKE.
• Saiba mais sobre as TPUs disponíveis no GKE.
• Leia a Visão geral do armazenamento do GKE.
• Saiba como configurar a Federação de Identidade da Carga de Trabalho.