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

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 vLLM en Google Kubernetes Engine (GKE).

En la solución de este documento, se supone que ya tienes tu modelo y 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 descarga de modelos para tus cargas de trabajo de IA en GKE.

Este documento está destinado 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 o 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 la nube que evalúan herramientas especializadas de carga de datos para cargas de trabajo de IA o AA

Para obtener más información sobre los roles comunes y las tareas de ejemplo a las que se hace referencia en Google Cloud el contenido, consulta Roles y tareas comunes del usuario 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 de modelos desde Cloud Storage a nodos de GPU o TPU.

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 aceleradores. Transmite pesos de modelos directamente desde el almacenamiento, como buckets de Cloud Storage, a la memoria de tu GPU o TPU. El streamer 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 principales 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 copia cero. Este enfoque permite que los tensores sean accesibles directamente 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 que está 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 con la autenticación integrada de Run:ai Model Streamer para Cloud Storage requiere vLLM versión 0.11.1 o posterior para GPUs y 0.18.0 o posterior para TPUs.

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 se debe descargar todo el archivo del modelo de varios gigabytes 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 del acelerador. Durante este período de carga, el acelerador costoso permanece inactivo, lo que es ineficiente y costoso.

En lugar del proceso de descarga y carga, el streamer de modelos transmite el modelo directamente desde Cloud Storage a la memoria de la GPU o TPU. El streamer 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 con vLLM en GKE para acelerar la carga de modelos mediante la transmisión de pesos de modelos directamente desde Cloud Storage a la memoria del acelerador, sin pasar por el disco local.

En el siguiente diagrama, se muestra esta arquitectura:

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 GPUs o TPUs: Ejecuta el servidor de inferencia vLLM.
• Servidor de inferencia vLLM: Configurado con la marca --load-format=runai_streamer, que habilita la funcionalidad del streamer de modelos.
• Run:ai Model Streamer: Cuando se inicia vLLM, el streamer de modelos lee los pesos 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 del acelerador del Pod de GKE, donde vLLM puede usarlos de inmediato para la inferencia.
• Rapid Cache de Cloud Storage (opcional): Si está habilitado, Rapid Cache 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 streamer.

Beneficios

• Tiempos de inicio en frío reducidos: El streamer de modelos reduce significativamente el tiempo que tardan en iniciarse los modelos. 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 las comparativas de Run:ai Model Streamer.
• Uso mejorado del acelerador: Si se minimizan las demoras en la carga de modelos, los aceleradores como las GPUs y las TPUs 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

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 Google Cloud consola.

En la Google Cloud consola, activa Cloud Shell.

Activa 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 gcloud auth login comando para autenticarte con Google Cloud.

Asigna roles de IAM

Asegúrate de que tu Google Cloud cuenta 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 por el ID del proyecto.

Prepara el entorno

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

Crea un clúster de GKE

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 las variables para el proyecto y el nombre del clúster:

   export PROJECT_ID=PROJECT_ID
   export CLUSTER_NAME=CLUSTER_NAME
   

   Reemplaza lo siguiente:

   • PROJECT_ID: Es el ID del Google Cloud proyecto. Para encontrar el ID del 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 vLLM en un paso posterior, Autopilot aprovisionará los nodos de GPU o TPU 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
      

Crea un grupo de nodos

Si creaste un clúster Standard, debes crear un grupo de nodos con GPUs o TPUs. Los clústeres de Autopilot aprovisionan nodos automáticamente según los requisitos de la carga de trabajo. Crea un grupo de nodos según el acelerador que deseas 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
    ```

Configura Workload Identity Federation for GKE

Configura Workload Identity Federation for GKE para permitir que tus cargas de trabajo de GKE accedan de forma segura al modelo en tu bucket de Cloud Storage.

1. Establece las 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 Google Cloud APIs.

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
   

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: Es el ID del Google Cloud proyecto.

      Se usarán PROJECT_NUMBER, PROJECT_ID, NAMESPACE y KSA_NAME para construir el identificador principal de Workload Identity Federation for GKE para 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 o TPUs 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 en su lugar, estás listo para implementar el servidor de inferencia vLLM, que usará esta cuenta de servicio para transmitir pesos de modelos con Run:ai Model Streamer.

Implementa vLLM con Run:ai Model Streamer

Implementa un Pod que ejecute el servidor compatible con vLLM OpenAI 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 para GPUs y 0.18.0 o posterior para TPUs.

En los siguientes manifiestos de muestra, se muestra cómo configurar vLLM con el streamer de modelos habilitado para un modelo de tamaño pequeño, como gemma-2-9b-it.

La marca --model-loader-extra-config={"distributed":true} habilita la carga distribuida de pesos de modelos y es una configuración recomendada para mejorar el rendimiento de carga de modelos desde el almacenamiento de objetos.

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

1. Selecciona el manifiesto de muestra según el acelerador que deseas usar:

   GPU

   En el siguiente manifiesto de muestra, se usa una sola GPU NVIDIA L4. Si usas un modelo grande que requiere varias GPUs, aumenta el valor --tensor-parallel-size a la cantidad requerida de GPUs.

   Guarda el siguiente manifiesto como vllm-deployment.yaml. El manifiesto está diseñado para brindar flexibilidad 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: Es tu espacio de nombres de Kubernetes.
   • KSA_NAME: Es el nombre de tu cuenta de servicio de Kubernetes.
   • BUCKET_NAME: Es el nombre de tu bucket de Cloud Storage.
   • PATH_TO_MODEL: Es la ruta de acceso al directorio de tu modelo dentro del bucket, por ejemplo, models/my-llama.

   TPU

   En el siguiente manifiesto de muestra, se usan nodos TPU v7x.

   Guarda el siguiente manifiesto 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
   

   Reemplaza lo siguiente:

   • NAMESPACE: Es tu espacio de nombres de Kubernetes.
   • KSA_NAME: Es el nombre de tu cuenta de servicio de Kubernetes.
   • BUCKET_NAME: Es el nombre de tu bucket de Cloud Storage.
   • PATH_TO_MODEL: Es la ruta de acceso al directorio de tu modelo dentro del bucket, por ejemplo, models/my-llama.
   • PATH_TO_CACHE: Es la ruta de acceso al directorio de caché de compilación de XLA dentro del bucket, por ejemplo, models/xla-cache.

2. Aplica el manifiesto para crear la Deployment:

   kubectl create -f vllm-deployment.yaml
   

Puedes generar más manifiestos de vLLM con la herramienta de inicio rápido de inferencia de GKE.

Verifica la Deployment

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
   

   Toma nota del nombre del Pod que comienza con vllm-streamer-deployment.

3. Para verificar si el streamer 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
   

Opcional: Aumenta el rendimiento con Rapid Cache

Rapid Cache de Cloud Storage puede acelerar aún más la carga de modelos 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 Rapid Cache para un bucket de Cloud Storage específico en una zona específica Google Cloud . 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 depende de si tus Pods se ejecutan en zonas predecibles.

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

• Para los clústeres regionales de GKE (Autopilot y Standard), en los que se pueden programar Pods en varias zonas, tienes las siguientes opciones:

  • Habilita el almacenamiento en caché en todas las zonas: Habilita Rapid 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 está habilitado el almacenamiento en caché. Para obtener más información, consulta los precios de Rapid 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 Rapid 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 Rapid 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

Limpia

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: La zona o región de tu clúster.

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

¿Qué sigue?

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