Prácticas recomendadas para optimizar la inferencia de modelos de lenguaje grandes con GPUs en Google Kubernetes Engine (GKE)

Google Kubernetes Engine (GKE) proporciona un control preciso para la inferencia de modelos de lenguaje extenso (LLM) con un rendimiento y un coste óptimos. En esta guía se describen las prácticas recomendadas para optimizar la inferencia y el servicio de LLMs abiertos con GPUs en GKE mediante los frameworks de servicio vLLM y Text Generation Inference (TGI).

Para ver una lista de comprobación resumida de todas las prácticas recomendadas, consulta el resumen de la lista de comprobación.

Objetivos

Esta guía está dirigida a clientes de IA generativa, usuarios nuevos o actuales de GKE, ingenieros de aprendizaje automático e ingenieros de LLMOps (DevOps) que quieran optimizar sus cargas de trabajo de LLM mediante GPUs con Kubernetes.

Al finalizar esta guía, podrás hacer lo siguiente:

  • Elige técnicas de optimización de LLMs posteriores al entrenamiento, como la cuantización, el paralelismo de tensores y la optimización de la memoria.
  • Sopesa las ventajas y desventajas generales al plantearte usar estas técnicas de optimización.
  • Despliega modelos LLM abiertos en GKE mediante frameworks de servicio como vLLM o TGI con los ajustes de optimización habilitados.

Descripción general de las técnicas de optimización del servicio de LLMs

A diferencia de las cargas de trabajo que no usan IA, las cargas de trabajo de LLMs suelen tener una latencia más alta y un rendimiento más bajo debido a que dependen de las operaciones de multiplicación de matrices. Para mejorar el rendimiento de la inferencia de los LLMs, puedes usar aceleradores de hardware especializados (por ejemplo, GPUs y TPUs) y frameworks de servicio optimizados.

Puedes aplicar una o varias de las siguientes prácticas recomendadas para reducir la latencia de la carga de trabajo de LLM y, al mismo tiempo, mejorar el rendimiento y la rentabilidad:

En los ejemplos de esta guía se usa el LLM Gemma 7B junto con los frameworks de servicio vLLM o TGI para aplicar estas prácticas recomendadas. Sin embargo, los conceptos y las funciones que se describen se pueden aplicar a la mayoría de los LLMs abiertos más populares.

Antes de empezar

Antes de probar los ejemplos de esta guía, completa estas tareas previas:

  1. Sigue las instrucciones de estas guías para acceder al modelo Gemma, preparar tu entorno y crear y configurar Google Cloud recursos:

    Asegúrate de guardar el token de acceso de Hugging Face en tu secreto de Kubernetes.

  2. Clona el repositorio de ejemplos https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/ en tu entorno de desarrollo local.

  3. Cambia el directorio de trabajo a /kubernetes-engine-samples/ai-ml/llm-serving-gemma/.

Práctica recomendada: cuantización

La cuantización es una técnica análoga a la compresión de imágenes con pérdidas que reduce el tamaño del modelo representando los pesos en formatos de menor precisión (8 o 4 bits), lo que disminuye los requisitos de memoria. Sin embargo, al igual que la compresión de imágenes, la cuantización implica un equilibrio: la reducción del tamaño del modelo puede provocar una disminución de la precisión.

Existen varios métodos de cuantización, cada uno con sus propias ventajas e inconvenientes. Algunos, como AWQ y GPTQ, requieren una cuantización previa y están disponibles en plataformas como Hugging Face o Kaggle. Por ejemplo, si aplicas GPTQ al modelo Llama-2 13B y AWQ al modelo Gemma 7B, puedes servir los modelos en una sola GPU L4 en lugar de en dos GPUs L4 sin cuantificación.

También puedes realizar la cuantización con herramientas como AutoAWQ y AutoGPTQ. Estos métodos pueden mejorar la latencia y el rendimiento. Por el contrario, las técnicas que usan EETQ y la biblioteca bitsandbytes para la cuantización no requieren modelos precuantizados, por lo que pueden ser una opción adecuada cuando no hay versiones precuantizadas disponibles.

La mejor técnica de cuantización que puedes usar depende de tus objetivos específicos y de la compatibilidad de la técnica con el framework de servicio que quieras usar. Para obtener más información, consulta la guía de cuantización de Hugging Face.

Selecciona una de estas pestañas para ver un ejemplo de aplicación de la cuantización con los frameworks TGI o vLLM:

TGI

GKE admite estas opciones de cuantización con TGI:

  • awq
  • gptq
  • eetq
  • bitsandbytes
  • bitsandbytes-nf4
  • bitsandbytes-fp4

Los métodos de cuantización AWQ y GPTQ requieren modelos precuantizados, mientras que la cuantización EETQ y bitsandbytes se pueden aplicar a cualquier modelo. Para obtener más información sobre estas opciones, consulta este artículo de Hugging Face.

Para usar la cuantización, define el parámetro -–quantize al iniciar el servidor del modelo.

El siguiente fragmento muestra cómo optimizar Gemma 7B con la cuantización bitsandbytes mediante TGI en GKE.

args:
- --model-id=$(MODEL_ID)
- --num-shard=2
- --quantize=bitsandbytes

Para aplicar esta configuración, usa el siguiente comando:

kubectl apply -f tgi/tgi-7b-bitsandbytes.yaml

vLLM

GKE admite estas opciones de cuantización con vLLM:

Para usar la cuantización de modelos con vLLM, los modelos deben estar cuantizados previamente. Cuando inicies el tiempo de ejecución, define el parámetro –quantization.

En el siguiente fragmento se muestra cómo optimizar el modelo Gemma 7B con la cuantización awq mediante vLLM en GKE:

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=1
- --quantization=awq
env:
- name: MODEL_ID
  value: google/gemma-7b-AWQ
resources:
  requests:
    nvidia.com/gpu: 1
  limits:
    nvidia.com/gpu: 1

Para aplicar esta configuración, usa el siguiente comando:

kubectl apply -f vllm/vllm-7b-awq.yaml

Mejorar la latencia mediante la cuantización de la caché de KV

Puedes usar la cuantización de caché de pares clave-valor FP8 E5M2 para reducir significativamente el espacio de memoria de la caché de pares clave-valor y mejorar la latencia, sobre todo en el caso de los tamaños de lote grandes. Sin embargo, esto reduce la precisión de la inferencia.

Para habilitar la cuantización de la caché de valores clave de FP8 E5M2, define el parámetro --kv-cache-dtype fp8_e5m2:

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=1
- --kv-cache-dtype=fp8_e5m2
- --max-model-len=1200
resources:
  requests:
    cpu: "2"
    memory: "25Gi"
    ephemeral-storage: "25Gi"
    nvidia.com/gpu: 1
  limits:
    cpu: "2"
    memory: "25Gi"
    ephemeral-storage: "25Gi"
    nvidia.com/gpu: 1

Para aplicar esta configuración, usa el siguiente comando:

kubectl apply -f vllm/vllm-7b-kvcache.yaml

Práctica recomendada: paralelismo de tensores

El paralelismo de tensores es una técnica que distribuye la carga computacional entre varias GPUs, lo que resulta esencial cuando ejecutas modelos grandes que superan la capacidad de memoria de una sola GPU. Este enfoque puede ser más rentable, ya que te permite usar varias GPUs asequibles en lugar de una sola GPU cara. También puede mejorar el rendimiento de la inferencia del modelo. El paralelismo de tensores aprovecha el hecho de que las operaciones de tensores se pueden realizar de forma independiente en fragmentos de datos más pequeños.

Para obtener más información sobre esta técnica, consulta la guía sobre paralelismo de tensores de Hugging Face.

Selecciona una de estas pestañas para ver un ejemplo de cómo aplicar el paralelismo de tensores con los frameworks TGI o vLLM:

TGI

Con TGI, el tiempo de ejecución de servicio usará todas las GPUs disponibles para el pod de forma predeterminada. Para definir el número de GPUs que quieres usar, especifica el parámetro --num-shard con el número de GPUs como valor.

Consulta la documentación de Hugging Face para ver la lista de modelos compatibles con el paralelismo de tensores.

El siguiente fragmento muestra cómo optimizar el modelo ajustado para instrucciones Gemma 7B usando el paralelismo de tensores y dos GPUs L4:

args:
- --model-id=$(MODEL_ID)
- --num-shard=2

Para aplicar esta configuración, usa el siguiente comando:

kubectl apply -f tgi/tgi-7b-it-tensorparallelism.yaml

En los clústeres de GKE Autopilot, al ejecutar este comando, se crea un pod con requisitos de recursos mínimos de 21 vCPUs y 78 GiB de memoria.

vLLM

vLLM admite la inferencia paralela de tensores distribuida. vLLM habilita la función de forma predeterminada si hay más de una GPU disponible.

El siguiente fragmento muestra cómo puedes optimizar el modelo ajustado para instrucciones Gemma 7B mediante el paralelismo de tensores y dos GPUs L4:

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=2

Para aplicar esta configuración, usa el siguiente comando:

kubectl apply -f vllm/vllm-7b-it-tensorparallelism.yaml

En los clústeres de GKE Autopilot, al ejecutar este comando, se crea un pod con requisitos de recursos mínimos de 21 vCPUs y 78 GiB de memoria.

Práctica recomendada: optimizar la memoria del modelo

Optimizar el uso de memoria de los LLMs es fundamental para que la inferencia sea eficiente. En esta sección se presentan estrategias de optimización de la capa de atención, como la atención paginada y la atención flash. Estas estrategias mejoran la eficiencia de la memoria, lo que permite secuencias de entrada más largas y reduce el tiempo de inactividad de la GPU. En esta sección también se describe cómo puedes ajustar los tamaños de entrada y salida del modelo para adaptarlos a las restricciones de memoria y optimizarlos para frameworks de servicio específicos.

Optimización de la capa de atención

Las capas de autoatención permiten a los modelos comprender el contexto en las tareas de procesamiento del lenguaje, ya que el significado de las palabras puede cambiar en función del contexto. Sin embargo, estas capas almacenan pesos de tokens de entrada, claves (K) y valores (V) en la vRAM de la GPU. Por lo tanto, a medida que se alarga la secuencia de entrada, el tamaño y el tiempo de computación aumentan de forma cuadrática.

El uso del almacenamiento en caché de clave-valor es especialmente útil cuando se trata de secuencias de entrada largas, donde la sobrecarga de la autoatención puede ser significativa. Este enfoque de optimización reduce el procesamiento computacional a una complejidad lineal.

Entre las técnicas específicas para optimizar los mecanismos de atención en los LLMs se incluyen las siguientes:

  • Atención paginada: Atención paginada mejora la gestión de la memoria en modelos grandes y secuencias de entrada largas mediante técnicas de paginación, similares a la memoria virtual del SO. De esta forma, se reduce la fragmentación y la duplicación en la caché de pares clave-valor, lo que permite usar secuencias de entrada más largas sin quedarse sin memoria de GPU.
  • Flash attention: Flash attention reduce los cuellos de botella de la memoria de la GPU minimizando las transferencias de datos entre la RAM de la GPU y la caché de nivel 1 durante la generación de tokens. De esta forma, se elimina el tiempo de inactividad de los núcleos de computación, lo que mejora significativamente el rendimiento de la inferencia y el entrenamiento de las GPUs.

Ajuste del tamaño de entrada y salida del modelo

Los requisitos de memoria dependen del tamaño de la entrada y la salida. Para generar resultados más largos y con más contexto, se necesitan más recursos, mientras que para generar resultados más cortos y con menos contexto, se pueden ahorrar costes usando una GPU más pequeña y barata.

Selecciona una de estas pestañas para ver un ejemplo de cómo ajustar los requisitos de memoria de entrada y salida del modelo en los frameworks TGI o vLLM:

TGI

El tiempo de ejecución de servicio de TGI comprueba los requisitos de memoria durante el inicio y no se inicia si el espacio de memoria máximo posible del modelo no cabe en la memoria de la GPU disponible. Esta comprobación elimina los fallos por falta de memoria (OOM) en cargas de trabajo que requieren mucha memoria.

GKE admite los siguientes parámetros de TGI para optimizar los requisitos de memoria del modelo:

En el siguiente fragmento se muestra cómo puedes servir un modelo ajustado para instrucciones de Gemma 7B con una sola GPU L4 y los siguientes ajustes de parámetros:--max-total-tokens=3072, --max-batch-prefill-tokens=512, --max-input-length=512

args:
- --model-id=$(MODEL_ID)
- --num-shard=1
- --max-total-tokens=3072 
- --max-batch-prefill-tokens=512
- --max-input-length=512
env:
- name: MODEL_ID
  value: google/gemma-7b

Para aplicar esta configuración, usa el siguiente comando:

kubectl apply -f tgi/tgi-7b-token.yaml

vLLM

En vLLM, configura la longitud del contexto del modelo, que influye directamente en el tamaño de la caché de KV y en los requisitos de RAM de la GPU. Las longitudes de contexto más cortas permiten usar GPUs más asequibles. El valor predeterminado es el número máximo de tokens que acepta el modelo. Limita la longitud máxima del contexto con --max-model-len MAX_MODEL_LEN si es necesario.

Por ejemplo, el modelo ajustado para instrucciones Gemma 7B, con una longitud de contexto predeterminada de 8192, supera la capacidad de memoria de una sola GPU NVIDIA L4. Para implementar en una L4, limita la longitud combinada de las peticiones y las salidas asignando a --max-model-len un valor inferior a 640. Este ajuste permite ejecutar el modelo en una sola GPU L4 a pesar de su gran longitud de contexto predeterminada.

Para implementar el límite de tokens modificado, usa el siguiente fragmento:

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=1
- --max-model-len=600

Para aplicar esta configuración, usa el siguiente comando:

kubectl apply -f vllm/vllm-7b-token.yaml

Resumen de la lista de comprobación

Objetivo de optimización Practicar
Latencia
  • Prioriza las GPUs potentes: considera la posibilidad de cambiar a GPUs con mayores capacidades computacionales y rendimiento de E/S de memoria.
  • Explora la cuantización: Las técnicas como AWQ pueden mejorar la latencia, pero ten en cuenta las posibles pérdidas de precisión.
Rendimiento
  • Escalar horizontalmente: aumenta el número de réplicas de servicio (pods) para distribuir la carga de trabajo.
  • Usa el paralelismo de tensores: En el caso de los modelos grandes que superen la capacidad de una sola GPU, usa el paralelismo de tensores para distribuir los cálculos entre varias GPUs. En el caso de los modelos más pequeños, te recomendamos que uses varias réplicas con paralelismo de tensores de `1` para evitar la sobrecarga.
  • Solicitudes por lotes y cuantificación: combina solicitudes y explora técnicas de cuantificación que mantengan una precisión aceptable.
Rentabilidad
  • Elige modelos más pequeños: Selecciona modelos de una familia que se ajusten a tus limitaciones de recursos y a tu presupuesto.
  • Aplicar cuantización: Usa la cuantización para reducir los requisitos de memoria, sobre todo cuando trabajes con modelos más grandes.
  • Limitar la longitud del contexto: limita la longitud del contexto para reducir aún más el uso de memoria y permite la ejecución en GPUs más pequeñas y rentables.

Siguientes pasos