Acelera la carga de modelos en GKE con Run:ai Model Streamer

Autopilot estándar

En este documento, se muestra cómo acelerar la carga de los pesos de modelos de IA grandes desde Cloud Storage con Run:ai Model Streamer y el servidor de inferencia de vLLM en Google Kubernetes Engine (GKE).

La solución que se describe en este documento supone que ya tienes tu modelo y tus pesos de IA en formato safetensors cargados en un bucket de Cloud Storage.

Si agregas la marca --load-format=runai_streamer a tu implementación de vLLM, puedes usar Run:ai Model Streamer para mejorar la eficiencia de la descarga de modelos para tus cargas de trabajo de IA en GKE.

Este documento está dirigido a los siguientes usuarios:

• Ingenieros de aprendizaje automático (AA) que necesitan cargar modelos de IA grandes desde el almacenamiento de objetos a nodos de GPU/TPU lo más rápido posible
• Administradores y operadores de plataformas que automatizan y optimizan la infraestructura de entrega de modelos en GKE.
• Arquitectos de Cloud que evalúan herramientas especializadas de carga de datos para cargas de trabajo de IA/AA

Para obtener más información sobre los roles comunes y las tareas de ejemplo a las que se hace referencia en el contenido de Google Cloud , consulta Roles y tareas comunes de los usuarios de GKE.

Descripción general

La solución que se describe en este documento usa tres componentes principales: Run:ai Model Streamer, vLLM y el formato de archivo safetensors para acelerar el proceso de carga de pesos del modelo desde Cloud Storage a los nodos de GPU.

Run:ai Model Streamer

Run:ai Model Streamer es un SDK de Python de código abierto que acelera la carga de modelos de IA grandes en las GPUs. Transmite los pesos del modelo directamente desde el almacenamiento, como los buckets de Cloud Storage, a la memoria de la GPU. El transmisor de modelos es especialmente adecuado para acceder a archivos safetensors ubicados en Cloud Storage.

safetensors

safetensors es un formato de archivo para almacenar tensores, las estructuras de datos centrales en los modelos de IA, de una manera que mejora la seguridad y la velocidad. safetensors se diseñó como una alternativa al formato pickle de Python y permite tiempos de carga rápidos a través de un enfoque de cero copias. Este enfoque permite que se pueda acceder directamente a los tensores desde la fuente sin necesidad de cargar primero todo el archivo en la memoria local.

vLLM

vLLM es una biblioteca de código abierto para la inferencia y la entrega de LLM. Es un servidor de inferencia de alto rendimiento optimizado para cargar rápidamente modelos de IA grandes. En este documento, vLLM es el motor principal que ejecuta tu modelo de IA en GKE y controla las solicitudes de inferencia entrantes. La compatibilidad integrada de autenticación de Run:ai Model Streamer para Cloud Storage requiere la versión 0.11.1 o posterior de vLLM.

Cómo Run:ai Model Streamer acelera la carga de modelos

Cuando inicias una aplicación de IA basada en LLM para la inferencia, a menudo se produce una demora significativa antes de que el modelo esté listo para usarse. Esta demora, conocida como inicio en frío, se produce porque el archivo del modelo completo de varios gigabytes se debe descargar desde una ubicación de almacenamiento, como un bucket de Cloud Storage, al disco local de tu máquina. Luego, el archivo se carga en la memoria de la GPU. Durante este período de carga, la GPU costosa permanece inactiva, lo que es ineficiente y costoso.

En lugar del proceso de descarga y, luego, carga, el transmisor de modelos transmite el modelo directamente desde tu Cloud Storage a la memoria de la GPU. El transmisor usa un backend de alto rendimiento para leer varias partes del modelo, llamadas tensores, en paralelo. Leer tensores de forma simultánea es mucho más rápido que cargar el archivo de forma secuencial.

Descripción general de la arquitectura

Run:ai Model Streamer se integra en vLLM en GKE para acelerar la carga de modelos transmitiendo los pesos del modelo directamente desde Cloud Storage a la memoria de la GPU, sin pasar por el disco local.

En el siguiente diagrama, se muestra esta arquitectura:

Arquitectura de Run:ai Model Streamer con vLLM y Cloud Storage.

Esta arquitectura incluye los siguientes componentes y flujo de trabajo:

• Bucket de Cloud Storage: Almacena los pesos del modelo de IA en formato safetensors.
• Pod de GKE con GPU: Ejecuta el servidor de inferencia de vLLM.
• Servidor de inferencia de vLLM: Se configura con la marca --load-format=runai_streamer, que habilita la funcionalidad del transmisor de modelos.
• Run:ai Model Streamer: Cuando se inicia vLLM, el transmisor de modelos lee las ponderaciones del modelo desde la ruta de acceso gs:// especificada en el bucket de Cloud Storage. En lugar de descargar archivos en el disco, transmite datos de tensores directamente a la memoria de la GPU del Pod de GKE, donde vLLM los puede usar de inmediato para la inferencia.
• Caché en cualquier lugar de Cloud Storage (opcional): Si está habilitada, la caché en cualquier lugar almacena en caché los datos del bucket en la misma zona que los nodos de GKE, lo que acelera aún más el acceso a los datos para el transmisor.

Beneficios

• Tiempos de inicio en frío reducidos: El transmisor de modelos reduce significativamente el tiempo que tardan los modelos en iniciarse. Carga los pesos del modelo hasta seis veces más rápido en comparación con los métodos convencionales. Para obtener más información, consulta Comparativas de Run:ai Model Streamer.
• Mayor uso de la GPU: Al minimizar las demoras en la carga del modelo, las GPUs pueden dedicar más tiempo a las tareas de inferencia reales, lo que aumenta la eficiencia general y la capacidad de procesamiento.
• Flujo de trabajo optimizado: La solución que se describe en este documento se integra con GKE, lo que permite que los servidores de inferencia, como vLLM o SGLang, accedan directamente a los modelos en los buckets de Cloud Storage.

Antes de comenzar

Asegúrate de completar los siguientes requisitos previos.

Selecciona o crea un proyecto y habilita las 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

Configura Cloud Shell

En este documento, se usan los comandos de Google Cloud CLI y kubectl para crear y administrar los recursos necesarios para esta solución. Para ejecutar estos comandos en Cloud Shell, haz clic en Activar Cloud Shell en la parte superior de la consola Google Cloud .

In the Google Cloud console, activate Cloud Shell.

Activate Cloud Shell

Como alternativa, puedes instalar e inicializar la CLI de gcloud en tu entorno de shell local para ejecutar los comandos. Si deseas usar una terminal de shell local, ejecuta el comando gcloud auth login para autenticarte con Google Cloud.

Asigna roles de IAM

Asegúrate de que tu cuenta de Google Cloud tenga los siguientes roles de IAM en tu proyecto para que puedas crear un clúster de GKE y administrar Cloud Storage:

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

Para otorgar estos roles, ejecuta los siguientes 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"

Reemplaza PROJECT_ID con el ID del proyecto.

Prepara el entorno

En esta sección, se te guía para configurar tu clúster de GKE y configurar los permisos para acceder a tu modelo en Cloud Storage.

Nota: Debes tener tu modelo de IA en formato safetensors almacenado en un bucket de Cloud Storage.

Crea un clúster de GKE con GPUs

El Run:ai Model Streamer se puede usar con clústeres de GKE Autopilot y Standard. Elige el modo de clúster que mejor se adapte a tus necesidades.

1. Establece variables para el nombre del proyecto y del clúster:

   export PROJECT_ID=PROJECT_ID
   export CLUSTER_NAME=CLUSTER_NAME
   

   Reemplaza lo siguiente:

   • PROJECT_ID: El ID de tu proyecto Google Cloud . Para encontrar el ID de tu proyecto, ejecuta el comando gcloud config get-value project.
   • CLUSTER_NAME: Es el nombre del clúster. Por ejemplo, run-ai-test

2. Crea un clúster de Autopilot o Standard:

   Autopilot

   Sigue estos pasos para crear un clúster de GKE Autopilot:

   1. Establece la región de tu clúster:

      export REGION=REGION
      

      Reemplaza REGION por la región en la que deseas crear el clúster. Para obtener un rendimiento óptimo, usa la misma región que tu bucket de Cloud Storage.

   2. Crea el clúster:

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

   Los clústeres de Autopilot aprovisionan nodos automáticamente según los requisitos de la carga de trabajo. Cuando implementes el servidor de vLLM en un paso posterior, Autopilot aprovisionará los nodos de GPU si es necesario. Para obtener más información, consulta Acerca de la creación automática de grupos de nodos.

   Estándar

   Sigue estos pasos para crear un clúster de GKE Standard:

   1. Establece la zona de tu clúster:

      export ZONE=ZONE
      

      Reemplaza ZONE por la zona en la que deseas crear el clúster. Para obtener un rendimiento óptimo, usa una zona en la misma región que tu bucket de Cloud Storage.

   2. Crea el clúster:

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

   3. Crea un grupo de nodos con una 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
      

Configura Workload Identity Federation for GKE

Configura la federación de identidades para cargas de trabajo de GKE para permitir que tus cargas de trabajo de GKE accedan de forma segura al modelo en tu bucket de Cloud Storage.

1. Establece variables para tu cuenta de servicio y espacio de nombres de Kubernetes:

   export KSA_NAME=KSA_NAME
   export NAMESPACE=NAMESPACE
   

   Reemplaza lo siguiente:

   • NAMESPACE: Es el espacio de nombres en el que deseas que se ejecuten tus cargas de trabajo. Asegúrate de usar el mismo espacio de nombres para crear todos los recursos de este documento.
   • KSA_NAME: Es el nombre de la cuenta de servicio de Kubernetes que tu Pod puede usar para autenticarse en las APIs de Google Cloud .

2. Crea un espacio de nombres de Kubernetes:

   kubectl create namespace $NAMESPACE
   

3. Crea una cuenta de servicio de Kubernetes (KSA):

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

   Precaución: Para cumplir con el principio de privilegio mínimo, evita reutilizar esta cuenta de servicio de Kubernetes en otros Pods dentro del mismo espacio de nombres. Cualquier carga de trabajo que use esta cuenta de servicio tendrá acceso de forma no intencional al bucket de Cloud Storage especificado.

4. Otorga a tu KSA los permisos necesarios:

   1. Configura las variables de entorno:

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

      Reemplaza lo siguiente:

      • BUCKET_NAME: Es el nombre de tu bucket de Cloud Storage que contiene tus archivos safetensors.
      • PROJECT_ID: El ID de tu proyecto Google Cloud .

      Los valores PROJECT_NUMBER, PROJECT_ID, NAMESPACE y KSA_NAME se usarán para construir el identificador principal de la federación de identidades para cargas de trabajo para GKE de tu proyecto en los siguientes pasos.

   2. Otorga el rol roles/storage.bucketViewer a tu KSA para ver objetos en tu bucket de 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. Otorga el rol roles/storage.objectUser a tu KSA para leer, escribir y borrar objetos en tu bucket de 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"
      

Ahora configuraste un clúster de GKE con GPUs y configuraste Workload Identity Federation for GKE, lo que otorga a una cuenta de servicio de Kubernetes los permisos necesarios para acceder a tu modelo de IA en Cloud Storage. Con el clúster y los permisos establecidos, ya puedes implementar el servidor de inferencia de vLLM, que usará esta cuenta de servicio para transmitir pesos del modelo con Run:ai Model Streamer.

Implementa vLLM con Run:ai Model Streamer

Implementa un Pod que ejecute el servidor compatible con OpenAI de vLLM y que esté configurado con la marca --load-format=runai_streamer para usar Run:ai Model Streamer. La versión de vLLM debe ser 0.11.1 o posterior.

En el siguiente manifiesto de muestra, se muestra una configuración de vLLM con el transmisor de modelos habilitado para un modelo de tamaño pequeño, como gemma-2-9b-it, con una sola GPU NVIDIA L4.

• Si usas un modelo grande que requiere varias GPUs, aumenta el valor de --tensor-parallel-size a la cantidad de GPUs requerida.
• La marca --model-loader-extra-config='{"distributed":true}' habilita la carga distribuida de los pesos del modelo y es un parámetro de configuración recomendado para mejorar el rendimiento de la carga del modelo desde el almacenamiento de objetos.

Para obtener más información, consulta Paralelismo de tensores y Parámetros ajustables.

1. Guarda el siguiente manifiesto como vllm-deployment.yaml. El manifiesto está diseñado para ser flexible en los clústeres de Autopilot y 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
   

   Reemplaza lo siguiente:

   • NAMESPACE: Tu espacio de nombres de Kubernetes.
   • KSA_NAME: Es el nombre de tu cuenta de servicio de Kubernetes.
   • BUCKET_NAME: el nombre de tu bucket de Cloud Storage.
   • PATH_TO_MODEL: Es la ruta de acceso al directorio del modelo dentro del bucket, por ejemplo, models/my-llama.

2. Aplica el manifiesto para crear el objeto Deployment:

   kubectl create -f vllm-deployment.yaml
   

Puedes generar más manifiestos de vLLM con la herramienta de Inicio rápido de GKE Inference.

Verifica la implementación

1. Verifica el estado del Deployment:

   kubectl get deployments -n NAMESPACE
   

2. Obtén el nombre del pod como se muestra a continuación:

   kubectl get pods -n NAMESPACE | grep vllm-streamer
   

   Anota el nombre del Pod que comienza con vllm-streamer-deployment.

3. Para verificar si el transmisor de modelos descarga el modelo y los pesos, consulta los registros del Pod:

   kubectl logs -f POD_NAME -n NAMESPACE
   

   Reemplaza POD_NAME por el nombre del Pod del paso anterior. Los registros de transmisión correctos son similares a los siguientes:

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

Éxito: Implementaste correctamente un servidor de inferencia de vLLM en GKE, aprovechando el Model Streamer de Run:ai para acelerar la carga del modelo directamente desde Cloud Storage.

Opcional: Mejora el rendimiento con Anywhere Cache

Anywhere Cache de Cloud Storage puede acelerar aún más la carga del modelo almacenando en caché los datos más cerca de tus nodos de GKE. El almacenamiento en caché es especialmente beneficioso cuando se escalan varios nodos en la misma zona.

Habilita Anywhere Cache para un bucket específico de Cloud Storage en una Google Cloud zona específica. Para mejorar el rendimiento, la zona de la caché debe coincidir con la zona en la que se ejecutan tus Pods de inferencia de GKE. Tu enfoque dependerá de si tus Pods se ejecutan en zonas predecibles.

• En el caso de los clústeres zonales de GKE Standard, en los que sabes en qué zona se ejecutarán tus Pods, habilita Anywhere Cache para esa zona específica.

• En el caso de los clústeres regionales de GKE (tanto Autopilot como Standard), en los que los Pods se pueden programar en varias zonas, tienes las siguientes opciones:

  • Habilita el almacenamiento en caché en todas las zonas: Habilita Anywhere Cache en cada zona dentro de la región del clúster. Esto garantiza que haya una caché disponible independientemente de dónde GKE programe tus Pods. Ten en cuenta que incurres en costos por cada zona en la que se habilita el almacenamiento en caché. Para obtener más información, consulta los precios de Anywhere Cache.
  • Coloca Pods en una zona específica: Usa una regla nodeSelector o nodeAffinity en el manifiesto de tu carga de trabajo para restringir tus Pods a una sola zona. Luego, puedes habilitar Anywhere Cache solo en esa zona. Este es un enfoque más rentable si tu carga de trabajo tolera estar restringida a una sola zona.

Para habilitar Anywhere Cache en la zona en la que reside tu clúster de GKE, ejecuta los siguientes 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

Realiza una limpieza

Para evitar que se apliquen cargos a tu Google Cloud cuenta por los recursos usados en este documento, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

Para borrar los recursos individuales, sigue estos pasos:

1. Borra el clúster de GKE. Esta acción quita todos los nodos y las cargas de trabajo.

   gcloud container clusters delete CLUSTER_NAME --location=ZONE_OR_REGION
   

   Reemplaza lo siguiente:

   • CLUSTER_NAME: El nombre de tu clúster.
   • ZONE_OR_REGION: Es la zona o región de tu clúster.

2. Inhabilita Anywhere Cache, si lo habilitaste, para evitar costos continuos. Para obtener más información, consulta Cómo inhabilitar una caché.

¿Qué sigue?

• Obtén más información para cargar modelos de Hugging Face en Cloud Storage (experimental)
• Obtén información para entregar LLMs en GKE con GPUs.
• Obtén más información sobre las GPUs disponibles en GKE.
• Lee la descripción general de GKE Storage.
• Obtén más información para configurar la federación de identidades para cargas de trabajo.