Accelera il caricamento dei modelli su GKE con Run:ai Model Streamer

Questo documento mostra come accelerare il caricamento dei pesi del modello di AI di grandi dimensioni da Cloud Storage utilizzando Run:ai Model Streamer con il server di inferenza vLLM su Google Kubernetes Engine (GKE).

La soluzione descritta in questo documento presuppone che il modello di AI e i relativi pesi in formato safetensors siano già caricati in un bucket Cloud Storage.

Se aggiungi il flag --load-format=runai_streamer al deployment di vLLM, puoi utilizzare Run:ai Model Streamer per migliorare l'efficienza del download dei modelli per i tuoi carichi di lavoro di AI su GKE.

Questo documento è destinato ai seguenti utenti:

  • Ingegneri di machine learning (ML) che devono caricare modelli di AI di grandi dimensioni dallo spazio di archiviazione di oggetti nei nodi GPU/TPU il più rapidamente possibile.
  • Amministratori e operatori della piattaforma che automatizzano e ottimizzano l'infrastruttura di pubblicazione dei modelli su GKE.
  • Architetti cloud che valutano strumenti di caricamento dei dati specializzati per i carichi di lavoro AI/ML.

Per scoprire di più sui ruoli comuni e sulle attività di esempio a cui viene fatto riferimento nei contenuti di Google Cloud , consulta Ruoli e attività comuni degli utenti GKE.

Panoramica

La soluzione descritta in questo documento utilizza tre componenti principali: Run:ai Model Streamer, vLLM e il formato file safetensors, per accelerare il processo di caricamento dei pesi del modello da Cloud Storage ai nodi GPU.

Run:ai Model Streamer

Run:ai Model Streamer è un SDK Python open source che velocizza il caricamento di modelli di AI di grandi dimensioni sulle GPU. Trasmette i pesi del modello direttamente dallo spazio di archiviazione, ad esempio dai bucket Cloud Storage, alla memoria della GPU. Lo streamer di modelli è particolarmente adatto per accedere ai file safetensors che si trovano in Cloud Storage.

safetensors

safetensors è un formato di file per archiviare i tensori, le strutture di dati principali nei modelli di AI, in modo da migliorare sia la sicurezza che la velocità. safetensors è progettato come alternativa al formato pickle di Python e consente tempi di caricamento rapidi grazie a un approccio di copia zero. Questo approccio consente di accedere direttamente ai tensori dall'origine senza dover caricare prima l'intero file nella memoria locale.

vLLM

vLLM è una libreria open source per l'inferenza e la distribuzione di LLM. Si tratta di un server di inferenza ad alte prestazioni ottimizzato per il caricamento rapido di modelli di AI di grandi dimensioni. In questo documento, vLLM è il motore principale che esegue il modello di AI su GKE e gestisce le richieste di inferenza in entrata. Il supporto dell'autenticazione integrata di Run:ai Model Streamer per Cloud Storage richiede vLLM versione 0.11.1 o successive.

In che modo Run:ai Model Streamer accelera il caricamento dei modelli

Quando avvii un'applicazione di AI basata su LLM per l'inferenza, spesso si verifica un ritardo significativo prima che il modello sia pronto per l'uso. Questo ritardo, noto come avvio a freddo, si verifica perché l'intero file del modello di diversi gigabyte deve essere scaricato da una posizione di archiviazione, ad esempio un bucket Cloud Storage, sul disco locale della macchina. Il file viene quindi caricato nella memoria della GPU. Durante questo periodo di caricamento, la GPU costosa è inattiva, il che è inefficiente e costoso.

Invece del processo di download e caricamento, lo streamer di modelli trasmette il modello direttamente da Cloud Storage alla memoria della GPU. Lo streamer utilizza un backend ad alte prestazioni per leggere in parallelo più parti del modello, chiamate tensori. La lettura simultanea dei tensori è molto più veloce del caricamento sequenziale del file.

Panoramica dell'architettura

Run:ai Model Streamer si integra con vLLM su GKE per accelerare il caricamento dei modelli trasmettendo i pesi del modello direttamente da Cloud Storage alla memoria GPU, bypassando il disco locale.

Il seguente diagramma mostra questa architettura:

Architettura di Run:ai Model Streamer che carica i pesi del modello da Cloud Storage a vLLM su GKE.
Esegui:architettura di Model Streamer con vLLM e Cloud Storage.

Questa architettura include i seguenti componenti e workflow:

  • Bucket Cloud Storage: archivia i pesi del modello di AI nel formato safetensors.
  • Pod GKE con GPU: esegue il server di inferenza vLLM.
  • Server di inferenza vLLM: configurato con il flag --load-format=runai_streamer, che attiva la funzionalità di streaming del modello.
  • Run:ai Model Streamer: quando viene avviato vLLM, Model Streamer legge i pesi del modello dal percorso gs:// specificato nel bucket Cloud Storage. Anziché scaricare i file sul disco, trasmette i dati dei tensori direttamente nella memoria GPU del pod GKE, dove sono immediatamente disponibili per vLLM per l'inferenza.
  • Cloud Storage Anywhere Cache (facoltativo): se abilitata, Anywhere Cache memorizza nella cache i dati del bucket nella stessa zona dei nodi GKE, accelerando ulteriormente l'accesso ai dati per lo streamer.

Vantaggi

  • Tempi di avvio a freddo ridotti: lo streamer di modelli riduce significativamente il tempo necessario per l'avvio dei modelli. Carica i pesi del modello fino a sei volte più velocemente rispetto ai metodi convenzionali. Per ulteriori informazioni, vedi Benchmark di Run:ai Model Streamer.
  • Utilizzo GPU migliorato:riducendo al minimo i ritardi di caricamento dei modelli, le GPU possono dedicare più tempo alle attività di inferenza effettive, aumentando l'efficienza e la capacità di elaborazione complessive.
  • Flusso di lavoro semplificato:la soluzione descritta in questo documento si integra con GKE, consentendo ai server di inferenza come vLLM o SGLang di accedere direttamente ai modelli nei bucket Cloud Storage.

Prima di iniziare

Assicurati di completare i seguenti prerequisiti.

Seleziona o crea un progetto e abilita le API

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

    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.

    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

  • Configura Cloud Shell

    Questo documento utilizza Google Cloud CLI e i comandi kubectl per creare e gestire le risorse richieste per questa soluzione. Puoi eseguire questi comandi in Cloud Shell facendo clic su Attiva Cloud Shell nella parte superiore della console Google Cloud .

    In the Google Cloud console, activate Cloud Shell.

    Activate Cloud Shell

    In alternativa, puoi installare e inizializzare gcloud CLI nel tuo ambiente shell locale per eseguire i comandi. Se vuoi utilizzare un terminale shell locale, esegui il comando gcloud auth login per autenticarti con Google Cloud.

    Concedi ruoli IAM

    Assicurati che il tuo Google Cloud account disponga dei seguenti ruoli IAM nel tuo progetto per poter creare un cluster GKE e gestire Cloud Storage:

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

    Per concedere questi ruoli, esegui i seguenti comandi:

    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"
    

    Sostituisci PROJECT_ID con l'ID progetto.

    prepara l'ambiente

    Questa sezione ti guida nella configurazione del cluster GKE e delle autorizzazioni per accedere al modello in Cloud Storage.

    Crea un cluster GKE con GPU

    Run:ai Model Streamer può essere utilizzato sia con i cluster GKE Autopilot sia con quelli Standard. Scegli la modalità cluster più adatta alle tue esigenze.

    1. Imposta le variabili per il nome del progetto e del cluster:

      export PROJECT_ID=PROJECT_ID
      export CLUSTER_NAME=CLUSTER_NAME
      

      Sostituisci quanto segue:

      • PROJECT_ID: il tuo ID progetto Google Cloud . Puoi trovare l'ID progetto eseguendo il comando gcloud config get-value project.
      • CLUSTER_NAME: il nome del cluster. Ad esempio, run-ai-test.
    2. Crea un cluster Autopilot o Standard:

      Autopilot

      Per creare un cluster GKE Autopilot:

      1. Imposta la regione per il cluster:

        export REGION=REGION
        

        Sostituisci REGION con la regione in cui vuoi creare il cluster. Per prestazioni ottimali, utilizza la stessa regione del bucket Cloud Storage.

      2. Crea il cluster:

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

      I cluster Autopilot eseguono automaticamente il provisioning dei nodi in base ai requisiti del carico di lavoro. Quando esegui il deployment del server vLLM in un passaggio successivo, Autopilot esegue il provisioning dei nodi GPU, se necessario. Per saperne di più, consulta Informazioni sulla creazione automatica di pool di nodi.

      Standard

      Per creare un cluster GKE Standard:

      1. Imposta la zona per il cluster:

        export ZONE=ZONE
        

        Sostituisci ZONE con la zona in cui vuoi creare il cluster. Per prestazioni ottimali, utilizza una zona nella stessa regione del bucket Cloud Storage.

      2. Crea il cluster:

        gcloud container clusters create $CLUSTER_NAME \
            --project=$PROJECT_ID \
            --zone=$ZONE \
            --workload-pool=$PROJECT_ID.svc.id.goog \
            --num-nodes=1
        
      3. Crea un node pool con una macchina 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
        

    Configurazione di Workload Identity Federation for GKE

    Configura la federazione delle identità per i carichi di lavoro per GKE per consentire ai tuoi workload GKE di accedere in modo sicuro al modello nel bucket Cloud Storage.

    1. Imposta le variabili per l'account di servizio e lo spazio dei nomi Kubernetes:

      export KSA_NAME=KSA_NAME
      export NAMESPACE=NAMESPACE
      

      Sostituisci quanto segue:

      • NAMESPACE: lo spazio dei nomi su cui vuoi eseguire i tuoi workload. Assicurati di utilizzare lo stesso spazio dei nomi per creare tutte le risorse in questo documento.
      • KSA_NAME: il nome del account di servizio Kubernetes che il tuo pod può utilizzare per l'autenticazione alle API Google Cloud .
    2. Crea uno spazio dei nomi Kubernetes:

      kubectl create namespace $NAMESPACE
      
    3. Crea un account di servizio Kubernetes (KSA):

      kubectl create serviceaccount $KSA_NAME \
          --namespace=$NAMESPACE
      
    4. Concedi al tuo KSA le autorizzazioni necessarie:

      1. Imposta le variabili di ambiente:

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

        Sostituisci quanto segue:

        • BUCKET_NAME: il nome del bucket Cloud Storage che contiene i file safetensors.
        • PROJECT_ID: il tuo ID progetto Google Cloud .

        PROJECT_NUMBER, PROJECT_ID, NAMESPACE e KSA_NAME verranno utilizzati per creare l'identificatore principale della federazione delle identità per i carichi di lavoro per GKE per il tuo progetto nei passaggi successivi.

      2. Concedi il ruolo roles/storage.bucketViewer al tuo KSA per visualizzare gli oggetti nel bucket 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. Concedi al tuo KSA il ruolo roles/storage.objectUser per leggere, scrivere ed eliminare oggetti nel tuo bucket 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"
        

    Ora hai configurato un cluster GKE con GPU e configurato Workload Identity Federation for GKE, concedendo a un service account Kubernetes le autorizzazioni necessarie per accedere al tuo modello di AI in Cloud Storage. Con il cluster e le autorizzazioni configurati, puoi eseguire il deployment del server di inferenza vLLM, che utilizzerà questo account di servizio per trasmettere in streaming i pesi del modello con Run:ai Model Streamer.

    Esegui il deployment di vLLM con Run:ai Model Streamer

    Esegui il deployment di un pod che esegue il server compatibile con vLLM OpenAI e configurato con il flag --load-format=runai_streamer per utilizzare Run:ai Model Streamer. La versione di vLLM deve essere 0.11.1 o successiva.

    Il seguente manifest di esempio mostra una configurazione vLLM con lo streamer di modelli abilitato per un modello di piccole dimensioni, come gemma-2-9b-it, che utilizza una singola GPU NVIDIA L4.

    • Se utilizzi un modello di grandi dimensioni che richiede più GPU, aumenta il valore di --tensor-parallel-size al numero di GPU richiesto.
    • Il flag --model-loader-extra-config='{"distributed":true}' consente il caricamento distribuito dei pesi del modello ed è un'impostazione consigliata per migliorare le prestazioni di caricamento del modello dallo spazio di archiviazione degli oggetti.

    Per ulteriori informazioni, consulta Parallelismo dei tensori e Parametri regolabili.

    1. Salva il seguente manifest come vllm-deployment.yaml. Il manifest è progettato per la flessibilità sia nei cluster Autopilot che in quelli 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
      

      Sostituisci quanto segue:

      • NAMESPACE: il tuo spazio dei nomi Kubernetes.
      • KSA_NAME: il nome del tuo account di servizio Kubernetes.
      • BUCKET_NAME: il nome del bucket Cloud Storage.
      • PATH_TO_MODEL: il percorso della directory del modello all'interno del bucket, ad esempio models/my-llama.
    2. Applica il manifest per creare il deployment:

      kubectl create -f vllm-deployment.yaml
      

    Puoi generare altri manifest vLLM utilizzando lo strumento GKE Inference Quickstart.

    Verifica il deployment

    1. Controlla lo stato del deployment:

      kubectl get deployments -n NAMESPACE
      
    2. Recupera il nome del pod:

      kubectl get pods -n NAMESPACE | grep vllm-streamer
      

      Prendi nota del nome del pod che inizia con vllm-streamer-deployment.

    3. Per verificare se lo streamer del modello scarica il modello e i pesi, visualizza i log del pod:

      kubectl logs -f POD_NAME -n NAMESPACE
      

      Sostituisci POD_NAME con il nome del pod del passaggio precedente. I log di streaming riusciti sono simili ai seguenti:

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

    (Facoltativo) Migliora il rendimento con Anywhere Cache

    Cloud Storage Anywhere Cache può accelerare ulteriormente il caricamento del modello memorizzando nella cache i dati più vicini ai nodi GKE. La memorizzazione nella cache è particolarmente utile quando si scalano più nodi nella stessa zona.

    Abiliti Anywhere Cache per un bucket Cloud Storage specifico in una zona Google Cloud specifica. Per migliorare le prestazioni, la zona della cache deve corrispondere a quella in cui vengono eseguiti i pod di inferenza GKE. Il tuo approccio dipende dal fatto che i pod vengano eseguiti in zone prevedibili.

    • Per i cluster di zona GKE Standard, in cui sai in quale zona verranno eseguiti i tuoi pod, abilita Anywhere Cache per quella zona specifica.

    • Per i cluster GKE regionali (sia Autopilot che Standard), in cui i pod possono essere pianificati in più zone, hai le seguenti opzioni:

      • Abilita la memorizzazione nella cache in tutte le zone: abilita Anywhere Cache in ogni zona all'interno della regione del cluster. In questo modo, una cache è disponibile indipendentemente da dove GKE pianifica i pod. Tieni presente che i costi vengono addebitati per ogni zona in cui è abilitata la memorizzazione nella cache. Per ulteriori informazioni, consulta la sezione Prezzi di Anywhere Cache.
      • Posiziona i pod in una zona specifica: utilizza una regola nodeSelector o nodeAffinity nel manifest del carico di lavoro per vincolare i pod a una singola zona. A questo punto, puoi attivare Anywhere Cache solo in quella zona. Si tratta di un approccio più conveniente se il tuo workload tollera di essere limitato a una singola zona.

    Per abilitare Anywhere Cache per la zona in cui si trova il cluster GKE, esegui questi comandi:

    # 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
    

    Esegui la pulizia

    Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo documento, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.

    Per eliminare le singole risorse:

    1. Elimina il cluster GKE. Questa azione rimuove tutti i nodi e i carichi di lavoro.

      gcloud container clusters delete CLUSTER_NAME --location=ZONE_OR_REGION
      

      Sostituisci quanto segue:

      • CLUSTER_NAME: il nome del tuo cluster.
      • ZONE_OR_REGION: la zona o la regione del cluster.
    2. Disattiva Anywhere Cache, se l'hai attivata, per evitare costi continui. Per saperne di più, consulta la sezione Disattivare una cache.

    Passaggi successivi