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

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 Cloud de Confiance by S3NS 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 que carrega os pesos do modelo do Cloud Storage para o vLLM no GKE.
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 o vLLM para inferência.
  • 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

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 Cloud de Confiance .

In the Cloud de Confiance 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 Cloud de Confiance by S3NS.

Conceda funções de IAM

Certifique-se de que a sua Cloud de Confiance 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.

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 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 Cloud de Confiance by S3NS projeto. Pode encontrar o ID do projeto executando 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 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 de 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 do 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.s3ns.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 de forma segura 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 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 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 Cloud de Confiance by S3NS 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.s3ns.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.s3ns.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 a 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 de GPUs necessário.
  • 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

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 vantajoso 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 Cloud de Confiance by S3NS 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 do facto 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 do Anywhere 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 Cloud de Confiance by S3NS 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 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?