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

Autopilot Padrão

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

A solução neste documento pressupõe que já tem o seu modelo de IA e ponderações no formato safetensors carregados num contentor do Cloud Storage.

Ao adicionar a flag --load-format=runai_streamer à sua implementação de vLLM, pode usar o Run:ai Model Streamer para melhorar a eficiência da transferência de modelos para as suas cargas de trabalho de IA no GKE.

Este documento destina-se aos seguintes utilizadores:

• Engenheiros de aprendizagem automática (ML) que precisam de carregar grandes modelos de IA do armazenamento de objetos para nós de GPU/TPU o mais rapidamente possível.
• Administradores e operadores de plataformas que automatizam e otimizam a infraestrutura de publicação de modelos no GKE.
• Arquitetos de nuvem que avaliam ferramentas de carregamento de dados especializadas para cargas de trabalho de IA/AA.

Para saber mais sobre as funções comuns e as tarefas de exemplo referenciadas no Google Cloud conteúdo, consulte o artigo Funções e tarefas comuns de utilizadores do GKE.

Vista geral

A solução descrita neste documento usa três componentes principais: o Run:ai Model Streamer, o vLLM e o formato de ficheiro safetensors, para acelerar o processo de carregamento de ponderações de modelos do Cloud Storage para nós de GPU.

Run:ai Model Streamer

O Run:ai Model Streamer é um SDK Python de código aberto que acelera o carregamento de grandes modelos de IA em GPUs. Transmite os pesos do modelo diretamente do armazenamento, como contentores do Cloud Storage, para a memória da sua GPU. O streamer de modelos é particularmente adequado para aceder a ficheiros safetensors localizados no Cloud Storage.

safetensors

safetensors é um formato de ficheiro para armazenar tensores, as estruturas de dados essenciais nos modelos de IA, de uma forma que melhora a segurança e a velocidade. O safetensors foi concebido como uma alternativa ao formato pickle do Python e permite tempos de carregamento rápidos através de uma abordagem de cópia zero. Esta abordagem permite que os tensores sejam diretamente acessíveis a partir da origem sem ter de carregar primeiro o ficheiro completo na memória local.

vLLM

O vLLM é uma biblioteca de código aberto para inferência e publicação de LLMs. É um servidor de inferência de alto desempenho otimizado para carregar rapidamente modelos de IA grandes. Neste documento, o vLLM é o motor principal que executa o seu modelo de IA no GKE e processa os pedidos de inferência recebidos. O suporte de autenticação integrado do Run:ai Model Streamer para o Cloud Storage requer a versão 0.11.1 ou posterior do vLLM.

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

Quando inicia uma aplicação de IA baseada em MDL/CE para inferência, ocorre frequentemente um atraso significativo antes de o modelo estar pronto para utilização. Este atraso, conhecido como início a frio, ocorre porque o ficheiro do modelo de vários gigabytes tem de ser transferido de uma localização de armazenamento, como um contentor do Cloud Storage, para o disco local da sua máquina. Em seguida, o ficheiro é carregado na memória da GPU. Durante este período de carregamento, a GPU cara está inativa, o que é ineficiente e dispendioso.

Em vez do processo de transferência e, em seguida, carregamento, o streamer de modelos faz stream do modelo diretamente do seu Cloud Storage para a memória da GPU. O streamer usa um back-end de alto desempenho para ler várias partes do modelo, denominadas tensores, em paralelo. A leitura de tensores em simultâneo é significativamente mais rápida do que o carregamento do ficheiro sequencialmente.

Vista geral da arquitetura

O Run:ai Model Streamer integra-se com o vLLM no GKE para acelerar o carregamento de modelos através da transmissão em stream dos pesos dos modelos diretamente do Cloud Storage para a memória da GPU, ignorando o disco local.

O diagrama seguinte mostra esta arquitetura:

Arquitetura do Run:ai Model Streamer com vLLM e Cloud Storage.

Esta arquitetura inclui os seguintes componentes e fluxo de trabalho:

• Contentor do Cloud Storage: armazena as ponderações do modelo de IA no formato safetensors.
• GKE Pod com GPU: executa o servidor de inferência vLLM.
• Servidor de inferência vLLM: configurado com a flag --load-format=runai_streamer, que ativa a funcionalidade de streaming de modelos.
• Run:ai Model Streamer: quando o vLLM é iniciado, o Model Streamer lê os pesos do modelo do caminho gs:// especificado no contentor do Cloud Storage. Em vez de transferir ficheiros para o disco, faz stream de dados de tensores diretamente para a memória da GPU do pod do GKE, onde ficam imediatamente disponíveis para a inferência do vLLM.
• Cache em qualquer lugar do Cloud Storage (opcional): se ativada, a cache em qualquer lugar armazena em cache os dados do contentor na mesma zona que os nós do GKE, acelerando ainda mais o acesso aos dados para o streamer.

Vantagens

• Tempos de início a frio reduzidos: o streamer de modelos reduz significativamente o tempo que os modelos demoram a iniciar. Carrega os pesos do modelo até seis vezes mais rápido em comparação com os métodos convencionais. Para mais informações, consulte o artigo Referências do Run:ai Model Streamer
• Utilização melhorada da GPU: ao minimizar os atrasos no carregamento de modelos, as GPUs podem dedicar mais tempo às tarefas de inferência reais, o que aumenta a eficiência geral e a capacidade de processamento.
• Fluxo de trabalho simplificado: a solução descrita neste documento integra-se com o GKE, o que permite que os servidores de inferência, como o vLLM ou o SGLang, acedam diretamente aos modelos em contentores do Cloud Storage.

Antes de começar

Certifique-se de que cumpre os seguintes pré-requisitos.

Selecione ou crie um projeto e ative as APIs

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.

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 Kubernetes Engine, Cloud Storage, Compute Engine, IAM 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 Kubernetes Engine, Cloud Storage, Compute Engine, IAM 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

Configure o Cloud Shell

Este documento usa a CLI Google Cloud e os comandos kubectl para criar e gerir os recursos necessários para esta solução. Pode executar estes comandos no Cloud Shell clicando em Ativar Cloud Shell na parte superior da consola Google Cloud .

In the Google Cloud console, activate Cloud Shell.

Activate Cloud Shell

Em alternativa, pode instalar e inicializar a CLI gcloud no seu ambiente de shell local para executar os comandos. Se quiser usar um terminal de shell local, execute o comando gcloud auth login para fazer a autenticação com Google Cloud.

Conceda funções de IAM

Certifique-se de que a sua Google Cloud conta tem as seguintes funções do IAM no seu projeto para poder criar um cluster do GKE e gerir o Cloud Storage:

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

Para conceder estas funções, 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.

Prepare o seu ambiente

Esta secção explica como configurar o cluster do GKE e configurar autorizações para aceder ao seu modelo no Cloud Storage.

Nota: tem de ter o seu modelo de IA no formato safetensors armazenado num contentor do Cloud Storage.

Crie um cluster do GKE com GPUs

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

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

   export PROJECT_ID=PROJECT_ID
   export CLUSTER_NAME=CLUSTER_NAME
   

   Substitua o seguinte:

   • PROJECT_ID: o ID do seu Google Cloud projeto. Pode encontrar o ID do projeto executando o comando gcloud config get-value project.
   • CLUSTER_NAME: o nome do seu cluster. Por exemplo, run-ai-test.

2. Crie um cluster do Autopilot ou um cluster padrão:

   Piloto automático

   Siga estes passos para criar um cluster do GKE Autopilot:

   1. Defina a região do cluster:

      export REGION=REGION
      

      Substitua REGION pela região onde quer criar o cluster. Para um desempenho ideal, use a mesma região que o seu contentor do Cloud Storage.

   2. Crie o cluster:

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

   Os clusters do Autopilot aprovisionam automaticamente nós com base nos requisitos da carga de trabalho. Quando implementar o servidor vLLM num passo posterior, o Autopilot aprovisiona os nós da GPU, se necessário. Para mais informações, consulte o artigo Acerca da criação automática de conjuntos de nós.

   Standard

   Siga estes passos para criar um cluster padrão do GKE:

   1. Defina a zona para o seu cluster:

      export ZONE=ZONE
      

      Substitua ZONE pela zona onde quer criar o cluster. Para um desempenho ideal, use uma zona na mesma região que o seu contentor do Cloud Storage.

   2. Crie o cluster:

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

   3. Crie um node pool com uma máquina G2 (GPU NVIDIA L4):

      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
      

Configure a Workload Identity Federation para o GKE

Configure a Workload Identity Federation para o GKE para permitir que as suas cargas de trabalho do GKE acedam em segurança ao modelo no seu contentor do Cloud Storage.

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

   export KSA_NAME=KSA_NAME
   export NAMESPACE=NAMESPACE
   

   Substitua o seguinte:

   • NAMESPACE: o espaço de nomes no qual quer que as suas cargas de trabalho sejam executadas. Certifique-se de que usa o mesmo espaço de nomes para criar todos os recursos neste documento.
   • KSA_NAME: o nome da conta de serviço do Kubernetes que o seu pod pode usar para autenticar em Google Cloud 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
   

   Aviso: para aderir ao princípio dos privilégios mínimos, evite reutilizar esta conta de serviço do Kubernetes noutros pods no mesmo espaço de nomes. Qualquer carga de trabalho que use esta conta de serviço terá acesso não intencional ao contentor do Cloud Storage especificado.

4. Conceda ao seu KSA as autorizaçõ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 o seguinte:

      • BUCKET_NAME: o nome do seu contentor do Cloud Storage que contém os ficheiros safetensors.
      • PROJECT_ID: o ID do seu Google Cloud projeto.

      Os valores PROJECT_NUMBER, PROJECT_ID, NAMESPACE e KSA_NAME vão ser usados para criar a Workload Identity Federation para o identificador principal do GKE para o seu projeto nos passos seguintes.

   2. Conceda a função roles/storage.bucketViewer ao seu KSA para ver objetos no seu contentor 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.svc.id.goog/subject/ns/$NAMESPACE/sa/$KSA_NAME" \
          --role="roles/storage.bucketViewer"
      

   3. Conceda a função roles/storage.objectUser à sua KSA para ler, escrever e eliminar objetos no seu contentor 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.svc.id.goog/subject/ns/$NAMESPACE/sa/$KSA_NAME" \
          --role="roles/storage.objectUser"
      

Configurou um cluster do GKE com GPUs e configurou a Workload Identity Federation para o GKE, concedendo a uma conta de serviço do Kubernetes as autorizações necessárias para aceder ao seu modelo de IA no Cloud Storage. Com o cluster e as autorizações implementados, está tudo pronto para implementar o servidor de inferência vLLM, que vai usar esta conta de serviço para fazer streaming de ponderações do modelo com o Run:ai Model Streamer.

Implemente o vLLM com o Run:ai Model Streamer

Implemente um pod que execute o servidor compatível com OpenAI vLLM e que esteja configurado com o sinalizador --load-format=runai_streamer para usar o Run:ai Model Streamer. A versão do vLLM tem de ser 0.11.1 ou posterior.

O manifesto de exemplo seguinte mostra uma configuração vLLM com o streamer de modelos ativado para um modelo de tamanho pequeno, como gemma-2-9b-it, usando uma única GPU NVIDIA L4.

• Se estiver a usar um modelo grande que requer várias GPUs, aumente o valor de --tensor-parallel-size para o número necessário de GPUs.
• A flag --model-loader-extra-config='{"distributed":true}' permite o carregamento distribuído de ponderações do modelo e é uma definição recomendada para melhorar o desempenho do carregamento do modelo a partir do armazenamento de objetos.

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

1. Guarde o seguinte manifesto como vllm-deployment.yaml. O manifesto foi concebido para ser flexível nos clusters do Autopilot e padrão.

       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 o seguinte:

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

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

   kubectl create -f vllm-deployment.yaml
   

Pode gerar mais manifestos vLLM com a ferramenta GKE Inference Quickstart.

Valide a implementação

1. Verifique o estado da implementação:

   kubectl get deployments -n NAMESPACE
   

2. Obtenha o nome do agrupamento:

   kubectl get pods -n NAMESPACE | grep vllm-streamer
   

   Tome nota do nome do Pod que começa com vllm-streamer-deployment.

3. Para verificar se o streamer de modelos transfere o modelo e os pesos, veja os registos do pod:

   kubectl logs -f POD_NAME -n NAMESPACE
   

   Substitua POD_NAME pelo nome do Pod do passo anterior. Os registos de streaming bem-sucedidos têm um aspeto semelhante ao seguinte:

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

Êxito: implementou com êxito um servidor de inferência vLLM no GKE, tirando partido do Run:ai Model Streamer para acelerar o carregamento de modelos diretamente do Cloud Storage.

Opcional: melhore o desempenho com a cache em qualquer lugar

A cache em qualquer lugar do Cloud Storage pode acelerar ainda mais o carregamento de modelos ao colocar dados em cache mais perto dos seus nós do GKE. O armazenamento em cache é especialmente benéfico quando dimensiona vários nós na mesma zona.

Ativa a cache em qualquer lugar para um contentor do Cloud Storage específico numa zona Google Cloud específica. Para melhorar o desempenho, a zona da cache tem de corresponder à zona onde os seus pods de inferência do GKE são executados. A sua abordagem depende de os seus pods serem executados em zonas previsíveis.

• Para clusters zonais do GKE Standard, onde sabe em que zona os seus pods vão ser executados, ative a cache em qualquer lugar para essa zona específica.

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

  • Ativar o armazenamento em cache em todas as zonas: ative a cache em qualquer lugar em todas as zonas na região do cluster. Isto garante que está disponível uma cache, independentemente do local onde o GKE agende os seus pods. Tenha em atenção que incorre em custos para cada zona onde o armazenamento em cache está ativado. Para mais informações, consulte os preços da cache em qualquer lugar.
  • Coloque os pods numa zona específica: use uma regra nodeSelector ou nodeAffinity no manifesto da carga de trabalho para restringir os pods a uma única zona. Em seguida, pode ativar a cache em qualquer lugar apenas nessa zona. Esta é uma abordagem mais rentável se a sua carga de trabalho tolerar a restrição a uma única zona.

Para ativar a cache em qualquer lugar para a zona onde 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 incorrer em custos na sua Google Cloud conta pelos recursos usados neste documento, elimine o projeto que contém os recursos ou mantenha o projeto e elimine os recursos individuais.

Para eliminar os recursos individuais, siga estes passos:

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

   gcloud container clusters delete CLUSTER_NAME --location=ZONE_OR_REGION
   

   Substitua o seguinte:

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

2. Desative a cache em qualquer lugar, se a tiver ativado, para evitar custos contínuos. Para mais informações, consulte o artigo Desative uma cache.

O que se segue?

• Saiba como carregar modelos do Hugging Face no Cloud Storage (experimental)
• Saiba como servir LLMs no GKE com GPUs.
• Saiba mais sobre as GPUs disponíveis no GKE.
• Leia a vista geral do armazenamento do GKE.
• Saiba como configurar a Workload Identity Federation.