Servir modelos abiertos con el contenedor premium de Hex-LLM en TPU de Cloud

Hex-LLM, un modelo de lenguaje extenso (LLM) de alta eficiencia que se sirve con XLA, es el framework de servicio de LLM de Vertex AI diseñado y optimizado para el hardware de TPU de Cloud. Hex-LLM combina tecnologías de servicio de LLMs, como el procesamiento por lotes continuo y PagedAttention, con optimizaciones de Vertex AI diseñadas para XLA y Cloud TPU. Es un servicio de LLM de alta eficiencia y bajo coste en TPU de Cloud para modelos de código abierto.

Hex-LLM está disponible en Model Garden a través del playground de modelos, el despliegue con un solo clic y el cuaderno.

Funciones

Hex-LLM se basa en proyectos de código abierto con las optimizaciones propias de Google para XLA y Cloud TPU. Hex-LLM consigue un alto rendimiento y una latencia baja al servir LLMs que se usan con frecuencia.

Hex-LLM incluye las siguientes optimizaciones:

  • Algoritmo de procesamiento por lotes continuo basado en tokens para ayudar a que los modelos utilicen el hardware al máximo con un gran número de solicitudes simultáneas.
  • Se ha reescrito por completo el kernel de atención, que se ha optimizado para XLA.
  • Estrategias de paralelismo de datos y de tensores flexibles y componibles con métodos de fragmentación de pesos altamente optimizados para ejecutar LLMs de forma eficiente en varios chips de TPU de Cloud.

Hex-LLM admite una amplia gama de LLMs densos y dispersos:

  • Gemma 2B y 7B
  • Gemma-2 9B y 27B
  • Llama 2 7B, 13B y 70B
  • Llama 3 8B y 70B
  • Llama-3.1 8B y 70B
  • Llama-3.2 1B y 3B
  • Llama-3.3 70B
  • Llama-Guard-3 1B y 8B
  • Llama-4-Scout-17B-16E
  • Mistral 7B
  • Mixtral 8x7B y 8x22B
  • Phi-3 mini y medium
  • Phi-4, Phi-4 reasoning y Phi-4 reasoning plus
  • Qwen-2 0.5B, 1.5B y 7B
  • Qwen-2.5 0.5B, 1.5B, 7B, 14B y 32B

Hex-LLM también ofrece varias funciones, como las siguientes:

  • Hex-LLM se incluye en un solo contenedor. Hex-LLM empaqueta el servidor de la API, el motor de inferencia y los modelos compatibles en una sola imagen de Docker para implementarse.
  • Compatible con el formato de los modelos de Hugging Face. Hex-LLM puede cargar un modelo de Hugging Face desde un disco local, Hugging Face Hub y un segmento de Cloud Storage.
  • Cuantización con bitsandbytes y AWQ.
  • Carga dinámica de LoRA. Hex-LLM puede cargar los pesos de LoRA leyendo el argumento de la solicitud durante el servicio.

Funciones avanzadas

Hex-LLM admite las siguientes funciones avanzadas:

  • Publicación multihost
  • Publicación desagregada [experimental]
  • Almacenamiento en caché de prefijos
  • Compatibilidad con la cuantización de 4 bits

Publicación multihost

Hex-LLM ahora admite el servicio de modelos con un slice de TPU de varios hosts. Esta función te permite servir modelos grandes que no se pueden cargar en una sola VM de TPU de host, que contiene como máximo ocho núcleos de la versión 5e.

Para habilitar esta función, define --num_hosts en los argumentos del contenedor Hex-LLM y --tpu_topology en la solicitud de subida del modelo del SDK de Vertex AI. En el siguiente ejemplo se muestra cómo desplegar el contenedor Hex-LLM con una topología de TPU 4x4 v5e que sirve el modelo Llama 3.1 70B bfloat16:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Meta-Llama-3.1-70B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=16",
    "--num_hosts=4",
    "--hbm_utilization_factor=0.9",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=["python", "-m", "hex_llm.server.api_server"],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    tpu_topology="4x4",
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

Para ver un tutorial completo sobre cómo desplegar el contenedor de Hex-LLM con una topología de TPU multihost, consulta el notebook Vertex AI Model Garden - Llama 3.1 (Deployment).

Por lo general, los únicos cambios que se deben hacer para habilitar el servicio multihost son los siguientes:

  1. Asigna el argumento --tensor_parallel_size al número total de núcleos de la topología de la TPU.
  2. Asigna el argumento --num_hosts al número de hosts de la topología de TPU.
  3. Define --tpu_topology con la API de subida de modelos del SDK de Vertex AI.

Publicación desagregada [experimental]

Hex-LLM ahora admite el servicio desagregado como función experimental. Solo se puede habilitar en la configuración de un solo host y el rendimiento está en proceso de ajuste.

El servicio desagregado es un método eficaz para equilibrar el tiempo hasta el primer token (TTFT) y el tiempo por token de salida (TPOT) de cada solicitud, así como el rendimiento general del servicio. Separa la fase de relleno automático y la fase de decodificación en cargas de trabajo diferentes para que no interfieran entre sí. Este método es especialmente útil en situaciones en las que se establecen requisitos de latencia estrictos.

Para habilitar esta función, define --disagg_topo en los argumentos del contenedor Hex-LLM. A continuación, se muestra un ejemplo de cómo implementar el contenedor Hex-LLM en una TPU v5e-8 que sirve el modelo Llama 3.1 8B bfloat16:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=2",
    "--disagg_topo=3,1",
    "--hbm_utilization_factor=0.9",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=["python", "-m", "hex_llm.server.api_server"],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

El argumento --disagg_topo acepta una cadena con el formato "number_of_prefill_workers,number_of_decode_workers". En el ejemplo anterior, se ha definido en "3,1" para configurar tres trabajadores de relleno automático y un trabajador de decodificación. Cada trabajador usa dos núcleos de TPU v5e.

Almacenamiento en caché de prefijos

El almacenamiento en caché de prefijos reduce el tiempo hasta el primer token (TTFT) de las peticiones que tienen contenido idéntico al principio, como preámbulos de toda la empresa, instrucciones del sistema comunes e historial de conversaciones de varias interacciones. En lugar de procesar los mismos tokens de entrada repetidamente, Hex-LLM puede conservar una caché temporal de los cálculos de los tokens de entrada procesados para mejorar el TTFT.

Para habilitar esta función, define --enable_prefix_cache_hbm en los argumentos del contenedor Hex-LLM. A continuación, se muestra un ejemplo de cómo desplegar el contenedor Hex-LLM en la TPU v5e-8 que sirve el modelo Llama 3.1 8B bfloat16:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=4",
    "--hbm_utilization_factor=0.9",
    "--enable_prefix_cache_hbm",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=["python", "-m", "hex_llm.server.api_server"],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

Hex-LLM usa el almacenamiento en caché de prefijos para optimizar el rendimiento de las peticiones que superan una longitud determinada (512 tokens de forma predeterminada, configurable mediante prefill_len_padding). Los aciertos de caché se producen en incrementos de este valor, lo que asegura que el recuento de tokens almacenados en caché sea siempre un múltiplo de prefill_len_padding. El campo cached_tokens de usage.prompt_tokens_details en la respuesta de la API de finalización de chat indica cuántos tokens de la petición eran aciertos de caché.

"usage": {
  "prompt_tokens": 643,
  "total_tokens": 743,
  "completion_tokens": 100,
  "prompt_tokens_details": {
    "cached_tokens": 512
  }
}

Prefill por fragmentos

La preparación por fragmentos divide la preparación de una solicitud en fragmentos más pequeños y combina la preparación y la decodificación en un paso por lotes. Hex-LLM implementa el prellenado en fragmentos para equilibrar el tiempo hasta el primer token (TTFT) y el tiempo por token de salida (TPOT), y mejora el rendimiento.

Para habilitar esta función, define --enable_chunked_prefill en los argumentos del contenedor Hex-LLM. A continuación, se muestra un ejemplo de cómo implementar el contenedor Hex-LLM en la TPU v5e-8 que sirve el modelo Llama 3.1 8B:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=4",
    "--hbm_utilization_factor=0.9",
    "--enable_chunked_prefill",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=["python", "-m", "hex_llm.server.api_server"],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

Compatibilidad con la cuantización de 4 bits

La cuantización es una técnica para reducir los costes computacionales y de memoria de la inferencia mediante la representación de los pesos o las activaciones con tipos de datos de baja precisión, como INT8 o INT4, en lugar de los tipos de datos habituales BF16 o FP32.

Hex-LLM admite la cuantización de solo pesos INT8. La asistencia ampliada incluye modelos con pesos INT4 cuantificados mediante la cuantificación de punto cero de AWQ. Hex-LLM admite variantes INT4 de las familias de modelos Mistral, Mixtral y Llama.

No se necesita ninguna otra marca para publicar modelos cuantificados.

Empezar a usar Model Garden

El contenedor de servicio de TPU de Cloud de Hex-LLM está integrado en Model Garden. Puedes acceder a esta tecnología de publicación a través del playground, la implementación con un solo clic y los ejemplos de cuadernos de Colab Enterprise para una gran variedad de modelos.

Usar el área de pruebas

El playground de Model Garden es un endpoint de Vertex AI predesplegado al que se puede acceder enviando solicitudes en la tarjeta del modelo.

  1. Introduce una petición y, si quieres, incluye argumentos para tu solicitud.

  2. Haz clic en ENVIAR para obtener la respuesta del modelo rápidamente.

Pruébalo con Gemma.

Usar el despliegue con un solo clic

Puedes desplegar un endpoint de Vertex AI personalizado con Hex-LLM mediante una tarjeta de modelo.

  1. Ve a la página de la tarjeta de modelo y haz clic en Desplegar.

  2. En la variante del modelo que quieras usar, selecciona el tipo de máquina TPU de Cloud v5e para la implementación.

  3. Haz clic en Implementar en la parte inferior para iniciar el proceso de implementación. Recibirás dos notificaciones por correo electrónico: una cuando se suba el modelo y otra cuando el endpoint esté listo.

Usar el cuaderno de Colab Enterprise

Para disfrutar de flexibilidad y personalización, puedes usar ejemplos de cuadernos de Colab Enterprise para desplegar un endpoint de Vertex AI con Hex-LLM mediante el SDK de Vertex AI para Python.

  1. Ve a la página de la tarjeta del modelo y haz clic en Abrir cuaderno.

  2. Selecciona el cuaderno Vertex Serving. El cuaderno se abrirá en Colab Enterprise.

  3. Ejecuta el cuaderno para desplegar un modelo con Hex-LLM y envía solicitudes de predicción al endpoint. El fragmento de código de la implementación es el siguiente:

hexllm_args = [
    f"--model=google/gemma-2-9b-it",
    f"--tensor_parallel_size=4",
    f"--hbm_utilization_factor=0.8",
    f"--max_running_seqs=512",
]
hexllm_envs = {
    "PJRT_DEVICE": "TPU",
    "MODEL_ID": "google/gemma-2-9b-it",
    "DEPLOY_SOURCE": "notebook",
}
model = aiplatform.Model.upload(
    display_name="gemma-2-9b-it",
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=[
        "python", "-m", "hex_llm.server.api_server"
    ],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=hexllm_envs,
    serving_container_shared_memory_size_mb=(16 * 1024),
    serving_container_deployment_timeout=7200,
)

endpoint = aiplatform.Endpoint.create(display_name="gemma-2-9b-it-endpoint")
model.deploy(
    endpoint=endpoint,
    machine_type="ct5lp-hightpu-4t",
    deploy_request_timeout=1800,
    service_account="<your-service-account>",
    min_replica_count=1,
    max_replica_count=1,
)

Estos son algunos ejemplos de cuadernos de Colab Enterprise:

Configurar argumentos de servidor y variables de entorno

Puedes definir los siguientes argumentos para iniciar el servidor Hex-LLM. Puedes adaptar los argumentos para que se ajusten mejor a tu caso práctico y a tus requisitos. Ten en cuenta que los argumentos están predefinidos para la implementación con un solo clic, lo que permite disfrutar de la experiencia de implementación más sencilla. Para personalizar los argumentos, puedes usar como referencia los ejemplos de cuadernos y definir los argumentos según corresponda.

Modelo

  • --model: el modelo que se va a cargar. Puedes especificar un ID de modelo de Hugging Face, una ruta de segmento de Cloud Storage (gs://my-bucket/my-model) o una ruta local. Se espera que los artefactos del modelo sigan el formato de Hugging Face y usen archivos safetensors para los pesos del modelo. BitsAndBytes Se admiten artefactos de modelos cuantificados int8 y AWQ para Llama, Gemma 2 y Mistral/Mixtral.
  • --tokenizer: el tokenizador que se va a cargar. Puede ser un ID de modelo de Hugging Face, una ruta de un contenedor de Cloud Storage (gs://my-bucket/my-model) o una ruta local. Si no se define este argumento, se utilizará el valor de --model de forma predeterminada.
  • --tokenizer_mode: el modo del tokenizador. Las opciones disponibles son ["auto", "slow"]. El valor predeterminado es "auto". Si se define como "auto", se usa el tokenizador rápido si está disponible. Los tokenizadores lentos se han escrito en Python y se proporcionan en la biblioteca Transformers, mientras que los tokenizadores rápidos, que ofrecen una mejora del rendimiento, se han escrito en Rust y se proporcionan en la biblioteca Tokenizers. Para obtener más información, consulta la documentación de Hugging Face.
  • --trust_remote_code: indica si se permiten archivos de código remoto definidos en los repositorios de modelos de Hugging Face. El valor predeterminado es False.
  • --load_format: formato de los puntos de control del modelo que se van a cargar. Las opciones disponibles son ["auto", "dummy"]. El valor predeterminado es "auto". Si se define como "auto", los pesos del modelo se cargan en formato safetensors. Si se define como "dummy", los pesos del modelo se inicializan de forma aleatoria. Definir este valor como "dummy" es útil para hacer experimentos.
  • --max_model_len: longitud máxima del contexto (longitud de entrada más longitud de salida) que se va a usar en el modelo. El valor predeterminado se lee del archivo de configuración del modelo en formato Hugging Face: config.json. Una longitud de contexto máxima mayor requiere más memoria de TPU.
  • --sliding_window: si se define, este argumento anula el tamaño de la ventana del modelo para la atención de ventana deslizante. Si se asigna un valor más alto a este argumento, el mecanismo de atención incluirá más tokens y se acercará al efecto de la autoatención estándar. Este argumento solo se puede usar de forma experimental. En los casos de uso generales, recomendamos usar el tamaño de ventana original del modelo.
  • --seed: la semilla para inicializar todos los generadores de números aleatorios. Cambiar este argumento puede afectar al resultado generado para la misma petición al cambiar los tokens que se muestrean como tokens siguientes. El valor predeterminado es 0.

Motor de inferencia

  • --num_hosts: número de hosts que se van a ejecutar. El valor predeterminado es 1. Para obtener más información, consulta la documentación sobre la configuración de la TPU v5e.
  • --disagg_topo: define el número de trabajadores de relleno automático y de decodificación con la función experimental de publicación desagregada. El valor predeterminado es None. El argumento sigue el formato "number_of_prefill_workers,number_of_decode_workers".
  • --data_parallel_size: número de réplicas paralelas de datos. El valor predeterminado es 1. Si se cambia de 1 a N, el rendimiento mejora aproximadamente en N y la latencia se mantiene igual.
  • --tensor_parallel_size: número de réplicas paralelas de tensores. El valor predeterminado es 1. Aumentar el número de réplicas paralelas de tensores suele mejorar la latencia, ya que acelera la multiplicación de matrices al reducir el tamaño de la matriz.
  • --worker_distributed_method: el método distribuido para iniciar el trabajador. Usa mp para el módulo multiprocessing o ray para la biblioteca Ray. El valor predeterminado es mp.
  • --enable_jit: indica si se debe habilitar el modo JIT (compilación justo a tiempo). El valor predeterminado es True. Si seleccionas --no-enable_jit, se inhabilitará. Si habilitas el modo JIT, el rendimiento de la inferencia mejora, pero se necesita más tiempo para la compilación inicial. Por lo general, las ventajas de rendimiento de la inferencia superan a la sobrecarga.
  • --warmup: indica si se debe calentar el servidor con solicitudes de muestra durante la inicialización. El valor predeterminado es True. Si activas --no-warmup, se inhabilita. Se recomienda usar el calentamiento, ya que las solicitudes iniciales activan una compilación más pesada y, por lo tanto, serán más lentas.
  • --max_prefill_seqs: número máximo de secuencias que se pueden programar para rellenar previamente por iteración. El valor predeterminado es 1. Cuanto mayor sea este valor, mayor será el rendimiento que pueda alcanzar el servidor, pero puede que la latencia se vea afectada.
  • --prefill_seqs_padding: el servidor rellena el tamaño del lote de relleno automático hasta un múltiplo de este valor. El valor predeterminado es 8. Si aumentas este valor, se reduce el tiempo de recompilación del modelo, pero se incrementa el tiempo de computación desperdiciado y la sobrecarga de inferencia. El ajuste óptimo depende del tráfico de solicitudes.
  • --prefill_len_padding: el servidor rellena la longitud de la secuencia hasta un múltiplo de este valor. El valor predeterminado es 512. Si aumentas este valor, se reduce el tiempo de recompilación del modelo, pero se incrementa la sobrecarga de computación e inferencia desperdiciada. El ajuste óptimo depende de la distribución de datos de las solicitudes.
  • --max_decode_seqs/--max_running_seqs: número máximo de secuencias que se pueden programar para decodificar por iteración. El valor predeterminado es 256. Cuanto mayor sea este valor, mayor será el rendimiento que pueda alcanzar el servidor, pero puede tener efectos adversos en la latencia.
  • --decode_seqs_padding: el servidor rellena el tamaño del lote de decodificación hasta un múltiplo de este valor. El valor predeterminado es 8. Si aumentas este valor, se reduce el tiempo de recompilación del modelo, pero se incrementa la sobrecarga de computación e inferencia desperdiciada. El ajuste óptimo depende del tráfico de solicitudes.
  • --decode_blocks_padding: el servidor rellena el número de bloques de memoria utilizados para la caché de clave-valor (caché KV) de una secuencia con un múltiplo de este valor durante la decodificación. El valor predeterminado es 128. Si aumentas este valor, se reduce el tiempo de recompilación del modelo, pero se incrementa la sobrecarga de computación e inferencia desperdiciada. El ajuste óptimo depende de la distribución de datos de las solicitudes.
  • --enable_prefix_cache_hbm: indica si se debe habilitar el almacenamiento en caché de prefijos en HBM. El valor predeterminado es False. Definir este argumento puede mejorar el rendimiento reutilizando los cálculos de los prefijos compartidos de solicitudes anteriores.
  • --enable_chunked_prefill: indica si se debe habilitar el relleno previo fragmentado. El valor predeterminado es False. Definir este argumento puede admitir una longitud de contexto mayor y mejorar el rendimiento.

Gestión de la memoria

  • --hbm_utilization_factor: porcentaje de memoria de alto ancho de banda (HBM) de TPU de Cloud libre que se puede asignar a la caché de KV después de cargar los pesos del modelo. El valor predeterminado es 0.9. Si asignas un valor más alto a este argumento, aumenta el tamaño de la caché de KV y puede mejorar el rendimiento, pero también aumenta el riesgo de que se agote la memoria HBM de la TPU de Cloud durante la inicialización y el tiempo de ejecución.
  • --num_blocks: número de bloques de dispositivos que se asignarán a la caché de KV. Si se define este argumento, el servidor ignora --hbm_utilization_factor. Si no se define este argumento, el servidor perfila el uso de HBM y calcula el número de bloques de dispositivos que se van a asignar en función de --hbm_utilization_factor. Si asignas un valor más alto a este argumento, se aumenta el tamaño de la caché de KV y se puede mejorar el rendimiento, pero también se incrementa el riesgo de que se agote la HBM de Cloud TPU durante la inicialización y el tiempo de ejecución.
  • --block_size: número de tokens almacenados en un bloque. Las opciones disponibles son [8, 16, 32, 2048, 8192]. El valor predeterminado es 32. Si asignas un valor más alto a este argumento, se reduce la sobrecarga en la gestión de bloques, pero se desperdicia más memoria. El impacto exacto en el rendimiento debe determinarse de forma empírica.

LoRA dinámico

  • --enable_lora: indica si se deben habilitar los adaptadores LoRA dinámicos para cargarlos desde Cloud Storage. El valor predeterminado es False. Esta función es compatible con la familia de modelos Llama.
  • --max_lora_rank: el rango máximo de LoRA admitido para los adaptadores de LoRA definidos en las solicitudes. El valor predeterminado es 16. Si asignas un valor más alto a este argumento, podrás usar adaptadores LoRA más flexibles con el servidor, pero aumentará la cantidad de HBM de TPU de Cloud asignada a los pesos de LoRA y se reducirá el rendimiento.
  • --enable_lora_cache: indica si se debe habilitar el almacenamiento en caché de los adaptadores LoRA dinámicos. El valor predeterminado es True. Si seleccionas --no-enable_lora_cache, se inhabilitará. El almacenamiento en caché mejora el rendimiento porque elimina la necesidad de volver a descargar los archivos de adaptador LoRA que se hayan usado anteriormente.
  • --max_num_mem_cached_lora: número máximo de adaptadores LoRA almacenados en la caché de memoria de TPU.El valor predeterminado es 16. Si asignas un valor mayor a este argumento, aumentará la probabilidad de que se produzca un acierto de caché, pero también aumentará la cantidad de uso de HBM de la TPU de Cloud.

También puedes configurar el servidor con las siguientes variables de entorno:

  • HEX_LLM_LOG_LEVEL: controla la cantidad de información de registro que se genera. El valor predeterminado es INFO. Asigna a esta opción uno de los niveles de registro estándar de Python definidos en el módulo de registro.
  • HEX_LLM_VERBOSE_LOG: indica si se debe habilitar la salida de registro detallada. Los valores permitidos son true y false. El valor predeterminado es false.

Ajustar los argumentos del servidor

Los argumentos del servidor están interrelacionados y tienen un efecto colectivo en el rendimiento del servicio. Por ejemplo, un valor más alto de --max_model_len=4096 provoca un mayor uso de la memoria de la TPU y, por lo tanto, requiere una asignación de memoria mayor y menos procesamiento por lotes. Además, algunos argumentos se determinan en función del caso práctico, mientras que otros se pueden ajustar. A continuación, se muestra un flujo de trabajo para configurar el servidor Hex-LLM.

  1. Determina la familia de modelos y la variante del modelo que te interesan. Por ejemplo, Llama 3.1 8B Instruct.
  2. Estima el límite inferior de la memoria de TPU necesaria en función del tamaño y la precisión del modelo: model_size * (num_bits / 8). En el caso de un modelo de 8B y una precisión bfloat16, el límite inferior de la memoria de TPU necesaria sería 8 * (16 / 8) = 16 GB.
  3. Estima el número de chips de la versión 5e de TPU que necesitas, donde cada chip de la versión 5e ofrece 16 GB: tpu_memory / 16. Para un modelo de 8B y una precisión bfloat16, necesitas más de un chip. De las configuraciones de 1, 4 y 8 chips, la más pequeña que ofrece más de un chip es la de 4 chips: ct5lp-hightpu-4t. Puedes configurar --tensor_parallel_size=4 posteriormente.
  4. Determina la longitud máxima del contexto (longitud de entrada + longitud de salida) del caso práctico previsto. Por ejemplo, 4096. Puedes configurar --max_model_len=4096 posteriormente.
  5. Ajusta la cantidad de memoria de TPU libre asignada a la caché de KV al valor máximo que se pueda alcanzar con las configuraciones del modelo, el hardware y el servidor (--hbm_utilization_factor). Empieza con 0.95. Despliega el servidor Hex-LLM y pruébalo con peticiones largas y una alta simultaneidad. Si el servidor se queda sin memoria, reduce el factor de utilización en consecuencia.

A continuación, se muestra un conjunto de argumentos de ejemplo para desplegar Llama 3.1 8B Instruct:

python -m hex_llm.server.api_server \
    --model=meta-llama/Llama-3.1-8B-Instruct \
    --tensor_parallel_size=4 \
    --max_model_len=4096
    --hbm_utilization_factor=0.95

A continuación, se muestra un ejemplo de conjunto de argumentos para desplegar Llama 3.1 70B Instruct AWQ en ct5lp-hightpu-4t:

python -m hex_llm.server.api_server \
    --model=hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4 \
    --tensor_parallel_size=4 \
    --max_model_len=4096
    --hbm_utilization_factor=0.45

Solicitar cuota de TPU de Cloud

En Model Garden, tu cuota predeterminada es de 32 chips de TPU de Cloud v5e en la región us-west1. Estas cuotas se aplican a los despliegues con un solo clic y a los despliegues de cuadernos de Colab Enterprise. Para solicitar un valor de cuota superior, consulta Solicitar un ajuste de cuota.